diff --git a/__init__.py b/__init__.py index c440aa5..7e4364d 100644 --- a/__init__.py +++ b/__init__.py @@ -11,3 +11,23 @@ NODE_CLASS_MAPPINGS.update(SETUP_CLASS) NODE_DISPLAY_NAME_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS.update(RUN_DISPLAY) NODE_DISPLAY_NAME_MAPPINGS.update(SETUP_DISPLAY) + +# Define web directory for JavaScript extensions +WEB_DIRECTORY = "./web" + +# 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 new file mode 100644 index 0000000..e8151aa --- /dev/null +++ b/dropbox_auth_manager.py @@ -0,0 +1,248 @@ +# dropbox_auth_manager.py +import os +import requests +import json + +class DropboxAuthManager: + def __init__(self, app_key=None, app_secret=None): + # Always store provided credentials + self.app_key = app_key + self.app_secret = app_secret + self.refresh_token = None + + # Try to load from file storage if no credentials provided + if app_key is None or app_secret is None: + self._try_load_from_file() + + def _get_credentials_file_path(self): + """Get path to credentials file""" + plugin_dir = os.path.dirname(os.path.abspath(__file__)) + return os.path.join(plugin_dir, ".dropbox_credentials.json") + + def _try_load_from_file(self): + """Try to load credentials from JSON file""" + try: + credentials_path = self._get_credentials_file_path() + if os.path.exists(credentials_path): + with open(credentials_path, 'r') as f: + creds = json.load(f) + + if self.app_key is None: + self.app_key = creds.get("app_key") + if self.app_secret is None: + self.app_secret = creds.get("app_secret") + if self.refresh_token is None: + self.refresh_token = creds.get("refresh_token") + + print(f"[DropboxAuthManager] Loaded credentials from file: {credentials_path}") + + except Exception as e: + print(f"[DropboxAuthManager] Could not load from credentials file: {e}") + + def is_connected(self): + # Load from file if we don't have credentials yet + if not (self.app_key and self.app_secret and self.refresh_token): + self._try_load_from_file() + return bool(self.app_key and self.app_secret and self.refresh_token) + + def store_tokens(self, app_key, app_secret, refresh_token): + """Store tokens in JSON file""" + try: + credentials_path = self._get_credentials_file_path() + + credentials = { + "app_key": app_key, + "app_secret": app_secret, + "refresh_token": refresh_token + } + + print(f"[DropboxAuthManager] Storing tokens in file: {credentials_path}") + with open(credentials_path, 'w') as f: + json.dump(credentials, f, indent=2) + + # Set appropriate file permissions (readable only by owner) + os.chmod(credentials_path, 0o600) + + # Update instance variables after successful storage + self.app_key = app_key + self.app_secret = app_secret + self.refresh_token = refresh_token + print(f"[DropboxAuthManager] Tokens stored successfully in credentials file") + + except Exception as e: + error_msg = f"Failed to store Dropbox tokens: {e}" + print(f"[DropboxAuthManager] {error_msg}") + raise RuntimeError(error_msg) + + 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.") + + data = { + "code": auth_code, + "grant_type": "authorization_code", + "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() + + creds = response.json() + refresh_token = creds.get("refresh_token") + if not refresh_token: + raise RuntimeError("Failed to obtain refresh token from Dropbox.") + + self.store_tokens(self.app_key, self.app_secret, refresh_token) + self.refresh_token = refresh_token + return True + + def exchange_auth_code_raw(self, auth_code, redirect_uri=None): + """Exchange auth code for tokens without storing them""" + if not (self.app_key and self.app_secret): + raise ValueError("App key and secret must be set before exchanging auth code.") + + data = { + "code": auth_code, + "grant_type": "authorization_code", + "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"} + + print(f"[DropboxAuthManager] Token exchange request data: {data}") + print(f"[DropboxAuthManager] Making request to: https://api.dropbox.com/oauth2/token") + + response = requests.post("https://api.dropbox.com/oauth2/token", headers=headers, data=data) + + print(f"[DropboxAuthManager] Response status: {response.status_code}") + print(f"[DropboxAuthManager] Response text: {response.text}") + + # If we get a 400 error and we're using a callback URL, try without it + # This handles cases where the auth code was obtained without a callback URL + if response.status_code == 400 and redirect_uri: + print(f"[DropboxAuthManager] 400 error with callback URL, trying without callback URL...") + + data_without_redirect = { + "code": auth_code, + "grant_type": "authorization_code", + "client_id": self.app_key, + "client_secret": self.app_secret + } + + print(f"[DropboxAuthManager] Retry request data (no callback URL): {data_without_redirect}") + + response = requests.post("https://api.dropbox.com/oauth2/token", headers=headers, data=data_without_redirect) + + print(f"[DropboxAuthManager] Retry response status: {response.status_code}") + print(f"[DropboxAuthManager] Retry response text: {response.text}") + + response.raise_for_status() + + return response.json() + + def get_access_token(self): + if not self.is_connected(): + raise RuntimeError("Dropbox not connected. Please run authentication first.") + + data = { + "refresh_token": self.refresh_token, + "grant_type": "refresh_token", + "client_id": self.app_key, + "client_secret": self.app_secret + } + 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() + + access_token = response.json().get("access_token") + if not access_token: + raise RuntimeError("Failed to retrieve access token.") + return access_token + + def reset(self, revoke_token=True): + """Clear stored tokens and optionally revoke authorization with Dropbox""" + + if revoke_token: + # Load credentials from file first if we don't have them in memory + if not (self.refresh_token and self.app_key and self.app_secret): + print("[DropboxAuthManager] Loading credentials for token revocation...") + self._try_load_from_file() + + # First, revoke the token with Dropbox if requested and possible + if self.refresh_token and self.app_key and self.app_secret: + try: + print("[DropboxAuthManager] Revoking token with Dropbox...") + + # Get current access token to revoke + access_token = self.get_access_token() + + # Revoke the token with Dropbox + revoke_data = { + "token": access_token + } + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json" + } + + response = requests.post( + "https://api.dropboxapi.com/2/auth/token/revoke", + headers=headers, + json=revoke_data + ) + + if response.status_code == 200: + print("[DropboxAuthManager] Token successfully revoked with Dropbox") + else: + print(f"[DropboxAuthManager] Token revocation failed: {response.status_code} - {response.text}") + + except Exception as e: + print(f"[DropboxAuthManager] Could not revoke token with Dropbox: {e}") + # Continue with local cleanup even if revocation fails + else: + print("[DropboxAuthManager] Skipping token revocation (revoke_token=False)") + + # Clear credentials file + try: + credentials_path = self._get_credentials_file_path() + if os.path.exists(credentials_path): + os.remove(credentials_path) + print(f"[DropboxAuthManager] Removed credentials file: {credentials_path}") + except Exception as e: + print(f"[DropboxAuthManager] Could not remove credentials file: {e}") + + # Clear instance variables + self.app_key = None + self.app_secret = None + self.refresh_token = None + + def get_oauth_url(self, redirect_uri=None, state=None, force_reapprove=False): + """Generate OAuth URL for initial authorization""" + if not self.app_key: + raise ValueError("App key must be set to generate OAuth URL.") + + 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}" + + # Force reapproval to ensure user sees authorization screen + if force_reapprove: + url += "&force_reapprove=true" + + return url + diff --git a/generate_base64.py b/generate_base64.py deleted file mode 100644 index 05ed4a9..0000000 --- a/generate_base64.py +++ /dev/null @@ -1,10 +0,0 @@ -# ComfyUI_DropSendNode/generate_base64.py - -import base64 - -app_key = input("Enter your App Key: ") -app_secret = input("Enter your App Secret: ") - -credentials = f"{app_key}:{app_secret}" -encoded = base64.b64encode(credentials.encode()).decode() -print(f"\nBasic {encoded}") diff --git a/get_refresh_token.py b/get_refresh_token.py deleted file mode 100644 index 1c71e68..0000000 --- a/get_refresh_token.py +++ /dev/null @@ -1,34 +0,0 @@ -# ComfyUI_DropSendNode/get_refresh_token.py - -import base64 -import requests -from .runpod_utils import get_redirect_uri, get_dropbox_authorization_url - -APP_KEY = input("Enter your APP KEY: ") -APP_SECRET = input("Enter your APP SECRET: ") - -redirect_uri = get_redirect_uri() -print(f"Using redirect_uri: {redirect_uri}") -print(f"Authorization URL: {get_dropbox_authorization_url(APP_KEY)}") -print() - -AUTH_CODE = input("Paste the Dropbox authorization code: ") - -auth_header = base64.b64encode(f"{APP_KEY}:{APP_SECRET}".encode()).decode() -res = requests.post( - "https://api.dropbox.com/oauth2/token", - data={ - "code": AUTH_CODE, - "grant_type": "authorization_code", - "redirect_uri": redirect_uri - }, - headers={ - "Authorization": f"Basic {auth_header}", - "Content-Type": "application/x-www-form-urlencoded" - } -) - -if res.ok: - print("\nYour refresh token is:\n", res.json().get("refresh_token")) -else: - print("Error:", res.text) diff --git a/oauth_handler.py b/oauth_handler.py new file mode 100644 index 0000000..ded42c1 --- /dev/null +++ b/oauth_handler.py @@ -0,0 +1,766 @@ +# 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] Processing callback for auth code exchange") + + # 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 + import os + + # 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") + + # 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. +
  2. Add them to your production environment (RunPod, Docker, etc.)
    3. +
  3. The ComfyUI DropSend node will detect and use them automatically
    4. +
  4. No additional setup required!
    5. +
