simbiot-harness/venv/lib/python3.12/site-packages/openai/auth/_x509.py

from __future__ import annotations

import re
import math
import time
import email.utils
from typing import Any, NoReturn, cast
from contextvars import ContextVar
from typing_extensions import TypeIs, override

import anyio
import httpx2

from .._utils import is_dict
from .._httpx2 import timeout_exceptions, normalize_httpx_url, _loaded_legacy_httpx
from ._workload import (
    TOKEN_EXCHANGE_GRANT_TYPE,
    WorkloadIdentity,
    X509WorkloadIdentity,
    _WorkloadIdentityAuth,
)
from .._constants import MAX_RETRY_DELAY, INITIAL_RETRY_DELAY, MAX_RETRY_AFTER_DELAY
from .._exceptions import OAuthError, OpenAIError, APITimeoutError, APIConnectionError

MTLS_API_BASE_URL = "https://mtls.api.openai.com/v1"
_X509_TOKEN_EXCHANGE_URL = "https://mtls.auth.openai.com/oauth/token"
_X509_SUBJECT_TOKEN_TYPE = "urn:openai:params:oauth:token-type:x509"
_MAX_EXCHANGE_RETRIES = 2
_REPLAY_POSITION_EXTENSION = "openai_x509_replay_position"
_REPLAY_FILE_POSITIONS_EXTENSION = "openai_x509_replay_file_positions"
_ALLOWED_IDENTITY_FIELDS = {"type", "identity_provider_id", "service_account_id", "refresh_buffer_seconds"}
_BEARER_ACCESS_TOKEN = re.compile(r"[A-Za-z0-9._~+/-]+=*")
_MTLS_REGIONAL_BASE_URLS = {
    "global": MTLS_API_BASE_URL,
    "us": "https://mtls-us.api.openai.com/v1",
    "eu": "https://mtls-eu.api.openai.com/v1",
}
_OPENAI_MTLS_HOSTS = {httpx2.URL(url).host for url in _MTLS_REGIONAL_BASE_URLS.values()}
_EXCHANGE_REQUEST_TIMEOUT: ContextVar[dict[str, float | None] | None] = ContextVar(
    "openai_x509_exchange_request_timeout", default=None
)


class _TransientTokenExchangeError(Exception):
    def __init__(self, error: OpenAIError) -> None:
        self.error = error


def validate_x509_api_url(url: httpx2.URL | str, *, expected_origin: httpx2.URL | None = None) -> None:
    normalized_url = normalize_httpx_url(url)
    if normalized_url.scheme != "https" or not normalized_url.host:
        raise OpenAIError("X.509 workload identity requires an absolute HTTPS API URL")

    if normalized_url.username or normalized_url.password:
        raise OpenAIError("X.509 workload identity API URLs cannot contain user credentials")

    if expected_origin is not None and (normalized_url.host, normalized_url.port) != (
        expected_origin.host,
        expected_origin.port,
    ):
        raise OpenAIError("X.509 workload identity requests must use the configured API origin")


def validate_x509_api_credentials(request: httpx2.Request) -> None:
    if any(
        header.lower().replace("_", "-") in ("api-key", "x-api-key", "proxy-authorization")
        for header in request.headers
    ):
        raise OpenAIError("X.509 workload identity requests cannot include API-key or proxy credentials")


def validate_x509_request_authority(request: httpx2.Request) -> None:
    host_headers = request.headers.get_list("host")
    if (
        len(host_headers) != 1
        or any(delimiter in host_headers[0] for delimiter in "/?#@\\")
        or any(header.startswith(":") for header in request.headers)
    ):
        raise OpenAIError("X.509 workload identity requests require exactly one valid Host authority")

    try:
        host_url = httpx2.URL(f"https://{host_headers[0]}")
    except httpx2.InvalidURL as error:
        raise OpenAIError("X.509 workload identity requests require a valid Host authority") from error

    if (
        host_url.username
        or host_url.password
        or host_url.path != "/"
        or host_url.query
        or host_url.fragment
        or (host_url.host, host_url.port) != (request.url.host, request.url.port)
    ):
        raise OpenAIError("X.509 workload identity Host authority must match the request URL")


