| 1 | """TOTP 2FA routes — setup, verify, validate at login, disable, regenerate backup codes.""" |
| 2 | |
| 3 | import os |
| 4 | from datetime import datetime |
| 5 | from datetime import timedelta |
| 6 | |
| 7 | import jwt |
| 8 | from fastapi import APIRouter |
| 9 | from fastapi import Depends |
| 10 | from fastapi import HTTPException |
| 11 | from fastapi import Security |
| 12 | from loguru import logger |
| 13 | |
| 14 | from app.auth.models.totp import TOTPBackupCodesResponse |
| 15 | from app.auth.models.totp import TOTPDisableRequest |
| 16 | from app.auth.models.totp import TOTPSetupResponse |
| 17 | from app.auth.models.totp import TOTPStatusResponse |
| 18 | from app.auth.models.totp import TOTPValidateRequest |
| 19 | from app.auth.models.totp import TOTPVerifyRequest |
| 20 | from app.auth.services.totp import disable_totp |
| 21 | from app.auth.services.totp import is_2fa_enabled |
| 22 | from app.auth.services.totp import regenerate_backup_codes |
| 23 | from app.auth.services.totp import setup_totp |
| 24 | from app.auth.services.totp import validate_totp |
| 25 | from app.auth.services.totp import verify_setup |
| 26 | from app.auth.services.universal import find_user |
| 27 | from app.auth.utils import AuthHandler |
| 28 | |
| 29 | ACCESS_TOKEN_EXPIRE_MINUTES = int(os.environ.get("ACCESS_TOKEN_EXPIRE_MINUTES", "1440")) |
| 30 | |
| 31 | totp_router = APIRouter() |
| 32 | auth_handler = AuthHandler() |
| 33 | |
| 34 | _jwt_secret = auth_handler.secret |
| 35 | |
| 36 | |
| 37 | # ── Helper: temp token for 2FA pending state ───────────────────────────────── |
| 38 | |
| 39 | |
| 40 | def _create_temp_token(username: str) -> str: |
| 41 | """Create a short-lived JWT for 2FA verification step (5 min).""" |
| 42 | payload = { |
| 43 | "sub": username, |
| 44 | "exp": datetime.utcnow() + timedelta(minutes=5), |
| 45 | "iat": datetime.utcnow(), |
| 46 | "type": "2fa_pending", |
| 47 | } |
| 48 | return jwt.encode(payload, _jwt_secret, algorithm="HS256") |
| 49 | |
| 50 | |
| 51 | def _decode_temp_token(token: str) -> str: |
| 52 | """Decode a 2FA temp token and return the username. Raises ValueError on failure.""" |
| 53 | try: |
| 54 | payload = jwt.decode(token, _jwt_secret, algorithms=["HS256"]) |
| 55 | if payload.get("type") != "2fa_pending": |
| 56 | raise ValueError("Invalid token type") |
| 57 | username = payload.get("sub") |
| 58 | if not username: |
| 59 | raise ValueError("No username in token") |
| 60 | return username |
| 61 | except jwt.ExpiredSignatureError: |
| 62 | raise ValueError("2FA verification token has expired. Please log in again.") |
| 63 | except jwt.InvalidTokenError: |
| 64 | raise ValueError("Invalid 2FA verification token.") |
| 65 | |
| 66 | |
| 67 | # ── Status ─────────────────────────────────────────────────────────────────── |
| 68 | |
| 69 | |
| 70 | @totp_router.get( |
| 71 | "/2fa/status", |
| 72 | response_model=TOTPStatusResponse, |
| 73 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 74 | ) |
| 75 | async def get_2fa_status( |
| 76 | token: str = Depends(AuthHandler().security), |
| 77 | ): |
| 78 | """Check if 2FA is enabled for the current user.""" |
| 79 | username, _ = auth_handler.decode_token(token) |
| 80 | user = await find_user(username) |
| 81 | if not user: |
| 82 | raise HTTPException(status_code=401, detail="User not found") |
| 83 | |
| 84 | enabled = await is_2fa_enabled(user.id) |
| 85 | return TOTPStatusResponse(enabled=enabled) |
| 86 | |
| 87 | |
| 88 | # ── Setup ──────────────────────────────────────────────────────────────────── |
| 89 | |
| 90 | |
| 91 | @totp_router.post( |
| 92 | "/2fa/setup", |
| 93 | response_model=TOTPSetupResponse, |
| 94 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 95 | ) |
| 96 | async def start_2fa_setup( |
| 97 | token: str = Depends(AuthHandler().security), |
| 98 | ): |
| 99 | """Generate a new TOTP secret, QR code, and backup codes. Does NOT activate until verified.""" |
| 100 | username, _ = auth_handler.decode_token(token) |
| 101 | user = await find_user(username) |
| 102 | if not user: |
| 103 | raise HTTPException(status_code=401, detail="User not found") |
| 104 | |
| 105 | try: |
| 106 | result = await setup_totp(user.id, user.username) |
| 107 | except ValueError as e: |
| 108 | raise HTTPException(status_code=400, detail=str(e)) |
| 109 | |
| 110 | return TOTPSetupResponse(**result) |
| 111 | |
| 112 | |
| 113 | # ── Verify setup (activate) ───────────────────────────────────────────────── |
| 114 | |
| 115 | |
| 116 | @totp_router.post( |
| 117 | "/2fa/verify-setup", |
| 118 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 119 | ) |
| 120 | async def verify_2fa_setup( |
| 121 | body: TOTPVerifyRequest, |
| 122 | token: str = Depends(AuthHandler().security), |
| 123 | ): |
| 124 | """Verify a TOTP code to activate 2FA.""" |
| 125 | username, _ = auth_handler.decode_token(token) |
| 126 | user = await find_user(username) |
| 127 | if not user: |
| 128 | raise HTTPException(status_code=401, detail="User not found") |
| 129 | |
| 130 | try: |
| 131 | await verify_setup(user.id, body.code) |
| 132 | except ValueError as e: |
| 133 | raise HTTPException(status_code=400, detail=str(e)) |
| 134 | |
| 135 | return {"message": "Two-factor authentication is now enabled.", "success": True} |
| 136 | |
| 137 | |
| 138 | # ── Disable ────────────────────────────────────────────────────────────────── |
| 139 | |
| 140 | |
| 141 | @totp_router.delete( |
| 142 | "/2fa/disable", |
| 143 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 144 | ) |
| 145 | async def disable_2fa( |
| 146 | body: TOTPDisableRequest, |
| 147 | token: str = Depends(AuthHandler().security), |
| 148 | ): |
| 149 | """Disable 2FA. Requires a valid TOTP code or backup code.""" |
| 150 | username, _ = auth_handler.decode_token(token) |
| 151 | user = await find_user(username) |
| 152 | if not user: |
| 153 | raise HTTPException(status_code=401, detail="User not found") |
| 154 | |
| 155 | if not body.code and not body.backup_code: |
| 156 | raise HTTPException(status_code=400, detail="Provide a TOTP code or backup code.") |
| 157 | |
| 158 | try: |
| 159 | await disable_totp(user.id, code=body.code, backup_code=body.backup_code) |
| 160 | except ValueError as e: |
| 161 | raise HTTPException(status_code=400, detail=str(e)) |
| 162 | |
| 163 | return {"message": "Two-factor authentication has been disabled.", "success": True} |
| 164 | |
| 165 | |
| 166 | # ── Validate at login ──────────────────────────────────────────────────────── |
| 167 | |
| 168 | |
| 169 | @totp_router.post("/2fa/validate") |
| 170 | async def validate_2fa_login(body: TOTPValidateRequest): |
| 171 | """ |
| 172 | Validate a TOTP code or backup code during login. |
| 173 | Accepts the temp_token issued by /auth/token when 2FA is required. |
| 174 | Returns a full access token on success. |
| 175 | """ |
| 176 | try: |
| 177 | username = _decode_temp_token(body.temp_token) |
| 178 | except ValueError as e: |
| 179 | raise HTTPException(status_code=401, detail=str(e)) |
| 180 | |
| 181 | user = await find_user(username) |
| 182 | if not user: |
| 183 | raise HTTPException(status_code=401, detail="User not found") |
| 184 | |
| 185 | if not body.code and not body.backup_code: |
| 186 | raise HTTPException(status_code=400, detail="Provide a TOTP code or backup code.") |
| 187 | |
| 188 | try: |
| 189 | await validate_totp(user.id, code=body.code, backup_code=body.backup_code) |
| 190 | except ValueError as e: |
| 191 | raise HTTPException(status_code=401, detail=str(e)) |
| 192 | |
| 193 | # Issue full access token |
| 194 | access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES) |
| 195 | access_token = await auth_handler.encode_token(user.username, access_token_expires) |
| 196 | logger.info(f"2FA login completed for {user.username}") |
| 197 | |
| 198 | return {"access_token": access_token, "token_type": "bearer"} |
| 199 | |
| 200 | |
| 201 | # ── Regenerate backup codes ────────────────────────────────────────────────── |
| 202 | |
| 203 | |
| 204 | @totp_router.post( |
| 205 | "/2fa/backup-codes/regenerate", |
| 206 | response_model=TOTPBackupCodesResponse, |
| 207 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 208 | ) |
| 209 | async def regenerate_2fa_backup_codes( |
| 210 | body: TOTPVerifyRequest, |
| 211 | token: str = Depends(AuthHandler().security), |
| 212 | ): |
| 213 | """Regenerate backup codes. Requires a valid TOTP code. Old codes are invalidated.""" |
| 214 | username, _ = auth_handler.decode_token(token) |
| 215 | user = await find_user(username) |
| 216 | if not user: |
| 217 | raise HTTPException(status_code=401, detail="User not found") |
| 218 | |
| 219 | try: |
| 220 | codes = await regenerate_backup_codes(user.id, body.code) |
| 221 | except ValueError as e: |
| 222 | raise HTTPException(status_code=400, detail=str(e)) |
| 223 | |
| 224 | return TOTPBackupCodesResponse(backup_codes=codes) |