Multi-Factor Authentication Integration

Introduction

Welcome to the third lesson of the "Multi-Factor Authentication (MFA) in FastAPI" course! In this lesson, we will focus on integrating MFA into a FastAPI application. Building on the foundational concepts from the previous lesson, we will explore the practical steps necessary to enhance the security of your application by implementing MFA. Let's dive in and see how we can secure our application with MFA!

Secure MFA Integration Patterns

Now that you're familiar with the core MFA functions like generate_secret, verify_token, and generate_backup_codes, let's focus on integrating these functions securely into your FastAPI application. Proper integration is crucial - even with well-implemented MFA functions, vulnerabilities can arise from insecure application design. We'll examine how to store MFA data securely, protect sensitive endpoints, and create a robust authentication flow. Security Reminder: Always encrypt MFA secrets at rest to prevent attackers from accessing them if your database is compromised. Important: Unlike passwords, MFA secrets must be encrypted (not hashed) because they must be retrieved for token verification. Encryption is reversible, which is necessary for TOTP to work. Always hash backup codes before storing them, just like passwords, so they cannot be used if leaked. Backup codes can be hashed because they're compared directly like passwords (the user enters the code, you hash it and compare to stored hash). Implement rate-limiting on all sensitive endpoints (such as /verify, /login, and /login/verify) to protect against brute-force attacks.

The Vulnerable Code

Let's examine a scenario where MFA is not properly integrated, leading to potential security vulnerabilities: Pythonfrom fastapi import APIRouter, HTTPException import pyotp router = APIRouter() # Example of a vulnerable MFA setup @router.post('/verify') async def verify_mfa(username: str, token: str): user = await User.find_by_username(username) totp = pyotp.TOTP(user.mfa_secret) is_valid = totp.verify(token) if is_valid: user.mfa_enabled = True await user.save() return {"success": True} else: return {"success": False, "error": "Invalid token"} This code contains several critical vulnerabilities: No authentication checks - anyone can access this endpoint No validation that the user exists before accessing properties No protection against brute force attacks The endpoint updates user settings without proper authorization

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Implementing MFA Setup

Let's create a secure route for setting up MFA, using the functions you already know:

Python

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
import pyotp
import qrcode
import io
import base64
import json

router = APIRouter()

# Route to set up MFA (requires authentication middleware)
@router.post('/setup')
async def mfa_setup(
    username: str,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user)  # Authentication dependency
):
    try:
        user = await get_user_by_id(current_user.id, db)
        if not user:
            raise HTTPException(status_code=404, detail="User not found")
        
        # Generate MFA secret using pyotp
        secret = pyotp.random_base32()
        
        # Encrypt the secret before storing
        encrypted_secret = encrypt_secret(secret)
        
        # Generate backup codes and hash them before storing
        backup_codes = generate_backup_codes()
        hashed_backup_codes = [
            {"hash": hash_backup_code(code), "used": False}
            for code in backup_codes
        ]
        
        # Set both mfa_secret and backup_codes before saving
        user.mfa_secret = encrypted_secret
        user.mfa_enabled = False
        user.backup_codes = json.dumps(hashed_backup_codes)
        await db.commit()
        
        # Generate QR code
        totp = pyotp.TOTP(secret)
        provisioning_uri = totp.provisioning_uri(
            name=user.username,
            issuer_name='My App'
        )
        
        qr = qrcode.QRCode(version=1, box_size=10, border=5)
        qr.add_data(provisioning_uri)
        qr.make(fit=True)
        
        img = qr.make_image(fill_color="black", back_color="white")
        img_buffer = io.BytesIO()
        img.save(img_buffer, format='PNG')
        img_buffer.seek(0)
        
        img_base64 = base64.b64encode(img_buffer.getvalue()).decode()
        qr_code_url = f"data:image/png;base64,{img_base64}"
        
        return {
            "success": True,
            "qr_code_url": qr_code_url,
            "backup_codes": backup_codes,  # Return plain codes to user only at creation
            "message": "Scan the QR code with your authenticator app, then verify with a token"
        }
    except Exception as error:
        raise HTTPException(status_code=500, detail="Failed to set up MFA")

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
import pyotp
import qrcode
import io
import base64
import json

