395 lines
20 KiB
Python
395 lines
20 KiB
Python
# ComfyUI_DropSendNode/setup_dropbox_node.py
|
|
|
|
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, get_server_base_url
|
|
|
|
|
|
class DropboxSetupNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Check if credentials are already stored (keyring or .env file)
|
|
try:
|
|
auth_manager = DropboxAuthManager()
|
|
is_connected = auth_manager.is_connected()
|
|
|
|
# Check keyring availability for smart defaults (non-blocking)
|
|
keyring_available = True # Assume available until proven otherwise
|
|
try:
|
|
# Only test keyring if we need to know for sure
|
|
keyring_available = auth_manager._ensure_keyring_setup()
|
|
if not keyring_available:
|
|
print(f"[DropboxSetup] Keyring not available in this environment")
|
|
except Exception as e:
|
|
print(f"[DropboxSetup] Keyring test failed: {e}")
|
|
keyring_available = False
|
|
|
|
# Also check for .env file credentials
|
|
if not is_connected:
|
|
node_dir = os.path.dirname(__file__)
|
|
env_path = os.path.join(node_dir, ".env")
|
|
if os.path.exists(env_path):
|
|
from dotenv import dotenv_values
|
|
env_vars = dotenv_values(env_path)
|
|
if all([env_vars.get("DROPBOX_APP_KEY"), env_vars.get("DROPBOX_APP_SECRET"), env_vars.get("DROPBOX_REFRESH_TOKEN")]):
|
|
is_connected = True
|
|
print(f"[DropboxSetup] Found credentials in .env file")
|
|
|
|
# Also check for display_only completion marker - but only if environment variables actually exist
|
|
if not is_connected:
|
|
node_dir = os.path.dirname(__file__)
|
|
display_marker_path = os.path.join(node_dir, ".dropbox_display_complete")
|
|
if os.path.exists(display_marker_path):
|
|
# For display_only, verify that the user actually set up environment variables
|
|
display_only_env_set = all([
|
|
os.getenv("DROPBOX_APP_KEY"),
|
|
os.getenv("DROPBOX_APP_SECRET"),
|
|
os.getenv("DROPBOX_REFRESH_TOKEN")
|
|
])
|
|
if display_only_env_set:
|
|
is_connected = True
|
|
print(f"[DropboxSetup] Found display_only completion marker with valid environment variables")
|
|
else:
|
|
print(f"[DropboxSetup] Found display_only completion marker but environment variables not set yet")
|
|
|
|
print(f"[DropboxSetup] INPUT_TYPES check - is_connected: {is_connected}")
|
|
except Exception as e:
|
|
print(f"[DropboxSetup] INPUT_TYPES error: {e}")
|
|
is_connected = False
|
|
|
|
# Base inputs that are always shown
|
|
inputs = {
|
|
"required": {
|
|
"dropbox_dest_folder": ("STRING", {"default": "/Apps/ComfyUI_Output_Files"}),
|
|
},
|
|
"optional": {}
|
|
}
|
|
|
|
if is_connected:
|
|
# If connected, only show reconnect option
|
|
inputs["optional"]["reconnect"] = ("BOOLEAN", {
|
|
"label": "🔄 Reconnect Dropbox (clears credentials & refreshes UI)",
|
|
"default": False
|
|
})
|
|
else:
|
|
# If not connected, show setup fields
|
|
inputs["required"].update({
|
|
"app_key": ("STRING", {"default": "", "multiline": False}),
|
|
"app_secret": ("STRING", {"default": "", "multiline": False}),
|
|
"auth_code": ("STRING", {"default": "", "multiline": False}),
|
|
})
|
|
inputs["optional"]["reconnect"] = ("BOOLEAN", {
|
|
"label": "Reset stored credentials",
|
|
"default": False
|
|
})
|
|
inputs["optional"]["auto_oauth"] = ("BOOLEAN", {
|
|
"label": "Automatic OAuth (no manual auth code needed)",
|
|
"default": True
|
|
})
|
|
# Smart default based on keyring availability
|
|
if keyring_available:
|
|
default_storage = "keyring"
|
|
storage_options = ["keyring", "env_file", "display_only"]
|
|
else:
|
|
default_storage = "display_only" # Better for RunPod/Docker
|
|
storage_options = ["display_only", "env_file", "keyring"]
|
|
|
|
inputs["optional"]["storage_method"] = (storage_options, {
|
|
"label": "Credential Storage Method",
|
|
"default": default_storage
|
|
})
|
|
|
|
return inputs
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "setup"
|
|
|
|
def setup(self, dropbox_dest_folder, app_key=None, app_secret=None, auth_code=None, reconnect=False, auto_oauth=True, storage_method="keyring"):
|
|
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" auto_oauth: {auto_oauth}")
|
|
print(f" storage_method: {storage_method}")
|
|
|
|
# Initialize auth manager
|
|
auth_manager = DropboxAuthManager()
|
|
print(f"[DropboxSetup] Auth manager initialized. Is connected: {auth_manager.is_connected()}")
|
|
|
|
# Handle reconnect/reset request
|
|
if reconnect:
|
|
print("[DropboxSetup] Reconnect requested - clearing credentials")
|
|
|
|
# Clear keyring credentials
|
|
auth_manager.reset()
|
|
|
|
# Also clear .env file if it exists
|
|
node_dir = os.path.dirname(__file__)
|
|
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)
|
|
|
|
# Also clear display_only marker if it exists
|
|
display_marker_path = os.path.join(node_dir, ".dropbox_display_complete")
|
|
if os.path.exists(display_marker_path):
|
|
print(f"[DropboxSetup] Removing display_only marker: {display_marker_path}")
|
|
os.remove(display_marker_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": "🔄 All 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 from all storage locations. ComfyUI will refresh to show auth fields..."
|
|
print(f"[DropboxSetup] {message}")
|
|
return {
|
|
"ui": {"text": [message]},
|
|
"result": (message,)
|
|
}
|
|
|
|
# Check if already connected (keyring has credentials)
|
|
if auth_manager.is_connected():
|
|
print("[DropboxSetup] Already connected - testing stored credentials")
|
|
# Test the stored credentials by getting an access token
|
|
try:
|
|
access_token = auth_manager.get_access_token()
|
|
message = "✅ Dropbox already connected using stored keyring credentials. Ready to upload files."
|
|
print(f"[DropboxSetup] {message}")
|
|
return {
|
|
"ui": {"text": [message]},
|
|
"result": (message,)
|
|
}
|
|
except Exception as e:
|
|
message = f"⚠️ Stored credentials found but invalid: {e}. Use 'reconnect' to reset."
|
|
print(f"[DropboxSetup] {message}")
|
|
return {
|
|
"ui": {"text": [message]},
|
|
"result": (message,)
|
|
}
|
|
|
|
# Check for display_only environment variables (user set them up after OAuth flow)
|
|
display_only_env_set = all([
|
|
os.getenv("DROPBOX_APP_KEY"),
|
|
os.getenv("DROPBOX_APP_SECRET"),
|
|
os.getenv("DROPBOX_REFRESH_TOKEN")
|
|
])
|
|
if display_only_env_set:
|
|
# Check if this came from display_only flow
|
|
node_dir = os.path.dirname(__file__)
|
|
display_marker_path = os.path.join(node_dir, ".dropbox_display_complete")
|
|
if os.path.exists(display_marker_path):
|
|
message = "✅ Dropbox connected using environment variables (display_only setup). Ready to upload files."
|
|
print(f"[DropboxSetup] {message}")
|
|
return {
|
|
"ui": {"text": [message]},
|
|
"result": (message,)
|
|
}
|
|
else:
|
|
# Legacy environment variables (not from display_only)
|
|
message = "✅ Dropbox credentials found in system environment variables. Ready to upload files."
|
|
print(f"[DropboxSetup] {message}")
|
|
return {
|
|
"ui": {"text": [message]},
|
|
"result": (message,)
|
|
}
|
|
|
|
# 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 ("⚠️ Detected RunPod secrets. Using those instead of keyring.",)
|
|
|
|
# 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 = "❌ Missing App Key or App Secret. Please provide both."
|
|
print(f"[DropboxSetup] {message}")
|
|
return (message,)
|
|
|
|
# If no auth code, generate OAuth URL
|
|
if not auth_code_clean:
|
|
print(f"[DropboxSetup] No auth code provided - generating OAuth URL")
|
|
auth_temp = DropboxAuthManager(app_key=app_key_clean)
|
|
|
|
if auto_oauth:
|
|
# Automatic OAuth flow with callback
|
|
session_id = str(uuid.uuid4())
|
|
base_url = get_server_base_url()
|
|
print(f"[DropboxSetup] Detected base URL: {base_url}")
|
|
callback_url = f"{base_url}/oauth/dropbox/callback"
|
|
print(f"[DropboxSetup] Using callback URL: {callback_url}")
|
|
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, storage_method=storage_method, dropbox_folder=dropbox_dest_folder, original_redirect_uri=callback_url)
|
|
|
|
try:
|
|
print(f"[DropboxSetup] Setting up automatic OAuth popup...")
|
|
# We'll use JavaScript to open a proper popup window
|
|
message = f"🚀 Automatic OAuth Ready!\n\n🖱️ Click the link below to open a small popup window for authorization:\n\n🔗 {oauth_url}\n\n✨ After authorization, the popup will close and ComfyUI will refresh automatically!"
|
|
print(f"[DropboxSetup] Session ID: {session_id}")
|
|
print(f"[DropboxSetup] Callback URL: {callback_url}")
|
|
print(f"[DropboxSetup] OAuth URL ready for popup: {oauth_url}")
|
|
except Exception as e:
|
|
print(f"[DropboxSetup] Error setting up OAuth: {e}")
|
|
message = f"🔗 Automatic OAuth Setup:\n\nPlease visit this URL to authorize:\n{oauth_url}\n\nAfter authorization, ComfyUI will refresh 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}")
|
|
|
|
# Use ComfyUI's dynamic return format for better UI integration
|
|
return {
|
|
"ui": {"text": [message]},
|
|
"result": (message,)
|
|
}
|
|
|
|
# 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
|
|
# Only pass redirect_uri if this was an automatic OAuth flow
|
|
if auto_oauth:
|
|
base_url = get_server_base_url()
|
|
callback_url = f"{base_url}/oauth/dropbox/callback"
|
|
print(f"[DropboxSetup] Using automatic OAuth with redirect_uri: {callback_url}")
|
|
result = auth_manager_setup.exchange_auth_code_raw(auth_code_clean, redirect_uri=callback_url)
|
|
else:
|
|
print(f"[DropboxSetup] Using manual OAuth (no redirect_uri)")
|
|
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 == "keyring":
|
|
# Store in system keyring (original behavior)
|
|
try:
|
|
auth_manager_setup.store_tokens(app_key_clean, app_secret_clean, refresh_token)
|
|
message = "✅ Dropbox connected successfully! Credentials stored securely in system keyring."
|
|
except RuntimeError as e:
|
|
if "Keyring not available" in str(e):
|
|
message = f"❌ All keyring backends failed in this environment.\n\n💡 Please use 'env_file' or 'display_only' storage method instead.\n\n🔧 To enable file-based keyring, ensure 'keyrings.alt' is installed.\n\nTechnical details: {e}"
|
|
else:
|
|
raise e
|
|
|
|
elif 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 and create completion marker
|
|
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")
|
|
|
|
message = f"""✅ Dropbox Connected Successfully!
|
|
|
|
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
|
📋 ENVIRONMENT VARIABLES - Copy & Paste Ready
|
|
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
|
|
|
DROPBOX_APP_KEY={app_key_clean}
|
|
|
|
DROPBOX_APP_SECRET={app_secret_clean}
|
|
|
|
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 keyring
|
|
auth_manager_setup.store_tokens(app_key_clean, app_secret_clean, refresh_token)
|
|
message = "✅ Dropbox connected successfully! Credentials stored securely in system keyring."
|
|
|
|
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("💡 Note: ComfyUI won't auto-refresh - refresh manually after copying!")
|
|
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"❌ 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"}
|