chore: import upstream snapshot with attribution

2026-07-13 13:03:19 +08:00
commit 91e75e620b
3227 changed files with 1307078 additions and 0 deletions
@@ -0,0 +1,5 @@
+"""Authentication module for CUA CLI."""
+
+from .store import CredentialStore, clear_credentials, get_api_key, save_api_key
+
+__all__ = ["CredentialStore", "get_api_key", "save_api_key", "clear_credentials"]
@@ -0,0 +1,177 @@
+"""Browser-based OAuth authentication for Cua CLI."""
+
+import asyncio
+import os
+import platform
+import subprocess
+import urllib.parse
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from threading import Thread
+from typing import Optional
+
+# Default URLs
+DEFAULT_WEBSITE_URL = "https://cua.ai"
+
+# Timeout for browser authentication (2 minutes)
+AUTH_TIMEOUT_SECONDS = 120
+
+
+@dataclass
+class AuthResult:
+    """Result from browser authentication containing token and workspace metadata."""
+
+    token: str
+    workspace_slug: str
+    workspace_name: str
+    org_name: str
+
+
+class CallbackHandler(BaseHTTPRequestHandler):
+    """HTTP request handler for OAuth callback."""
+
+    token: Optional[str] = None
+    workspace_slug: Optional[str] = None
+    workspace_name: Optional[str] = None
+    org_name: Optional[str] = None
+    error: Optional[str] = None
+
+    def do_GET(self) -> None:
+        """Handle GET request from OAuth callback."""
+        parsed = urllib.parse.urlparse(self.path)
+        params = urllib.parse.parse_qs(parsed.query)
+
+        if "token" in params:
+            CallbackHandler.token = params["token"][0]
+            CallbackHandler.workspace_slug = params.get("workspace_slug", [""])[0]
+            CallbackHandler.workspace_name = params.get("workspace_name", [""])[0]
+            CallbackHandler.org_name = params.get("org_name", [""])[0]
+            self._send_response(
+                200,
+                "<html><body><h1>Authentication successful!</h1>"
+                "<p>You can close this window and return to the terminal.</p></body></html>",
+            )
+        elif "error" in params:
+            CallbackHandler.error = params.get("error_description", params["error"])[0]
+            self._send_response(
+                400,
+                f"<html><body><h1>Authentication failed</h1>"
+                f"<p>{CallbackHandler.error}</p></body></html>",
+            )
+        else:
+            self._send_response(
+                400,
+                "<html><body><h1>Invalid callback</h1>"
+                "<p>Missing token parameter.</p></body></html>",
+            )
+
+    def _send_response(self, status: int, body: str) -> None:
+        """Send an HTTP response."""
+        self.send_response(status)
+        self.send_header("Content-Type", "text/html")
+        self.end_headers()
+        self.wfile.write(body.encode())
+
+    def log_message(self, format: str, *args) -> None:
+        """Suppress default logging."""
+        pass
+
+
+def open_browser(url: str) -> bool:
+    """Open a URL in the default browser.
+
+    Args:
+        url: The URL to open
+
+    Returns:
+        True if the browser was opened successfully
+    """
+    system = platform.system()
+
+    try:
+        if system == "Darwin":
+            subprocess.run(["open", url], check=True)
+        elif system == "Windows":
+            subprocess.run(["cmd", "/c", "start", url], check=True, shell=True)
+        else:  # Linux and others
+            subprocess.run(["xdg-open", url], check=True)
+        return True
+    except (subprocess.SubprocessError, FileNotFoundError):
+        return False
+
+
+async def authenticate_via_browser(
+    website_url: Optional[str] = None,
+    workspace_slug: Optional[str] = None,
+) -> AuthResult:
+    """Authenticate via browser OAuth flow.
+
+    Starts a local HTTP server to receive the callback, opens the browser
+    to the authentication page, and waits for the token.
+
+    Args:
+        website_url: Base URL for the authentication page. Defaults to CUA_WEBSITE_URL
+                    env var or https://cua.ai
+        workspace_slug: Optional workspace slug to pre-select in the browser
+
+    Returns:
+        AuthResult with token and workspace metadata
+
+    Raises:
+        TimeoutError: If authentication times out
+        RuntimeError: If authentication fails
+    """
+    website_url = website_url or os.environ.get("CUA_WEBSITE_URL", DEFAULT_WEBSITE_URL)
+
+    # Reset handler state
+    CallbackHandler.token = None
+    CallbackHandler.workspace_slug = None
+    CallbackHandler.workspace_name = None
+    CallbackHandler.org_name = None
+    CallbackHandler.error = None
+
+    # Start local server on dynamic port
+    server = HTTPServer(("localhost", 0), CallbackHandler)
+    port = server.server_address[1]
+    callback_url = f"http://localhost:{port}"
+
+    # Build auth URL
+    encoded_callback = urllib.parse.quote(callback_url, safe="")
+    auth_url = f"{website_url}/cli-auth?callback_url={encoded_callback}"
+    if workspace_slug:
+        auth_url += f"&workspace_slug={urllib.parse.quote(workspace_slug, safe='')}"
+
+    # Start server in background thread
+    server_thread = Thread(target=server.handle_request, daemon=True)
+    server_thread.start()
+
+    # Open browser
+    print("Opening browser for authentication...")
+    if not open_browser(auth_url):
+        print("Could not open browser automatically.")
+        print(f"Please visit: {auth_url}")
+
+    # Wait for callback with timeout
+    try:
+        for _ in range(AUTH_TIMEOUT_SECONDS):
+            if CallbackHandler.token or CallbackHandler.error:
+                break
+            await asyncio.sleep(1)
+        else:
+            raise TimeoutError("Authentication timed out. Please try again.")
+
+        if CallbackHandler.error:
+            raise RuntimeError(f"Authentication failed: {CallbackHandler.error}")
+
+        if not CallbackHandler.token:
+            raise RuntimeError("Authentication failed: No token received")
+
+        return AuthResult(
+            token=CallbackHandler.token,
+            workspace_slug=CallbackHandler.workspace_slug or "",
+            workspace_name=CallbackHandler.workspace_name or "",
+            org_name=CallbackHandler.org_name or "",
+        )
+
+    finally:
+        server.server_close()
@@ -0,0 +1,293 @@
+"""SQLite-based credential storage for Cua CLI."""
+
+import os
+import sqlite3
+from pathlib import Path
+from typing import Optional
+
+# Default storage location
+CUA_DIR = Path.home() / ".cua"
+CREDENTIALS_DB = CUA_DIR / "credentials.db"
+
+# Key names
+API_KEY_NAME = "api_key"
+
+# Workspace key patterns
+ACTIVE_WORKSPACE_KEY = "active_workspace"
+
+
+class CredentialStore:
+    """SQLite-based credential store with WAL mode for concurrent access."""
+
+    def __init__(self, db_path: Path | None = None):
+        """Initialize the credential store.
+
+        Args:
+            db_path: Path to the SQLite database. Defaults to ~/.cua/credentials.db
+        """
+        self.db_path = db_path or CREDENTIALS_DB
+        self._ensure_db()
+
+    def _ensure_db(self) -> None:
+        """Ensure the database and table exist."""
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+
+        conn = sqlite3.connect(self.db_path)
+        try:
+            # Enable WAL mode for better concurrent access
+            conn.execute("PRAGMA journal_mode=WAL")
+
+            # Create key-value table
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS kv (
+                    key TEXT PRIMARY KEY,
+                    value TEXT NOT NULL
+                )
+            """)
+            conn.commit()
+        finally:
+            conn.close()
+
+    def get(self, key: str) -> Optional[str]:
+        """Get a value from the store.
+
+        Args:
+            key: The key to look up
+
+        Returns:
+            The value, or None if not found
+        """
+        conn = sqlite3.connect(self.db_path)
+        try:
+            cursor = conn.execute("SELECT value FROM kv WHERE key = ?", (key,))
+            row = cursor.fetchone()
+            return row[0] if row else None
+        finally:
+            conn.close()
+
+    def set(self, key: str, value: str) -> None:
+        """Set a value in the store.
+
+        Args:
+            key: The key to set
+            value: The value to store
+        """
+        conn = sqlite3.connect(self.db_path)
+        try:
+            conn.execute(
+                "INSERT OR REPLACE INTO kv (key, value) VALUES (?, ?)",
+                (key, value),
+            )
+            conn.commit()
+        finally:
+            conn.close()
+
+    def delete(self, key: str) -> bool:
+        """Delete a value from the store.
+
+        Args:
+            key: The key to delete
+
+        Returns:
+            True if the key was deleted, False if it didn't exist
+        """
+        conn = sqlite3.connect(self.db_path)
+        try:
+            cursor = conn.execute("DELETE FROM kv WHERE key = ?", (key,))
+            conn.commit()
+            return cursor.rowcount > 0
+        finally:
+            conn.close()
+
+    def delete_by_prefix(self, prefix: str) -> int:
+        """Delete all keys matching a prefix.
+
+        Args:
+            prefix: The key prefix to match (uses LIKE prefix%)
+
+        Returns:
+            Number of rows deleted
+        """
+        conn = sqlite3.connect(self.db_path)
+        try:
+            cursor = conn.execute("DELETE FROM kv WHERE key LIKE ?", (prefix + "%",))
+            conn.commit()
+            return cursor.rowcount
+        finally:
+            conn.close()
+
+    def list_by_prefix(self, prefix: str) -> list[tuple[str, str]]:
+        """List all key-value pairs matching a prefix.
+
+        Args:
+            prefix: The key prefix to match (uses LIKE prefix%)
+
+        Returns:
+            List of (key, value) tuples
+        """
+        conn = sqlite3.connect(self.db_path)
+        try:
+            cursor = conn.execute("SELECT key, value FROM kv WHERE key LIKE ?", (prefix + "%",))
+            return cursor.fetchall()
+        finally:
+            conn.close()
+
+    def clear(self) -> None:
+        """Clear all stored credentials."""
+        conn = sqlite3.connect(self.db_path)
+        try:
+            conn.execute("DELETE FROM kv")
+            conn.commit()
+        finally:
+            conn.close()
+
+
+# Module-level convenience functions
+_store: Optional[CredentialStore] = None
+
+
+def _get_store() -> CredentialStore:
+    """Get the global credential store instance."""
+    global _store
+    if _store is None:
+        _store = CredentialStore()
+    return _store
+
+
+def save_workspace(slug: str, api_key: str, name: str, org: str) -> None:
+    """Save workspace credentials and metadata."""
+    store = _get_store()
+    store.set(f"workspace:{slug}:api_key", api_key)
+    store.set(f"workspace:{slug}:name", name)
+    store.set(f"workspace:{slug}:org", org)
+
+
+def get_active_workspace() -> Optional[str]:
+    """Get the slug of the currently active workspace."""
+    return _get_store().get(ACTIVE_WORKSPACE_KEY)
+
+
+def set_active_workspace(slug: str) -> None:
+    """Set the active workspace by slug."""
+    _get_store().set(ACTIVE_WORKSPACE_KEY, slug)
+
+
+def list_workspaces() -> list[dict]:
+    """List all authenticated workspaces.
+
+    Returns:
+        List of dicts with keys: slug, name, org, is_active
+    """
+    store = _get_store()
+    active = store.get(ACTIVE_WORKSPACE_KEY)
+    rows = store.list_by_prefix("workspace:")
+
+    # Group by slug — extract from keys like workspace:<slug>:api_key
+    slugs: dict[str, dict] = {}
+    for key, value in rows:
+        parts = key.split(":")
+        if len(parts) != 3:
+            continue
+        slug = parts[1]
+        field = parts[2]
+        if slug not in slugs:
+            slugs[slug] = {"slug": slug, "name": "", "org": "", "is_active": False}
+        if field == "name":
+            slugs[slug]["name"] = value
+        elif field == "org":
+            slugs[slug]["org"] = value
+
+    # Only include workspaces that have an api_key stored
+    result = []
+    for key, _value in rows:
+        parts = key.split(":")
+        if len(parts) == 3 and parts[2] == "api_key" and parts[1] in slugs:
+            ws = slugs[parts[1]]
+            ws["is_active"] = ws["slug"] == active
+            result.append(ws)
+
+    result.sort(key=lambda ws: ws["slug"])
+    return result
+
+
+def get_workspace_api_key(slug: str) -> Optional[str]:
+    """Get the API key for a specific workspace."""
+    return _get_store().get(f"workspace:{slug}:api_key")
+
+
+def delete_workspace(slug: str) -> None:
+    """Delete all stored data for a workspace."""
+    store = _get_store()
+    store.delete(f"workspace:{slug}:api_key")
+    store.delete(f"workspace:{slug}:name")
+    store.delete(f"workspace:{slug}:org")
+
+
+def clear_all_workspaces() -> None:
+    """Delete all workspace keys and the active workspace marker."""
+    store = _get_store()
+    store.delete_by_prefix("workspace:")
+    store.delete(ACTIVE_WORKSPACE_KEY)
+
+
+def get_api_key() -> Optional[str]:
+    """Get the stored API key.
+
+    Priority:
+    1. CUA_API_KEY environment variable
+    2. Active workspace's API key
+    3. Legacy 'api_key' key (backward compat)
+
+    Returns:
+        The API key, or None if not found
+    """
+    # Environment variable takes precedence
+    env_key = os.environ.get("CUA_API_KEY")
+    if env_key:
+        return env_key
+
+    # Active workspace
+    store = _get_store()
+    active_slug = store.get(ACTIVE_WORKSPACE_KEY)
+    if active_slug:
+        ws_key = store.get(f"workspace:{active_slug}:api_key")
+        if ws_key:
+            return ws_key
+
+    # Legacy fallback
+    return store.get(API_KEY_NAME)
+
+
+def save_api_key(api_key: str) -> None:
+    """Save an API key to the credential store (legacy path)."""
+    _get_store().set(API_KEY_NAME, api_key)
+
+
+def clear_credentials() -> None:
+    """Clear all stored credentials including workspaces and legacy key."""
+    store = _get_store()
+    clear_all_workspaces()
+    store.delete(API_KEY_NAME)
+
+
+def clear_legacy_credentials() -> None:
+    """Clear only the legacy API key (not workspace data)."""
+    _get_store().delete(API_KEY_NAME)
+
+
+def require_api_key() -> str:
+    """Get the API key, raising an error if not found.
+
+    Returns:
+        The API key
+
+    Raises:
+        RuntimeError: If no API key is configured
+    """
+    api_key = get_api_key()
+    if not api_key:
+        raise RuntimeError(
+            "No API key configured. Run 'cua auth login' to authenticate, "
+            "or set the CUA_API_KEY environment variable."
+        )
+    return api_key