"""TOTP service: secret management, recovery-code lifecycle, setup/enable/disable (M4-T05). State-machine summary --------------------- DISABLED (default) totp_secret=None, totp_enabled=False, no recovery codes PENDING (after setup, before enable) totp_secret=, totp_enabled=False, recovery codes stored as hashes - Re-calling setup replaces the pending secret and regenerates recovery codes. ENABLED (after enable) totp_secret=, totp_enabled=True, recovery codes stored as hashes After disable: totp_secret=None, totp_enabled=False, recovery codes deleted → back to DISABLED Recovery-code timing -------------------- 1. setup → generate 10 plaintext codes, persist their Argon2 hashes immediately, return plaintext to caller (ONLY time). 2. enable → verify 6-digit TOTP code; if ok, flip totp_enabled=True. Recovery codes are already in the DB from step 1. 3. disable → clear secret, clear enabled flag, delete all recovery codes. Security note: ``totp_secret`` is stored as plaintext (same posture as other secrets in this project: protected by file-system permissions on the SQLite database). Recovery codes are stored as Argon2 hashes only. """ from __future__ import annotations import logging import secrets import string import pyotp from sqlalchemy import delete, select from sqlalchemy.orm import Session from app.models.auth import AuthUser, RecoveryCode from app.services.auth import hash_password, verify_password logger = logging.getLogger(__name__) # Number of recovery codes to generate per setup _RECOVERY_CODE_COUNT = 10 # Alphabet for each 4-character segment: lowercase letters + digits, no ambiguous chars _CODE_ALPHABET = string.ascii_lowercase + string.digits # --------------------------------------------------------------------------- # Internal helpers # --------------------------------------------------------------------------- def _generate_recovery_code() -> str: """Return a single recovery code in ``xxxx-xxxx`` format.""" segment = lambda: "".join(secrets.choice(_CODE_ALPHABET) for _ in range(4)) # noqa: E731 return f"{segment()}-{segment()}" def _verify_totp_code(secret: str, code: str) -> bool: """Verify a 6-digit TOTP code against the given base-32 secret (±1 window).""" return pyotp.TOTP(secret).verify(code, valid_window=1) def _delete_recovery_codes(db: Session, *, user_id: int) -> None: """Delete ALL recovery codes for a user (called on setup re-run and disable).""" db.execute(delete(RecoveryCode).where(RecoveryCode.user_id == user_id)) # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- def setup( db: Session, *, user: AuthUser, issuer: str, ) -> tuple[str, str, list[str]]: """Generate a new pending TOTP secret and recovery codes. The secret is stored in ``user.totp_secret`` (``totp_enabled`` stays ``False``). Any pre-existing pending secret and recovery codes are replaced atomically. Returns ------- (secret, otpauth_uri, plaintext_recovery_codes) ``plaintext_recovery_codes`` are returned **once** here and MUST NOT be stored or returned elsewhere. """ # Generate new TOTP secret new_secret = pyotp.random_base32() otpauth_uri = pyotp.TOTP(new_secret).provisioning_uri( name=user.username, issuer_name=issuer, ) # Generate plaintext recovery codes plaintext_codes = [_generate_recovery_code() for _ in range(_RECOVERY_CODE_COUNT)] # --- Persist atomically --- # 1. Delete any old pending recovery codes (idempotent on re-setup) assert user.id is not None # mypy guard; always set after DB insertion _delete_recovery_codes(db, user_id=user.id) # 2. Update the secret (pending — totp_enabled stays False) user.totp_secret = new_secret # 3. Persist new recovery codes as Argon2 hashes for plaintext in plaintext_codes: db.add(RecoveryCode(user_id=user.id, code_hash=hash_password(plaintext))) db.commit() db.refresh(user) logger.info("TOTP setup (pending) for user '%s'; %d recovery codes generated.", user.username, _RECOVERY_CODE_COUNT) return new_secret, otpauth_uri, plaintext_codes def enable(db: Session, *, user: AuthUser, code: str) -> bool: """Enable TOTP after the user proves they can generate the correct code. Parameters ---------- code: The current 6-digit TOTP code from the user's authenticator app. Returns ------- True on success; False if the code is invalid or no secret is pending. """ if not user.totp_secret: logger.info("TOTP enable rejected for '%s': no pending secret.", user.username) return False if not _verify_totp_code(user.totp_secret, code): logger.info("TOTP enable rejected for '%s': wrong code.", user.username) return False user.totp_enabled = True db.commit() db.refresh(user) logger.info("TOTP enabled for user '%s'.", user.username) return True def disable( db: Session, *, user: AuthUser, password: str | None = None, code: str | None = None, ) -> bool: """Disable TOTP. Caller must supply exactly one of ``password`` or ``code``. On success: ``totp_enabled=False``, ``totp_secret=None``, all recovery codes deleted. Returns ------- True on success; False if neither credential is valid. """ if not password and not code: logger.info("TOTP disable rejected for '%s': no credential provided.", user.username) return False if password: if not verify_password(password, user.password_hash): logger.info("TOTP disable rejected for '%s': wrong password.", user.username) return False elif code: if not user.totp_secret or not _verify_totp_code(user.totp_secret, code): logger.info("TOTP disable rejected for '%s': wrong TOTP code.", user.username) return False # Clear TOTP state assert user.id is not None _delete_recovery_codes(db, user_id=user.id) user.totp_enabled = False user.totp_secret = None db.commit() db.refresh(user) logger.info("TOTP disabled for user '%s'.", user.username) return True def verify_totp_code(user: AuthUser, code: str) -> bool: """Verify a 6-digit TOTP code for a user. Returns True if the code is valid (±1 time window), False otherwise. The user must have a ``totp_secret`` set; returns False if not. This is the public entry-point for T06 login two-factor verification. """ if not user.totp_secret: return False return _verify_totp_code(user.totp_secret, code) def verify_recovery_code(db: Session, *, user: AuthUser, code: str) -> bool: """Verify and consume a one-time recovery code. Finds the first unused recovery code whose hash matches ``code``, marks it as consumed (sets ``used_at``), and returns ``True``. Returns ``False`` if no matching unused code exists. This function is provided for T06 (login two-factor) but lives here so the TOTP service owns all recovery-code logic. """ from datetime import UTC, datetime unused = db.execute( select(RecoveryCode).where( RecoveryCode.user_id == user.id, RecoveryCode.used_at.is_(None), ) ).scalars().all() for rc in unused: if verify_password(code, rc.code_hash): rc.used_at = datetime.now(UTC) db.commit() logger.info( "Recovery code consumed for user '%s' (id=%d).", user.username, rc.id ) return True return False