DataOps-platform/app/core/system/oidc.py

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