PyPI - codex-backend-sdk - Versions diffs - 0.1.1__py3-none-any.whl - Mend

codex-backend-sdk 0.1.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (9) hide show

codex_backend_sdk/__init__.py +72 -0
codex_backend_sdk/codex_client.py +1084 -0
codex_backend_sdk/oauth.py +299 -0
codex_backend_sdk/pkce.py +29 -0
codex_backend_sdk/storage.py +166 -0
codex_backend_sdk-0.1.1.dist-info/METADATA +369 -0
codex_backend_sdk-0.1.1.dist-info/RECORD +9 -0
codex_backend_sdk-0.1.1.dist-info/WHEEL +4 -0
codex_backend_sdk-0.1.1.dist-info/licenses/LICENSE +21 -0

codex_backend_sdk/oauth.py ADDED Viewed

@@ -0,0 +1,299 @@
+"""
+ChatGPT OAuth flow — mirrors codex-rs/login/src/server.rs.
+Flow:
+  1. Generate PKCE + state
+  2. Build authorization URL → open browser
+  3. Listen on localhost:1455 for the OAuth callback
+  4. Exchange authorization code for tokens
+  5. Optionally exchange id_token for an OpenAI API key
+  6. Persist tokens to ~/.codex/auth.json
+"""
+import threading
+import urllib.parse
+import webbrowser
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import Optional
+import requests
+from .pkce import PkceCodes, generate_pkce, generate_state
+from .storage import TokenStore, save_tokens
+ISSUER = "https://auth.openai.com"
+CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
+CALLBACK_PORT = 1455
+REDIRECT_URI = f"http://localhost:{CALLBACK_PORT}/auth/callback"
+SCOPES = "openid profile email offline_access api.connectors.read api.connectors.invoke"
+ORIGINATOR = "codex_cli_rs"
+# Result container shared between the HTTP handler and the caller
+_oauth_result: dict = {}
+_oauth_event = threading.Event()   # set when code received
+_server_done_event = threading.Event()  # set when /success is served
+def _build_authorize_url(
+    pkce: PkceCodes,
+    state: str,
+    workspace_id: Optional[str] = None,
+    scopes: Optional[str] = None,
+) -> str:
+    params = {
+        "response_type": "code",
+        "client_id": CLIENT_ID,
+        "redirect_uri": REDIRECT_URI,
+        "scope": scopes or SCOPES,
+        "code_challenge": pkce.code_challenge,
+        "code_challenge_method": "S256",
+        "id_token_add_organizations": "true",
+        "codex_cli_simplified_flow": "true",
+        "state": state,
+        "originator": ORIGINATOR,
+    }
+    if workspace_id:
+        params["allowed_workspace_id"] = workspace_id
+    # Use quote (RFC 3986, spaces → %20) not quote_plus (spaces → +),
+    # matching the urlencoding::encode behaviour in codex-rs/login/src/server.rs.
+    qs = "&".join(f"{k}={urllib.parse.quote(v, safe='')}" for k, v in params.items())
+    return f"{ISSUER}/oauth/authorize?{qs}"
+def _make_handler(pkce: PkceCodes, state: str):
+    """Factory that returns an HTTP handler class closed over pkce/state."""
+    class _CallbackHandler(BaseHTTPRequestHandler):
+        def log_message(self, fmt, *args):  # suppress default access log
+            pass
+        def do_GET(self):
+            parsed = urllib.parse.urlparse(self.path)
+            params = dict(urllib.parse.parse_qsl(parsed.query))
+            if parsed.path == "/auth/callback":
+                self._handle_callback(params)
+            elif parsed.path == "/success":
+                self._send_html(200, (
+                    "<html><body style='font-family:sans-serif;padding:2em'>"
+                    "<h2>✓ Login successful</h2>"
+                    "<p>You can close this tab and return to the terminal.</p>"
+                    "</body></html>"
+                ))
+                _server_done_event.set()
+            else:
+                self._send_text(404, "Not found")
+        def _handle_callback(self, params: dict):
+            if params.get("state") != state:
+                self._send_text(400, "State mismatch — possible CSRF attack.")
+                _oauth_result["error"] = "state_mismatch"
+                _oauth_event.set()
+                return
+            if error := params.get("error"):
+                desc = params.get("error_description", error)
+                self._send_html(400, f"<h2>Login failed</h2><p>{desc}</p>")
+                _oauth_result["error"] = desc
+                _oauth_event.set()
+                return
+            code = params.get("code", "")
+            if not code:
+                self._send_text(400, "Missing authorization code.")
+                _oauth_result["error"] = "missing_code"
+                _oauth_event.set()
+                return
+            # Redirect browser to /success while we process in background
+            self.send_response(302)
+            self.send_header("Location", f"http://localhost:{CALLBACK_PORT}/success")
+            self.end_headers()
+            _oauth_result["code"] = code
+            _oauth_event.set()
+        def _send_text(self, status: int, body: str):
+            data = body.encode()
+            self.send_response(status)
+            self.send_header("Content-Type", "text/plain; charset=utf-8")
+            self.send_header("Content-Length", str(len(data)))
+            self.end_headers()
+            self.wfile.write(data)
+        def _send_html(self, status: int, body: str):
+            data = body.encode()
+            self.send_response(status)
+            self.send_header("Content-Type", "text/html; charset=utf-8")
+            self.send_header("Content-Length", str(len(data)))
+            self.end_headers()
+            self.wfile.write(data)
+    return _CallbackHandler
+def _exchange_code(code: str, pkce: PkceCodes) -> dict:
+    """POST authorization code to token endpoint, return raw JSON."""
+    resp = requests.post(
+        f"{ISSUER}/oauth/token",
+        headers={"Content-Type": "application/x-www-form-urlencoded"},
+        data={
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": REDIRECT_URI,
+            "client_id": CLIENT_ID,
+            "code_verifier": pkce.code_verifier,
+        },
+        timeout=30,
+    )
+    resp.raise_for_status()
+    return resp.json()
+def obtain_api_key(id_token: str, access_token: Optional[str] = None) -> Optional[str]:
+    """
+    Exchange an id_token (or access_token fallback) for an OpenAI API key.
+    Returns the API key string, or None on failure.
+    Some accounts (personal API Platform orgs) don't embed organization_id in
+    the id_token. We first try with the id_token; if that fails with an org
+    error, we retry using the access_token (already scoped to api.openai.com/v1).
+    """
+    def _exchange(subject: str, token_type: str) -> Optional[str]:
+        resp = requests.post(
+            f"{ISSUER}/oauth/token",
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            data={
+                "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+                "client_id": CLIENT_ID,
+                "requested_token": "openai-api-key",
+                "subject_token": subject,
+                "subject_token_type": token_type,
+            },
+            timeout=30,
+        )
+        if resp.ok:
+            return resp.json().get("access_token")
+        short_type = token_type.split(":")[-1]
+        print(f"[auth] {short_type} exchange → {resp.status_code}: {resp.text}")
+        return None
+    try:
+        key = _exchange(id_token, "urn:ietf:params:oauth:token-type:id_token")
+        if key:
+            return key
+        if access_token:
+            print("[auth] retrying with access_token…")
+            key = _exchange(access_token, "urn:ietf:params:oauth:token-type:access_token")
+            if key:
+                return key
+    except Exception as exc:
+        print(f"[auth] API key exchange error: {exc}")
+    return None
+def refresh_access_token(refresh_token: str) -> dict:
+    """
+    Use a refresh token to get a new access_token (and optionally a new refresh_token).
+    Returns the raw JSON response from the token endpoint.
+    """
+    resp = requests.post(
+        f"{ISSUER}/oauth/token",
+        headers={"Content-Type": "application/json"},
+        json={
+            "client_id": CLIENT_ID,
+            "grant_type": "refresh_token",
+            "refresh_token": refresh_token,
+            "scope": "openid profile email offline_access",
+        },
+        timeout=30,
+    )
+    resp.raise_for_status()
+    return resp.json()
+def run_oauth_flow(
+    *,
+    open_browser: bool = True,
+    request_api_key: bool = True,
+    persist: bool = True,
+    workspace_id: Optional[str] = None,
+    scopes: Optional[str] = None,
+) -> TokenStore:
+    """
+    Run the full OAuth authorization-code + PKCE flow.
+    Opens the browser (unless open_browser=False), waits for the callback,
+    exchanges the code for tokens, optionally fetches an API key,
+    and saves everything to ~/.codex/auth.json.
+    Returns a TokenStore with all credentials.
+    """
+    global _oauth_result, _oauth_event, _server_done_event
+    _oauth_result = {}
+    _oauth_event = threading.Event()
+    _server_done_event = threading.Event()
+    pkce = generate_pkce()
+    state = generate_state()
+    handler_cls = _make_handler(pkce, state)
+    try:
+        server = HTTPServer(("127.0.0.1", CALLBACK_PORT), handler_cls)
+    except OSError as exc:
+        raise RuntimeError(
+            f"Cannot bind to localhost:{CALLBACK_PORT} — "
+            "another process may still be holding that port from a previous login attempt. "
+            f"(OS error: {exc})"
+        ) from exc
+    server.timeout = 1  # allow periodic _oauth_event checks
+    auth_url = _build_authorize_url(pkce, state, workspace_id=workspace_id, scopes=scopes)
+    print(f"\n[auth] Opening browser for login:\n  {auth_url}\n")
+    if open_browser:
+        webbrowser.open(auth_url)
+    else:
+        print("[auth] open_browser=False — open the URL above manually.")
+    print("[auth] Waiting for OAuth callback on localhost:1455 …")
+    # Keep serving until /success is served (happy path) or callback signals an error
+    while not _server_done_event.is_set():
+        server.handle_request()
+        if _oauth_event.is_set() and _oauth_result.get("error"):
+            break
+    server.server_close()
+    if error := _oauth_result.get("error"):
+        raise RuntimeError(f"OAuth callback error: {error}")
+    code = _oauth_result["code"]
+    print("[auth] Authorization code received — exchanging for tokens …")
+    raw = _exchange_code(code, pkce)
+    access_token = raw["access_token"]
+    refresh_token = raw["refresh_token"]
+    id_token = raw["id_token"]
+    api_key: Optional[str] = None
+    if request_api_key:
+        print("[auth] Exchanging for OpenAI API key …")
+        api_key = obtain_api_key(id_token, access_token)
+        if api_key:
+            print(f"[auth] API key obtained: {api_key[:8]}…")
+    store = TokenStore.from_exchange(
+        access_token=access_token,
+        refresh_token=refresh_token,
+        id_token=id_token,
+        api_key=api_key,
+    )
+    if persist:
+        save_tokens(store)
+        print("[auth] Tokens saved to ~/.codex/auth.json")
+    return store

