handled email device and backup tokens

Author: dfm88

fastapi_2fa/api/deps/users.py

@@ -8,7 +8,7 @@
 from fastapi_2fa.core.config import settings
 from fastapi_2fa.crud.users import user_crud
 from fastapi_2fa.models.users import User
-from fastapi_2fa.schemas.token_schema import TokenPayload
+from fastapi_2fa.schemas.jwt_token_schema import JwtTokenPayload
 
 reuseable_oauth = OAuth2PasswordBearer(
     tokenUrl=f"{settings.API_V1_STR}/auth/login", scheme_name="JWT"
@@ -25,7 +25,7 @@ async def _get_user_from_jwt(
 ):
     try:
         payload = jwt.decode(token=token, key=key, algorithms=[settings.ALGORITHM])
-        token_data = TokenPayload(**payload)
+        token_data = JwtTokenPayload(**payload)
     except jwt.ExpiredSignatureError:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,

fastapi_2fa/api/endpoints/api_v1/auth.py

@@ -9,17 +9,18 @@
 from sqlalchemy.orm import Session
 
 from fastapi_2fa.api.deps.db import get_db
-from fastapi_2fa.api.deps.users import (get_authenticated_user,
-                                        get_authenticated_user_pre_tfa)
+from fastapi_2fa.api.deps.users import get_authenticated_user
 from fastapi_2fa.core import security
 from fastapi_2fa.core.config import settings
 from fastapi_2fa.core.enums import DeviceTypeEnum
-from fastapi_2fa.core.utils import send_backup_tokens
-from fastapi_2fa.core.two_factor_auth import verify_token
+from fastapi_2fa.core.two_factor_auth import send_tfa_token
+from fastapi_2fa.core.utils import send_mail_backup_tokens
 from fastapi_2fa.crud.device import device_crud
 from fastapi_2fa.crud.users import user_crud
 from fastapi_2fa.models.users import User
-from fastapi_2fa.schemas.token_schema import TokenPayload, TokenSchema, PreTfaTokenSchema
+from fastapi_2fa.schemas.jwt_token_schema import (JwtTokenPayload,
+                                                  JwtTokenSchema,
+                                                  PreTfaJwtTokenSchema)
 from fastapi_2fa.schemas.user_schema import UserCreate, UserOut
 
 auth_router = APIRouter()
@@ -53,7 +54,7 @@ async def signup(
                 device=user_data.device,
                 user=user
             )
-            send_backup_tokens(user=user, device=device)
+            send_mail_backup_tokens(user=user, device=device)
 
             if device.device_type == DeviceTypeEnum.CODE_GENERATOR:
                 return StreamingResponse(content=qr_code, media_type="image/png")
@@ -76,7 +77,7 @@ async def signup(
     "/login",
     summary="Create access and refresh tokens for user",
     status_code=status.HTTP_200_OK,
-    response_model=TokenSchema | PreTfaTokenSchema,
+    response_model=JwtTokenSchema | PreTfaJwtTokenSchema,
 )
 async def login(
     response: Response,
@@ -94,44 +95,26 @@ async def login(
             detail="Incorrect email or password",
         )
 
-    # verify 2 factor authentication
-    if user_crud.is_tfa_enabled(user=user):
+    # handle users with tfa enabled
+    if user.tfa_enabled:
+        send_tfa_token(
+            user=user,
+            device_type=user.device.device_type
+        )
         response.status_code = status.HTTP_202_ACCEPTED
-        return PreTfaTokenSchema(
+        return PreTfaJwtTokenSchema(
             access_token=security.create_pre_tfa_token(user.id),
             refresh_token=None,
         )
 
     # create access and refresh tokens
-    return TokenSchema(
+    return JwtTokenSchema(
         access_token=security.create_jwt_access_token(user.id),
         refresh_token=security.create_jwt_refresh_token(user.id),
     )
 
 
-@auth_router.post(
-    "/login/tfa",
-    summary="Verify two factor authentication token",
-    response_model=TokenSchema,
-)
-async def login_tfa(
-    tfa_token: str,
-    db: Session = Depends(get_db),
-    user: User = Depends(get_authenticated_user_pre_tfa),
-) -> Any:
-    if verify_token(user=user, token=tfa_token):
-        return TokenSchema(
-            access_token=security.create_jwt_access_token(user.id),
-            refresh_token=security.create_jwt_refresh_token(user.id),
-        )
-    
-    raise HTTPException(
-        status_code=status.HTTP_403_FORBIDDEN,
-        detail="TOTP token mismatch"
-    )
-
-
-@auth_router.post(
+@auth_router.get(
     "/test-token", summary="Test if the access token is ok", response_model=UserOut
 )
 async def test_token(user: User = Depends(get_authenticated_user)):
@@ -141,7 +124,7 @@ async def test_token(user: User = Depends(get_authenticated_user)):
     return user
 
 
-@auth_router.post("/refresh", summary="Refresh token", response_model=TokenSchema)
+@auth_router.post("/refresh", summary="Refresh token", response_model=JwtTokenSchema)
 async def refresh_token(
     db: Session = Depends(get_db), refresh_token: str = Body(embed=True, title='refresh token')
 ):
@@ -151,7 +134,7 @@ async def refresh_token(
             key=settings.JWT_SECRET_KEY_REFRESH,
             algorithms=[settings.ALGORITHM],
         )
-        token_data = TokenPayload(**payload)
+        token_data = JwtTokenPayload(**payload)
     except (jwt.JWTError, ValidationError):
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,

fastapi_2fa/api/endpoints/api_v1/two_factor_auth.py

@@ -5,18 +5,123 @@
 from sqlalchemy.orm import Session
 
 from fastapi_2fa.api.deps.db import get_db
-from fastapi_2fa.api.deps.users import get_authenticated_user
+from fastapi_2fa.api.deps.users import (get_authenticated_user,
+                                        get_authenticated_user_pre_tfa)
+from fastapi_2fa.core import security
 from fastapi_2fa.core.enums import DeviceTypeEnum
-from fastapi_2fa.core.utils import send_backup_tokens
+from fastapi_2fa.core.two_factor_auth import (qr_code_from_key,
+                                              verify_backup_token,
+                                              verify_token)
+from fastapi_2fa.core.utils import send_mail_backup_tokens
+from fastapi_2fa.crud.backup_token import backup_token_crud
 from fastapi_2fa.crud.device import device_crud
 from fastapi_2fa.crud.users import user_crud
 from fastapi_2fa.models.users import User
 from fastapi_2fa.schemas.device_schema import DeviceCreate
+from fastapi_2fa.schemas.jwt_token_schema import JwtTokenSchema
 from fastapi_2fa.schemas.user_schema import UserUpdate
 
 tfa_router = APIRouter()
 
 
+@tfa_router.post(
+    "/login_tfa",
+    summary="Verify two factor authentication token",
+    response_model=JwtTokenSchema,
+)
+async def login_tfa(
+    tfa_token: str,
+    db: Session = Depends(get_db),
+    user: User = Depends(get_authenticated_user_pre_tfa),
+) -> Any:
+    if verify_token(user=user, token=tfa_token):
+        return JwtTokenSchema(
+            access_token=security.create_jwt_access_token(user.id),
+            refresh_token=security.create_jwt_refresh_token(user.id),
+        )
+
+    raise HTTPException(
+        status_code=status.HTTP_403_FORBIDDEN,
+        detail="TOTP token mismatch"
+    )
+
+
+@tfa_router.post(
+    "/recover_tfa",
+    summary="Checks and consumes one of the user's backup tokens re initializing",
+    response_model=JwtTokenSchema,
+)
+async def recover_tfa(
+    tfa_backup_token: str,
+    db: Session = Depends(get_db),
+    user: User = Depends(get_authenticated_user_pre_tfa),
+) -> Any:
+    if backup_tokens := await backup_token_crud.get_user_backup_tokens(
+        db=db,
+        user=user
+    ):
+        matched_bkp_token = verify_backup_token(
+            backup_tokens=backup_tokens,
+            tfa_backup_token=tfa_backup_token
+        )
+
+        if matched_bkp_token:
+            print('..consuming backup token')
+            await backup_token_crud.remove(
+                db=db, id=matched_bkp_token.id
+            )
+            return JwtTokenSchema(
+                access_token=security.create_jwt_access_token(user.id),
+                refresh_token=security.create_jwt_refresh_token(user.id),
+            )
+
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="TOTP backup token not found"
+        )
+
+    # user has elapsed all backup tokens
+    raise HTTPException(
+        status_code=status.HTTP_404_NOT_FOUND,
+        detail=f"User {user.email} has elapsed his backup tokens, "
+               "please contact the system administrator"
+    )
+
+
+@tfa_router.get(
+    "/get_my_qrcode",
+    summary="Returns authenticated user's qr_code "
+            "if user's device is of type 'code_generator'",
+    responses={
+        200: {
+            "content": {"image/png": {}},
+            "description": "Returns no content or a qr code "
+                           "if tfas is enabled and device_type "
+                           "is 'code_generator'",
+        }
+    },
+)
+async def get_my_qrcode(
+    user: User = Depends(get_authenticated_user),
+) -> Any:
+    if (
+        user.tfa_enabled and
+        user_crud.device.device_type == DeviceTypeEnum.CODE_GENERATOR
+    ):
+        qr_code = qr_code_from_key(
+            encoded_key=user.device.key,
+            user_email=user.email
+        )
+        return StreamingResponse(content=qr_code, media_type="image/png")
+
+    # user has elapsed all backup tokens
+    raise HTTPException(
+        status_code=status.HTTP_400_BAD_REQUEST,
+        detail=f"User {user.email} has not tfa enabled or "
+               "has not a 'code_generator' device"
+    )
+
+
 @tfa_router.put(
     "/enable_tfa",
     summary="Enable two factor authentication for registered user",
@@ -33,8 +138,8 @@ async def enable_tfa(
     db: Session = Depends(get_db),
     user: User = Depends(get_authenticated_user),
 ) -> Any:
-    if not user_crud(transaction=True).is_tfa_enabled(user):
-        user = await user_crud.update(
+    if not user.tfa_enabled:
+        user = await user_crud(transaction=True).update(
             db=db,
             db_obj=user,
             obj_in=UserUpdate(tfa_enabled=True)
@@ -45,7 +150,7 @@ async def enable_tfa(
             user=user
         )
 
-        send_backup_tokens(user=user, device=device)
+        send_mail_backup_tokens(user=user, device=device)
 
         if device.device_type == DeviceTypeEnum.CODE_GENERATOR:
             return StreamingResponse(content=qr_code, media_type="image/png")

fastapi_2fa/core/two_factor_auth.py

@@ -7,6 +7,10 @@
 from fernet import Fernet
 
 from fastapi_2fa.core.config import settings
+from fastapi_2fa.core.enums import DeviceTypeEnum
+from fastapi_2fa.core.utils import send_mail_totp_token
+from fastapi_2fa.models.device import BackupToken
+from fastapi_2fa.models.users import User
 
 ENCODING = 'utf-8'
 
@@ -47,15 +51,52 @@ def get_fake_otp_tokens(
         yield random_otp
 
 
-def verify_token(user: User, token:str) -> bool:
+
+def get_current_totp(user: User) -> pyotp.TOTP:
     assert user.tfa_enabled is True, 'User does not have TFA enabled'
     assert user.device is not None, 'User has no associated device'
     decoded_key = _fernet_decode(value=user.device.key)
-    totp = pyotp.TOTP(decoded_key)
-    result =  totp.verify(token, valid_window=settings.TOTP_TOKEN_TOLERANCE)
+    return pyotp.TOTP(decoded_key)
+
+
+def verify_token(user: User, token: str) -> bool:
+    totp: pyotp.TOTP = get_current_totp(user=user)
+    result = totp.verify(token, valid_window=settings.TOTP_TOKEN_TOLERANCE)
     return result
 
 
+def verify_backup_token(
+    backup_tokens: list[BackupToken], tfa_backup_token: str
+) -> BackupToken | None:
+    """Checks between "backup_tokens" if there is the "tfa_backup_token"
+
+    Args:
+        backup_tokens list[BackupToken]
+        tfa_backup_token (str)
+
+    Returns:
+        BackupToken | None: the matched BackupToken or None if no backup tokens matched
+    """
+    for bkp_token in backup_tokens:
+        if bkp_token.token == tfa_backup_token:
+            # consume backup token and return access jwt
+            print('Found backup token match..')
+            return bkp_token
+
+
+def send_tfa_token(
+    user: User,
+    device_type: DeviceTypeEnum
+) -> None:
+    if device_type == DeviceTypeEnum.EMAIL:
+        totp: pyotp.TOTP = get_current_totp(user=user)
+        current_totp = totp.now()
+        send_mail_totp_token(
+            user=user,
+            token=current_totp
+        )
+
+
 def qr_code_from_key(encoded_key: str, user_email: str):
     decoded_key = _fernet_decode(value=encoded_key)
     qrcode_key = pyotp.TOTP(

fastapi_2fa/core/utils.py

@@ -22,7 +22,7 @@ def __repr__(self) -> str:
         )
 
 
-def send_backup_tokens(user: User, device: Device):
+def send_mail_backup_tokens(user: User, device: Device):
     email_ob = Email(
         to_=[user.email],
         from_=settings.FAKE_EMAIL_SENDER,
@@ -31,7 +31,7 @@ def send_backup_tokens(user: User, device: Device):
     send_email(email=email_ob)
 
 
-def send_totp_token(user: User, token: str):
+def send_mail_totp_token(user: User, token: str):
     email_ob = Email(
         to_=[user.email],
         from_=settings.FAKE_EMAIL_SENDER,

fastapi_2fa/crud/backup_token.py

@@ -0,0 +1,26 @@
+from sqlalchemy import select
+from sqlalchemy.orm import Session, joinedload
+
+from fastapi_2fa.crud.base_crud import CrudBase
+from fastapi_2fa.models.device import BackupToken, Device
+from fastapi_2fa.models.users import User
+from fastapi_2fa.schemas.backup_token_schema import (BackupTokenCreate,
+                                                     BackupTokenUpdate)
+
+
+class BackupTokenCrud(CrudBase[Device, BackupTokenCreate, BackupTokenUpdate]):
+
+    @staticmethod
+    async def get_user_backup_tokens(db: Session, user: User) -> list:
+        if not user.device:
+            return []
+        result = await db.execute(
+            select(Device).where(
+                Device.id == user.device.id
+            ).options(joinedload(Device.backup_tokens))
+        )
+        device_with_bkp_tokens: Device = result.scalar()
+        return device_with_bkp_tokens.backup_tokens
+
+
+backup_token_crud = BackupTokenCrud(model=BackupToken)