src/Agentic_System/test/_oauth_server.py

"""
OAuth helper: starts the Google MCP server with the OAuth callback server,
waits for the user to authorize, then exchanges the code for credentials.

Usage: python _oauth_server.py
  - Opens the auth URL in your browser
  - Waits for the callback on localhost:8000
  - Saves credentials to ~/.google_workspace_mcp/credentials/
  - Press Ctrl+C to stop
"""
import asyncio
import os
import sys
import webbrowser

# Add google-mcp-server to path
MCP_SERVER_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "google-mcp-server")
sys.path.insert(0, MCP_SERVER_DIR)

from dotenv import load_dotenv
load_dotenv(os.path.join(MCP_SERVER_DIR, ".env"))
load_dotenv(os.path.join(os.path.dirname(__file__), ".env"))
load_dotenv(os.path.join(os.path.dirname(__file__), "../../.env"))

from auth.google_auth import (
    create_oauth_flow,
    load_client_secrets_from_env,
    check_client_secrets,
)
from auth.scopes import get_scopes_for_tools
from auth.oauth_callback_server import ensure_oauth_callback_available, cleanup_oauth_callback_server
from auth.credential_store import get_credential_store

EMAIL = os.getenv("USER_GOOGLE_EMAIL", "aiwithjawadsaghir@gmail.com")
REDIRECT_URI = "http://localhost:8000/oauth2callback"

async def main():
    print("=" * 60)
    print("Google OAuth Re-Authentication Helper")
    print("=" * 60)

    # Check client secrets
    err = check_client_secrets()
    if err:
        print(f"ERROR: {err}")
        return

    # Start the OAuth callback server
    print(f"Starting OAuth callback server on {REDIRECT_URI}...")
    success, error_msg = ensure_oauth_callback_available("http://localhost", 8000)
    if not success:
        print(f"ERROR: Could not start callback server: {error_msg}")
        return
    print("Callback server running.")

    # Get scopes for gmail
    scopes = get_scopes_for_tools(["gmail"])
    print(f"Requesting scopes: {len(scopes)} scopes for Gmail")

    # Create OAuth flow
    flow = create_oauth_flow(scopes=list(scopes), redirect_uri=REDIRECT_URI)
    auth_url, _ = flow.authorization_url(
        access_type="offline",
        prompt="consent",
    )

    print()
    print("Opening auth URL in your browser...")
    print(f"URL: {auth_url[:100]}...")
    print()
    print("After authorizing in the browser, the callback will be received here.")
    print("Waiting for callback... (Press Ctrl+C to cancel)")
    webbrowser.open(auth_url)

    # Wait for the callback by polling the callback server
    # The callback server stores the auth code in shared state
    # We just need to keep the server alive
    try:
        while True:
            await asyncio.sleep(1)
    except KeyboardInterrupt:
        print("\nShutting down...")
    finally:
        cleanup_oauth_callback_server()

if __name__ == "__main__":
    asyncio.run(main())