+
+ +
+ Copy the credentials above, then 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) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 24e2bf4..a41eb3f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ requests python-dotenv watchdog +dropbox diff --git a/setup_dropbox_node.py b/setup_dropbox_node.py index dd9cea3..ed86c81 100644 --- a/setup_dropbox_node.py +++ b/setup_dropbox_node.py @@ -2,18 +2,35 @@ import os import requests +import webbrowser +import uuid +import json from dotenv import load_dotenv, dotenv_values -from .runpod_utils import get_redirect_uri +import urllib.parse +from .dropbox_auth_manager import DropboxAuthManager +from .oauth_handler import OAuthCallbackHandler, get_server_base_url + class DropboxSetupNode: @classmethod def INPUT_TYPES(cls): + # Always show all fields for simplicity - no dynamic field hiding return { "required": { - "app_key": ("STRING", {"default": ""}), - "app_secret": ("STRING", {"default": ""}), - "auth_code": ("STRING", {"default": ""}), + "app_key": ("STRING", {"default": "", "multiline": False}), + "app_secret": ("STRING", {"default": "", "multiline": False}), + "auth_code": ("STRING", {"default": "", "multiline": False}), "dropbox_dest_folder": ("STRING", {"default": "/Apps/ComfyUI_Output_Files"}), + }, + "optional": { + "reconnect": ("BOOLEAN", { + "label": "Reset stored credentials", + "default": False + }), + "storage_method": (["env_file", "display_only"], { + "label": "Credential Storage Method", + "default": "env_file" + }) } } @@ -21,54 +38,210 @@ class DropboxSetupNode: OUTPUT_NODE = True FUNCTION = "setup" - def setup(self, app_key, app_secret, auth_code, dropbox_dest_folder): - node_dir = os.path.dirname(__file__) - env_path = os.path.join(node_dir, ".env") + def setup(self, dropbox_dest_folder, app_key=None, app_secret=None, auth_code=None, reconnect=False, storage_method="env_file"): + try: + print(f"[DropboxSetup] Called with:") + print(f" app_key: '{app_key}' (type: {type(app_key)}, bool: {bool(app_key)})") + print(f" app_secret: '{app_secret}' (type: {type(app_secret)}, bool: {bool(app_secret)})") + print(f" auth_code: '{auth_code}' (type: {type(auth_code)}, bool: {bool(auth_code)})") + print(f" reconnect: {reconnect}") + print(f" storage_method: {storage_method}") + + # Initialize auth manager + auth_manager = DropboxAuthManager() + print(f"[DropboxSetup] Auth manager initialized") + + # Handle reconnect/reset request + if reconnect: + print("[DropboxSetup] Reconnect requested - clearing all credentials") + + # Clear all credentials (skip token revocation to avoid delays) + print("[DropboxSetup] Clearing all stored credentials") + auth_manager.reset(revoke_token=False) + + # Manually clear all credential files + node_dir = os.path.dirname(__file__) + + # Clear .env file if it exists + env_path = os.path.join(node_dir, ".env") + if os.path.exists(env_path): + print(f"[DropboxSetup] Removing .env file: {env_path}") + os.remove(env_path) + + # Send WebSocket message to trigger ComfyUI refresh after clearing credentials + try: + from server import PromptServer + message_data = { + "type": "dropbox_reconnect_complete", + "success": True, + "message": "Credentials cleared - ComfyUI will refresh to show auth fields" + } + PromptServer.instance.send_sync("dropbox_reconnect_complete", message_data) + print(f"[DropboxSetup] Sent WebSocket notification for reconnect completion") + except Exception as e: + print(f"[DropboxSetup] Warning: Could not send WebSocket notification: {e}") + + message = "Dropbox credentials cleared. ComfyUI will refresh to show auth fields..." + print(f"[DropboxSetup] {message}") + return { + "ui": {"text": [message]}, + "result": (message,) + } + + # Check for environment variables (from any source) + env_vars_set = all([ + os.getenv("DROPBOX_APP_KEY"), + os.getenv("DROPBOX_APP_SECRET"), + os.getenv("DROPBOX_REFRESH_TOKEN") + ]) + if env_vars_set: + message = "Dropbox credentials found in system environment variables. Ready to upload files." + print(f"[DropboxSetup] {message}") + return { + "ui": {"text": [message]}, + "result": (message,) + } - # 1. Skip setup if generic environment variables already exist - general_env_set = all([ - os.getenv("DROPBOX_APP_KEY"), - os.getenv("DROPBOX_APP_SECRET"), - os.getenv("DROPBOX_REFRESH_TOKEN") - ]) - if general_env_set: - return ("⚠️ Dropbox credentials already set in system environment variables. Skipping .env creation.",) + # Check for RunPod secrets + runpod_env_set = all([ + os.getenv("RUNPOD_SECRET_DROPBOX_ACCESS_TOKEN"), + os.getenv("RUNPOD_SECRET_DROPBOX_REFRESH_TOKEN") + ]) + if runpod_env_set: + return ("Warning: Detected RunPod secrets. Using those instead.",) - # 2. Skip setup if RunPod secrets are defined - runpod_env_set = all([ - os.getenv("RUNPOD_SECRET_DROPBOX_ACCESS_TOKEN"), - os.getenv("RUNPOD_SECRET_DROPBOX_REFRESH_TOKEN") - ]) - if runpod_env_set: - return ("⚠️ Detected RunPod secrets (RUNPOD_SECRET_DROPBOX_ACCESS_TOKEN, RUNPOD_SECRET_DROPBOX_REFRESH_TOKEN). Skipping .env creation.",) + # New setup flow using DropboxAuthManager + print(f"[DropboxSetup] Starting new setup flow") + + # Clean up the inputs first + app_key_clean = app_key.strip() if app_key else "" + app_secret_clean = app_secret.strip() if app_secret else "" + auth_code_clean = auth_code.strip() if auth_code else "" + + print(f"[DropboxSetup] Cleaned inputs:") + print(f" app_key_clean: '{app_key_clean}' (len: {len(app_key_clean)})") + print(f" app_secret_clean: '{app_secret_clean}' (len: {len(app_secret_clean)})") + print(f" auth_code_clean: '{auth_code_clean}' (len: {len(auth_code_clean)})") + + # Check if we have app credentials + if not app_key_clean or not app_secret_clean: + message = "Error: Missing App Key or App Secret. Please provide both." + print(f"[DropboxSetup] {message}") + return (message,) + + # If no auth code, generate OAuth URL for manual code flow + if not auth_code_clean: + print(f"[DropboxSetup] No auth code provided - generating OAuth URL for manual flow") + auth_temp = DropboxAuthManager(app_key=app_key_clean) + + # Manual OAuth flow without redirect_uri (Dropbox will display the code) + oauth_url = auth_temp.get_oauth_url(force_reapprove=True) + + try: + print(f"[DropboxSetup] Setting up manual OAuth popup...") + # JavaScript will automatically open a popup window + message = f"Dropbox OAuth Ready!\n\nClick the link below to authorize with Dropbox:\n\n{oauth_url}\n\nA popup window will open. After authorization, Dropbox will show your auth code.\nCopy the code and paste it into the 'auth_code' field above, then run this node again." + print(f"[DropboxSetup] OAuth URL ready for popup: {oauth_url}") + except Exception as e: + print(f"[DropboxSetup] Error setting up OAuth: {e}") + message = f"Dropbox Authorization:\n\nPlease visit this URL to authorize:\n{oauth_url}\n\nAfter authorization, Dropbox will display your auth code. Copy it and paste into the 'auth_code' field above." + + print(f"[DropboxSetup] OAuth URL: {oauth_url}") + + # Use ComfyUI's dynamic return format for better UI integration + return { + "ui": {"text": [message]}, + "result": (message,) + } - # 3. Run auth_code flow if no env config is present - if not all([app_key, app_secret, auth_code]): - return ("❌ Missing App Key, Secret, or Authorization Code.",) + # Exchange auth code for refresh token using DropboxAuthManager + print(f"[DropboxSetup] Attempting to exchange auth code") + auth_manager_setup = DropboxAuthManager(app_key_clean, app_secret_clean) + + # Get the tokens without storing them yet + # Use manual OAuth flow (auth codes from manual copy/paste) + print(f"[DropboxSetup] Using manual OAuth flow") + result = auth_manager_setup.exchange_auth_code_raw(auth_code_clean) + refresh_token = result.get("refresh_token") + + print(f"[DropboxSetup] Auth code exchange successful") + print(f"[DropboxSetup] Using storage method: {storage_method}") + + # Handle different storage methods + if storage_method == "env_file": + # Store in .env file + 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={app_key_clean}\n") + f.write(f"DROPBOX_APP_SECRET={app_secret_clean}\n") + f.write(f"DROPBOX_REFRESH_TOKEN={refresh_token}\n") + f.write(f"DROPBOX_FOLDER={dropbox_dest_folder}\n") + message = "Dropbox connected successfully! Credentials saved to .env file." + + elif storage_method == "display_only": + # Display credentials for manual copying + + message = f"""Dropbox Connected Successfully! - data = { - "code": auth_code, - "grant_type": "authorization_code", - "client_id": app_key, - "client_secret": app_secret, - "redirect_uri": get_redirect_uri() - } - headers = {"Content-Type": "application/x-www-form-urlencoded"} +===================================================================== +ENVIRONMENT VARIABLES - Copy & Paste Ready +===================================================================== - resp = requests.post("https://api.dropbox.com/oauth2/token", headers=headers, data=data) - resp.raise_for_status() - refresh_token = resp.json().get("refresh_token") +DROPBOX_APP_KEY={app_key_clean} - # Write .env with fallback credentials - with open(env_path, "w") as f: - f.write(f"DROPBOX_APP_KEY={app_key}\n") - f.write(f"DROPBOX_APP_SECRET={app_secret}\n") - f.write(f"DROPBOX_REFRESH_TOKEN={refresh_token}\n") - f.write(f"DROPBOX_FOLDER={dropbox_dest_folder}\n") +DROPBOX_APP_SECRET={app_secret_clean} - load_dotenv(env_path, override=True) - return ("✅ .env file created with Dropbox credentials.",) +DROPBOX_REFRESH_TOKEN={refresh_token} + +DROPBOX_FOLDER={dropbox_dest_folder} + +===================================================================== +Perfect for RunPod, Docker, and production environments! +These credentials are ready to use immediately. +=====================================================================""" + + else: + # Fallback to env_file storage + 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={app_key_clean}\n") + f.write(f"DROPBOX_APP_SECRET={app_secret_clean}\n") + f.write(f"DROPBOX_REFRESH_TOKEN={refresh_token}\n") + f.write(f"DROPBOX_FOLDER={dropbox_dest_folder}\n") + message = "Dropbox connected successfully! Credentials saved to .env file." + + print(f"[DropboxSetup] {message}") + + # For display_only, ensure credentials are prominently shown in console + if storage_method == "display_only": + print("\n" + "="*80) + print("DROPBOX CREDENTIALS READY FOR PRODUCTION - COPY FROM CONSOLE") + print("="*80) + print(f"DROPBOX_APP_KEY={app_key_clean}") + print(f"DROPBOX_APP_SECRET={app_secret_clean}") + print(f"DROPBOX_REFRESH_TOKEN={refresh_token}") + print(f"DROPBOX_FOLDER={dropbox_dest_folder}") + print("="*80) + print("Copy the lines above to your environment variables!") + print("Perfect for RunPod, Docker, and production environments!") + print("="*80 + "\n") + + # Use ComfyUI's dynamic return format for better UI integration + return { + "ui": {"text": [message]}, + "result": (message,) + } + + except Exception as e: + message = f"Error: Setup failed: {e}" + print(f"[DropboxSetup] ERROR: {message}") + return { + "ui": {"text": [message]}, + "result": (message,) + } # Required mappings for ComfyUI NODE_CLASS_MAPPINGS = {"DropboxSetupNode": DropboxSetupNode} -NODE_DISPLAY_NAME_MAPPINGS = {"DropboxSetupNode": "📦⚙️ Dropbox AutoUploader Setup"} +NODE_DISPLAY_NAME_MAPPINGS = {"DropboxSetupNode": "Dropbox AutoUploader Setup"} diff --git a/tests/test_dropbox_auth.py b/tests/test_dropbox_auth.py new file mode 100644 index 0000000..e10e407 --- /dev/null +++ b/tests/test_dropbox_auth.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +""" +Simple non-interactive test script for DropboxAuthManager + +Usage: + python test_dropbox_auth_simple.py + +This script tests the DropboxAuthManager without requiring user input. +""" + +import sys +import os +# Add parent directory to path so we can import from the main module +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from dropbox_auth_manager import DropboxAuthManager + +def test_basic_functionality(): + """Test basic DropboxAuthManager functionality""" + print("=" * 50) + print("Testing DropboxAuthManager Basic Functionality") + print("=" * 50) + + try: + # Test initialization without credentials + print("\n1. Testing initialization without credentials...") + auth = DropboxAuthManager() + print(f"PASS: DropboxAuthManager initialized") + print(f" is_connected(): {auth.is_connected()}") + + # Test initialization with test credentials + print("\n2. Testing initialization with test credentials...") + test_auth = DropboxAuthManager(app_key="test_key_123", app_secret="test_secret_456") + print(f"PASS: DropboxAuthManager initialized with test credentials") + print(f" app_key: {test_auth.app_key}") + print(f" is_connected(): {test_auth.is_connected()}") + + # Test OAuth URL generation + print("\n3. Testing OAuth URL generation...") + oauth_url = test_auth.get_oauth_url() + print(f"PASS: Basic OAuth URL: {oauth_url[:80]}...") + + # Test OAuth URL with manual flow (no redirect_uri needed) + oauth_url_manual = test_auth.get_oauth_url(state="test-session-123", force_reapprove=True) + print(f"PASS: Manual OAuth URL: {oauth_url_manual[:100]}...") + + # Test that manual OAuth URL is properly formatted + if "response_type=code" in oauth_url_manual and "token_access_type=offline" in oauth_url_manual: + print("PASS: OAuth URL properly formatted for manual flow") + else: + print("FAIL: OAuth URL missing required parameters") + + # Test that state is included when provided + if "state=test-session-123" in oauth_url_manual: + print("PASS: state parameter properly included in OAuth URL") + else: + print("FAIL: state parameter not found in OAuth URL") + + print("\n" + "=" * 50) + print("PASS: All basic functionality tests passed!") + print("=" * 50) + + except Exception as e: + print(f"FAIL: Error during basic functionality test: {e}") + raise + +def test_error_handling(): + """Test error handling scenarios""" + print("\n" + "=" * 50) + print("Testing Error Handling") + print("=" * 50) + + try: + # Test OAuth URL generation without app_key + print("\n1. Testing OAuth URL without app_key...") + auth_no_key = DropboxAuthManager() + try: + auth_no_key.get_oauth_url() + print("FAIL: Should have raised ValueError") + except ValueError as e: + print(f"PASS: Correctly raised ValueError: {e}") + + # Test auth code exchange without credentials + print("\n2. Testing auth code exchange without credentials...") + try: + auth_no_key.exchange_auth_code_raw("test_code") + print("FAIL: Should have raised ValueError") + except ValueError as e: + print(f"PASS: Correctly raised ValueError: {e}") + + print("\n" + "=" * 50) + print("PASS: All error handling tests passed!") + print("=" * 50) + + except Exception as e: + print(f"FAIL: Unexpected error during error handling test: {e}") + raise + +def test_token_exchange_retry_logic(): + """Test the token exchange retry logic (without actual API calls)""" + print("\n" + "=" * 50) + print("Testing Token Exchange Retry Logic") + print("=" * 50) + + # This test just verifies the method exists and can be called with parameters + # It won't make actual API calls since we don't have real credentials + try: + test_auth = DropboxAuthManager(app_key="test_key", app_secret="test_secret") + + print("\n1. Testing exchange_auth_code_raw method signature...") + + # Test method exists and accepts parameters (will fail with API call, but that's expected) + try: + test_auth.exchange_auth_code_raw("fake_code") + except Exception as e: + if "400" in str(e) or "requests" in str(e) or "HTTPSConnectionPool" in str(e): + print(f"PASS: Method executed correctly (expected API failure): {type(e).__name__}") + else: + print(f"FAIL: Unexpected error type: {e}") + + # Test method with redirect_uri (for backward compatibility) + try: + test_auth.exchange_auth_code_raw("fake_code", redirect_uri="http://localhost:8188/oauth/dropbox/callback") + except Exception as e: + if "400" in str(e) or "requests" in str(e) or "HTTPSConnectionPool" in str(e): + print(f"PASS: Method with redirect_uri executed correctly (expected API failure): {type(e).__name__}") + else: + print(f"FAIL: Unexpected error type: {e}") + + print("\n" + "=" * 50) + print("PASS: Token exchange retry logic tests completed!") + print("=" * 50) + + except Exception as e: + print(f"FAIL: Error during retry logic test: {e}") + raise + +def main(): + """Run all tests""" + print("Starting DropboxAuthManager Tests") + print("This script tests the auth manager without making real API calls") + + try: + test_basic_functionality() + test_error_handling() + test_token_exchange_retry_logic() + + print("\nAll tests completed successfully!") + print("To test with real Dropbox credentials, use the interactive flow in ComfyUI") + + except Exception as e: + print(f"\nTest failed: {e}") + return 1 + + return 0 + +if __name__ == "__main__": + exit(main()) \ No newline at end of file diff --git a/web/dropbox_oauth.js b/web/dropbox_oauth.js new file mode 100644 index 0000000..2a15311 --- /dev/null +++ b/web/dropbox_oauth.js @@ -0,0 +1,194 @@ +// dropbox_oauth.js - ComfyUI extension for seamless Dropbox OAuth flow + +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +// Track OAuth popup windows +let oauthPopup = null; + +app.registerExtension({ + name: "DropSendNode.OAuthHandler", + + async setup() { + console.log("[DropSendNode] OAuth extension loaded"); + + // Listen for WebSocket messages about OAuth completion + api.addEventListener("dropbox_oauth_complete", (event) => { + console.log("[DropSendNode] Received OAuth completion notification via WebSocket", event.detail); + + const data = event.detail; + if (data.success) { + console.log("[DropSendNode] OAuth successful"); + this.handleOAuthSuccess(); + } else { + console.error("[DropSendNode] OAuth failed:", data.message); + } + }); + + // Listen for WebSocket messages about reconnect completion + api.addEventListener("dropbox_reconnect_complete", (event) => { + console.log("[DropSendNode] Received reconnect completion notification via WebSocket", event.detail); + + const data = event.detail; + if (data.success) { + console.log("[DropSendNode] Reconnect successful"); + this.handleReconnectSuccess(); + } + }); + + // Listen for postMessage from OAuth popup window + window.addEventListener("message", (event) => { + if (event.data && event.data.type === 'dropbox_oauth_complete') { + console.log("[DropSendNode] Received OAuth completion notification via postMessage", event.data); + + if (event.data.success) { + console.log("[DropSendNode] OAuth successful"); + this.handleOAuthSuccess(); + } else { + console.error("[DropSendNode] OAuth failed:", event.data.message); + } + } + }); + }, + + handleOAuthSuccess() { + // Don't auto-close popup - let user close manually after seeing success message + console.log("[DropSendNode] OAuth success detected - popup will stay open for user to close manually"); + + // Reset any reconnect checkboxes + this.resetReconnectFields(); + + // Note: Manual refresh required if user wants to update UI + console.log("[DropSendNode] OAuth success handled - manual refresh available"); + }, + + resetReconnectFields() { + try { + // Find all DropboxSetupNode instances and reset their reconnect fields + const nodes = app.graph.nodes.filter(node => node.type === "DropboxSetupNode"); + nodes.forEach(node => { + if (node.widgets) { + const reconnectWidget = node.widgets.find(w => w.name === "reconnect"); + if (reconnectWidget && reconnectWidget.value === true) { + console.log("[DropSendNode] Resetting reconnect field to false"); + reconnectWidget.value = false; + if (node.onWidgetChanged) { + node.onWidgetChanged("reconnect", false); + } + } + } + }); + } catch (e) { + console.log("[DropSendNode] Could not reset reconnect fields:", e); + } + }, + + handleReconnectSuccess() { + // For reconnect, we DO want auto-refresh to show the auth fields again + console.log("[DropSendNode] Reconnect success - refreshing interface to show auth fields"); + + // Reset any cached state that might prevent popup detection + this.resetReconnectFields(); + + // Refresh ComfyUI interface to show auth fields after reconnect + setTimeout(() => { + console.log("[DropSendNode] Refreshing ComfyUI interface after reconnect"); + window.location.reload(); + }, 1000); // Short delay for reconnect + }, + + // Override node execution to detect and handle OAuth URLs + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "DropboxSetupNode") { + console.log("[DropSendNode] Registered DropboxSetupNode enhancement"); + + // Add custom OAuth popup handling + const originalExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function(message) { + console.log("[DropSendNode] Node executed with message:", message); + + // Check if this execution includes OAuth URL generation + if (message && typeof message === 'object' && message.text && message.text[0]) { + const text = message.text[0]; + console.log("[DropSendNode] Message text content:", text.substring(0, 100) + "..."); + + if (text.includes('Dropbox OAuth Ready!')) { + console.log("[DropSendNode] Detected Dropbox OAuth ready message"); + + // Extract OAuth URL from the message + const urlMatch = text.match(/(https:\/\/www\.dropbox\.com\/oauth2\/authorize[^\s]+)/); + if (urlMatch) { + const oauthUrl = urlMatch[1]; + console.log("[DropSendNode] Extracted OAuth URL:", oauthUrl); + + // Auto-open the popup after a brief delay + setTimeout(() => { + console.log("[DropSendNode] Auto-opening OAuth popup..."); + window.openDropboxOAuth(oauthUrl); + }, 500); + } else { + console.log("[DropSendNode] No OAuth URL found in message"); + } + } else { + console.log("[DropSendNode] Message does not contain 'Dropbox OAuth Ready!' trigger"); + } + } else { + console.log("[DropSendNode] Message structure not recognized for OAuth detection"); + } + + if (originalExecuted) { + return originalExecuted.apply(this, arguments); + } + }; + } + } +}); + +// Utility function to open OAuth in a small popup window +window.openDropboxOAuth = function(url) { + console.log("[DropSendNode] Opening OAuth popup:", url); + + // Close existing popup if open + if (oauthPopup && !oauthPopup.closed) { + oauthPopup.close(); + } + + // Open small popup window + const popup = window.open( + url, + 'dropbox_oauth', + 'width=500,height=700,scrollbars=yes,resizable=yes,status=no,location=no,toolbar=no,menubar=no' + ); + + if (popup) { + oauthPopup = popup; + + // Focus the popup + popup.focus(); + + console.log("[DropSendNode] OAuth popup opened successfully"); + + // Monitor popup closure with more detailed logging + const checkClosed = setInterval(() => { + if (popup.closed) { + console.log("[DropSendNode] OAuth popup closed by user or completed"); + clearInterval(checkClosed); + oauthPopup = null; + } + }, 500); // Check more frequently + + // Also monitor popup URL changes (if possible) + try { + popup.addEventListener('beforeunload', () => { + console.log("[DropSendNode] OAuth popup is navigating"); + }); + } catch (e) { + console.log("[DropSendNode] Cannot monitor popup navigation (cross-origin)"); + } + + return popup; + } else { + console.error("[DropSendNode] Could not open OAuth popup - popup blocked?"); + return null; + } +}; \ No newline at end of file