app/core/access/token_crypto.py

"""Author RAG Chatbot SaaS — Cryptographic Token Utilities.

All JWT and HMAC operations are centralized here.
RULE: Never call jose/jwt directly outside this module.
RULE: Never log token values — only log token IDs or claims.
"""

import hashlib
import hmac
import json
import time
import uuid
from base64 import urlsafe_b64decode, urlsafe_b64encode
from datetime import datetime, timedelta, timezone
from typing import Any

import structlog
from jose import JWTError, jwt

from app.config import get_settings
from app.exceptions.auth import ExpiredTokenError, InvalidTokenError

logger = structlog.get_logger(__name__)
cfg = get_settings()


def _jwt_signing_key() -> str:
    """Return the key used to sign JWTs."""
    if cfg.JWT_ALGORITHM.upper() == "HS256":
        return cfg.SECRET_KEY
    if not cfg.JWT_PRIVATE_KEY:
        raise InvalidTokenError("JWT signing key is not configured")
    return cfg.JWT_PRIVATE_KEY


def _jwt_verification_key() -> str:
    """Return the key used to verify JWTs."""
    if cfg.JWT_ALGORITHM.upper() == "HS256":
        return cfg.SECRET_KEY
    if not cfg.JWT_PUBLIC_KEY:
        raise InvalidTokenError("JWT verification key is not configured")
    return cfg.JWT_PUBLIC_KEY


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

def create_access_token(subject: str, extra_claims: dict[str, Any] | None = None) -> str:
    """Create a short-lived RS256 JWT access token.

    Args:
        subject: User ID (UUID string) to set as the 'sub' claim.
        extra_claims: Optional additional claims to embed.

    Returns:
        Signed JWT string.
    """
    now = datetime.now(timezone.utc)
    expire = now + timedelta(minutes=cfg.ACCESS_TOKEN_EXPIRE_MINUTES)
    payload = {
        "sub": subject,
        "jti": uuid.uuid4().hex,  # R-004/R-005: JWT ID for blacklisting
        "iat": now,
        "exp": expire,
        "type": "access",
        **(extra_claims or {}),
    }
    return jwt.encode(payload, _jwt_signing_key(), algorithm=cfg.JWT_ALGORITHM)


def create_refresh_token(subject: str) -> str:
    """Create a long-lived RS256 JWT refresh token.

    Args:
        subject: User ID (UUID string).

    Returns:
        Signed JWT string.
    """
    now = datetime.now(timezone.utc)
    expire = now + timedelta(days=cfg.REFRESH_TOKEN_EXPIRE_DAYS)
    payload = {
        "sub": subject,
        "jti": uuid.uuid4().hex,  # R-004/R-005: JWT ID for blacklisting
        "iat": now,
        "exp": expire,
        "type": "refresh",
    }
    return jwt.encode(payload, _jwt_signing_key(), algorithm=cfg.JWT_ALGORITHM)


def decode_jwt(token: str) -> dict[str, Any]:
    """Decode and validate a JWT. Raises on any failure.

    Args:
        token: Raw JWT string.

    Returns:
        Decoded payload dict.

    Raises:
        ExpiredTokenError: If the token is expired.
        InvalidTokenError: If the token is malformed or signature is invalid.
    """
    try:
        payload = jwt.decode(token, _jwt_verification_key(), algorithms=[cfg.JWT_ALGORITHM])
        return payload
    except JWTError as e:
        error_msg = str(e).lower()
        if "expired" in error_msg:
            raise ExpiredTokenError()
        raise InvalidTokenError(f"Token validation failed: {e}")


# ─── HMAC Subscription Tokens ─────────────────────────────────────────────────

def _datetime_to_unix(dt: datetime) -> int:
    """Convert datetime to unix timestamp, treating naive values as UTC."""
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return int(dt.timestamp())


def create_subscription_token(
    author_id: str,
    grant_id: str,
    granted_at: datetime,
    expires_at: datetime,
) -> str:
    """Create a tamper-proof subscription token using HMAC-SHA256.

    The token is a URL-safe Base64 encoded JSON payload with an HMAC signature.
    Modifying any byte of the payload invalidates the signature.

    Args:
        author_id: UUID of the author being granted access.
        grant_id: UUID of the ClientAccess record.
        granted_at: When this grant was made.
        expires_at: When this grant expires.

    Returns:
        Signed token string (URL-safe, no slashes or plus signs).
    """
    payload = {
        "author_id": author_id,
        "grant_id": grant_id,
        "granted_at": _datetime_to_unix(granted_at),
        "expires_at": _datetime_to_unix(expires_at),
    }
    payload_bytes = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode()
    payload_b64 = urlsafe_b64encode(payload_bytes).rstrip(b"=")

    signature = _compute_hmac(payload_b64)
    token = payload_b64 + b"." + signature

    return token.decode()


def validate_subscription_token_signature(token: str) -> dict[str, Any]:
    """Validate HMAC signature and decode subscription token payload.

    Uses constant-time comparison to prevent timing attacks.

    Args:
        token: Raw subscription token string.

    Returns:
        Decoded payload dict.

    Raises:
        InvalidTokenError: If signature is invalid or token is malformed.
        ExpiredTokenError: If the token has passed its expires_at.
    """
    try:
        parts = token.encode().split(b".")
        if len(parts) != 2:
            raise InvalidTokenError("Malformed subscription token")

        payload_b64, provided_sig = parts
        expected_sig = _compute_hmac(payload_b64)

        # Constant-time comparison — prevents timing attacks
        if not hmac.compare_digest(expected_sig, provided_sig):
            logger.warning("Subscription token HMAC mismatch — possible tampering detected")
            raise InvalidTokenError("Invalid subscription token signature")

        # Decode payload
        padding = b"=" * (4 - len(payload_b64) % 4)
        payload = json.loads(urlsafe_b64decode(payload_b64 + padding))

        # Check expiry
        if int(time.time()) > payload["expires_at"]:
            raise ExpiredTokenError("Subscription has expired")

        return payload

    except (InvalidTokenError, ExpiredTokenError):
        raise
    except Exception as e:
        raise InvalidTokenError(f"Could not parse subscription token: {e}")


def hash_subscription_token(token: str) -> str:
    """Create a SHA-256 hash of the token for secure DB storage.

    We store the hash, never the raw token.

    Args:
        token: Raw subscription token string.

    Returns:
        Hex string of SHA-256 hash.
    """
    return hashlib.sha256(token.encode()).hexdigest()


def _compute_hmac(payload_b64: bytes) -> bytes:
    """Compute URL-safe Base64 encoded HMAC-SHA256 of payload.

    Args:
        payload_b64: URL-safe Base64 encoded payload bytes.

    Returns:
        URL-safe Base64 encoded signature bytes.
    """
    signature = hmac.new(
        (cfg.SUBSCRIPTION_SECRET or cfg.SECRET_KEY).encode(),
        payload_b64,
        hashlib.sha256,
    ).digest()
    return urlsafe_b64encode(signature).rstrip(b"=")