diff --git a/backend/.gitignore b/backend/.gitignore index 11f9256fe..da641cf7d 100644 --- a/backend/.gitignore +++ b/backend/.gitignore @@ -4,4 +4,5 @@ _old uploads .ipynb_checkpoints *.db -_test \ No newline at end of file +_test +Pipfile \ No newline at end of file diff --git a/backend/apps/web/main.py b/backend/apps/web/main.py index 03273f8bc..cb53e9e2f 100644 --- a/backend/apps/web/main.py +++ b/backend/apps/web/main.py @@ -1,8 +1,9 @@ -from fastapi import FastAPI, Request, Depends, HTTPException +from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware - +from starlette.middleware.authentication import AuthenticationMiddleware from apps.web.routers import auths, users, chats, modelfiles, utils from config import WEBUI_VERSION, WEBUI_AUTH +from apps.web.middlewares.auth import BearerTokenAuthBackend, on_auth_error app = FastAPI() @@ -18,11 +19,12 @@ app.add_middleware( app.include_router(auths.router, prefix="/auths", tags=["auths"]) + +app.add_middleware(AuthenticationMiddleware, backend=BearerTokenAuthBackend(), on_error=on_auth_error) + app.include_router(users.router, prefix="/users", tags=["users"]) app.include_router(chats.router, prefix="/chats", tags=["chats"]) app.include_router(modelfiles.router, prefix="/modelfiles", tags=["modelfiles"]) - - app.include_router(utils.router, prefix="/utils", tags=["utils"]) diff --git a/backend/apps/web/middlewares/auth.py b/backend/apps/web/middlewares/auth.py new file mode 100644 index 000000000..433cdfef6 --- /dev/null +++ b/backend/apps/web/middlewares/auth.py @@ -0,0 +1,27 @@ +from apps.web.models.users import Users +from fastapi import Request, status +from starlette.authentication import ( + AuthCredentials, AuthenticationBackend, AuthenticationError, +) +from starlette.requests import HTTPConnection +from utils.utils import verify_token +from starlette.responses import JSONResponse +from constants import ERROR_MESSAGES + +class BearerTokenAuthBackend(AuthenticationBackend): + + async def authenticate(self, conn: HTTPConnection): + if "Authorization" not in conn.headers: + return + data = verify_token(conn) + if data != None and 'email' in data: + user = Users.get_user_by_email(data['email']) + if user is None: + raise AuthenticationError('Invalid credentials') + return AuthCredentials([user.role]), user + else: + raise AuthenticationError('Invalid credentials') + +def on_auth_error(request: Request, exc: Exception): + print('Authentication failed: ', exc) + return JSONResponse({"detail": ERROR_MESSAGES.INVALID_TOKEN}, status_code=status.HTTP_401_UNAUTHORIZED) \ No newline at end of file diff --git a/backend/apps/web/routers/chats.py b/backend/apps/web/routers/chats.py index 798972e7c..e67776b3c 100644 --- a/backend/apps/web/routers/chats.py +++ b/backend/apps/web/routers/chats.py @@ -1,5 +1,5 @@ -from fastapi import Response -from fastapi import Depends, FastAPI, HTTPException, status + +from fastapi import Depends, Request, HTTPException, status from datetime import datetime, timedelta from typing import List, Union, Optional @@ -30,17 +30,8 @@ router = APIRouter() @router.get("/", response_model=List[ChatTitleIdResponse]) -async def get_user_chats(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) - - if user: - return Chats.get_chat_lists_by_user_id(user.id, skip, limit) - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, - ) +async def get_user_chats(request:Request, skip: int = 0, limit: int = 50): + return Chats.get_chat_lists_by_user_id(request.user.id, skip, limit) ############################ @@ -49,20 +40,11 @@ async def get_user_chats(skip: int = 0, limit: int = 50, cred=Depends(bearer_sch @router.get("/all", response_model=List[ChatResponse]) -async def get_all_user_chats(cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) - - if user: - return [ +async def get_all_user_chats(request:Request,): + return [ ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) - for chat in Chats.get_all_chats_by_user_id(user.id) + for chat in Chats.get_all_chats_by_user_id(request.user.id) ] - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, - ) ############################ @@ -71,18 +53,9 @@ async def get_all_user_chats(cred=Depends(bearer_scheme)): @router.post("/new", response_model=Optional[ChatResponse]) -async def create_new_chat(form_data: ChatForm, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) - - if user: - chat = Chats.insert_new_chat(user.id, form_data) - return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, - ) +async def create_new_chat(form_data: ChatForm,request:Request): + chat = Chats.insert_new_chat(request.user.id, form_data) + return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) ############################ @@ -91,25 +64,14 @@ async def create_new_chat(form_data: ChatForm, cred=Depends(bearer_scheme)): @router.get("/{id}", response_model=Optional[ChatResponse]) -async def get_chat_by_id(id: str, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) +async def get_chat_by_id(id: str, request:Request): + chat = Chats.get_chat_by_id_and_user_id(id, request.user.id) - if user: - chat = Chats.get_chat_by_id_and_user_id(id, user.id) - - if chat: - return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.NOT_FOUND, - ) + if chat: + return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, - ) + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, + detail=ERROR_MESSAGES.NOT_FOUND) ############################ @@ -118,27 +80,18 @@ async def get_chat_by_id(id: str, cred=Depends(bearer_scheme)): @router.post("/{id}", response_model=Optional[ChatResponse]) -async def update_chat_by_id(id: str, form_data: ChatForm, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) - - if user: - chat = Chats.get_chat_by_id_and_user_id(id, user.id) - if chat: +async def update_chat_by_id(id: str, form_data: ChatForm, request:Request): + chat = Chats.get_chat_by_id_and_user_id(id, request.user.id) + if chat: updated_chat = {**json.loads(chat.chat), **form_data.chat} chat = Chats.update_chat_by_id(id, updated_chat) return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) - else: + else: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.ACCESS_PROHIBITED, ) - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, - ) ############################ @@ -147,15 +100,6 @@ async def update_chat_by_id(id: str, form_data: ChatForm, cred=Depends(bearer_sc @router.delete("/{id}", response_model=bool) -async def delete_chat_by_id(id: str, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) - - if user: - result = Chats.delete_chat_by_id_and_user_id(id, user.id) - return result - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, - ) +async def delete_chat_by_id(id: str, request: Request): + result = Chats.delete_chat_by_id_and_user_id(id, request.user.id) + return result \ No newline at end of file diff --git a/backend/utils/utils.py b/backend/utils/utils.py index 62e6958fb..97fd1f6f8 100644 --- a/backend/utils/utils.py +++ b/backend/utils/utils.py @@ -55,13 +55,13 @@ def extract_token_from_auth_header(auth_header: str): def verify_token(request): try: - bearer = request.headers["authorization"] - if bearer: - token = bearer[len("Bearer ") :] - decoded = jwt.decode( + authorization = request.headers["authorization"] + if authorization: + _, token = authorization.split() + decoded_token = jwt.decode( token, JWT_SECRET_KEY, options={"verify_signature": False} ) - return decoded + return decoded_token else: return None except Exception as e: