implement seamless automatic OAuth flow with callback handling

__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()

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
 

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"""
+            <html><body>
+                <h2>❌ Authorization Failed</h2>
+                <p>{error_msg}</p>
+                <p>You can close this window and try again in ComfyUI.</p>
+            </body></html>
+            """, content_type='text/html', status=400)
+        
+        if not auth_code:
+            print(f"[OAuthCallback] ERROR: No authorization code received")
+            return web.Response(text="""
+            <html><body>
+                <h2>❌ No Authorization Code</h2>
+                <p>No authorization code was received. Please try again.</p>
+                <p>You can close this window and try again in ComfyUI.</p>
+            </body></html>
+            """, 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"""
+                <html><body>
+                    <h2>✅ Authorization Successful!</h2>
+                    <p>{message}</p>
+                    <p>You can now close this window and return to ComfyUI.</p>
+                    <script>
+                        setTimeout(() => window.close(), 3000);
+                    </script>
+                </body></html>
+                """, content_type='text/html')
+            else:
+                return web.Response(text=f"""
+                <html><body>
+                    <h2>❌ Setup Failed</h2>
+                    <p>{message}</p>
+                    <p>You can close this window and try again in ComfyUI.</p>
+                </body></html>
+                """, 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"""
+            <html><body>
+                <h2>🔄 Authorization Code Received</h2>
+                <p>Please return to ComfyUI and manually enter this code:</p>
+                <p><strong>{auth_code}</strong></p>
+                <p>You can close this window after copying the code.</p>
+            </body></html>
+            """, content_type='text/html')
+            
+    except Exception as e:
+        print(f"[OAuthCallback] Unexpected error: {e}")
+        return web.Response(text=f"""
+        <html><body>
+            <h2>❌ Callback Error</h2>
+            <p>An unexpected error occurred: {e}</p>
+            <p>You can close this window and try again in ComfyUI.</p>
+        </body></html>
+        """, content_type='text/html', status=500)

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}")