Create authentication.py

private_gpt/server/utils/authentication.py

@@ -0,0 +1,101 @@
+from datetime import timedelta, datetime
+from typing import Annotated
+from fastapi import APIRouter, Depends, HTTPException
+from starlette import status
+from passlib.context import CryptContext
+from fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer
+from jose import jwt, JWTError
+from pydantic import BaseModel
+
+router = APIRouter(
+    prefix='/v1/auth',
+    tags=['auth']
+)
+
+SECRET_KEY = '1b510738bd0cd2d5757a27935aa4355b04efcd268a180fafd5fb7d8de60cc73a'
+ALGORITHM = 'HS256'
+
+bcrypt_context = CryptContext(schemes=['bcrypt'], deprecated='auto')
+oauth2_bearer = OAuth2PasswordBearer(tokenUrl='/v1/auth/token')
+
+hardcoded_users = [
+    {   "id": 1, 
+        "username": "test",
+        "password_hash": bcrypt_context.hash("secret"),
+        "role": "user"
+    },
+
+    {   "id": 2, 
+        "username": "admin",
+        "password_hash": bcrypt_context.hash("admin"),
+        "role": "admin"
+    },
+  
+  
+]
+
+
+class CreateUserRequest(BaseModel):
+    username: str
+    password: str
+
+class Token(BaseModel):
+    access_token: str
+    token_type: str
+
+def authenticate_user(username: str, password: str, role: str):
+    for user in hardcoded_users:
+        if user["username"] == username:
+            stored_password_hash = user.get("password_hash")
+            stored_role = user.get("role")
+            if (
+                stored_password_hash
+                and stored_role
+                and bcrypt_context.verify(password, stored_password_hash)
+            ):
+                return {"username": username, "id": user["id"], "role": stored_role}
+    return None
+
+
+
+@router.post("/", status_code=status.HTTP_201_CREATED)
+async def create_user(create_user_request: CreateUserRequest):
+    # This function is not necessary for hardcoded users, as users are predefined
+    pass
+
+@router.post("/token", response_model=Token)
+async def login_for_access_token(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()]
+):
+    user = authenticate_user(form_data.username, form_data.password, role="user")
+    if not user:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Could not validate user.",
+        )
+
+    token = create_access_token(user["username"], user["id"], user["role"], timedelta(minutes=10080))
+
+    return Token(access_token=token, token_type="bearer")
+
+
+
+def create_access_token(username: str, user_id: int, role: str, expires_delta: timedelta):
+    encode = {'sub': username, 'id': user_id, 'role': role}  
+    expires = datetime.utcnow() + expires_delta
+    encode.update({'exp': expires})
+    return jwt.encode(encode, SECRET_KEY, algorithm=ALGORITHM)
+
+
+async def get_current_user(token: Annotated[str, Depends(oauth2_bearer)]):
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=ALGORITHM)
+        username: str = payload.get('sub')
+        user_id: int = payload.get('id')
+        role: str = payload.get('role')  # Add this line to get the role
+        if username is None or user_id is None:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail='Could not validate user.')
+        return {'username': username, 'id': user_id, 'role': role}  # Include the role in the returned dictionary
+    except JWTError:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate user.")
+