From 71c9369b593e9162e4e6129da170f2d637befe7f Mon Sep 17 00:00:00 2001 From: adamhaley Date: Fri, 15 Aug 2025 17:37:30 -0700 Subject: [PATCH] implement seamless automatic OAuth flow with callback handling --- __init__.py | 17 +++++ dropbox_auth_manager.py | 20 +++++- oauth_handler.py | 148 ++++++++++++++++++++++++++++++++++++++++ setup_dropbox_node.py | 47 ++++++++++--- 4 files changed, 219 insertions(+), 13 deletions(-) create mode 100644 oauth_handler.py diff --git a/__init__.py b/__init__.py index c440aa5..7f364db 100644 --- a/__init__.py +++ b/__init__.py @@ -11,3 +11,20 @@ NODE_CLASS_MAPPINGS.update(SETUP_CLASS) NODE_DISPLAY_NAME_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS.update(RUN_DISPLAY) NODE_DISPLAY_NAME_MAPPINGS.update(SETUP_DISPLAY) + +# Register OAuth callback route with ComfyUI server +def register_oauth_routes(): + """Register OAuth callback routes with ComfyUI's PromptServer""" + try: + from server import PromptServer + from .oauth_handler import handle_oauth_callback + + # Add the OAuth callback route + PromptServer.instance.app.router.add_get('/oauth/dropbox/callback', handle_oauth_callback) + print("[DropSendNode] Registered OAuth callback route: /oauth/dropbox/callback") + + except Exception as e: + print(f"[DropSendNode] Warning: Could not register OAuth routes: {e}") + +# Register routes when the node is loaded +register_oauth_routes() diff --git a/dropbox_auth_manager.py b/dropbox_auth_manager.py index 34af144..0f49279 100644 --- a/dropbox_auth_manager.py +++ b/dropbox_auth_manager.py @@ -28,7 +28,7 @@ class DropboxAuthManager: except KeyringError as e: raise RuntimeError(f"Failed to store Dropbox tokens securely: {e}") - def exchange_auth_code(self, auth_code): + def exchange_auth_code(self, auth_code, redirect_uri=None): if not (self.app_key and self.app_secret): raise ValueError("App key and secret must be set before exchanging auth code.") @@ -38,6 +38,11 @@ class DropboxAuthManager: "client_id": self.app_key, "client_secret": self.app_secret } + + # Include redirect_uri if it was used in the authorization request + if redirect_uri: + data["redirect_uri"] = redirect_uri + headers = {"Content-Type": "application/x-www-form-urlencoded"} response = requests.post("https://api.dropbox.com/oauth2/token", headers=headers, data=data) response.raise_for_status() @@ -85,9 +90,18 @@ class DropboxAuthManager: # If passwords don't exist, that's fine pass - def get_oauth_url(self): + def get_oauth_url(self, redirect_uri=None, state=None): """Generate OAuth URL for initial authorization""" if not self.app_key: raise ValueError("App key must be set to generate OAuth URL.") - return f"https://www.dropbox.com/oauth2/authorize?response_type=code&client_id={self.app_key}&token_access_type=offline" + + url = f"https://www.dropbox.com/oauth2/authorize?response_type=code&client_id={self.app_key}&token_access_type=offline" + + if redirect_uri: + url += f"&redirect_uri={redirect_uri}" + + if state: + url += f"&state={state}" + + return url diff --git a/oauth_handler.py b/oauth_handler.py new file mode 100644 index 0000000..4f46512 --- /dev/null +++ b/oauth_handler.py @@ -0,0 +1,148 @@ +# oauth_handler.py - OAuth callback handling for DropSendNode + +import asyncio +from aiohttp import web +from .dropbox_auth_manager import DropboxAuthManager + +# Global state to track pending OAuth flows +pending_oauth_sessions = {} + +class OAuthCallbackHandler: + def __init__(self): + self.session_id = None + self.app_key = None + self.app_secret = None + self.completion_callback = None + + def start_oauth_session(self, session_id, app_key, app_secret, completion_callback=None): + """Start a new OAuth session and store the credentials for callback processing""" + print(f"[OAuthHandler] Starting OAuth session: {session_id}") + + self.session_id = session_id + self.app_key = app_key + self.app_secret = app_secret + self.completion_callback = completion_callback + + # Store in global pending sessions + pending_oauth_sessions[session_id] = { + 'app_key': app_key, + 'app_secret': app_secret, + 'completion_callback': completion_callback, + 'handler': self + } + + print(f"[OAuthHandler] OAuth session stored for: {session_id}") + + async def process_callback(self, auth_code, redirect_uri=None): + """Process the OAuth callback with the received auth code""" + try: + print(f"[OAuthHandler] Processing callback for session: {self.session_id}") + print(f"[OAuthHandler] Auth code received: {auth_code[:10]}...") + print(f"[OAuthHandler] Redirect URI: {redirect_uri}") + + # Use DropboxAuthManager to exchange the code + auth_manager = DropboxAuthManager(self.app_key, self.app_secret) + auth_manager.exchange_auth_code(auth_code, redirect_uri=redirect_uri) + + print(f"[OAuthHandler] Auth code exchange successful") + + # Call completion callback if provided + if self.completion_callback: + await self.completion_callback(True, "✅ Dropbox connected successfully!") + + return True, "✅ Dropbox connected successfully! Credentials stored securely." + + except Exception as e: + error_msg = f"❌ OAuth callback failed: {e}" + print(f"[OAuthHandler] ERROR: {error_msg}") + + if self.completion_callback: + await self.completion_callback(False, error_msg) + + return False, error_msg + +# OAuth callback route handler +async def handle_oauth_callback(request): + """Handle incoming OAuth callback from Dropbox""" + try: + # Extract parameters from callback + auth_code = request.query.get('code') + error = request.query.get('error') + session_id = request.query.get('state') # We'll use state parameter to track sessions + + print(f"[OAuthCallback] Received callback - code: {bool(auth_code)}, error: {error}, state: {session_id}") + + if error: + error_msg = f"OAuth authorization failed: {error}" + print(f"[OAuthCallback] ERROR: {error_msg}") + return web.Response(text=f""" + +