def _validate_transport_request(
    request: httpx2.Request,
    *,
    expected_origin: httpx2.URL,
    expected_authorization: str | None,
    token_exchange: bool,
) -> None:
    validate_x509_api_url(request.url, expected_origin=expected_origin)
    validate_x509_request_authority(request)

    target = request.extensions.get("target")
    if target is not None and target != request.url.raw_path:
        raise OpenAIError("X.509 workload identity request target must match the request URL")

    sni_hostname = request.extensions.get("sni_hostname")
    if request.url.host in _OPENAI_MTLS_HOSTS and sni_hostname is not None:
        if not isinstance(sni_hostname, str) or sni_hostname.lower() != expected_origin.host.lower():
            raise OpenAIError("X.509 workload identity TLS hostname must match the configured origin")

    if token_exchange:
        if str(request.url) != _X509_TOKEN_EXCHANGE_URL:
            raise OpenAIError("X.509 token exchange requests must use the pinned authentication URL")
        if any(
            header.lower().replace("_", "-")
            in ("authorization", "proxy-authorization", "cookie", "x-api-key", "api-key")
            for header in request.headers
        ):
            raise OpenAIError("X.509 token exchange requests cannot include API credentials")
    else:
        validate_x509_api_credentials(request)
        if request.headers.get("Authorization") != expected_authorization:
            raise OpenAIError("X.509 workload identity authorization cannot be changed by HTTP request hooks")


class _SyncX509ForwardingTransport(httpx2.BaseTransport):
    def __init__(
        self,
        *,
        http_client: httpx2.Client,
        expected_origin: httpx2.URL,
        expected_authorization: str | None,
        token_exchange: bool,
    ) -> None:
        self._http_client = http_client
        self._expected_origin = expected_origin
        self._expected_authorization = expected_authorization
        self._token_exchange = token_exchange

    @override
    def handle_request(self, request: httpx2.Request) -> httpx2.Response:
        _validate_transport_request(
            request,
            expected_origin=self._expected_origin,
            expected_authorization=self._expected_authorization,
            token_exchange=self._token_exchange,
        )
        if self._http_client.is_closed:
            raise RuntimeError("Cannot send a request, as the client has been closed.")
        return self._http_client._transport_for_url(request.url).handle_request(request)

    @override
    def close(self) -> None:
        # The caller owns the selected connection pool and proxy transports.
        return None


class _AsyncX509ForwardingTransport(httpx2.AsyncBaseTransport):
    def __init__(
        self,
        *,
        http_client: httpx2.AsyncClient,
        expected_origin: httpx2.URL,
        expected_authorization: str | None,
        token_exchange: bool,
    ) -> None:
        self._http_client = http_client
        self._expected_origin = expected_origin
        self._expected_authorization = expected_authorization
        self._token_exchange = token_exchange

    @override
    async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
        _validate_transport_request(
            request,
            expected_origin=self._expected_origin,
            expected_authorization=self._expected_authorization,
            token_exchange=self._token_exchange,
        )
        if self._http_client.is_closed:
            raise RuntimeError("Cannot send a request, as the client has been closed.")
        return await self._http_client._transport_for_url(request.url).handle_async_request(request)

    @override
    async def aclose(self) -> None:
        # The caller owns the selected connection pool and proxy transports.
        return None


def _scoped_sync_client(
    http_client: httpx2.Client,
    *,
    expected_origin: httpx2.URL,
    expected_authorization: str | None = None,
    token_exchange: bool = False,
) -> httpx2.Client:
    transport = _SyncX509ForwardingTransport(
        http_client=http_client,
        expected_origin=expected_origin,
        expected_authorization=expected_authorization,
        token_exchange=token_exchange,
    )
    legacy_httpx = _loaded_legacy_httpx()
    client_type = httpx2.Client
    if legacy_httpx is not None and not isinstance(cast(object, http_client), httpx2.Client):
        client_type = legacy_httpx.Client
    scoped_client = client_type(
        transport=transport,
        timeout=http_client.timeout,
        event_hooks=None if token_exchange else http_client.event_hooks,
        default_encoding=http_client._default_encoding,
        trust_env=False,
    )
    if not token_exchange:
        scoped_client._cookies = http_client.cookies
    return scoped_client


