trycua--cua/libs/python/cua-cli/cua_cli/auth/store.py

"""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