jason/google-mcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

working non loop server but staged earlier than before

Browse Source
This commit is contained in:
Taylor Wilsdon
2025-04-27 15:33:48 -04:00
parent d0978cb2f6
commit 43ef646aa4
7 changed files with 50 additions and 641 deletions
Download Patch File Download Diff File Expand all files Collapse all files

192
auth/callback_server.py
View File

@@ -1,192 +0,0 @@
# auth/callback_server.py
import http.server
import logging
import socketserver
import threading
import urllib.parse
import webbrowser
from typing import Callable, Optional, Dict, Any
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class OAuthCallbackHandler(http.server.BaseHTTPRequestHandler):
"""Handler for OAuth callback requests."""
# Class variable to store the callback function
callback_function: Optional[Callable] = None
def do_GET(self):
"""Handle GET requests to the callback endpoint."""
try:
# Parse the URL and extract query parameters
parsed_url = urllib.parse.urlparse(self.path)
query_params = urllib.parse.parse_qs(parsed_url.query)
# Check if we're handling the OAuth callback
if parsed_url.path == '/callback':
# Extract authorization code and state
code = query_params.get('code', [''])[0]
state = query_params.get('state', [''])[0]
logger.info(f"Received OAuth callback with code: {code[:10]}... and state: {state}")
# Show success page to the user
self.send_response(200)
self.send_header('Content-Type', 'text/html')
self.end_headers()
html_content = """
<!DOCTYPE html>
<html>
<head>
<title>Google OAuth - Success</title>
<style>
body {
font-family: Arial, sans-serif;
max-width: 600px;
margin: 40px auto;
padding: 20px;
text-align: center;
}
.success {
color: #4CAF50;
font-size: 24px;
margin-bottom: 20px;
}
.info {
color: #555;
margin-bottom: 30px;
}
.close-button {
background-color: #4CAF50;
color: white;
padding: 10px 20px;
border: none;
border-radius: 4px;
cursor: pointer;
font-size: 16px;
}
</style>
</head>
<body>
<div class="success">✅ Authentication Successful!</div>
<div class="info">
You have successfully authenticated with Google.
You can now close this window and return to your application.
</div>
<button class="close-button" onclick="window.close()">Close Window</button>
<script>
// Auto-close after 10 seconds
setTimeout(function() {
window.close();
}, 10000);
</script>
</body>
</html>
"""
self.wfile.write(html_content.encode())
# Call the callback function if provided
if OAuthCallbackHandler.callback_function and code:
threading.Thread(
target=OAuthCallbackHandler.callback_function,
args=(code, state),
daemon=True
).start()
# Signal the server to shutdown after handling the request
threading.Thread(
target=self.server.shutdown,
daemon=True
).start()
else:
# Handle other paths with a 404 response
self.send_response(404)
self.end_headers()
self.wfile.write(b"Not Found")
except Exception as e:
logger.error(f"Error handling callback request: {e}")
self.send_response(500)
self.end_headers()
self.wfile.write(f"Internal Server Error: {str(e)}".encode())
def log_message(self, format, *args):
"""Override to use our logger instead of printing to stderr."""
logger.info(f"{self.address_string()} - {format%args}")
class OAuthCallbackServer:
"""Server to handle OAuth callbacks."""
def __init__(self,
port: int = 8080,
callback: Optional[Callable] = None,
auto_open_browser: bool = True):
"""
Initialize the callback server.
Args:
port: Port to listen on (default: 8080)
callback: Function to call with the code and state
auto_open_browser: Whether to automatically open the browser
"""
self.port = port
self.server = None
self.server_thread = None
self.auto_open_browser = auto_open_browser
# Set the callback function
OAuthCallbackHandler.callback_function = callback
def start(self) -> None:
"""Start the callback server in a separate thread."""
if self.server:
logger.warning("Server is already running")
return
try:
# Create and start the server
self.server = socketserver.TCPServer(('localhost', self.port), OAuthCallbackHandler)
logger.info(f"Starting OAuth callback server on port {self.port}")
# Run the server in a separate thread
self.server_thread = threading.Thread(
target=self.server.serve_forever,
daemon=True
)
self.server_thread.start()
logger.info(f"OAuth callback server is running on http://localhost:{self.port}")
except Exception as e:
logger.error(f"Failed to start callback server: {e}")
raise
def stop(self) -> None:
"""Stop the callback server."""
if self.server:
logger.info("Stopping OAuth callback server")
self.server.shutdown()
self.server.server_close()
self.server = None
self.server_thread = None
else:
logger.warning("Server is not running")
def open_browser(self, url: str) -> bool:
"""Open the default web browser to the given URL."""
if not self.auto_open_browser:
return False
try:
logger.info(f"Opening browser to: {url}")
webbrowser.open(url)
return True
except Exception as e:
logger.error(f"Failed to open browser: {e}")
return False

