"""Author RAG Chatbot SaaS — TOTP Two-Factor Authentication. Used exclusively for SuperAdmin accounts. RULE: SuperAdmin login requires TOTP verification after password check. Uses pyotp (RFC 6238 compliant) — works with Google Authenticator. """ import pyotp import qrcode import io import base64 import structlog from app.config import get_settings logger = structlog.get_logger(__name__) cfg = get_settings() def generate_totp_secret() -> str: """Generate a new TOTP secret for a SuperAdmin account. Returns: Base32-encoded secret string (store encrypted in DB). """ return pyotp.random_base32() def get_totp_uri(secret: str, email: str) -> str: """Build the otpauth:// URI for QR code generation. Args: secret: The TOTP secret (Base32 encoded). email: SuperAdmin email address (used as account label). Returns: otpauth:// URI string compatible with authenticator apps. """ totp = pyotp.TOTP(secret) return totp.provisioning_uri(name=email, issuer_name=cfg.SUPERADMIN_2FA_ISSUER) def generate_qr_code_base64(secret: str, email: str) -> str: """Generate a QR code image as a Base64 data URI. Args: secret: The TOTP secret. email: SuperAdmin email for QR label. Returns: Base64-encoded PNG image as data URI string. """ uri = get_totp_uri(secret, email) img = qrcode.make(uri) buffer = io.BytesIO() img.save(buffer, format="PNG") b64 = base64.b64encode(buffer.getvalue()).decode() return f"data:image/png;base64,{b64}" def verify_totp(secret: str, code: str) -> bool: """Verify a TOTP code against the stored secret. Allows 1 time-step window (±30 seconds) to handle clock drift. Args: secret: The stored TOTP secret. code: 6-digit code from the authenticator app. Returns: True if code is valid, False otherwise. """ if not code or len(code) != 6 or not code.isdigit(): return False totp = pyotp.TOTP(secret) is_valid = totp.verify(code, valid_window=1) if not is_valid: logger.warning("Invalid TOTP code attempt") return is_valid