Budget_Tool_Martinhal/backend/app/mfa.py

"""MFA with an authenticator app (TOTP, RFC 6238: 6 digits, 30 seconds, SHA-1).

Works with Microsoft Authenticator, Google Authenticator, Authy, 1Password and similar apps.
"""
import base64
import hashlib
import hmac
import secrets
import struct
import time
from datetime import datetime, timedelta, timezone
from urllib.parse import quote

import segno
from cryptography.fernet import Fernet, InvalidToken

from .config import settings
from .models import User

ISSUER = "Budget & Capex"
STEP = 30
WINDOW = 1  # also accept the code just before and just after, for clocks that are slightly off
MAX_FAILED = 5
LOCK_MINUTES = 5
RECOVERY_CODES = 8


def _fernet() -> Fernet:
    key = hashlib.sha256(("mfa-secret:" + settings.secret_key).encode()).digest()
    return Fernet(base64.urlsafe_b64encode(key))


def new_secret() -> str:
    return base64.b32encode(secrets.token_bytes(20)).decode().rstrip("=")


def seal(secret: str) -> str:
    return _fernet().encrypt(secret.encode()).decode()


def unseal(token: str | None) -> str | None:
    try:
        return _fernet().decrypt(token.encode()).decode() if token else None
    except InvalidToken:  # SECRET_KEY changed: an administrator has to reset this user's MFA
        return None


def code_at(secret: str, step: int) -> str:
    key = base64.b32decode(secret + "=" * (-len(secret) % 8))
    digest = hmac.new(key, struct.pack(">Q", step), hashlib.sha1).digest()
    offset = digest[-1] & 0x0F
    value = struct.unpack(">I", digest[offset:offset + 4])[0] & 0x7FFFFFFF
    return f"{value % 1_000_000:06d}"


def matching_step(secret: str, code: str, last_step: int | None = None, at: float | None = None) -> int | None:
    """The time step the code belongs to, or None. A step already used is refused (no replay)."""
    code = "".join(ch for ch in code if ch.isdigit())
    if len(code) != 6:
        return None
    now = int((at or time.time()) // STEP)
    for step in range(now - WINDOW, now + WINDOW + 1):
        if hmac.compare_digest(code_at(secret, step), code) and (last_step is None or step > last_step):
            return step
    return None


def otpauth_uri(secret: str, email: str) -> str:
    label = quote(f"{ISSUER}:{email}")
    return f"otpauth://totp/{label}?secret={secret}&issuer={quote(ISSUER)}&algorithm=SHA1&digits=6&period={STEP}"


def qr_svg(uri: str) -> str:
    return segno.make(uri, error="m").svg_inline(scale=5, dark="#000", light="#fff", border=2)


def _hash(code: str) -> str:
    return hashlib.sha256(code.replace("-", "").strip().lower().encode()).hexdigest()


def new_recovery_codes() -> tuple[list[str], list[str]]:
    """Codes to show once, and their hashes to keep."""
    alphabet = "abcdefghjkmnpqrstuvwxyz23456789"
    codes = ["".join(secrets.choice(alphabet) for _ in range(4)) + "-" + "".join(secrets.choice(alphabet) for _ in range(4))
             for _ in range(RECOVERY_CODES)]
    return codes, [_hash(c) for c in codes]


def locked(user: User) -> bool:
    until = user.mfa_locked_until
    if until and until.tzinfo is None:
        until = until.replace(tzinfo=timezone.utc)
    return bool(until and until > datetime.now(timezone.utc))


def check(user: User, code: str) -> str | None:
    """Checks an authenticator code or a recovery code. Returns 'totp', 'recovery' or None, and updates
    the user's counters (a recovery code is used up; too many wrong codes lock MFA for a few minutes)."""
    if locked(user):
        return None
    code = (code or "").strip()
    secret = unseal(user.mfa_secret)
    kind = None
    if secret and (step := matching_step(secret, code, user.mfa_last_step)) is not None:
        user.mfa_last_step, kind = step, "totp"
    elif any(ch.isalpha() for ch in code) and _hash(code) in (user.mfa_recovery or []):
        user.mfa_recovery = [h for h in user.mfa_recovery if h != _hash(code)]
        kind = "recovery"
    if kind:
        user.mfa_failed, user.mfa_locked_until = 0, None
    else:
        user.mfa_failed = (user.mfa_failed or 0) + 1
        if user.mfa_failed >= MAX_FAILED:
            user.mfa_failed, user.mfa_locked_until = 0, datetime.now(timezone.utc) + timedelta(minutes=LOCK_MINUTES)
    return kind


def disable(user: User) -> None:
    user.mfa_enabled, user.mfa_secret, user.mfa_pending_secret = False, None, None
    user.mfa_enabled_at, user.mfa_last_step, user.mfa_recovery = None, None, None
    user.mfa_failed, user.mfa_locked_until = 0, None