2025-04-27 16:59:35 -04:00
# auth/callback_server.py
import http . server
import logging
2025-05-10 17:57:25 -04:00
import os
2025-04-27 16:59:35 -04:00
import socketserver
import threading
import urllib . parse
import webbrowser
2025-05-10 17:57:25 -04:00
from typing import Callable , Optional , Dict , Any , Literal
from oauthlib . oauth2 . rfc6749 . errors import InsecureTransportError
2025-04-27 16:59:35 -04:00
# Configure logging
logging . basicConfig ( level = logging . INFO )
logger = logging . getLogger ( __name__ )
2025-05-10 17:57:25 -04:00
def get_response_html (
title : str ,
status : Literal [ " success " , " error " ] ,
message : str ,
show_close_button : bool = True ,
auto_close_seconds : int = 0
) - > str :
""" Generate HTML response for OAuth callback.
Args:
title: Page title
status: ' success ' or ' error '
message: Message to display to the user
show_close_button: Whether to show a close button
auto_close_seconds: Auto-close after this many seconds (0 to disable)
Returns:
HTML content as a string
"""
color = " #4CAF50 " if status == " success " else " #d32f2f "
close_button = """
<button class= " button " onclick= " window.close() " >Close Window</button>
""" if show_close_button else " "
auto_close_script = f """
<script>
setTimeout(function() {{ window.close(); }} , { auto_close_seconds * 1000 } );
</script>
""" if auto_close_seconds > 0 else " "
return f """ <!DOCTYPE html>
<html>
<head>
<title> { title } </title>
<style>
body {{
font-family: -apple-system, BlinkMacSystemFont, " Segoe UI " , Roboto, Helvetica, Arial, sans-serif;
max-width: 500px;
margin: 40px auto;
padding: 20px;
text-align: center;
color: #333;
}}
.status {{
color: { color } ;
font-size: 24px;
margin-bottom: 20px;
}}
.message {{
margin-bottom: 30px;
line-height: 1.5;
}}
.button {{
background-color: { color } ;
color: white;
padding: 10px 20px;
border: none;
border-radius: 4px;
cursor: pointer;
font-size: 16px;
}}
</style>
</head>
<body>
2025-05-10 18:50:01 -04:00
<div class= " status " > { title } </div>
2025-05-10 17:57:25 -04:00
<div class= " message " > { message } </div>
{ close_button }
{ auto_close_script }
</body>
</html> """
2025-04-27 16:59:35 -04:00
class OAuthCallbackHandler ( http . server . BaseHTTPRequestHandler ) :
""" Handler for OAuth callback requests. """
2025-05-10 17:57:25 -04:00
# Class variables to store callback functions by state
callback_registry : Dict [ str , Callable ] = { }
@classmethod
def register_callback ( cls , state : str , callback : Callable ) - > None :
""" Register a callback function for a specific state parameter. """
logger . info ( f " Registering callback for state: { state } " )
cls . callback_registry [ state ] = callback
@classmethod
def unregister_callback ( cls , state : str ) - > None :
""" Unregister a callback function for a specific state parameter. """
if state in cls . callback_registry :
logger . info ( f " Unregistering callback for state: { state } " )
del cls . callback_registry [ state ]
2025-04-27 16:59:35 -04:00
def do_GET ( self ) :
""" Handle GET requests to the callback endpoint. """
request_thread_id = threading . get_ident ( )
logger . info ( f " [Handler { request_thread_id } ] GET request received for path: { self . path } " )
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 ( )
2025-05-10 17:57:25 -04:00
html_content = get_response_html (
title = " Authentication Successful " ,
status = " success " ,
message = " You have successfully authenticated with Google. You can now close this window and return to your application. " ,
show_close_button = True ,
auto_close_seconds = 10
)
2025-04-27 16:59:35 -04:00
self . wfile . write ( html_content . encode ( ) )
2025-05-10 17:57:25 -04:00
try :
# Ensure OAUTHLIB_INSECURE_TRANSPORT is set
if ' OAUTHLIB_INSECURE_TRANSPORT ' not in os . environ :
logger . warning ( " OAUTHLIB_INSECURE_TRANSPORT not set. Setting it for localhost development. " )
os . environ [ ' OAUTHLIB_INSECURE_TRANSPORT ' ] = ' 1 '
# Call the appropriate callback function based on state
if state in OAuthCallbackHandler . callback_registry and code :
logger . info ( f " [Handler { request_thread_id } ] Found callback for state: { state } " )
callback_function = OAuthCallbackHandler . callback_registry [ state ]
logger . info ( f " [Handler { request_thread_id } ] Preparing to call callback function in new thread. " )
callback_thread = threading . Thread (
target = callback_function ,
args = ( code , state ) ,
daemon = True
)
callback_thread . start ( )
logger . info ( f " [Handler { request_thread_id } ] Callback function thread started (ID: { callback_thread . ident } ). " )
# Unregister the callback after it's been called
OAuthCallbackHandler . unregister_callback ( state )
else :
logger . warning ( f " [Handler { request_thread_id } ] No callback registered for state: { state } or no code received. " )
except InsecureTransportError as e :
logger . error ( f " [Handler { request_thread_id } ] InsecureTransportError: { e } . Ensure OAUTHLIB_INSECURE_TRANSPORT is set for localhost development. " )
self . send_response ( 400 )
self . send_header ( ' Content-Type ' , ' text/html ' )
self . end_headers ( )
error_html = get_response_html (
title = " OAuth Error: Insecure Transport " ,
status = " error " ,
message = " The OAuth flow requires HTTPS or explicit allowance of HTTP for localhost development. Please ensure OAUTHLIB_INSECURE_TRANSPORT is set in your environment. " ,
show_close_button = False
2025-04-27 16:59:35 -04:00
)
2025-05-10 17:57:25 -04:00
self . wfile . write ( error_html . encode ( ) )
return
2025-04-27 16:59:35 -04:00
2025-05-10 17:57:25 -04:00
# Note: We no longer shut down the server after handling a callback
# This allows it to handle multiple auth flows over time
2025-04-27 16:59:35 -04:00
else :
# Handle other paths with a 404 response
self . send_response ( 404 )
2025-05-10 17:57:25 -04:00
self . send_header ( ' Content-Type ' , ' text/html ' )
2025-04-27 16:59:35 -04:00
self . end_headers ( )
2025-05-10 17:57:25 -04:00
error_html = get_response_html (
title = " Not Found " ,
status = " error " ,
message = " The requested resource was not found. " ,
show_close_button = False
)
self . wfile . write ( error_html . encode ( ) )
2025-04-27 16:59:35 -04:00
except Exception as e :
logger . error ( f " [Handler { request_thread_id } ] Error handling callback request: { e } " , exc_info = True )
try :
self . send_response ( 500 )
2025-05-10 17:57:25 -04:00
self . send_header ( ' Content-Type ' , ' text/html ' )
2025-04-27 16:59:35 -04:00
self . end_headers ( )
2025-05-10 17:57:25 -04:00
error_html = get_response_html (
title = " Internal Server Error " ,
status = " error " ,
message = f " An error occurred while processing your request: { str ( e ) } " ,
show_close_button = False
)
self . wfile . write ( error_html . encode ( ) )
2025-04-27 16:59:35 -04:00
except Exception as send_error :
2025-05-10 17:57:25 -04:00
logger . error ( f " [Handler { request_thread_id } ] Error sending 500 response: { send_error } " )
2025-04-27 16:59:35 -04:00
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. """
2025-05-10 17:57:25 -04:00
def __init__ ( self ,
port : int = 8080 ,
2025-04-27 16:59:35 -04:00
auto_open_browser : bool = True ) :
"""
Initialize the callback server.
Args:
port: Port to listen on (default: 8080)
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
2025-05-10 17:57:25 -04:00
def start ( self ) - > Dict [ str , Any ] :
"""
Start the callback server in a separate thread.
Returns:
Dict containing server status and port information:
{
' success ' : bool,
' port ' : int,
' message ' : str
}
"""
2025-04-27 16:59:35 -04:00
if self . server :
logger . warning ( " Server is already running " )
2025-05-10 17:57:25 -04:00
return { ' success ' : False , ' port ' : self . port , ' message ' : ' Server is already running ' }
2025-04-27 16:59:35 -04:00
2025-05-10 17:57:25 -04:00
original_port = self . port
max_port = 8090 # Try ports 8080-8090
def serve ( ) :
thread_id = threading . get_ident ( )
logger . info ( f " [Server Thread { thread_id } ] Starting serve_forever loop. " )
try :
self . server . serve_forever ( )
except Exception as serve_e :
logger . error ( f " [Server Thread { thread_id } ] Exception in serve_forever: { serve_e } " , exc_info = True )
finally :
logger . info ( f " [Server Thread { thread_id } ] serve_forever loop finished. " )
# Ensure server_close is called even if shutdown wasn't clean
2025-04-27 16:59:35 -04:00
try :
2025-05-10 17:57:25 -04:00
if self . server :
self . server . server_close ( )
except Exception as close_e :
logger . error ( f " [Server Thread { thread_id } ] Error during server_close: { close_e } " )
2025-04-27 16:59:35 -04:00
2025-05-10 17:57:25 -04:00
try :
while self . port < = max_port :
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 } " )
if self . port != original_port :
logger . info ( f " Successfully reassigned from port { original_port } to { self . port } " )
# Start the server thread
self . server_thread = threading . Thread ( target = serve , daemon = True )
self . server_thread . start ( )
logger . info ( f " OAuth callback server thread started (ID: { self . server_thread . ident } ) on http://localhost: { self . port } " )
return {
' success ' : True ,
' port ' : self . port ,
' message ' : f " Server started successfully on port { self . port } "
}
except OSError as e :
if e . errno == 48 : # Address already in use
logger . warning ( f " Port { self . port } is already in use, trying next port " )
self . port + = 1
if self . port > max_port :
error_msg = f " Failed to find available port in range { original_port } - { max_port } "
logger . error ( error_msg )
return { ' success ' : False , ' port ' : None , ' message ' : error_msg }
continue
else :
logger . error ( f " Failed to start server: { e } " )
return { ' success ' : False , ' port ' : None , ' message ' : str ( e ) }
2025-04-27 16:59:35 -04:00
except Exception as e :
2025-05-10 17:57:25 -04:00
error_msg = f " Failed to start callback server: { e } "
logger . error ( error_msg )
return { ' success ' : False , ' port ' : None , ' message ' : error_msg }
2025-04-27 16:59:35 -04:00
def stop ( self ) - > None :
""" Stop the callback server. """
if self . server :
logger . info ( " Stopping OAuth callback server " )
# shutdown() signals serve_forever to stop
self . server . shutdown ( )
logger . info ( " Server shutdown() called. " )
# Wait briefly for the server thread to finish
if self . server_thread :
2025-05-10 17:57:25 -04:00
self . server_thread . join ( timeout = 2.0 ) # Wait up to 2 seconds
if self . server_thread . is_alive ( ) :
logger . warning ( " Server thread did not exit cleanly after shutdown. " )
2025-04-27 16:59:35 -04:00
# server_close() is now called in the 'finally' block of the serve() function
self . server = None
self . server_thread = None
logger . info ( " Server resources released. " )
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
