Merge pull request #52 from igorbenav/token-blacklist

Token blacklist

Author: igorbenav

README.md

@@ -45,48 +45,48 @@ services:
     volumes:
       - redis-data:/data
 
-#   #-------- uncomment to create first superuser --------
-#   create_superuser:
-#     build: 
-#       context: .
-#       dockerfile: Dockerfile
-#     env_file:
-#       - ./src/.env
-#     depends_on:
-#       - db
-#       - web
-#     command: python -m src.scripts.create_first_superuser
-#     volumes:
-#       - ./src:/code/src
+  # #-------- uncomment to create first superuser --------
+  # create_superuser:
+  #   build: 
+  #     context: .
+  #     dockerfile: Dockerfile
+  #   env_file:
+  #     - ./src/.env
+  #   depends_on:
+  #     - db
+  #     - web
+  #   command: python -m src.scripts.create_first_superuser
+  #   volumes:
+  #     - ./src:/code/src
 
-#   #-------- uncomment to run tests --------
-#   # pytest:
-#   #   build: 
-#   #     context: .
-#   #     dockerfile: Dockerfile 
-#   #   env_file:
-#   #     - ./src/.env 
-#   #   depends_on:
-#   #     - db
-#   #     - create_superuser
-#   #     - redis
-#   #   command: python -m pytest
-#   #   volumes:
-#   #     - ./src:/code/src
+  #-------- uncomment to run tests --------
+  # pytest:
+  #   build: 
+  #     context: .
+  #     dockerfile: Dockerfile 
+  #   env_file:
+  #     - ./src/.env 
+  #   depends_on:
+  #     - db
+  #     - create_superuser
+  #     - redis
+  #   command: python -m pytest
+  #   volumes:
+  #     - ./src:/code/src
 
-#   #-------- uncomment to create first tier --------
-#   create_tier:
-#     build: 
-#       context: .
-#       dockerfile: Dockerfile
-#     env_file:
-#       - ./src/.env
-#     depends_on:
-#       - db
-#       - web
-#     command: python -m src.scripts.create_first_tier
-#     volumes:
-#       - ./src:/code/src
+  # #-------- uncomment to create first tier --------
+  # create_tier:
+  #   build: 
+  #     context: .
+  #     dockerfile: Dockerfile
+  #   env_file:
+  #     - ./src/.env
+  #   depends_on:
+  #     - db
+  #     - web
+  #   command: python -m src.scripts.create_first_tier
+  #   volumes:
+  #     - ./src:/code/src
 
 volumes:
   postgres-data:

docker-compose.yml

@@ -11,18 +11,18 @@
     Request
 )
 
+from app.api.exceptions import credentials_exception, privileges_exception
 from app.core.database import async_get_db
+from app.core.logger import logging
 from app.core.models import TokenData
 from app.core.rate_limit import is_rate_limited
-from app.core.logger import logging
-from app.models.user import User
-from app.api.exceptions import credentials_exception, privileges_exception
-from app.crud.crud_users import crud_users
-from app.crud.crud_tier import crud_tiers
+from app.core.security import verify_token
 from app.crud.crud_rate_limit import crud_rate_limits
+from app.crud.crud_tier import crud_tiers
+from app.crud.crud_users import crud_users
+from app.models.user import User
 from app.schemas.rate_limit import sanitize_path
 
-
 logger = logging.getLogger(__name__)
 
 DEFAULT_LIMIT = settings.DEFAULT_RATE_LIMIT_LIMIT
@@ -31,7 +31,7 @@
 async def get_current_user(
     token: Annotated[str, Depends(oauth2_scheme)],
     db: Annotated[AsyncSession, Depends(async_get_db)]
-) -> User:
+) -> dict:
     try:
         payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
         username_or_email: str = payload.get("sub")
@@ -53,10 +53,29 @@ async def get_current_user(
     raise credentials_exception
 
 
+async def get_current_user(
+    token: Annotated[str, Depends(oauth2_scheme)],
+    db: Annotated[AsyncSession, Depends(async_get_db)]
+) -> dict:
+    token_data = await verify_token(token, db)
+    if token_data is None:
+        raise credentials_exception
+
+    if "@" in token_data.username_or_email:
+        user = await crud_users.get(db=db, email=token_data.username_or_email, is_deleted=False)
+    else: 
+        user = await crud_users.get(db=db, username=token_data.username_or_email, is_deleted=False)
+    
+    if user:
+        return user
+
+    raise credentials_exception
+
+
 async def get_optional_user(
     request: Request,
     db: AsyncSession = Depends(async_get_db)
-) -> User | None:
+) -> dict | None:
     token = request.headers.get("Authorization")
     if not token:
         return None