❌ Authorization Failed

+

{error_msg}

+

You can close this window and try again in ComfyUI.

+ + """, content_type='text/html', status=400) + + if not auth_code: + print(f"[OAuthCallback] ERROR: No authorization code received") + return web.Response(text=""" + +

❌ No Authorization Code

+

No authorization code was received. Please try again.

+

You can close this window and try again in ComfyUI.

+ + """, content_type='text/html', status=400) + + # Find the pending OAuth session + if session_id and session_id in pending_oauth_sessions: + session_data = pending_oauth_sessions[session_id] + handler = session_data['handler'] + + # Process the callback with the redirect URI + callback_url = "http://localhost:8188/oauth/dropbox/callback" + success, message = await handler.process_callback(auth_code, redirect_uri=callback_url) + + # Clean up the session + del pending_oauth_sessions[session_id] + + if success: + return web.Response(text=f""" + +

✅ Authorization Successful!

+

{message}

+

You can now close this window and return to ComfyUI.

+ + + """, content_type='text/html') + else: + return web.Response(text=f""" + +

❌ Setup Failed

+

{message}

+

You can close this window and try again in ComfyUI.

+ + """, content_type='text/html', status=500) + + else: + # Fallback for sessions without state tracking + print("[OAuthCallback] No session found, attempting generic processing") + return web.Response(text=f""" + +

🔄 Authorization Code Received

+

Please return to ComfyUI and manually enter this code:

+

{auth_code}

+

You can close this window after copying the code.

+ + """, content_type='text/html') + + except Exception as e: + print(f"[OAuthCallback] Unexpected error: {e}") + return web.Response(text=f""" + +

❌ Callback Error

+

An unexpected error occurred: {e}