def _scoped_async_client(
    http_client: httpx2.AsyncClient,
    *,
    expected_origin: httpx2.URL,
    expected_authorization: str | None = None,
    token_exchange: bool = False,
) -> httpx2.AsyncClient:
    transport = _AsyncX509ForwardingTransport(
        http_client=http_client,
        expected_origin=expected_origin,
        expected_authorization=expected_authorization,
        token_exchange=token_exchange,
    )
    legacy_httpx = _loaded_legacy_httpx()
    client_type = httpx2.AsyncClient
    if legacy_httpx is not None and not isinstance(cast(object, http_client), httpx2.AsyncClient):
        client_type = legacy_httpx.AsyncClient
    scoped_client = client_type(
        transport=transport,
        timeout=http_client.timeout,
        event_hooks=None if token_exchange else http_client.event_hooks,
        default_encoding=http_client._default_encoding,
        trust_env=False,
    )
    if not token_exchange:
        scoped_client._cookies = http_client.cookies
    return scoped_client


def _as_finite_float(value: object) -> float | None:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    try:
        value = float(value)
    except OverflowError:
        return None
    return value if math.isfinite(value) else None


def is_x509_workload_identity(
    identity: WorkloadIdentity | X509WorkloadIdentity | None,
) -> TypeIs[X509WorkloadIdentity]:
    return identity is not None and identity.get("type") == "x509"


def x509_data_residency_base_url(
    base_url: httpx2.URL | str | None,
    data_residency: str | None,
    workload_identity: WorkloadIdentity | X509WorkloadIdentity | None,
) -> httpx2.URL | str | None:
    if data_residency is None or not is_x509_workload_identity(workload_identity):
        return base_url
    if data_residency not in _MTLS_REGIONAL_BASE_URLS:
        raise OpenAIError("X.509 workload identity requires a supported regional mTLS endpoint")
    return _MTLS_REGIONAL_BASE_URLS[data_residency]


def x509_safe_environment_headers(
    headers: dict[str, str], workload_identity: X509WorkloadIdentity | None
) -> dict[str, str]:
    if workload_identity is None:
        return headers
    return {name: value for name, value in headers.items() if name.lower() != "authorization"}


def _validate_identity(identity: X509WorkloadIdentity) -> None:
    if "provider" in identity or "client_id" in identity:
        raise OpenAIError("X.509 workload identity does not accept a subject-token provider or client ID")

    if set(identity) - _ALLOWED_IDENTITY_FIELDS:
        raise OpenAIError("X.509 workload identity accepts only identity IDs and an optional refresh buffer")

    if any(
        not isinstance(identity.get(field), str) or not identity.get(field)
        for field in (
            "identity_provider_id",
            "service_account_id",
        )
    ):
        raise OpenAIError("X.509 workload identity requires identity-provider and service-account IDs")

    refresh_buffer = cast(object, identity.get("refresh_buffer_seconds"))
    if refresh_buffer is not None:
        refresh_buffer_value = _as_finite_float(refresh_buffer)
        if refresh_buffer_value is None or refresh_buffer_value < 0:
            raise OpenAIError("X.509 workload identity requires a finite, non-negative refresh buffer")


def _exchange_payload(identity: X509WorkloadIdentity) -> dict[str, str]:
    return {
        "grant_type": TOKEN_EXCHANGE_GRANT_TYPE,
        "subject_token_type": _X509_SUBJECT_TOKEN_TYPE,
        "identity_provider_id": identity["identity_provider_id"],
        "service_account_id": identity["service_account_id"],
    }


def _token_exchange_request(
    identity: X509WorkloadIdentity,
    *,
    http_client: httpx2.Client | httpx2.AsyncClient,
) -> httpx2.Request:
    request_type = httpx2.Request
    legacy_httpx = _loaded_legacy_httpx()
    if legacy_httpx is not None and not isinstance(cast(object, http_client), (httpx2.Client, httpx2.AsyncClient)):
        request_type = cast(type[httpx2.Request], cast(Any, legacy_httpx).Request)

    configured_timeout = _EXCHANGE_REQUEST_TIMEOUT.get()
    timeout = {
        phase: min(value, 10.0) if value is not None else 10.0
        for phase, value in (configured_timeout or httpx2.Timeout(10.0).as_dict()).items()
    }
    return request_type(
        "POST",
        _X509_TOKEN_EXCHANGE_URL,
        json=_exchange_payload(identity),
        extensions={"timeout": timeout},
    )