router = APIRouter()

# Route to set up MFA (requires authentication middleware)
@router.post('/setup')
async def mfa_setup(
    username: str,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user)  # Authentication dependency
):
    try:
        user = await get_user_by_id(current_user.id, db)
        if not user:
            raise HTTPException(status_code=404, detail="User not found")
        
        # Generate MFA secret using pyotp
        secret = pyotp.random_base32()
        
        # Encrypt the secret before storing
        encrypted_secret = encrypt_secret(secret)
        
        # Generate backup codes and hash them before storing
        backup_codes = generate_backup_codes()
        hashed_backup_codes = [
            {"hash": hash_backup_code(code), "used": False}
            for code in backup_codes
        ]
        
        # Set both mfa_secret and backup_codes before saving
        user.mfa_secret = encrypted_secret
        user.mfa_enabled = False
        user.backup_codes = json.dumps(hashed_backup_codes)
        await db.commit()
        
        # Generate QR code
        totp = pyotp.TOTP(secret)
        provisioning_uri = totp.provisioning_uri(
            name=user.username,
            issuer_name='My App'
        )
        
        qr = qrcode.QRCode(version=1, box_size=10, border=5)
        qr.add_data(provisioning_uri)
        qr.make(fit=True)
        
        img = qr.make_image(fill_color="black", back_color="white")
        img_buffer = io.BytesIO()
        img.save(img_buffer, format='PNG')
        img_buffer.seek(0)
        
        img_base64 = base64.b64encode(img_buffer.getvalue()).decode()
        qr_code_url = f"data:image/png;base64,{img_base64}"
        
        return {
            "success": True,
            "qr_code_url": qr_code_url,
            "backup_codes": backup_codes,  # Return plain codes to user only at creation
            "message": "Scan the QR code with your authenticator app, then verify with a token"
        }
    except Exception as error:
        raise HTTPException(status_code=500, detail="Failed to set up MFA")

This secure implementation includes:

Authentication dependency to ensure only logged-in users can access it
Error handling for failed operations
Proper user verification
Temporary storage of the secret until verified
Encryption of secrets and hashing of backup codes with usage tracking
(In production) Add rate-limiting middleware to this endpoint

Verifying MFA Tokens Securely

Next, we'll create a secure route to verify the MFA token provided by the user:

Python

@router.post('/verify')
async def mfa_verify(
    token: str,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user)  # Authentication dependency
):
    try:
        if not token:
            raise HTTPException(status_code=400, detail="Token is required")
        
        user = await get_user_by_id(current_user.id, db)
        if not user or not user.mfa_secret:
            raise HTTPException(status_code=400, detail="MFA not set up")
        
        # Decrypt the secret before verification
        decrypted_secret = decrypt_secret(user.mfa_secret)
        
        # Verify the token
        totp = pyotp.TOTP(decrypted_secret)
        is_valid = totp.verify(token)
        
        if is_valid:
            # Only now do we enable MFA for the user
            user.mfa_enabled = True
            await db.commit()
            
            # Generate a new session token or refresh existing session
            session_token = generate_session_token(user)
            
            return {
                "success": True,
                "message": "MFA enabled successfully",
                "session_token": session_token
            }
        else:
            raise HTTPException(status_code=401, detail="Invalid token")
    except HTTPException:
        raise
    except Exception as error:
        raise HTTPException(status_code=500, detail="Failed to verify token")

@router.post('/verify')
async def mfa_verify(
    token: str,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user)  # Authentication dependency
):
    try:
        if not token:
            raise HTTPException(status_code=400, detail="Token is required")
        
        user = await get_user_by_id(current_user.id, db)
        if not user or not user.mfa_secret:
            raise HTTPException(status_code=400, detail="MFA not set up")
        
        # Decrypt the secret before verification
        decrypted_secret = decrypt_secret(user.mfa_secret)
        
        # Verify the token
        totp = pyotp.TOTP(decrypted_secret)
        is_valid = totp.verify(token)
        
        if is_valid:
            # Only now do we enable MFA for the user
            user.mfa_enabled = True
            await db.commit()
            
            # Generate a new session token or refresh existing session
            session_token = generate_session_token(user)
            
            return {
                "success": True,
                "message": "MFA enabled successfully",
                "session_token": session_token
            }
        else:
            raise HTTPException(status_code=401, detail="Invalid token")
    except HTTPException:
        raise
    except Exception as error:
        raise HTTPException(status_code=500, detail="Failed to verify token")

