Unzip budget-app v7 into repo root

budget-app/backend/app/mfa.py

@@ -0,0 +1,121 @@
+"""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