def _retry_delay(response: httpx2.Response | None, attempt: int) -> float | None:
    if response is not None:
        should_retry = response.headers.get("x-should-retry")
        if response.status_code in (400, 401, 403) or should_retry == "false":
            return None
        if should_retry != "true" and response.status_code not in (408, 409, 429) and response.status_code < 500:
            return None

        retry_after_ms = response.headers.get("retry-after-ms")
        retry_after = response.headers.get("retry-after")
        if retry_after_ms is not None:
            try:
                millisecond_delay = float(retry_after_ms) / 1000
            except ValueError:
                pass
            else:
                if math.isfinite(millisecond_delay) and 0 <= millisecond_delay <= MAX_RETRY_AFTER_DELAY:
                    return millisecond_delay
                if millisecond_delay > MAX_RETRY_AFTER_DELAY:
                    return None
        if retry_after is not None:
            try:
                delay = float(retry_after)
            except ValueError:
                try:
                    parsed = email.utils.parsedate_tz(retry_after)
                    delay = float(email.utils.mktime_tz(parsed) - time.time()) if parsed is not None else -1
                except (OverflowError, OSError, ValueError):
                    delay = -1

            if math.isfinite(delay) and 0 <= delay <= MAX_RETRY_AFTER_DELAY:
                return delay
            if delay > MAX_RETRY_AFTER_DELAY:
                return None

    return float(min(INITIAL_RETRY_DELAY * 2**attempt, MAX_RETRY_DELAY))


def _is_replayable_request(request: httpx2.Request) -> bool:
    if isinstance(getattr(request, "_content", None), bytes):
        return True

    stream = request.stream
    fields = getattr(stream, "fields", None)
    if isinstance(fields, list):
        file_positions: list[tuple[object, int]] = []
        for field in cast(list[object], fields):
            file = getattr(field, "file", None)
            if file is None or isinstance(file, (str, bytes)):
                continue
            seekable = getattr(file, "seekable", None)
            seek = getattr(file, "seek", None)
            tell = getattr(file, "tell", None)
            try:
                if not callable(seekable) or not seekable() or not callable(seek) or not callable(tell):
                    return False
                position = tell()
            except (OSError, ValueError):
                return False
            if not isinstance(position, int):
                return False
            file_positions.append((file, position))
        request.extensions[_REPLAY_FILE_POSITIONS_EXTENSION] = file_positions
        return True

    source = getattr(stream, "_stream", stream)
    seekable = getattr(source, "seekable", None)
    seek = getattr(source, "seek", None)
    tell = getattr(source, "tell", None)
    try:
        if not callable(seekable) or not seekable() or not callable(seek) or not callable(tell):
            return False
        request.extensions[_REPLAY_POSITION_EXTENSION] = tell()
    except (OSError, ValueError):
        return False
    return True


def _transport_errors() -> tuple[type[Exception], ...]:
    legacy_httpx = _loaded_legacy_httpx()
    if legacy_httpx is None:
        return (httpx2.TransportError,)
    legacy_transport_error = cast(type[Exception], getattr(legacy_httpx, "TransportError", httpx2.TransportError))
    return (httpx2.TransportError, legacy_transport_error)


def _raise_transport_error(error: Exception, *, request: httpx2.Request) -> NoReturn:
    if isinstance(error, timeout_exceptions()):
        raise APITimeoutError(request=request) from error
    raise APIConnectionError(request=request) from error