codex_backend_sdk/pkce.py ADDED Viewed

@@ -0,0 +1,29 @@
+"""PKCE (Proof Key for Code Exchange) helpers — mirrors codex-rs/login/src/pkce.rs."""
+import base64
+import hashlib
+import os
+from dataclasses import dataclass
+@dataclass
+class PkceCodes:
+    code_verifier: str
+    code_challenge: str
+def generate_pkce() -> PkceCodes:
+    """Generate a PKCE pair using S256 method."""
+    raw = os.urandom(64)
+    code_verifier = base64.urlsafe_b64encode(raw).rstrip(b"=").decode()
+    digest = hashlib.sha256(code_verifier.encode()).digest()
+    code_challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
+    return PkceCodes(code_verifier=code_verifier, code_challenge=code_challenge)
+def generate_state() -> str:
+    """Generate a cryptographically random state parameter."""
+    raw = os.urandom(32)
+    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()

codex_backend_sdk/storage.py ADDED Viewed

@@ -0,0 +1,166 @@
+"""Token persistence — mirrors codex-rs/login/src/auth/storage.rs.
+Stores tokens in ~/.codex/auth.json (same format as the Codex CLI).
+"""
+import base64
+import json
+import os
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+REFRESH_EXPIRY_MARGIN_S = 5 * 60   # refresh if exp within 5 minutes
+REFRESH_INTERVAL_S = 55 * 60       # refresh at least every 55 minutes
+def _codex_home() -> Path:
+    if env := os.environ.get("CODEX_HOME"):
+        return Path(env)
+    return Path.home() / ".codex"
+def _auth_path() -> Path:
+    return _codex_home() / "auth.json"
+def _decode_jwt_payload(token: str) -> dict:
+    """Decode the payload section of a JWT without verification."""
+    parts = token.split(".")
+    if len(parts) < 2:
+        return {}
+    payload_b64 = parts[1]
+    # Add padding if needed
+    padding = 4 - len(payload_b64) % 4
+    if padding != 4:
+        payload_b64 += "=" * padding
+    try:
+        raw = base64.urlsafe_b64decode(payload_b64)
+        return json.loads(raw)
+    except Exception:
+        return {}
+def _extract_openai_claims(token: str) -> dict:
+    """Extract the nested 'https://api.openai.com/auth' claims from a JWT."""
+    payload = _decode_jwt_payload(token)
+    return payload.get("https://api.openai.com/auth", {})
+@dataclass
+class TokenStore:
+    access_token: str
+    refresh_token: str
+    id_token_raw: str
+    account_id: Optional[str] = None
+    openai_api_key: Optional[str] = None
+    last_refresh: Optional[str] = None
+    # Decoded claims from id_token for convenience
+    email: Optional[str] = None
+    plan_type: Optional[str] = None
+    @classmethod
+    def from_exchange(
+        cls,
+        access_token: str,
+        refresh_token: str,
+        id_token: str,
+        api_key: Optional[str] = None,
+    ) -> "TokenStore":
+        claims = _extract_openai_claims(id_token)
+        access_claims = _extract_openai_claims(access_token)
+        return cls(
+            access_token=access_token,
+            refresh_token=refresh_token,
+            id_token_raw=id_token,
+            account_id=claims.get("chatgpt_account_id"),
+            openai_api_key=api_key,
+            last_refresh=datetime.now(timezone.utc).isoformat(),
+            email=_decode_jwt_payload(id_token).get("email"),
+            plan_type=access_claims.get("chatgpt_plan_type"),
+        )
+def token_needs_refresh(store: TokenStore) -> bool:
+    """
+    Return True if the access token should be refreshed before use.
+    Mirrors openai-oauth auth.ts shouldRefreshAccessToken:
+      - exp claim within REFRESH_EXPIRY_MARGIN_S seconds → refresh
+      - last_refresh older than REFRESH_INTERVAL_S → refresh
+    """
+    now = datetime.now(timezone.utc).timestamp()
+    claims = _decode_jwt_payload(store.access_token)
+    exp = claims.get("exp")
+    if isinstance(exp, (int, float)):
+        if exp <= now + REFRESH_EXPIRY_MARGIN_S:
+            return True
+        # Token has a valid exp and it's not near expiry — trust it
+        return False
+    # No exp claim — fall back to last_refresh timestamp
+    if store.last_refresh:
+        try:
+            refreshed_at = datetime.fromisoformat(store.last_refresh)
+            if (datetime.now(timezone.utc) - refreshed_at).total_seconds() >= REFRESH_INTERVAL_S:
+                return True
+        except Exception:
+            pass
+    return False
+def save_tokens(store: TokenStore) -> None:
+    path = _auth_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "auth_mode": "chatgpt",
+        "openai_api_key": store.openai_api_key,
+        "tokens": {
+            "access_token": store.access_token,
+            "refresh_token": store.refresh_token,
+            "id_token": store.id_token_raw,
+            "account_id": store.account_id,
+        },
+        "last_refresh": store.last_refresh,
+    }
+    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+    # Restrict permissions on Unix
+    try:
+        os.chmod(path, 0o600)
+    except OSError:
+        pass
+def load_tokens() -> Optional[TokenStore]:
+    path = _auth_path()
+    if not path.exists():
+        return None
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except Exception:
+        return None
+    tokens = data.get("tokens") or {}
+    access_token = tokens.get("access_token", "")
+    refresh_token = tokens.get("refresh_token", "")
+    id_token = tokens.get("id_token", "")
+    if not access_token:
+        return None
+    return TokenStore(
+        access_token=access_token,
+        refresh_token=refresh_token,
+        id_token_raw=id_token,
+        account_id=tokens.get("account_id"),
+        openai_api_key=data.get("openai_api_key"),
+        last_refresh=data.get("last_refresh"),
+        email=_decode_jwt_payload(id_token).get("email"),
+        plan_type=_extract_openai_claims(access_token).get("chatgpt_plan_type"),
+    )