diff --git a/.gitignore b/.gitignore index 69c5ec7..85b5779 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ __pycache__/ *.pyo *.pdb *.pyc +*.log # Virtual environments (optional, depending on setup) venv/ diff --git a/__init__.py b/__init__.py index 7e4364d..aba2fe7 100644 --- a/__init__.py +++ b/__init__.py @@ -1,33 +1,27 @@ # ComfyUI_DropSendNode/__init__.py # Register Dropbox Setup + Run-only nodes - -from .run_dropbox_node import NODE_CLASS_MAPPINGS as RUN_CLASS, NODE_DISPLAY_NAME_MAPPINGS as RUN_DISPLAY -from .setup_dropbox_node import NODE_CLASS_MAPPINGS as SETUP_CLASS, NODE_DISPLAY_NAME_MAPPINGS as SETUP_DISPLAY - +from .dropsend_uploader_node import NODE_CLASS_MAPPINGS as RUN_CLASS, NODE_DISPLAY_NAME_MAPPINGS as RUN_DISPLAY +from .dropsend_setup_node import NODE_CLASS_MAPPINGS as SETUP_CLASS, NODE_DISPLAY_NAME_MAPPINGS as SETUP_DISPLAY NODE_CLASS_MAPPINGS = {} NODE_CLASS_MAPPINGS.update(RUN_CLASS) 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() +register_oauth_routes() \ No newline at end of file diff --git a/dropbox_auth_manager.py b/dropbox_auth_manager.py index e8151aa..28a82c1 100644 --- a/dropbox_auth_manager.py +++ b/dropbox_auth_manager.py @@ -128,7 +128,7 @@ class DropboxAuthManager: 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 + # If 400 error, 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...") diff --git a/dropbox_upload.py b/dropbox_upload.py index 75b81ce..4a06030 100644 --- a/dropbox_upload.py +++ b/dropbox_upload.py @@ -1,20 +1,30 @@ +# dropbox_upload.py import os +import hashlib import dropbox import requests from dropbox.exceptions import ApiError +import logging + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler('dropsend.log'), + logging.StreamHandler() + ] +) +logger = logging.getLogger(__name__) def get_token(): - # 1. Prioritize direct access_token from environment variables access_token = os.getenv("DROPBOX_ACCESS_TOKEN") or os.getenv("RUNPOD_SECRET_DROPBOX_ACCESS_TOKEN") if access_token: return access_token - # 2. Try refresh_token flow from environment variables or RunPod secrets refresh_token = os.getenv("DROPBOX_REFRESH_TOKEN") or os.getenv("RUNPOD_SECRET_DROPBOX_REFRESH_TOKEN") client_id = os.getenv("DROPBOX_APP_KEY") or os.getenv("RUNPOD_SECRET_DROPBOX_APP_KEY") client_secret = os.getenv("DROPBOX_APP_SECRET") or os.getenv("RUNPOD_SECRET_DROPBOX_APP_SECRET") - # 3. Check if all required credentials are present missing = [] if not refresh_token: missing.append("DROPBOX_REFRESH_TOKEN or RUNPOD_SECRET_DROPBOX_REFRESH_TOKEN") @@ -23,9 +33,8 @@ def get_token(): if not client_secret: missing.append("DROPBOX_APP_SECRET or RUNPOD_SECRET_DROPBOX_APP_SECRET") if missing: - raise Exception(f"❌ Missing Dropbox credentials: {', '.join(missing)}") + raise Exception(f"Missing Dropbox credentials: {', '.join(missing)}") - # 4. Perform OAuth refresh token flow try: response = requests.post( "https://api.dropbox.com/oauth2/token", @@ -40,29 +49,40 @@ def get_token(): response.raise_for_status() return response.json()["access_token"] except requests.exceptions.RequestException as e: - raise Exception(f"❌ Failed to refresh Dropbox token: {str(e)}") + raise Exception(f"Failed to refresh Dropbox token: {str(e)}") def upload_to_dropbox(local_file_path, dropbox_folder="/ComfyUI_Output_Files"): try: - # Initialize Dropbox client access_token = get_token() dbx = dropbox.Dropbox(access_token) - # Ensure the Dropbox folder exists try: dbx.files_get_metadata(dropbox_folder) except ApiError as e: if e.error.is_path() and e.error.get_path().is_not_found(): dbx.files_create_folder_v2(dropbox_folder) - print(f"📁 Created Dropbox folder: {dropbox_folder}") + logger.info(f"📦📁 Created Dropbox folder: {dropbox_folder}") else: - raise Exception(f"❌ Failed to check/create Dropbox folder: {str(e)}") + raise Exception(f"Failed to check/create Dropbox folder: {str(e)}") - # Upload the file + # Compute local file hash + with open(local_file_path, "rb") as f: + local_hash = hashlib.sha256(f.read()).hexdigest() + + # Upload file file_name = os.path.basename(local_file_path) dropbox_path = f"{dropbox_folder}/{file_name}" with open(local_file_path, "rb") as f: dbx.files_upload(f.read(), dropbox_path, mode=dropbox.files.WriteMode.overwrite) - print(f"📦📤 Uploaded to Dropbox: {dropbox_path}") + + # Verify upload integrity + metadata, response = dbx.files_download(dropbox_path) + remote_hash = hashlib.sha256(response.content).hexdigest() + if local_hash != remote_hash: + raise Exception(f"Upload integrity check failed for {dropbox_path}") + + logger.info(f"📦✅ Upload Verified: {dropbox_path}") + return True # Indicate successful upload and verification except Exception as e: - raise Exception(f"❌ Upload failed: {str(e)}") + logger.error(f"Upload failed: {str(e)}") + raise Exception(f"Upload failed: {str(e)}") \ No newline at end of file diff --git a/setup_dropbox_node.py b/dropsend_setup_node.py similarity index 79% rename from setup_dropbox_node.py rename to dropsend_setup_node.py index ed86c81..3d4ecf4 100644 --- a/setup_dropbox_node.py +++ b/dropsend_setup_node.py @@ -9,7 +9,7 @@ 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 - +from cryptography.fernet import Fernet class DropboxSetupNode: @classmethod @@ -29,7 +29,11 @@ class DropboxSetupNode: }), "storage_method": (["env_file", "display_only"], { "label": "Credential Storage Method", - "default": "env_file" + "default": "display_only" + }), + "encryption_key_method": (["off", "Display Only", "save to .env"], { + "label": "Encryption Key Method", + "default": "Display Only" }) } } @@ -38,7 +42,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, storage_method="env_file"): + def setup(self, dropbox_dest_folder, app_key=None, app_secret=None, auth_code=None, reconnect=False, storage_method="display_only", encryption_key_method="Display Only"): try: print(f"[DropboxSetup] Called with:") print(f" app_key: '{app_key}' (type: {type(app_key)}, bool: {bool(app_key)})") @@ -46,6 +50,7 @@ class DropboxSetupNode: print(f" auth_code: '{auth_code}' (type: {type(auth_code)}, bool: {bool(auth_code)})") print(f" reconnect: {reconnect}") print(f" storage_method: {storage_method}") + print(f" encryption_key_method: {encryption_key_method}") # Initialize auth manager auth_manager = DropboxAuthManager() @@ -135,7 +140,7 @@ class DropboxSetupNode: 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) + oauth_url = auth_temp.get_oauth_url(require_reapprove=True) try: print(f"[DropboxSetup] Setting up manual OAuth popup...") @@ -144,7 +149,7 @@ class DropboxSetupNode: 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." + message = f"Dropbox Authorization:\n\nPlease visit this URL to authorize:\n\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}") @@ -167,67 +172,70 @@ class DropboxSetupNode: print(f"[DropboxSetup] Auth code exchange successful") print(f"[DropboxSetup] Using storage method: {storage_method}") - # Handle different storage methods + # Generate encryption key only if encryption_key_method is not "off" + encryption_key = None + if encryption_key_method != "off": + encryption_key = Fernet.generate_key().decode() + + # Prepare .env lines + env_lines = [] if storage_method == "env_file": - # Store in .env file + env_lines.append(f"DROPBOX_APP_KEY={app_key_clean}") + env_lines.append(f"DROPBOX_APP_SECRET={app_secret_clean}") + env_lines.append(f"DROPBOX_REFRESH_TOKEN={refresh_token}") + env_lines.append(f"DROPBOX_FOLDER={dropbox_dest_folder}") + if encryption_key_method == "save to .env" and encryption_key: + env_lines.append(f"COMFYUI_ENCRYPTION_KEY={encryption_key}") + + # Write to .env if there are lines to write + if env_lines: 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") + f.write("\n".join(env_lines) + "\n") + + # Prepare display lines + display_lines = [] + if storage_method == "display_only": + display_lines.append(f"DROPBOX_APP_KEY={app_key_clean}") + display_lines.append(f"DROPBOX_APP_SECRET={app_secret_clean}") + display_lines.append(f"DROPBOX_REFRESH_TOKEN={refresh_token}") + display_lines.append(f"DROPBOX_FOLDER={dropbox_dest_folder}") + if encryption_key_method == "Display Only" and encryption_key: + display_lines.append(f"COMFYUI_ENCRYPTION_KEY={encryption_key}") + + # Build message + message = "" + if storage_method == "env_file": message = "Dropbox connected successfully! Credentials saved to .env file." - elif storage_method == "display_only": - # Display credentials for manual copying - - message = f"""Dropbox Connected Successfully! + message = """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} +""" + "\n".join(display_lines) + """ ===================================================================== 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": + if display_lines: 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}") + for line in display_lines: + print(line) print("="*80) print("Copy the lines above to your environment variables!") print("Perfect for RunPod, Docker, and production environments!") print("="*80 + "\n") + print(f"[DropboxSetup] {message}") + # Use ComfyUI's dynamic return format for better UI integration return { "ui": {"text": [message]}, @@ -244,4 +252,4 @@ These credentials are ready to use immediately. # Required mappings for ComfyUI NODE_CLASS_MAPPINGS = {"DropboxSetupNode": DropboxSetupNode} -NODE_DISPLAY_NAME_MAPPINGS = {"DropboxSetupNode": "Dropbox AutoUploader Setup"} +NODE_DISPLAY_NAME_MAPPINGS = {"DropboxSetupNode": "📦⚙️ DropSend - Setup Node"} \ No newline at end of file diff --git a/dropsend_uploader_node.py b/dropsend_uploader_node.py new file mode 100644 index 0000000..a5f24b7 --- /dev/null +++ b/dropsend_uploader_node.py @@ -0,0 +1,168 @@ +# ComfyUI_DropSendNode/dropsend_uploader_node.py + +import os +import threading +import logging +import time +from dotenv import load_dotenv, dotenv_values +from watchdog.observers import Observer +from .monitor_output import start_monitoring, watcher_observer, stop_queue_processor +from .encrypt_file import FileEncryptHandler, ENCRYPT_EXTENSIONS, stop_queue_processor as stop_encrypt_queue_processor + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler('dropsend.log'), + logging.StreamHandler() + ] +) +logger = logging.getLogger(__name__) + +# Define encrypt_observer at module level +encrypt_observer = None + +class DropSendRunNode: + @classmethod + def INPUT_TYPES(cls): + node_dir = os.path.dirname(__file__) + env_path = os.path.join(node_dir, ".env") + + default_watch = os.path.join(os.getcwd(), "output") + default_dest = "/Apps/ComfyUI_Output_Files" + default_encrypt = False + default_delete_enc = False + default_subfolder_monitor = True + default_run_process = True + + if os.path.exists(env_path): + cfg = dotenv_values(env_path) + default_dest = cfg.get("DROPBOX_FOLDER", default_dest) + default_encrypt = cfg.get("ENABLE_ENCRYPTION", "False").lower() == "true" + default_delete_enc = cfg.get("POST_DELETE_ENC", "False").lower() == "true" + default_subfolder_monitor = cfg.get("SUBFOLDER_MONITOR", "True").lower() == "true" + default_run_process = cfg.get("RUN_PROCESS", "True").lower() == "true" + + return { + "required": { + "watch_folder": ("STRING", {"default": default_watch}), + "dropbox_dest_folder": ("STRING", {"default": default_dest}), + "enable_encryption": ("BOOLEAN", {"default": default_encrypt}), + "Post_Delete_Enc": ("BOOLEAN", {"default": default_delete_enc}), + "Subfolder_Monitor": ("BOOLEAN", {"default": default_subfolder_monitor}), + "run_process": ("BOOLEAN", {"default": default_run_process, "label": "Run Process"}), + } + } + + RETURN_TYPES = ("STRING",) + OUTPUT_NODE = True + FUNCTION = "start" + + def start(self, watch_folder, dropbox_dest_folder, enable_encryption, Post_Delete_Enc, Subfolder_Monitor, run_process): + global encrypt_observer + logger.info(f"Starting DropSendRunNode: watch_folder={watch_folder}, dropbox_dest_folder={dropbox_dest_folder}, encryption={enable_encryption}, Post_Delete_Enc={Post_Delete_Enc}, Subfolder_Monitor={Subfolder_Monitor}, run_process={run_process}") + + # Handle stopping the process if run_process is False + if not run_process: + logger.info("Stopping DropSendRunNode monitoring") + if watcher_observer and watcher_observer.is_alive(): + watcher_observer.stop() + watcher_observer.join() + logger.info("Upload watcher stopped") + if encrypt_observer and encrypt_observer.is_alive(): + encrypt_observer.stop() + encrypt_observer.join() + logger.info("Encryption watcher stopped") + # Stop queue processors + stop_queue_processor() # Stop upload queue processor + stop_encrypt_queue_processor() # Stop encryption queue processor + # Print confirmation message to console + stop_message = f""" +===================================================================== +📦🛑 DropSend - AutoUploader - RunProcess STOPPED +===================================================================== +All monitoring, uploading, and encryption processes for {watch_folder} have been stopped. +Set 'run_process' to True and run the node again to resume. +===================================================================== +""" + print(stop_message) + logger.info(stop_message) + return (f"All monitoring, uploading, and encryption stopped for {watch_folder}",) + + # Validate watch_folder + if not os.path.exists(watch_folder): + logger.error(f"Watch folder does not exist: {watch_folder}") + raise ValueError(f"Watch folder does not exist: {watch_folder}") + if not os.path.isdir(watch_folder): + logger.error(f"Watch folder is not a directory: {watch_folder}") + raise ValueError(f"Watch folder is not a directory: {watch_folder}") + + # Validate Dropbox credentials + from .dropbox_upload import get_token + try: + get_token() + except Exception as e: + logger.error(f"Invalid Dropbox credentials: {str(e)}") + raise ValueError(f"Invalid Dropbox credentials: {str(e)}") + + # Validate encryption key if enabled + if enable_encryption: + ENCRYPT_KEY = os.getenv("COMFYUI_ENCRYPTION_KEY") or os.getenv("RUNPOD_SECRET_COMFYUI_ENCRYPTION_KEY") + if not ENCRYPT_KEY: + logger.error("Encryption enabled but no encryption key found") + raise ValueError(f"Encryption enabled but no encryption key found in 'COMFYUI_ENCRYPTION_KEY' or 'RUNPOD_SECRET_COMFYUI_ENCRYPTION_KEY'") + + # Persist settings + node_dir = os.path.dirname(__file__) + env_path = os.path.join(node_dir, ".env") + + existing = {} + if os.path.exists(env_path): + existing = dotenv_values(env_path) + + existing["DROPBOX_FOLDER"] = dropbox_dest_folder + existing["ENABLE_ENCRYPTION"] = str(enable_encryption) + existing["POST_DELETE_ENC"] = str(Post_Delete_Enc) + existing["SUBFOLDER_MONITOR"] = str(Subfolder_Monitor) + existing["RUN_PROCESS"] = str(run_process) + with open(env_path, "w") as f: + for k, v in existing.items(): + f.write(f"{k}={v}\n") + + load_dotenv(env_path, override=True) + + # Start processes only if run_process is True + if run_process: + # Start folder watcher + thread = threading.Thread( + target=start_monitoring, + args=(watch_folder, dropbox_dest_folder, enable_encryption, Post_Delete_Enc, Subfolder_Monitor), + daemon=True + ) + thread.start() + + # Start encryption watcher if enabled + if enable_encryption: + encrypt_handler = FileEncryptHandler(watch_folder, False, Subfolder_Monitor) + encrypt_observer = Observer() + encrypt_observer.schedule(encrypt_handler, watch_folder, recursive=Subfolder_Monitor) + encrypt_observer.start() + logger.info(f"Starting encryption monitor for {watch_folder}") + + def keep_encrypt_alive(): + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + if encrypt_observer: + encrypt_observer.stop() + if encrypt_observer: + encrypt_observer.join() + + threading.Thread(target=keep_encrypt_alive, daemon=True).start() + + return (f"📦✅ .env updated: DROPBOX_FOLDER={dropbox_dest_folder}, ENABLE_ENCRYPTION={enable_encryption}, POST_DELETE_ENC={Post_Delete_Enc}, SUBFOLDER_MONITOR={Subfolder_Monitor}, RUN_PROCESS={run_process}",) + +# Required mappings for ComfyUI +NODE_CLASS_MAPPINGS = {"DropSendRunNode": DropSendRunNode} +NODE_DISPLAY_NAME_MAPPINGS = {"DropSendRunNode": "📦📤 DropSend - AutoUploader"} \ No newline at end of file diff --git a/encrypt_file.py b/encrypt_file.py new file mode 100644 index 0000000..a980790 --- /dev/null +++ b/encrypt_file.py @@ -0,0 +1,94 @@ +# ComfyUI_DropSendNode/encrypt_file.py +import os +import time +import logging +from queue import Queue +from watchdog.observers import Observer +from watchdog.events import FileSystemEventHandler +from cryptography.fernet import Fernet +import threading + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler('dropsend.log'), + logging.StreamHandler() + ] +) +logger = logging.getLogger(__name__) + +ENCRYPT_EXTENSIONS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.mp4', '.avi', '.mov') +_stop_queue_processor = False # Stop signal for queue processor + +def stop_queue_processor(): + """Signal the queue processor to stop.""" + global _stop_queue_processor + _stop_queue_processor = True + logger.info("Signaled encryption queue processor to stop") + +class FileEncryptHandler(FileSystemEventHandler): + def __init__(self, watch_dir="/workspace/ComfyUI/output", delete_original=False, subfolder_monitor=True): + self.watch_dir = watch_dir + self.delete_original = delete_original + self.subfolder_monitor = subfolder_monitor + self.file_queue = Queue() + self.start_queue_processor() + + def on_created(self, event): + if not event.is_directory and event.src_path.lower().endswith(ENCRYPT_EXTENSIONS): + logger.info(f"Detected new file to encrypt: {event.src_path}") + self.file_queue.put(event.src_path) + + def start_queue_processor(self): + def process_queue(): + while not _stop_queue_processor: + if not self.file_queue.empty(): + file_path = self.file_queue.get() + try: + time.sleep(1) # Ensure file is fully written + size1 = os.path.getsize(file_path) + time.sleep(0.1) + size2 = os.path.getsize(file_path) + if size1 != size2: + logger.warning(f"File {file_path} still writing, requeueing...") + self.file_queue.put(file_path) + continue + + ENCRYPT_KEY = os.getenv("COMFYUI_ENCRYPTION_KEY") or os.getenv("RUNPOD_SECRET_COMFYUI_ENCRYPTION_KEY") + if not ENCRYPT_KEY: + logger.error("Encryption key not found in 'COMFYUI_ENCRYPTION_KEY' or 'RUNPOD_SECRET_COMFYUI_ENCRYPTION_KEY'") + continue + fernet = Fernet(ENCRYPT_KEY.encode()) + + with open(file_path, 'rb') as f: + file_data = f.read() + + encrypted_data = fernet.encrypt(file_data) + base_name = os.path.splitext(file_path)[0] + enc_path = base_name + '.enc' + with open(enc_path, 'wb') as f: + f.write(encrypted_data) + logger.info(f"Created encrypted file: {enc_path}") + + # Defer deletion until after upload verification (handled in monitor_output.py) + except Exception as e: + logger.error(f"Error encrypting {file_path}: {e}") + time.sleep(0.1) # Small delay to avoid CPU overload + logger.info("Encryption queue processor stopped") + + threading.Thread(target=process_queue, daemon=True).start() + +if __name__ == "__main__": + watch_dir = os.getenv("WATCH_DIR", "/workspace/ComfyUI/output") + os.makedirs(watch_dir, exist_ok=True) + observer = Observer() + observer.schedule(FileEncryptHandler(watch_dir), path=watch_dir, recursive=True) + observer.start() + logger.info(f"Monitoring {watch_dir} for new files to encrypt...") + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + observer.stop() + observer.join() \ No newline at end of file diff --git a/monitor_output.py b/monitor_output.py index b89a001..9c27f93 100644 --- a/monitor_output.py +++ b/monitor_output.py @@ -3,14 +3,27 @@ import os import time import threading +import logging +from queue import Queue from watchdog.observers import Observer from watchdog.events import FileSystemEventHandler from .dropbox_upload import upload_to_dropbox +from .encrypt_file import ENCRYPT_EXTENSIONS + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler('dropsend.log'), + logging.StreamHandler() + ] +) +logger = logging.getLogger(__name__) -# Global reference to track and restart the monitor watcher_observer = None watcher_handler = None - +encryption_enabled = False +_stop_queue_processor = False # Stop signal for queue processor def wait_for_complete_write(file_path, timeout=10): last_size = -1 @@ -29,52 +42,76 @@ def wait_for_complete_write(file_path, timeout=10): last_size = current_size except FileNotFoundError: pass - time.sleep(0.5) checks += 1 - print(f"⚠️ Timeout: File may still be writing: {file_path}") + logger.warning(f"Timeout: File may still be writing: {file_path}") return False +def stop_queue_processor(): + """Signal the queue processor to stop.""" + global _stop_queue_processor + _stop_queue_processor = True + logger.info("Signaled upload queue processor to stop") class NewFileHandler(FileSystemEventHandler): - def __init__(self, dropbox_dest_folder): + def __init__(self, dropbox_dest_folder, delete_enc=False): super().__init__() self.dropbox_dest_folder = dropbox_dest_folder + self.delete_enc = delete_enc + self.file_queue = Queue() + self.start_queue_processor() def on_created(self, event): if event.is_directory: return - file_path = event.src_path - print(f"📦✅ Detected new file: {file_path}") - - if not wait_for_complete_write(file_path): - print(f"⚠️ Skipping upload: File not stable yet: {file_path}") + if encryption_enabled and not file_path.lower().endswith('.enc'): + logger.info(f"Skipping non-.enc file (encryption enabled): {file_path}") return + if not encryption_enabled and file_path.lower().endswith('.enc'): + logger.info(f"Skipping .enc file (encryption disabled): {file_path}") + return + logger.info(f"Detected new file: {file_path}") + self.file_queue.put(file_path) - try: - upload_to_dropbox(file_path, self.dropbox_dest_folder) - except Exception as e: - print(f"📦❌ Upload failed: {e}") + def start_queue_processor(self): + def process_queue(): + while not _stop_queue_processor: + if not self.file_queue.empty(): + file_path = self.file_queue.get() + if wait_for_complete_write(file_path): + try: + if upload_to_dropbox(file_path, self.dropbox_dest_folder): + if self.delete_enc and file_path.lower().endswith('.enc'): + os.remove(file_path) + logger.info(f"Deleted .enc file after upload: {file_path}") + except Exception as e: + logger.error(f"Upload failed: {e}") + self.file_queue.put(file_path) # Requeue on failure + time.sleep(0.1) # Small delay to avoid CPU overload + logger.info("Upload queue processor stopped") + threading.Thread(target=process_queue, daemon=True).start() -def start_monitoring(watch_folder, dropbox_dest_folder): - global watcher_observer, watcher_handler +def start_monitoring(watch_folder, dropbox_dest_folder, enable_encryption=False, delete_enc=False, subfolder_monitor=True): + global watcher_observer, watcher_handler, encryption_enabled, _stop_queue_processor + encryption_enabled = enable_encryption + _stop_queue_processor = False # Reset stop signal when starting - # Stop and restart watcher if already running if watcher_observer and watcher_observer.is_alive(): - print("🔁 Restarting monitor to update destination folder.") + logger.info("Restarting monitor to update settings.") watcher_observer.stop() watcher_observer.join() - watcher_handler = NewFileHandler(dropbox_dest_folder) + watcher_handler = NewFileHandler(dropbox_dest_folder, delete_enc) watcher_observer = Observer() - watcher_observer.schedule(watcher_handler, watch_folder, recursive=False) + watcher_observer.schedule(watcher_handler, watch_folder, recursive=subfolder_monitor) watcher_observer.start() - print(f"🔎🗂️ Now watching: {watch_folder}") - print(f"📦📥 Upload target: {dropbox_dest_folder}") + logger.info(f"Now watching: {watch_folder} (Subfolder_Monitor: {'Yes' if subfolder_monitor else 'No'})") + logger.info(f"Upload target: {dropbox_dest_folder}") + logger.info(f"Encryption enabled: {'Yes' if encryption_enabled else 'No'}") def keep_alive(): try: @@ -84,4 +121,4 @@ def start_monitoring(watch_folder, dropbox_dest_folder): watcher_observer.stop() watcher_observer.join() - threading.Thread(target=keep_alive, daemon=True).start() + threading.Thread(target=keep_alive, daemon=True).start() \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index a41eb3f..b412a52 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,4 @@ requests python-dotenv watchdog dropbox +cryptography diff --git a/run_dropbox_node.py b/run_dropbox_node.py deleted file mode 100644 index 7ce1383..0000000 --- a/run_dropbox_node.py +++ /dev/null @@ -1,64 +0,0 @@ -# ComfyUI_DropSendNode/run_dropbox_node.py - -import os -import threading -from dotenv import load_dotenv, dotenv_values -from .monitor_output import start_monitoring - -class DropSendRunNode: - @classmethod - def INPUT_TYPES(cls): - node_dir = os.path.dirname(__file__) - env_path = os.path.join(node_dir, ".env") - - default_watch = os.path.join(os.getcwd(), "output") - default_dest = "/Apps/ComfyUI_Output_Files" - - # if .env exists, pull last-used folder - if os.path.exists(env_path): - cfg = dotenv_values(env_path) - default_dest = cfg.get("DROPBOX_FOLDER", default_dest) - - return { - "required": { - "watch_folder": ("STRING", {"default": default_watch}), - "dropbox_dest_folder": ("STRING", {"default": default_dest}), - } - } - - RETURN_TYPES = ("STRING",) - OUTPUT_NODE = True - FUNCTION = "start" - - def start(self, watch_folder, dropbox_dest_folder): - # persist new destination into .env alongside secrets - node_dir = os.path.dirname(__file__) - env_path = os.path.join(node_dir, ".env") - - existing = {} - if os.path.exists(env_path): - existing = dotenv_values(env_path) - - existing["DROPBOX_FOLDER"] = dropbox_dest_folder - with open(env_path, "w") as f: - for k, v in existing.items(): - f.write(f"{k}={v}\n") - - load_dotenv(env_path, override=True) - - # start (or restart) the folder watcher - thread = threading.Thread( - target=start_monitoring, - args=(watch_folder, dropbox_dest_folder), - daemon=True - ) - thread.start() - - return ( - f"📦✅ .env updated: DROPBOX_FOLDER={dropbox_dest_folder}", - f"🔎📁 Monitoring Folder: {watch_folder} → {dropbox_dest_folder}" - ) - -# Required mappings for ComfyUI -NODE_CLASS_MAPPINGS = {"DropSendRunNode": DropSendRunNode} -NODE_DISPLAY_NAME_MAPPINGS = {"DropSendRunNode": "📦📤 Dropbox AutoUploader Run"}