class _X509WorkloadIdentityAuth(_WorkloadIdentityAuth[X509WorkloadIdentity]):
    def __init__(self, *, workload_identity: X509WorkloadIdentity, max_retries: int) -> None:
        _validate_identity(workload_identity)
        super().__init__(workload_identity=workload_identity, token_exchange_url=_X509_TOKEN_EXCHANGE_URL)
        self._max_exchange_retries = min(max(max_retries, 0), _MAX_EXCHANGE_RETRIES)
        self._follow_redirects = False

    @override
    def _handle_token_response(self, response: httpx2.Response) -> dict[str, Any]:
        if response.status_code not in (400, 401, 403):
            token_data = super()._handle_token_response(response)
            response_body = response.json()
            token_type = response_body.get("token_type")
            if "token_type" in response_body and (not isinstance(token_type, str) or token_type.lower() != "bearer"):
                raise OpenAIError("X.509 token exchange response must use the Bearer token type")
            if _BEARER_ACCESS_TOKEN.fullmatch(token_data["access_token"]) is None:
                raise OpenAIError("X.509 token exchange response did not include a valid Bearer access_token")
            return token_data

        try:
            response_body = response.json() if response.content else None
        except ValueError:
            response_body = None

        oauth_error = response_body.get("error") if is_dict(response_body) else None
        safe_body = {"error": oauth_error} if isinstance(oauth_error, str) else None
        raise OAuthError(response=response, body=safe_body)

    @override
    def _validate_expires_in(self, expires_in: object) -> float:
        expires_in_value = _as_finite_float(expires_in)
        if expires_in_value is None or expires_in_value <= 0:
            raise OpenAIError("X.509 token exchange response did not include a positive, finite expires_in")
        return expires_in_value

    @override
    def _can_retry_request(self, request: httpx2.Request) -> bool:
        return _is_replayable_request(request)

    @override
    def _prepare_retry_request(self, request: httpx2.Request) -> None:
        file_positions = request.extensions.get(_REPLAY_FILE_POSITIONS_EXTENSION)
        if isinstance(file_positions, list):
            for file, file_position in cast(list[tuple[object, int]], file_positions):
                seek = getattr(file, "seek", None)
                if callable(seek):
                    seek(file_position)
            return

        position = request.extensions.get(_REPLAY_POSITION_EXTENSION)
        if not isinstance(position, int):
            return
        source = getattr(request.stream, "_stream", request.stream)
        seek = getattr(source, "seek", None)
        if callable(seek):
            seek(position)

    def _usable_token_after_transient_failure(self) -> str | None:
        with self._lock:
            if self._token_unusable():
                return None
            self._cached_token_refresh_at_monotonic = time.monotonic() + INITIAL_RETRY_DELAY
            return self._cached_token

    @override
    def _perform_refresh(self) -> None:
        try:
            super()._perform_refresh()
        except (APIConnectionError, _TransientTokenExchangeError):
            if self._usable_token_after_transient_failure() is None:
                raise

    def _handle_exchange_response(self, response: httpx2.Response) -> dict[str, Any]:
        try:
            return self._handle_token_response(response)
        except OpenAIError as error:
            if (
                response.status_code in (408, 409, 429)
                or response.status_code >= 500
                or (response.status_code not in (400, 401, 403) and response.headers.get("x-should-retry") == "true")
            ):
                raise _TransientTokenExchangeError(error) from error
            raise


class SyncX509WorkloadIdentityAuth(_X509WorkloadIdentityAuth):
    _http_client: httpx2.Client

    def __init__(
        self, *, workload_identity: X509WorkloadIdentity, http_client: httpx2.Client, max_retries: int
    ) -> None:
        super().__init__(workload_identity=workload_identity, max_retries=max_retries)
        self._http_client = http_client

    def send_api_request(
        self,
        request: httpx2.Request,
        *,
        expected_origin: httpx2.URL,
        expected_authorization: str | None,
        stream: bool,
        **kwargs: Any,
    ) -> httpx2.Response:
        with _scoped_sync_client(
            self._http_client, expected_origin=expected_origin, expected_authorization=expected_authorization
        ) as scoped_client:
            kwargs.setdefault("auth", None)
            return scoped_client.send(request, stream=stream, **kwargs)

    def get_token_for_request(self, request: httpx2.Request) -> str:
        timeout_token = _EXCHANGE_REQUEST_TIMEOUT.set(request.extensions.get("timeout"))
        try:
            try:
                return self.get_token()
            except (APIConnectionError, _TransientTokenExchangeError) as error:
                token = self._usable_token_after_transient_failure()
                if token is None:
                    if isinstance(error, _TransientTokenExchangeError):
                        raise error.error from None
                    raise
                return token
        finally:
            _EXCHANGE_REQUEST_TIMEOUT.reset(timeout_token)

    @override
    def _fetch_token_from_exchange(self) -> dict[str, Any]:
        for attempt in range(self._max_exchange_retries + 1):
            exchange_request = _token_exchange_request(self.workload_identity, http_client=self._http_client)
            try:
                with _scoped_sync_client(
                    self._http_client,
                    expected_origin=httpx2.URL(_X509_TOKEN_EXCHANGE_URL),
                    token_exchange=True,
                ) as scoped_client:
                    response = scoped_client.send(
                        exchange_request,
                        auth=None,
                        follow_redirects=False,
                    )
            except _transport_errors() as error:
                if attempt >= self._max_exchange_retries:
                    _raise_transport_error(error, request=exchange_request)
                delay = _retry_delay(None, attempt)
            else:
                delay = _retry_delay(response, attempt)
                if attempt >= self._max_exchange_retries or delay is None:
                    return self._handle_exchange_response(response)

            if delay is not None:
                time.sleep(delay)

        raise AssertionError("X.509 token exchange retry loop exhausted unexpectedly")


