From a01b112f7fb018131a77929bf8900c9878047c3e Mon Sep 17 00:00:00 2001 From: Anuraag Jain Date: Thu, 28 Dec 2023 22:15:54 +0200 Subject: [PATCH 1/5] feat(auth): add auth middleware - refactored chat routes to use request.user instead of doing authentication in every route --- backend/.gitignore | 3 +- backend/apps/web/main.py | 10 +-- backend/apps/web/middlewares/auth.py | 27 +++++++ backend/apps/web/routers/chats.py | 102 ++++++--------------------- backend/utils/utils.py | 10 +-- 5 files changed, 63 insertions(+), 89 deletions(-) create mode 100644 backend/apps/web/middlewares/auth.py diff --git a/backend/.gitignore b/backend/.gitignore index 11f9256f..da641cf7 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 03273f8b..cb53e9e2 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 00000000..433cdfef --- /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 798972e7..e67776b3 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 62e6958f..97fd1f6f 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: From bdd153d8f5d91487f8c7eca77c340c62b1073626 Mon Sep 17 00:00:00 2001 From: Anuraag Jain Date: Sat, 30 Dec 2023 12:53:33 +0200 Subject: [PATCH 2/5] refac: use dependencies to verify token - feat: added new util to get the current user when needed. Middleware was adding authentication logic to all the routes. let's revisit if we can move the non-auth endpoints to a separate route. - refac: update the routes to use new helpers for verification and retrieving user - chore: added black for local formatting of py code --- backend/apps/ollama/main.py | 10 +- backend/apps/web/main.py | 29 +++-- backend/apps/web/middlewares/auth.py | 27 ---- backend/apps/web/models/users.py | 10 -- backend/apps/web/routers/auths.py | 26 ++-- backend/apps/web/routers/chats.py | 56 +++++---- backend/apps/web/routers/modelfiles.py | 164 +++++++++---------------- backend/apps/web/routers/users.py | 57 +++------ backend/requirements.txt | 2 + backend/utils/utils.py | 37 +++--- 10 files changed, 167 insertions(+), 251 deletions(-) delete mode 100644 backend/apps/web/middlewares/auth.py diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py index 64c6361e..7e138c39 100644 --- a/backend/apps/ollama/main.py +++ b/backend/apps/ollama/main.py @@ -8,7 +8,7 @@ import json from apps.web.models.users import Users from constants import ERROR_MESSAGES -from utils.utils import extract_token_from_auth_header +from utils.utils import decode_token from config import OLLAMA_API_BASE_URL, WEBUI_AUTH app = Flask(__name__) @@ -34,8 +34,12 @@ def proxy(path): # Basic RBAC support if WEBUI_AUTH: if "Authorization" in headers: - token = extract_token_from_auth_header(headers["Authorization"]) - user = Users.get_user_by_token(token) + _, credentials = headers["Authorization"].split() + token_data = decode_token(credentials) + if token_data is None or "email" not in token_data: + return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401 + + user = Users.get_user_by_email(token_data["email"]) if user: # Only user and admin roles can access if user.role in ["user", "admin"]: diff --git a/backend/apps/web/main.py b/backend/apps/web/main.py index cb53e9e2..4519048b 100644 --- a/backend/apps/web/main.py +++ b/backend/apps/web/main.py @@ -1,9 +1,10 @@ -from fastapi import FastAPI +from fastapi import FastAPI, Depends +from fastapi.routing import APIRoute 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 +from utils.utils import verify_auth_token app = FastAPI() @@ -17,14 +18,26 @@ app.add_middleware( allow_headers=["*"], ) - 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( + users.router, + prefix="/users", + tags=["users"], + dependencies=[Depends(verify_auth_token)], +) +app.include_router( + chats.router, + prefix="/chats", + tags=["chats"], + dependencies=[Depends(verify_auth_token)], +) +app.include_router( + modelfiles.router, + prefix="/modelfiles", + tags=["modelfiles"], + dependencies=[Depends(verify_auth_token)], +) 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 deleted file mode 100644 index 433cdfef..00000000 --- a/backend/apps/web/middlewares/auth.py +++ /dev/null @@ -1,27 +0,0 @@ -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/models/users.py b/backend/apps/web/models/users.py index 782b7f47..669ab7dc 100644 --- a/backend/apps/web/models/users.py +++ b/backend/apps/web/models/users.py @@ -3,8 +3,6 @@ from peewee import * from playhouse.shortcuts import model_to_dict from typing import List, Union, Optional import time - -from utils.utils import decode_token from utils.misc import get_gravatar_url from apps.web.internal.db import DB @@ -83,14 +81,6 @@ class UsersTable: except: return None - def get_user_by_token(self, token: str) -> Optional[UserModel]: - data = decode_token(token) - - if data != None and "email" in data: - return self.get_user_by_email(data["email"]) - else: - return None - def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]: return [ UserModel(**model_to_dict(user)) diff --git a/backend/apps/web/routers/auths.py b/backend/apps/web/routers/auths.py index 27d6a3b6..5193afc6 100644 --- a/backend/apps/web/routers/auths.py +++ b/backend/apps/web/routers/auths.py @@ -20,7 +20,7 @@ from apps.web.models.users import Users from utils.utils import ( get_password_hash, - bearer_scheme, + get_current_user, create_token, ) from utils.misc import get_gravatar_url @@ -35,22 +35,14 @@ router = APIRouter() @router.get("/", response_model=UserResponse) -async def get_session_user(cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) - if user: - return { - "id": user.id, - "email": user.email, - "name": user.name, - "role": user.role, - "profile_image_url": user.profile_image_url, - } - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, - ) +async def get_session_user(user=Depends(get_current_user)): + return { + "id": user.id, + "email": user.email, + "name": user.name, + "role": user.role, + "profile_image_url": user.profile_image_url, + } ############################ diff --git a/backend/apps/web/routers/chats.py b/backend/apps/web/routers/chats.py index e67776b3..54b529e6 100644 --- a/backend/apps/web/routers/chats.py +++ b/backend/apps/web/routers/chats.py @@ -1,8 +1,7 @@ - from fastapi import Depends, Request, HTTPException, status from datetime import datetime, timedelta from typing import List, Union, Optional - +from utils.utils import get_current_user from fastapi import APIRouter from pydantic import BaseModel import json @@ -30,8 +29,10 @@ router = APIRouter() @router.get("/", response_model=List[ChatTitleIdResponse]) -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) +async def get_user_chats( + user=Depends(get_current_user), skip: int = 0, limit: int = 50 +): + return Chats.get_chat_lists_by_user_id(user.id, skip, limit) ############################ @@ -40,11 +41,11 @@ async def get_user_chats(request:Request, skip: int = 0, limit: int = 50): @router.get("/all", response_model=List[ChatResponse]) -async def get_all_user_chats(request:Request,): +async def get_all_user_chats(user=Depends(get_current_user)): return [ - ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) - for chat in Chats.get_all_chats_by_user_id(request.user.id) - ] + ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) + for chat in Chats.get_all_chats_by_user_id(user.id) + ] ############################ @@ -53,8 +54,8 @@ async def get_all_user_chats(request:Request,): @router.post("/new", response_model=Optional[ChatResponse]) -async def create_new_chat(form_data: ChatForm,request:Request): - chat = Chats.insert_new_chat(request.user.id, form_data) +async def create_new_chat(form_data: ChatForm, user=Depends(get_current_user)): + chat = Chats.insert_new_chat(user.id, form_data) return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) @@ -64,14 +65,15 @@ async def create_new_chat(form_data: ChatForm,request:Request): @router.get("/{id}", response_model=Optional[ChatResponse]) -async def get_chat_by_id(id: str, request:Request): - chat = Chats.get_chat_by_id_and_user_id(id, request.user.id) +async def get_chat_by_id(id: str, user=Depends(get_current_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) + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND + ) ############################ @@ -80,18 +82,20 @@ async def get_chat_by_id(id: str, request:Request): @router.post("/{id}", response_model=Optional[ChatResponse]) -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) +async def update_chat_by_id( + id: str, form_data: ChatForm, user=Depends(get_current_user) +): + chat = Chats.get_chat_by_id_and_user_id(id, user.id) if chat: - updated_chat = {**json.loads(chat.chat), **form_data.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)}) + chat = Chats.update_chat_by_id(id, updated_chat) + return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)}) else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.ACCESS_PROHIBITED, - ) + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=ERROR_MESSAGES.ACCESS_PROHIBITED, + ) ############################ @@ -100,6 +104,6 @@ async def update_chat_by_id(id: str, form_data: ChatForm, request:Request): @router.delete("/{id}", response_model=bool) -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 +async def delete_chat_by_id(id: str, user=Depends(get_current_user)): + result = Chats.delete_chat_by_id_and_user_id(id, user.id) + return result diff --git a/backend/apps/web/routers/modelfiles.py b/backend/apps/web/routers/modelfiles.py index dd1f6cc5..c54ef4a2 100644 --- a/backend/apps/web/routers/modelfiles.py +++ b/backend/apps/web/routers/modelfiles.py @@ -1,4 +1,3 @@ -from fastapi import Response from fastapi import Depends, FastAPI, HTTPException, status from datetime import datetime, timedelta from typing import List, Union, Optional @@ -16,9 +15,7 @@ from apps.web.models.modelfiles import ( ModelfileResponse, ) -from utils.utils import ( - bearer_scheme, -) +from utils.utils import bearer_scheme, get_current_user from constants import ERROR_MESSAGES router = APIRouter() @@ -30,16 +27,7 @@ router = APIRouter() @router.get("/", response_model=List[ModelfileResponse]) async def get_modelfiles(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) - - if user: - return Modelfiles.get_modelfiles(skip, limit) - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, - ) + return Modelfiles.get_modelfiles(skip, limit) ############################ @@ -48,36 +36,28 @@ async def get_modelfiles(skip: int = 0, limit: int = 50, cred=Depends(bearer_sch @router.post("/create", response_model=Optional[ModelfileResponse]) -async def create_new_modelfile(form_data: ModelfileForm, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) +async def create_new_modelfile( + form_data: ModelfileForm, user=Depends(get_current_user) +): + if user.role != "admin": + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=ERROR_MESSAGES.ACCESS_PROHIBITED, + ) - if user: - # Admin Only - if user.role == "admin": - modelfile = Modelfiles.insert_new_modelfile(user.id, form_data) + modelfile = Modelfiles.insert_new_modelfile(user.id, form_data) - if modelfile: - return ModelfileResponse( - **{ - **modelfile.model_dump(), - "modelfile": json.loads(modelfile.modelfile), - } - ) - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.DEFAULT(), - ) - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.ACCESS_PROHIBITED, - ) + if modelfile: + return ModelfileResponse( + **{ + **modelfile.model_dump(), + "modelfile": json.loads(modelfile.modelfile), + } + ) else: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, + detail=ERROR_MESSAGES.DEFAULT(), ) @@ -87,31 +67,20 @@ async def create_new_modelfile(form_data: ModelfileForm, cred=Depends(bearer_sch @router.post("/", response_model=Optional[ModelfileResponse]) -async def get_modelfile_by_tag_name( - form_data: ModelfileTagNameForm, cred=Depends(bearer_scheme) -): - token = cred.credentials - user = Users.get_user_by_token(token) +async def get_modelfile_by_tag_name(form_data: ModelfileTagNameForm): + modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name) - if user: - modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name) - - if modelfile: - return ModelfileResponse( - **{ - **modelfile.model_dump(), - "modelfile": json.loads(modelfile.modelfile), - } - ) - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.NOT_FOUND, - ) + if modelfile: + return ModelfileResponse( + **{ + **modelfile.model_dump(), + "modelfile": json.loads(modelfile.modelfile), + } + ) else: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, + detail=ERROR_MESSAGES.NOT_FOUND, ) @@ -122,44 +91,34 @@ async def get_modelfile_by_tag_name( @router.post("/update", response_model=Optional[ModelfileResponse]) async def update_modelfile_by_tag_name( - form_data: ModelfileUpdateForm, cred=Depends(bearer_scheme) + form_data: ModelfileUpdateForm, user=Depends(get_current_user) ): - token = cred.credentials - user = Users.get_user_by_token(token) + if user.role != "admin": + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=ERROR_MESSAGES.ACCESS_PROHIBITED, + ) + modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name) + if modelfile: + updated_modelfile = { + **json.loads(modelfile.modelfile), + **form_data.modelfile, + } - if user: - if user.role == "admin": - modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name) - if modelfile: - updated_modelfile = { - **json.loads(modelfile.modelfile), - **form_data.modelfile, - } + modelfile = Modelfiles.update_modelfile_by_tag_name( + form_data.tag_name, updated_modelfile + ) - modelfile = Modelfiles.update_modelfile_by_tag_name( - form_data.tag_name, updated_modelfile - ) - - return ModelfileResponse( - **{ - **modelfile.model_dump(), - "modelfile": json.loads(modelfile.modelfile), - } - ) - 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.ACCESS_PROHIBITED, - ) + return ModelfileResponse( + **{ + **modelfile.model_dump(), + "modelfile": json.loads(modelfile.modelfile), + } + ) else: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, + detail=ERROR_MESSAGES.ACCESS_PROHIBITED, ) @@ -170,22 +129,13 @@ async def update_modelfile_by_tag_name( @router.delete("/delete", response_model=bool) async def delete_modelfile_by_tag_name( - form_data: ModelfileTagNameForm, cred=Depends(bearer_scheme) + form_data: ModelfileTagNameForm, user=Depends(get_current_user) ): - token = cred.credentials - user = Users.get_user_by_token(token) - - if user: - if user.role == "admin": - result = Modelfiles.delete_modelfile_by_tag_name(form_data.tag_name) - return result - else: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.ACCESS_PROHIBITED, - ) - else: + if user.role != "admin": raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, + detail=ERROR_MESSAGES.ACCESS_PROHIBITED, ) + + result = Modelfiles.delete_modelfile_by_tag_name(form_data.tag_name) + return result diff --git a/backend/apps/web/routers/users.py b/backend/apps/web/routers/users.py index 08437bd3..f587dcf2 100644 --- a/backend/apps/web/routers/users.py +++ b/backend/apps/web/routers/users.py @@ -10,11 +10,7 @@ import uuid from apps.web.models.users import UserModel, UserRoleUpdateForm, Users -from utils.utils import ( - get_password_hash, - bearer_scheme, - create_token, -) +from utils.utils import get_current_user from constants import ERROR_MESSAGES router = APIRouter() @@ -25,23 +21,13 @@ router = APIRouter() @router.get("/", response_model=List[UserModel]) -async def get_users(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) - - if user: - if user.role == "admin": - return Users.get_users(skip, limit) - else: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=ERROR_MESSAGES.ACCESS_PROHIBITED, - ) - else: +async def get_users(skip: int = 0, limit: int = 50, user=Depends(get_current_user)): + if user.role != "admin": raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, + status_code=status.HTTP_403_FORBIDDEN, + detail=ERROR_MESSAGES.ACCESS_PROHIBITED, ) + return Users.get_users(skip, limit) ############################ @@ -50,26 +36,19 @@ async def get_users(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme)) @router.post("/update/role", response_model=Optional[UserModel]) -async def update_user_role(form_data: UserRoleUpdateForm, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) +async def update_user_role( + form_data: UserRoleUpdateForm, user=Depends(get_current_user) +): + if user.role != "admin": + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=ERROR_MESSAGES.ACCESS_PROHIBITED, + ) - if user: - if user.role == "admin": - if user.id != form_data.id: - return Users.update_user_role_by_id(form_data.id, form_data.role) - else: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=ERROR_MESSAGES.ACTION_PROHIBITED, - ) - else: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=ERROR_MESSAGES.ACCESS_PROHIBITED, - ) + if user.id != form_data.id: + return Users.update_user_role_by_id(form_data.id, form_data.role) else: raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, + status_code=status.HTTP_403_FORBIDDEN, + detail=ERROR_MESSAGES.ACTION_PROHIBITED, ) diff --git a/backend/requirements.txt b/backend/requirements.txt index 2644d559..6da59fb6 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -18,3 +18,5 @@ bcrypt PyJWT pyjwt[crypto] + +black \ No newline at end of file diff --git a/backend/utils/utils.py b/backend/utils/utils.py index 97fd1f6f..c8991152 100644 --- a/backend/utils/utils.py +++ b/backend/utils/utils.py @@ -1,7 +1,9 @@ -from fastapi.security import HTTPBasicCredentials, HTTPBearer +from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials +from fastapi import HTTPException, status, Depends +from apps.web.models.users import Users from pydantic import BaseModel from typing import Union, Optional - +from constants import ERROR_MESSAGES from passlib.context import CryptContext from datetime import datetime, timedelta import requests @@ -53,16 +55,23 @@ def extract_token_from_auth_header(auth_header: str): return auth_header[len("Bearer ") :] -def verify_token(request): - try: - authorization = request.headers["authorization"] - if authorization: - _, token = authorization.split() - decoded_token = jwt.decode( - token, JWT_SECRET_KEY, options={"verify_signature": False} +def verify_auth_token(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())): + data = decode_token(auth_token.credentials) + if data != None and "email" in data: + user = Users.get_user_by_email(data["email"]) + if user is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=ERROR_MESSAGES.INVALID_TOKEN, ) - return decoded_token - else: - return None - except Exception as e: - return None + return + else: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=ERROR_MESSAGES.UNAUTHORIZED, + ) + + +def get_current_user(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())): + data = decode_token(auth_token.credentials) + return Users.get_user_by_email(data["email"]) From 08c0d7a9ec99dceeee0174da6b34f42a13457c0a Mon Sep 17 00:00:00 2001 From: Anuraag Jain Date: Sat, 30 Dec 2023 13:00:21 +0200 Subject: [PATCH 3/5] fix: merge conflicts --- backend/apps/web/routers/auths.py | 14 ++++++++------ backend/apps/web/routers/users.py | 31 +++++++++++-------------------- 2 files changed, 19 insertions(+), 26 deletions(-) diff --git a/backend/apps/web/routers/auths.py b/backend/apps/web/routers/auths.py index 34dab9b2..7ff645dc 100644 --- a/backend/apps/web/routers/auths.py +++ b/backend/apps/web/routers/auths.py @@ -23,6 +23,7 @@ from utils.utils import ( get_password_hash, get_current_user, create_token, + verify_auth_token, ) from utils.misc import get_gravatar_url from constants import ERROR_MESSAGES @@ -35,7 +36,7 @@ router = APIRouter() ############################ -@router.get("/", response_model=UserResponse) +@router.get("/", response_model=UserResponse, dependencies=[Depends(verify_auth_token)]) async def get_session_user(user=Depends(get_current_user)): return { "id": user.id, @@ -51,11 +52,12 @@ async def get_session_user(user=Depends(get_current_user)): ############################ -@router.post("/update/password", response_model=bool) -async def update_password(form_data: UpdatePasswordForm, cred=Depends(bearer_scheme)): - token = cred.credentials - session_user = Users.get_user_by_token(token) - +@router.post( + "/update/password", response_model=bool, dependencies=[Depends(verify_auth_token)] +) +async def update_password( + form_data: UpdatePasswordForm, session_user=Depends(get_current_user) +): if session_user: user = Auths.authenticate_user(session_user.email, form_data.password) diff --git a/backend/apps/web/routers/users.py b/backend/apps/web/routers/users.py index 2c33c158..950b23fa 100644 --- a/backend/apps/web/routers/users.py +++ b/backend/apps/web/routers/users.py @@ -62,34 +62,25 @@ async def update_user_role( @router.delete("/{user_id}", response_model=bool) -async def delete_user_by_id(user_id: str, cred=Depends(bearer_scheme)): - token = cred.credentials - user = Users.get_user_by_token(token) +async def delete_user_by_id(user_id: str, user=Depends(get_current_user)): + if user.role == "admin": + if user.id != user_id: + result = Auths.delete_auth_by_id(user_id) - if user: - if user.role == "admin": - if user.id != user_id: - result = Auths.delete_auth_by_id(user_id) - - if result: - return True - else: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=ERROR_MESSAGES.DELETE_USER_ERROR, - ) + if result: + return True else: raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=ERROR_MESSAGES.ACTION_PROHIBITED, + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=ERROR_MESSAGES.DELETE_USER_ERROR, ) else: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail=ERROR_MESSAGES.ACCESS_PROHIBITED, + detail=ERROR_MESSAGES.ACTION_PROHIBITED, ) else: raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=ERROR_MESSAGES.INVALID_TOKEN, + status_code=status.HTTP_403_FORBIDDEN, + detail=ERROR_MESSAGES.ACCESS_PROHIBITED, ) From 60c801ee2238555245b2ad64cc902d8f56a4ad8a Mon Sep 17 00:00:00 2001 From: Anuraag Jain Date: Sat, 30 Dec 2023 13:08:03 +0200 Subject: [PATCH 4/5] remove unused import in main --- backend/apps/web/main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/backend/apps/web/main.py b/backend/apps/web/main.py index 4519048b..b19a7c1f 100644 --- a/backend/apps/web/main.py +++ b/backend/apps/web/main.py @@ -1,7 +1,6 @@ from fastapi import FastAPI, Depends from fastapi.routing import APIRoute 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 utils.utils import verify_auth_token From 77323d9b25dd25a0ad9b87a7701e15f45579c3e8 Mon Sep 17 00:00:00 2001 From: Anuraag Jain Date: Mon, 1 Jan 2024 10:55:50 +0200 Subject: [PATCH 5/5] refac: remove the verify_token and use get-current user for auth+user --- backend/apps/web/main.py | 22 +++------------------- backend/apps/web/routers/auths.py | 13 +++---------- backend/apps/web/routers/chats.py | 1 + backend/apps/web/routers/modelfiles.py | 8 +++----- backend/utils/utils.py | 9 ++------- 5 files changed, 12 insertions(+), 41 deletions(-) diff --git a/backend/apps/web/main.py b/backend/apps/web/main.py index b19a7c1f..b3d90381 100644 --- a/backend/apps/web/main.py +++ b/backend/apps/web/main.py @@ -3,7 +3,6 @@ from fastapi.routing import APIRoute from fastapi.middleware.cors import CORSMiddleware from apps.web.routers import auths, users, chats, modelfiles, utils from config import WEBUI_VERSION, WEBUI_AUTH -from utils.utils import verify_auth_token app = FastAPI() @@ -19,24 +18,9 @@ app.add_middleware( app.include_router(auths.router, prefix="/auths", tags=["auths"]) -app.include_router( - users.router, - prefix="/users", - tags=["users"], - dependencies=[Depends(verify_auth_token)], -) -app.include_router( - chats.router, - prefix="/chats", - tags=["chats"], - dependencies=[Depends(verify_auth_token)], -) -app.include_router( - modelfiles.router, - prefix="/modelfiles", - tags=["modelfiles"], - dependencies=[Depends(verify_auth_token)], -) +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/routers/auths.py b/backend/apps/web/routers/auths.py index 7ff645dc..24e9f426 100644 --- a/backend/apps/web/routers/auths.py +++ b/backend/apps/web/routers/auths.py @@ -19,12 +19,7 @@ from apps.web.models.auths import ( from apps.web.models.users import Users -from utils.utils import ( - get_password_hash, - get_current_user, - create_token, - verify_auth_token, -) +from utils.utils import get_password_hash, get_current_user, create_token from utils.misc import get_gravatar_url from constants import ERROR_MESSAGES @@ -36,7 +31,7 @@ router = APIRouter() ############################ -@router.get("/", response_model=UserResponse, dependencies=[Depends(verify_auth_token)]) +@router.get("/", response_model=UserResponse) async def get_session_user(user=Depends(get_current_user)): return { "id": user.id, @@ -52,9 +47,7 @@ async def get_session_user(user=Depends(get_current_user)): ############################ -@router.post( - "/update/password", response_model=bool, dependencies=[Depends(verify_auth_token)] -) +@router.post("/update/password", response_model=bool) async def update_password( form_data: UpdatePasswordForm, session_user=Depends(get_current_user) ): diff --git a/backend/apps/web/routers/chats.py b/backend/apps/web/routers/chats.py index 794eb7f7..0eec4553 100644 --- a/backend/apps/web/routers/chats.py +++ b/backend/apps/web/routers/chats.py @@ -108,6 +108,7 @@ async def delete_chat_by_id(id: str, user=Depends(get_current_user)): result = Chats.delete_chat_by_id_and_user_id(id, user.id) return result + ############################ # DeleteAllChats ############################ diff --git a/backend/apps/web/routers/modelfiles.py b/backend/apps/web/routers/modelfiles.py index c54ef4a2..841d534d 100644 --- a/backend/apps/web/routers/modelfiles.py +++ b/backend/apps/web/routers/modelfiles.py @@ -5,8 +5,6 @@ from typing import List, Union, Optional from fastapi import APIRouter from pydantic import BaseModel import json - -from apps.web.models.users import Users from apps.web.models.modelfiles import ( Modelfiles, ModelfileForm, @@ -15,7 +13,7 @@ from apps.web.models.modelfiles import ( ModelfileResponse, ) -from utils.utils import bearer_scheme, get_current_user +from utils.utils import get_current_user from constants import ERROR_MESSAGES router = APIRouter() @@ -26,7 +24,7 @@ router = APIRouter() @router.get("/", response_model=List[ModelfileResponse]) -async def get_modelfiles(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme)): +async def get_modelfiles(skip: int = 0, limit: int = 50, user=Depends(get_current_user)): return Modelfiles.get_modelfiles(skip, limit) @@ -67,7 +65,7 @@ async def create_new_modelfile( @router.post("/", response_model=Optional[ModelfileResponse]) -async def get_modelfile_by_tag_name(form_data: ModelfileTagNameForm): +async def get_modelfile_by_tag_name(form_data: ModelfileTagNameForm, user=Depends(get_current_user)): modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name) if modelfile: diff --git a/backend/utils/utils.py b/backend/utils/utils.py index c8991152..f98644f5 100644 --- a/backend/utils/utils.py +++ b/backend/utils/utils.py @@ -55,7 +55,7 @@ def extract_token_from_auth_header(auth_header: str): return auth_header[len("Bearer ") :] -def verify_auth_token(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())): +def get_current_user(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())): data = decode_token(auth_token.credentials) if data != None and "email" in data: user = Users.get_user_by_email(data["email"]) @@ -64,14 +64,9 @@ def verify_auth_token(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBea status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.INVALID_TOKEN, ) - return + return user else: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.UNAUTHORIZED, ) - - -def get_current_user(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())): - data = decode_token(auth_token.credentials) - return Users.get_user_by_email(data["email"])