machinepainting-ComfyUI_DriveSendNode/drivesend_uploader_node.py

# drivesend_uploader_node.py
# DriveSend AutoUploader Node - monitors output folder and uploads to Google Drive

import os
import threading
import logging
import time
from pathlib import Path
from dotenv import dotenv_values, load_dotenv
from watchdog.observers import Observer
from .monitor_output import start_monitoring, stop_monitoring, watcher_observer, stop_queue_processor
from .encrypt_file import FileEncryptHandler, ENCRYPT_EXTENSIONS, get_encryption_key
from .encrypt_file import stop_queue_processor as stop_encrypt_queue_processor
from .encrypt_file import reset_queue_processor as reset_encrypt_queue_processor
from .safe_paths import resolve_safe_watch_folder

# logging is configured centrally in __init__.py via a FileHandler on
# the root logger. Use a named logger here so records carry the module
# name into drivesend.log.
logger = logging.getLogger(__name__)

NODE_DIR = Path(__file__).parent
encrypt_observer = None


# ---------------------------------------------------------------------------
# Secure persistence for non-credential uploader settings
# ---------------------------------------------------------------------------
# The uploader persists watch_folder, auth_method, and behavior toggles
# to .env so the next session can defaults them in INPUT_TYPES. This is
# unconditional (not gated by COMFYUI_DRIVESEND_ALLOW_SETUP) because
# none of these values are credentials. Race-free 0o600 open keeps the
# file from being world/group-readable in the open->chmod window.

_UPLOADER_KEYS = {
    "DRIVESEND_WATCH_FOLDER",
    "DRIVESEND_AUTH_METHOD",
    "ENABLE_ENCRYPTION",
    "POST_DELETE_ENC",
    "SUBFOLDER_MONITOR",
    "RUN_PROCESS",
}