@@ -66,7 +85,11 @@ async def get_optional_user(
         if token_type.lower() != 'bearer' or not token_value:
             return None
 
-        return await get_current_user(token_value, db)
+        token_data = await verify_token(token_value, db)
+        if token_data is None:
+            return None
+
+        return await get_current_user(token_value, is_deleted=False, db=db)
 
     except HTTPException as http_exc:
         if http_exc.status_code != 401:
@@ -75,10 +98,10 @@ async def get_optional_user(
 
     except Exception as exc:
         logger.error(f"Unexpected error in get_optional_user: {exc}")
-        return None    
+        return None
 
 
-async def get_current_superuser(current_user: Annotated[User, Depends(get_current_user)]) -> User:
+async def get_current_superuser(current_user: Annotated[User, Depends(get_current_user)]) -> dict:
     if not current_user["is_superuser"]:
         raise privileges_exception
 

src/app/api/dependencies.py

@@ -1,6 +1,7 @@
 from fastapi import APIRouter
 
 from app.api.v1.login import router as login_router
+from app.api.v1.logout import router as logout_router
 from app.api.v1.users import router as users_router
 from app.api.v1.posts import router as posts_router
 from app.api.v1.tasks import router as tasks_router
@@ -9,6 +10,7 @@
 
 router = APIRouter(prefix="/v1")
 router.include_router(login_router)
+router.include_router(logout_router)
 router.include_router(users_router)
 router.include_router(posts_router)
 router.include_router(tasks_router)

src/app/api/v1/__init__.py

@@ -18,7 +18,6 @@ async def login_for_access_token(
     form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
     db: Annotated[AsyncSession, Depends(async_get_db)]
 ):
-    
     user = await authenticate_user(
         username_or_email=form_data.username, 
         password=form_data.password, 

src/app/api/v1/login.py

@@ -0,0 +1,35 @@
+from datetime import datetime
+
+from fastapi import APIRouter, Depends, HTTPException, status
+from sqlalchemy.ext.asyncio import AsyncSession
+from jose import jwt, JWTError
+
+from app.core.security import oauth2_scheme, SECRET_KEY, ALGORITHM
+from app.core.database import async_get_db
+from app.crud.crud_token_blacklist import crud_token_blacklist
+from app.schemas.token_blacklist import TokenBlacklistCreate
+
+router = APIRouter(tags=["login"])
+
+@router.post("/logout")
+async def logout(
+    token: str = Depends(oauth2_scheme),
+    db: AsyncSession = Depends(async_get_db)
+):
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        expires_at = datetime.fromtimestamp(payload.get("exp"))
+        await crud_token_blacklist.create(
+            db, 
+            object=TokenBlacklistCreate(
+                **{"token": token, "expires_at": expires_at}
+            )
+        )
+        return {"message": "Logged out successfully"}
+    
+    except JWTError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid token",
+            headers={"WWW-Authenticate": "Bearer"},
+        )

src/app/api/v1/logout.py

@@ -16,7 +16,7 @@ class Token(BaseModel):
 
 
 class TokenData(BaseModel):
-    username_or_email: str | None = None
+    username_or_email: str
 
 
 class UUIDModel(BaseModel):

src/app/core/models.py

@@ -2,12 +2,13 @@
 
 from sqlalchemy.ext.asyncio import AsyncSession
 from passlib.context import CryptContext
-from jose import jwt
+from jose import jwt, JWTError
 from fastapi.security import OAuth2PasswordBearer
 
-from app.crud.crud_users import crud_users
 from app.core.config import settings
-
+from app.core.models import TokenData
+from app.crud.crud_token_blacklist import crud_token_blacklist
+from app.crud.crud_users import crud_users
 
 SECRET_KEY = settings.SECRET_KEY
 ALGORITHM = settings.ALGORITHM
@@ -45,3 +46,33 @@ async def create_access_token(data: dict, expires_delta: timedelta | None = None
     to_encode.update({"exp": expire})
     encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
     return encoded_jwt
+
+async def verify_token(token: str, db: AsyncSession) -> TokenData | None:
+    """
+    Verify a JWT token and return TokenData if valid.
+
+    Parameters
+    ----------
+    token: str 
+        The JWT token to be verified.
+    db: AsyncSession 
+        Database session for performing database operations.
+
+    Returns
+    -------
+    TokenData | None
+        TokenData instance if the token is valid, None otherwise.
+    """
+    is_blacklisted = await crud_token_blacklist.exists(db, token=token)
+    if is_blacklisted:
+        return None
+
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        username_or_email: str = payload.get("sub")
+        if username_or_email is None:
+            return None
+        return TokenData(username_or_email=username_or_email)
+    
+    except JWTError:
+        return None

src/app/core/security.py

@@ -0,0 +1,6 @@
+from app.crud.crud_base import CRUDBase
+from app.models.token_blacklist import TokenBlacklist
+from app.schemas.token_blacklist import TokenBlacklistCreate, TokenBlacklistUpdate
+
+CRUDTokenBlacklist = CRUDBase[TokenBlacklist, TokenBlacklistCreate, TokenBlacklistUpdate, TokenBlacklistUpdate, None]
+crud_token_blacklist = CRUDTokenBlacklist(TokenBlacklist)

src/app/crud/crud_token_blacklist.py

@@ -4,7 +4,7 @@
 
 from app.core.database import Base
 
-def _extract_matching_columns_from_schema(model: Type[Base], schema: Union[Type[BaseModel], List, None]) -> List[Any]:
+def _extract_matching_columns_from_schema(model: Type[Base], schema: Union[Type[BaseModel], list, None]) -> List[Any]:
     """
     Retrieves a list of ORM column objects from a SQLAlchemy model that match the field names in a given Pydantic schema.
 
@@ -46,7 +46,7 @@ def _extract_matching_columns_from_kwargs(model: Type[Base], kwargs: dict) -> Li
     return column_list
 
 
-def _extract_matching_columns_from_column_names(model: Type[Base], column_names: List) -> List[Any]:
+def _extract_matching_columns_from_column_names(model: Type[Base], column_names: list) -> List[Any]:
     column_list = []
     for column_name in column_names:
         if hasattr(model, column_name):