totp token verify

Author: dfm88

env/.env

@@ -16,7 +16,7 @@ TFA_TOKEN_LENGTH=6
 # default tolerance = 30 sec
 # this number is multiplied for 30 sec to increas it.
 # -->MAX = 10 => 5 minutes
-TOTP_TOKEN_DURATION=2
+TOTP_TOKEN_TOLERANCE=2
 TOTP_ISSUER_NAME=fastapi_2fa
 
 ACCESS_TOKEN_EXPIRE_MINUTES=30

fastapi_2fa/api/deps/users.py

@@ -20,7 +20,8 @@ async def _get_user_from_jwt(
     token: str,
     db: Session,
     expire_err_message: str = "Token expired",
-    jwt_err_message: str = "Could not validate credentials",
+    jwt_err_message: str = "Could not validate credentials, "
+                           "if TFA is enabled, please confirm token first",
 ):
     try:
         payload = jwt.decode(token=token, key=key, algorithms=[settings.ALGORITHM])

fastapi_2fa/api/endpoints/api_v1/auth.py

@@ -15,10 +15,11 @@
 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.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
+from fastapi_2fa.schemas.token_schema import TokenPayload, TokenSchema, PreTfaTokenSchema
 from fastapi_2fa.schemas.user_schema import UserCreate, UserOut
 
 auth_router = APIRouter()
@@ -75,7 +76,7 @@ async def signup(
     "/login",
     summary="Create access and refresh tokens for user",
     status_code=status.HTTP_200_OK,
-    response_model=TokenSchema,
+    response_model=TokenSchema | PreTfaTokenSchema,
 )
 async def login(
     response: Response,
@@ -96,7 +97,7 @@ async def login(
     # verify 2 factor authentication
     if user_crud.is_tfa_enabled(user=user):
         response.status_code = status.HTTP_202_ACCEPTED
-        return TokenSchema(
+        return PreTfaTokenSchema(
             access_token=security.create_pre_tfa_token(user.id),
             refresh_token=None,
         )
@@ -110,16 +111,24 @@ async def login(
 
 @auth_router.post(
     "/login/tfa",
-    summary="Verify two factor authenticazion token",
-    response_model=UserOut,
+    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:
-    print(f"{tfa_token}")
-    return user
+    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(

fastapi_2fa/core/config.py

@@ -28,8 +28,8 @@ class BaseConfig(BaseSettings):
     # default tolerance = 30 sec
     # this number is multiplied for 30 sec to increas it.
     # -->MAX = 10 => 5 minutes
-    TOTP_TOKEN_DURATION: int = Field(
-        default=os.environ.get("TOTP_TOKEN_DURATION", 2),
+    TOTP_TOKEN_TOLERANCE: int = Field(
+        default=os.environ.get("TOTP_TOKEN_TOLERANCE", 2),
         gt=0,
         le=10
     )

fastapi_2fa/core/two_factor_auth.py

@@ -46,10 +46,20 @@ def get_fake_otp_tokens(
         )
         yield random_otp
 
+
+def verify_token(user: User, token:str) -> bool:
+    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 result
+
+
 def qr_code_from_key(encoded_key: str, user_email: str):
     decoded_key = _fernet_decode(value=encoded_key)
     qrcode_key = pyotp.TOTP(
-        decoded_key, interval=settings.TOTP_TOKEN_DURATION
+        decoded_key, interval=settings.TOTP_TOKEN_TOLERANCE
     ).provisioning_uri(user_email, issuer_name=settings.TOTP_ISSUER_NAME)
     img = qrcode.make(qrcode_key)
     buffer = io.BytesIO()

fastapi_2fa/models/device.py

@@ -41,6 +41,7 @@ def __repr__(self) -> str:
         return (
             f"<{self.__class__.__name__}("
             f"id={self.id}, "
+            f"token={self.token}, "
             f"device={self.device}, "
             f")>"
         )

fastapi_2fa/schemas/token_schema.py

@@ -3,13 +3,13 @@
 
 class TokenSchema(BaseModel):
     access_token: str
-    refresh_token: str | None
+    refresh_token: str
 
 
 class TokenPayload(BaseModel):
     sub: int = None
     exp: int = None
 
 
-class PreTfaTokenSchema(BaseModel):
-    pre_tfa_token: str
+class PreTfaTokenSchema(TokenSchema):
+    refresh_token: None