def _persist_uploader_settings(updates):
    """Merge non-credential uploader settings into .env with 0o600."""
    env_path = NODE_DIR / ".env"
    existing = {}
    if env_path.exists():
        existing = dict(dotenv_values(env_path))
    merged = dict(existing)
    for k, v in updates.items():
        merged[k] = str(v)
    if merged:
        fd = os.open(str(env_path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        with os.fdopen(fd, "w") as f:
            for k, v in merged.items():
                f.write(f"{k}={v}\n")
        load_dotenv(env_path, override=True)


class DriveSendAutoUploaderNode:
    """ComfyUI node for automatic Google Drive uploads."""

    CATEGORY = "DriveSend"
    FUNCTION = "start"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("status",)
    OUTPUT_NODE = True

    @classmethod
    def INPUT_TYPES(cls):
        # Default-from-.env so the user does not have to retype settings
        # across sessions.
        env_path = NODE_DIR / ".env"
        cfg = {}
        if env_path.exists():
            cfg = dict(dotenv_values(env_path))

        default_watch = cfg.get("DRIVESEND_WATCH_FOLDER") or os.path.join(os.getcwd(), "output")
        default_auth = cfg.get("DRIVESEND_AUTH_METHOD", "oauth")
        if default_auth not in ("oauth", "service_account"):
            default_auth = "oauth"
        default_encrypt = cfg.get("ENABLE_ENCRYPTION", "True").lower() == "true"
        default_delete_enc = cfg.get("POST_DELETE_ENC", "False").lower() == "true"
        default_subfolder = cfg.get("SUBFOLDER_MONITOR", "True").lower() == "true"
        default_run = cfg.get("RUN_PROCESS", "True").lower() == "true"

        return {
            "required": {
                "watch_folder": ("STRING", {"default": default_watch}),
                "auth_method": (["oauth", "service_account"], {"default": default_auth}),
                "enable_encryption": ("BOOLEAN", {"default": default_encrypt}),
                "Post_Delete_Enc": ("BOOLEAN", {"default": default_delete_enc}),
                "Subfolder_Monitor": ("BOOLEAN", {"default": default_subfolder}),
                "run_process": ("BOOLEAN", {"default": default_run, "label": "Run Process"}),
            }
        }

    def start(self, watch_folder, auth_method, enable_encryption, Post_Delete_Enc,
              Subfolder_Monitor, run_process):
        """
        Start or stop the Google Drive upload monitor.

        When encryption is enabled, two watchers run:
        1. Encrypt watcher: creates .enc copies of new images (preserves originals)
        2. Upload watcher: uploads .enc files to Google Drive
        """
        global encrypt_observer

        logger.info(
            "Starting DriveSend AutoUploader: watch_folder=%s, auth_method=%s, "
            "encryption=%s, Post_Delete_Enc=%s, Subfolder_Monitor=%s, run_process=%s",
            watch_folder, auth_method, enable_encryption, Post_Delete_Enc,
            Subfolder_Monitor, run_process,
        )

        # ----- Stop path -----
        if not run_process:
            logger.info("Stopping DriveSend AutoUploader monitoring")

            if watcher_observer and watcher_observer.is_alive():
                stop_monitoring()
                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_processor()
            stop_encrypt_queue_processor()

            stop_message = (
                "=====================================================================\n"
                "🚙🛑 DriveSend - AutoUploader - STOPPED\n"
                "=====================================================================\n"
                f"All monitoring, uploading, and encryption processes for {watch_folder} have been stopped.\n"
                "Set 'run_process' to True and run the node again to resume.\n"
                "====================================================================="
            )
            print(stop_message)
            logger.info(stop_message)
            return (f"All monitoring stopped for {watch_folder}",)

        # ----- Validate watch_folder against the allowed roots -----
        try:
            watch_folder = resolve_safe_watch_folder(watch_folder)
        except ValueError as e:
            logger.error(str(e))
            return (f"Error: {e}",)

        watch_path = Path(watch_folder)
        if not watch_path.exists():
            logger.error(f"Watch folder does not exist: {watch_folder}")
            return (f"Error: Watch folder does not exist: {watch_folder}",)
        if not watch_path.is_dir():
            logger.error(f"Watch folder is not a directory: {watch_folder}")
            return (f"Error: Watch folder is not a directory: {watch_folder}",)

        # ----- folder_id required -----
        folder_id = os.environ.get("GOOGLE_DRIVE_FOLDER_ID")
        if not 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 / your secrets manager\n"
                "  2. Restart your pod / container so the env vars load",
            )

        # ----- 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]):
                token_file = NODE_DIR / "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 = NODE_DIR / "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\n"
                    "  2. Place service_account.json in the node folder\n\n"
                    "Note: service accounts only work with Google Workspace.\n"
                    "For personal Gmail, use OAuth instead.",
                )

        # ----- Encryption key requirement -----
        if enable_encryption:
            enc_key = get_encryption_key()
            if not enc_key:
                logger.error("Encryption enabled but no encryption key found")
                return (
                    "Error: Encryption enabled but no key found.\n\n"
                    "Set COMFYUI_ENCRYPTION_KEY environment variable\n"
                    "or run DriveSend Setup with encryption enabled.",
                )

        # ----- Persist (non-credential) uploader settings -----
        _persist_uploader_settings({
            "DRIVESEND_WATCH_FOLDER": watch_folder,
            "DRIVESEND_AUTH_METHOD": auth_method,
            "ENABLE_ENCRYPTION": enable_encryption,
            "POST_DELETE_ENC": Post_Delete_Enc,
            "SUBFOLDER_MONITOR": Subfolder_Monitor,
            "RUN_PROCESS": run_process,
        })

        # ----- Reset queue processors before starting -----
        from .monitor_output import reset_queue_processor as reset_upload_queue
        reset_upload_queue()
        reset_encrypt_queue_processor()

        # ----- Start upload monitor -----
        try:
            start_monitoring(
                watch_folder=watch_folder,
                folder_id=folder_id,
                auth_method=auth_method,
                enable_encryption=enable_encryption,
                delete_enc=Post_Delete_Enc,
                subfolder_monitor=Subfolder_Monitor,
            )
        except Exception as e:
            logger.error(f"Failed to start upload monitor: {e}")
            return (f"Error starting upload monitor: {e}",)

        # ----- Start encryption watcher if enabled -----
        if enable_encryption:
            if encrypt_observer and encrypt_observer.is_alive():
                encrypt_observer.stop()
                encrypt_observer.join()

            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()

        # ----- Status message (no credentials) -----
        status_lines = [
            "DriveSend AutoUploader started.",
            "",
            f"Watching: {watch_folder}",
            f"Drive folder: {folder_id[:20]}...",
            f"Auth method: {auth_method}",
            f"Encryption: {'Enabled (creating .enc copies)' if enable_encryption else 'Disabled'}",
            f"Subfolder monitoring: {'Enabled' if Subfolder_Monitor else 'Disabled'}",
            f"Delete .enc after upload: {'Yes' if Post_Delete_Enc else 'No'}",
        ]

        if enable_encryption:
            status_lines.extend([
                "",
                "Note: original files are preserved for ComfyUI naming.",
                "You must manually delete files from the output folder.",
            ])

        status_lines.extend([
            "",
            "New files will be uploaded automatically.",
            "Set run_process to False to stop.",
        ])

        banner = "\n".join(status_lines)
        print(banner)
        logger.info(banner)
        return (banner,)


NODE_CLASS_MAPPINGS = {
    "DriveSendAutoUploader": DriveSendAutoUploaderNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DriveSendAutoUploader": "🚙📤 DriveSend - AutoUploader"
}