Files
ITWorx_MCP_Hub-Public/services/control_api/auth/provider.py
T

220 lines
8.2 KiB
Python

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