stravinsky 0.1.2__py3-none-any.whl

mcp_bridge/__init__.py

@@ -0,0 +1,5 @@
+# Stravinsky MCP Bridge
+# Provides MCP tools for OAuth-authenticated access to Gemini and OpenAI models
+
+__version__ = "0.1.0"
+__author__ = "David Andrews"

mcp_bridge/auth/__init__.py

@@ -0,0 +1,32 @@
+# Authentication module
+from .token_store import TokenStore, TokenData
+from .oauth import (
+    perform_oauth_flow as gemini_oauth_flow,
+    refresh_access_token as gemini_refresh_token,
+    ANTIGRAVITY_CLIENT_ID,
+    ANTIGRAVITY_SCOPES,
+    ANTIGRAVITY_HEADERS,
+)
+from .openai_oauth import (
+    perform_oauth_flow as openai_oauth_flow,
+    refresh_access_token as openai_refresh_token,
+    CLIENT_ID as OPENAI_CLIENT_ID,
+    OPENAI_CALLBACK_PORT,
+)
+
+__all__ = [
+    # Token Store
+    "TokenStore",
+    "TokenData",
+    # Gemini OAuth
+    "gemini_oauth_flow",
+    "gemini_refresh_token",
+    "ANTIGRAVITY_CLIENT_ID",
+    "ANTIGRAVITY_SCOPES",
+    "ANTIGRAVITY_HEADERS",
+    # OpenAI OAuth
+    "openai_oauth_flow",
+    "openai_refresh_token",
+    "OPENAI_CLIENT_ID",
+    "OPENAI_CALLBACK_PORT",
+]

mcp_bridge/auth/cli.py

