"""OIDC Authorization Code + PKCE protocol helpers with strict token validation.""" from __future__ import annotations import base64 import hashlib import ipaddress import json import secrets import socket from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any from urllib.parse import urlencode, urlsplit import jwt import requests import urllib3 from app.core.system.enterprise_identity import ( IdentityPolicyError, IdentityUpstreamError, IdpConfig, ) def _digest(value: str) -> str: return hashlib.sha256(value.encode()).hexdigest() @dataclass(frozen=True) class AuthorizationStart: url: str state: str nonce: str code_verifier: str code_challenge: str code_challenge_method: str = "S256" class OidcClient: def __init__(self, repository: Any, *, clock: Callable[[], datetime] | None = None, flow_lifetime: timedelta = timedelta(minutes=5)) -> None: self.repository = repository self.clock = clock or (lambda: datetime.now(UTC)) self.flow_lifetime = flow_lifetime def begin(self, config: IdpConfig, redirect_uri: str) -> AuthorizationStart: config.validate() if redirect_uri not in config.redirect_uris: raise IdentityPolicyError("redirect URI is not allowlisted") state, nonce, verifier = (secrets.token_urlsafe(32) for _ in range(3)) challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode() self.repository.put_flow({"state_hash": _digest(state), "provider_uid": config.provider_uid, "provider_version": config.version, "nonce_hash": _digest(nonce), "verifier_hash": _digest(verifier), "redirect_uri": redirect_uri, "expires_at": self.clock() + self.flow_lifetime, "consumed_at": None}) params = {"response_type": "code", "client_id": config.client_id, "redirect_uri": redirect_uri, "scope": "openid profile", "state": state, "nonce": nonce, "code_challenge": challenge, "code_challenge_method": "S256"} return AuthorizationStart(f"{config.authorization_endpoint}?{urlencode(params)}", state, nonce, verifier, challenge) def verify_callback(self, config: IdpConfig, *, state: str, redirect_uri: str, code_verifier: str, id_token: str, jwks: Mapping[str, Any], commit: bool = True) -> dict[str, Any]: flow = self.repository.consume_flow(_digest(state), commit=commit) if not flow: raise IdentityPolicyError("state is invalid or already consumed") if self.clock() >= flow["expires_at"] or flow["provider_uid"] != config.provider_uid or flow["provider_version"] != config.version: raise IdentityPolicyError("state is expired or bound to another provider") if redirect_uri != flow["redirect_uri"] or redirect_uri not in config.redirect_uris or _digest(code_verifier) != flow["verifier_hash"]: raise IdentityPolicyError("callback binding or PKCE verifier is invalid") try: header = jwt.get_unverified_header(id_token) algorithm = header.get("alg") if algorithm not in config.algorithms or algorithm == "none": raise IdentityPolicyError("ID token algorithm is not allowed") candidates = [key for key in jwks.get("keys", ()) if key.get("kid") == header.get("kid") and key.get("alg", algorithm) == algorithm] if len(candidates) != 1: raise IdentityPolicyError("ID token signing key is ambiguous or missing") public_key = jwt.algorithms.get_default_algorithms()[algorithm].from_jwk(candidates[0]) claims = jwt.decode(id_token, public_key, algorithms=[algorithm], audience=config.client_id, issuer=config.issuer, options={"verify_exp": False, "require": ["iss", "aud", "sub", "nonce", "iat", "exp"]}) except IdentityPolicyError: raise except jwt.PyJWTError as exc: raise IdentityPolicyError("ID token validation failed") from exc if _digest(str(claims["nonce"])) != flow["nonce_hash"]: raise IdentityPolicyError("ID token nonce mismatch") now = int(self.clock().timestamp()) if int(claims["exp"]) <= now: raise IdentityPolicyError("ID token expired") if int(claims["iat"]) > now + 30: raise IdentityPolicyError("ID token issued in the future") return dict(claims) class OidcTransport: """Bounded HTTPS transport whose production connection is pinned to validated IPs.""" def __init__(self, http: Any | None = None, *, connect_timeout: float = 3.0, read_timeout: float = 7.0, max_response_bytes: int = 1024 * 1024, resolver: Callable[..., Any] = socket.getaddrinfo, pool_factory: Callable[..., Any] = urllib3.HTTPSConnectionPool, allowed_hosts: set[str] | None = None) -> None: self.http = http self.connect_timeout = connect_timeout self.read_timeout = read_timeout self.timeout = (connect_timeout, read_timeout) self.max_response_bytes = max_response_bytes self.resolver = resolver self.pool_factory = pool_factory self.allowed_hosts = {host.lower() for host in allowed_hosts} if allowed_hosts else None def _validate_destination(self, url: str) -> tuple[Any, tuple[str, ...]]: parsed = urlsplit(url) if parsed.scheme != "https" or not parsed.hostname: raise IdentityPolicyError("OIDC destination must be an absolute HTTPS URL") hostname = parsed.hostname.lower() if self.allowed_hosts is not None and hostname not in self.allowed_hosts: raise IdentityPolicyError("enterprise identity provider host is not server-allowlisted") try: addresses = { result[4][0] for result in self.resolver(hostname, parsed.port or 443, type=socket.SOCK_STREAM) } except (OSError, socket.gaierror) as exc: raise IdentityUpstreamError("enterprise identity provider DNS resolution failed") from exc if not addresses: raise IdentityUpstreamError("enterprise identity provider DNS returned no addresses") try: parsed_addresses = [ipaddress.ip_address(address) for address in addresses] except ValueError as exc: raise IdentityUpstreamError("enterprise identity provider DNS response is invalid") from exc if any(not address.is_global for address in parsed_addresses): raise IdentityPolicyError("enterprise identity provider resolved to a non-public address") return parsed, tuple(sorted(str(address) for address in parsed_addresses)) def _decode_json_response(self, *, status: int, headers: Mapping[str, Any], chunks: Any) -> Mapping[str, Any]: if 300 <= status < 400: raise IdentityUpstreamError("enterprise identity provider redirect is forbidden") if status < 200 or status >= 300: raise IdentityUpstreamError("enterprise identity provider returned an HTTP error") content_type = str(headers.get("Content-Type", "")).lower() if "application/json" not in content_type: raise IdentityUpstreamError("enterprise identity provider response must be JSON") content_length = headers.get("Content-Length") if content_length: try: parsed_length = int(content_length) except (TypeError, ValueError) as exc: raise IdentityUpstreamError("enterprise identity provider response length is invalid") from exc if parsed_length < 0: raise IdentityUpstreamError("enterprise identity provider response length is invalid") if parsed_length > self.max_response_bytes: raise IdentityUpstreamError("enterprise identity provider response is too large") body = bytearray() try: for chunk in chunks: if chunk: body.extend(chunk) if len(body) > self.max_response_bytes: raise IdentityUpstreamError("enterprise identity provider response is too large") except (requests.RequestException, urllib3.exceptions.HTTPError, OSError) as exc: raise IdentityUpstreamError("enterprise identity provider response read failed") from exc try: payload = json.loads(body.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise IdentityUpstreamError("enterprise identity provider returned invalid JSON") from exc if not isinstance(payload, Mapping): raise IdentityUpstreamError("enterprise identity provider JSON object is required") return payload def _pinned_request_json(self, method: str, parsed: Any, addresses: tuple[str, ...], data: Mapping[str, Any] | None) -> Mapping[str, Any]: hostname = parsed.hostname.lower() port = parsed.port or 443 request_target = parsed.path or "/" if parsed.query: request_target += f"?{parsed.query}" host_header = hostname if port == 443 else f"{hostname}:{port}" body = urlencode(data).encode("utf-8") if data is not None else None headers = {"Host": host_header, "Accept": "application/json"} if body is not None: headers["Content-Type"] = "application/x-www-form-urlencoded" last_error: Exception | None = None for address in addresses: pool = self.pool_factory( host=address, port=port, server_hostname=hostname, assert_hostname=hostname, cert_reqs="CERT_REQUIRED", ca_certs=requests.certs.where(), timeout=urllib3.Timeout(connect=self.connect_timeout, read=self.read_timeout), retries=False, maxsize=1, block=True, ) try: response = pool.urlopen( method.upper(), request_target, body=body, headers=headers, redirect=False, retries=False, assert_same_host=False, preload_content=False, decode_content=True, timeout=urllib3.Timeout(connect=self.connect_timeout, read=self.read_timeout), ) except (urllib3.exceptions.HTTPError, OSError) as exc: last_error = exc pool.close() continue try: return self._decode_json_response( status=int(response.status), headers=response.headers, chunks=response.stream(amt=64 * 1024, decode_content=True), ) finally: response.release_conn() pool.close() raise IdentityUpstreamError("enterprise identity provider connection failed") from last_error def _request_json(self, method: str, url: str, **kwargs: Any) -> Mapping[str, Any]: parsed, addresses = self._validate_destination(url) if self.http is None: return self._pinned_request_json(method, parsed, addresses, kwargs.get("data")) try: response = getattr(self.http, method)(url, timeout=self.timeout, allow_redirects=False, stream=True, **kwargs) except requests.RequestException as exc: raise IdentityUpstreamError("enterprise identity provider request failed") from exc try: try: status = int(getattr(response, "status_code", 0)) except (TypeError, ValueError) as exc: raise IdentityUpstreamError("enterprise identity provider HTTP status is invalid") from exc return self._decode_json_response( status=status, headers=getattr(response, "headers", {}), chunks=response.iter_content(chunk_size=64 * 1024), ) finally: close = getattr(response, "close", None) if callable(close): close() def exchange_code(self, config: IdpConfig, *, code: str, redirect_uri: str, verifier: str) -> str: payload = self._request_json( "post", config.token_endpoint, data={"grant_type": "authorization_code", "code": code, "redirect_uri": redirect_uri, "client_id": config.client_id, "client_secret": config.resolve_secret(), "code_verifier": verifier}, ) id_token = payload.get("id_token") if not isinstance(id_token, str) or not id_token: raise IdentityPolicyError("token endpoint omitted ID token") return id_token def fetch_jwks(self, config: IdpConfig) -> Mapping[str, Any]: payload = self._request_json("get", config.jwks_uri) if not isinstance(payload, Mapping) or not isinstance(payload.get("keys"), list): raise IdentityPolicyError("JWKS response is invalid") return payload