diff --git a/.github/workflows/publish_action.yml b/.github/workflows/publish_action.yml new file mode 100644 index 0000000..be0b02f --- /dev/null +++ b/.github/workflows/publish_action.yml @@ -0,0 +1,20 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + paths: + - "pyproject.toml" + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/drivesend_setup_node.py b/drivesend_setup_node.py index 8c2166b..409e92a 100644 --- a/drivesend_setup_node.py +++ b/drivesend_setup_node.py @@ -277,5 +277,5 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { - "DriveSendSetup": "DriveSend Setup" + "DriveSendSetup": "🚙⚙️ DriveSend - Setup Node" } diff --git a/encrypt_file.py b/encrypt_file.py index f57f780..7966c45 100644 --- a/encrypt_file.py +++ b/encrypt_file.py @@ -1,12 +1,27 @@ -""" -File encryption module for DriveSend -Uses Fernet (AES-128) symmetric encryption -""" +# encrypt_file.py +# DriveSend encryption module - creates .enc copies of files (preserves originals for ComfyUI) import os -import hashlib -from pathlib import Path +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('drivesend.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 get_encryption_key(): @@ -33,151 +48,88 @@ def get_encryption_key(): ) -def encrypt_file(input_path, output_path=None, key=None): - """ - Encrypt a file using Fernet (AES-128). - - Args: - input_path: Path to the file to encrypt - output_path: Path for the encrypted file (default: input_path + '.enc') - key: Encryption key (default: from environment variable) - - Returns: - dict with encryption result: - - success: bool - - encrypted_path: str (if successful) - - checksum: str (SHA256 of original file) - - error: str (if failed) - """ - try: - input_path = Path(input_path) - - if not input_path.exists(): - return {"success": False, "error": f"File not found: {input_path}"} - - # Get encryption key - if key is None: - key = get_encryption_key() - - if not key: - return {"success": False, "error": "No encryption key found. Set COMFYUI_ENCRYPTION_KEY environment variable."} - - # Ensure key is bytes - if isinstance(key, str): - key = key.encode() - - # Create Fernet cipher - try: - fernet = Fernet(key) - except Exception as e: - return {"success": False, "error": f"Invalid encryption key: {e}"} - - # Read file - with open(input_path, 'rb') as f: - data = f.read() - - # Calculate checksum of original file - checksum = hashlib.sha256(data).hexdigest() - - # Encrypt - encrypted_data = fernet.encrypt(data) - - # Determine output path - if output_path is None: - output_path = Path(str(input_path) + '.enc') - else: - output_path = Path(output_path) - - # Write encrypted file - with open(output_path, 'wb') as f: - f.write(encrypted_data) - - return { - "success": True, - "encrypted_path": str(output_path), - "original_path": str(input_path), - "checksum": checksum, - "original_size": len(data), - "encrypted_size": len(encrypted_data) - } - - except Exception as e: - return {"success": False, "error": str(e)} +def stop_queue_processor(): + """Signal the queue processor to stop.""" + global _stop_queue_processor + _stop_queue_processor = True + logger.info("[DriveSend] Signaled encryption queue processor to stop") -def decrypt_file(input_path, output_path=None, key=None): +def reset_queue_processor(): + """Reset the stop signal for queue processor.""" + global _stop_queue_processor + _stop_queue_processor = False + + +class FileEncryptHandler(FileSystemEventHandler): """ - Decrypt a file using Fernet (AES-128). - - Args: - input_path: Path to the encrypted file - output_path: Path for the decrypted file (default: removes .enc extension) - key: Encryption key (default: from environment variable) - - Returns: - dict with decryption result: - - success: bool - - decrypted_path: str (if successful) - - checksum: str (SHA256 of decrypted file) - - error: str (if failed) + Watches for new image/video files and creates encrypted .enc copies. + Original files are preserved so ComfyUI maintains its sequential naming. """ + + 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"[DriveSend] 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"[DriveSend] File {file_path} still writing, requeueing...") + self.file_queue.put(file_path) + continue + + ENCRYPT_KEY = get_encryption_key() + if not ENCRYPT_KEY: + logger.error("[DriveSend] Encryption key not found. Set COMFYUI_ENCRYPTION_KEY environment variable.") + 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"[DriveSend] Created encrypted file: {enc_path}") + + # Original file is kept so ComfyUI maintains sequential naming + # User must manually delete originals from output folder + except Exception as e: + logger.error(f"[DriveSend] Error encrypting {file_path}: {e}") + time.sleep(0.1) # Small delay to avoid CPU overload + logger.info("[DriveSend] 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"[DriveSend] Monitoring {watch_dir} for new files to encrypt...") try: - input_path = Path(input_path) - - if not input_path.exists(): - return {"success": False, "error": f"File not found: {input_path}"} - - # Get encryption key - if key is None: - key = get_encryption_key() - - if not key: - return {"success": False, "error": "No encryption key found. Set COMFYUI_ENCRYPTION_KEY environment variable."} - - # Ensure key is bytes - if isinstance(key, str): - key = key.encode() - - # Create Fernet cipher - try: - fernet = Fernet(key) - except Exception as e: - return {"success": False, "error": f"Invalid encryption key: {e}"} - - # Read encrypted file - with open(input_path, 'rb') as f: - encrypted_data = f.read() - - # Decrypt - try: - decrypted_data = fernet.decrypt(encrypted_data) - except Exception as e: - return {"success": False, "error": f"Decryption failed (wrong key?): {e}"} - - # Calculate checksum of decrypted file - checksum = hashlib.sha256(decrypted_data).hexdigest() - - # Determine output path - if output_path is None: - # Remove .enc extension if present - if input_path.suffix == '.enc': - output_path = input_path.with_suffix('') - else: - output_path = Path(str(input_path) + '.decrypted') - else: - output_path = Path(output_path) - - # Write decrypted file - with open(output_path, 'wb') as f: - f.write(decrypted_data) - - return { - "success": True, - "decrypted_path": str(output_path), - "encrypted_path": str(input_path), - "checksum": checksum, - "decrypted_size": len(decrypted_data) - } - - except Exception as e: - return {"success": False, "error": str(e)} + while True: + time.sleep(1) + except KeyboardInterrupt: + observer.stop() + observer.join() diff --git a/monitor_output.py b/monitor_output.py index e75ad0b..07b5a6c 100644 --- a/monitor_output.py +++ b/monitor_output.py @@ -1,45 +1,35 @@ -""" -Output Monitor Module -Watches a directory for new files and triggers uploads -""" +# monitor_output.py +# DriveSend upload monitor - watches for files and uploads to Google Drive import os import time import threading -import queue import logging -from pathlib import Path +from queue import Queue from watchdog.observers import Observer from watchdog.events import FileSystemEventHandler - from .gdrive_upload import upload_file -from .encrypt_file import encrypt_file, get_encryption_key +from .encrypt_file import ENCRYPT_EXTENSIONS logging.basicConfig( level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s' + format='%(asctime)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler('drivesend.log'), + logging.StreamHandler() + ] ) logger = logging.getLogger(__name__) -# Supported file extensions -SUPPORTED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.webp', '.gif', '.mp4', '.avi', '.mov'} - -# Global stop signal for queue processor -_stop_queue_processor = False +watcher_observer = None +watcher_handler = None +encryption_enabled = False +_stop_queue_processor = False # Stop signal for queue processor +_drive_service = None # Cached Drive service def wait_for_complete_write(file_path, timeout=10): - """ - Wait for a file to be completely written before processing. - Checks if file size is stable for 2 consecutive checks (1 second). - - Args: - file_path: Path to the file - timeout: Maximum time to wait in seconds - - Returns: - bool: True if file is ready, False if timeout - """ + """Wait for a file to be completely written before processing.""" last_size = -1 stable_count = 0 checks = 0 @@ -49,7 +39,7 @@ def wait_for_complete_write(file_path, timeout=10): current_size = os.path.getsize(file_path) if current_size == last_size: stable_count += 1 - if stable_count >= 2: # File size stable for 2 checks (1 second) + if stable_count >= 2: return True else: stable_count = 0 @@ -59,239 +49,187 @@ def wait_for_complete_write(file_path, timeout=10): time.sleep(0.5) checks += 1 - logger.warning(f"Timeout: File may still be writing: {file_path}") + logger.warning(f"[DriveSend] Timeout: File may still be writing: {file_path}") return False -class UploadHandler(FileSystemEventHandler): - """Handles file system events and queues files for upload.""" +def stop_queue_processor(): + """Signal the queue processor to stop.""" + global _stop_queue_processor + _stop_queue_processor = True + logger.info("[DriveSend] Signaled upload queue processor to stop") + + +def reset_queue_processor(): + """Reset the stop signal for queue processor.""" + global _stop_queue_processor + _stop_queue_processor = False + + +class NewFileHandler(FileSystemEventHandler): + """ + Handles new file events and queues them for upload. + When encryption is enabled, only uploads .enc files. + When encryption is disabled, uploads image/video files directly. + """ - def __init__(self, upload_queue, recursive=True, enable_encryption=False): + def __init__(self, folder_id, auth_method='oauth', delete_enc=False): super().__init__() - self.upload_queue = upload_queue - self.recursive = recursive - self.enable_encryption = enable_encryption - self.processed_files = set() - + self.folder_id = folder_id + self.auth_method = auth_method + 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 - file_path = Path(event.src_path) - - # Handle encryption mode file filtering - if self.enable_encryption: - # When encryption is enabled, only process .enc files - if file_path.suffix.lower() != '.enc': - return - else: - # When encryption is disabled, skip .enc files and check supported extensions - if file_path.suffix.lower() == '.enc': - return - if file_path.suffix.lower() not in SUPPORTED_EXTENSIONS: - return - - # Avoid processing the same file multiple times - if str(file_path) in self.processed_files: + # When encryption is enabled, only process .enc files + if encryption_enabled and not file_path.lower().endswith('.enc'): + logger.debug(f"[DriveSend] Skipping non-.enc file (encryption enabled): {file_path}") return - self.processed_files.add(str(file_path)) + # When encryption is disabled, skip .enc files and only process supported extensions + if not encryption_enabled: + if file_path.lower().endswith('.enc'): + logger.debug(f"[DriveSend] Skipping .enc file (encryption disabled): {file_path}") + return + if not file_path.lower().endswith(ENCRYPT_EXTENSIONS): + logger.debug(f"[DriveSend] Skipping unsupported file type: {file_path}") + return - # Add to queue (wait_for_complete_write happens in worker) - self.upload_queue.put(str(file_path)) - print(f"[DriveSend] Queued: {file_path.name}") + logger.info(f"[DriveSend] Detected new file: {file_path}") + self.file_queue.put(file_path) - -class OutputMonitor: - """Monitors a directory and uploads new files to Google Drive.""" - - def __init__( - self, - watch_dir, - folder_id=None, - recursive=True, - enable_encryption=False, - post_delete_enc=False, - auth_method='service_account', - owner_email=None, - **auth_kwargs - ): - self.watch_dir = Path(watch_dir) - self.folder_id = folder_id or os.environ.get('GOOGLE_DRIVE_FOLDER_ID') - self.recursive = recursive - self.enable_encryption = enable_encryption - self.post_delete_enc = post_delete_enc - self.auth_method = auth_method - self.owner_email = owner_email or os.environ.get('GOOGLE_OWNER_EMAIL') - self.auth_kwargs = auth_kwargs - - self.upload_queue = queue.Queue() - self.observer = None - self.upload_thread = None - self.running = False - self.service = None - - def _upload_worker(self): - """Worker thread that processes the upload queue.""" - global _stop_queue_processor - from .gdrive_auth_manager import get_drive_service - - # Get authenticated service - try: - self.service = get_drive_service(self.auth_method, **self.auth_kwargs) - except Exception as e: - print(f"[DriveSend] Failed to authenticate: {e}") - return - - while self.running and not _stop_queue_processor: - try: - # Get file from queue with timeout - file_path = self.upload_queue.get(timeout=1) - except queue.Empty: - continue + def start_queue_processor(self): + def process_queue(): + global _drive_service - # Wait for file to be completely written - if not wait_for_complete_write(file_path): - print(f"[DriveSend] Warning: File may be incomplete: {file_path}") - # Still try to upload, but warn + # Initialize Drive service once + if _drive_service is None: + try: + from .gdrive_auth_manager import get_drive_service + _drive_service = get_drive_service(self.auth_method) + logger.info("[DriveSend] Drive service initialized") + except Exception as e: + logger.error(f"[DriveSend] Failed to initialize Drive service: {e}") + return - try: - file_to_upload = file_path - encrypted_file = None - - # Encrypt if enabled (and file is not already .enc) - if self.enable_encryption and not file_path.lower().endswith('.enc'): - key = get_encryption_key() - if key: - result = encrypt_file(file_path, key=key) - if result['success']: - file_to_upload = result['output_path'] - encrypted_file = file_to_upload - print(f"[DriveSend] Encrypted: {Path(file_path).name} -> {Path(file_to_upload).name}") - else: - print(f"[DriveSend] Encryption failed: {result.get('error')}") - else: - print("[DriveSend] Warning: Encryption enabled but no key found") - - # Upload file with owner_email for ownership transfer - result = upload_file( - file_to_upload, - folder_id=self.folder_id, - service=self.service, - auth_method=self.auth_method, - owner_email=self.owner_email - ) - - if result['success']: - print(f"[DriveSend] ✓ Uploaded: {Path(file_to_upload).name}") - - # Clean up encrypted file if requested - if encrypted_file and self.post_delete_enc: + 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: - os.remove(encrypted_file) - print(f"[DriveSend] Deleted encrypted file: {Path(encrypted_file).name}") + result = upload_file( + file_path, + folder_id=self.folder_id, + service=_drive_service + ) + if result.get('success'): + logger.info(f"[DriveSend] ✓ Uploaded: {os.path.basename(file_path)}") + # Delete .enc file after successful upload if requested + if self.delete_enc and file_path.lower().endswith('.enc'): + os.remove(file_path) + logger.info(f"[DriveSend] Deleted .enc file after upload: {file_path}") + else: + logger.error(f"[DriveSend] Upload failed: {result.get('error')}") + self.file_queue.put(file_path) # Requeue on failure except Exception as e: - print(f"[DriveSend] Failed to delete encrypted file: {e}") - else: - print(f"[DriveSend] ✗ Upload failed: {result.get('error')}") - # Requeue on failure - self.upload_queue.put(file_path) - - except Exception as e: - print(f"[DriveSend] Error processing {file_path}: {e}") - # Requeue on failure - self.upload_queue.put(file_path) - - finally: - self.upload_queue.task_done() + logger.error(f"[DriveSend] Upload failed: {e}") + self.file_queue.put(file_path) # Requeue on failure + time.sleep(0.1) # Small delay to avoid CPU overload + logger.info("[DriveSend] Upload queue processor stopped") + + threading.Thread(target=process_queue, daemon=True).start() + + +def start_monitoring(watch_folder, folder_id, auth_method='oauth', enable_encryption=False, + delete_enc=False, subfolder_monitor=True): + """ + Start monitoring a folder for new files to upload. - def start(self): - """Start monitoring the directory.""" - global _stop_queue_processor - - if self.running: - print("[DriveSend] Monitor already running") - return - - if not self.watch_dir.exists(): - print(f"[DriveSend] Watch directory does not exist: {self.watch_dir}") - return - - self.running = True - _stop_queue_processor = False # Reset stop signal - - # Start upload worker thread - self.upload_thread = threading.Thread(target=self._upload_worker, daemon=True) - self.upload_thread.start() - - # Set up file system observer - handler = UploadHandler(self.upload_queue, self.recursive, self.enable_encryption) - self.observer = Observer() - self.observer.schedule(handler, str(self.watch_dir), recursive=self.recursive) - self.observer.start() - - print(f"[DriveSend] Started monitoring: {self.watch_dir}") - if self.recursive: - print("[DriveSend] Recursive monitoring enabled") - if self.enable_encryption: - print("[DriveSend] Encryption enabled") - if self.owner_email: - print(f"[DriveSend] Ownership will transfer to: {self.owner_email}") - elif self.auth_method == 'service_account': - print("[DriveSend] WARNING: No owner_email set - uploads may fail due to quota!") + Args: + watch_folder: Directory to monitor + folder_id: Google Drive folder ID + auth_method: 'oauth' or 'service_account' + enable_encryption: If True, only upload .enc files (created by separate encrypt watcher) + delete_enc: Delete .enc files after successful upload + subfolder_monitor: Monitor subfolders recursively + """ + global watcher_observer, watcher_handler, encryption_enabled, _stop_queue_processor, _drive_service - def stop(self): - """Stop monitoring the directory.""" - global _stop_queue_processor - - if not self.running: - return - - self.running = False - _stop_queue_processor = True # Signal queue processor to stop - - if self.observer: - self.observer.stop() - self.observer.join(timeout=5) - self.observer = None - - if self.upload_thread: - self.upload_thread.join(timeout=5) - self.upload_thread = None - - print("[DriveSend] Stopped monitoring") + encryption_enabled = enable_encryption + _stop_queue_processor = False # Reset stop signal when starting + _drive_service = None # Reset service to force re-auth + + if watcher_observer and watcher_observer.is_alive(): + logger.info("[DriveSend] Restarting monitor to update settings.") + watcher_observer.stop() + watcher_observer.join() + + watcher_handler = NewFileHandler(folder_id, auth_method, delete_enc) + watcher_observer = Observer() + watcher_observer.schedule(watcher_handler, watch_folder, recursive=subfolder_monitor) + watcher_observer.start() + + logger.info(f"[DriveSend] Now watching: {watch_folder} (Subfolder_Monitor: {'Yes' if subfolder_monitor else 'No'})") + logger.info(f"[DriveSend] Upload target: Google Drive folder {folder_id}") + logger.info(f"[DriveSend] Encryption mode: {'Yes - uploading .enc files only' if encryption_enabled else 'No - uploading images directly'}") + + def keep_alive(): + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + watcher_observer.stop() + watcher_observer.join() + + threading.Thread(target=keep_alive, daemon=True).start() + + +def stop_monitoring(): + """Stop the upload monitor.""" + global watcher_observer, _drive_service - def is_running(self): - """Check if the monitor is running.""" - return self.running - - -# Global monitor instance -_monitor = None - - -def get_monitor(): - """Get the global monitor instance.""" - global _monitor - return _monitor + stop_queue_processor() + + if watcher_observer and watcher_observer.is_alive(): + watcher_observer.stop() + watcher_observer.join() + logger.info("[DriveSend] Upload monitor stopped") + + _drive_service = None # Clear cached service +# Legacy function names for compatibility def start_monitor(**kwargs): - """Start a new global monitor instance.""" - global _monitor + """Alias for start_monitoring.""" + # Map new param names to old function + watch_dir = kwargs.pop('watch_dir', kwargs.pop('watch_folder', None)) + recursive = kwargs.pop('recursive', kwargs.pop('subfolder_monitor', True)) + post_delete_enc = kwargs.pop('post_delete_enc', kwargs.pop('delete_enc', False)) - if _monitor and _monitor.is_running(): - _monitor.stop() - - _monitor = OutputMonitor(**kwargs) - _monitor.start() - return _monitor + return start_monitoring( + watch_folder=watch_dir, + subfolder_monitor=recursive, + delete_enc=post_delete_enc, + **kwargs + ) def stop_monitor(): - """Stop the global monitor instance.""" - global _monitor + """Alias for stop_monitoring.""" + return stop_monitoring() + + +def get_monitor(): + """Get the current watcher observer.""" + global watcher_observer - if _monitor: - _monitor.stop() - _monitor = None + class MonitorWrapper: + def is_running(self): + return watcher_observer and watcher_observer.is_alive() + + return MonitorWrapper() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..628040a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "comfyui-drivesend-node" +description = "Automatically upload ComfyUI output files to Google Drive with optional encryption. Supports OAuth and service accounts, real-time monitoring, cloud setups (RunPod), and includes cross-platform decryption scripts." +version = "1.0.0" +license = { file = "LICENSE" } +dependencies = ["requests", "python-dotenv", "watchdog", "cryptography", "google-auth", "google-auth-oauthlib", "google-api-python-client"] + +[project.urls] +Repository = "https://github.com/machinepainting/ComfyUI_DriveSendNode" + +[tool.comfy] +PublisherId = "machinepainting" +DisplayName = "DriveSend Node" +Icon = ""