M4-T05: add TOTP service and setup/enable/disable/status API

app/services/totp.py

@@ -0,0 +1,221 @@
+"""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=<base32>, totp_enabled=False, recovery codes stored as hashes
+  - Re-calling setup replaces the pending secret and regenerates recovery codes.
+
+ENABLED (after enable)
+  totp_secret=<base32>, 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_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