@@ -0,0 +1,208 @@
+"""
+Authentication CLI for Claude Superagent MCP Bridge
+
+Handles OAuth authentication for Gemini and OpenAI.
+"""
+
+import argparse
+import sys
+import time
+
+from .token_store import TokenStore
+from .oauth import perform_oauth_flow as gemini_oauth, refresh_access_token as gemini_refresh
+from .openai_oauth import (
+    perform_oauth_flow as openai_oauth,
+    refresh_access_token as openai_refresh,
+)
+
+
+def cmd_login(provider: str, token_store: TokenStore) -> int:
+    """
+    Perform OAuth login for a provider.
+    
+    For Gemini: Uses Google OAuth with Antigravity credentials
+    For OpenAI: Uses OpenAI OAuth (ChatGPT Plus/Pro subscription)
+    """
+    if provider == "gemini":
+        print(f"Starting Google OAuth for Gemini...")
+        
+        try:
+            result = gemini_oauth()
+            
+            expires_at = int(time.time()) + result.tokens.expires_in
+            
+            token_store.set_token(
+                provider="gemini",
+                access_token=result.tokens.access_token,
+                refresh_token=result.tokens.refresh_token,
+                expires_at=expires_at,
+            )
+            
+            print(f"\n✓ Successfully authenticated as: {result.user_info.email}")
+            print(f"  Token expires in: {result.tokens.expires_in // 60} minutes")
+            return 0
+            
+        except Exception as e:
+            print(f"\n✗ OAuth failed: {e}", file=sys.stderr)
+            return 1
+            
+    elif provider == "openai":
+        print(f"Starting OpenAI OAuth for ChatGPT Plus/Pro...")
+        print("Note: Requires ChatGPT Plus/Pro subscription and port 1455 available")
+        
+        try:
+            result = openai_oauth()
+            
+            expires_at = int(time.time()) + result.expires_in
+            
+            token_store.set_token(
+                provider="openai",
+                access_token=result.access_token,
+                refresh_token=result.refresh_token,
+                expires_at=expires_at,
+            )
+            
+            print(f"\n✓ Successfully authenticated with OpenAI")
+            print(f"  Token expires in: {result.expires_in // 60} minutes")
+            return 0
+            
+        except Exception as e:
+            print(f"\n✗ OAuth failed: {e}", file=sys.stderr)
+            print("\nTroubleshooting:")
+            print("  - Ensure you have a ChatGPT Plus/Pro subscription")
+            print("  - Stop Codex CLI if running: killall codex")
+            print("  - Or use Codex CLI directly: codex login")
+            return 1
+        
+    else:
+        print(f"Unknown provider: {provider}", file=sys.stderr)
+        print("Supported providers: gemini, openai")
+        return 1
+
+
+def cmd_logout(provider: str, token_store: TokenStore) -> int:
+    """Remove stored tokens for a provider."""
+    token_store.delete_token(provider)
+    print(f"✓ Logged out from {provider}")
+    return 0
+
+
+def cmd_status(token_store: TokenStore) -> int:
+    """Show authentication status for all providers."""
+    print("Authentication Status:\n")
+    
+    for provider in ["gemini", "openai"]:
+        has_token = token_store.has_valid_token(provider)
+        status = "✓ Authenticated" if has_token else "✗ Not authenticated"
+        print(f"  {provider.capitalize()}: {status}")
+        
+        if has_token:
+            token = token_store.get_token(provider)
+            if token and token.get("expires_at"):
+                expires = token["expires_at"]
+                remaining = expires - int(time.time())
+                if remaining > 0:
+                    hours = remaining // 3600
+                    minutes = (remaining % 3600) // 60
+                    print(f"    Expires in: {hours}h {minutes}m")
+                else:
+                    print(f"    Token expired")
+    
+    print()
+    return 0
+
+
+def cmd_refresh(provider: str, token_store: TokenStore) -> int:
+    """Refresh access token for a provider."""
+    token = token_store.get_token(provider)
+    if not token or not token.get("refresh_token"):
+        print(f"No refresh token found for {provider}")
+        print(f"Please run: python -m mcp_bridge.auth.cli login {provider}")
+        return 1
+    
+    try:
+        print(f"Refreshing {provider} token...")
+        
+        if provider == "gemini":
+            result = gemini_refresh(token["refresh_token"])
+        elif provider == "openai":
+            result = openai_refresh(token["refresh_token"])
+        else:
+            print(f"Refresh not supported for {provider}")
+            return 1
+        
+        expires_at = int(time.time()) + result.expires_in
+        
+        token_store.set_token(
+            provider=provider,
+            access_token=result.access_token,
+            refresh_token=result.refresh_token or token["refresh_token"],
+            expires_at=expires_at,
+        )
+        
+        print(f"✓ Token refreshed, expires in {result.expires_in // 60} minutes")
+        return 0
+        
+    except Exception as e:
+        print(f"✗ Refresh failed: {e}", file=sys.stderr)
+        return 1
+
+
+def main():
+    """CLI entry point."""
+    parser = argparse.ArgumentParser(
+        description="Stravinsky Authentication CLI",
+        prog="python -m mcp_bridge.auth.cli",
+    )
+    
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+    
+    # login command
+    login_parser = subparsers.add_parser("login", help="Authenticate with a provider")
+    login_parser.add_argument(
+        "provider",
+        choices=["gemini", "openai"],
+        help="Provider to authenticate with",
+    )
+    
+    # logout command
+    logout_parser = subparsers.add_parser("logout", help="Remove stored credentials")
+    logout_parser.add_argument(
+        "provider",
+        choices=["gemini", "openai"],
+        help="Provider to logout from",
+    )
+    
+    # status command
+    subparsers.add_parser("status", help="Show authentication status")
+    
+    # refresh command
+    refresh_parser = subparsers.add_parser("refresh", help="Refresh access token")
+    refresh_parser.add_argument(
+        "provider",
+        choices=["gemini", "openai"],
+        help="Provider to refresh token for",
+    )
+    
+    args = parser.parse_args()
+    
+    if not args.command:
+        parser.print_help()
+        return 1
+    
+    token_store = TokenStore()
+    
+    if args.command == "login":
+        return cmd_login(args.provider, token_store)
+    elif args.command == "logout":
+        return cmd_logout(args.provider, token_store)
+    elif args.command == "status":
+        return cmd_status(token_store)
+    elif args.command == "refresh":
+        return cmd_refresh(args.provider, token_store)
+    
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

mcp_bridge/auth/oauth.py

