workspace-mcp 0.2.0__py3-none-any.whl

auth/oauth_callback_server.py

@@ -0,0 +1,241 @@
+"""
+Transport-aware OAuth callback handling.
+
+In streamable-http mode: Uses the existing FastAPI server
+In stdio mode: Starts a minimal HTTP server just for OAuth callbacks
+"""
+
+import asyncio
+import logging
+import threading
+import time
+from typing import Optional, Dict, Any
+import socket
+
+from fastapi import FastAPI, Request
+import uvicorn
+
+from auth.google_auth import handle_auth_callback, CONFIG_CLIENT_SECRETS_PATH
+from auth.scopes import OAUTH_STATE_TO_SESSION_ID_MAP, SCOPES
+from auth.oauth_responses import create_error_response, create_success_response, create_server_error_response
+
+logger = logging.getLogger(__name__)
+
+class MinimalOAuthServer:
+    """
+    Minimal HTTP server for OAuth callbacks in stdio mode.
+    Only starts when needed and uses the same port (8000) as streamable-http mode.
+    """
+
+    def __init__(self, port: int = 8000, base_uri: str = "http://localhost"):
+        self.port = port
+        self.base_uri = base_uri
+        self.app = FastAPI()
+        self.server = None
+        self.server_thread = None
+        self.is_running = False
+
+        # Setup the callback route
+        self._setup_callback_route()
+
+    def _setup_callback_route(self):
+        """Setup the OAuth callback route."""
+
+        @self.app.get("/oauth2callback")
+        async def oauth_callback(request: Request):
+            """Handle OAuth callback - same logic as in core/server.py"""
+            state = request.query_params.get("state")
+            code = request.query_params.get("code")
+            error = request.query_params.get("error")
+
+            if error:
+                error_message = f"Authentication failed: Google returned an error: {error}. State: {state}."
+                logger.error(error_message)
+                return create_error_response(error_message)
+
+            if not code:
+                error_message = "Authentication failed: No authorization code received from Google."
+                logger.error(error_message)
+                return create_error_response(error_message)
+
+            try:
+                logger.info(f"OAuth callback: Received code (state: {state}). Attempting to exchange for tokens.")
+
+                mcp_session_id: Optional[str] = OAUTH_STATE_TO_SESSION_ID_MAP.pop(state, None)
+                if mcp_session_id:
+                    logger.info(f"OAuth callback: Retrieved MCP session ID '{mcp_session_id}' for state '{state}'.")
+                else:
+                    logger.warning(f"OAuth callback: No MCP session ID found for state '{state}'. Auth will not be tied to a specific session.")
+
+                # Exchange code for credentials
+                verified_user_id, credentials = handle_auth_callback(
+                    client_secrets_path=CONFIG_CLIENT_SECRETS_PATH,
+                    scopes=SCOPES,
+                    authorization_response=str(request.url),
+                    redirect_uri=f"{self.base_uri}:{self.port}/oauth2callback",
+                    session_id=mcp_session_id
+                )
+
+                log_session_part = f" (linked to session: {mcp_session_id})" if mcp_session_id else ""
+                logger.info(f"OAuth callback: Successfully authenticated user: {verified_user_id} (state: {state}){log_session_part}.")
+
+                # Return success page using shared template
+                return create_success_response(verified_user_id)
+
+            except Exception as e:
+                error_message_detail = f"Error processing OAuth callback (state: {state}): {str(e)}"
+                logger.error(error_message_detail, exc_info=True)
+                return create_server_error_response(str(e))
+
+    def start(self) -> bool:
+        """
+        Start the minimal OAuth server.
+
+        Returns:
+            True if started successfully, False otherwise
+        """
+        if self.is_running:
+            logger.info("Minimal OAuth server is already running")
+            return True
+
+        # Check if port is available
+        # Extract hostname from base_uri (e.g., "http://localhost" -> "localhost")
+        try:
+            from urllib.parse import urlparse
+            parsed_uri = urlparse(self.base_uri)
+            hostname = parsed_uri.hostname or 'localhost'
+        except Exception:
+            hostname = 'localhost'
+            
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.bind((hostname, self.port))
+        except OSError:
+            logger.error(f"Port {self.port} is already in use on {hostname}. Cannot start minimal OAuth server.")
+            return False
+
+        def run_server():
+            """Run the server in a separate thread."""
+            try:
+                config = uvicorn.Config(
+                    self.app,
+                    host=hostname,
+                    port=self.port,
+                    log_level="warning",
+                    access_log=False
+                )
+                self.server = uvicorn.Server(config)
+                asyncio.run(self.server.serve())
+
+            except Exception as e:
+                logger.error(f"Minimal OAuth server error: {e}", exc_info=True)
+
+        # Start server in background thread
+        self.server_thread = threading.Thread(target=run_server, daemon=True)
+        self.server_thread.start()
+
+        # Wait for server to start
+        max_wait = 3.0
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            try:
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                    result = s.connect_ex((hostname, self.port))
+                    if result == 0:
+                        self.is_running = True
+                        logger.info(f"Minimal OAuth server started on {hostname}:{self.port}")
+                        return True
+            except Exception:
+                pass
+            time.sleep(0.1)
+
+        logger.error(f"Failed to start minimal OAuth server on {hostname}:{self.port}")
+        return False
+
+    def stop(self):
+        """Stop the minimal OAuth server."""
+        if not self.is_running:
+            return
+
+        try:
+            if self.server:
+                if hasattr(self.server, 'should_exit'):
+                    self.server.should_exit = True
+
+            if self.server_thread and self.server_thread.is_alive():
+                self.server_thread.join(timeout=3.0)
+
+            self.is_running = False
+            logger.info(f"Minimal OAuth server stopped")
+
+        except Exception as e:
+            logger.error(f"Error stopping minimal OAuth server: {e}", exc_info=True)
+
+
+# Global instance for stdio mode
+_minimal_oauth_server: Optional[MinimalOAuthServer] = None
+
+def get_oauth_redirect_uri(transport_mode: str = "stdio", port: int = 8000, base_uri: str = "http://localhost") -> str:
+    """
+    Get the appropriate OAuth redirect URI based on transport mode.
+
+    Args:
+        transport_mode: "stdio" or "streamable-http"
+        port: Port number (default 8000)
+        base_uri: Base URI (default "http://localhost")
+
+    Returns:
+        OAuth redirect URI
+    """
+    return f"{base_uri}:{port}/oauth2callback"
+
+def ensure_oauth_callback_available(transport_mode: str = "stdio", port: int = 8000, base_uri: str = "http://localhost") -> bool:
+    """
+    Ensure OAuth callback endpoint is available for the given transport mode.
+
+    For streamable-http: Assumes the main server is already running
+    For stdio: Starts a minimal server if needed
+
+    Args:
+        transport_mode: "stdio" or "streamable-http"
+        port: Port number (default 8000)
+        base_uri: Base URI (default "http://localhost")
+
+    Returns:
+        True if callback endpoint is available, False otherwise
+    """
+    global _minimal_oauth_server
+
+    if transport_mode == "streamable-http":
+        # In streamable-http mode, the main FastAPI server should handle callbacks
+        logger.debug("Using existing FastAPI server for OAuth callbacks (streamable-http mode)")
+        return True
+
+    elif transport_mode == "stdio":
+        # In stdio mode, start minimal server if not already running
+        if _minimal_oauth_server is None:
+            logger.info(f"Creating minimal OAuth server instance for {base_uri}:{port}")
+            _minimal_oauth_server = MinimalOAuthServer(port, base_uri)
+
+        if not _minimal_oauth_server.is_running:
+            logger.info("Starting minimal OAuth server for stdio mode")
+            result = _minimal_oauth_server.start()
+            if result:
+                logger.info(f"Minimal OAuth server successfully started on {base_uri}:{port}")
+            else:
+                logger.error(f"Failed to start minimal OAuth server on {base_uri}:{port}")
+            return result
+        else:
+            logger.info("Minimal OAuth server is already running")
+            return True
+
+    else:
+        logger.error(f"Unknown transport mode: {transport_mode}")
+        return False
+
+def cleanup_oauth_callback_server():
+    """Clean up the minimal OAuth server if it was started."""
+    global _minimal_oauth_server
+    if _minimal_oauth_server:
+        _minimal_oauth_server.stop()
+        _minimal_oauth_server = None

