""" auth.py — Password hashing, session cookie management, and TOTP helpers for multi-user auth. Session cookie format: base64url(json_payload) + "." + hmac_sha256(base64url, secret)[:32] Payload: {"uid": "...", "un": "...", "role": "...", "iat": epoch} """ from __future__ import annotations import base64 import hashlib import hmac import json import time from dataclasses import dataclass from io import BytesIO import pyotp import qrcode from argon2 import PasswordHasher from argon2.exceptions import InvalidHashError, VerificationError, VerifyMismatchError _ph = PasswordHasher() _COOKIE_SEP = "." # ── Password hashing ────────────────────────────────────────────────────────── def hash_password(password: str) -> str: return _ph.hash(password) def verify_password(password: str, hash: str) -> bool: try: return _ph.verify(hash, password) except (VerifyMismatchError, VerificationError, InvalidHashError): return False # ── User dataclass ──────────────────────────────────────────────────────────── @dataclass class CurrentUser: id: str username: str role: str # 'admin' | 'user' is_active: bool = True @property def is_admin(self) -> bool: return self.role == "admin" # Synthetic admin user for API key auth — no DB lookup needed SYNTHETIC_API_ADMIN = CurrentUser( id="api-key-admin", username="api-key", role="admin", ) # ── Session cookie ──────────────────────────────────────────────────────────── def create_session_cookie(user: dict, secret: str) -> str: payload = json.dumps( {"uid": user["id"], "un": user["username"], "role": user["role"], "iat": int(time.time())}, separators=(",", ":"), ) b64 = base64.urlsafe_b64encode(payload.encode()).rstrip(b"=").decode() sig = hmac.new(secret.encode(), b64.encode(), hashlib.sha256).hexdigest()[:32] return f"{b64}{_COOKIE_SEP}{sig}" def decode_session_cookie(cookie: str, secret: str) -> CurrentUser | None: try: b64, sig = cookie.rsplit(_COOKIE_SEP, 1) expected = hmac.new(secret.encode(), b64.encode(), hashlib.sha256).hexdigest()[:32] if not hmac.compare_digest(sig, expected): return None padding = 4 - len(b64) % 4 payload = json.loads(base64.urlsafe_b64decode(b64 + "=" * padding).decode()) return CurrentUser(id=payload["uid"], username=payload["un"], role=payload["role"]) except Exception: return None # ── TOTP helpers ────────────────────────────────────────────────────────────── def generate_totp_secret() -> str: return pyotp.random_base32() def verify_totp(secret: str, code: str) -> bool: return pyotp.TOTP(secret).verify(code, valid_window=1) def make_totp_provisioning_uri(secret: str, username: str, issuer: str = "oAI-Web") -> str: return pyotp.TOTP(secret).provisioning_uri(username, issuer_name=issuer) def make_totp_qr_png_b64(provisioning_uri: str) -> str: img = qrcode.make(provisioning_uri) buf = BytesIO() img.save(buf, format="PNG") return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()