# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import base64 import contextlib import errno import ipaddress import json import math import os import pathlib import sys import tempfile import threading import time from collections import namedtuple from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Generic, Never, SupportsIndex, TypeVar, cast from urllib.parse import urlparse import grpc import httpx from ._proto import ( datamodel_pb2, openshell_pb2, openshell_pb2_grpc, ) from .errors import _error_mapping_channel from .mutations import DeletionOutcome, DeletionResult _ClientCallDetailsBase = namedtuple( "_ClientCallDetailsBase", ("method", "timeout", "metadata", "credentials", "wait_for_ready", "compression"), ) _OAUTH_MAX_RESPONSE_BYTES = 1 << 20 _PAGER_MAX_CONSUMED_TOKENS = 10_000 _PAGER_MAX_CONSUMED_TOKEN_BYTES = 1 << 20 T = TypeVar("T") @dataclass(frozen=True) class Page(Generic[T]): """One response page from a list operation.""" items: builtins.list[T] next_page_token: str class Pager(Generic[T]): """Lazy, single-pass iterator over the continuation-token contract. The repeated-token guard has bounded memory and raises ``SandboxError`` if the traversal exceeds that guard's token-count or byte budget. """ def __init__(self, fetch: Callable[[str], Page[T]], page_token: str = "") -> None: self._fetch = fetch self._page_token: str | None = page_token self._consumed_page_tokens: set[str] = set() self._consumed_page_token_bytes = 0 def __iter__(self) -> Pager[T]: return self def _validate_page_token_budget(self, page_token: str) -> int: if not page_token: return 0 token_bytes = len(page_token.encode("utf-8")) if ( len(self._consumed_page_tokens) >= _PAGER_MAX_CONSUMED_TOKENS or self._consumed_page_token_bytes + token_bytes > _PAGER_MAX_CONSUMED_TOKEN_BYTES ): raise SandboxError("pager continuation token history limit exceeded") return token_bytes def __next__(self) -> Page[T]: if self._page_token is None: raise StopIteration page_token = self._page_token token_bytes = self._validate_page_token_budget(page_token) page = self._fetch(page_token) if page_token: self._consumed_page_tokens.add(page_token) self._consumed_page_token_bytes += token_bytes next_page_token = page.next_page_token if next_page_token and next_page_token in self._consumed_page_tokens: raise SandboxError("pager received a repeated continuation token") self._page_token = next_page_token or None return page def all(self) -> builtins.list[T]: """Consume the pager and collect every remaining item.""" return [item for page in self for item in page.items] def _workspace_scope(workspace: str) -> datamodel_pb2.WorkspaceSelector: if not workspace: raise ValueError("workspace must be non-empty") return datamodel_pb2.WorkspaceSelector(workspace=workspace) def _all_workspaces_scope() -> datamodel_pb2.WorkspaceSelector: return datamodel_pb2.WorkspaceSelector(all_workspaces=datamodel_pb2.AllWorkspaces()) def _service_exposure_messages( exposures: Sequence[ServiceExposure] | None, ) -> list[openshell_pb2.SandboxServiceExposure]: return [ openshell_pb2.SandboxServiceExposure( service=exposure.service, target_port=exposure.target_port, ) for exposure in exposures or () ] class _ClientCallDetails(_ClientCallDetailsBase, grpc.ClientCallDetails): pass if TYPE_CHECKING: import builtins from collections.abc import Callable, Iterator, Mapping, Sequence @dataclass(frozen=True) class TlsConfig: """Channel TLS material. All three fields are optional so callers can pick the trust profile: - Full mTLS: pass all three (server trusts client identity). - CA-only: pass `ca_path` (custom CA, no client identity). - System roots: pass no fields (`TlsConfig()`) — uses the OS trust store. Useful for OIDC gateways behind a public CA. `cert_path` and `key_path` must be set together or not at all. """ ca_path: pathlib.Path | None = None cert_path: pathlib.Path | None = None key_path: pathlib.Path | None = None def __post_init__(self) -> None: if (self.cert_path is None) != (self.key_path is None): raise ValueError("TlsConfig: cert_path and key_path must be set together") class _BearerAuthInterceptor( grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor, grpc.StreamUnaryClientInterceptor, grpc.StreamStreamClientInterceptor, ): """Add `authorization: Bearer ` to every outgoing RPC. Implemented as an interceptor (not call credentials) so it works on both plaintext and TLS channels without needing `grpc.composite_channel_credentials`. The token provider is invoked per call, so callers can swap tokens at runtime by mutating shared state or returning a fresh value from the callable. """ def __init__(self, token_provider: Callable[[], str]) -> None: self._token_provider = token_provider def _attach(self, details: grpc.ClientCallDetails) -> grpc.ClientCallDetails: original_metadata = getattr(details, "metadata", None) metadata = list(original_metadata) if original_metadata else [] metadata.append(("authorization", f"Bearer {self._token_provider()}")) return _ClientCallDetails( getattr(details, "method", None), getattr(details, "timeout", None), metadata, getattr(details, "credentials", None), getattr(details, "wait_for_ready", None), getattr(details, "compression", None), ) def intercept_unary_unary(self, continuation, client_call_details, request): return continuation(self._attach(client_call_details), request) def intercept_unary_stream(self, continuation, client_call_details, request): return continuation(self._attach(client_call_details), request) def intercept_stream_unary( self, continuation, client_call_details, request_iterator ): return continuation(self._attach(client_call_details), request_iterator) def intercept_stream_stream( self, continuation, client_call_details, request_iterator ): return continuation(self._attach(client_call_details), request_iterator) def _normalize_bearer( bearer: str | Callable[[], str] | None, ) -> Callable[[], str] | None: if bearer is None: return None if callable(bearer): return cast("Callable[[], str]", bearer) token = bearer return lambda: token def _is_loopback_host(hostname: str | None) -> bool: if hostname is None: return False if hostname.lower() == "localhost": return True with contextlib.suppress(ValueError): return ipaddress.ip_address(hostname).is_loopback return False def _is_local_grpc_endpoint(endpoint: str) -> bool: if endpoint.startswith("unix:"): return True return _is_loopback_host(urlparse(f"//{endpoint}").hostname) def _validate_oauth_url(name: str, raw: str) -> str: """Validate an OAuth endpoint without reflecting attacker-controlled URLs.""" parsed = urlparse(raw) if not parsed.scheme or not parsed.hostname or parsed.username or parsed.password: raise SandboxError(f"invalid OAuth {name} URL") if parsed.fragment: raise SandboxError(f"OAuth {name} URL must not contain a fragment") if parsed.scheme == "https": return raw if parsed.scheme == "http" and _is_loopback_host(parsed.hostname): return raw raise SandboxError( f"OAuth {name} URL must use HTTPS (HTTP is allowed only for loopback hosts)" ) def _oauth_json_request( client: httpx.Client, method: str, url: str, *, kind: str, headers: Mapping[str, str], data: Mapping[str, str] | None = None, ) -> tuple[int, dict[str, object]]: """Read a bounded JSON response without retaining arbitrary error bodies.""" with client.stream(method, url, headers=headers, data=data) as response: if response.status_code != 200: return response.status_code, {} content = bytearray() for chunk in response.iter_bytes(): content.extend(chunk) if len(content) > _OAUTH_MAX_RESPONSE_BYTES: raise SandboxError(f"OAuth {kind} response is too large") try: payload = json.loads(content) except (UnicodeDecodeError, json.JSONDecodeError): raise SandboxError(f"OAuth {kind} returned invalid JSON") from None if not isinstance(payload, dict): raise SandboxError(f"OAuth {kind} returned invalid JSON") return 200, payload class ClientCredentialsAuth: """Renewable OAuth 2.0 client-credentials bearer provider. Tokens and client secrets remain in memory. The first RPC lazily performs discovery and an exchange; later RPCs share the cached token until it is within 30 seconds of expiry. Concurrent callers share one exchange. """ def __init__( self, *, client_secret: str | Callable[[], str], issuer: str | None = None, client_id: str | None = None, scopes: Sequence[str] | None = None, audience: str | None = None, insecure: bool = False, timeout: float = 30.0, _transport: httpx.BaseTransport | None = None, ) -> None: if not isinstance(client_secret, str) and not callable(client_secret): raise TypeError("client_secret must be a string or zero-argument callable") if isinstance(client_secret, str) and not client_secret: raise SandboxError("OAuth client secret must not be empty") self._secret = client_secret self._issuer = issuer self._client_id = client_id self._scopes = tuple(scopes) if scopes is not None else None self._audience = audience self._insecure = insecure self._timeout = timeout self._transport = _transport self._lock = threading.Lock() self._access_token: str | None = None self._expires_at = 0.0 self._token_endpoint: str | None = None def __repr__(self) -> str: return ( "ClientCredentialsAuth(issuer=" f"{self._issuer!r}, client_id={self._client_id!r}, " f"scopes={self._scopes!r}, audience={self._audience!r})" ) def _apply_gateway_metadata( self, metadata: Mapping[str, object], *, insecure: bool = False ) -> None: """Fill omitted fields and apply active-gateway transport settings.""" if self._issuer is None: value = metadata.get("oidc_issuer") self._issuer = value if isinstance(value, str) and value else None if self._client_id is None: value = metadata.get("oidc_client_id") self._client_id = value if isinstance(value, str) and value else None if self._audience is None: value = metadata.get("oidc_audience") self._audience = value if isinstance(value, str) and value else None if self._scopes is None: value = metadata.get("oidc_scopes") if isinstance(value, str): self._scopes = tuple(value.split()) # Preserve an explicit provider-level opt-in while also honoring the # active-gateway construction flag used by the existing OIDC refresher. self._insecure = self._insecure or insecure def __call__(self) -> str: if ( self._access_token is not None and time.time() + _OIDC_TOKEN_EXPIRY_GRACE_SECONDS < self._expires_at ): return self._access_token with self._lock: if ( self._access_token is not None and time.time() + _OIDC_TOKEN_EXPIRY_GRACE_SECONDS < self._expires_at ): return self._access_token token, expires_at = self._exchange() self._access_token = token self._expires_at = expires_at return token def _exchange(self) -> tuple[str, float]: if not self._issuer or not self._client_id: raise SandboxError("OAuth client credentials require issuer and client_id") issuer = _validate_oauth_url("issuer", self._issuer).rstrip("/") headers = {"accept": "application/json"} try: with httpx.Client( verify=not self._insecure, follow_redirects=False, timeout=self._timeout, transport=self._transport, ) as client: if self._token_endpoint is None: status, discovery = _oauth_json_request( client, "GET", f"{issuer}/.well-known/openid-configuration", kind="discovery", headers=headers, ) if status != 200: raise SandboxError(f"OAuth discovery failed with HTTP {status}") discovered = discovery.get("issuer") if ( not isinstance(discovered, str) or discovered.rstrip("/") != issuer ): raise SandboxError("OAuth discovery issuer mismatch") endpoint = discovery.get("token_endpoint") if not isinstance(endpoint, str) or not endpoint: raise SandboxError( "OAuth discovery document is missing token_endpoint" ) self._token_endpoint = _validate_oauth_url( "token endpoint", endpoint ) secret = self._resolve_secret() data = { "grant_type": "client_credentials", "client_id": self._client_id, "client_secret": secret, } if self._scopes: data["scope"] = " ".join(self._scopes) if self._audience: data["audience"] = self._audience status, payload = _oauth_json_request( client, "POST", self._token_endpoint, kind="client credentials response", data=data, headers=headers, ) except SandboxError: raise except httpx.HTTPError: raise SandboxError("OAuth client credentials request failed") from None if status != 200: raise SandboxError( f"OAuth client credentials exchange failed with HTTP {status}" ) access_token = payload.get("access_token") expires_in = payload.get("expires_in") if not isinstance(access_token, str) or not access_token: raise SandboxError( "OAuth client credentials response is missing access_token" ) if ( isinstance(expires_in, bool) or not isinstance(expires_in, (int, float)) or not math.isfinite(expires_in) or expires_in <= 0 ): raise SandboxError( "OAuth client credentials response requires a positive finite expires_in" ) return access_token, time.time() + float(expires_in) def _resolve_secret(self) -> str: if isinstance(self._secret, str): return self._secret try: value = self._secret() except Exception: raise SandboxError("OAuth client secret supplier failed") from None if not isinstance(value, str) or not value: raise SandboxError("OAuth client secret supplier returned an empty secret") return value @dataclass(frozen=True) class SandboxStatusRef: phase: int current_policy_version: int exit_code: int | None = None restart_count: int = 0 next_restart_at_ms: int | None = None main_process_started_at_ms: int | None = None @dataclass(frozen=True) class ServiceExposure: """A loopback HTTP service to expose during sandbox creation.""" target_port: int service: str = "" class _ImmutableLabels(dict[str, str]): """A read-only, copy- and pickle-safe label mapping.""" def _deny_mutation(self, *args: object, **kwargs: object) -> Never: del args, kwargs raise TypeError("sandbox labels are immutable") __setitem__ = _deny_mutation __delitem__ = _deny_mutation clear = _deny_mutation pop = _deny_mutation popitem = _deny_mutation setdefault = _deny_mutation update = _deny_mutation __ior__ = _deny_mutation def __deepcopy__(self, memo: dict[int, object]) -> _ImmutableLabels: del memo return type(self)(self) def __reduce_ex__(self, protocol: SupportsIndex, /) -> str | tuple[Any, ...]: del protocol return type(self), (dict(self),) @dataclass(frozen=True) class SandboxWorkloadTemplateProvenanceRef: name: str resource_version: str @dataclass(frozen=True) class SandboxRef: id: str name: str workspace: str status: SandboxStatusRef # Excluded from equality/hash to preserve the original identity while the # immutable mapping remains safe for deepcopy, pickle, and asdict. labels: Mapping[str, str] = field(default_factory=_ImmutableLabels, compare=False) created_from_workload_template: SandboxWorkloadTemplateProvenanceRef | None = None # Populated by create operations. The empty key identifies the unnamed service. service_urls: Mapping[str, str] = field( default_factory=_ImmutableLabels, compare=False ) def __post_init__(self) -> None: object.__setattr__(self, "labels", _ImmutableLabels(self.labels)) object.__setattr__(self, "service_urls", _ImmutableLabels(self.service_urls)) @property def phase(self) -> int: return self.status.phase @property def current_policy_version(self) -> int: return self.status.current_policy_version @dataclass(frozen=True) class ExecChunk: stream: str data: bytes @dataclass(frozen=True) class ExecResult: exit_code: int stdout: str stderr: str class SandboxError(RuntimeError): pass class SandboxSession: def __init__(self, client: SandboxClient, sandbox: SandboxRef) -> None: self._client = client self.sandbox = sandbox self._workspace = sandbox.workspace @property def id(self) -> str: return self.sandbox.id def exec( self, command: Sequence[str], *, stream_output: bool = False, workdir: str | None = None, env: Mapping[str, str] | None = None, stdin: bytes | None = None, timeout_seconds: int | None = None, no_login_shell: bool = False, ) -> ExecResult: return self._client.exec( self.sandbox.name, command, workspace=self._workspace, stream_output=stream_output, workdir=workdir, env=env, stdin=stdin, timeout_seconds=timeout_seconds, no_login_shell=no_login_shell, ) def exec_python( self, function: Callable[..., object], *, args: Sequence[object] = (), kwargs: Mapping[str, object] | None = None, stream_output: bool = False, workdir: str | None = None, env: Mapping[str, str] | None = None, timeout_seconds: int | None = None, ) -> ExecResult: return self._client.exec_python( self.sandbox.name, function, workspace=self._workspace, args=args, kwargs=kwargs, stream_output=stream_output, workdir=workdir, env=env, timeout_seconds=timeout_seconds, ) def delete(self, *, allow_missing: bool = False) -> DeletionResult: return self._client.delete( self.sandbox.name, workspace=self._workspace, allow_missing=allow_missing ) def stop(self) -> SandboxRef: self.sandbox = self._client.stop(self.sandbox.name, workspace=self._workspace) return self.sandbox def start(self) -> SandboxRef: self.sandbox = self._client.start(self.sandbox.name, workspace=self._workspace) return self.sandbox class SandboxClient: """gRPC client for sandbox CRUD and command execution.""" def __init__( self, endpoint: str, *, tls: TlsConfig | None = None, bearer_token: str | Callable[[], str] | None = None, client_credentials: ClientCredentialsAuth | None = None, timeout: float = 30.0, cluster_name: str | None = None, _bearer_close: Callable[[], None] | None = None, ) -> None: """Create a SandboxClient. Args: endpoint: host:port for the gateway gRPC service. tls: mTLS material. None for a plaintext channel. bearer_token: OIDC access token, or a zero-arg callable returning the current token (called per RPC; supports runtime refresh). Combines with `tls` — pass both when the gateway uses mTLS for transport identity and OIDC for user identity. client_credentials: renewable OAuth client-credentials provider. Mutually exclusive with `bearer_token`. A non-loopback endpoint requires `tls` so the acquired bearer is never sent in cleartext. timeout: default per-call timeout in seconds. cluster_name: optional friendly name for error messages. _bearer_close: internal — wired by `from_active_cluster` when an `_OidcRefresher` owns the bearer callable, so `close()` can release the refresher's HTTP client. Public callers should not pass this; they own the lifecycle of any callable they supplied as `bearer_token`. """ if bearer_token is not None and client_credentials is not None: raise SandboxError( "bearer_token and client_credentials are mutually exclusive" ) if ( client_credentials is not None and tls is None and not _is_local_grpc_endpoint(endpoint) ): raise SandboxError( "OAuth client credentials require TLS for non-loopback gateway endpoints" ) self._endpoint = endpoint self._timeout = timeout self._cluster_name = cluster_name self._bearer_close = _bearer_close if tls is None: self._channel = grpc.insecure_channel(endpoint) else: # Build credentials from whatever subset of mTLS material the # caller supplied. None for `root_certificates` makes gRPC use # the system trust store, which is what we want for OIDC # gateways behind a public CA. credentials = grpc.ssl_channel_credentials( root_certificates=(tls.ca_path.read_bytes() if tls.ca_path else None), private_key=(tls.key_path.read_bytes() if tls.key_path else None), certificate_chain=( tls.cert_path.read_bytes() if tls.cert_path else None ), ) self._channel = grpc.secure_channel(endpoint, credentials) provider = _normalize_bearer(client_credentials or bearer_token) if provider is not None: self._channel = grpc.intercept_channel( self._channel, _BearerAuthInterceptor(provider), ) self._stub = openshell_pb2_grpc.OpenShellStub( _error_mapping_channel(self._channel) ) @classmethod def from_active_cluster( cls, *, cluster: str | None = None, timeout: float = 30.0, auto_refresh: bool = True, write_back: bool = True, insecure: bool = False, client_credentials: ClientCredentialsAuth | None = None, ) -> SandboxClient: """Construct a `SandboxClient` from the active gateway's on-disk state. Args: cluster: explicit gateway name; otherwise reads `$OPENSHELL_GATEWAY` or `~/.config/openshell/active_gateway`. timeout: per-call gRPC timeout in seconds. auto_refresh: when True (default) and the gateway uses OIDC, lazily refresh the access token via the IdP's token endpoint if the cached `oidc_token.json` is near expiry. Matches the lazy-refresh patterns used by `google-auth` and `botocore`. Set False to keep the SDK as a read-only consumer of the CLI's cache (fail closed on expiry). write_back: when True (default, and `auto_refresh=True`), atomically persist refreshed bundles back to `oidc_token.json` so other processes — including the Rust CLI — see the rotation. Required for IdPs with refresh-token rotation enabled (Keycloak, Entra in strict mode): an in-memory-only refresh would leave the on-disk `refresh_token` pointing at an invalidated value, and any other process starting from that disk state would fail on its first refresh. Set False only when you know the SDK is the sole consumer of this gateway directory. insecure: when True, disables TLS certificate verification for OIDC discovery and refresh calls. Mirrors the Rust CLI's `--insecure` flag for issuers behind self-signed certs. Off by default. client_credentials: renewable OAuth client-credentials provider. Omitted issuer, client ID, audience, and scopes are filled from the registered gateway metadata. This provider does not read or write `oidc_token.json`. Remote plaintext gateways are rejected. """ cluster_name = cluster or _resolve_active_cluster() gateway_dir = _xdg_config_home() / "openshell" / "gateways" / cluster_name metadata_path = gateway_dir / "metadata.json" try: metadata = json.loads(metadata_path.read_text(encoding="utf-8")) except FileNotFoundError: raise SandboxError(f"gateway '{cluster_name}' not found") from None if "gateway_endpoint" not in metadata: raise SandboxError(f"gateway '{cluster_name}' metadata missing endpoint") parsed = urlparse(metadata["gateway_endpoint"]) host = parsed.hostname or "127.0.0.1" port = parsed.port or (443 if parsed.scheme == "https" else 80) endpoint = f"{host}:{port}" # TLS transport. Mirror crates/openshell-tui/src/lib.rs # `build_oidc_channel` — for an https gateway, always build a # secure channel and pick the strongest available trust profile. tls: TlsConfig | None = None if parsed.scheme == "https": mtls_dir = gateway_dir / "mtls" ca = mtls_dir / "ca.crt" if (mtls_dir / "ca.crt").exists() else None cert = mtls_dir / "tls.crt" if (mtls_dir / "tls.crt").exists() else None key = mtls_dir / "tls.key" if (mtls_dir / "tls.key").exists() else None if ca is not None and cert is not None and key is not None: # Full mTLS. tls = TlsConfig(ca_path=ca, cert_path=cert, key_path=key) elif ca is not None: # CA-only trust (no client identity). tls = TlsConfig(ca_path=ca) else: # System roots (e.g. OIDC gateway behind a public CA). tls = TlsConfig() # OIDC bearer. Mirror the Rust CLI/TUI: the gateway metadata's # `auth_mode` is authoritative — a stale oidc_token.json next to # a non-OIDC gateway should NOT cause us to attach a bearer. bearer_token: Callable[[], str] | None = None bearer_close: Callable[[], None] | None = None if client_credentials is not None: if metadata.get("auth_mode") != "oidc": raise SandboxError( f"gateway '{cluster_name}' is not configured for OIDC" ) client_credentials._apply_gateway_metadata(metadata, insecure=insecure) elif metadata.get("auth_mode") == "oidc": bearer_token, bearer_close = _make_cluster_bearer_provider( gateway_dir, cluster_name, auto_refresh=auto_refresh, write_back=write_back, insecure=insecure, ) return cls( endpoint, tls=tls, bearer_token=bearer_token, client_credentials=client_credentials, timeout=timeout, cluster_name=cluster_name, _bearer_close=bearer_close, ) def close(self) -> None: """Release the gRPC channel and any bearer-auth resources. Idempotent. If `from_active_cluster` wired up an OIDC refresher for this client, the refresher's underlying httpx.Client is closed here too — otherwise long-lived services that churn clients would leak sockets / file descriptors until GC. """ self._channel.close() if self._bearer_close is not None: with contextlib.suppress(Exception): self._bearer_close() self._bearer_close = None def __enter__(self) -> SandboxClient: return self def __exit__(self, *args: object) -> None: self.close() def health(self) -> openshell_pb2.HealthResponse: return self._stub.Health(openshell_pb2.HealthRequest(), timeout=self._timeout) def create( self, *, workspace: str, spec: openshell_pb2.SandboxSpec | None = None, name: str | None = None, labels: Mapping[str, str] | None = None, service_exposures: Sequence[ServiceExposure] | None = None, ) -> SandboxRef: request_spec = spec if spec is not None else _default_spec() response = self._stub.CreateSandbox( openshell_pb2.CreateSandboxRequest( workspace_scope=_workspace_scope(workspace), spec=request_spec, name=name or "", labels=dict(labels) if labels else {}, service_exposures=_service_exposure_messages(service_exposures), ), timeout=self._timeout, ) sandbox_ref = _sandbox_ref(response.sandbox, response.service_urls) if sandbox_ref.id == "": raise SandboxError("CreateSandbox returned empty sandbox id") return sandbox_ref def create_from_template( self, *, workspace: str, workload_template: str, spec: openshell_pb2.SandboxSpec | None = None, name: str | None = None, labels: Mapping[str, str] | None = None, service_exposures: Sequence[ServiceExposure] | None = None, ) -> SandboxRef: if not workload_template.strip(): raise SandboxError("workload_template is required") request_spec = spec if spec is not None else openshell_pb2.SandboxSpec() response = self._stub.CreateSandbox( openshell_pb2.CreateSandboxRequest( workspace_scope=_workspace_scope(workspace), spec=request_spec, name=name or "", labels=dict(labels) if labels else {}, workload_template=workload_template, service_exposures=_service_exposure_messages(service_exposures), ), timeout=self._timeout, ) sandbox_ref = _sandbox_ref(response.sandbox, response.service_urls) if sandbox_ref.id == "": raise SandboxError("CreateSandbox returned empty sandbox id") return sandbox_ref def create_session( self, *, workspace: str, spec: openshell_pb2.SandboxSpec | None = None, name: str | None = None, labels: Mapping[str, str] | None = None, service_exposures: Sequence[ServiceExposure] | None = None, ) -> SandboxSession: return SandboxSession( self, self.create( workspace=workspace, spec=spec, name=name, labels=labels, service_exposures=service_exposures, ), ) def create_session_from_template( self, *, workspace: str, workload_template: str, spec: openshell_pb2.SandboxSpec | None = None, name: str | None = None, labels: Mapping[str, str] | None = None, service_exposures: Sequence[ServiceExposure] | None = None, ) -> SandboxSession: return SandboxSession( self, self.create_from_template( workspace=workspace, workload_template=workload_template, spec=spec, name=name, labels=labels, service_exposures=service_exposures, ), ) def sandbox_templates(self) -> SandboxTemplateClient: return SandboxTemplateClient(self._channel, timeout=self._timeout) def get(self, name: str, *, workspace: str) -> SandboxRef: response = self._stub.GetSandbox( openshell_pb2.GetSandboxRequest( workspace_scope=_workspace_scope(workspace), name=name, ), timeout=self._timeout, ) return _sandbox_ref(response.sandbox) def get_session(self, name: str, *, workspace: str) -> SandboxSession: return SandboxSession(self, self.get(name, workspace=workspace)) def list( self, *, workspace: str, page_size: int = 100, page_token: str = "", label_selector: str | None = None, ) -> Pager[SandboxRef]: def fetch(token: str) -> Page[SandboxRef]: response = self._stub.ListSandboxes( openshell_pb2.ListSandboxesRequest( workspace_scope=_workspace_scope(workspace), page_size=page_size, page_token=token, label_selector=label_selector or "", ), timeout=self._timeout, ) return Page( items=[_sandbox_ref(item) for item in response.sandboxes], next_page_token=getattr(response, "next_page_token", ""), ) return Pager(fetch, page_token) def list_all( self, *, workspace: str, page_size: int = 100, page_token: str = "", label_selector: str | None = None, ) -> builtins.list[SandboxRef]: return self.list( workspace=workspace, page_size=page_size, page_token=page_token, label_selector=label_selector, ).all() def list_for_all_workspaces( self, *, page_size: int = 100, page_token: str = "", label_selector: str | None = None, ) -> Pager[SandboxRef]: def fetch(token: str) -> Page[SandboxRef]: response = self._stub.ListSandboxes( openshell_pb2.ListSandboxesRequest( workspace_scope=_all_workspaces_scope(), page_size=page_size, page_token=token, label_selector=label_selector or "", ), timeout=self._timeout, ) return Page( items=[_sandbox_ref(item) for item in response.sandboxes], next_page_token=getattr(response, "next_page_token", ""), ) return Pager(fetch, page_token) def list_all_for_all_workspaces( self, *, page_size: int = 100, page_token: str = "", label_selector: str | None = None, ) -> builtins.list[SandboxRef]: return self.list_for_all_workspaces( page_size=page_size, page_token=page_token, label_selector=label_selector, ).all() def list_ids( self, *, workspace: str, page_size: int = 100, label_selector: str | None = None, ) -> builtins.list[str]: return [ item.id for item in self.list_all( workspace=workspace, page_size=page_size, label_selector=label_selector, ) ] def list_ids_for_all_workspaces( self, *, page_size: int = 100, label_selector: str | None = None, ) -> builtins.list[str]: return [ item.id for item in self.list_all_for_all_workspaces( page_size=page_size, label_selector=label_selector, ) ] def delete( self, name: str, *, workspace: str, allow_missing: bool = False ) -> DeletionResult: response = self._stub.DeleteSandbox( openshell_pb2.DeleteSandboxRequest( workspace_scope=_workspace_scope(workspace), name=name, allow_missing=allow_missing, ), timeout=self._timeout, ) return DeletionResult( DeletionOutcome(response.outcome), response.sandbox_id or None ) def stop(self, name: str, *, workspace: str) -> SandboxRef: response = self._stub.StopSandbox( openshell_pb2.StopSandboxRequest( workspace_scope=_workspace_scope(workspace), name=name, ), timeout=self._timeout, ) return _sandbox_ref(response.sandbox) def start(self, name: str, *, workspace: str) -> SandboxRef: response = self._stub.StartSandbox( openshell_pb2.StartSandboxRequest( workspace_scope=_workspace_scope(workspace), name=name, ), timeout=self._timeout, ) return _sandbox_ref(response.sandbox) def wait_deleted( self, name: str, *, workspace: str, timeout_seconds: float = 60.0, expected_sandbox_id: str | None = None, ) -> None: deadline = time.time() + timeout_seconds while time.time() < deadline: try: current = self.get(name, workspace=workspace) if ( expected_sandbox_id is not None and current.id != expected_sandbox_id ): return except grpc.RpcError as exc: if getattr(exc, "code", lambda: None)() == grpc.StatusCode.NOT_FOUND: return raise time.sleep(1) raise SandboxError(f"sandbox {name} was not deleted within timeout") def wait_ready( self, name: str, *, workspace: str, timeout_seconds: float = 300.0 ) -> SandboxRef: return self._wait_for_phase( name, workspace=workspace, target_phase=openshell_pb2.SANDBOX_PHASE_READY, target_name="ready", timeout_seconds=timeout_seconds, ) def wait_stopped( self, name: str, *, workspace: str, timeout_seconds: float = 300.0 ) -> SandboxRef: return self._wait_for_phase( name, workspace=workspace, target_phase=openshell_pb2.SANDBOX_PHASE_STOPPED, target_name="stopped", timeout_seconds=timeout_seconds, ) def _wait_for_phase( self, name: str, *, workspace: str, target_phase: int, target_name: str, timeout_seconds: float, ) -> SandboxRef: deadline = time.time() + timeout_seconds while time.time() < deadline: sandbox = self.get(name, workspace=workspace) if sandbox.status.phase == target_phase: return sandbox if ( target_phase == openshell_pb2.SANDBOX_PHASE_READY and sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_COMPLETED ): return sandbox if ( target_phase == openshell_pb2.SANDBOX_PHASE_READY and sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_STOPPED ): raise SandboxError(f"sandbox {name} stopped before becoming ready") if sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_ERROR: raise SandboxError(f"sandbox {name} entered error phase") time.sleep(1) raise SandboxError(f"sandbox {name} was not {target_name} within timeout") def exec_stream( self, sandbox: str, command: Sequence[str], *, workspace: str, workdir: str | None = None, env: Mapping[str, str] | None = None, stdin: bytes | None = None, timeout_seconds: int | None = None, no_login_shell: bool = False, ) -> Iterator[ExecChunk | ExecResult]: if not command: raise SandboxError("command must not be empty") request = openshell_pb2.ExecSandboxRequest( workspace_scope=_workspace_scope(workspace), sandbox=sandbox, command=list(command), workdir=workdir or "", environment=dict(env or {}), stdin=stdin or b"", no_login_shell=no_login_shell, ) if timeout_seconds: request.execution_timeout.seconds = timeout_seconds # Use whichever is larger: the default client timeout or the command # timeout plus headroom for SSH setup / teardown overhead. grpc_deadline = self._timeout if timeout_seconds and timeout_seconds + 10 > grpc_deadline: grpc_deadline = timeout_seconds + 10 stream = self._stub.ExecSandbox(request, timeout=grpc_deadline) stdout_parts: list[bytes] = [] stderr_parts: list[bytes] = [] exit_code: int | None = None for event in stream: payload = event.WhichOneof("payload") if payload == "stdout": data = bytes(event.stdout.data) stdout_parts.append(data) yield ExecChunk(stream="stdout", data=data) elif payload == "stderr": data = bytes(event.stderr.data) stderr_parts.append(data) yield ExecChunk(stream="stderr", data=data) elif payload == "exit": exit_code = int(event.exit.exit_code) if exit_code is None: raise SandboxError("ExecSandbox stream ended without an exit event") yield ExecResult( exit_code=exit_code, stdout=b"".join(stdout_parts).decode("utf-8", errors="replace"), stderr=b"".join(stderr_parts).decode("utf-8", errors="replace"), ) def exec( self, sandbox: str, command: Sequence[str], *, workspace: str, stream_output: bool = False, workdir: str | None = None, env: Mapping[str, str] | None = None, stdin: bytes | None = None, timeout_seconds: int | None = None, no_login_shell: bool = False, ) -> ExecResult: result: ExecResult | None = None for item in self.exec_stream( sandbox, command, workspace=workspace, workdir=workdir, env=env, stdin=stdin, timeout_seconds=timeout_seconds, no_login_shell=no_login_shell, ): if stream_output and isinstance(item, ExecChunk): if item.stream == "stdout": sys.stdout.buffer.write(item.data) sys.stdout.flush() else: sys.stderr.buffer.write(item.data) sys.stderr.flush() if isinstance(item, ExecResult): result = item if result is None: raise SandboxError("ExecSandbox did not return a result") return result def exec_python( self, sandbox: str, function: Callable[..., object], *, workspace: str, args: Sequence[object] = (), kwargs: Mapping[str, object] | None = None, stream_output: bool = False, workdir: str | None = None, env: Mapping[str, str] | None = None, timeout_seconds: int | None = None, ) -> ExecResult: exec_env = dict(env or {}) exec_env["OPENSHELL_PYFUNC_B64"] = _serialize_python_callable( function, args=args, kwargs=kwargs, ) return self.exec( sandbox, [_SANDBOX_PYTHON_BIN, "-c", _PYTHON_CLOUDPICKLE_BOOTSTRAP], workspace=workspace, stream_output=stream_output, workdir=workdir, env=exec_env, timeout_seconds=timeout_seconds, ) class SandboxTemplateClient: """gRPC client for reusable sandbox template lifecycle operations.""" def __init__(self, channel: grpc.Channel, *, timeout: float = 30.0) -> None: self._stub = openshell_pb2_grpc.OpenShellStub(_error_mapping_channel(channel)) self._timeout = timeout @classmethod def from_sandbox_client(cls, client: SandboxClient) -> SandboxTemplateClient: return client.sandbox_templates() def create( self, *, workspace: str, template: openshell_pb2.SandboxWorkloadTemplate | None = None, name: str | None = None, image: str | None = None, labels: Mapping[str, str] | None = None, annotations: Mapping[str, str] | None = None, environment: Mapping[str, str] | None = None, cpu: str | None = None, memory: str | None = None, gpu_count: int | None = None, gpu: bool = False, driver_config: Mapping[str, Any] | None = None, ) -> openshell_pb2.SandboxWorkloadTemplate: if template is None: if name is None or not name.strip(): raise SandboxError("name is required when template is omitted") template = _sandbox_workload_template( name=name, image=image, labels=labels, annotations=annotations, environment=environment, cpu=cpu, memory=memory, gpu_count=gpu_count, gpu=gpu, driver_config=driver_config, ) elif ( any(value is not None for value in (name, image, cpu, memory, gpu_count)) or any( bool(value) for value in (labels, annotations, environment, driver_config) ) or gpu ): raise SandboxError( "template cannot be combined with template builder fields" ) response = self._stub.CreateSandboxTemplate( openshell_pb2.CreateSandboxTemplateRequest( workspace_scope=_workspace_scope(workspace), template=template, ), timeout=self._timeout, ) return response.template def get( self, name: str, *, workspace: str, ) -> openshell_pb2.SandboxWorkloadTemplate: response = self._stub.GetSandboxTemplate( openshell_pb2.GetSandboxTemplateRequest( workspace_scope=_workspace_scope(workspace), name=name ), timeout=self._timeout, ) return response.template def list( self, *, workspace: str, page_size: int = 100, page_token: str = "", label_selector: str = "", ) -> Pager[openshell_pb2.SandboxWorkloadTemplate]: def fetch(token: str) -> Page[openshell_pb2.SandboxWorkloadTemplate]: response = self._stub.ListSandboxTemplates( openshell_pb2.ListSandboxTemplatesRequest( workspace_scope=_workspace_scope(workspace), page_size=page_size, page_token=token, label_selector=label_selector, ), timeout=self._timeout, ) return Page( items=list(response.templates), next_page_token=getattr(response, "next_page_token", ""), ) return Pager(fetch, page_token) def list_all( self, *, workspace: str, page_size: int = 100, page_token: str = "", label_selector: str = "", ) -> builtins.list[openshell_pb2.SandboxWorkloadTemplate]: return self.list( workspace=workspace, page_size=page_size, page_token=page_token, label_selector=label_selector, ).all() def list_for_all_workspaces( self, *, page_size: int = 100, page_token: str = "", label_selector: str = "", ) -> Pager[openshell_pb2.SandboxWorkloadTemplate]: def fetch(token: str) -> Page[openshell_pb2.SandboxWorkloadTemplate]: response = self._stub.ListSandboxTemplates( openshell_pb2.ListSandboxTemplatesRequest( workspace_scope=_all_workspaces_scope(), page_size=page_size, page_token=token, label_selector=label_selector, ), timeout=self._timeout, ) return Page( items=list(response.templates), next_page_token=getattr(response, "next_page_token", ""), ) return Pager(fetch, page_token) def list_all_for_all_workspaces( self, *, page_size: int = 100, page_token: str = "", label_selector: str = "", ) -> builtins.list[openshell_pb2.SandboxWorkloadTemplate]: return self.list_for_all_workspaces( page_size=page_size, page_token=page_token, label_selector=label_selector, ).all() def delete( self, name: str, *, workspace: str, allow_missing: bool = False ) -> DeletionResult: response = self._stub.DeleteSandboxTemplate( openshell_pb2.DeleteSandboxTemplateRequest( workspace_scope=_workspace_scope(workspace), name=name, allow_missing=allow_missing, ), timeout=self._timeout, ) return DeletionResult(DeletionOutcome(response.outcome)) @dataclass(frozen=True) class WorkspaceRef: name: str phase: str labels: dict[str, str] def _workspace_ref(ws: datamodel_pb2.Workspace) -> WorkspaceRef: meta = ws.metadata return WorkspaceRef( name=meta.name, phase=datamodel_pb2.WorkspacePhase.Name(ws.status.phase), labels=dict(meta.labels), ) class WorkspaceClient: """gRPC client for workspace lifecycle operations.""" def __init__(self, channel: grpc.Channel, *, timeout: float = 30.0) -> None: self._stub = openshell_pb2_grpc.OpenShellStub(_error_mapping_channel(channel)) self._timeout = timeout @classmethod def from_sandbox_client(cls, client: SandboxClient) -> WorkspaceClient: return cls(client._channel, timeout=client._timeout) def create( self, name: str, *, labels: Mapping[str, str] | None = None, ) -> WorkspaceRef: response = self._stub.CreateWorkspace( openshell_pb2.CreateWorkspaceRequest( name=name, labels=dict(labels) if labels else {}, ), timeout=self._timeout, ) return _workspace_ref(response.workspace) def get(self, name: str) -> WorkspaceRef: response = self._stub.GetWorkspace( openshell_pb2.GetWorkspaceRequest(name=name), timeout=self._timeout, ) return _workspace_ref(response.workspace) def list( self, *, page_size: int = 100, page_token: str = "", label_selector: str | None = None, ) -> Pager[WorkspaceRef]: def fetch(token: str) -> Page[WorkspaceRef]: response = self._stub.ListWorkspaces( openshell_pb2.ListWorkspacesRequest( page_size=page_size, page_token=token, label_selector=label_selector or "", ), timeout=self._timeout, ) return Page( items=[_workspace_ref(ws) for ws in response.workspaces], next_page_token=getattr(response, "next_page_token", ""), ) return Pager(fetch, page_token) def list_all( self, *, page_size: int = 100, page_token: str = "", label_selector: str | None = None, ) -> builtins.list[WorkspaceRef]: return self.list( page_size=page_size, page_token=page_token, label_selector=label_selector, ).all() def delete(self, name: str, *, allow_missing: bool = False) -> DeletionResult: response = self._stub.DeleteWorkspace( openshell_pb2.DeleteWorkspaceRequest( name=name, allow_missing=allow_missing ), timeout=self._timeout, ) return DeletionResult(DeletionOutcome(response.outcome)) class Sandbox: """Context-managed sandbox session bound to one sandbox id.""" def __init__( self, *, workspace: str, cluster: str | None = None, sandbox: str | SandboxRef | None = None, delete_on_exit: bool = True, spec: openshell_pb2.SandboxSpec | None = None, name: str | None = None, labels: Mapping[str, str] | None = None, workload_template: str | None = None, timeout: float = 30.0, ready_timeout_seconds: float = 120.0, auto_refresh: bool = True, write_back: bool = True, insecure: bool = False, client_credentials: ClientCredentialsAuth | None = None, ) -> None: """Bind a Sandbox context to the active gateway. OIDC kwargs (`auto_refresh`, `write_back`, `insecure`, and `client_credentials`) forward directly to `SandboxClient.from_active_cluster` and have the same semantics. They're surfaced on `Sandbox` so callers using the higher-level wrapper get parity with `SandboxClient` for OIDC-protected gateways (e.g. passing `insecure=True` for a self-signed dev IdP). Non-OIDC gateways ignore them. """ self._workspace = workspace self._cluster = cluster self._sandbox_input = sandbox self._delete_on_exit = delete_on_exit self._spec = spec self._name = name # Copy so later caller mutation cannot change what gets sent on enter. self._labels = dict(labels) if labels is not None else None self._workload_template = workload_template self._timeout = timeout self._ready_timeout_seconds = ready_timeout_seconds self._auto_refresh = auto_refresh self._write_back = write_back self._insecure = insecure self._client_credentials = client_credentials self._client: SandboxClient | None = None self._session: SandboxSession | None = None @property def id(self) -> str: if self._session is None: raise SandboxError("sandbox context has not been entered") return self._session.id @property def sandbox(self) -> SandboxRef: if self._session is None: raise SandboxError("sandbox context has not been entered") return self._session.sandbox def __enter__(self) -> Sandbox: # Creation metadata cannot be applied when attaching to an existing # sandbox; reject it before opening a connection. if self._sandbox_input is not None and ( self._name is not None or self._labels is not None or self._workload_template is not None ): raise SandboxError( "name, labels, and workload_template cannot be set when attaching to an existing sandbox" ) client = SandboxClient.from_active_cluster( cluster=self._cluster, timeout=self._timeout, auto_refresh=self._auto_refresh, write_back=self._write_back, insecure=self._insecure, client_credentials=self._client_credentials, ) self._client = client if self._sandbox_input is None and self._workload_template is not None: self._session = client.create_session_from_template( workspace=self._workspace, workload_template=self._workload_template, spec=self._spec, name=self._name, labels=self._labels, ) elif self._sandbox_input is None: self._session = client.create_session( workspace=self._workspace, spec=self._spec, name=self._name, labels=self._labels, ) elif isinstance(self._sandbox_input, SandboxRef): self._session = SandboxSession(client, self._sandbox_input) else: self._session = client.get_session( self._sandbox_input, workspace=self._workspace ) self._workspace = getattr(self._session, "_workspace", self._workspace) ready = client.wait_ready( self._session.sandbox.name, workspace=self._workspace, timeout_seconds=self._ready_timeout_seconds, ) self._session = SandboxSession(client, ready) return self def __exit__(self, *args: object) -> None: try: if ( self._delete_on_exit and self._session is not None and self._client is not None ): result = self._session.delete(allow_missing=True) if result.outcome == DeletionOutcome.ACCEPTED: self._client.wait_deleted( self._session.sandbox.name, workspace=self._workspace, expected_sandbox_id=result.sandbox_id, ) elif result.outcome not in ( DeletionOutcome.COMPLETED, DeletionOutcome.ALREADY_ABSENT, ): raise SandboxError( f"unsupported deletion outcome: {result.outcome}" ) finally: if self._client is not None: self._client.close() self._session = None self._client = None def exec( self, command: Sequence[str], *, stream_output: bool = False, workdir: str | None = None, env: Mapping[str, str] | None = None, stdin: bytes | None = None, timeout_seconds: int | None = None, no_login_shell: bool = False, ) -> ExecResult: if self._session is None: raise SandboxError("sandbox context has not been entered") return self._session.exec( command, stream_output=stream_output, workdir=workdir, env=env, stdin=stdin, timeout_seconds=timeout_seconds, no_login_shell=no_login_shell, ) def exec_python( self, function: Callable[..., object], *, args: Sequence[object] = (), kwargs: Mapping[str, object] | None = None, stream_output: bool = False, workdir: str | None = None, env: Mapping[str, str] | None = None, timeout_seconds: int | None = None, ) -> ExecResult: if self._session is None: raise SandboxError("sandbox context has not been entered") return self._session.exec_python( function, args=args, kwargs=kwargs, stream_output=stream_output, workdir=workdir, env=env, timeout_seconds=timeout_seconds, ) _PYTHON_CLOUDPICKLE_BOOTSTRAP = ( "import base64,cloudpickle,os;" "payload=base64.b64decode(os.environ['OPENSHELL_PYFUNC_B64']);" "func,args,kwargs=cloudpickle.loads(payload);" "result=func(*args,**kwargs);" "print(result) if result is not None else None" ) _SANDBOX_PYTHON_BIN = "python" def _serialize_python_callable( function: Callable[..., object], *, args: Sequence[object], kwargs: Mapping[str, object] | None, ) -> str: try: import cloudpickle except ImportError as exc: # pragma: no cover - import error path raise SandboxError("cloudpickle is required for exec_python") from exc payload = cloudpickle.dumps((function, tuple(args), dict(kwargs or {}))) return base64.b64encode(payload).decode("ascii") def _sandbox_ref( sandbox: openshell_pb2.Sandbox, service_urls: Mapping[str, str] | None = None, ) -> SandboxRef: status = sandbox.status if sandbox.HasField("status") else None provenance = ( SandboxWorkloadTemplateProvenanceRef( name=sandbox.created_from_workload_template.name, resource_version=sandbox.created_from_workload_template.resource_version, ) if sandbox.HasField("created_from_workload_template") else None ) return SandboxRef( id=sandbox.metadata.id if sandbox.metadata else "", name=sandbox.metadata.name if sandbox.metadata else "", workspace=sandbox.metadata.workspace if sandbox.metadata else "", status=SandboxStatusRef( phase=status.phase if status else 0, current_policy_version=status.current_policy_version if status else 0, exit_code=status.exit_code if status is not None and status.HasField("exit_code") else None, restart_count=status.restart_count if status is not None else 0, next_restart_at_ms=status.next_restart_time.ToMilliseconds() if status is not None and status.HasField("next_restart_time") else None, main_process_started_at_ms=status.main_process_started_time.ToMilliseconds() if status is not None and status.HasField("main_process_started_time") else None, ), labels=sandbox.metadata.labels if sandbox.metadata else {}, created_from_workload_template=provenance, service_urls=service_urls or {}, ) def _default_spec() -> openshell_pb2.SandboxSpec: # Omit the policy field so the sandbox container discovers its policy # from /etc/openshell/policy.yaml (baked into the image at build time). # This avoids duplicating policy defaults between the SDK and the # container image and ensures sandboxes get the full dev-sandbox-policy # (including network_policies) out of the box. return openshell_pb2.SandboxSpec() def _sandbox_workload_template( *, name: str, image: str | None = None, labels: Mapping[str, str] | None = None, annotations: Mapping[str, str] | None = None, environment: Mapping[str, str] | None = None, cpu: str | None = None, memory: str | None = None, gpu_count: int | None = None, gpu: bool = False, driver_config: Mapping[str, Any] | None = None, ) -> openshell_pb2.SandboxWorkloadTemplate: if gpu_count is not None and gpu_count <= 0: raise SandboxError("gpu_count must be greater than zero") template = openshell_pb2.SandboxWorkloadTemplate() template.metadata.name = name # The gateway permits an empty workload so it can apply the default image # at sandbox create time, but the workload message itself is required. template.spec.workload.SetInParent() if labels: template.metadata.labels.update(dict(labels)) if annotations: template.metadata.annotations.update(dict(annotations)) if image is not None: template.spec.workload.image = image if environment: template.spec.workload.environment.update(dict(environment)) if cpu is not None: template.spec.workload.resources.cpu = cpu if memory is not None: template.spec.workload.resources.memory = memory if gpu or gpu_count is not None: template.spec.workload.resources.gpu.SetInParent() if gpu_count is not None: template.spec.workload.resources.gpu.count = gpu_count if driver_config: template.spec.driver_config.update(dict(driver_config)) return template def _xdg_config_home() -> pathlib.Path: configured = os.environ.get("XDG_CONFIG_HOME") if configured: return pathlib.Path(configured) return pathlib.Path.home() / ".config" # Re-check the cached token roughly 30 seconds before the issuer's # stated expiry, to leave room for in-flight RPCs and clock skew. This # matches `openshell-bootstrap::oidc_token::is_token_expired`. _OIDC_TOKEN_EXPIRY_GRACE_SECONDS = 30 _IS_WINDOWS = os.name == "nt" _WINDOWS_REPLACE_RETRYABLE_ERRORS = frozenset({5, 32}) _WINDOWS_REPLACE_TIMEOUT_SECONDS = 0.25 _WINDOWS_REPLACE_INITIAL_DELAY_SECONDS = 0.005 _WINDOWS_REPLACE_MAX_DELAY_SECONDS = 0.05 _WINDOWS_REPLACE_LOCK = threading.Lock() def _atomic_replace(source: pathlib.Path, destination: pathlib.Path) -> None: """Atomically replace a file, retrying transient Windows sharing errors.""" if not _IS_WINDOWS: source.replace(destination) return # Serialize writers in this process. The retry still handles other # processes (including the Rust CLI) and filesystem scanners that briefly # open the destination without delete sharing. with _WINDOWS_REPLACE_LOCK: deadline = time.monotonic() + _WINDOWS_REPLACE_TIMEOUT_SECONDS delay = _WINDOWS_REPLACE_INITIAL_DELAY_SECONDS while True: try: source.replace(destination) return except PermissionError as error: winerror = getattr(error, "winerror", None) retryable = winerror in _WINDOWS_REPLACE_RETRYABLE_ERRORS or ( winerror is None and error.errno == errno.EACCES ) if not retryable or time.monotonic() >= deadline: raise time.sleep(delay) delay = min(delay * 2, _WINDOWS_REPLACE_MAX_DELAY_SECONDS) def _read_oidc_token_bundle(gateway_dir: pathlib.Path) -> dict | None: """Read and parse `oidc_token.json` for a gateway. Returns the parsed dict, or `None` if the file is absent or unreadable. See `openshell-bootstrap::oidc_token::store_oidc_token` for the writer. """ token_path = gateway_dir / "oidc_token.json" try: return json.loads(token_path.read_text(encoding="utf-8")) except FileNotFoundError: return None except (OSError, UnicodeDecodeError, json.JSONDecodeError): return None def _normalize_issuer(bundle: dict) -> str | None: """Return the bundle's issuer with a trailing slash stripped. Used to detect whether the issuer changed when adopting a bundle re-read from disk, so a cached token endpoint computed for the old issuer can be invalidated. Trailing-slash differences are treated as equal, matching `_discover_token_endpoint`'s normalization. """ issuer = bundle.get("issuer") return issuer.rstrip("/") if isinstance(issuer, str) else None def _load_cluster_bearer_token(gateway_dir: pathlib.Path) -> str | None: """Read a single (possibly expired) access token from disk. Lower-level helper used by both the legacy single-shot path and the refreshing provider. Returns the raw access_token string or None. """ bundle = _read_oidc_token_bundle(gateway_dir) if bundle is None: return None access_token = bundle.get("access_token") if not isinstance(access_token, str) or not access_token: return None return access_token def _make_fail_closed_bearer_provider( gateway_dir: pathlib.Path, cluster_name: str, ) -> Callable[[], str]: """Per-RPC provider that re-reads `oidc_token.json` but does NOT refresh. Raises `SandboxError` when the token is missing or expired. Available as an opt-out from the default `_OidcRefresher` for callers (e.g. tests) that want to assert expiry behavior or that don't want the SDK to make outbound HTTP calls to the IdP. """ def provider() -> str: bundle = _read_oidc_token_bundle(gateway_dir) if bundle is None: raise SandboxError( f"OIDC token for gateway '{cluster_name}' is missing or " f"unreadable. Re-authenticate with: openshell gateway login" ) access_token = bundle.get("access_token") if not isinstance(access_token, str) or not access_token: raise SandboxError( f"OIDC token for gateway '{cluster_name}' has no access " f"token. Re-authenticate with: openshell gateway login" ) expires_at = bundle.get("expires_at") if isinstance(expires_at, int): now = int(time.time()) if now + _OIDC_TOKEN_EXPIRY_GRACE_SECONDS >= expires_at: raise SandboxError( f"OIDC token for gateway '{cluster_name}' has expired. " f"Re-authenticate with: openshell gateway login" ) return access_token return provider class _InvalidGrantError(SandboxError): """Refresh failed with OAuth2 `invalid_grant` (RFC 6749 §5.2). The refresh_token was rejected — expired, revoked, or rotated out by a concurrent refresh in another process. Subclass of `SandboxError` so an uncaught instance still surfaces the standard re-authenticate hint; `current_access_token` catches it to retry once with a peer-rotated bundle before giving up. """ class _OidcRefresher: """Thread-safe in-process OAuth2 refresh for a gateway's `oidc_token.json`. Mirrors the lazy-refresh pattern used by `google-auth`'s `Credentials` and `botocore`'s `SSOTokenProvider`. Uses `httpx` for transport so we can pin `follow_redirects=False` and the TLS verification policy explicitly — same posture as the Rust CLI's use of `reqwest` with `Policy::none()` and an opt-in `danger_accept_invalid_certs` (see `crates/openshell-cli/src/oidc_auth.rs::http_client`). The OAuth2 refresh-token grant itself (RFC 6749 §6) is a single form-encoded POST, handled inline rather than via an OAuth2 library. Properties: - **Lazy**: check expiry on every RPC; refresh only when stale. - **Lock-coordinated**: concurrent RPCs share a single refresh, not one per call. Plain `threading.Lock` (no separate worker thread, unlike `google-auth`'s `RefreshThreadManager` — sufficient for our use case). - **Disk-aware**: before refreshing, re-read `oidc_token.json` — the CLI or another process may have already rotated the bundle. - **Discovery-validated**: fetches the OIDC discovery document from `/.well-known/openid-configuration`, rejects responses whose `issuer` field doesn't match the configured one (preventing SSRF / misdirection of the refresh_token to an attacker-controlled endpoint). Mirrors the Rust CLI's `discover()` function. - **Redirect-hardened**: `follow_redirects=False` on the underlying `httpx.Client` so a 3xx during discovery or refresh is treated as a failure rather than silently chasing the redirect to an arbitrary host. - **Write-back by default**: when `auto_refresh=True`, refreshed bundles are atomically persisted to `oidc_token.json` at mode 0600 so other processes (Rust CLI, TUI, other Python clients) see the rotated `refresh_token`. Required for IdPs that invalidate the old `refresh_token` on rotation (Keycloak with rotation enabled, Entra in strict mode); without write-back, a second process would `invalid_grant` on next refresh. - **`insecure=True` flag** disables TLS certificate verification for the discovery and refresh calls. Matches the Rust CLI's `--insecure` plumbing for OIDC issuers behind self-signed certs. - **Refresh failures** surface as `SandboxError` with a "re-authenticate with: openshell gateway login" hint. """ def __init__( self, gateway_dir: pathlib.Path, cluster_name: str, *, write_back: bool = True, insecure: bool = False, ) -> None: self._gateway_dir = gateway_dir self._cluster_name = cluster_name self._write_back = write_back self._lock = threading.Lock() self._bundle: dict | None = None self._token_endpoint: str | None = None # Single httpx.Client serves both discovery (unauthenticated # GET) and refresh (POST with form-encoded body) — they share # the same security posture. # # - follow_redirects=False: a 3xx during discovery would # otherwise steer us to an attacker-controlled token # endpoint. Matches `reqwest::redirect::Policy::none()` in # the Rust CLI's `oidc_auth.rs::http_client`. # - verify=not insecure: opt-in TLS-verification disable for # self-signed issuers. Matches the Rust CLI's `--insecure` # flag plumbing. # # We don't use authlib's `OAuth2Client` here because it # auto-injects an Authorization header on every request from # its stored token, which would break the unauthenticated # discovery GET. The refresh_token grant for a public client # is a single form-encoded POST — small enough to spell out # directly with httpx, and easier to test. self._http = httpx.Client( follow_redirects=False, verify=not insecure, timeout=15.0, ) def close(self) -> None: self._http.close() def __del__(self) -> None: # Best-effort cleanup; tests + short-lived callers may not call # close() explicitly. with contextlib.suppress(Exception): self._http.close() def current_access_token(self) -> str: """Return a non-expired access token, refreshing if needed.""" with self._lock: if self._bundle is None: self._bundle = _read_oidc_token_bundle(self._gateway_dir) if self._bundle is None: raise SandboxError( f"OIDC token for gateway '{self._cluster_name}' is " f"missing or unreadable. Re-authenticate with: " f"openshell gateway login" ) if self._is_fresh(self._bundle): return self._bundle["access_token"] # Cached bundle is stale. Before refreshing, re-read disk — # another process (CLI, TUI, another SDK client) may have # rotated the bundle while we were idle. Adopt the disk # bundle when it was refreshed more recently than ours, EVEN # WHEN its access token is also stale: otherwise we'd refresh # with our in-memory refresh_token, which a rotating IdP may # have already invalidated when the other process refreshed # (Keycloak with rotation, Entra in strict mode). # # "More recently" is judged by `expires_at`: a refresh issues # a new access token with a forward expiry alongside the # (possibly rotated) refresh_token, so the bundle with the # later expiry carries the newest refresh_token. This also # preserves the write_back=False case, where our in-memory # bundle has already rotated past the on-disk one and must # NOT be clobbered by the older disk copy. disk = _read_oidc_token_bundle(self._gateway_dir) if disk is not None and self._expiry(disk) > self._expiry(self._bundle): # If the issuer changed under us, the cached token # endpoint no longer applies — force re-discovery. if _normalize_issuer(disk) != _normalize_issuer(self._bundle): self._token_endpoint = None self._bundle = disk if self._is_fresh(disk): return disk["access_token"] # Truly stale; refresh against the IdP using the freshest # bundle we have (disk if it was newer, else in-memory). try: self._bundle = self._refresh(self._bundle) except _InvalidGrantError as exc: # We lost a cross-process rotation race: between our disk # re-read above and our refresh POST, a peer (CLI, TUI, # another SDK client) rotated the refresh_token and the # IdP invalidated ours. This is the residual window that # neither google-auth nor botocore close without an OS # file lock. Rather than lock, recover: re-read disk once # and, if a peer wrote a *different* refresh_token, retry # with it before surfacing a re-authenticate error. self._bundle = self._recover_from_invalid_grant(self._bundle, exc) if self._write_back: self._write_to_disk(self._bundle) return self._bundle["access_token"] def _recover_from_invalid_grant( self, attempted: dict, exc: _InvalidGrantError ) -> dict: """Re-read disk after an `invalid_grant` and retry once if a peer rotated the refresh_token. `attempted` is the bundle whose refresh_token the IdP just rejected. If disk now holds a different refresh_token, a concurrent process won the rotation race — adopt and reuse it (returning early if it is already fresh, otherwise refreshing with it). If disk offers nothing new, the rejection is genuine: re-raise so the caller sees the re-authenticate hint. """ disk = _read_oidc_token_bundle(self._gateway_dir) if disk is None or disk.get("refresh_token") == attempted.get("refresh_token"): # No peer rotation — the refresh_token really is dead. raise exc if _normalize_issuer(disk) != _normalize_issuer(attempted): self._token_endpoint = None if self._is_fresh(disk): return disk # Single retry with the peer's rotated token; a second # _InvalidGrantError here propagates (no further retry). return self._refresh(disk) @staticmethod def _is_fresh(bundle: dict) -> bool: access_token = bundle.get("access_token") if not isinstance(access_token, str) or not access_token: return False exp = bundle.get("expires_at") if not isinstance(exp, int): # No expiry info — treat as fresh (matches the # `is_token_expired` semantics in the Rust CLI). return True return int(time.time()) + _OIDC_TOKEN_EXPIRY_GRACE_SECONDS < exp @staticmethod def _expiry(bundle: dict) -> float: """Access-token expiry as a comparable number. A bundle without an `expires_at` is treated as non-expiring (`+inf`) — consistent with `_is_fresh`, which treats a missing expiry as always fresh. Used to decide which of two stale bundles (in-memory vs. on-disk) was refreshed more recently and therefore holds the newest refresh_token. """ exp = bundle.get("expires_at") return float(exp) if isinstance(exp, int) else float("inf") def _discover_token_endpoint(self, bundle: dict) -> str: if self._token_endpoint is not None: return self._token_endpoint issuer = bundle.get("issuer") if not isinstance(issuer, str) or not issuer: raise SandboxError( f"OIDC bundle for gateway '{self._cluster_name}' has no " f"`issuer`; cannot refresh. Re-authenticate with: openshell " f"gateway login" ) normalized_issuer = issuer.rstrip("/") discovery_url = f"{normalized_issuer}/.well-known/openid-configuration" try: resp = self._http.get(discovery_url) except httpx.HTTPError as e: raise SandboxError( f"OIDC discovery failed for gateway " f"'{self._cluster_name}': {e}. Re-authenticate with: " f"openshell gateway login" ) from e # follow_redirects=False means a 3xx surfaces as a non-2xx # status; treat any non-success as a discovery failure rather # than silently following. if not 200 <= resp.status_code < 300: raise SandboxError( f"OIDC discovery failed for gateway " f"'{self._cluster_name}': HTTP {resp.status_code} " f"from {discovery_url}. Re-authenticate with: openshell " f"gateway login" ) try: disco = resp.json() except ValueError as e: raise SandboxError( f"OIDC discovery returned invalid JSON for gateway " f"'{self._cluster_name}': {e}" ) from e # Critical: validate that the discovery document's `issuer` # matches the configured one. Without this, a misdirected or # malicious discovery response could steer the refresh_token # POST to an attacker-controlled endpoint. discovered_issuer = disco.get("issuer", "") if not isinstance(discovered_issuer, str) or ( discovered_issuer.rstrip("/") != normalized_issuer ): raise SandboxError( f"OIDC discovery issuer mismatch for gateway " f"'{self._cluster_name}': expected '{normalized_issuer}', " f"got '{discovered_issuer}'." ) endpoint = disco.get("token_endpoint") if not isinstance(endpoint, str) or not endpoint: raise SandboxError( f"OIDC discovery for gateway '{self._cluster_name}' did " f"not include a token_endpoint." ) self._token_endpoint = endpoint return endpoint def _refresh(self, bundle: dict) -> dict: refresh_token = bundle.get("refresh_token") if not isinstance(refresh_token, str) or not refresh_token: raise SandboxError( f"OIDC token for gateway '{self._cluster_name}' has no " f"refresh token. Re-authenticate with: openshell gateway " f"login" ) token_endpoint = self._discover_token_endpoint(bundle) client_id = bundle.get("client_id", "openshell-cli") # RFC 6749 §6: refresh_token grant. Form-encoded POST with # grant_type, refresh_token, and (for a public client) client_id. # No Authorization header (token_endpoint_auth_method="none"). try: resp = self._http.post( token_endpoint, data={ "grant_type": "refresh_token", "refresh_token": refresh_token, "client_id": client_id, }, ) except httpx.HTTPError as e: raise SandboxError( f"OIDC token refresh failed for gateway " f"'{self._cluster_name}': {type(e).__name__}: {e}. " f"Re-authenticate with: openshell gateway login" ) from e if resp.status_code != 200: # Include the IdP's error body for diagnostics — RFC 6749 # mandates a JSON body like {"error":"invalid_grant", ...} # on failure, which is the most useful signal to surface. error_code = None with contextlib.suppress(Exception): body = resp.json() if isinstance(body, dict): error_code = body.get("error") detail = "" with contextlib.suppress(Exception): detail = f": {resp.text[:200]}" message = ( f"OIDC token refresh failed for gateway " f"'{self._cluster_name}': HTTP {resp.status_code}" f"{detail}. Re-authenticate with: openshell gateway " f"login" ) # `invalid_grant` specifically means the refresh_token was # rejected — distinguished from transport/5xx errors so the # caller can retry once with a peer-rotated bundle (a lost # cross-process rotation race) before surfacing the failure. if error_code == "invalid_grant": raise _InvalidGrantError(message) raise SandboxError(message) try: token = resp.json() except ValueError as e: raise SandboxError( f"OIDC refresh response for gateway '{self._cluster_name}' " f"is not JSON: {e}" ) from e access_token = token.get("access_token") if not isinstance(access_token, str) or not access_token: raise SandboxError( f"OIDC refresh response for gateway '{self._cluster_name}' " f"is missing access_token." ) expires_at = token.get("expires_at") if expires_at is None: expires_in = token.get("expires_in") if isinstance(expires_in, (int, float)): expires_at = int(time.time()) + int(expires_in) return { "access_token": access_token, # Refresh-token rotation: some IdPs (Keycloak with rotation # enabled, Entra in strict mode) reissue and invalidate the # old one. Honor the new value when present. "refresh_token": token.get("refresh_token", refresh_token), "expires_at": int(expires_at) if expires_at is not None else None, "issuer": bundle.get("issuer", ""), "client_id": client_id, } def _write_to_disk(self, bundle: dict) -> None: """Atomic-replace `oidc_token.json` with the refreshed bundle. Strips `None` values to match the Rust writer's `skip_serializing_if = "Option::is_none"` behavior so a Python- written file is byte-identical in shape to what the CLI writes. Uses `tempfile.mkstemp` (PID + random suffix) so two writers racing on the same gateway directory don't share a tmp file and trample each other's content. Each writer gets a unique path; `.replace()` is atomic per-writer, and POSIX rename semantics ensure the final `oidc_token.json` is always complete-and-readable to anyone observing. """ path = self._gateway_dir / "oidc_token.json" serializable = {k: v for k, v in bundle.items() if v is not None} payload = json.dumps(serializable, indent=2) # mkstemp creates the file with mode 0600 already on POSIX # (it uses O_CREAT | O_EXCL with restrictive umask), so chmod # is a belt-and-braces step for filesystems that don't honor # the initial mode. fd, tmp_name = tempfile.mkstemp( prefix=".oidc_token.", suffix=".tmp", dir=str(self._gateway_dir), ) tmp_path = pathlib.Path(tmp_name) try: with os.fdopen(fd, "w", encoding="utf-8") as f: f.write(payload) with contextlib.suppress(OSError): tmp_path.chmod(0o600) _atomic_replace(tmp_path, path) except BaseException: # Clean up our tmp on failure so we don't leave orphaned # `.oidc_token..tmp` files lying around. The replace # already moved the file on the success path. with contextlib.suppress(OSError): tmp_path.unlink() raise def _make_cluster_bearer_provider( gateway_dir: pathlib.Path, cluster_name: str, *, auto_refresh: bool = True, write_back: bool = True, insecure: bool = False, ) -> tuple[Callable[[], str], Callable[[], None] | None]: """Build a per-RPC token provider for a gateway directory. Returns `(token_provider, close_fn_or_none)`. `close_fn` is non-None only when an `_OidcRefresher` was constructed; callers that own the provider's lifecycle (e.g. `SandboxClient.close()`) should invoke it during teardown so the underlying httpx.Client is released rather than relying on `__del__`. With `auto_refresh=True` (the default), returns an `_OidcRefresher`- backed callable that lazily refreshes against the IdP's token endpoint when the cached bundle is stale. This mirrors the lazy-refresh pattern used by `google.oauth2.credentials.Credentials` and `botocore.tokens.SSOTokenProvider` and lets long-running scripts survive token rotation without intervention. With `auto_refresh=False`, falls back to the read-only / fail-closed behavior: the SDK consumes whatever the CLI most recently wrote and raises `SandboxError` when the token expires. Useful for tests or callers that don't want the SDK to make outbound HTTP calls to the IdP. No close_fn is returned in this case. `write_back=True` (only meaningful when `auto_refresh=True`) makes the refresher atomically persist the rotated bundle back to `oidc_token.json` so other processes — including the Rust CLI — see the new token. Defaults to True because OIDC providers with refresh-token rotation (Keycloak, Entra) invalidate the old refresh_token on rotation; an in-memory-only refresh would leave the on-disk bundle pointing at an invalidated value, and any other process starting from that disk state would fail on its first refresh. `insecure=True` disables TLS certificate verification for both the OIDC discovery document fetch and the refresh-token POST. Mirrors the Rust CLI's `--insecure` flag for OIDC issuers behind self-signed certs. """ if not auto_refresh: return _make_fail_closed_bearer_provider(gateway_dir, cluster_name), None refresher = _OidcRefresher( gateway_dir, cluster_name, write_back=write_back, insecure=insecure, ) return refresher.current_access_token, refresher.close def _resolve_active_cluster() -> str: env_gateway = os.environ.get("OPENSHELL_GATEWAY") if env_gateway: return env_gateway active_file = _xdg_config_home() / "openshell" / "active_gateway" try: value = active_file.read_text(encoding="utf-8").strip() except FileNotFoundError: raise SandboxError("no active gateway configured") from None if value == "": raise SandboxError("no active gateway configured") return value