This implementation includes:

Input validation for the token
Authentication checks
Verification that MFA is set up
Proper error handling
Session regeneration after successful verification
Decryption of secrets before use
(In production) Add rate-limiting middleware to this endpoint

Implementing Login with MFA

When MFA is enabled, the login process becomes a two-step procedure. In this flow, a temporary token (often called a temp_token) is issued after the user successfully enters their username and password, but before they complete MFA verification.

A temp_token is a short-lived, signed token (such as a JWT or a random string stored server-side) that securely encodes the user's identity and the fact that they have passed the first authentication step. It is used to track the login session between the password and MFA steps, ensuring that only the correct user can proceed to MFA verification. Typically, a temp_token should expire quickly—commonly within 5 to 10 minutes—to minimize the risk of misuse.

Python

from datetime import datetime, timedelta
from jose import jwt

@router.post('/login')
async def login(username: str, password: str, db: AsyncSession = Depends(get_db)):
    try:
        user = await get_user_by_username(username, db)
        
        # Verify username and password
        if not user or not verify_password(password, user.password_hash):
            raise HTTPException(status_code=401, detail="Invalid credentials")
        
        # Check if MFA is enabled
        if user.mfa_enabled:
            # Generate a temporary login token that will be used to complete login
            temp_token_data = {
                'user_id': user.id,
                'exp': datetime.utcnow() + timedelta(minutes=10),
                'type': 'temp_login'
            }
            temp_token = jwt.encode(temp_token_data, JWT_SECRET_KEY, algorithm='HS256')
            
            return {
                "success": True,
                "require_mfa": True,
                "temp_token": temp_token,
                "message": "Please enter your MFA code to complete login"
            }
        else:
            # Regular login without MFA
            session_token = generate_session_token(user)
            return {"success": True, "session_token": session_token}
    except HTTPException:
        raise
    except Exception as error:
        raise HTTPException(status_code=500, detail="Login failed")

# Second step for MFA verification
@router.post('/login/verify')
async def login_verify_mfa(
    temp_token: str, 
    mfa_code: str, 
    db: AsyncSession = Depends(get_db)
):
    try:
        # Verify the temporary token
        try:
            decoded = jwt.decode(temp_token, JWT_SECRET_KEY, algorithms=['HS256'])
            if decoded.get('type') != 'temp_login':
                raise HTTPException(status_code=401, detail="Invalid session")
            user_id = decoded['user_id']
        except jwt.InvalidTokenError:
            raise HTTPException(status_code=401, detail="Invalid session")
        
        user = await get_user_by_id(user_id, db)
        if not user or not user.mfa_secret:
            raise HTTPException(status_code=400, detail="User not found")
        
        # Decrypt the secret before verification
        decrypted_secret = decrypt_secret(user.mfa_secret)
        
        # Convert stored hashed backup codes to set
        hashed_backup_codes = set()
        if user.backup_codes:
            backup_codes = json.loads(user.backup_codes)
            hashed_backup_codes = {code['hash'] for code in backup_codes if not code.get('used')}
        
        # Verify using both token and backup codes
        totp = pyotp.TOTP(decrypted_secret)
        is_valid = totp.verify(mfa_code)
        
        # Check backup codes if TOTP fails
        if not is_valid and hashed_backup_codes:
            hashed_input = hash_backup_code(mfa_code)
            if hashed_input in hashed_backup_codes:
                # Mark backup code as used
                backup_codes = json.loads(user.backup_codes)
                for code in backup_codes:
                    if code['hash'] == hashed_input:
                        code['used'] = True
                        break
                user.backup_codes = json.dumps(backup_codes)
                await db.commit()
                is_valid = True
        
        if is_valid:
            # Complete login
            session_token = generate_session_token(user)
            return {"success": True, "session_token": session_token}
        else:
            raise HTTPException(status_code=401, detail="Invalid MFA code")
    except HTTPException:
        raise
    except Exception as error:
        raise HTTPException(status_code=500, detail="Verification failed")

