app/core/security.py

"""Author RAG — Core Security Utilities.

JWT creation/verification, TOTP, HMAC widget tokens, bcrypt password hashing.
"""

import hashlib
import hmac
import base64
import json
from datetime import datetime, timedelta, timezone
from typing import Any

import bcrypt
import pyotp
from jose import JWTError, jwt

from app.config import get_settings

cfg = get_settings()


# ─── Password ────────────────────────────────────────────────────────────────

def hash_password(plain: str) -> str:
    """Hash a plaintext password with bcrypt."""
    return bcrypt.hashpw(plain.encode(), bcrypt.gensalt()).decode()


def verify_password(plain: str, hashed: str) -> bool:
    """Timing-safe bcrypt password verification."""
    return bcrypt.checkpw(plain.encode(), hashed.encode())


# ─── JWT ─────────────────────────────────────────────────────────────────────

def create_jwt(data: dict[str, Any], expires_minutes: int | None = None) -> str:
    """Create a signed JWT access token."""
    payload = data.copy()
    expire = datetime.now(timezone.utc) + timedelta(
        minutes=expires_minutes or cfg.ACCESS_TOKEN_EXPIRE_MINUTES
    )
    payload["exp"] = expire
    return jwt.encode(payload, cfg.SECRET_KEY, algorithm=cfg.JWT_ALGORITHM)


def verify_jwt(token: str) -> dict:
    """Verify and decode a JWT. Raises JWTError on failure."""
    return jwt.decode(token, cfg.SECRET_KEY, algorithms=[cfg.JWT_ALGORITHM])


# ─── TOTP ─────────────────────────────────────────────────────────────────────

def generate_totp_secret() -> str:
    """Generate a new base32 TOTP secret."""
    return pyotp.random_base32()


def verify_totp(token: str, secret: str) -> bool:
    """Verify a 6-digit TOTP code against the secret."""
    totp = pyotp.TOTP(secret)
    return totp.verify(token, valid_window=1)


def get_totp_provisioning_uri(secret: str, account_name: str) -> str:
    """Get the otpauth:// URI for QR code generation."""
    totp = pyotp.TOTP(secret)
    return totp.provisioning_uri(
        name=account_name,
        issuer_name=cfg.SUPERADMIN_2FA_ISSUER,
    )


# ─── HMAC Widget Tokens ───────────────────────────────────────────────────────

def sign_widget_token(author_slug: str, domain: str) -> str:
    """Generate a signed widget token for embedding on author's site.

    Format: base64url(payload).base64url(hmac)
    """
    payload = json.dumps(
        {"author_slug": author_slug, "domain": domain},
        separators=(",", ":"),
    )
    encoded = base64.urlsafe_b64encode(payload.encode()).decode().rstrip("=")
    sig = _hmac_sign(encoded)
    return f"{encoded}.{sig}"


def verify_widget_token(token: str) -> dict:
    """Verify a widget token and return the payload.

    Raises ValueError if invalid.
    """
    try:
        parts = token.split(".")
        if len(parts) != 2:
            raise ValueError("Malformed token")
        encoded_payload, provided_sig = parts
        expected_sig = _hmac_sign(encoded_payload)
        if not hmac.compare_digest(expected_sig, provided_sig):
            raise ValueError("Invalid signature")
        padded = encoded_payload + "=" * (4 - len(encoded_payload) % 4)
        return json.loads(base64.urlsafe_b64decode(padded).decode())
    except ValueError:
        raise
    except Exception as e:
        raise ValueError(f"Token verification failed: {e}") from e


def _hmac_sign(payload: str) -> str:
    secret = (cfg.SUBSCRIPTION_SECRET or cfg.SECRET_KEY).encode()
    sig = hmac.new(secret, payload.encode(), hashlib.sha256).digest()
    return base64.urlsafe_b64encode(sig).decode().rstrip("=")