add auth

trauma/__init__.py

@@ -12,12 +12,17 @@ from trauma.core.wrappers import TraumaResponseWrapper, ErrorTraumaResponse
 def create_app() -> FastAPI:
     app = FastAPI()
 
+    from trauma.api.account import account_router
+    app.include_router(account_router, tags=["account"])
+
     from trauma.api.chat import chat_router
     app.include_router(chat_router, tags=['chat'])
 
     from trauma.api.message import message_router
     app.include_router(message_router, tags=['message'])
 
+    from trauma.api.security import security_router
+    app.include_router(security_router, tags=['security'])
 
     app.add_middleware(
         CORSMiddleware,

trauma/api/account/__init__.py

@@ -0,0 +1,7 @@
+from fastapi import APIRouter
+
+account_router = APIRouter(
+    prefix='/api/account'
+)
+
+from . import views

trauma/api/account/db_requests.py

@@ -0,0 +1,17 @@
+import asyncio
+
+from trauma.api.account.model import AccountModel
+from trauma.core.config import settings
+
+
+async def get_all_model_obj(page_size: int, page_index: int) -> tuple[list[AccountModel], int]:
+    skip = page_size * page_index
+    objects, total_count = await asyncio.gather(
+        settings.DB_CLIENT.accounts
+        .find()
+        .skip(skip)
+        .limit(page_size)
+        .to_list(length=page_size),
+        settings.DB_CLIENT.accounts.count_documents({})
+    )
+    return objects, total_count

trauma/api/account/dto.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+from pydantic import BaseModel
+
+
+class AccessToken(BaseModel):
+    type: str = "Bearer"
+    value: str
+
+
+class AccountType(Enum):
+    User = 1
+    Admin = 2

trauma/api/account/model.py

@@ -0,0 +1,34 @@
+from datetime import datetime
+
+from passlib.context import CryptContext
+from pydantic import field_validator, Field, EmailStr
+
+from trauma.api.account.dto import AccountType
+from trauma.core.database import MongoBaseModel
+
+
+class AccountModel(MongoBaseModel):
+    email: EmailStr
+    password: str | None = Field(exclude=True, default=None)
+    accountType: AccountType = AccountType.User
+
+    datetimeInserted: datetime = Field(default_factory=datetime.now)
+    datetimeUpdated: datetime = Field(default_factory=datetime.now)
+
+    @field_validator("datetimeUpdated", mode="before", check_fields=False)
+    @classmethod
+    def validate_datetimeUpdated(cls, v):
+        return v or datetime.now()
+
+    @field_validator('password', mode='before', check_fields=False)
+    @classmethod
+    def set_password_hash(cls, v):
+        if not v.startswith("$2b$"):
+            return CryptContext(schemes=["bcrypt"], deprecated="auto").hash(v)
+        return v
+
+    class Config:
+        arbitrary_types_allowed = True
+        json_encoders = {
+            datetime: lambda dt: dt.isoformat()
+        }

trauma/api/account/schemas.py

@@ -0,0 +1,14 @@
+from pydantic import BaseModel
+
+from trauma.api.account.model import AccountModel
+from trauma.api.common.dto import Paging
+from trauma.core.wrappers import TraumaResponseWrapper
+
+
+class AccountWrapper(TraumaResponseWrapper[AccountModel]):
+    pass
+
+
+class AllAccountsResponse(BaseModel):
+    paging: Paging
+    data: list[AccountModel]

trauma/api/account/views.py

@@ -0,0 +1,33 @@
+from typing import Optional
+
+from fastapi import Depends, Query
+
+from trauma.api.account import account_router
+from trauma.api.account.db_requests import get_all_model_obj
+from trauma.api.account.dto import AccountType
+from trauma.api.account.model import AccountModel
+from trauma.api.account.schemas import AccountWrapper, AllAccountsResponse
+from trauma.api.common.dto import Paging
+from trauma.core.security import PermissionDependency
+from trauma.core.wrappers import TraumaResponseWrapper
+
+
+@account_router.get('/all')
+async def get_all_accounts(
+        pageSize: Optional[int] = Query(10, description="Number of countries to return per page"),
+        pageIndex: Optional[int] = Query(0, description="Page index to retrieve"),
+        _: AccountModel = Depends(PermissionDependency([AccountType.Admin]))
+) -> TraumaResponseWrapper[AllAccountsResponse]:
+    countries, total_count = await get_all_model_obj(pageSize, pageIndex)
+    response = AllAccountsResponse(
+        paging=Paging(pageSize=pageSize, pageIndex=pageIndex, totalCount=total_count),
+        data=countries
+    )
+    return TraumaResponseWrapper(data=response)
+
+
+@account_router.get('')
+async def get_account(
+        account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
+) -> AccountWrapper:
+    return AccountWrapper(data=account)

trauma/api/chat/db_requests.py

@@ -2,6 +2,8 @@ import asyncio
 
 from fastapi import HTTPException
 
+from trauma.api.account.dto import AccountType
+from trauma.api.account.model import AccountModel
 from trauma.api.chat.model import ChatModel
 from trauma.api.chat.schemas import CreateChatRequest, ChatTitleRequest
 from trauma.api.message.dto import Author
@@ -9,37 +11,43 @@ from trauma.api.message.model import MessageModel
 from trauma.core.config import settings
 
 
-async def get_chat_obj(chat_id: str) -> ChatModel:
+async def get_chat_obj(chat_id: str, account: AccountModel) -> ChatModel:
     chat = await settings.DB_CLIENT.chats.find_one({"id": chat_id})
     if not chat:
         raise HTTPException(status_code=404, detail="Chat not found")
     chat = ChatModel.from_mongo(chat)
+    if chat.account.id != account.id and account.accountType != AccountType.Admin:
+        raise HTTPException(status_code=403, detail="Not authorized")
     return chat
 
 
-async def create_chat_obj(chat_request: CreateChatRequest) -> ChatModel:
-    chat = ChatModel(model=chat_request.model)
+async def create_chat_obj(chat_request: CreateChatRequest, account: AccountModel) -> ChatModel:
+    chat = ChatModel(model=chat_request.model, account=account)
     await settings.DB_CLIENT.chats.insert_one(chat.to_mongo())
     return chat
 
 
-async def delete_chat_obj(chat_id: str) -> None:
+async def delete_chat_obj(chat_id: str, account: AccountModel) -> None:
     chat = await settings.DB_CLIENT.chats.find_one({"id": chat_id})
     if not chat:
         raise HTTPException(status_code=404, detail="Chat not found")
     chat = ChatModel.from_mongo(chat)
+    if chat.account.id != account.id and account.accountType != AccountType.Admin:
+        raise HTTPException(status_code=403, detail="Not authorized")
     await settings.DB_CLIENT.chats.delete_one({"id": chat_id})
 
 
-async def update_chat_obj_title(chatId: str, chat_request: ChatTitleRequest) -> ChatModel:
-    chat = await settings.DB_CLIENT.chats.find_one({"id": chatId})
+async def update_chat_obj_title(chat_id: str, chat_request: ChatTitleRequest, account: AccountModel) -> ChatModel:
+    chat = await settings.DB_CLIENT.chats.find_one({"id": chat_id})
     if not chat:
         raise HTTPException(status_code=404, detail="Chat not found")
 
     chat = ChatModel.from_mongo(chat)
+    if chat.account.id != account.id and account.accountType != AccountType.Admin:
+        raise HTTPException(status_code=403, detail="Not authorized")
 
     chat.title = chat_request.title
-    await settings.DB_CLIENT.chats.update_one({"id": chatId}, {"$set": chat.to_mongo()})
+    await settings.DB_CLIENT.chats.update_one({"id": chat_id}, {"$set": chat.to_mongo()})
     return chat
 
 

trauma/api/chat/model.py

@@ -2,6 +2,7 @@ from datetime import datetime
 
 from pydantic import Field
 
+from trauma.api.account.model import AccountModel
 from trauma.api.chat.dto import ModelType, EntityData
 from trauma.core.database import MongoBaseModel
 
@@ -10,5 +11,6 @@ class ChatModel(MongoBaseModel):
     title: str = 'New Chat'
     model: ModelType
     entityData: EntityData = EntityData()
+    account: AccountModel
     datetimeInserted: datetime = Field(default_factory=datetime.now)
     datetimeUpdated: datetime = Field(default_factory=datetime.now)

trauma/api/chat/views.py

@@ -1,7 +1,9 @@
 from typing import Optional
 
-from fastapi import Query
+from fastapi import Query, Depends
 
+from trauma.api.account.dto import AccountType
+from trauma.api.account.model import AccountModel
 from trauma.api.chat import chat_router
 from trauma.api.chat.db_requests import (get_chat_obj,
                                          create_chat_obj,
@@ -10,6 +12,7 @@ from trauma.api.chat.db_requests import (get_chat_obj,
                                          get_all_chats_obj, save_intro_message)
 from trauma.api.chat.schemas import ChatWrapper, AllChatWrapper, CreateChatRequest, AllChatResponse, ChatTitleRequest
 from trauma.api.common.dto import Paging
+from trauma.core.security import PermissionDependency
 from trauma.core.wrappers import TraumaResponseWrapper
 
 
@@ -17,6 +20,7 @@ from trauma.core.wrappers import TraumaResponseWrapper
 async def get_all_chats(
         pageSize: Optional[int] = Query(10, description="Number of objects to return per page"),
         pageIndex: Optional[int] = Query(0, description="Page index to retrieve"),
+        _: AccountModel = Depends(PermissionDependency([AccountType.Admin]))
 ) -> AllChatWrapper:
     chats, total_count = await get_all_chats_obj(pageSize, pageIndex)
     response = AllChatResponse(
@@ -28,30 +32,36 @@ async def get_all_chats(
 
 @chat_router.get('/{chatId}')
 async def get_chat(
-        chatId: str
+        chatId: str,
+        account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
 ) -> ChatWrapper:
-    chat = await get_chat_obj(chatId)
+    chat = await get_chat_obj(chatId, account)
     return ChatWrapper(data=chat)
 
 
 @chat_router.delete('/{chatId}')
-async def delete_chat(chatId: str) -> TraumaResponseWrapper:
-    await delete_chat_obj(chatId)
+async def delete_chat(
+        chatId: str,
+        account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
+) -> TraumaResponseWrapper:
+    await delete_chat_obj(chatId, account)
     return TraumaResponseWrapper()
 
 
 @chat_router.patch('/{chatId}/title')
 async def update_chat_title(
-        chatId: str, chat: ChatTitleRequest
+        chatId: str, chat: ChatTitleRequest,
+        account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
 ) -> ChatWrapper:
-    chat = await update_chat_obj_title(chatId, chat)
+    chat = await update_chat_obj_title(chatId, chat, account)
     return ChatWrapper(data=chat)
 
 
 @chat_router.post('')
 async def create_chat(
-        chat_data: CreateChatRequest
+        chat_data: CreateChatRequest,
+        account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
 ) -> ChatWrapper:
-    chat = await create_chat_obj(chat_data)
+    chat = await create_chat_obj(chat_data, account)
     await save_intro_message(chat.id)
     return ChatWrapper(data=chat)

trauma/api/common/db_requests.py

@@ -0,0 +1,17 @@
+from fastapi import HTTPException
+from pydantic import EmailStr
+
+from trauma.core.config import settings
+
+
+async def check_unique_fields_existence(name: str,
+                                        new_value: EmailStr | str,
+                                        current_value: str | None = None) -> None:
+    if new_value == current_value or not new_value:
+        return
+    account = await settings.DB_CLIENT.accounts.find_one(
+        {name: str(new_value)},
+        collation={"locale": "en", "strength": 2}
+    )
+    if account:
+        raise HTTPException(status_code=400, detail=f'Account with specified {name} already exists.')

trauma/api/message/db_requests.py

@@ -2,6 +2,8 @@ import asyncio
 
 from fastapi import HTTPException
 
+from trauma.api.account.dto import AccountType
+from trauma.api.account.model import AccountModel
 from trauma.api.chat.model import ChatModel
 from trauma.api.data.model import EntityModel
 from trauma.api.message.dto import Author
@@ -24,7 +26,7 @@ async def create_message_obj(
     return message, chat
 
 
-async def get_all_chat_messages_obj(chat_id: str) -> tuple[list[MessageModel], ChatModel]:
+async def get_all_chat_messages_obj(chat_id: str, account: AccountModel) -> tuple[list[MessageModel], ChatModel]:
     messages, chat = await asyncio.gather(
         settings.DB_CLIENT.messages.find({"chatId": chat_id}).to_list(length=None),
         settings.DB_CLIENT.chats.find_one({"id": chat_id})
@@ -35,6 +37,8 @@ async def get_all_chat_messages_obj(chat_id: str) -> tuple[list[MessageModel], C
         raise HTTPException(status_code=404, detail="Chat not found")
 
     chat = ChatModel.from_mongo(chat)
+    if chat.account.id != account.id and account.accountType != AccountType.Admin:
+        raise HTTPException(status_code=404, detail="Not Authorized.")
     return messages, chat
 
 

trauma/api/message/views.py

@@ -1,3 +1,7 @@
+from fastapi import Depends
+
+from trauma.api.account.dto import AccountType
+from trauma.api.account.model import AccountModel
 from trauma.api.common.dto import Paging
 from trauma.api.message import message_router
 from trauma.api.message.ai.engine import search_entities
@@ -7,14 +11,16 @@ from trauma.api.message.schemas import (AllMessageWrapper,
                                         CreateMessageRequest,
                                         CreateMessageResponse)
 from trauma.api.message.utils import transform_messages_to_openai
+from trauma.core.security import PermissionDependency
 from trauma.core.wrappers import TraumaResponseWrapper
 
 
 @message_router.get('/{chatId}/all')
 async def get_all_chat_messages(
-        chatId: str
+        chatId: str,
+        account: AccountModel = Depends(PermissionDependency([AccountType.Admin]))
 ) -> AllMessageWrapper:
-    messages, _ = await get_all_chat_messages_obj(chatId)
+    messages, _ = await get_all_chat_messages_obj(chatId, account)
     response = AllMessageResponse(
         paging=Paging(pageSize=len(messages), pageIndex=0, totalCount=len(messages)),
         data=messages
@@ -26,8 +32,9 @@ async def get_all_chat_messages(
 async def create_message(
         chatId: str,
         message_data: CreateMessageRequest,
+        account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
 ) -> TraumaResponseWrapper[CreateMessageResponse]:
-    messages, chat = await get_all_chat_messages_obj(chatId)
+    messages, chat = await get_all_chat_messages_obj(chatId, account)
     message_history = transform_messages_to_openai(messages)
     response = await search_entities(message_data.text, message_history, chat)
     return TraumaResponseWrapper(data=response)

trauma/api/security/__init__.py

@@ -0,0 +1,7 @@
+from fastapi import APIRouter
+
+security_router = APIRouter(
+    prefix='/api/security'
+)
+
+from . import views

trauma/api/security/db_requests.py

@@ -0,0 +1,32 @@
+import asyncio
+
+from fastapi import HTTPException
+
+from trauma.api.account.model import AccountModel
+from trauma.api.common.db_requests import check_unique_fields_existence
+from trauma.api.security.schemas import RegisterAccountRequest, LoginAccountRequest
+from trauma.core.config import settings
+from trauma.core.security import verify_password
+
+
+async def save_account(data: RegisterAccountRequest) -> AccountModel:
+    await check_unique_fields_existence("email", data.email)
+    account = AccountModel(
+        **data.model_dump()
+    )
+    await settings.DB_CLIENT.accounts.insert_one(account.to_mongo())
+    return account
+
+
+async def authenticate_account(data: LoginAccountRequest) -> AccountModel:
+    account = await settings.DB_CLIENT.accounts.find_one(
+        {"email": data.email},
+        collation={"locale": "en", "strength": 2})
+    if account is None:
+        raise HTTPException(status_code=404, detail="Invalid email or password.")
+
+    account = AccountModel.from_mongo(account)
+
+    if not verify_password(data.password, account.password):
+        raise HTTPException(status_code=401, detail="Invalid email or password.")
+    return account

trauma/api/security/schemas.py

@@ -0,0 +1,28 @@
+from pydantic import BaseModel, EmailStr
+
+from trauma.api.account.dto import AccessToken
+from trauma.api.account.model import AccountModel
+from trauma.core.wrappers import TraumaResponseWrapper
+
+
+class RegisterAccountRequest(BaseModel):
+    email: EmailStr
+    password: str
+
+
+class RegisterAccountWrapper(TraumaResponseWrapper[AccountModel]):
+    pass
+
+
+class LoginAccountRequest(BaseModel):
+    email: EmailStr
+    password: str
+
+
+class LoginAccountResponse(BaseModel):
+    accessToken: AccessToken
+    account: AccountModel
+
+
+class LoginAccountWrapper(TraumaResponseWrapper[LoginAccountResponse]):
+    pass

trauma/api/security/views.py

@@ -0,0 +1,26 @@
+from trauma.api.account.dto import AccessToken
+from trauma.api.security import security_router
+from trauma.api.security.db_requests import authenticate_account, save_account
+from trauma.api.security.schemas import (RegisterAccountRequest,
+                                         RegisterAccountWrapper,
+                                         LoginAccountResponse,
+                                         LoginAccountWrapper,
+                                         LoginAccountRequest)
+from trauma.core.security import create_access_token
+
+
+@security_router.post('/register')
+async def register_user(data: RegisterAccountRequest) -> RegisterAccountWrapper:
+    account = await save_account(data)
+    return RegisterAccountWrapper(data=account)
+
+
+@security_router.post('/login')
+async def login(data: LoginAccountRequest) -> LoginAccountWrapper:
+    account = await authenticate_account(data)
+    access_token = create_access_token(account.email, str(account.id))
+    response = LoginAccountResponse(
+        accessToken=AccessToken(value=access_token),
+        account=account,
+    )
+    return LoginAccountWrapper(data=response)

trauma/core/security.py

@@ -6,6 +6,7 @@ from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 from jose import jwt, JWTError
 from passlib.context import CryptContext
 
+from trauma.api.account.dto import AccountType
 from trauma.api.account.model import AccountModel
 from trauma.core.config import settings
 
@@ -29,36 +30,34 @@ def create_access_token(email: str, account_id: str):
 
 
 class PermissionDependency:
-    def __init__(self, is_public: bool = False):
-        self.is_public = is_public
+    def __init__(self, account_types: list[AccountType]):
+        self.account_types = account_types
 
     def __call__(
             self, credentials: HTTPAuthorizationCredentials | None = Depends(HTTPBearer(auto_error=False))
     ) -> AccountModel | None:
-        if credentials is None:
-            if self.is_public:
-                return None
-            else:
-                raise HTTPException(status_code=403, detail="Permission denied")
         try:
             account_id = self.authenticate_jwt_token(credentials.credentials)
             account_data = anyio.from_thread.run(self.get_account_by_id, account_id)
             self.check_account_health(account_data)
-            return account_data
+            return AccountModel.from_mongo(account_data)
 
         except JWTError:
             raise HTTPException(status_code=403, detail="Permission denied")
 
-    async def get_account_by_id(self, account_id: str) -> AccountModel:
+    @staticmethod
+    async def get_account_by_id(account_id: str) -> dict:
         account = await settings.DB_CLIENT.accounts.find_one({"id": account_id})
-        return AccountModel.from_mongo(account)
+        return account
 
-    @staticmethod
-    def check_account_health(account: AccountModel):
+    def check_account_health(self, account: dict):
         if not account:
             raise HTTPException(status_code=403, detail="Permission denied")
+        if account['accountType'] not in [i.value for i in self.account_types]:
+            raise HTTPException(status_code=403, detail="Permission denied")
 
-    def authenticate_jwt_token(self, token: str) -> str:
+    @staticmethod
+    def authenticate_jwt_token(token: str) -> str:
         payload = jwt.decode(token,
                              settings.SECRET_KEY,
                              algorithms="HS256",