"""OIDC discovery, key rotation, and asymmetric JWT validation.""" from __future__ import annotations import json import time from datetime import UTC, datetime, timedelta from typing import Any from urllib.parse import urlencode, urljoin import httpx import jwt from pydantic import AnyHttpUrl, BaseModel, ConfigDict, TypeAdapter, field_validator from .config import OIDCSettings from .errors import InvalidToken, ProviderUnavailable, TokenExchangeError from .models import TokenSet, ValidatedClaims class DiscoveryDocument(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) issuer: str authorization_endpoint: AnyHttpUrl token_endpoint: AnyHttpUrl jwks_uri: AnyHttpUrl id_token_signing_alg_values_supported: tuple[str, ...] = ("RS256",) @field_validator("id_token_signing_alg_values_supported", mode="before") @classmethod def normalize_algorithms(cls, value: Any) -> tuple[str, ...]: if value is None: return ("RS256",) return tuple(value) @field_validator("issuer") @classmethod def validate_issuer_url(cls, value: str) -> str: parsed = TypeAdapter(AnyHttpUrl).validate_python(value) if parsed.scheme != "https" or parsed.username is not None or parsed.password is not None: raise ValueError("OIDC issuer must be an HTTPS URL without userinfo") return value class OIDCProviderClient: """Provider client with bounded discovery/JWKS caches and one rotation retry.""" def __init__( self, settings: OIDCSettings, *, transport: httpx.BaseTransport | None = None, timeout_seconds: float = 5.0, cache_ttl_seconds: int = 300, ) -> None: self.settings = settings self._client = httpx.Client(transport=transport, timeout=timeout_seconds) self._cache_ttl_seconds = max(30, min(cache_ttl_seconds, 900)) self._discovery: DiscoveryDocument | None = None self._discovery_loaded = 0.0 self._jwks: dict[str, dict[str, Any]] = {} self._jwks_loaded = 0.0 def close(self) -> None: self._client.close() def _request_json(self, method: str, url: str, **kwargs: Any) -> dict[str, Any]: try: response = self._client.request(method, url, **kwargs) response.raise_for_status() value = response.json() except (httpx.HTTPError, ValueError, json.JSONDecodeError) as exc: raise ProviderUnavailable() from exc if not isinstance(value, dict): raise ProviderUnavailable() return value def discover(self, *, force: bool = False) -> DiscoveryDocument: now = time.monotonic() if ( self._discovery is not None and not force and now - self._discovery_loaded < self._cache_ttl_seconds ): return self._discovery configured_issuer = str(self.settings.issuer) discovery_base = configured_issuer.rstrip("/") + "/" raw = self._request_json("GET", urljoin(discovery_base, ".well-known/openid-configuration")) try: document = DiscoveryDocument.model_validate(raw) except ValueError as exc: raise ProviderUnavailable() from exc if document.issuer != configured_issuer: raise ProviderUnavailable() issuer_origin = discovery_base.split("/", 3)[:3] for endpoint in ( document.authorization_endpoint, document.token_endpoint, document.jwks_uri, ): if str(endpoint).split("/", 3)[:3] != issuer_origin: raise ProviderUnavailable() self._discovery = document self._discovery_loaded = now return document def _load_jwks(self, *, force: bool = False) -> dict[str, dict[str, Any]]: now = time.monotonic() if self._jwks and not force and now - self._jwks_loaded < self._cache_ttl_seconds: return self._jwks document = self.discover() raw = self._request_json("GET", str(document.jwks_uri)) keys = raw.get("keys") if not isinstance(keys, list): raise ProviderUnavailable() parsed: dict[str, dict[str, Any]] = {} for key in keys: if ( isinstance(key, dict) and isinstance(key.get("kid"), str) and key.get("kty") in {"RSA", "EC", "OKP"} ): parsed[key["kid"]] = key if not parsed: raise ProviderUnavailable() self._jwks = parsed self._jwks_loaded = now return parsed def build_authorization_url(self, *, state: str, nonce: str, code_challenge: str) -> str: document = self.discover() query = urlencode( { "response_type": "code", "client_id": self.settings.client_id, "redirect_uri": str(self.settings.redirect_uri), "scope": "openid profile email groups", "state": state, "nonce": nonce, "code_challenge": code_challenge, "code_challenge_method": "S256", } ) return f"{document.authorization_endpoint}?{query}" def exchange_code(self, *, code: str, code_verifier: str) -> TokenSet: document = self.discover() data = { "grant_type": "authorization_code", "code": code, "redirect_uri": str(self.settings.redirect_uri), "client_id": self.settings.client_id, "code_verifier": code_verifier, } secret = self.settings.client_secret() if secret is not None: data["client_secret"] = secret try: response = self._client.post(str(document.token_endpoint), data=data) response.raise_for_status() return TokenSet.model_validate(response.json()) except (httpx.HTTPError, ValueError, json.JSONDecodeError) as exc: raise TokenExchangeError() from exc def validate_token(self, token: str, *, now: datetime | None = None) -> ValidatedClaims: if not token or len(token) > 65536: raise InvalidToken() try: header = jwt.get_unverified_header(token) except jwt.PyJWTError as exc: raise InvalidToken() from exc algorithm = header.get("alg") document = self.discover() kid = header.get("kid") if not isinstance(algorithm, str) or algorithm not in self.settings.allowed_algorithms: raise InvalidToken() if ( not isinstance(kid, str) or not kid or algorithm not in document.id_token_signing_alg_values_supported ): raise InvalidToken() keys = self._load_jwks() key = keys.get(kid) if key is None: keys = self._load_jwks(force=True) key = keys.get(kid) if key is None: raise InvalidToken() try: signing_key = jwt.PyJWK.from_dict(key).key claims = jwt.decode( token, signing_key, algorithms=list(self.settings.allowed_algorithms), audience=self.settings.audience, issuer=str(document.issuer), leeway=self.settings.clock_skew_seconds, options={ "require": ["sub", "iss", "aud", "exp", "iat"], "verify_exp": False, "verify_iat": False, "verify_nbf": False, }, ) validated = ValidatedClaims.model_validate(claims) except (jwt.PyJWTError, ValueError, TypeError) as exc: raise InvalidToken() from exc reference = now or datetime.now(UTC) if reference.tzinfo is None or reference.utcoffset() is None: raise InvalidToken() skew = timedelta(seconds=self.settings.clock_skew_seconds) if validated.expires_at + skew <= reference: raise InvalidToken() if validated.issued_at - skew > reference: raise InvalidToken() if validated.not_before is not None and validated.not_before - skew > reference: raise InvalidToken() return validated