flowdeck/app/services/two_factor.py

"""FlowDeck — TOTP 2FA + backup codes (v7.2.0).

Secrets are Fernet-encrypted at rest (same construction as SSO secrets).
Login flow: ``POST /auth/local-login`` returns ``2fa_required`` + a short-lived
signed ``pending`` token; ``POST /auth/local-verify`` exchanges it for a
session. See ``docs/V72_Enterprise_SCIM_2FA.md``.
"""
from __future__ import annotations

import hashlib
import json
import secrets

from app.db import get_conn

BACKUP_CODE_COUNT = 10


def _fernet():
    import base64

    from cryptography.fernet import Fernet

    from app.config import settings
    key = hashlib.sha256((settings.app_secret_key or "flowdeck").encode()).digest()
    return Fernet(base64.urlsafe_b64encode(key))


def is_enabled(user_id: int) -> bool:
    with get_conn() as conn:
        row = conn.execute("SELECT totp_secret_enc FROM users WHERE id=?",
                           (user_id,)).fetchone()
    if not row or not row["totp_secret_enc"]:
        return False
    try:
        return bool(_fernet().decrypt(row["totp_secret_enc"].encode()).decode())
    except Exception:  # noqa: BLE001
        return False


def setup_secret(user_id: int) -> dict:
    """Create a new TOTP secret (not yet active until verified)."""
    import pyotp
    secret = pyotp.random_base32()
    with get_conn() as conn:
        row = conn.execute("SELECT login, email FROM users WHERE id=?", (user_id,)).fetchone()
    label = (row["email"] or row["login"]) if row else f"user{user_id}"
    uri = pyotp.totp.TOTP(secret).provisioning_uri(name=label, issuer_name="FlowDeck")
    return {"secret": secret, "otpauth_url": uri}


def activate_secret(user_id: int, secret: str, code: str) -> list[str]:
    """Verify ``code`` against ``secret``; on success store + return backup codes."""
    import pyotp
    if not pyotp.TOTP(secret).verify(code, valid_window=1):
        raise ValueError("invalid code")
    codes = [secrets.token_hex(4) for _ in range(BACKUP_CODE_COUNT)]
    hashes = [hashlib.sha256(c.encode()).hexdigest() for c in codes]
    with get_conn() as conn:
        conn.execute("UPDATE users SET totp_secret_enc=?, totp_backup_hashes=? WHERE id=?",
                     (_fernet().encrypt(secret.encode()).decode(),
                      json.dumps(hashes), user_id))
        conn.commit()
    return codes


def verify_code(user_id: int, code: str) -> bool:
    """Check a TOTP code or consume a backup code."""
    code = (code or "").strip().replace(" ", "")
    if not code:
        return False
    with get_conn() as conn:
        row = conn.execute("SELECT totp_secret_enc, totp_backup_hashes FROM users WHERE id=?",
                           (user_id,)).fetchone()
    if not row or not row["totp_secret_enc"]:
        return False
    try:
        secret = _fernet().decrypt(row["totp_secret_enc"].encode()).decode()
    except Exception:  # noqa: BLE001
        return False
    import pyotp
    if secret and pyotp.TOTP(secret).verify(code, valid_window=1):
        return True
    # backup codes (single use)
    try:
        hashes = json.loads(row["totp_backup_hashes"] or "[]")
    except (TypeError, json.JSONDecodeError):
        hashes = []
    digest = hashlib.sha256(code.encode()).hexdigest()
    if digest in hashes:
        hashes.remove(digest)
        with get_conn() as conn:
            conn.execute("UPDATE users SET totp_backup_hashes=? WHERE id=?",
                         (json.dumps(hashes), user_id))
            conn.commit()
        return True
    return False


def disable(user_id: int) -> None:
    with get_conn() as conn:
        conn.execute("UPDATE users SET totp_secret_enc='', totp_backup_hashes='[]'"
                     " WHERE id=?", (user_id,))
        conn.commit()


def remaining_backup_codes(user_id: int) -> int:
    with get_conn() as conn:
        row = conn.execute("SELECT totp_backup_hashes FROM users WHERE id=?",
                           (user_id,)).fetchone()
    try:
        return len(json.loads(row["totp_backup_hashes"] or "[]")) if row else 0
    except (TypeError, json.JSONDecodeError):
        return 0


# ── pending 2FA challenge (signed, 5 min) ──────────────────────────────────

def mint_pending(user_id: int) -> str:
    from itsdangerous import URLSafeTimedSerializer

    from app.config import settings
    ser = URLSafeTimedSerializer(settings.app_secret_key, salt="totp-pending")
    return ser.dumps({"user_id": user_id})


def redeem_pending(token: str, max_age: int = 300) -> int | None:
    from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer

    from app.config import settings
    ser = URLSafeTimedSerializer(settings.app_secret_key, salt="totp-pending")
    try:
        payload = ser.loads(token, max_age=max_age)
        return int(payload.get("user_id", 0)) or None
    except (BadSignature, SignatureExpired, ValueError):
        return None