@@ -0,0 +1,418 @@
+"""
+Antigravity OAuth 2.0 Implementation with PKCE
+
+Implements Google OAuth for Antigravity (Gemini) authentication.
+Based on the original TypeScript implementation from Stravinsky.
+"""
+
+import base64
+import hashlib
+import json
+import os
+import secrets
+import threading
+import webbrowser
+from dataclasses import dataclass
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from typing import Any
+from urllib.parse import parse_qs, urlencode, urlparse
+
+import httpx
+
+
+# OAuth 2.0 Client Credentials (from constants.ts)
+ANTIGRAVITY_CLIENT_ID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
+ANTIGRAVITY_CLIENT_SECRET = "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"
+
+# OAuth Callback
+ANTIGRAVITY_CALLBACK_PORT = 51121
+
+# OAuth Scopes
+ANTIGRAVITY_SCOPES = [
+    "https://www.googleapis.com/auth/cloud-platform",
+    "https://www.googleapis.com/auth/userinfo.email",
+    "https://www.googleapis.com/auth/userinfo.profile",
+    "https://www.googleapis.com/auth/cclog",
+    "https://www.googleapis.com/auth/experimentsandconfigs",
+]
+
+# API Endpoints
+ANTIGRAVITY_ENDPOINTS = [
+    "https://daily-cloudcode-pa.sandbox.googleapis.com",  # dev
+    "https://autopush-cloudcode-pa.sandbox.googleapis.com",  # staging
+    "https://cloudcode-pa.googleapis.com",  # prod
+]
+
+# Default Project ID
+ANTIGRAVITY_DEFAULT_PROJECT_ID = "rising-fact-p41fc"
+
+# API Version
+ANTIGRAVITY_API_VERSION = "v1internal"
+
+# Request Headers
+ANTIGRAVITY_HEADERS = {
+    "User-Agent": "google-api-nodejs-client/9.15.1",
+    "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
+    "Client-Metadata": json.dumps({
+        "ideType": "IDE_UNSPECIFIED",
+        "platform": "PLATFORM_UNSPECIFIED",
+        "pluginType": "GEMINI",
+    }),
+}
+
+# Google OAuth endpoints
+GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
+GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
+GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v1/userinfo"
+
+# Token refresh buffer (60 seconds before expiry)
+TOKEN_REFRESH_BUFFER_MS = 60_000
+
+
+@dataclass
+class PKCEPair:
+    """PKCE verifier and challenge pair."""
+    verifier: str
+    challenge: str
+    method: str = "S256"
+
+
+@dataclass
+class TokenResult:
+    """OAuth token exchange result."""
+    access_token: str
+    refresh_token: str
+    expires_in: int
+    token_type: str
+
+
+@dataclass
+class UserInfo:
+    """User info from Google."""
+    email: str
+    name: str | None = None
+    picture: str | None = None
+
+
+@dataclass
+class OAuthResult:
+    """Complete OAuth flow result."""
+    tokens: TokenResult
+    user_info: UserInfo
+    verifier: str
+
+
+def generate_pkce_pair() -> PKCEPair:
+    """
+    Generate PKCE verifier and challenge pair.
+    
+    Uses SHA-256 for the challenge (S256 method).
+    """
+    # Generate 32-byte random verifier
+    verifier_bytes = secrets.token_bytes(32)
+    verifier = base64.urlsafe_b64encode(verifier_bytes).rstrip(b"=").decode("ascii")
+    
+    # Generate challenge from verifier using SHA-256
+    challenge_bytes = hashlib.sha256(verifier.encode("ascii")).digest()
+    challenge = base64.urlsafe_b64encode(challenge_bytes).rstrip(b"=").decode("ascii")
+    
+    return PKCEPair(verifier=verifier, challenge=challenge, method="S256")
+
+
+def encode_state(verifier: str, project_id: str | None = None) -> str:
+    """Encode OAuth state as URL-safe base64 JSON."""
+    state = {"verifier": verifier}
+    if project_id:
+        state["projectId"] = project_id
+    return base64.urlsafe_b64encode(json.dumps(state).encode()).decode()
+
+
+def decode_state(encoded: str) -> dict[str, Any]:
+    """Decode OAuth state from base64 JSON."""
+    # Handle both base64url and standard base64
+    encoded = encoded.replace("-", "+").replace("_", "/")
+    # Add padding
+    padding = 4 - len(encoded) % 4
+    if padding != 4:
+        encoded += "=" * padding
+    
+    try:
+        decoded = base64.b64decode(encoded).decode("utf-8")
+        return json.loads(decoded)
+    except Exception:
+        return {}
+
+
+def build_auth_url(
+    port: int,
+    project_id: str | None = None,
+    client_id: str = ANTIGRAVITY_CLIENT_ID,
+) -> tuple[str, str]:
+    """
+    Build Google OAuth authorization URL with PKCE.
+    
+    Returns:
+        Tuple of (auth_url, pkce_verifier)
+    """
+    pkce = generate_pkce_pair()
+    
+    redirect_uri = f"http://localhost:{port}/oauth-callback"
+    state = encode_state(pkce.verifier, project_id)
+    
+    params = {
+        "client_id": client_id,
+        "redirect_uri": redirect_uri,
+        "response_type": "code",
+        "scope": " ".join(ANTIGRAVITY_SCOPES),
+        "state": state,
+        "code_challenge": pkce.challenge,
+        "code_challenge_method": "S256",
+        "access_type": "offline",
+        "prompt": "consent",
+    }
+    
+    url = f"{GOOGLE_AUTH_URL}?{urlencode(params)}"
+    return url, pkce.verifier
+
+
+def exchange_code(
+    code: str,
+    verifier: str,
+    port: int,
+    client_id: str = ANTIGRAVITY_CLIENT_ID,
+    client_secret: str = ANTIGRAVITY_CLIENT_SECRET,
+) -> TokenResult:
+    """
+    Exchange authorization code for tokens.
+    
+    Args:
+        code: Authorization code from OAuth callback
+        verifier: PKCE verifier from initial auth request
+        port: Callback server port
+        
+    Returns:
+        Token exchange result with access and refresh tokens.
+    """
+    redirect_uri = f"http://localhost:{port}/oauth-callback"
+    
+    data = {
+        "client_id": client_id,
+        "client_secret": client_secret,
+        "code": code,
+        "grant_type": "authorization_code",
+        "redirect_uri": redirect_uri,
+        "code_verifier": verifier,
+    }
+    
+    with httpx.Client() as client:
+        response = client.post(
+            GOOGLE_TOKEN_URL,
+            data=data,
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+        )
+        
+        if response.status_code != 200:
+            raise Exception(f"Token exchange failed: {response.status_code} - {response.text}")
+        
+        result = response.json()
+        
+        return TokenResult(
+            access_token=result["access_token"],
+            refresh_token=result.get("refresh_token", ""),
+            expires_in=result.get("expires_in", 3600),
+            token_type=result.get("token_type", "Bearer"),
+        )
+
+
+def refresh_access_token(
+    refresh_token: str,
+    client_id: str = ANTIGRAVITY_CLIENT_ID,
+    client_secret: str = ANTIGRAVITY_CLIENT_SECRET,
+) -> TokenResult:
+    """
+    Refresh an access token using a refresh token.
+    
+    Args:
+        refresh_token: Valid refresh token
+        
+    Returns:
+        New token result.
+    """
+    data = {
+        "client_id": client_id,
+        "client_secret": client_secret,
+        "refresh_token": refresh_token,
+        "grant_type": "refresh_token",
+    }
+    
+    with httpx.Client() as client:
+        response = client.post(
+            GOOGLE_TOKEN_URL,
+            data=data,
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+        )
+        
+        if response.status_code != 200:
+            raise Exception(f"Token refresh failed: {response.status_code} - {response.text}")
+        
+        result = response.json()
+        
+        return TokenResult(
+            access_token=result["access_token"],
+            refresh_token=refresh_token,  # Keep original refresh token
+            expires_in=result.get("expires_in", 3600),
+            token_type=result.get("token_type", "Bearer"),
+        )
+
+
+def fetch_user_info(access_token: str) -> UserInfo:
+    """
+    Fetch user info from Google's userinfo API.
+    
+    Args:
+        access_token: Valid access token
+        
+    Returns:
+        User info with email, name, and picture.
+    """
+    with httpx.Client() as client:
+        response = client.get(
+            f"{GOOGLE_USERINFO_URL}?alt=json",
+            headers={"Authorization": f"Bearer {access_token}"},
+        )
+        
+        if response.status_code != 200:
+            raise Exception(f"Failed to fetch user info: {response.status_code}")
+        
+        data = response.json()
+        
+        return UserInfo(
+            email=data.get("email", ""),
+            name=data.get("name"),
+            picture=data.get("picture"),
+        )
+
+
+class OAuthCallbackHandler(BaseHTTPRequestHandler):
+    """HTTP request handler for OAuth callback."""
+    
+    callback_result: dict[str, Any] = {}
+    server_ready = threading.Event()
+    
+    def log_message(self, format, *args):
+        """Suppress HTTP server logs."""
+        pass
+    
+    def do_GET(self):
+        """Handle OAuth callback GET request."""
+        parsed = urlparse(self.path)
+        
+        if parsed.path == "/oauth-callback":
+            params = parse_qs(parsed.query)
+            
+            OAuthCallbackHandler.callback_result = {
+                "code": params.get("code", [""])[0],
+                "state": params.get("state", [""])[0],
+                "error": params.get("error", [None])[0],
+            }
+            
+            if OAuthCallbackHandler.callback_result["code"] and not OAuthCallbackHandler.callback_result["error"]:
+                body = b"<html><body><h1>Login successful!</h1><p>You can close this window.</p></body></html>"
+            else:
+                body = b"<html><body><h1>Login failed</h1><p>Please check the CLI output.</p></body></html>"
+            
+            self.send_response(200)
+            self.send_header("Content-Type", "text/html")
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+            
+            # Signal completion
+            OAuthCallbackHandler.server_ready.set()
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+
+def perform_oauth_flow(
+    project_id: str | None = None,
+    timeout: int = 300,
+) -> OAuthResult:
+    """
+    Perform full OAuth flow with browser-based authentication.
+    
+    1. Start local callback server
+    2. Open browser for Google auth
+    3. Wait for callback
+    4. Exchange code for tokens
+    5. Fetch user info
+    
+    Args:
+        project_id: Optional project ID for state
+        timeout: Timeout in seconds (default 5 minutes)
+        
+    Returns:
+        Complete OAuth result with tokens and user info.
+    """
+    # Reset callback state
+    OAuthCallbackHandler.callback_result = {}
+    OAuthCallbackHandler.server_ready.clear()
+    
+    # Start callback server
+    server = HTTPServer(("localhost", 0), OAuthCallbackHandler)
+    port = server.server_address[1]
+    
+    server_thread = threading.Thread(target=lambda: server.handle_request())
+    server_thread.daemon = True
+    server_thread.start()
+    
+    try:
+        # Build auth URL and open browser
+        auth_url, verifier = build_auth_url(port, project_id)
+        
+        print(f"\nOpening browser for Google authentication...")
+        print(f"If browser doesn't open, visit:\n{auth_url}\n")
+        
+        webbrowser.open(auth_url)
+        
+        # Wait for callback
+        if not OAuthCallbackHandler.server_ready.wait(timeout):
+            raise Exception("OAuth callback timeout")
+        
+        result = OAuthCallbackHandler.callback_result
+        
+        if result.get("error"):
+            raise Exception(f"OAuth error: {result['error']}")
+        
+        if not result.get("code"):
+            raise Exception("No authorization code received")
+        
+        # Verify state
+        state = decode_state(result["state"])
+        if state.get("verifier") != verifier:
+            raise Exception("PKCE verifier mismatch - possible security issue")
+        
+        # Exchange code for tokens
+        tokens = exchange_code(result["code"], verifier, port)
+        
+        # Fetch user info
+        user_info = fetch_user_info(tokens.access_token)
+        
+        print(f"✓ Authenticated as: {user_info.email}")
+        
+        return OAuthResult(tokens=tokens, user_info=user_info, verifier=verifier)
+        
+    finally:
+        server.server_close()
+
+
+if __name__ == "__main__":
+    # Test the OAuth flow
+    try:
+        result = perform_oauth_flow()
+        print(f"\nAccess Token: {result.tokens.access_token[:20]}...")
+        print(f"Refresh Token: {result.tokens.refresh_token[:20]}...")
+        print(f"Expires In: {result.tokens.expires_in}s")
+        print(f"User: {result.user_info.email}")
+    except Exception as e:
+        print(f"OAuth failed: {e}")