SSO MCP Server

"""JWKS client for fetching and caching signing keys.

Fetches JSON Web Key Sets (JWKS) from OAuth authorization servers
and caches them for efficient token validation.
"""

from __future__ import annotations

import time
from typing import TYPE_CHECKING

import httpx
import jwt
from jwt import PyJWK, PyJWKSet

from sso_mcp_server import get_logger
from sso_mcp_server.auth.exceptions import InvalidTokenError, JWKSFetchError

if TYPE_CHECKING:
    from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey

_logger = get_logger("jwks_client")

# Default timeout for HTTP requests
HTTP_TIMEOUT = 10.0


class JWKSClient:
    """Client for fetching and caching JWKS from OAuth issuers.

    Implements caching with configurable TTL and automatic key rotation
    handling when a key ID is not found in the cached JWKS.

    Attributes:
        cache_ttl: Time-to-live for cached JWKS in seconds.
    """

    def __init__(self, cache_ttl: int = 3600) -> None:
        """Initialize JWKS client.

        Args:
            cache_ttl: Cache time-to-live in seconds (default: 1 hour).
        """
        self._cache_ttl = cache_ttl
        self._cache: dict[str, _CachedJWKS] = {}
        self._http_client: httpx.AsyncClient | None = None

        _logger.debug("jwks_client_initialized", cache_ttl=cache_ttl)

    async def _get_http_client(self) -> httpx.AsyncClient:
        """Get or create HTTP client (lazy initialization)."""
        if self._http_client is None:
            self._http_client = httpx.AsyncClient(timeout=HTTP_TIMEOUT)
        return self._http_client

    async def close(self) -> None:
        """Close the HTTP client."""
        if self._http_client is not None:
            await self._http_client.aclose()
            self._http_client = None

    async def get_signing_key(self, token: str) -> RSAPublicKey:
        """Get the signing key for a JWT token.

        Extracts the key ID (kid) and issuer (iss) from the token header,
        fetches the JWKS from the issuer (with caching), and returns the
        matching public key.

        Args:
            token: The JWT token string.

        Returns:
            The RSA public key for verifying the token signature.

        Raises:
            InvalidTokenError: If the token is malformed or missing required claims.
            JWKSFetchError: If the JWKS cannot be fetched from the issuer.
        """
        # Decode token header without verification to get kid and iss
        try:
            unverified_header = jwt.get_unverified_header(token)
            unverified_payload = jwt.decode(token, options={"verify_signature": False})
        except jwt.exceptions.DecodeError as e:
            _logger.warning("failed_to_decode_token_header", error=str(e))
            raise InvalidTokenError("Malformed JWT token") from e

        kid = unverified_header.get("kid")
        issuer = unverified_payload.get("iss")

        if not kid:
            raise InvalidTokenError("Token missing 'kid' (key ID) in header")
        if not issuer:
            raise InvalidTokenError("Token missing 'iss' (issuer) claim")

        _logger.debug("extracting_signing_key", kid=kid, issuer=issuer)

        # Get JWKS for this issuer
        jwks = await self._get_jwks(issuer)

        # Find the key matching the kid
        key = self._find_key_by_kid(jwks, kid, issuer)

        return key.key

    async def _get_jwks(self, issuer: str, force_refresh: bool = False) -> PyJWKSet:
        """Get JWKS for an issuer (with caching).

        Args:
            issuer: The token issuer URL.
            force_refresh: Force fetching fresh JWKS ignoring cache.

        Returns:
            The PyJWKSet for the issuer.

        Raises:
            JWKSFetchError: If the JWKS cannot be fetched.
        """
        # Check cache
        if not force_refresh and issuer in self._cache:
            cached = self._cache[issuer]
            if not cached.is_expired():
                _logger.debug("jwks_cache_hit", issuer=issuer)
                return cached.jwks

        _logger.debug("jwks_cache_miss", issuer=issuer, force_refresh=force_refresh)

        # Fetch JWKS
        jwks = await self._fetch_jwks(issuer)

        # Cache it
        self._cache[issuer] = _CachedJWKS(jwks=jwks, ttl=self._cache_ttl)

        return jwks

    async def _fetch_jwks(self, issuer: str) -> PyJWKSet:
        """Fetch JWKS from an issuer's well-known endpoint.

        First fetches the OpenID Connect discovery document to get the jwks_uri,
        then fetches the JWKS from that URI.

        Args:
            issuer: The token issuer URL.

        Returns:
            The PyJWKSet containing the signing keys.

        Raises:
            JWKSFetchError: If the JWKS cannot be fetched.
        """
        client = await self._get_http_client()

        # Construct OpenID Connect discovery URL
        discovery_url = self._get_discovery_url(issuer)
        _logger.debug("fetching_oidc_discovery", url=discovery_url)

        try:
            # Fetch discovery document
            response = await client.get(discovery_url)
            response.raise_for_status()
            discovery = response.json()

            jwks_uri = discovery.get("jwks_uri")
            if not jwks_uri:
                raise JWKSFetchError(issuer, "Discovery document missing 'jwks_uri'")

            _logger.debug("fetching_jwks", jwks_uri=jwks_uri)

            # Fetch JWKS
            response = await client.get(jwks_uri)
            response.raise_for_status()
            jwks_data = response.json()

            return PyJWKSet.from_dict(jwks_data)

        except httpx.HTTPStatusError as e:
            _logger.error("jwks_http_error", issuer=issuer, status=e.response.status_code)
            raise JWKSFetchError(issuer, f"HTTP {e.response.status_code}") from e
        except httpx.RequestError as e:
            _logger.error("jwks_request_error", issuer=issuer, error=str(e))
            raise JWKSFetchError(issuer, str(e)) from e
        except (KeyError, ValueError) as e:
            _logger.error("jwks_parse_error", issuer=issuer, error=str(e))
            raise JWKSFetchError(issuer, f"Invalid JWKS format: {e}") from e

    def _get_discovery_url(self, issuer: str) -> str:
        """Construct the OpenID Connect discovery URL for an issuer.

        Handles both issuers with and without path components per RFC 8414.

        Args:
            issuer: The token issuer URL.

        Returns:
            The discovery document URL.
        """
        # Remove trailing slash
        issuer = issuer.rstrip("/")

        # For Azure AD v2.0, the issuer typically ends with /v2.0
        # The discovery URL is at /.well-known/openid-configuration
        return f"{issuer}/.well-known/openid-configuration"

    def _find_key_by_kid(self, jwks: PyJWKSet, kid: str, issuer: str) -> PyJWK:
        """Find a key in the JWKS by key ID.

        If the key is not found, attempts to refresh the JWKS once
        to handle key rotation.

        Args:
            jwks: The JWKS to search.
            kid: The key ID to find.
            issuer: The issuer (for error messages and refresh).

        Returns:
            The matching PyJWK.

        Raises:
            InvalidTokenError: If no matching key is found.
        """
        for key in jwks.keys:
            if key.key_id == kid:
                return key

        _logger.warning("key_not_found_in_jwks", kid=kid, issuer=issuer)
        raise InvalidTokenError(f"No key found for kid '{kid}'")

    async def get_signing_key_with_retry(self, token: str) -> RSAPublicKey:
        """Get signing key with automatic retry on cache miss.

        If the key is not found in the cached JWKS, this method will
        force a refresh of the JWKS and try again once. This handles
        key rotation scenarios.

        Args:
            token: The JWT token string.

        Returns:
            The RSA public key for verifying the token signature.
        """
        try:
            return await self.get_signing_key(token)
        except InvalidTokenError as e:
            if "No key found for kid" in str(e):
                # Key not found - might be key rotation, try refreshing JWKS
                _logger.info("retrying_jwks_fetch_after_key_not_found")

                # Get issuer from token
                unverified_payload = jwt.decode(token, options={"verify_signature": False})
                issuer = unverified_payload.get("iss")

                if issuer:
                    # Force refresh JWKS
                    await self._get_jwks(issuer, force_refresh=True)
                    # Retry
                    return await self.get_signing_key(token)
            raise

    def clear_cache(self) -> None:
        """Clear the JWKS cache."""
        self._cache.clear()
        _logger.debug("jwks_cache_cleared")


class _CachedJWKS:
    """Internal class for caching JWKS with expiry."""

    def __init__(self, jwks: PyJWKSet, ttl: int) -> None:
        """Initialize cached JWKS.

        Args:
            jwks: The JWKS to cache.
            ttl: Time-to-live in seconds.
        """
        self.jwks = jwks
        self.expires_at = time.time() + ttl

    def is_expired(self) -> bool:
        """Check if the cache entry has expired."""
        return time.time() >= self.expires_at