gdmcode 0.1.0__py3-none-any.whl

src/auth.py

@@ -0,0 +1,411 @@
+"""Credential management for gdm — keychain storage + interactive login.
+
+Providers:
+  - xAI Grok: API key (no OAuth, just a string)
+  - Google Gemini: API key OR OAuth device flow (RFC 8628)
+  - OpenAI Codex: API key
+
+Storage priority (highest → lowest):
+  1. Environment variables (XAI_API_KEY, GEMINI_API_KEY, OPENAI_API_KEY)
+  2. System keychain (via `keyring` — macOS Keychain, Windows Credential Store, etc.)
+  3. ~/.config/gdm/config.toml [api] section
+
+Usage::
+
+    from src.auth import CredentialStore
+    store = CredentialStore()
+    store.set("grok", "xai-abc123")
+    key = store.get("grok")  # "xai-abc123"
+"""
+from __future__ import annotations
+
+import logging
+import os
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from src.exceptions import ConfigError
+
+if TYPE_CHECKING:
+    pass
+
+__all__ = [
+    "CredentialStore",
+    "GdmCredentials",
+    "login_interactive",
+    "logout",
+]
+
+log = logging.getLogger(__name__)
+
+_KEYRING_SERVICE = "gdm-code"
+_KEYRING_KEYS: dict[str, str] = {
+    "grok": "xai_api_key",
+    "gemini": "gemini_api_key",
+    "gemini_refresh": "gemini_refresh_token",
+    "codex": "openai_api_key",
+    "proxy": "proxy_token",
+}
+
+# Gemini OAuth constants (bring-your-own GCP project recommended for now).
+# Users can set GDM_GOOGLE_CLIENT_ID / GDM_GOOGLE_CLIENT_SECRET env vars.
+_GEMINI_AUTH_URI = "https://accounts.google.com/o/oauth2/device/code"
+_GEMINI_TOKEN_URI = "https://oauth2.googleapis.com/token"
+_GEMINI_SCOPE = "https://www.googleapis.com/auth/generative-language"
+_GEMINI_DEVICE_POLL_INTERVAL = 5  # seconds between polls
+
+
+# ---------------------------------------------------------------------------
+# GdmCredentials
+# ---------------------------------------------------------------------------
+
+@dataclass
+class GdmCredentials:
+    """All resolved credentials for a session."""
+
+    xai_api_key: str | None = None
+    gemini_api_key: str | None = None
+    openai_api_key: str | None = None
+    proxy_token: str | None = None
+
+    @property
+    def has_grok(self) -> bool:
+        return bool(self.xai_api_key)
+
+    @property
+    def has_gemini(self) -> bool:
+        return bool(self.gemini_api_key)
+
+    @property
+    def has_codex(self) -> bool:
+        return bool(self.openai_api_key)
+
+    @property
+    def any_key(self) -> bool:
+        return self.has_grok or self.has_gemini or self.has_codex
+
+
+# ---------------------------------------------------------------------------
+# CredentialStore
+# ---------------------------------------------------------------------------
+
+class CredentialStore:
+    """Unified credential store with keychain backend.
+
+    Falls back gracefully if `keyring` is not installed — writes/reads from
+    ~/.config/gdm/.credentials (plain text, user-readable only).
+    """
+
+    def __init__(self) -> None:
+        self._keyring_available = _check_keyring()
+        self._fallback_path = Path.home() / ".config" / "gdm" / ".credentials"
+
+    # ── Public API ────────────────────────────────────────────────────────
+
+    def set(self, provider: str, value: str) -> None:
+        """Persist a credential for `provider`."""
+        key = _KEYRING_KEYS.get(provider, provider)
+        if self._keyring_available:
+            import keyring  # type: ignore[import]
+            keyring.set_password(_KEYRING_SERVICE, key, value)
+            log.debug("Stored %r in system keychain", key)
+        else:
+            self._set_fallback(key, value)
+            log.debug("Stored %r in fallback credentials file", key)
+
+    def get(self, provider: str) -> str | None:
+        """Retrieve a credential for `provider`, or None if not found."""
+        key = _KEYRING_KEYS.get(provider, provider)
+        if self._keyring_available:
+            import keyring  # type: ignore[import]
+            try:
+                return keyring.get_password(_KEYRING_SERVICE, key) or None
+            except Exception as exc:  # noqa: BLE001
+                log.warning("keyring.get_password failed: %s", exc)
+        return self._get_fallback(key)
+
+    def delete(self, provider: str) -> None:
+        """Remove a stored credential."""
+        key = _KEYRING_KEYS.get(provider, provider)
+        if self._keyring_available:
+            import keyring  # type: ignore[import]
+            try:
+                keyring.delete_password(_KEYRING_SERVICE, key)
+            except Exception:  # noqa: BLE001
+                pass
+        self._delete_fallback(key)
+
+    def load_all(self) -> GdmCredentials:
+        """Return all credentials, checking env vars first, then keychain."""
+        return GdmCredentials(
+            xai_api_key=(
+                os.environ.get("XAI_API_KEY")
+                or self.get("grok")
+                or None
+            ),
+            gemini_api_key=(
+                os.environ.get("GEMINI_API_KEY")
+                or self.get("gemini")
+                or None
+            ),
+            openai_api_key=(
+                os.environ.get("OPENAI_API_KEY")
+                or self.get("codex")
+                or None
+            ),
+            proxy_token=(
+                os.environ.get("GDM_PROXY_TOKEN")
+                or self.get("proxy")
+                or None
+            ),
+        )
+
+    # ── Fallback file (when keyring not available) ────────────────────────
+
+    def _set_fallback(self, key: str, value: str) -> None:
+        self._fallback_path.parent.mkdir(parents=True, exist_ok=True)
+        data = self._read_fallback_file()
+        data[key] = value
+        lines = "\n".join(f"{k}={v}" for k, v in data.items())
+        self._fallback_path.write_text(lines, encoding="utf-8")
+        # Restrict permissions to owner-only on POSIX.
+        try:
+            os.chmod(self._fallback_path, 0o600)
+        except OSError as exc:
+            log.warning(
+                "chmod failed for credential file %s: %s — file may be world-readable",
+                self._fallback_path,
+                exc,
+            )
+
+    def _get_fallback(self, key: str) -> str | None:
+        data = self._read_fallback_file()
+        return data.get(key)
+
+    def _delete_fallback(self, key: str) -> None:
+        data = self._read_fallback_file()
+        data.pop(key, None)
+        lines = "\n".join(f"{k}={v}" for k, v in data.items())
+        self._fallback_path.write_text(lines, encoding="utf-8")
+
+    def _read_fallback_file(self) -> dict[str, str]:
+        if not self._fallback_path.exists():
+            return {}
+        try:
+            result: dict[str, str] = {}
+            for line in self._fallback_path.read_text(encoding="utf-8").splitlines():
+                if "=" in line:
+                    k, _, v = line.partition("=")
+                    result[k.strip()] = v.strip()
+            return result
+        except OSError:
+            return {}
+
+
+# ---------------------------------------------------------------------------
+# Interactive login / logout
+# ---------------------------------------------------------------------------
+
+def login_interactive(provider: str, store: CredentialStore | None = None) -> None:
+    """Interactively prompt the user to log in for a provider.
+
+    Supports:
+      - "grok" → paste API key
+      - "gemini" → paste API key OR OAuth device flow
+      - "codex" → paste API key
+
+    Args:
+        provider: one of "grok", "gemini", "codex", "all"
+        store: CredentialStore instance (creates a new one if None)
+    """
+    from rich.console import Console
+    from rich.prompt import Prompt
+
+    console = Console()
+    s = store or CredentialStore()
+
+    providers_to_login = (
+        ["grok", "gemini", "codex"] if provider == "all" else [provider.lower()]
+    )
+
+    for prov in providers_to_login:
+        match prov:
+            case "grok":
+                _login_grok(s, console)
+            case "gemini":
+                _login_gemini(s, console)
+            case "codex":
+                _login_codex(s, console)
+            case _:
+                console.print(f"[red]Unknown provider: {prov!r}[/red]")
+                console.print("Valid providers: grok, gemini, codex, all")
+
+
+def logout(provider: str, store: CredentialStore | None = None) -> None:
+    """Remove stored credentials for a provider.
+
+    Args:
+        provider: one of "grok", "gemini", "codex", "all"
+        store: CredentialStore instance (creates a new one if None)
+    """
+    from rich.console import Console
+    console = Console()
+    s = store or CredentialStore()
+
+    if provider == "all":
+        for prov in ("grok", "gemini", "gemini_refresh", "codex"):
+            s.delete(prov)
+        console.print("[green]✓[/green] Logged out from all providers.")
+    else:
+        s.delete(provider)
+        if provider == "gemini":
+            s.delete("gemini_refresh")
+        console.print(f"[green]✓[/green] Logged out from [bold]{provider}[/bold].")
+
+
+# ---------------------------------------------------------------------------
+# Per-provider login helpers
+# ---------------------------------------------------------------------------
+
+def _login_grok(store: CredentialStore, console: object) -> None:
+    from rich.console import Console
+    from rich.prompt import Prompt
+    c: Console = console  # type: ignore[assignment]
+    c.print("\n[bold cyan]xAI Grok login[/bold cyan]")
+    c.print("Get your API key at: [link=https://console.x.ai]https://console.x.ai[/link]")
+    key = Prompt.ask("Paste your XAI_API_KEY", password=True, console=c)
+    if key.strip():
+        store.set("grok", key.strip())
+        c.print("[green]✓[/green] Grok API key saved to keychain.")
+    else:
+        c.print("[yellow]No key entered — skipped.[/yellow]")
+
+
+def _login_codex(store: CredentialStore, console: object) -> None:
+    from rich.console import Console
+    from rich.prompt import Prompt
+    c: Console = console  # type: ignore[assignment]
+    c.print("\n[bold cyan]OpenAI / Codex login[/bold cyan]")
+    c.print("Get your API key at: [link=https://platform.openai.com/api-keys]platform.openai.com/api-keys[/link]")
+    key = Prompt.ask("Paste your OPENAI_API_KEY", password=True, console=c)
+    if key.strip():
+        store.set("codex", key.strip())
+        c.print("[green]✓[/green] OpenAI API key saved to keychain.")
+    else:
+        c.print("[yellow]No key entered — skipped.[/yellow]")
+
+
+def _login_gemini(store: CredentialStore, console: object) -> None:
+    from rich.console import Console
+    from rich.prompt import Prompt
+    c: Console = console  # type: ignore[assignment]
+    c.print("\n[bold cyan]Google Gemini login[/bold cyan]")
+    c.print("1. API key (quick) — no OAuth needed")
+    c.print("2. OAuth device flow — uses Google account, no key required\n")
+    choice = Prompt.ask("Choose login method", choices=["1", "2"], default="1", console=c)
+
+    if choice == "1":
+        c.print("Get your API key at: [link=https://aistudio.google.com/app/apikey]aistudio.google.com/app/apikey[/link]")
+        key = Prompt.ask("Paste your GEMINI_API_KEY", password=True, console=c)
+        if key.strip():
+            store.set("gemini", key.strip())
+            c.print("[green]✓[/green] Gemini API key saved to keychain.")
+        else:
+            c.print("[yellow]No key entered — skipped.[/yellow]")
+    else:
+        _login_gemini_oauth(store, c)
+
+
+def _login_gemini_oauth(store: CredentialStore, console: object) -> None:
+    """Complete Gemini OAuth device authorization flow (RFC 8628)."""
+    from rich.console import Console
+    c: Console = console  # type: ignore[assignment]
+
+    client_id = os.environ.get("GDM_GOOGLE_CLIENT_ID")
+    client_secret = os.environ.get("GDM_GOOGLE_CLIENT_SECRET")
+
+    if not client_id or not client_secret:
+        c.print(
+            "[yellow]Warning:[/yellow] Gemini OAuth requires a GCP OAuth app.\n"
+            "Set GDM_GOOGLE_CLIENT_ID and GDM_GOOGLE_CLIENT_SECRET, or use method 1 (API key).\n"
+            "See: https://console.cloud.google.com/apis/credentials"
+        )
+        return
+
+    try:
+        import requests
+    except ImportError:
+        c.print("[red]`requests` not installed. Run: pip install requests[/red]")
+        return
+
+    # Step 1: Request device code.
+    resp = requests.post(
+        _GEMINI_AUTH_URI,
+        data={"client_id": client_id, "scope": _GEMINI_SCOPE},
+        timeout=20,
+    )
+    if resp.status_code != 200:
+        c.print(f"[red]Device code request failed: {resp.status_code} {resp.text}[/red]")
+        return
+
+    data = resp.json()
+    device_code: str = data["device_code"]
+    user_code: str = data["user_code"]
+    verification_url: str = data["verification_url"]
+    expires_in: int = data.get("expires_in", 1800)
+    interval: int = data.get("interval", _GEMINI_DEVICE_POLL_INTERVAL)
+
+    c.print(f"\n[bold]Visit:[/bold] {verification_url}")
+    c.print(f"[bold]Enter code:[/bold] {user_code}\n")
+    c.print(f"Waiting for authorization (expires in {expires_in}s)...")
+
+    # Step 2: Poll for token.
+    deadline = time.time() + expires_in
+    while time.time() < deadline:
+        time.sleep(interval)
+        token_resp = requests.post(
+            _GEMINI_TOKEN_URI,
+            data={
+                "client_id": client_id,
+                "client_secret": client_secret,
+                "device_code": device_code,
+                "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+            },
+            timeout=20,
+        )
+        token_data = token_resp.json()
+        error = token_data.get("error")
+        if error == "authorization_pending":
+            continue
+        if error == "slow_down":
+            interval += 5
+            continue
+        if "access_token" in token_data:
+            access_token: str = token_data["access_token"]
+            refresh_token: str | None = token_data.get("refresh_token")
+            store.set("gemini", access_token)
+            if refresh_token:
+                store.set("gemini_refresh", refresh_token)
+            c.print("[green]✓[/green] Gemini OAuth complete. Access token saved to keychain.")
+            return
+        # Unrecoverable error.
+        c.print(f"[red]OAuth error: {token_data.get('error_description', error)}[/red]")
+        return
+
+    c.print("[red]Authorization timed out. Please try again.[/red]")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _check_keyring() -> bool:
+    """Return True if keyring is available and functional."""
+    try:
+        import keyring  # noqa: F401
+        return True
+    except ImportError:
+        log.info("keyring not installed — using fallback credential file")
+        return False