| 1 | """TOTP 2FA business logic — setup, verification, backup codes, brute-force protection.""" |
| 2 | |
| 3 | import base64 |
| 4 | |
| 5 | # ── Encryption key for TOTP secrets ────────────────────────────────────────── |
| 6 | import hashlib |
| 7 | import io |
| 8 | import os |
| 9 | import secrets |
| 10 | import string |
| 11 | import time |
| 12 | from typing import Optional |
| 13 | |
| 14 | import bcrypt |
| 15 | import pyotp |
| 16 | import qrcode |
| 17 | from cryptography.fernet import Fernet |
| 18 | from loguru import logger |
| 19 | from sqlalchemy.ext.asyncio import AsyncSession |
| 20 | from sqlalchemy.orm.attributes import flag_modified |
| 21 | from sqlmodel import select |
| 22 | |
| 23 | from app.auth.models.totp import UserTOTP |
| 24 | from app.auth.utils import AuthHandler |
| 25 | from app.db.db_session import async_engine |
| 26 | |
| 27 | # Prefer a dedicated TOTP_ENCRYPTION_KEY (proper Fernet key, separate from JWT). |
| 28 | # Falls back to a key derived from JWT_SECRET so existing deployments that |
| 29 | # haven't set TOTP_ENCRYPTION_KEY continue to decrypt stored TOTP secrets. |
| 30 | _totp_enc_key = os.environ.get("TOTP_ENCRYPTION_KEY") |
| 31 | if _totp_enc_key: |
| 32 | # Trim accidental whitespace/newlines that .env editors sometimes append |
| 33 | _fernet_key = _totp_enc_key.strip().encode() |
| 34 | _fernet_key_source = "TOTP_ENCRYPTION_KEY" |
| 35 | else: |
| 36 | _fernet_key = base64.urlsafe_b64encode(hashlib.sha256(AuthHandler.secret.encode()).digest()) |
| 37 | _fernet_key_source = "JWT_SECRET (derived fallback)" |
| 38 | |
| 39 | try: |
| 40 | _fernet = Fernet(_fernet_key) |
| 41 | except (ValueError, TypeError) as e: |
| 42 | # Surface a clear, actionable error instead of the obscure |
| 43 | # `binascii.Error: Incorrect padding` chain. See issue #838. |
| 44 | if _fernet_key_source == "TOTP_ENCRYPTION_KEY": |
| 45 | raise RuntimeError( |
| 46 | "TOTP_ENCRYPTION_KEY is malformed. Expected a 32-byte url-safe base64 key " |
| 47 | "(typically 44 characters ending with '='). Generate a valid one with:\n" |
| 48 | ' python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())"\n' |
| 49 | f"Then paste the value into your .env without surrounding quotes or trailing whitespace. " |
| 50 | f"(underlying error: {e})", |
| 51 | ) from e |
| 52 | # Fallback path failed — would be a bug, not a misconfiguration |
| 53 | raise RuntimeError( |
| 54 | f"Failed to initialise TOTP Fernet from derived JWT_SECRET key. (underlying error: {e})", |
| 55 | ) from e |
| 56 | |
| 57 | # ── Brute-force protection ─────────────────────────────────────────────────── |
| 58 | # Per-user: {user_id: (fail_count, first_fail_time)} |
| 59 | _2fa_attempts: dict[int, tuple[int, float]] = {} |
| 60 | _MAX_ATTEMPTS = 5 |
| 61 | _LOCKOUT_SECONDS = 900 # 15 min |
| 62 | |
| 63 | |
| 64 | def _check_rate_limit(user_id: int) -> None: |
| 65 | """Raise ValueError if user has too many recent 2FA failures.""" |
| 66 | entry = _2fa_attempts.get(user_id) |
| 67 | if entry is None: |
| 68 | return |
| 69 | fail_count, first_fail_time = entry |
| 70 | if time.time() - first_fail_time > _LOCKOUT_SECONDS: |
| 71 | # Window expired, reset |
| 72 | _2fa_attempts.pop(user_id, None) |
| 73 | return |
| 74 | if fail_count >= _MAX_ATTEMPTS: |
| 75 | remaining = int(_LOCKOUT_SECONDS - (time.time() - first_fail_time)) |
| 76 | raise ValueError(f"Too many failed 2FA attempts. Try again in {remaining} seconds.") |
| 77 | |
| 78 | |
| 79 | def _record_failure(user_id: int) -> None: |
| 80 | entry = _2fa_attempts.get(user_id) |
| 81 | if entry is None or time.time() - entry[1] > _LOCKOUT_SECONDS: |
| 82 | _2fa_attempts[user_id] = (1, time.time()) |
| 83 | else: |
| 84 | _2fa_attempts[user_id] = (entry[0] + 1, entry[1]) |
| 85 | |
| 86 | |
| 87 | def _clear_failures(user_id: int) -> None: |
| 88 | _2fa_attempts.pop(user_id, None) |
| 89 | |
| 90 | |
| 91 | # ── Encryption helpers ─────────────────────────────────────────────────────── |
| 92 | |
| 93 | |
| 94 | def _encrypt_secret(secret: str) -> str: |
| 95 | return _fernet.encrypt(secret.encode()).decode() |
| 96 | |
| 97 | |
| 98 | def _decrypt_secret(enc: str) -> str: |
| 99 | return _fernet.decrypt(enc.encode()).decode() |
| 100 | |
| 101 | |
| 102 | # ── Backup codes ───────────────────────────────────────────────────────────── |
| 103 | |
| 104 | _BACKUP_CODE_COUNT = 8 |
| 105 | _BACKUP_CODE_LENGTH = 10 |
| 106 | |
| 107 | |
| 108 | def _generate_backup_codes() -> tuple[list[str], list[dict]]: |
| 109 | """Generate backup codes. Returns (plaintext_list, hashed_list_for_db).""" |
| 110 | alphabet = string.ascii_uppercase + string.digits |
| 111 | codes = [] |
| 112 | hashed = [] |
| 113 | for _ in range(_BACKUP_CODE_COUNT): |
| 114 | code = "".join(secrets.choice(alphabet) for _ in range(_BACKUP_CODE_LENGTH)) |
| 115 | codes.append(code) |
| 116 | hashed_code = bcrypt.hashpw(code.encode("utf-8"), bcrypt.gensalt()).decode("utf-8") |
| 117 | hashed.append({"hash": hashed_code, "used": False}) |
| 118 | return codes, hashed |
| 119 | |
| 120 | |
| 121 | def _verify_backup_code(backup_code: str, stored_codes: list[dict]) -> int: |
| 122 | """Verify a backup code. Returns the index if valid, -1 otherwise.""" |
| 123 | for i, entry in enumerate(stored_codes): |
| 124 | if entry.get("used"): |
| 125 | continue |
| 126 | if bcrypt.checkpw(backup_code.upper().strip().encode("utf-8"), entry["hash"].encode("utf-8")): |
| 127 | return i |
| 128 | return -1 |
| 129 | |
| 130 | |
| 131 | # ── QR code generation ─────────────────────────────────────────────────────── |
| 132 | |
| 133 | |
| 134 | def _generate_qr_data_uri(otpauth_url: str) -> str: |
| 135 | """Generate a QR code as a base64 data URI.""" |
| 136 | img = qrcode.make(otpauth_url, box_size=6, border=2) |
| 137 | buf = io.BytesIO() |
| 138 | img.save(buf, format="PNG") |
| 139 | b64 = base64.b64encode(buf.getvalue()).decode() |
| 140 | return f"data:image/png;base64,{b64}" |
| 141 | |
| 142 | |
| 143 | # ── DB operations ──────────────────────────────────────────────────────────── |
| 144 | |
| 145 | |
| 146 | async def get_user_totp(user_id: int) -> Optional[UserTOTP]: |
| 147 | async with AsyncSession(async_engine) as session: |
| 148 | result = await session.execute(select(UserTOTP).where(UserTOTP.user_id == user_id)) |
| 149 | return result.scalars().first() |
| 150 | |
| 151 | |
| 152 | async def is_2fa_enabled(user_id: int) -> bool: |
| 153 | totp = await get_user_totp(user_id) |
| 154 | return totp is not None and totp.enabled |
| 155 | |
| 156 | |
| 157 | async def setup_totp(user_id: int, username: str) -> dict: |
| 158 | """Generate a new TOTP secret and backup codes. Does NOT activate until verified.""" |
| 159 | secret = pyotp.random_base32() |
| 160 | enc_secret = _encrypt_secret(secret) |
| 161 | plaintext_codes, hashed_codes = _generate_backup_codes() |
| 162 | |
| 163 | totp = pyotp.TOTP(secret) |
| 164 | issuer = "CoPilot" |
| 165 | otpauth_url = totp.provisioning_uri(name=username, issuer_name=issuer) |
| 166 | qr_data_uri = _generate_qr_data_uri(otpauth_url) |
| 167 | |
| 168 | async with AsyncSession(async_engine) as session: |
| 169 | result = await session.execute(select(UserTOTP).where(UserTOTP.user_id == user_id)) |
| 170 | existing = result.scalars().first() |
| 171 | if existing and existing.enabled: |
| 172 | raise ValueError("2FA is already enabled. Disable it first to reconfigure.") |
| 173 | |
| 174 | if existing: |
| 175 | existing.secret_enc = enc_secret |
| 176 | existing.backup_codes = hashed_codes |
| 177 | existing.enabled = False |
| 178 | existing.last_used_at = None |
| 179 | else: |
| 180 | entry = UserTOTP( |
| 181 | user_id=user_id, |
| 182 | secret_enc=enc_secret, |
| 183 | enabled=False, |
| 184 | backup_codes=hashed_codes, |
| 185 | last_used_at=None, |
| 186 | ) |
| 187 | session.add(entry) |
| 188 | |
| 189 | await session.commit() |
| 190 | |
| 191 | logger.info(f"2FA setup initiated for user_id={user_id}") |
| 192 | return { |
| 193 | "secret": secret, |
| 194 | "otpauth_url": otpauth_url, |
| 195 | "qr_data_uri": qr_data_uri, |
| 196 | "backup_codes": plaintext_codes, |
| 197 | } |
| 198 | |
| 199 | |
| 200 | async def verify_setup(user_id: int, code: str) -> bool: |
| 201 | """Verify a TOTP code to activate 2FA.""" |
| 202 | _check_rate_limit(user_id) |
| 203 | |
| 204 | async with AsyncSession(async_engine) as session: |
| 205 | result = await session.execute(select(UserTOTP).where(UserTOTP.user_id == user_id)) |
| 206 | entry = result.scalars().first() |
| 207 | if entry is None: |
| 208 | raise ValueError("No 2FA setup found. Call setup first.") |
| 209 | if entry.enabled: |
| 210 | raise ValueError("2FA is already enabled.") |
| 211 | |
| 212 | secret = _decrypt_secret(entry.secret_enc) |
| 213 | totp = pyotp.TOTP(secret) |
| 214 | |
| 215 | # valid_window=1 → accept ±1 step (±30s drift tolerance) |
| 216 | if not totp.verify(code, valid_window=1): |
| 217 | _record_failure(user_id) |
| 218 | raise ValueError("Invalid verification code. Check your authenticator app and device clock.") |
| 219 | |
| 220 | entry.enabled = True |
| 221 | entry.last_used_at = int(time.time()) // 30 # current TOTP counter |
| 222 | await session.commit() |
| 223 | |
| 224 | _clear_failures(user_id) |
| 225 | logger.info(f"2FA activated for user_id={user_id}") |
| 226 | return True |
| 227 | |
| 228 | |
| 229 | async def validate_totp(user_id: int, code: Optional[str] = None, backup_code: Optional[str] = None) -> bool: |
| 230 | """Validate a TOTP code or backup code during login.""" |
| 231 | _check_rate_limit(user_id) |
| 232 | |
| 233 | if not code and not backup_code: |
| 234 | raise ValueError("Provide either a TOTP code or a backup code.") |
| 235 | |
| 236 | async with AsyncSession(async_engine) as session: |
| 237 | result = await session.execute(select(UserTOTP).where(UserTOTP.user_id == user_id)) |
| 238 | entry = result.scalars().first() |
| 239 | if entry is None or not entry.enabled: |
| 240 | raise ValueError("2FA is not enabled for this user.") |
| 241 | |
| 242 | # Try TOTP code first |
| 243 | if code: |
| 244 | secret = _decrypt_secret(entry.secret_enc) |
| 245 | totp = pyotp.TOTP(secret) |
| 246 | current_counter = int(time.time()) // 30 |
| 247 | |
| 248 | if not totp.verify(code, valid_window=1): |
| 249 | _record_failure(user_id) |
| 250 | raise ValueError("Invalid authentication code.") |
| 251 | |
| 252 | # Replay prevention: reject if same counter as last use |
| 253 | if entry.last_used_at is not None and current_counter <= entry.last_used_at: |
| 254 | _record_failure(user_id) |
| 255 | raise ValueError("This code has already been used. Wait for a new code.") |
| 256 | |
| 257 | entry.last_used_at = current_counter |
| 258 | await session.commit() |
| 259 | _clear_failures(user_id) |
| 260 | return True |
| 261 | |
| 262 | # Try backup code |
| 263 | if backup_code: |
| 264 | codes = list(entry.backup_codes) |
| 265 | idx = _verify_backup_code(backup_code, codes) |
| 266 | if idx < 0: |
| 267 | _record_failure(user_id) |
| 268 | raise ValueError("Invalid backup code.") |
| 269 | |
| 270 | codes[idx]["used"] = True |
| 271 | entry.backup_codes = codes |
| 272 | flag_modified(entry, "backup_codes") |
| 273 | await session.commit() |
| 274 | _clear_failures(user_id) |
| 275 | logger.info(f"2FA backup code used for user_id={user_id} (code index {idx})") |
| 276 | return True |
| 277 | |
| 278 | raise ValueError("Provide either a TOTP code or a backup code.") |
| 279 | |
| 280 | |
| 281 | async def disable_totp(user_id: int, code: Optional[str] = None, backup_code: Optional[str] = None) -> bool: |
| 282 | """Disable 2FA. Requires a valid TOTP code or backup code.""" |
| 283 | # Validate the code first |
| 284 | await validate_totp(user_id, code=code, backup_code=backup_code) |
| 285 | |
| 286 | async with AsyncSession(async_engine) as session: |
| 287 | result = await session.execute(select(UserTOTP).where(UserTOTP.user_id == user_id)) |
| 288 | entry = result.scalars().first() |
| 289 | if entry: |
| 290 | await session.delete(entry) |
| 291 | await session.commit() |
| 292 | |
| 293 | logger.info(f"2FA disabled for user_id={user_id}") |
| 294 | return True |
| 295 | |
| 296 | |
| 297 | async def regenerate_backup_codes(user_id: int, code: str) -> list[str]: |
| 298 | """Regenerate backup codes. Requires a valid TOTP code.""" |
| 299 | _check_rate_limit(user_id) |
| 300 | |
| 301 | async with AsyncSession(async_engine) as session: |
| 302 | result = await session.execute(select(UserTOTP).where(UserTOTP.user_id == user_id)) |
| 303 | entry = result.scalars().first() |
| 304 | if entry is None or not entry.enabled: |
| 305 | raise ValueError("2FA is not enabled.") |
| 306 | |
| 307 | # Verify current TOTP code |
| 308 | secret = _decrypt_secret(entry.secret_enc) |
| 309 | totp = pyotp.TOTP(secret) |
| 310 | if not totp.verify(code, valid_window=1): |
| 311 | _record_failure(user_id) |
| 312 | raise ValueError("Invalid authentication code.") |
| 313 | |
| 314 | plaintext_codes, hashed_codes = _generate_backup_codes() |
| 315 | entry.backup_codes = hashed_codes |
| 316 | await session.commit() |
| 317 | |
| 318 | _clear_failures(user_id) |
| 319 | logger.info(f"2FA backup codes regenerated for user_id={user_id}") |
| 320 | return plaintext_codes |