Spaces:
Running
Running
| """ | |
| Self-contained Google OAuth re-authentication helper. | |
| Starts a minimal callback server on localhost:8000, creates an OAuth flow | |
| with proper state management, opens the auth URL in your browser, and waits | |
| for the callback to exchange the code for credentials. | |
| Usage: | |
| python _do_oauth.py | |
| The script will: | |
| 1. Start a callback server on http://localhost:8000/oauth2callback | |
| 2. Open the Google auth page in your browser | |
| 3. Wait for you to authorize | |
| 4. Exchange the code for credentials and save them | |
| 5. Exit | |
| """ | |
| import asyncio | |
| import json | |
| import os | |
| import sys | |
| import secrets | |
| import webbrowser | |
| import logging | |
| from urllib.parse import parse_qs, urlparse | |
| # Add google-mcp-server to path so we can use its auth modules | |
| MCP_SERVER_DIR = os.path.abspath( | |
| 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")) | |
| logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") | |
| logger = logging.getLogger(__name__) | |
| # Google OAuth libraries | |
| from google_auth_oauthlib.flow import Flow | |
| # MCP server auth modules | |
| from auth.google_auth import ( | |
| load_client_secrets_from_env, | |
| get_user_info, | |
| ) | |
| from auth.credential_store import get_credential_store | |
| from auth.scopes import get_scopes_for_tools | |
| EMAIL = os.getenv("USER_GOOGLE_EMAIL", "aiwithjawadsaghir@gmail.com") | |
| REDIRECT_URI = "http://localhost:8000/oauth2callback" | |
| PORT = 8000 | |
| # Will be set when the callback is received | |
| _auth_result = {"code": None, "error": None, "state": None} | |
| _auth_event = asyncio.Event() | |
| def create_callback_app(): | |
| """Create a minimal ASGI app with just the OAuth callback route.""" | |
| from starlette.applications import Starlette | |
| from starlette.requests import Request | |
| from starlette.responses import HTMLResponse | |
| from starlette.routing import Route | |
| async def oauth_callback(request: Request): | |
| code = request.query_params.get("code") | |
| error = request.query_params.get("error") | |
| state = request.query_params.get("state") | |
| if error: | |
| _auth_result["error"] = error | |
| _auth_event.set() | |
| return HTMLResponse( | |
| f"<html><body><h2>Authentication Failed</h2><p>{error}</p>" | |
| "<p>You can close this window.</p></body></html>", | |
| status_code=400, | |
| ) | |
| if not code: | |
| _auth_result["error"] = "No authorization code received" | |
| _auth_event.set() | |
| return HTMLResponse( | |
| "<html><body><h2>Error</h2><p>No authorization code received.</p>" | |
| "<p>You can close this window.</p></body></html>", | |
| status_code=400, | |
| ) | |
| _auth_result["code"] = code | |
| _auth_result["state"] = state | |
| _auth_event.set() | |
| return HTMLResponse( | |
| "<html><body>" | |
| "<h2 style='color: green;'>✔ Authentication Successful!</h2>" | |
| "<p>You can close this window and return to the terminal.</p>" | |
| "</body></html>" | |
| ) | |
| return Starlette(routes=[Route("/oauth2callback", oauth_callback)]) | |
| async def main(): | |
| print("=" * 60) | |
| print(" Google OAuth Re-Authentication") | |
| print("=" * 60) | |
| # Load client config | |
| env_config = load_client_secrets_from_env() | |
| if not env_config: | |
| print("ERROR: No OAuth client credentials found.") | |
| print("Set GOOGLE_OAUTH_CLIENT_ID and GOOGLE_OAUTH_CLIENT_SECRET in .env") | |
| return | |
| client_config = env_config["web"] | |
| client_id = client_config["client_id"] | |
| client_secret = client_config["client_secret"] | |
| print(f" Client ID: {client_id[:30]}...") | |
| print(f" Email: {EMAIL}") | |
| # Get scopes for ALL services (so the token works for everything) | |
| all_tools = ["gmail", "calendar", "drive", "docs", "sheets", "slides"] | |
| scopes = list(get_scopes_for_tools(all_tools)) | |
| print(f" Scopes: {len(scopes)} scopes for {', '.join(all_tools)}") | |
| # Allow HTTP for localhost | |
| os.environ["OAUTHLIB_INSECURE_TRANSPORT"] = "1" | |
| # Create OAuth flow | |
| state = secrets.token_hex(16) | |
| flow = Flow.from_client_config( | |
| {"web": { | |
| "client_id": client_id, | |
| "client_secret": client_secret, | |
| "auth_uri": "https://accounts.google.com/o/oauth2/auth", | |
| "token_uri": "https://oauth2.googleapis.com/token", | |
| }}, | |
| scopes=scopes, | |
| redirect_uri=REDIRECT_URI, | |
| state=state, | |
| ) | |
| auth_url, _ = flow.authorization_url( | |
| access_type="offline", | |
| prompt="consent", | |
| ) | |
| # Start callback server | |
| import uvicorn | |
| app = create_callback_app() | |
| config = uvicorn.Config(app, host="0.0.0.0", port=PORT, log_level="warning") | |
| server = uvicorn.Server(config) | |
| print() | |
| print(f" Callback server starting on http://localhost:{PORT}/oauth2callback") | |
| print() | |
| print(" Opening Google authorization page in your browser...") | |
| print(" (If it doesn't open, copy the URL below)") | |
| print() | |
| print(f" {auth_url[:120]}...") | |
| print() | |
| print(" Waiting for authorization...") | |
| print("=" * 60) | |
| # Start server in background | |
| server_task = asyncio.create_task(server.serve()) | |
| # Wait briefly for server to start, then open browser | |
| await asyncio.sleep(1) | |
| webbrowser.open(auth_url) | |
| # Wait for callback | |
| try: | |
| await asyncio.wait_for(_auth_event.wait(), timeout=300) # 5 min timeout | |
| except asyncio.TimeoutError: | |
| print("\nERROR: Timed out waiting for authorization (5 minutes).") | |
| server.should_exit = True | |
| await server_task | |
| return | |
| # Process result | |
| if _auth_result["error"]: | |
| print(f"\nERROR: Authorization failed: {_auth_result['error']}") | |
| server.should_exit = True | |
| await server_task | |
| return | |
| code = _auth_result["code"] | |
| print(f"\n Authorization code received!") | |
| # Exchange code for credentials | |
| try: | |
| flow.fetch_token(code=code) | |
| credentials = flow.credentials | |
| print(f" Token exchange successful!") | |
| # Get user info | |
| user_info = get_user_info(credentials) | |
| user_email = user_info.get("email", EMAIL) if user_info else EMAIL | |
| print(f" Authenticated as: {user_email}") | |
| # Save credentials | |
| store = get_credential_store() | |
| store.store_credential(user_email, credentials) | |
| print(f" Credentials saved!") | |
| print() | |
| print("=" * 60) | |
| print(f" SUCCESS! Credentials saved for {user_email}") | |
| print(f" Token expires: {credentials.expiry}") | |
| print(f" Has refresh token: {bool(credentials.refresh_token)}") | |
| print("=" * 60) | |
| except Exception as e: | |
| print(f"\nERROR: Token exchange failed: {e}") | |
| # Shutdown server | |
| server.should_exit = True | |
| await server_task | |
| if __name__ == "__main__": | |
| asyncio.run(main()) | |