oauth_callback_server.py

"""
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
import socket
import urllib.request
import uvicorn

from fastapi import FastAPI, Request
from fastapi.responses import FileResponse, JSONResponse
from typing import Optional
from urllib.parse import urlparse

from auth.scopes import SCOPES, get_current_scopes  # noqa
from auth.oauth_responses import (
    create_error_response,
    create_success_response,
    create_server_error_response,
)
from auth.google_auth import handle_auth_callback, check_client_secrets
from auth.oauth_config import get_oauth_redirect_uri

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

        # CLI auth completion signaling
        self.auth_completed = threading.Event()
        self.auth_result: Optional[dict] = None  # {"success": bool, "user_id": str|None, "error": str|None}
        self._auth_lock = threading.Lock()

        # Setup the callback route
        self._setup_callback_route()
        # Setup attachment serving route
        self._setup_attachment_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"""
            code = request.query_params.get("code")
            error = request.query_params.get("error")

            if error:
                error_message = (
                    f"Authentication failed: Google returned an error: {error}."
                )
                logger.error(error_message)
                self.auth_result = {"success": False, "user_id": None, "error": error_message}
                self.auth_completed.set()
                return create_error_response(error_message)

            if not code:
                error_message = (
                    "Authentication failed: No authorization code received from Google."
                )
                logger.error(error_message)
                self.auth_result = {"success": False, "user_id": None, "error": error_message}
                self.auth_completed.set()
                return create_error_response(error_message)

            try:
                # Check if we have credentials available (environment variables or file)
                error_message = check_client_secrets()
                if error_message:
                    self.auth_result = {"success": False, "user_id": None, "error": error_message}
                    self.auth_completed.set()
                    return create_server_error_response(error_message)

                logger.info(
                    "OAuth callback: Received authorization code. Attempting to exchange for tokens."
                )

                # Session ID tracking removed - not needed

                # Exchange code for credentials
                redirect_uri = get_oauth_redirect_uri()
                verified_user_id, credentials = handle_auth_callback(
                    scopes=get_current_scopes(),
                    authorization_response=str(request.url),
                    redirect_uri=redirect_uri,
                    session_id=None,
                )

                logger.info(
                    f"OAuth callback: Successfully authenticated user: {verified_user_id}."
                )

                # Signal completion for CLI auth flow
                self.auth_result = {"success": True, "user_id": verified_user_id, "error": None}
                self.auth_completed.set()

                # Return success page using shared template
                return create_success_response(verified_user_id)

            except Exception as e:
                logger.error(f"Error processing OAuth callback: {e}", exc_info=True)
                generic_error = "An unexpected error occurred while processing authentication. Please try again."
                self.auth_result = {"success": False, "user_id": None, "error": generic_error}
                self.auth_completed.set()
                return create_server_error_response(generic_error)

    def _setup_attachment_route(self):
        """Setup the attachment serving route."""
        from core.attachment_storage import get_attachment_storage

        @self.app.get("/attachments/{file_id}")
        async def serve_attachment(file_id: str, request: Request):
            """Serve a stored attachment file."""
            storage = get_attachment_storage()
            metadata = storage.get_attachment_metadata(file_id)

            if not metadata:
                return JSONResponse(
                    {"error": "Attachment not found or expired"}, status_code=404
                )

            file_path = storage.get_attachment_path(file_id)
            if not file_path:
                return JSONResponse(
                    {"error": "Attachment file not found"}, status_code=404
                )

            return FileResponse(
                path=str(file_path),
                filename=metadata["filename"],
                media_type=metadata["mime_type"],
            )

    def start(self) -> tuple[bool, str]:
        """
        Start the minimal OAuth server.

        Returns:
            Tuple of (success: bool, error_message: str)
        """
        if self.is_running:
            logger.info("Minimal OAuth server is already running")
            return True, ""

        # Extract hostname from base_uri (e.g., "http://localhost" -> "localhost")
        try:
            parsed_uri = urlparse(self.base_uri)
            hostname = parsed_uri.hostname or "localhost"
        except Exception:
            hostname = "localhost"

        _startup_error = [None]  # mutable container for thread communication

        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:
                _startup_error[0] = e
                logger.error(f"Minimal OAuth server error: {e}", exc_info=True)
                self.is_running = False

        # Start server in background thread
        self.server_thread = threading.Thread(target=run_server, daemon=True)
        self.server_thread.start()

        # Wait for server to start — verify with an actual HTTP request to the
        # callback route so we confirm route registration, not just TCP binding.
        # A missing-code response (400) or any non-404 proves the route exists.
        max_wait = 5.0
        start_time = time.time()
        probe_url = f"http://{hostname}:{self.port}/oauth2callback"
        while time.time() - start_time < max_wait:
            if _startup_error[0]:
                error_msg = f"OAuth server failed to start: {_startup_error[0]}"
                logger.error(error_msg)
                return False, error_msg
            try:
                resp = urllib.request.urlopen(probe_url, timeout=0.5)
                # Any 2xx/3xx means route is up
                if resp.status < 500:
                    self.is_running = True
                    logger.info(
                        f"Minimal OAuth server started on {hostname}:{self.port}"
                    )
                    return True, ""
            except urllib.error.HTTPError as http_err:
                # 4xx responses (e.g. 400 missing code, 422 validation) confirm
                # the route is registered and the server is handling requests.
                if http_err.code != 404:
                    self.is_running = True
                    logger.info(
                        f"Minimal OAuth server started on {hostname}:{self.port}"
                    )
                    return True, ""
            except Exception:
                pass
            time.sleep(0.1)

        error_msg = (
            f"Failed to start minimal OAuth server on {hostname}:{self.port}"
            f" - callback route did not respond within {max_wait}s"
        )
        logger.error(error_msg)
        return False, error_msg

    def reset_auth_state(self):
        """Reset auth completion state so wait_for_auth blocks for a fresh callback."""
        with self._auth_lock:
            self.auth_completed.clear()
            self.auth_result = None

    def wait_for_auth(self, timeout: float = 300) -> Optional[dict]:
        """
        Block until OAuth callback is received or timeout.

        Args:
            timeout: Maximum seconds to wait (default 5 minutes)

        Returns:
            Auth result dict {"success": bool, "user_id": str|None, "error": str|None}
            or None if timed out
        """
        completed = self.auth_completed.wait(timeout=timeout)
        if completed:
            return self.auth_result
        return None

    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("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 ensure_oauth_callback_available(
    transport_mode: str = "stdio", port: int = 8000, base_uri: str = "http://localhost"
) -> tuple[bool, str]:
    """
    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:
        Tuple of (success: bool, error_message: str)
    """
    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")
            success, error_msg = _minimal_oauth_server.start()
            if success:
                logger.info(
                    f"Minimal OAuth server successfully started on {base_uri}:{port}"
                )
                return True, ""
            else:
                logger.error(
                    f"Failed to start minimal OAuth server on {base_uri}:{port}: {error_msg}"
                )
                return False, error_msg
        else:
            logger.info("Minimal OAuth server is already running")
            return True, ""

    else:
        error_msg = f"Unknown transport mode: {transport_mode}"
        logger.error(error_msg)
        return False, error_msg


def set_cli_oauth_server(server: MinimalOAuthServer) -> None:
    """Register a MinimalOAuthServer as the global instance (used by CLI mode)."""
    global _minimal_oauth_server
    _minimal_oauth_server = server


def get_cli_oauth_port() -> int:
    """
    Find an available port for the CLI OAuth callback server.
    Tries ports 8000-8009, falls back to OS-assigned port.

    Returns:
        Available port number
    """
    for port in range(8000, 8010):
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind(("localhost", port))
                return port
        except OSError:
            continue

    # Fallback: let OS assign a port
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("localhost", 0))
        return s.getsockname()[1]


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