app/core/access/totp.py

"""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