Update private_gpt/server/utils/authentication.py

private_gpt/server/utils/authentication.py

@@ -6,6 +6,7 @@ from passlib.context import CryptContext
 from fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer
 from jose import jwt, JWTError
 from pydantic import BaseModel
+import uuid
 import os 
 router = APIRouter(
     prefix='/v1/auth',
@@ -19,14 +20,18 @@ bcrypt_context = CryptContext(schemes=['bcrypt'], deprecated='auto')
 oauth2_bearer = OAuth2PasswordBearer(tokenUrl='/v1/auth/token')
 
 hardcoded_users = [
-    {   "id": 1, 
+    {   "id": str(uuid.uuid4()), 
+        "first_name": "Ibraaheem",
+        "last_name": "Akbar",
         "username": "test",
         "password_hash": bcrypt_context.hash(os.environ.get("USER_HASH")),
         "role": "user"
     },
 
-    {   "id": 2, 
+    {   "id": str(uuid.uuid4()), 
         "username": "admin",
+        "first_name": "John",
+        "last_name": "Doe",
         "password_hash": bcrypt_context.hash(os.environ.get("ADMIN_HASH")),
         "role": "admin"
     },
@@ -35,9 +40,12 @@ hardcoded_users = [
 ]
 
 
+
 class CreateUserRequest(BaseModel):
     username: str
     password: str
+    first_name: str
+    last_name: str
 
 class Token(BaseModel):
     access_token: str
@@ -53,15 +61,33 @@ def authenticate_user(username: str, password: str, role: str):
                 and stored_role
                 and bcrypt_context.verify(password, stored_password_hash)
             ):
-                return {"username": username, "id": user["id"], "role": stored_role}
+                # Include 'first_name' and 'last_name' in the user dictionary
+                user_data = {
+                    "username": username,
+                    "id": user["id"],
+                    "role": stored_role,
+                    "first_name": user.get("first_name", ""),
+                    "last_name": user.get("last_name", ""),
+                }
+                return user_data
     return None
 
 
 
-@router.post("/", status_code=status.HTTP_201_CREATED)
+
+@router.post("/create_user", 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
+    user_id = str(uuid.uuid4())
+    user_data = {
+        "id": user_id,
+        "first_name": create_user_request.first_name,
+        "last_name": create_user_request.last_name,
+        "username": create_user_request.username,
+        "password_hash": bcrypt_context.hash(create_user_request.password),
+        "role": "user"
+    }
+    hardcoded_users.append(user_data)
+    return {"message": "User created successfully"}
 
 @router.post("/token", response_model=Token)
 async def login_for_access_token(
@@ -74,14 +100,14 @@ async def login_for_access_token(
             detail="Could not validate user.",
         )
 
-    token = create_access_token(user["username"], user["id"], user["role"], timedelta(minutes=10080))
+    token = create_access_token(user["username"], user["id"], user["role"], user["first_name"], user["last_name"], 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}  
+def create_access_token(username: str, user_id: int, role: str, first_name: str, last_name: str, expires_delta: timedelta):
+    encode = {'sub': username, 'id': user_id, 'role': role, 'first_name': first_name, 'last_name': last_name}
     expires = datetime.utcnow() + expires_delta
     encode.update({'exp': expires})
     return jwt.encode(encode, SECRET_KEY, algorithm=ALGORITHM)
@@ -93,9 +119,11 @@ async def get_current_user(token: Annotated[str, Depends(oauth2_bearer)]):
         username: str = payload.get('sub')
         user_id: int = payload.get('id')
         role: str = payload.get('role')  # Add this line to get the role
+        first_name: str = payload.get('first_name')
+        last_name: str = payload.get('last_name')
         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
+        return {'username': username, 'id': user_id, 'role': role, 'first_name': first_name, 'last_name': last_name}
     except JWTError:
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate user.")