muslin_connect_hub_backend/utils/totp.py

import base64
import hashlib
import hmac
import io
import secrets
import struct
import time
from typing import Optional
from urllib.parse import quote

try:
    import pyotp  # type: ignore
except ImportError:
    pyotp = None

try:
    import qrcode  # type: ignore
    import qrcode.image.svg  # type: ignore
except ImportError:
    qrcode = None


BASE32_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"


def generate_totp_secret(length: int = 32) -> str:
    """Generate a cryptographically secure RFC 4648 Base32 secret for TOTP."""
    if pyotp is not None:
        return pyotp.random_base32(length=length)
    return "".join(secrets.choice(BASE32_ALPHABET) for _ in range(length))


def _decode_base32_secret(secret: str) -> bytes:
    cleaned = "".join(secret.strip().split()).upper()
    padding = (8 - len(cleaned) % 8) % 8
    cleaned += "=" * padding
    return base64.b32decode(cleaned, casefold=True)


def get_totp_code(
    secret: str,
    for_time: Optional[int] = None,
    interval: int = 30,
    digits: int = 6,
) -> str:
    """Compute the 6-digit RFC 6238 TOTP code for a given timestamp."""
    if not secret:
        return ""
    ts = int(for_time if for_time is not None else time.time())
    if pyotp is not None:
        try:
            totp = pyotp.TOTP(secret, digits=digits, interval=interval)
            return totp.at(ts)
        except Exception:
            pass

    key = _decode_base32_secret(secret)
    counter = int(ts // interval)
    msg = struct.pack(">Q", counter)
    digest = hmac.new(key, msg, hashlib.sha1).digest()
    offset = digest[-1] & 0x0F
    binary = struct.unpack(">I", digest[offset : offset + 4])[0] & 0x7FFFFFFF
    otp = binary % (10 ** digits)
    return f"{otp:0{digits}d}"


def verify_totp_code(
    secret: str,
    code: str,
    valid_window: int = 1,
    interval: int = 30,
    digits: int = 6,
) -> bool:
    """Verify a TOTP code within the allowed time-step window (RFC 6238)."""
    if not secret or not code:
        return False
    normalized = "".join(str(code).strip().split()).replace("-", "")
    if not normalized.isdigit() or len(normalized) != digits:
        return False

    if pyotp is not None:
        try:
            totp = pyotp.TOTP(secret, digits=digits, interval=interval)
            return bool(totp.verify(normalized, valid_window=valid_window))
        except Exception:
            pass

    now = int(time.time())
    for offset in range(-valid_window, valid_window + 1):
        candidate = get_totp_code(
            secret=secret,
            for_time=now + (offset * interval),
            interval=interval,
            digits=digits,
        )
        if candidate and hmac.compare_digest(candidate, normalized):
            return True
    return False


def build_otpauth_uri(
    secret: str,
    account_name: str,
    issuer: str = "Muslim Connect Hub",
) -> str:
    """Build a standard otpauth:// URI for Google/Microsoft Authenticator."""
    if pyotp is not None:
        try:
            return pyotp.TOTP(secret).provisioning_uri(
                name=account_name,
                issuer_name=issuer,
            )
        except Exception:
            pass
    label = f"{quote(issuer)}:{quote(account_name)}"
    issuer_param = quote(issuer)
    return (
        f"otpauth://totp/{label}"
        f"?secret={secret}&issuer={issuer_param}&algorithm=SHA1&digits=6&period=30"
    )


def generate_qr_data_url(otpauth_uri: str) -> str:
    """
    Generate a base64 Data URL (PNG or SVG) for the given otpauth:// URI.
    """
    if qrcode is not None:
        try:
            qr = qrcode.QRCode(
                version=1,
                error_correction=qrcode.constants.ERROR_CORRECT_M,
                box_size=8,
                border=2,
            )
            qr.add_data(otpauth_uri)
            qr.make(fit=True)
            img = qr.make_image(fill_color="black", back_color="white")
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            b64 = base64.b64encode(buf.getvalue()).decode("ascii")
            return f"data:image/png;base64,{b64}"
        except Exception:
            try:
                img = qrcode.make(otpauth_uri, image_factory=qrcode.image.svg.SvgImage)
                buf = io.BytesIO()
                img.save(buf)
                b64 = base64.b64encode(buf.getvalue()).decode("ascii")
                return f"data:image/svg+xml;base64,{b64}"
            except Exception:
                pass

    # Fallback SVG placeholder if qrcode package is unavailable
    svg = (
        '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 200 200" width="200" height="200">'
        '<rect width="200" height="200" fill="#ffffff" stroke="#e2e8f0" stroke-width="4"/>'
        '<rect x="20" y="20" width="50" height="50" fill="none" stroke="#0f172a" stroke-width="8"/>'
        '<rect x="35" y="35" width="20" height="20" fill="#0f172a"/>'
        '<rect x="130" y="20" width="50" height="50" fill="none" stroke="#0f172a" stroke-width="8"/>'
        '<rect x="145" y="35" width="20" height="20" fill="#0f172a"/>'
        '<rect x="20" y="130" width="50" height="50" fill="none" stroke="#0f172a" stroke-width="8"/>'
        '<rect x="35" y="145" width="20" height="20" fill="#0f172a"/>'
        '<rect x="85" y="85" width="30" height="30" fill="#0f172a"/>'
        '</svg>'
    )
    b64 = base64.b64encode(svg.encode("utf-8")).decode("ascii")
    return f"data:image/svg+xml;base64,{b64}"