+

You can close this window and try again in ComfyUI.

+ + """, content_type='text/html', status=500) \ No newline at end of file diff --git a/setup_dropbox_node.py b/setup_dropbox_node.py index 7f4db42..0162abf 100644 --- a/setup_dropbox_node.py +++ b/setup_dropbox_node.py @@ -3,9 +3,11 @@ import os import requests import webbrowser +import uuid from dotenv import load_dotenv, dotenv_values import urllib.parse from .dropbox_auth_manager import DropboxAuthManager +from .oauth_handler import OAuthCallbackHandler class DropboxSetupNode: @@ -45,6 +47,10 @@ class DropboxSetupNode: "label": "Reset stored credentials", "default": False }) + inputs["optional"]["auto_oauth"] = ("BOOLEAN", { + "label": "Automatic OAuth (no manual auth code needed)", + "default": True + }) return inputs @@ -52,7 +58,7 @@ class DropboxSetupNode: OUTPUT_NODE = True FUNCTION = "setup" - def setup(self, dropbox_dest_folder, app_key=None, app_secret=None, auth_code=None, reconnect=False): + def setup(self, dropbox_dest_folder, app_key=None, app_secret=None, auth_code=None, reconnect=False, auto_oauth=True): try: print(f"[DropboxSetup] Called with:") print(f" app_key: '{app_key}' (type: {type(app_key)}, bool: {bool(app_key)})") @@ -135,16 +141,37 @@ class DropboxSetupNode: if not auth_code_clean: print(f"[DropboxSetup] No auth code provided - generating OAuth URL") auth_temp = DropboxAuthManager(app_key=app_key_clean) - oauth_url = auth_temp.get_oauth_url() - # Automatically open browser for better UX - try: - print(f"[DropboxSetup] Opening browser automatically...") - webbrowser.open(oauth_url) - message = f"🌐 Browser opened for Dropbox authorization!\n\n📋 After authorizing, paste the code in the auth_code field:\n{oauth_url}" - except Exception as e: - print(f"[DropboxSetup] Could not auto-open browser: {e}") - message = f"📋 Visit this URL to authorize, then paste the auth code:\n{oauth_url}" + if auto_oauth: + # Automatic OAuth flow with callback + session_id = str(uuid.uuid4()) + callback_url = "http://localhost:8188/oauth/dropbox/callback" + oauth_url = auth_temp.get_oauth_url(redirect_uri=callback_url, state=session_id) + + # Set up OAuth session for callback handling + oauth_handler = OAuthCallbackHandler() + oauth_handler.start_oauth_session(session_id, app_key_clean, app_secret_clean) + + try: + print(f"[DropboxSetup] Opening browser for automatic OAuth flow...") + webbrowser.open(oauth_url) + message = f"🚀 Automatic OAuth Started!\n\n🌐 Browser opened for authorization. After you authorize in Dropbox, this will complete automatically.\n\n⏳ Waiting for authorization..." + print(f"[DropboxSetup] Session ID: {session_id}") + print(f"[DropboxSetup] Callback URL: {callback_url}") + except Exception as e: + print(f"[DropboxSetup] Could not auto-open browser: {e}") + message = f"🔗 Automatic OAuth Setup:\n\nPlease visit this URL to authorize:\n{oauth_url}\n\nAfter authorization, setup will complete automatically!" + else: + # Manual OAuth flow (original behavior) + oauth_url = auth_temp.get_oauth_url() + + try: + print(f"[DropboxSetup] Opening browser for manual OAuth flow...") + webbrowser.open(oauth_url) + message = f"🌐 Browser opened for Dropbox authorization!\n\n📋 After authorizing, paste the code in the auth_code field:\n{oauth_url}" + except Exception as e: + print(f"[DropboxSetup] Could not auto-open browser: {e}") + message = f"📋 Visit this URL to authorize, then paste the auth code:\n{oauth_url}" print(f"[DropboxSetup] OAuth URL: {oauth_url}")