onehouse-python-sdk 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

onehouse_python_sdk/__init__.py

@@ -6,4 +6,4 @@ from .lakebase.client import LakebaseClient
 from .resources.client import OnehouseResources
 
 __all__ = ["SqlClient", "LakebaseClient", "OnehouseResources", "exceptions"]
-__version__ = "0.1.0"
+__version__ = "0.4.0"

onehouse_python_sdk/_oauth.py

@@ -0,0 +1,293 @@
+"""Loopback PKCE sign-in and token-endpoint calls, shared by every OAuth login.
+
+Standard library only, and outside every subsystem, so both the control-plane
+login (``resources.oauth``, external-api's OAuth server) and the data-plane
+AI-Gateway login (``ai_gateway``, Keycloak) use it without importing each other.
+
+The browser only ever carries an authorization code, which is worthless
+without the PKCE verifier that never leaves this process (RFC 7636, RFC 8252).
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import http.server
+import json
+import secrets
+import ssl
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+import webbrowser
+from dataclasses import dataclass
+from typing import Callable, Mapping, Optional
+
+from .exceptions import OnehouseSdkError
+
+LOGIN_TIMEOUT_SECONDS = 300
+_HTTP_TIMEOUT_SECONDS = 20
+
+_DONE_PAGE = b"""<!doctype html><meta charset="utf-8"><title>Onehouse</title>
+<body style="font:16px system-ui;padding:3rem;max-width:32rem">
+<h2>Signed in</h2><p>You can close this tab and return to your terminal.</p></body>"""
+
+_FAILED_PAGE = b"""<!doctype html><meta charset="utf-8"><title>Onehouse</title>
+<body style="font:16px system-ui;padding:3rem;max-width:32rem">
+<h2>Sign-in did not complete</h2><p>Return to your terminal for the reason.</p></body>"""
+
+
+class OAuthFlowError(OnehouseSdkError):
+    """A sign-in or token-endpoint call did not produce a usable token."""
+
+
+@dataclass(frozen=True)
+class AuthorizationResult:
+    """What the browser step returns: redeem ``code`` with ``verifier`` at the same
+    ``redirect_uri``, byte for byte."""
+
+    code: str
+    redirect_uri: str
+    verifier: str
+    iss: Optional[str] = None
+
+
+def _b64url(raw: bytes) -> str:
+    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
+
+
+def new_verifier() -> str:
+    """43 characters of base64url (RFC 7636 4.1)."""
+    return _b64url(secrets.token_bytes(32))
+
+
+def challenge_for(verifier: str) -> str:
+    """S256 only: ``plain`` would put the verifier in the browser's history."""
+    return _b64url(hashlib.sha256(verifier.encode("ascii")).digest())
+
+
+def default_ssl_context() -> ssl.SSLContext:
+    """The OS trust store when ``truststore`` is installed, so a TLS-intercepting
+    corporate proxy's CA is trusted; then certifi; then Python's defaults.
+    The same order as ``lakebase.auth._ssl``, which this module may not import."""
+    try:
+        import truststore
+
+        return truststore.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+    except ImportError:
+        pass
+    try:
+        import certifi
+
+        return ssl.create_default_context(cafile=certifi.where())
+    except ImportError:
+        return ssl.create_default_context()
+
+
+class _Catcher(http.server.BaseHTTPRequestHandler):
+    """Serves one useful request: the authorization server's redirect back."""
+
+    callback_path = ""
+    expected_state = ""
+    result: dict = {}
+    # A connection that sends nothing (a browser preconnect, a port scan) would
+    # otherwise block handle_request() past the login deadline.
+    timeout = 5
+
+    def do_GET(self) -> None:  # noqa: N802 - name fixed by BaseHTTPRequestHandler
+        parsed = urllib.parse.urlparse(self.path)
+        if parsed.path != type(self).callback_path:
+            self.send_response(404)
+            self.end_headers()
+            return
+        params = urllib.parse.parse_qs(parsed.query)
+        # State before anything else: a callback this process did not start is not
+        # ours to act on, and must not end the wait.
+        got_state = (params.get("state") or [""])[0]
+        if not secrets.compare_digest(got_state, type(self).expected_state):
+            self._respond(400, _FAILED_PAGE)
+            return
+        error = (params.get("error") or [""])[0]
+        if error:
+            description = (params.get("error_description") or [""])[0]
+            type(self).result = {"error": f"{error}: {description}" if description else error}
+            self._respond(400, _FAILED_PAGE)
+            return
+        code = (params.get("code") or [""])[0]
+        if not code:
+            type(self).result = {"error": "the authorization server returned no code"}
+            self._respond(400, _FAILED_PAGE)
+            return
+        type(self).result = {"code": code, "iss": (params.get("iss") or [None])[0]}
+        self._respond(200, _DONE_PAGE)
+
+    def _respond(self, status: int, body: bytes) -> None:
+        self.send_response(status)
+        self.send_header("Content-Type", "text/html; charset=utf-8")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def log_message(self, *_args) -> None:
+        """The default access log would print the callback URL, code included."""
+
+
+class LoopbackServer:
+    """A one-shot HTTP server on 127.0.0.1, on an OS-chosen port unless one is given.
+
+    127.0.0.1, never ``localhost``: redirect matching is byte for byte, and
+    ``localhost`` resolves through DNS (RFC 8252 7.3, 8.3).
+    """
+
+    def __init__(self, callback_path: str, state: str, port: int = 0) -> None:
+        self._callback_path = callback_path
+        self._handler = type(
+            "_Bound",
+            (_Catcher,),
+            {"callback_path": callback_path, "expected_state": state, "result": {}},
+        )
+        self._httpd = http.server.HTTPServer(("127.0.0.1", port), self._handler)
+        self._httpd.timeout = 0.5
+
+    @property
+    def port(self) -> int:
+        return self._httpd.server_address[1]
+
+    @property
+    def redirect_uri(self) -> str:
+        return f"http://127.0.0.1:{self.port}{self._callback_path}"
+
+    def wait(self, timeout: float) -> dict:
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            try:
+                self._httpd.handle_request()
+            except OSError:
+                # A connection that timed out or reset mid-read is not an outcome.
+                continue
+            if self._handler.result:
+                return self._handler.result
+        return {"error": f"timed out after {int(timeout)}s waiting for the browser"}
+
+    def close(self) -> None:
+        self._httpd.server_close()
+
+
+def authorize(
+    authorization_endpoint: str,
+    *,
+    client_id: str,
+    callback_path: str,
+    scope: Optional[str] = None,
+    extra_params: Optional[Mapping[str, str]] = None,
+    open_browser: Callable[[str], object] = webbrowser.open,
+    announce: Callable[[str], None] = print,
+    timeout: float = LOGIN_TIMEOUT_SECONDS,
+    port: int = 0,
+) -> AuthorizationResult:
+    """Run the browser step of an authorization-code + PKCE sign-in.
+
+    ``port`` pins the loopback port, which only an SSH forward needs: the tunnel
+    must name a port before the sign-in starts.
+    """
+    verifier = new_verifier()
+    state = _b64url(secrets.token_bytes(16))
+    try:
+        server = LoopbackServer(callback_path, state, port)
+    except OSError as err:
+        where = f"port {port}" if port else "a local port"
+        raise OAuthFlowError(f"could not bind {where} for the sign-in ({err})") from err
+    try:
+        params = {
+            "response_type": "code",
+            "client_id": client_id,
+            "redirect_uri": server.redirect_uri,
+            "state": state,
+            "code_challenge": challenge_for(verifier),
+            "code_challenge_method": "S256",
+        }
+        if scope:
+            params["scope"] = scope
+        params.update(extra_params or {})
+        separator = "&" if "?" in authorization_endpoint else "?"
+        url = f"{authorization_endpoint}{separator}{urllib.parse.urlencode(params)}"
+        announce(f"Opening your browser to sign in.\nIf it does not open, visit:\n  {url}")
+        open_browser(url)
+        outcome = server.wait(timeout)
+        redirect_uri = server.redirect_uri
+    finally:
+        server.close()
+    if "error" in outcome:
+        raise OAuthFlowError(outcome["error"])
+    return AuthorizationResult(
+        code=outcome["code"], redirect_uri=redirect_uri, verifier=verifier, iss=outcome.get("iss")
+    )
+
+
+@dataclass(frozen=True)
+class HttpResult:
+    status: int
+    body: dict
+
+
+def _read_json(raw: bytes) -> dict:
+    try:
+        parsed = json.loads(raw.decode("utf-8")) if raw else {}
+    except (ValueError, UnicodeDecodeError):
+        return {}
+    return parsed if isinstance(parsed, dict) else {}
+
+
+def _send(request: urllib.request.Request, timeout: float) -> HttpResult:
+    try:
+        with urllib.request.urlopen(
+            request, timeout=timeout, context=default_ssl_context()
+        ) as resp:
+            return HttpResult(resp.status, _read_json(resp.read()))
+    except urllib.error.HTTPError as err:
+        # An OAuth error response is a 4xx with a JSON body worth reading.
+        return HttpResult(err.code, _read_json(err.read()))
+    except (urllib.error.URLError, OSError) as err:
+        raise OAuthFlowError(f"could not reach {request.full_url} ({err})") from err
+
+
+def post_form(url: str, data: Mapping[str, str], timeout: float = _HTTP_TIMEOUT_SECONDS) -> HttpResult:
+    """POST ``application/x-www-form-urlencoded``, as every token endpoint takes it."""
+    request = urllib.request.Request(
+        url,
+        data=urllib.parse.urlencode(dict(data)).encode("ascii"),
+        headers={
+            "Content-Type": "application/x-www-form-urlencoded",
+            "Accept": "application/json",
+        },
+        method="POST",
+    )
+    return _send(request, timeout)
+
+
+def get_json(url: str, timeout: float = _HTTP_TIMEOUT_SECONDS) -> HttpResult:
+    request = urllib.request.Request(url, headers={"Accept": "application/json"}, method="GET")
+    return _send(request, timeout)
+
+
+def oauth_error(result: HttpResult) -> str:
+    """The RFC 6749 5.2 ``error`` code of a failed call, or its HTTP status."""
+    error = result.body.get("error")
+    return str(error) if error else f"HTTP {result.status}"
+
+
+__all__ = [
+    "LOGIN_TIMEOUT_SECONDS",
+    "AuthorizationResult",
+    "HttpResult",
+    "LoopbackServer",
+    "OAuthFlowError",
+    "authorize",
+    "challenge_for",
+    "default_ssl_context",
+    "get_json",
+    "new_verifier",
+    "oauth_error",
+    "post_form",
+]

