# 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, storage_method="keyring", dropbox_folder="/Apps/ComfyUI_Output_Files", original_redirect_uri=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 self.storage_method = storage_method self.dropbox_folder = dropbox_folder # Store in global pending sessions pending_oauth_sessions[session_id] = { 'app_key': app_key, 'app_secret': app_secret, 'completion_callback': completion_callback, 'storage_method': storage_method, 'dropbox_folder': dropbox_folder, 'original_redirect_uri': original_redirect_uri, # Store the original redirect URI 'handler': self } print(f"[OAuthHandler] OAuth session stored for: {session_id} with storage method: {storage_method}") 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 print(f"[OAuthHandler] Creating DropboxAuthManager for token exchange...") auth_manager = DropboxAuthManager(self.app_key, self.app_secret) print(f"[OAuthHandler] Exchanging auth code for tokens...") result = auth_manager.exchange_auth_code_raw(auth_code, redirect_uri=redirect_uri) refresh_token = result.get("refresh_token") print(f"[OAuthHandler] Auth code exchange successful - got refresh token: {bool(refresh_token)}") print(f"[OAuthHandler] Using storage method: {self.storage_method}") # Handle different storage methods success_message = None if self.storage_method == "env_file": # Store in .env file import os node_dir = os.path.dirname(__file__) env_path = os.path.join(node_dir, ".env") with open(env_path, "w") as f: f.write(f"DROPBOX_APP_KEY={self.app_key}\n") f.write(f"DROPBOX_APP_SECRET={self.app_secret}\n") f.write(f"DROPBOX_REFRESH_TOKEN={refresh_token}\n") f.write(f"DROPBOX_FOLDER={self.dropbox_folder}\n") success_message = "Dropbox connected successfully! Credentials saved to .env file." elif self.storage_method == "display_only": # Display credentials for manual copying and create completion marker import os node_dir = os.path.dirname(__file__) display_marker_path = os.path.join(node_dir, ".dropbox_display_complete") with open(display_marker_path, "w") as f: f.write("display_only_setup_completed") # Print credentials to console for easy copy/paste print("\n" + "=" * 80) print("DROPBOX CREDENTIALS READY FOR PRODUCTION - COPY FROM CONSOLE") print("=" * 80) print(f"DROPBOX_APP_KEY={self.app_key}") print(f"DROPBOX_APP_SECRET={self.app_secret}") print(f"DROPBOX_REFRESH_TOKEN={refresh_token}") print(f"DROPBOX_FOLDER={self.dropbox_folder}") print("=" * 80) print("Copy the lines above to your environment variables!") print("Perfect for RunPod, Docker, and production environments!") print("=" * 80 + "\n") success_message = f"""Dropbox Connected Successfully! ===================================================================== ENVIRONMENT VARIABLES - Copy & Paste Ready ===================================================================== DROPBOX_APP_KEY={self.app_key} DROPBOX_APP_SECRET={self.app_secret} DROPBOX_REFRESH_TOKEN={refresh_token} DROPBOX_FOLDER={self.dropbox_folder} ===================================================================== Perfect for RunPod, Docker, and production environments! These credentials are ready to use immediately. =====================================================================""" else: # Fallback to env_file import os node_dir = os.path.dirname(__file__) env_path = os.path.join(node_dir, ".env") with open(env_path, "w") as f: f.write(f"DROPBOX_APP_KEY={self.app_key}\n") f.write(f"DROPBOX_APP_SECRET={self.app_secret}\n") f.write(f"DROPBOX_REFRESH_TOKEN={refresh_token}\n") f.write(f"DROPBOX_FOLDER={self.dropbox_folder}\n") success_message = "Dropbox connected successfully! Credentials saved to .env file." # Send WebSocket message to trigger ComfyUI refresh (except for display_only) if self.storage_method != "display_only": try: from server import PromptServer message_data = { "type": "dropbox_oauth_complete", "session_id": self.session_id, "success": True, "message": success_message } PromptServer.instance.send_sync("dropbox_oauth_complete", message_data) print(f"[OAuthHandler] Sent WebSocket notification for OAuth completion") except Exception as e: print(f"[OAuthHandler] Warning: Could not send WebSocket notification: {e}") else: print(f"[OAuthHandler] Skipping WebSocket notification for display_only method - no auto-refresh") # Call completion callback if provided if self.completion_callback: await self.completion_callback(True, success_message) return True, success_message except Exception as e: error_msg = f"Error: 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 def create_display_only_success_page(session_id, app_key, app_secret, refresh_token, dropbox_folder): """Create a success page with textarea fields for display_only storage method""" from aiohttp import web return web.Response(text=f""" Dropbox Authorization Successful
✅

