stratonext 0.0.1__py3-none-any.whl

stratonext/auth.py

@@ -0,0 +1,287 @@
+"""OAuth2 user-login helpers for the stratonext CLI (AWS Cognito).
+
+The CLI authenticates as a human user via Authorization Code + PKCE (RFC 7636)
+with a localhost redirect. Cognito does not support the Device Authorization
+Grant (RFC 8628), so the CLI opens a browser tab and listens on
+http://localhost:19191 for the redirect (see ADR-031, ADR-035). Tokens are
+cached in ~/.stratonext/credentials with refresh-token support.
+
+Endpoints are resolved from the Cognito OIDC issuer (STRATONEXT_AUTH_AUTHORITY)
+via its /.well-known/openid-configuration document. Scopes are built from the
+platform resource-server identifier (STRATONEXT_AUTH_AUDIENCE) because Cognito
+custom scopes must be prefixed with it, e.g.
+"https://platform.prod.stratonext.com/organization:read".
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import http.server
+import json
+import os
+import secrets
+import sys
+import time
+import urllib.parse
+import webbrowser
+from dataclasses import dataclass
+from pathlib import Path
+
+import httpx
+
+from .settings import resolve
+
+CREDENTIALS_PATH = Path.home() / ".stratonext" / "credentials"
+
+# Refresh this many seconds before actual expiry, same margin AgentCallerTokenProvider uses.
+EXPIRY_MARGIN_SECONDS = 60
+
+# Localhost redirect for the Authorization Code + PKCE flow. Must match the CLI
+# client's callback URL in Cognito (cognito-stack.ts).
+REDIRECT_PORT = 19191
+REDIRECT_URI = f"http://localhost:{REDIRECT_PORT}"
+LOGIN_TIMEOUT = 300
+
+# In-process cache of OIDC discovery documents, keyed by issuer.
+_discovery_cache: dict[str, dict] = {}
+
+# Production defaults — the CLI targets prod unless the matching STRATONEXT_AUTH_*
+# / STRATONEXT_API_URL env var overrides it (set those for local or non-prod dev).
+DEFAULT_AUTHORITY = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_26SWBDh3Q"
+DEFAULT_CLI_CLIENT_ID = "60evptk8ab01298r41na2et8p8"
+DEFAULT_PLATFORM_AUDIENCE = "https://platform.prod.stratonext.com"
+
+
+def _user_scopes(audience: str, *, readonly: bool = False) -> str:
+    """Cognito user-login scopes — OIDC scopes plus resource-server-prefixed org scopes."""
+    scopes = ["openid", f"{audience}/organization:read"]
+    if not readonly:
+        scopes.append(f"{audience}/organization:write")
+    return " ".join(scopes)
+
+
+async def _discover(authority: str, http_client: httpx.AsyncClient) -> dict:
+    """Fetch and cache the Cognito OIDC discovery document for an issuer."""
+    if authority in _discovery_cache:
+        return _discovery_cache[authority]
+    resp = await http_client.get(f"{authority.rstrip('/')}/.well-known/openid-configuration")
+    resp.raise_for_status()
+    doc = resp.json()
+    _discovery_cache[authority] = doc
+    return doc
+
+
+@dataclass(frozen=True)
+class UserAuthConfig:
+    """Public-client config for the OAuth2 Authorization Code + PKCE flow (RFC 7636)."""
+
+    authority: str
+    client_id: str
+    audience: str
+    scopes: str
+
+    @classmethod
+    def from_env(cls, *, readonly: bool = False) -> UserAuthConfig:
+        # All values default to prod; override with the matching env var for dev.
+        audience = resolve("STRATONEXT_AUTH_AUDIENCE", DEFAULT_PLATFORM_AUDIENCE)
+        return cls(
+            authority=resolve("STRATONEXT_AUTH_AUTHORITY", DEFAULT_AUTHORITY),
+            client_id=resolve("STRATONEXT_AUTH_CLIENT_ID", DEFAULT_CLI_CLIENT_ID),
+            audience=audience,
+            scopes=_user_scopes(audience, readonly=readonly),
+        )
+
+
+async def get_access_token_user(
+    config: UserAuthConfig,
+    *,
+    http_client: httpx.AsyncClient | None = None,
+    credentials_path: Path = CREDENTIALS_PATH,
+) -> str:
+    """User auth — check cache, try refresh, then full Authorization Code + PKCE flow."""
+    creds = _read_credentials(credentials_path)
+
+    if creds and creds.get("access_token") and creds.get("expires_at", 0) > time.time():
+        return creds["access_token"]
+
+    if creds and creds.get("refresh_token"):
+        try:
+            return await _refresh_access_token(config, creds["refresh_token"], http_client, credentials_path)
+        except httpx.HTTPStatusError:
+            pass  # refresh expired or revoked — fall through to full login
+
+    return await initiate_login(config, http_client=http_client, credentials_path=credentials_path)
+
+
+async def initiate_login(
+    config: UserAuthConfig,
+    *,
+    http_client: httpx.AsyncClient | None = None,
+    credentials_path: Path = CREDENTIALS_PATH,
+) -> str:
+    """Run Authorization Code + PKCE: open browser, capture localhost redirect, exchange code."""
+    verifier = secrets.token_urlsafe(64)
+    challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode()
+    state = secrets.token_urlsafe(16)
+
+    owns_client = http_client is None
+    http_client = http_client or httpx.AsyncClient()
+    try:
+        discovery = await _discover(config.authority, http_client)
+        authorize_url = discovery["authorization_endpoint"] + "?" + urllib.parse.urlencode(
+            {
+                "response_type": "code",
+                "client_id": config.client_id,
+                "redirect_uri": REDIRECT_URI,
+                "scope": config.scopes,
+                "state": state,
+                "code_challenge": challenge,
+                "code_challenge_method": "S256",
+                # Skip the Cognito hosted-UI provider chooser and go straight to Google
+                # (the CLI client only allows Google sign-in).
+                "identity_provider": "Google",
+            }
+        )
+
+        print(f"\nOpening browser to log in:\n  {authorize_url}", file=sys.stderr)
+        webbrowser.open(authorize_url)
+        print("Waiting for authentication...", file=sys.stderr)
+
+        params = await _await_redirect(state)
+        if "error" in params:
+            raise RuntimeError(f"Login failed: {params['error'][0]}")
+
+        token_resp = await http_client.post(
+            discovery["token_endpoint"],
+            data={
+                "grant_type": "authorization_code",
+                "client_id": config.client_id,
+                "code": params["code"][0],
+                "redirect_uri": REDIRECT_URI,
+                "code_verifier": verifier,
+            },
+        )
+        token_resp.raise_for_status()
+        body = token_resp.json()
+        _write_credentials(body, credentials_path, scopes=config.scopes)
+        print("Logged in.", file=sys.stderr)
+        return body["access_token"]
+    finally:
+        if owns_client:
+            await http_client.aclose()
+
+
+def clear_credentials(credentials_path: Path = CREDENTIALS_PATH) -> None:
+    """Remove cached user credentials (logout)."""
+    if credentials_path.exists():
+        credentials_path.unlink()
+
+
+async def _await_redirect(state: str) -> dict[str, list[str]]:
+    """Serve http://localhost:19191 until the OAuth redirect arrives; return its query params."""
+    import asyncio
+
+    return await asyncio.get_event_loop().run_in_executor(None, _serve_until_redirect, state)
+
+
+def _serve_until_redirect(state: str) -> dict[str, list[str]]:
+    captured: dict[str, list[str]] = {}
+
+    class _Handler(http.server.BaseHTTPRequestHandler):
+        def do_GET(self) -> None:  # noqa: N802 (stdlib naming)
+            params = urllib.parse.parse_qs(urllib.parse.urlparse(self.path).query)
+            if "code" in params or "error" in params:
+                captured.update(params)
+            self.send_response(200)
+            self.send_header("Content-Type", "text/html")
+            self.end_headers()
+            self.wfile.write(b"<html><body>Login complete. You can close this tab.</body></html>")
+
+        def log_message(self, *args: object) -> None:  # silence request logging
+            pass
+
+    server = http.server.HTTPServer(("localhost", REDIRECT_PORT), _Handler)
+    server.timeout = LOGIN_TIMEOUT
+    try:
+        deadline = time.time() + LOGIN_TIMEOUT
+        while not captured and time.time() < deadline:
+            server.handle_request()  # blocks for one request (favicon, etc. are ignored)
+    finally:
+        server.server_close()
+
+    if not captured:
+        raise RuntimeError("Login timed out. Run 'stratonext login' to try again.")
+    if captured.get("state", [None])[0] != state:
+        raise RuntimeError("Login failed: OAuth state mismatch (possible CSRF).")
+    return captured
+
+
+async def _refresh_access_token(
+    config: UserAuthConfig,
+    refresh_token: str,
+    http_client: httpx.AsyncClient | None,
+    credentials_path: Path,
+) -> str:
+    owns_client = http_client is None
+    http_client = http_client or httpx.AsyncClient()
+    try:
+        token_endpoint = (await _discover(config.authority, http_client))["token_endpoint"]
+        resp = await http_client.post(
+            token_endpoint,
+            data={
+                "grant_type": "refresh_token",
+                "client_id": config.client_id,
+                "refresh_token": refresh_token,
+            },
+        )
+        resp.raise_for_status()
+        body = resp.json()
+    finally:
+        if owns_client:
+            await http_client.aclose()
+
+    # Cognito refresh responses omit the refresh_token — preserve the existing one.
+    _write_credentials(body, credentials_path, refresh_token=refresh_token)
+    return body["access_token"]
+
+
+def token_scopes(access_token: str) -> str | None:
+    """Read the space-separated `scope` claim from a JWT access token (no signature check)."""
+    try:
+        payload = access_token.split(".")[1]
+        payload += "=" * (-len(payload) % 4)  # restore base64 padding
+        return json.loads(base64.urlsafe_b64decode(payload)).get("scope")
+    except (IndexError, ValueError, json.JSONDecodeError):
+        return None
+
+
+def _read_credentials(credentials_path: Path) -> dict | None:
+    if not credentials_path.exists():
+        return None
+    try:
+        return json.loads(credentials_path.read_text())
+    except (json.JSONDecodeError, OSError):
+        return None
+
+
+def _write_credentials(
+    token_response: dict,
+    credentials_path: Path,
+    *,
+    scopes: str | None = None,
+    refresh_token: str | None = None,
+) -> None:
+    existing = _read_credentials(credentials_path) or {}
+    data = {
+        "access_token": token_response["access_token"],
+        # Refresh responses don't repeat the refresh_token; fall back to the caller's / cached one.
+        "refresh_token": token_response.get("refresh_token") or refresh_token or existing.get("refresh_token"),
+        "expires_at": time.time() + token_response.get("expires_in", 3600) - EXPIRY_MARGIN_SECONDS,
+        # Preserve scopes from the original login; refresh responses don't repeat them.
+        "scopes": scopes or existing.get("scopes"),
+        "user_id": existing.get("user_id"),
+    }
+    credentials_path.parent.mkdir(parents=True, exist_ok=True)
+    credentials_path.write_text(json.dumps(data))
+    credentials_path.chmod(0o600)