onehouse_python_sdk/_session.py

@@ -0,0 +1,169 @@
+"""Where a signed-in session outlives the process, and the lock that guards refresh.
+
+Standard library only, and outside every subsystem, so the control-plane login
+and the AI-Gateway login share one store type. The CLI supplies a keychain store
+with the same interface; this module ships the ``0600`` file store it falls back
+to, so the SDK core never depends on ``keyring``.
+
+Refresh tokens rotate. Presenting a rotated one again reads as theft and revokes
+the whole chain, so two processes must never refresh from the same stored token.
+Callers take :meth:`SessionStore.lock`, re-read, refresh only if still needed,
+and write the new token before releasing.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import sys
+import tempfile
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Iterator, Optional, Protocol
+
+try:  # pragma: no cover - the platform decides which one exists
+    import fcntl
+except ImportError:  # pragma: no cover
+    fcntl = None  # type: ignore[assignment]
+try:  # pragma: no cover
+    import msvcrt
+except ImportError:  # pragma: no cover
+    msvcrt = None  # type: ignore[assignment]
+
+_LOCK_TIMEOUT_SECONDS = 30.0
+_LOCK_POLL_SECONDS = 0.05
+
+
+@dataclass
+class Session:
+    """``expires_at`` is the access token's expiry in epoch seconds."""
+
+    refresh_token: str
+    access_token: Optional[str] = None
+    expires_at: float = 0.0
+
+    def access_token_valid(self, skew_seconds: float, now: Optional[float] = None) -> bool:
+        current = time.time() if now is None else now
+        return bool(self.access_token) and self.expires_at - skew_seconds > current
+
+
+class SessionStore(Protocol):
+    def get(self) -> Optional[Session]: ...
+
+    def put(self, session: Session) -> None: ...
+
+    def delete(self) -> None: ...
+
+    def lock(self) -> contextlib.AbstractContextManager: ...
+
+
+class FileLock:
+    """An exclusive advisory lock on ``path``, held across processes.
+
+    ``fcntl.flock`` on POSIX, ``msvcrt.locking`` on Windows. Blocks up to
+    ``timeout`` seconds, then raises ``TimeoutError`` rather than hang a command.
+    """
+
+    def __init__(self, path: Path, timeout: float = _LOCK_TIMEOUT_SECONDS) -> None:
+        self._path = path
+        self._timeout = timeout
+        self._fd: Optional[int] = None
+
+    def __enter__(self) -> "FileLock":
+        self._path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
+        fd = os.open(str(self._path), os.O_RDWR | os.O_CREAT, 0o600)
+        deadline = time.monotonic() + self._timeout
+        while True:
+            try:
+                _try_lock(fd)
+                break
+            except OSError:
+                if time.monotonic() >= deadline:
+                    os.close(fd)
+                    raise TimeoutError(f"timed out waiting for the lock on {self._path}")
+                time.sleep(_LOCK_POLL_SECONDS)
+        self._fd = fd
+        return self
+
+    def __exit__(self, *_exc) -> None:
+        if self._fd is not None:
+            try:
+                _unlock(self._fd)
+            finally:
+                os.close(self._fd)
+                self._fd = None
+
+
+def _try_lock(fd: int) -> None:
+    if fcntl is not None:
+        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+    elif msvcrt is not None:  # pragma: no cover - Windows
+        msvcrt.locking(fd, msvcrt.LK_NBLCK, 1)
+    else:  # pragma: no cover
+        raise OSError("no file locking available on this platform")
+
+
+def _unlock(fd: int) -> None:
+    if fcntl is not None:
+        fcntl.flock(fd, fcntl.LOCK_UN)
+    elif msvcrt is not None:  # pragma: no cover - Windows
+        msvcrt.locking(fd, msvcrt.LK_UNLCK, 1)
+
+
+def write_private_file(path: Path, text: str) -> None:
+    """Write ``text`` so no other user can read it, replacing ``path`` atomically.
+
+    On Windows ``chmod`` cannot express ``0600``; the file then relies on the
+    user-profile ACL, as ``~/.onehouse/credentials`` already does.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
+    fd, tmp = tempfile.mkstemp(dir=str(path.parent), prefix=f".{path.name}.")
+    try:
+        if sys.platform != "win32":
+            os.fchmod(fd, 0o600)
+        with os.fdopen(fd, "w", encoding="utf-8") as handle:
+            handle.write(text)
+        os.replace(tmp, str(path))
+    except BaseException:
+        with contextlib.suppress(OSError):
+            os.unlink(tmp)
+        raise
+
+
+class FileSessionStore:
+    """A session in one ``0600`` JSON file, locked through a sibling ``.lock`` file."""
+
+    def __init__(self, path: Path) -> None:
+        self.path = Path(path).expanduser()
+
+    def get(self) -> Optional[Session]:
+        try:
+            payload = json.loads(self.path.read_text(encoding="utf-8"))
+        except FileNotFoundError:
+            return None
+        except (OSError, ValueError):
+            return None
+        if not isinstance(payload, dict) or not payload.get("refresh_token"):
+            return None
+        return Session(
+            refresh_token=str(payload["refresh_token"]),
+            access_token=payload.get("access_token") or None,
+            expires_at=float(payload.get("expires_at") or 0.0),
+        )
+
+    def put(self, session: Session) -> None:
+        write_private_file(self.path, json.dumps(asdict(session)))
+
+    def delete(self) -> None:
+        with contextlib.suppress(FileNotFoundError):
+            self.path.unlink()
+
+    @contextlib.contextmanager
+    def lock(self) -> Iterator[None]:
+        with FileLock(self.path.with_name(self.path.name + ".lock")):
+            yield
+
+
+__all__ = ["FileLock", "FileSessionStore", "Session", "SessionStore", "write_private_file"]

onehouse_python_sdk/ai_gateway/__init__.py

@@ -0,0 +1,29 @@
+"""AI Gateway tokens: discovery, service-principal minting and the human Keycloak session.
+
+A data-plane subsystem. Standard library only, and ``resources`` must not import
+it, so ``oh login`` never pulls in gateway code. The gateway runs in the customer's
+data plane behind its internal load balancer, so everything here needs VPN or the
+customer network.
+"""
+
+from .discovery import GatewayAuthMetadata, discover
+from .exceptions import AiGatewayError, AuthDisabledError
+from .tokens import (
+    KEYCLOAK_CLIENT_ID,
+    KeycloakSessionTokenSource,
+    ServicePrincipalTokenSource,
+    login,
+    logout,
+)
+
+__all__ = [
+    "KEYCLOAK_CLIENT_ID",
+    "AiGatewayError",
+    "AuthDisabledError",
+    "GatewayAuthMetadata",
+    "KeycloakSessionTokenSource",
+    "ServicePrincipalTokenSource",
+    "discover",
+    "login",
+    "logout",
+]

onehouse_python_sdk/ai_gateway/discovery.py

@@ -0,0 +1,143 @@
+"""Find the gateway's authorization server from the gateway itself.
+
+The Keycloak token URL is discovered, never derived from the gateway host or the
+link uid, and never configured: the same chain an MCP client follows.
+
+1. RFC 9728 protected-resource metadata from the gateway, path form first:
+   ``resource`` is the ``aud`` the edge requires, ``authorization_servers[0]``
+   the issuer. Images from before ENG-50210 serve only the bare path.
+2. RFC 8414 metadata for that issuer, inserted form first, then OIDC discovery.
+   The returned ``issuer`` must equal the one asked for (RFC 8414 3.3).
+"""
+
+from __future__ import annotations
+
+import threading
+import urllib.parse
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from .. import _oauth
+from .exceptions import AiGatewayError, AuthDisabledError
+
+_PROTECTED_RESOURCE_PATHS = (
+    "/.well-known/oauth-protected-resource/mcp",
+    "/.well-known/oauth-protected-resource",
+)
+_VPN_HINT = "AI Gateway hosts resolve only on the VPN or the customer network"
+
+_cache: Dict[str, "GatewayAuthMetadata"] = {}
+_cache_lock = threading.Lock()
+
+
+@dataclass(frozen=True)
+class GatewayAuthMetadata:
+    gateway_url: str
+    resource: str
+    issuer: str
+    token_endpoint: str
+    authorization_endpoint: Optional[str]
+    end_session_endpoint: Optional[str]
+
+
+def normalize_gateway_url(gateway_url: str) -> str:
+    parts = urllib.parse.urlsplit(gateway_url.strip())
+    if parts.scheme != "https" or not parts.netloc:
+        raise AiGatewayError(
+            f"AI Gateway URL {gateway_url!r} must be an https URL, such as "
+            "https://ai-gateway-<id>.dp.onehouse.ai"
+        )
+    return f"https://{parts.netloc.lower()}"
+
+
+def _get(url: str) -> _oauth.HttpResult:
+    try:
+        return _oauth.get_json(url)
+    except _oauth.OAuthFlowError as err:
+        raise AiGatewayError(f"{err}. {_VPN_HINT}.") from err
+
+
+def _is_auth_disabled(result: _oauth.HttpResult) -> bool:
+    error = result.body.get("error")
+    return result.status == 404 and isinstance(error, dict) and error.get("code") == "auth_disabled"
+
+
+def _protected_resource(gateway_url: str) -> dict:
+    last: Optional[_oauth.HttpResult] = None
+    for path in _PROTECTED_RESOURCE_PATHS:
+        result = _get(gateway_url + path)
+        if _is_auth_disabled(result):
+            raise AuthDisabledError(
+                f"{gateway_url} does not require authentication, so no token is needed"
+            )
+        if result.status == 200:
+            return result.body
+        last = result
+    raise AiGatewayError(
+        f"{gateway_url} published no protected-resource metadata "
+        f"({_oauth.oauth_error(last) if last else 'no response'})"
+    )
+
+
+def _authorization_server_metadata(issuer: str) -> dict:
+    parts = urllib.parse.urlsplit(issuer)
+    origin = f"{parts.scheme}://{parts.netloc}"
+    candidates = (
+        f"{origin}/.well-known/oauth-authorization-server{parts.path.rstrip('/')}",
+        f"{issuer.rstrip('/')}/.well-known/openid-configuration",
+    )
+    for url in candidates:
+        result = _get(url)
+        if result.status != 200:
+            continue
+        if result.body.get("issuer") != issuer:
+            raise AiGatewayError(
+                f"Authorization-server metadata at {url} names issuer "
+                f"{result.body.get('issuer')!r}, not {issuer!r}"
+            )
+        return result.body
+    raise AiGatewayError(f"No authorization-server metadata found for {issuer}")
+
+
+def discover(gateway_url: str, *, use_cache: bool = True) -> GatewayAuthMetadata:
+    """The gateway's audience and Keycloak endpoints, cached per gateway URL."""
+    gateway = normalize_gateway_url(gateway_url)
+    if use_cache:
+        with _cache_lock:
+            cached = _cache.get(gateway)
+        if cached is not None:
+            return cached
+
+    resource_metadata = _protected_resource(gateway)
+    resource = resource_metadata.get("resource")
+    servers = resource_metadata.get("authorization_servers") or []
+    if not isinstance(resource, str) or not resource or not servers:
+        raise AiGatewayError(f"{gateway} published incomplete protected-resource metadata")
+    # The audience must be this gateway's, or a token minted for it is useless here.
+    if urllib.parse.urlsplit(resource).netloc.lower() != urllib.parse.urlsplit(gateway).netloc:
+        raise AiGatewayError(f"{gateway} names a resource on another host: {resource!r}")
+    issuer = str(servers[0])
+    as_metadata = _authorization_server_metadata(issuer)
+    token_endpoint = as_metadata.get("token_endpoint")
+    if not token_endpoint:
+        raise AiGatewayError(f"Authorization server {issuer} published no token_endpoint")
+
+    metadata = GatewayAuthMetadata(
+        gateway_url=gateway,
+        resource=resource,
+        issuer=issuer,
+        token_endpoint=token_endpoint,
+        authorization_endpoint=as_metadata.get("authorization_endpoint"),
+        end_session_endpoint=as_metadata.get("end_session_endpoint"),
+    )
+    with _cache_lock:
+        _cache[gateway] = metadata
+    return metadata
+
+
+def clear_cache() -> None:
+    with _cache_lock:
+        _cache.clear()
+
+
+__all__ = ["GatewayAuthMetadata", "clear_cache", "discover", "normalize_gateway_url"]

onehouse_python_sdk/ai_gateway/exceptions.py

@@ -0,0 +1,11 @@
+"""Exceptions raised by the AI Gateway subsystem."""
+
+from onehouse_python_sdk.exceptions import OnehouseSdkError
+
+
+class AiGatewayError(OnehouseSdkError):
+    """Discovery, minting or the sign-in failed."""
+
+
+class AuthDisabledError(AiGatewayError):
+    """The gateway has authentication off, so no token is needed."""