Dropbox Connected Successfully!

Your credentials are ready for production use. Copy them to your environment variables.

🔐 Environment Variables

📋 How to Use These Credentials

  1. Copy each environment variable above
  2. Add them to your production environment (RunPod, Docker, etc.)
  3. The ComfyUI DropSend node will detect and use them automatically
  4. No additional setup required!
Copy the credentials above, then manually refresh ComfyUI to hide the auth fields. Close this window when done.
""", content_type='text/html') def get_server_base_url(request=None): """Dynamically detect the server base URL for OAuth callbacks""" import os import socket try: if request: # Try to get from the incoming request scheme = 'https' if request.secure else 'http' host = request.host base_url = f"{scheme}://{host}" print(f"[OAuthHandler] Detected server URL from request: {base_url}") return base_url else: # Fallback: try to detect from ComfyUI server instance from server import PromptServer if hasattr(PromptServer, 'instance') and PromptServer.instance: # Get the port from ComfyUI server port = getattr(PromptServer.instance, 'port', 8188) # Try to detect if we're on RunPod or similar cloud service hostname = socket.gethostname() print(f"[OAuthHandler] Detected hostname: {hostname}") # Check for RunPod-specific hostname patterns if 'runpod' in hostname.lower() or hostname.endswith('.runpod.net'): # For RunPod, try to get the public URL runpod_public_ip = os.getenv('RUNPOD_PUBLIC_IP') if runpod_public_ip: base_url = f"https://{runpod_public_ip}:{port}" print(f"[OAuthHandler] Detected RunPod URL: {base_url}") return base_url # Check for common cloud service environment variables public_url = os.getenv('COMFYUI_PUBLIC_URL') or os.getenv('PUBLIC_URL') if public_url: base_url = public_url.rstrip('/') print(f"[OAuthHandler] Using configured PUBLIC_URL: {base_url}") return base_url # Default fallback to localhost base_url = f"http://localhost:{port}" print(f"[OAuthHandler] Fallback to localhost: {base_url}") return base_url except Exception as e: print(f"[OAuthHandler] Error detecting server URL: {e}") import traceback traceback.print_exc() # Final fallback base_url = "http://localhost:8188" print(f"[OAuthHandler] Using final fallback: {base_url}") return base_url # 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'] # Use the original redirect URI that was used for authorization (not the detected one) original_callback_url = session_data.get('original_redirect_uri') if original_callback_url: print(f"[OAuthCallback] Using stored redirect URI: {original_callback_url}") callback_url = original_callback_url else: # Fallback to dynamic detection base_url = get_server_base_url(request) callback_url = f"{base_url}/oauth/dropbox/callback" print(f"[OAuthCallback] Using detected redirect URI: {callback_url}") success, message = await handler.process_callback(auth_code, redirect_uri=callback_url) # Clean up the session del pending_oauth_sessions[session_id] if success: # Check if this is a display_only flow and show appropriate success page storage_method = session_data.get('storage_method', 'keyring') if storage_method == 'display_only': # Extract refresh token from the success message (token exchange already happened in process_callback) if 'DROPBOX_REFRESH_TOKEN=' in message: import re refresh_token_match = re.search(r'DROPBOX_REFRESH_TOKEN=([^\s\n]+)', message) if refresh_token_match: refresh_token = refresh_token_match.group(1) app_key = session_data['app_key'] app_secret = session_data['app_secret'] dropbox_folder = session_data['dropbox_folder'] # Special success page with textarea fields for credentials (stays open for copying) return create_display_only_success_page(session_id, app_key, app_secret, refresh_token, dropbox_folder) # If we couldn't extract the refresh token, fall back to error handling return web.Response(text=f"""

❌ Display Setup Error

Could not extract refresh token from OAuth response.

Please close this window and try again in ComfyUI.

Debug info: {message[:200]}...

""", content_type='text/html', status=500) else: # Standard success page for env_file storage method (auto-closes quickly) return web.Response(text=f""" Dropbox Authorization Successful
✓

Authorization Successful!

{message}

""", content_type='text/html') else: return web.Response(text=f""" Dropbox Authorization Failed
❌

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""" Dropbox Authorization Code
🔄

Authorization Code Received

Please return to ComfyUI and manually enter this code:

{auth_code}

📋 Select the code above and copy it to ComfyUI

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