Merge pull request #159 from igorbenav/blacklist-refresh-token

blacklist both tokens for logout

Author: igorbenav

src/app/api/v1/logout.py

@@ -1,20 +1,31 @@
-from fastapi import APIRouter, Depends, Response
+from fastapi import APIRouter, Depends, Response, Cookie
 from jose import JWTError
 from sqlalchemy.ext.asyncio import AsyncSession
+from typing import Optional
 
 from ...core.db.database import async_get_db
 from ...core.exceptions.http_exceptions import UnauthorizedException
-from ...core.security import blacklist_token, oauth2_scheme
+from ...core.security import blacklist_tokens, oauth2_scheme
 
 router = APIRouter(tags=["login"])
 
 
 @router.post("/logout")
 async def logout(
-    response: Response, access_token: str = Depends(oauth2_scheme), db: AsyncSession = Depends(async_get_db)
+    response: Response,
+    access_token: str = Depends(oauth2_scheme),
+    refresh_token: Optional[str] = Cookie(None, alias="refresh_token"),
+    db: AsyncSession = Depends(async_get_db)
 ) -> dict[str, str]:
     try:
-        await blacklist_token(token=access_token, db=db)
+        if not refresh_token:
+            raise UnauthorizedException("Refresh token not found")
+            
+        await blacklist_tokens(
+            access_token=access_token,
+            refresh_token=refresh_token,
+            db=db
+        )
         response.delete_cookie(key="refresh_token")
 
         return {"message": "Logged out successfully"}

src/app/core/security.py

@@ -1,3 +1,4 @@
+from enum import Enum
 from datetime import UTC, datetime, timedelta
 from typing import Any, Literal
 
@@ -11,6 +12,7 @@
 from .db.crud_token_blacklist import crud_token_blacklist
 from .schemas import TokenBlacklistCreate, TokenData
 
+
 SECRET_KEY = settings.SECRET_KEY
 ALGORITHM = settings.ALGORITHM
 ACCESS_TOKEN_EXPIRE_MINUTES = settings.ACCESS_TOKEN_EXPIRE_MINUTES
@@ -19,6 +21,10 @@
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/v1/login")
 
 
+class TokenType(str, Enum):
+    ACCESS = "access"
+    REFRESH = "refresh"
+
 async def verify_password(plain_password: str, hashed_password: str) -> bool:
     correct_password: bool = bcrypt.checkpw(plain_password.encode(), hashed_password.encode())
     return correct_password
@@ -50,7 +56,7 @@ async def create_access_token(data: dict[str, Any], expires_delta: timedelta | N
         expire = datetime.now(UTC).replace(tzinfo=None) + expires_delta
     else:
         expire = datetime.now(UTC).replace(tzinfo=None) + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
-    to_encode.update({"exp": expire})
+    to_encode.update({"exp": expire, "token_type": TokenType.ACCESS})
     encoded_jwt: str = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
     return encoded_jwt
 
@@ -61,18 +67,20 @@ async def create_refresh_token(data: dict[str, Any], expires_delta: timedelta |
         expire = datetime.now(UTC).replace(tzinfo=None) + expires_delta
     else:
         expire = datetime.now(UTC).replace(tzinfo=None) + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
-    to_encode.update({"exp": expire})
+    to_encode.update({"exp": expire, "token_type": TokenType.REFRESH})
     encoded_jwt: str = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
     return encoded_jwt
 
 
-async def verify_token(token: str, db: AsyncSession) -> TokenData | None:
+async def verify_token(token: str, expected_token_type: TokenType, db: AsyncSession) -> TokenData | None:
     """Verify a JWT token and return TokenData if valid.
 
     Parameters
     ----------
     token: str
         The JWT token to be verified.
+    expected_token_type: TokenType
+        The expected type of token (access or refresh)
     db: AsyncSession
         Database session for performing database operations.
 
@@ -88,15 +96,47 @@ async def verify_token(token: str, db: AsyncSession) -> TokenData | None:
     try:
         payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
         username_or_email: str = payload.get("sub")
-        if username_or_email is None:
+        token_type: str = payload.get("token_type")
+        
+        if username_or_email is None or token_type != expected_token_type:
             return None
+            
         return TokenData(username_or_email=username_or_email)
 
     except JWTError:
         return None
 
 
+async def blacklist_tokens(access_token: str, refresh_token: str, db: AsyncSession) -> None:
+    """Blacklist both access and refresh tokens.
+
+    Parameters
+    ----------
+    access_token: str
+        The access token to blacklist
+    refresh_token: str
+        The refresh token to blacklist
+    db: AsyncSession
+        Database session for performing database operations.
+    """
+    for token in [access_token, refresh_token]:
+        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
+            )
+        )
+
 async def blacklist_token(token: str, db: AsyncSession) -> None:
     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}))
+    await crud_token_blacklist.create(
+        db, 
+        object=TokenBlacklistCreate(
+            token=token,
+            expires_at=expires_at
+        )
+    )