class AsyncX509WorkloadIdentityAuth(_X509WorkloadIdentityAuth):
    _http_client: httpx2.AsyncClient

    def __init__(
        self, *, workload_identity: X509WorkloadIdentity, http_client: httpx2.AsyncClient, max_retries: int
    ) -> None:
        super().__init__(workload_identity=workload_identity, max_retries=max_retries)
        self._http_client = http_client
        self._async_lock = anyio.Lock()

    async def send_api_request(
        self,
        request: httpx2.Request,
        *,
        expected_origin: httpx2.URL,
        expected_authorization: str | None,
        stream: bool,
        **kwargs: Any,
    ) -> httpx2.Response:
        async with _scoped_async_client(
            self._http_client, expected_origin=expected_origin, expected_authorization=expected_authorization
        ) as scoped_client:
            kwargs.setdefault("auth", None)
            return await scoped_client.send(request, stream=stream, **kwargs)

    async def get_token_for_request(self, request: httpx2.Request) -> str:
        timeout_token = _EXCHANGE_REQUEST_TIMEOUT.set(request.extensions.get("timeout"))
        try:
            try:
                return await self.get_token_async()
            except (APIConnectionError, _TransientTokenExchangeError) as error:
                token = self._usable_token_after_transient_failure()
                if token is None:
                    if isinstance(error, _TransientTokenExchangeError):
                        raise error.error from None
                    raise
                return token
        finally:
            _EXCHANGE_REQUEST_TIMEOUT.reset(timeout_token)

    @override
    async def get_token_async(self) -> str:
        async with self._async_lock:
            with self._lock:
                if not self._token_unusable() and not self._needs_refresh():
                    return cast(str, self._cached_token)

            try:
                token_data = await self._fetch_token_from_exchange_async()
            except (APIConnectionError, _TransientTokenExchangeError):
                token = self._usable_token_after_transient_failure()
                if token is None:
                    raise
                return token
            self._store_token(token_data)
            with self._lock:
                return cast(str, self._cached_token)

    async def _fetch_token_from_exchange_async(self) -> dict[str, Any]:
        for attempt in range(self._max_exchange_retries + 1):
            exchange_request = _token_exchange_request(self.workload_identity, http_client=self._http_client)
            try:
                async with _scoped_async_client(
                    self._http_client,
                    expected_origin=httpx2.URL(_X509_TOKEN_EXCHANGE_URL),
                    token_exchange=True,
                ) as scoped_client:
                    response = await scoped_client.send(
                        exchange_request,
                        auth=None,
                        follow_redirects=False,
                    )
            except _transport_errors() as error:
                if attempt >= self._max_exchange_retries:
                    _raise_transport_error(error, request=exchange_request)
                delay = _retry_delay(None, attempt)
            else:
                delay = _retry_delay(response, attempt)
                if attempt >= self._max_exchange_retries or delay is None:
                    return self._handle_exchange_response(response)

            if delay is not None:
                await anyio.sleep(delay)

        raise AssertionError("X.509 token exchange retry loop exhausted unexpectedly")