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