from datetime import datetime, timedelta
from jose import jwt

@router.post('/login')
async def login(username: str, password: str, db: AsyncSession = Depends(get_db)):
    try:
        user = await get_user_by_username(username, db)
        
        # Verify username and password
        if not user or not verify_password(password, user.password_hash):
            raise HTTPException(status_code=401, detail="Invalid credentials")
        
        # Check if MFA is enabled
        if user.mfa_enabled:
            # Generate a temporary login token that will be used to complete login
            temp_token_data = {
                'user_id': user.id,
                'exp': datetime.utcnow() + timedelta(minutes=10),
                'type': 'temp_login'
            }
            temp_token = jwt.encode(temp_token_data, JWT_SECRET_KEY, algorithm='HS256')
            
            return {
                "success": True,
                "require_mfa": True,
                "temp_token": temp_token,
                "message": "Please enter your MFA code to complete login"
            }
        else:
            # Regular login without MFA
            session_token = generate_session_token(user)
            return {"success": True, "session_token": session_token}
    except HTTPException:
        raise
    except Exception as error:
        raise HTTPException(status_code=500, detail="Login failed")

# Second step for MFA verification
@router.post('/login/verify')
async def login_verify_mfa(
    temp_token: str, 
    mfa_code: str, 
    db: AsyncSession = Depends(get_db)
):
    try:
        # Verify the temporary token
        try:
            decoded = jwt.decode(temp_token, JWT_SECRET_KEY, algorithms=['HS256'])
            if decoded.get('type') != 'temp_login':
                raise HTTPException(status_code=401, detail="Invalid session")
            user_id = decoded['user_id']
        except jwt.InvalidTokenError:
            raise HTTPException(status_code=401, detail="Invalid session")
        
        user = await get_user_by_id(user_id, db)
        if not user or not user.mfa_secret:
            raise HTTPException(status_code=400, detail="User not found")
        
        # Decrypt the secret before verification
        decrypted_secret = decrypt_secret(user.mfa_secret)
        
        # Convert stored hashed backup codes to set
        hashed_backup_codes = set()
        if user.backup_codes:
            backup_codes = json.loads(user.backup_codes)
            hashed_backup_codes = {code['hash'] for code in backup_codes if not code.get('used')}
        
        # Verify using both token and backup codes
        totp = pyotp.TOTP(decrypted_secret)
        is_valid = totp.verify(mfa_code)
        
        # Check backup codes if TOTP fails
        if not is_valid and hashed_backup_codes:
            hashed_input = hash_backup_code(mfa_code)
            if hashed_input in hashed_backup_codes:
                # Mark backup code as used
                backup_codes = json.loads(user.backup_codes)
                for code in backup_codes:
                    if code['hash'] == hashed_input:
                        code['used'] = True
                        break
                user.backup_codes = json.dumps(backup_codes)
                await db.commit()
                is_valid = True
        
        if is_valid:
            # Complete login
            session_token = generate_session_token(user)
            return {"success": True, "session_token": session_token}
        else:
            raise HTTPException(status_code=401, detail="Invalid MFA code")
    except HTTPException:
        raise
    except Exception as error:
        raise HTTPException(status_code=500, detail="Verification failed")

This two-step process ensures:

Users with MFA enabled must provide a valid MFA code
Temporary tokens control the flow between login steps
Both TOTP and backup codes are supported
Used backup codes are invalidated
Secrets are decrypted before use, and backup codes are always hashed
(In production) Add rate-limiting middleware to these endpoints

Introduction

Secure MFA Integration Patterns

The Vulnerable Code

Join the 1M+ learners on CodeSignal

Secure User Model Modifications

Implementing MFA Setup

Verifying MFA Tokens Securely

Implementing Login with MFA

Conclusion and Next Steps