implement guided auth flow for both chat based and server callback flows
This commit is contained in:
Taylor Wilsdon
1
auth/__init__.py
Normal file
1
auth/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
# Make the auth directory a Python package
|
||||
BIN
auth/__pycache__/__init__.cpython-312.pyc
Normal file
BIN
auth/__pycache__/__init__.cpython-312.pyc
Normal file
Binary file not shown.
BIN
auth/__pycache__/callback_server.cpython-312.pyc
Normal file
BIN
auth/__pycache__/callback_server.cpython-312.pyc
Normal file
Binary file not shown.
BIN
auth/__pycache__/google_auth.cpython-312.pyc
Normal file
BIN
auth/__pycache__/google_auth.cpython-312.pyc
Normal file
Binary file not shown.
BIN
auth/__pycache__/oauth_manager.cpython-312.pyc
Normal file
BIN
auth/__pycache__/oauth_manager.cpython-312.pyc
Normal file
Binary file not shown.
192
auth/callback_server.py
Normal file
192
auth/callback_server.py
Normal file
@@ -0,0 +1,192 @@
|
||||
# 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
|
||||
182
auth/oauth_manager.py
Normal file
182
auth/oauth_manager.py
Normal file
@@ -0,0 +1,182 @@
|
||||
# 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}"
|
||||
Reference in New Issue
Block a user