""" DriveSend AutoUploader Node Monitors output folder and uploads new files to Google Drive """ import os from pathlib import Path from .monitor_output import start_monitor, stop_monitor, get_monitor class DriveSendAutoUploaderNode: """ComfyUI node for automatic Google Drive uploads.""" CATEGORY = "DriveSend" FUNCTION = "run" RETURN_TYPES = ("STRING",) RETURN_NAMES = ("status",) OUTPUT_NODE = True @classmethod def INPUT_TYPES(cls): return { "required": { "auth_method": (["oauth", "service_account"], {"default": "oauth"}), "run_process": ("BOOLEAN", {"default": True}), "enable_encryption": ("BOOLEAN", {"default": False}), "Subfolder_Monitor": ("BOOLEAN", {"default": True}), "Post_Delete_Enc": ("BOOLEAN", {"default": False}), }, "optional": { "watch_folder": ("STRING", { "default": "", "multiline": False, "tooltip": "Folder to monitor (leave blank for ComfyUI output folder)" }), } } def run(self, auth_method, run_process, enable_encryption, Subfolder_Monitor, Post_Delete_Enc, watch_folder=""): """ Start or stop the Google Drive upload monitor. Args: auth_method: 'oauth' or 'service_account' run_process: True to start, False to stop enable_encryption: Encrypt files before upload Subfolder_Monitor: Monitor subfolders recursively Post_Delete_Enc: Delete .enc files after upload watch_folder: Folder to monitor (default: ComfyUI output) """ # Stop if requested if not run_process: monitor = get_monitor() if monitor and monitor.is_running(): stop_monitor() return ("⏹️ DriveSend monitor stopped.",) else: return ("ℹ️ Monitor was not running.",) # Determine watch folder if watch_folder: watch_path = Path(watch_folder) else: # Default to ComfyUI output folder # Try common locations possible_paths = [ Path("/workspace/ComfyUI/output"), # RunPod Path("./output"), # Relative Path(__file__).parent.parent.parent / "output", # Relative to node ] watch_path = None for p in possible_paths: if p.exists(): watch_path = p break if not watch_path: return ("❌ Error: Could not find ComfyUI output folder. Please specify watch_folder.",) if not watch_path.exists(): return (f"❌ Error: Watch folder does not exist: {watch_path}",) # Get folder ID from environment effective_folder_id = os.environ.get('GOOGLE_DRIVE_FOLDER_ID') if not effective_folder_id: return ( "❌ Error: GOOGLE_DRIVE_FOLDER_ID not set.\n\n" "Run the DriveSend Setup node first, then:\n" "1. Copy credentials to RunPod Secrets\n" "2. Restart pod with environment variables loaded", ) # Check for encryption key if encryption is enabled if enable_encryption: enc_key = ( os.environ.get('COMFYUI_ENCRYPTION_KEY') or os.environ.get('comfyui_encryption_key') or os.environ.get('DROPSEND_ENCRYPTION_KEY') or os.environ.get('DRIVESEND_ENCRYPTION_KEY') or os.environ.get('RUNPOD_SECRET_COMFYUI_ENCRYPTION_KEY') ) if not enc_key: return ( "❌ Error: Encryption enabled but no key found.\n\n" "Set COMFYUI_ENCRYPTION_KEY environment variable\n" "or run DriveSend Setup with encryption enabled.", ) # Validate auth method requirements if auth_method == 'oauth': client_id = os.environ.get('GOOGLE_CLIENT_ID') client_secret = os.environ.get('GOOGLE_CLIENT_SECRET') refresh_token = os.environ.get('GOOGLE_REFRESH_TOKEN') if not all([client_id, client_secret, refresh_token]): # Check for token.json file as fallback token_file = Path(__file__).parent / 'token.json' if not token_file.exists(): return ( "❌ Error: OAuth credentials not found.\n\n" "Set these environment variables:\n" "- GOOGLE_CLIENT_ID\n" "- GOOGLE_CLIENT_SECRET\n" "- GOOGLE_REFRESH_TOKEN\n\n" "Or run DriveSend Setup to authorize.", ) elif auth_method == 'service_account': sa_json = os.environ.get('GOOGLE_SERVICE_ACCOUNT_JSON') sa_file = Path(__file__).parent / 'service_account.json' if not sa_json and not sa_file.exists(): return ( "❌ Error: Service account credentials not found.\n\n" "Either:\n" "1. Set GOOGLE_SERVICE_ACCOUNT_JSON environment variable (base64)\n" "2. Place service_account.json in the node folder\n\n" "⚠️ NOTE: Service accounts only work with Google Workspace (paid).\n" "For personal Gmail, use OAuth instead.", ) # Start the monitor try: monitor = start_monitor( watch_dir=str(watch_path), folder_id=effective_folder_id, recursive=Subfolder_Monitor, enable_encryption=enable_encryption, post_delete_enc=Post_Delete_Enc, auth_method=auth_method ) status_lines = [ "✅ DriveSend monitor started!", "", f"📁 Watching: {watch_path}", f"☁️ Uploading to: Google Drive folder {effective_folder_id[:20]}...", f"🔐 Auth method: {auth_method}", f"🔒 Encryption: {'Enabled' if enable_encryption else 'Disabled'}", f"📂 Subfolder monitoring: {'Enabled' if Subfolder_Monitor else 'Disabled'}", "", "New files will be uploaded automatically.", "Set run_process to False to stop.", ] return ("\n".join(status_lines),) except Exception as e: return (f"❌ Error starting monitor: {e}",) # Node registration NODE_CLASS_MAPPINGS = { "DriveSendAutoUploader": DriveSendAutoUploaderNode } NODE_DISPLAY_NAME_MAPPINGS = { "DriveSendAutoUploader": "DriveSend AutoUploader" }