43
auth/google_auth.py
View File

@@ -3,7 +3,7 @@
import os
import json
import logging
from typing import List, Optional, Tuple, Dict, Any, Callable
from typing import List, Optional, Tuple, Dict, Any
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import Flow, InstalledAppFlow
@@ -11,16 +11,13 @@ from google.auth.transport.requests import Request
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from auth.callback_server import OAuthCallbackServer
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Constants
DEFAULT_CREDENTIALS_DIR = ".credentials"
DEFAULT_REDIRECT_URI = "http://localhost:8080/callback"
DEFAULT_SERVER_PORT = 8080
DEFAULT_REDIRECT_URI = "http://localhost:8080/callback" # Example, should be configurable
# --- Helper Functions ---
@@ -95,14 +92,7 @@ def load_client_secrets(client_secrets_path: str) -> Dict[str, Any]:
# --- Core OAuth Logic ---
def start_auth_flow(
client_secrets_path: str,
scopes: List[str],
redirect_uri: str = DEFAULT_REDIRECT_URI,
auto_handle_callback: bool = False,
callback_function: Optional[Callable] = None,
port: int = DEFAULT_SERVER_PORT
) -> Tuple[str, str]:
def start_auth_flow(client_secrets_path: str, scopes: List[str], redirect_uri: str = DEFAULT_REDIRECT_URI) -> Tuple[str, str]:
"""
Initiates the OAuth 2.0 flow and returns the authorization URL and state.
@@ -110,30 +100,23 @@ def start_auth_flow(
client_secrets_path: Path to the Google client secrets JSON file.
scopes: List of OAuth scopes required.
redirect_uri: The URI Google will redirect to after authorization.
auto_handle_callback: Whether to automatically handle the callback by
starting a local server on the specified port.
callback_function: Function to call with the code and state when received.
port: Port to run the callback server on, if auto_handle_callback is True.
Returns:
A tuple containing the authorization URL and the state parameter.
"""
try:
# Create and start the callback server if auto_handle_callback is enabled
server = None
if auto_handle_callback:
logger.info("Starting OAuth callback server")
server = OAuthCallbackServer(port=port, callback=callback_function, auto_open_browser=False)
server.start()
# Load client secrets using the helper
# Note: Flow.from_client_secrets_file handles loading internally,
# but loading separately allows access to client_id/secret if needed elsewhere.
# client_config = load_client_secrets(client_secrets_path) # Keep if needed
# Set up the OAuth flow
flow = Flow.from_client_secrets_file(
client_secrets_path,
scopes=scopes,
redirect_uri=redirect_uri
)
# Indicate that the user needs *offline* access to retrieve a refresh token.
# Indicate that the user *offline* access to retrieve a refresh token.
# 'prompt': 'consent' ensures the user sees the consent screen even if
# they have previously granted permissions, which is useful for getting
# a refresh token again if needed.
@@ -142,19 +125,11 @@ def start_auth_flow(
prompt='consent'
)
logger.info(f"Generated authorization URL. State: {state}")
# Auto-open the browser if requested
if auto_handle_callback and server:
server.open_browser(authorization_url)
return authorization_url, state
except Exception as e:
logger.error(f"Error starting OAuth flow: {e}")
# If we created a server, shut it down
if auto_handle_callback and server:
server.stop()
raise # Re-raise the exception for the caller to handle
raise # Re-raise the exception for the caller to handle
def handle_auth_callback(
client_secrets_path: str,

182
auth/oauth_manager.py
View File

@@ -1,182 +0,0 @@
# auth/oauth_manager.py
import logging
import os
import threading
import time
from typing import Dict, Optional, Callable, Any, Tuple
from auth.callback_server import OAuthCallbackServer
from auth.google_auth import start_auth_flow, handle_auth_callback, get_credentials, get_user_info
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Track active OAuth flows
active_flows: Dict[str, Dict[str, Any]] = {}
def start_oauth_flow(
user_id: str,
scopes: list,
client_secrets_path: str = 'client_secret.json',
port: int = 8080
) -> str:
"""
Start an OAuth flow with automatic callback handling.
Args:
user_id: The unique identifier (e.g., email address) for the user.
scopes: List of OAuth scopes required.
client_secrets_path: Path to the Google client secrets JSON file.
port: Port to run the callback server on.
Returns:
A string with instructions for the user, including the authentication URL.
"""
logger.info(f"Starting OAuth flow for user {user_id} with scopes {scopes}")
# Cleanup any previous flow for this user
if user_id in active_flows:
stop_oauth_flow(user_id)
# Create a callback function for this user
def handle_code(code: str, state: str) -> None:
try:
logger.info(f"Received authorization code for user {user_id}")
# Construct full callback URL
redirect_uri = f"http://localhost:{port}/callback"
full_callback_url = f"{redirect_uri}?code={code}&state={state}"
# Exchange code for credentials
authenticated_user, credentials = handle_auth_callback(
client_secrets_path=client_secrets_path,
scopes=scopes,
authorization_response=full_callback_url,
redirect_uri=redirect_uri
)
# Update flow status
active_flows[user_id]["status"] = "authenticated"
active_flows[user_id]["authenticated_user"] = authenticated_user
active_flows[user_id]["credentials"] = credentials
logger.info(f"Authentication successful for user {authenticated_user}")
except Exception as e:
logger.error(f"Error handling OAuth callback for user {user_id}: {e}")
active_flows[user_id]["status"] = "error"
active_flows[user_id]["error"] = str(e)
# Start a callback server
callback_server = OAuthCallbackServer(
port=port,
callback=handle_code,
auto_open_browser=True # Auto-open browser for better UX
)
try:
# Start the server
callback_server.start()
# Generate the authorization URL
redirect_uri = f"http://localhost:{port}/callback"
auth_url, state = start_auth_flow(
client_secrets_path=client_secrets_path,
scopes=scopes,
redirect_uri=redirect_uri,
auto_handle_callback=False # We're handling it ourselves
)
# Store flow information
active_flows[user_id] = {
"status": "pending",
"start_time": time.time(),
"scopes": scopes,
"server": callback_server,
"auth_url": auth_url,
"state": state
}
# Return instructions for the user
return (
f"Authentication required. Please visit this URL to authorize access: "
f"{auth_url}\n\n"
f"A browser window should open automatically. After authorizing, you'll be "
f"redirected to a success page.\n\n"
f"If the browser doesn't open automatically, copy and paste the URL into your browser. "
f"You can also check the status of your authentication by using:\n\n"
f"check_auth_status\n"
f"user_id: {user_id}"
)
except Exception as e:
logger.error(f"Error starting OAuth flow for user {user_id}: {e}")
# Clean up the server if it was started
if "server" in active_flows.get(user_id, {}):
active_flows[user_id]["server"].stop()
del active_flows[user_id]
raise
def check_auth_status(user_id: str) -> str:
"""
Check the status of an active OAuth flow.
Args:
user_id: The unique identifier for the user.
Returns:
A string describing the current status.
"""
if user_id not in active_flows:
return f"No active authentication flow found for user {user_id}."
flow = active_flows[user_id]
status = flow.get("status", "unknown")
if status == "authenticated":
authenticated_user = flow.get("authenticated_user", "unknown")
return (
f"Authentication successful for user {authenticated_user}. "
f"You can now use the Google Calendar tools."
)
elif status == "error":
error = flow.get("error", "Unknown error")
return f"Authentication failed: {error}"
elif status == "pending":
elapsed = int(time.time() - flow.get("start_time", time.time()))
auth_url = flow.get("auth_url", "")
return (
f"Authentication pending for {elapsed} seconds. "
f"Please complete the authorization at: {auth_url}"
)
else:
return f"Unknown authentication status: {status}"
def stop_oauth_flow(user_id: str) -> str:
"""
Stop an active OAuth flow and clean up resources.
Args:
user_id: The unique identifier for the user.
Returns:
A string describing the result.
"""
if user_id not in active_flows:
return f"No active authentication flow found for user {user_id}."
try:
# Stop the callback server
if "server" in active_flows[user_id]:
active_flows[user_id]["server"].stop()
# Remove the flow
final_status = active_flows[user_id].get("status", "unknown")
del active_flows[user_id]
return f"Authentication flow for user {user_id} stopped. Final status: {final_status}."
except Exception as e:
logger.error(f"Error stopping OAuth flow for user {user_id}: {e}")
return f"Error stopping authentication flow: {e}"