Cleanup token auth

app/api/userlogout.py

@@ -1,14 +1,21 @@
+from typing import Any
 from fastapi import APIRouter, Depends, HTTPException
 from ..utils.db import tinydb_helper  # Ensure this import is correct based on our project structure
-from ..dependencies import oauth2_scheme
+from ..dependencies import get_current_user, oauth2_scheme
 
 router = APIRouter()
 
 @router.post("/user/logout")
-async def user_logout(token: str = Depends(oauth2_scheme)):
+async def user_logout(token: str = Depends(oauth2_scheme), current_user: Any = Depends(get_current_user)):
     try:
+        # Assuming `get_current_user` now also ensures and returns the full payload including `user_id`
+        user_id = current_user["user_id"]
         # Invalidate the token by removing it from the database
+        if not tinydb_helper.query_token(user_id, token):
+            raise HTTPException(status_code=404, detail="Token not found.")
         tinydb_helper.remove_token_by_value(token)
+        if tinydb_helper.query_token(user_id, token):
+            raise HTTPException(status_code=404, detail="Logout unsuccessful.")
         return {"message": "User logged out successfully"}
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Error during logout: {str(e)}")

app/dependencies.py

@@ -1,28 +1,33 @@
 from fastapi import Depends, HTTPException, status
 from fastapi.security import OAuth2PasswordBearer
-from jose import jwt, JWTError
-from .utils.db import tinydb_helper  # Ensure correct import path
-from .utils.jwt_utils import SECRET_KEY, ALGORITHM  # Ensure these are defined in our jwt_utils.py
+from jose import jwt, JWTError  # Ensure this is correctly imported
+from .utils.db import tinydb_helper  # Ensure this instance is correctly initialized elsewhere
+from .utils.jwt_utils import SECRET_KEY, ALGORITHM, decode_jwt
 
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 
-def decode_access_token(token: str, credentials_exception) -> dict:
-    try:
-        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
-        user_id: str = payload.get("sub")
-        name: str = payload.get("name")
-        role: str = payload.get("role")
-        if user_id is None or name is None or role is None:
-            raise credentials_exception
-        return {"user_id": user_id, "name": name, "role": role}
-    except jwt.PyJWTError:
-        raise credentials_exception
-
-async def get_current_user(token: str = Depends(oauth2_scheme)) -> dict:
+async def get_current_user(token: str = Depends(oauth2_scheme)):
     credentials_exception = HTTPException(
         status_code=status.HTTP_401_UNAUTHORIZED,
         detail="Could not validate credentials",
         headers={"WWW-Authenticate": "Bearer"},
     )
-    return decode_access_token(token, credentials_exception)
+
+    # Utilize the centralized JWT decoding and catch any JWT-related errors
+    try:
+        payload = decode_jwt(token)
+    except JWTError:
+        raise credentials_exception
+
+    user_id: str = payload.get("sub")
+    if user_id is None:
+        raise credentials_exception
+
+    # Verify if the token is stored and valid
+    if not tinydb_helper.query_token(user_id, token):
+        raise credentials_exception
+    
+    # Payload is already obtained and validated, so just return it or its specific parts as needed
+    return {"user_id": user_id, "name": payload.get("name"), "role": payload.get("role")}
+
 

app/utils/db.py

@@ -20,14 +20,18 @@ class TinyDBHelper:
     def query_token(self, user_id: str, token: str) -> bool:
         """Query to check if the token exists and is valid."""
         User = Query()
-        # Assuming our tokens table contains 'user_id', 'token', and 'expires_at'
         result = self.tokens_table.search((User.user_id == user_id) & (User.token == token))
-        # Optionally, check if the token is expired
+
+        # Check if the result is empty (i.e., no matching token found)
+        if not result:
+            return False
+
+        # Check if the token is expired
         expires_at = datetime.fromisoformat(result[0]['expires_at'])
         if datetime.utcnow() > expires_at:
             return False
         
-        return bool(result)
+        return True
 
     def remove_token_by_value(self, token: str):
         """Remove a token based on its value."""

app/utils/jwt_utils.py

@@ -1,5 +1,6 @@
 from datetime import datetime, timedelta
 from jose import JWTError, jwt
+from fastapi import HTTPException
 from typing import Any, Union
 #from tinydb import TinyDB, Query
 #from tinydb.storages import MemoryStorage
@@ -10,23 +11,20 @@ SECRET_KEY = "a_very_secret_key"
 ALGORITHM = "HS256"
 ACCESS_TOKEN_EXPIRE_MINUTES = 30  # The expiration time for the access token
 
-# db = TinyDB(storage=MemoryStorage)
-# tokens_table = db.table('tokens')
-
-# def insert_token(user_id: str, token: str, expires_in: timedelta):
-#     expiration = datetime.utcnow() + expires_in
-#     tokens_table.insert({'user_id': user_id, 'token': token, 'expires_at': expiration.isoformat()})
-
-# def validate_token(user_id: str, token: str) -> bool:
-#     User = Query()
-#     result = tokens_table.search((User.user_id == user_id) & (User.token == token))
-#     if not result:
-#         return False
-#     # Check token expiration
-#     expires_at = datetime.fromisoformat(result[0]['expires_at'])
-#     if datetime.utcnow() > expires_at:
-#         return False
-#     return True
+def encode_jwt(data: dict):
+    # Encode a JWT token
+    return jwt.encode(data, SECRET_KEY, algorithm=ALGORITHM)
+
+def decode_jwt(token: str):
+    try:
+        # Decode a JWT token
+        return jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+    except JWTError as e:
+        # Handle specific JWT errors (e.g., token expired, invalid token)
+        raise HTTPException(status_code=401, detail="Token is invalid or expired")
+    except Exception as e:
+        # Handle unexpected errors
+        raise HTTPException(status_code=500, detail="An error occurred while decoding token")
 
 def create_access_token(data: dict, expires_delta: Union[timedelta, None] = None) -> str:
     """