auth/oauth_responses.py

@@ -0,0 +1,223 @@
+"""
+Shared OAuth callback response templates.
+
+Provides reusable HTML response templates for OAuth authentication flows
+to eliminate duplication between server.py and oauth_callback_server.py.
+"""
+
+from fastapi.responses import HTMLResponse
+from typing import Optional
+
+
+def create_error_response(error_message: str, status_code: int = 400) -> HTMLResponse:
+    """
+    Create a standardized error response for OAuth failures.
+
+    Args:
+        error_message: The error message to display
+        status_code: HTTP status code (default 400)
+
+    Returns:
+        HTMLResponse with error page
+    """
+    content = f"""
+        <html>
+        <head><title>Authentication Error</title></head>
+        <body style="font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; max-width: 600px; margin: 40px auto; padding: 20px; text-align: center;">
+            <h2 style="color: #d32f2f;">Authentication Error</h2>
+            <p>{error_message}</p>
+            <p>Please ensure you grant the requested permissions. You can close this window and try again.</p>
+            <script>setTimeout(function() {{ window.close(); }}, 10000);</script>
+        </body>
+        </html>
+    """
+    return HTMLResponse(content=content, status_code=status_code)
+
+
+def create_success_response(verified_user_id: Optional[str] = None) -> HTMLResponse:
+    """
+    Create a standardized success response for OAuth authentication.
+
+    Args:
+        verified_user_id: The authenticated user's email (optional)
+
+    Returns:
+        HTMLResponse with success page
+    """
+    # Handle the case where no user ID is provided
+    user_display = verified_user_id if verified_user_id else "Google User"
+
+    content = f"""<html>
+<head>
+    <title>Authentication Successful</title>
+    <style>
+        * {{
+            margin: 0;
+            padding: 0;
+            box-sizing: border-box;
+        }}
+
+        body {{
+            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
+            background: linear-gradient(135deg,#0f172a,#1e293b,#334155);
+            min-height: 100vh;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            color: #1a1a1a;
+            -webkit-font-smoothing: antialiased;
+            -moz-osx-font-smoothing: grayscale;
+        }}
+
+        .container {{
+            background: rgba(255, 255, 255, 0.95);
+            backdrop-filter: blur(10px);
+            padding: 60px;
+            border-radius: 20px;
+            box-shadow: 0 30px 60px rgba(0, 0, 0, 0.12);
+            text-align: center;
+            max-width: 480px;
+            width: 90%;
+            transform: translateY(-20px);
+            animation: slideUp 0.6s ease-out;
+        }}
+
+        @keyframes slideUp {{
+            from {{
+                opacity: 0;
+                transform: translateY(0);
+            }}
+            to {{
+                opacity: 1;
+                transform: translateY(-20px);
+            }}
+        }}
+
+        .icon {{
+            width: 80px;
+            height: 80px;
+            margin: 0 auto 30px;
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            border-radius: 50%;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            font-size: 40px;
+            color: white;
+            animation: pulse 2s ease-in-out infinite;
+        }}
+
+        @keyframes pulse {{
+            0%, 100% {{
+                transform: scale(1);
+            }}
+            50% {{
+                transform: scale(1.05);
+            }}
+        }}
+
+        h1 {{
+            font-size: 28px;
+            font-weight: 600;
+            margin-bottom: 20px;
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            -webkit-background-clip: text;
+            -webkit-text-fill-color: transparent;
+            background-clip: text;
+        }}
+
+        .message {{
+            font-size: 16px;
+            line-height: 1.6;
+            color: #4a5568;
+            margin-bottom: 20px;
+        }}
+
+        .user-id {{
+            font-weight: 600;
+            color: #667eea;
+            padding: 4px 12px;
+            background: rgba(102, 126, 234, 0.1);
+            border-radius: 6px;
+            display: inline-block;
+            margin: 0 4px;
+        }}
+
+        .button {{
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            padding: 16px 40px;
+            border: none;
+            border-radius: 30px;
+            font-size: 16px;
+            font-weight: 500;
+            cursor: pointer;
+            transition: all 0.3s ease;
+            margin-top: 30px;
+            display: inline-block;
+            text-decoration: none;
+            box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3);
+        }}
+
+        .button:hover {{
+            transform: translateY(-2px);
+            box-shadow: 0 7px 20px rgba(102, 126, 234, 0.4);
+        }}
+
+        .button:active {{
+            transform: translateY(0);
+        }}
+
+        .auto-close {{
+            font-size: 13px;
+            color: #a0aec0;
+            margin-top: 30px;
+            opacity: 0.8;
+        }}
+    </style>
+    <script>
+        setTimeout(function() {{
+            window.close();
+        }}, 10000);
+    </script>
+</head>
+<body>
+    <div class="container">
+        <div class="icon">✓</div>
+        <h1>Authentication Successful</h1>
+        <div class="message">
+            You've been authenticated as <span class="user-id">{user_display}</span>
+        </div>
+        <div class="message">
+            Your credentials have been securely saved. You can now close this window and retry your original command.
+        </div>
+        <button class="button" onclick="window.close()">Close Window</button>
+        <div class="auto-close">This window will close automatically in 10 seconds</div>
+    </div>
+</body>
+</html>"""
+    return HTMLResponse(content=content)
+
+
+def create_server_error_response(error_detail: str) -> HTMLResponse:
+    """
+    Create a standardized server error response for OAuth processing failures.
+
+    Args:
+        error_detail: The detailed error message
+
+    Returns:
+        HTMLResponse with server error page
+    """
+    content = f"""
+        <html>
+        <head><title>Authentication Processing Error</title></head>
+        <body style="font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; max-width: 600px; margin: 40px auto; padding: 20px; text-align: center;">
+            <h2 style="color: #d32f2f;">Authentication Processing Error</h2>
+            <p>An unexpected error occurred while processing your authentication: {error_detail}</p>
+            <p>Please try again. You can close this window.</p>
+            <script>setTimeout(function() {{ window.close(); }}, 10000);</script>
+        </body>
+        </html>
+    """
+    return HTMLResponse(content=content, status_code=500)

auth/scopes.py

@@ -0,0 +1,108 @@
+"""
+Google Workspace OAuth Scopes
+
+This module centralizes OAuth scope definitions for Google Workspace integration.
+Separated from service_decorator.py to avoid circular imports.
+"""
+import logging
+from typing import Dict
+
+logger = logging.getLogger(__name__)
+
+# Temporary map to associate OAuth state with MCP session ID
+# This should ideally be a more robust cache in a production system (e.g., Redis)
+OAUTH_STATE_TO_SESSION_ID_MAP: Dict[str, str] = {}
+
+# Individual OAuth Scope Constants
+USERINFO_EMAIL_SCOPE = 'https://www.googleapis.com/auth/userinfo.email'
+OPENID_SCOPE = 'openid'
+CALENDAR_READONLY_SCOPE = 'https://www.googleapis.com/auth/calendar.readonly'
+CALENDAR_EVENTS_SCOPE = 'https://www.googleapis.com/auth/calendar.events'
+
+# Google Drive scopes
+DRIVE_READONLY_SCOPE = 'https://www.googleapis.com/auth/drive.readonly'
+DRIVE_FILE_SCOPE = 'https://www.googleapis.com/auth/drive.file'
+
+# Google Docs scopes
+DOCS_READONLY_SCOPE = 'https://www.googleapis.com/auth/documents.readonly'
+DOCS_WRITE_SCOPE = 'https://www.googleapis.com/auth/documents'
+
+# Gmail API scopes
+GMAIL_READONLY_SCOPE = 'https://www.googleapis.com/auth/gmail.readonly'
+GMAIL_SEND_SCOPE = 'https://www.googleapis.com/auth/gmail.send'
+GMAIL_COMPOSE_SCOPE = 'https://www.googleapis.com/auth/gmail.compose'
+GMAIL_MODIFY_SCOPE = 'https://www.googleapis.com/auth/gmail.modify'
+GMAIL_LABELS_SCOPE = 'https://www.googleapis.com/auth/gmail.labels'
+
+# Google Chat API scopes
+CHAT_READONLY_SCOPE = 'https://www.googleapis.com/auth/chat.messages.readonly'
+CHAT_WRITE_SCOPE = 'https://www.googleapis.com/auth/chat.messages'
+CHAT_SPACES_SCOPE = 'https://www.googleapis.com/auth/chat.spaces'
+
+# Google Sheets API scopes
+SHEETS_READONLY_SCOPE = 'https://www.googleapis.com/auth/spreadsheets.readonly'
+SHEETS_WRITE_SCOPE = 'https://www.googleapis.com/auth/spreadsheets'
+
+# Google Forms API scopes
+FORMS_BODY_SCOPE = 'https://www.googleapis.com/auth/forms.body'
+FORMS_BODY_READONLY_SCOPE = 'https://www.googleapis.com/auth/forms.body.readonly'
+FORMS_RESPONSES_READONLY_SCOPE = 'https://www.googleapis.com/auth/forms.responses.readonly'
+
+# Google Slides API scopes
+SLIDES_SCOPE = 'https://www.googleapis.com/auth/presentations'
+SLIDES_READONLY_SCOPE = 'https://www.googleapis.com/auth/presentations.readonly'
+
+# Base OAuth scopes required for user identification
+BASE_SCOPES = [
+    USERINFO_EMAIL_SCOPE,
+    OPENID_SCOPE
+]
+
+# Service-specific scope groups
+DOCS_SCOPES = [
+    DOCS_READONLY_SCOPE,
+    DOCS_WRITE_SCOPE
+]
+
+CALENDAR_SCOPES = [
+    CALENDAR_READONLY_SCOPE,
+    CALENDAR_EVENTS_SCOPE
+]
+
+DRIVE_SCOPES = [
+    DRIVE_READONLY_SCOPE,
+    DRIVE_FILE_SCOPE
+]
+
+GMAIL_SCOPES = [
+    GMAIL_READONLY_SCOPE,
+    GMAIL_SEND_SCOPE,
+    GMAIL_COMPOSE_SCOPE,
+    GMAIL_MODIFY_SCOPE,
+    GMAIL_LABELS_SCOPE
+]
+
+CHAT_SCOPES = [
+    CHAT_READONLY_SCOPE,
+    CHAT_WRITE_SCOPE,
+    CHAT_SPACES_SCOPE
+]
+
+SHEETS_SCOPES = [
+    SHEETS_READONLY_SCOPE,
+    SHEETS_WRITE_SCOPE
+]
+
+FORMS_SCOPES = [
+    FORMS_BODY_SCOPE,
+    FORMS_BODY_READONLY_SCOPE,
+    FORMS_RESPONSES_READONLY_SCOPE
+]
+
+SLIDES_SCOPES = [
+    SLIDES_SCOPE,
+    SLIDES_READONLY_SCOPE
+]
+
+# Combined scopes for all supported Google Workspace operations
+SCOPES = list(set(BASE_SCOPES + CALENDAR_SCOPES + DRIVE_SCOPES + GMAIL_SCOPES + DOCS_SCOPES + CHAT_SCOPES + SHEETS_SCOPES + FORMS_SCOPES + SLIDES_SCOPES))