From 4f85fcadfae1dfcb183452760d10b323b72ad7e7 Mon Sep 17 00:00:00 2001 From: MachinePainting Date: Sat, 9 May 2026 07:52:42 -0700 Subject: [PATCH] v1.1.0: Security hardening parity with DropSend (untested) Mirrors the security architecture of ComfyUI_DropSendNode v1.1.x to close the same vulnerability classes. The actual Google OAuth grant path is unchanged from v1.0.1 (same Flow.from_client_config, same OOB redirect, same fetch_token call), so existing users' OAuth setup continues to work. What changed is the wrapper around it. Note: this release has not been runtime-tested end to end; the security logic is a direct port of the tested DropSend implementation. Setup Node refactor: - Removed google_client_id, google_client_secret, auth_code, and the service-account JSON from INPUT_TYPES. They no longer enter PromptServer.history, workflow JSON, PNG metadata, localStorage, or copy-pasted nodes. - Added a browser-only modal (web/drivesend_oauth.js) and a same-origin POST route /drivesend/setup/stash for entry. Hardening: origin check, live-WebSocket gate, per-IP rate limit, 32 KB body cap, JSON content type lockdown, 60-second TTL, 32-entry capacity, lock-protected access. - Added COMFYUI_DRIVESEND_ALLOW_SETUP=1 opt-in gate. Without it the node refuses to write or clear credentials, blocking remote workflow submitters from hijacking or wiping the operator's Drive connection. - Removed plaintext stdout prints of GOOGLE_CLIENT_ID / GOOGLE_CLIENT_SECRET / GOOGLE_REFRESH_TOKEN / COMFYUI_ENCRYPTION_KEY. Replaced with a redacted (last-4-chars) banner. - WebSocket credential delivery is sid-targeted, never broadcast. Refuses to send when client_id is None and surfaces a clear error. JS handler verifies client_id matches api.clientId before rendering. - Service-account JSON is entered via a textarea in the same modal, base64-encoded into GOOGLE_SERVICE_ACCOUNT_JSON for delivery. File-system hardening: - .env and token.json now created via os.open(O_WRONLY|O_CREAT|O_TRUNC, 0o600), race-free, instead of open('w') with default umask (0o644). - Added safe_paths.py: clamps watch_folder to ComfyUI's output tree (with COMFYUI_DRIVESEND_ALLOWED_WATCH_PATHS opt-in for additional roots), rejects symlinks in event paths, validates Drive folder IDs. - AutoUploader now runs every event path through safe_paths.is_safe_event_path. Logging: - Centralized logging in __init__.py with race-free 0o600 creation of drivesend.log on the root logger. Removed module-level logging.basicConfig duplicates from monitor_output.py, encrypt_file.py, drivesend_uploader_node.py. - Verified: no logger or print call in the codebase emits an access token, refresh token, client secret, service-account JSON body, or encryption key value. .gitignore additions: - token.json, client_secret.json, service_account.json (Google credential caches, previously not ignored). - Editor backups: *.bak, *.backup, *~, *.swp, *.swo, .idea/. - Encrypted artifacts: *.enc. - Stray credential material: *.key, *.pem, *.p12, *.pfx. - .claude/, .pytest_cache/, .coverage, htmlcov/, build artifacts, Thumbs.db, Desktop.ini. README rewrite: - Linear three-step structure (Google Cloud setup, Setup Node, AutoUploader). - Cloud / Local / Service-Account sections for the Setup Node. - Threat model section documenting all enforced protections. - "Setup is disabled" and "Browser delivery refused" troubleshooting. - No em dashes per project style. Tests: - tests/test_security.py with five regression tests covering: send_sync is never called with sid=None for credential events; payload echoes client_id for JS-side verification; .env file mode is exactly 0o600; client_id pattern rejects malformed input; stash entries expire after TTL. Tests run without google-auth, googleapiclient, or watchdog installed via ghost modules. Version: 1.0.1 -> 1.1.0 --- .gitignore | 41 ++ README.md | 554 +++++++++++-------- __init__.py | 72 ++- drivesend_setup_node.py | 1046 ++++++++++++++++++++++++++++-------- drivesend_uploader_node.py | 275 ++++++---- encrypt_file.py | 34 +- gdrive_auth_manager.py | 13 +- monitor_output.py | 22 +- pyproject.toml | 2 +- safe_paths.py | 139 +++++ tests/__init__.py | 0 tests/test_security.py | 292 ++++++++++ web/drivesend_oauth.js | 571 ++++++++++++++++++++ 13 files changed, 2436 insertions(+), 625 deletions(-) create mode 100644 safe_paths.py create mode 100644 tests/__init__.py create mode 100644 tests/test_security.py create mode 100644 web/drivesend_oauth.js diff --git a/.gitignore b/.gitignore index 85b5779..c90f6f8 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,11 @@ env/ .env .env.* +# Google Drive credential files (refresh token, client secret, service account key) +token.json +client_secret.json +service_account.json + # MacOS system files .DS_Store @@ -33,3 +38,39 @@ node_modules/ output/ temp/ logs/ + +# Editor/IDE backups (may capture secrets from edited files like .env) +*.bak +*.backup +*~ +*.swp +*.swo +.idea/ + +# Encrypted artifacts from local test runs +*.enc + +# Stray credential material +*.key +*.pem +*.p12 +*.pfx + +# Claude Code session data +.claude/ + +# Python test artifacts +.pytest_cache/ +.coverage +htmlcov/ + +# Python build artifacts +*.egg-info/ +build/ +dist/ +pip-log.txt +pip-wheel-metadata/ + +# Windows OS files +Thumbs.db +Desktop.ini diff --git a/README.md b/README.md index bc03b4f..8b9063d 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ -# ComfyUI DriveSend Node (ComfyUI >> Google Drive) +# ComfyUI DriveSend Node (ComfyUI to Google Drive) -Automatically upload your ComfyUI output files to Google Drive with optional encryption. Set it and forget it! +Automatically upload your ComfyUI output files to Google Drive with optional encryption. Set it and forget it. > **Prefer Dropbox?** Check out [DropSend Node](https://github.com/machinepainting/ComfyUI_DropSendNode) @@ -8,7 +8,7 @@ Automatically upload your ComfyUI output files to Google Drive with optional enc --- -## ๐Ÿ”„ How It Works +## How It Works ``` โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” @@ -39,31 +39,21 @@ Automatically upload your ComfyUI output files to Google Drive with optional enc โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ ``` -**Encryption is completely optional.** Leave `enable_encryption` off and files upload directly as-is. +Encryption is optional. Leave `enable_encryption` off and files upload as-is. --- -## ๐Ÿ“ค Features - -### DriveSend AutoUploader Node -- Real-time folder monitoring with optional subfolder support -- Supports: `.png`, `.jpg`, `.jpeg`, `.webp`, `.gif`, `.mp4`, `.avi`, `.mov` -- Optional AES-128 encryption before upload -- SHA256 checksum verification -- Queue system for reliable uploads +## Two Nodes ### DriveSend Setup Node -- OAuth 2.0 authentication for personal Gmail (free) -- Outputs credentials for RunPod environment variables -- Optional encryption key generation +Runs once. Exchanges your Google OAuth credentials for a refresh token, then either writes them to `.env` (local installs) or shows them in a one-shot browser panel (cloud installs) so you can copy them into your platform's secrets manager. Also supports Google Workspace service accounts. -### Decryption Scripts (Local Use) -- Cross-platform scripts for macOS, Windows, and Linux -- Decrypt `.enc` files after downloading from Google Drive +### DriveSend AutoUploader Node +Runs every workflow. Watches your ComfyUI output folder, optionally encrypts each new file, and uploads it to Google Drive. Supports `.png`, `.jpg`, `.jpeg`, `.webp`, `.gif`, `.mp4`, `.avi`, `.mov`. Includes SHA256 verification, queue-based retries, and optional subfolder monitoring. --- -## ๐Ÿ’พ Installation +## Installation ```bash cd ComfyUI/custom_nodes/ @@ -75,305 +65,409 @@ Restart ComfyUI after installation. --- -## ๐Ÿ”ง Google Cloud Setup +## Step 1: Set Up Your Google Cloud Project -### Step 1: Create Google Cloud Project +This part takes about 5 minutes the first time and never again. -1. Go to [Google Cloud Console](https://console.cloud.google.com/) -2. Sign in with your **personal Gmail account** -3. Click **Select a project** (next to "Google Cloud" logo) โ†’ **New Project** -4. Name it (e.g., `ComfyUI-DriveSend`) โ†’ Click **Create** -5. Make sure your new project is selected +### A) Create the project -### Step 2: Enable Google Drive API +1. Go to [Google Cloud Console](https://console.cloud.google.com/) and sign in with your Gmail. +2. Click **Select a project** (next to "Google Cloud") then **New Project**. +3. Name it (e.g. `ComfyUI-DriveSend`) and click **Create**. +4. Make sure the new project is selected. -1. Go to **APIs & Services** โ†’ **Library** (left sidebar) -2. Search for **Google Drive API** -3. Click on it โ†’ Click **Enable** +### B) Enable the Google Drive API -### Step 3: Configure Google Auth Platform +1. Open **APIs & Services** then **Library** (left sidebar). +2. Search for **Google Drive API** and click **Enable**. -1. Go to **APIs & Services** โ†’ **OAuth consent screen** (left sidebar) -2. Click **Get Started** +### C) Configure the OAuth consent screen + +1. Open **APIs & Services** then **OAuth consent screen**. +2. Click **Get Started**. 3. Fill in: - **App name:** `DriveSend` - - **User support email:** Select your email -4. Click **Next** โ†’ Select **External** โ†’ Click **Next** -5. Enter your email for contact information โ†’ Click **Next** -6. Check the agreement box โ†’ Click **Create** + - **User support email:** your email +4. Click **Next**, choose **External**, click **Next**. +5. Enter your email for contact info, click **Next**. +6. Check the agreement box and click **Create**. -### Step 4: Create OAuth Client +### D) Create the OAuth client -1. Click **Create OAuth Client** -2. **Application type:** Select **Desktop app** -3. **Name:** `DriveSend` -4. Click **Create** +1. Click **Create OAuth Client**. +2. **Application type:** **Desktop app**. +3. **Name:** `DriveSend`. +4. Click **Create**. +5. Copy and save these two values now: + - **Client ID** (ends with `.apps.googleusercontent.com`) + - **Client Secret** (starts with `GOCSPX-`) +6. Click **Done**. -โš ๏ธ **IMPORTANT:** Copy and save these values NOW: -- **Client ID** (ends with `.apps.googleusercontent.com`) -- **Client Secret** (starts with `GOCSPX-`) -- Click **Download JSON** as backup +### E) Create a Drive folder and grab its ID -5. Click **Done** +1. Go to [Google Drive](https://drive.google.com/) and create a folder (e.g. `ComfyUI_Uploads`). +2. Open the folder and copy the ID from the URL: -### Step 5: Create a Google Drive Folder - -1. Go to [Google Drive](https://drive.google.com/) -2. Create a new folder (e.g., `ComfyUI_Uploads`) -3. Open the folder and copy the **Folder ID** from the URL: - ``` - https://drive.google.com/drive/folders/XXXXXXXXXXXXXXXXX - โ””โ”€โ”€ This is your Folder ID - ``` +``` +https://drive.google.com/drive/folders/XXXXXXXXXXXXXXXXX + ^^^^^^^^^^^^^^^^^ + this is the folder_id +``` --- -## ๐Ÿƒโ€โ™‚๏ธ DriveSend Setup (RunPod & Cloud Users) +## Step 2: Run the Setup Node -### Step 6: Run Setup Node +The Setup Node is gated behind an environment variable. This prevents a remote workflow from hijacking or wiping your credentials. Pick the section that matches where you run ComfyUI. -1. Add **DriveSend Setup Node** to your ComfyUI workflow -2. Configure: - - `auth_method`: **oauth** - - `folder_id`: Paste your Folder ID - - `client_id`: Paste your Client ID - - `client_secret`: Paste your Client Secret - - `auth_code`: Leave blank (for now) - - `storage_method`: **display_only** - - `encryption_key_method`: **Display Only** (or **off** if not using encryption) +### A) Cloud (RunPod and similar) -3. Run the node โ†’ Check the terminal for a URL -4. **Copy and open the URL** to open Google authorization -5. Sign in and **accept access** for DriveSend -6. Click **Continue** when prompted โ†’ Copy the **authorization code** -7. Paste code into `auth_code` field -8. Run the node again - -### Step 7: Copy Credentials from Console - -The console displays: +**1. Set the gate.** Add this to your RunPod template's **Environment Variables**: ``` -============================================================ -[DriveSend Setup] Configuration Complete! -============================================================ - -GOOGLE_DRIVE_FOLDER_ID=your_folder_id -GOOGLE_CLIENT_ID=your_client_id.apps.googleusercontent.com -GOOGLE_CLIENT_SECRET=GOCSPX-your_secret -GOOGLE_REFRESH_TOKEN=your_refresh_token -COMFYUI_ENCRYPTION_KEY=your_key (if encryption enabled) - -============================================================ +COMFYUI_DRIVESEND_ALLOW_SETUP=1 ``` -**Copy only the value after the `=` sign** (do not include the "VARIABLE_NAME="). +Start (or restart) the pod. To verify, open a pod shell and run: -### Step 8: Create RunPod Secrets +```bash +echo $COMFYUI_DRIVESEND_ALLOW_SETUP # should print: 1 +``` -1. Go to [RunPod.io](https://www.runpod.io) โ†’ Click **Secrets** in sidebar -2. Create these secrets (Make sure the name of the secret is EXACTLY as shown below): +**2. Add the Setup Node to your workflow.** Configure: -| Secret Name | Secret Value | -|-------------|--------------| -| `GOOGLE_DRIVE_FOLDER_ID` | Your folder ID | -| `GOOGLE_CLIENT_ID` | Your client ID | -| `GOOGLE_CLIENT_SECRET` | Your client secret | -| `GOOGLE_REFRESH_TOKEN` | Your refresh token | -| `COMFYUI_ENCRYPTION_KEY` | Your encryption key *(optional - only if using encryption)* | +| Field | Value | +|---|---| +| `auth_method` | `oauth` | +| `folder_id` | Your Drive folder ID (from Step 1E) | +| `storage_method` | `display_only` (values shown in browser, never saved on the pod) | +| `encryption_key_method` | `display_only` (or `off` if you do not want encryption) | +| `reconnect` | leave off | -### Step 9: Create a Custom Template (Recommended) +**3. Click "Set credentials..."** on the node. A modal opens with three password fields: -Creating a template saves your environment variables so you don't have to add them every time. +| Field | What to paste | +|---|---| +| Google OAuth Client ID | Your Client ID (from Step 1D) | +| Google OAuth Client Secret | Your Client Secret (from Step 1D) | +| Auth Code | Leave blank for now | -1. Go to [RunPod.io](https://www.runpod.io) โ†’ Click **My Templates** in sidebar -2. Click **+ New Template** -3. Configure your template settings (container image, GPU, etc.) -4. Click **Environment Variables** dropdown -5. Click **+ Add Environment Variable** for each: +Click **Save**. The modal closes. -| Key | Value (click ๐Ÿ—๏ธ and select secret) | -|-----|-------------------------------------| +**4. Click Queue (Run).** The node prints an OAuth URL. The URL is delivered three ways, use whichever works: + +- A browser popup opens automatically (preferred) +- The ComfyUI terminal prints a `DRIVESEND - GOOGLE DRIVE AUTHORIZATION REQUIRED` banner with the URL (Cmd or Ctrl click to open) +- Wire a `Show Text` node to the Setup Node's `STRING` output + +**5. Authorize at Google.** Sign in, click Allow, and copy the authorization code Google displays. + +**6. Paste the auth code.** A second modal auto-opens for this. Paste your Client ID, Client Secret, and the new Auth Code. Click **Save**. + +**7. Click Queue again.** A **DriveSend Credentials** panel appears in your browser with these values: + +| Value | Description | +|---|---| +| `GOOGLE_DRIVE_FOLDER_ID` | Your Drive folder ID | +| `GOOGLE_CLIENT_ID` | From Google Cloud Console | +| `GOOGLE_CLIENT_SECRET` | From Google Cloud Console | +| `GOOGLE_REFRESH_TOKEN` | Long-lived token, just generated | +| `COMFYUI_ENCRYPTION_KEY` | Only if you chose to generate one | + +Each value has a **Copy** button. There is also a **Copy all as NAME=value** button. + +The panel is browser-only. Closing it discards the data and nothing is saved on the pod. If you close it before copying, just rerun the Setup Node to regenerate the values. + +**8. Save the values to RunPod Secrets.** In RunPod, click **Secrets** and create a secret for each value. Names must match exactly: + +| Secret Name | Value | +|---|---| +| `GOOGLE_DRIVE_FOLDER_ID` | from the panel | +| `GOOGLE_CLIENT_ID` | from the panel | +| `GOOGLE_CLIENT_SECRET` | from the panel | +| `GOOGLE_REFRESH_TOKEN` | from the panel | +| `COMFYUI_ENCRYPTION_KEY` | from the panel (only if using encryption) | + +**9. Add the secrets to your pod template.** Click **My Templates**, edit your template, and under **Environment Variables** add: + +| Key | Value | +|---|---| | `GOOGLE_DRIVE_FOLDER_ID` | `{{ RUNPOD_SECRET_GOOGLE_DRIVE_FOLDER_ID }}` | | `GOOGLE_CLIENT_ID` | `{{ RUNPOD_SECRET_GOOGLE_CLIENT_ID }}` | | `GOOGLE_CLIENT_SECRET` | `{{ RUNPOD_SECRET_GOOGLE_CLIENT_SECRET }}` | | `GOOGLE_REFRESH_TOKEN` | `{{ RUNPOD_SECRET_GOOGLE_REFRESH_TOKEN }}` | -| `COMFYUI_ENCRYPTION_KEY` | `{{ RUNPOD_SECRET_COMFYUI_ENCRYPTION_KEY }}` *(optional)* | +| `COMFYUI_ENCRYPTION_KEY` | `{{ RUNPOD_SECRET_COMFYUI_ENCRYPTION_KEY }}` (optional) | -6. Click **Save Template** -7. When deploying a new pod, click **Change Template** โ†’ select your custom template. Done! +Save the template. The next time you deploy a pod from this template, ComfyUI will see the credentials automatically. -> โš ๏ธ **Why this matters:** If you skip creating a custom template, you'll have to manually add all 5 environment variables every single time you start a new pod. Save yourself the headache โ€” make the template once and reuse it. +> You can now remove `COMFYUI_DRIVESEND_ALLOW_SETUP=1` from the template unless you plan to re-run the Setup Node from this pod. -> **๐Ÿ’ก New to templates?** Take a template you already use, copy its container/docker settings into a new custom template, add the environment variables above, and save. That's it. - -> **Already have a pod running?** Terminate it โ†’ deploy a new pod using your custom template. This is the cleanest way to pick up the new environment variables. - -### Step 10: Test Upload - -1. Open ComfyUI on your RunPod -2. Add **DriveSend AutoUploader** node to a workflow -3. Set: - - `auth_method`: **oauth** - - `run_process`: **True** - - `enable_encryption`: **True** or **False** (your choice) - - `Post_Delete_Enc`: **False** - - `Subfolder_Monitor`: **True** - - `Run_Process`: **True** -4. Run a workflow that generates an image -5. Check your Google Drive folder! +**10. Deploy a fresh pod from your template.** Skip to **Step 3**. --- -## ๐Ÿ’ป DriveSend Setup (Local Users) +### B) Local install -1. Add **DriveSend Setup Node** to your workflow +**1. Set the gate in the same shell that will launch ComfyUI:** + +```bash +# macOS / Linux +export COMFYUI_DRIVESEND_ALLOW_SETUP=1 + +# Windows (Command Prompt) +set COMFYUI_DRIVESEND_ALLOW_SETUP=1 + +# Windows (PowerShell) +$env:COMFYUI_DRIVESEND_ALLOW_SETUP = "1" +``` + +Verify (each command should print `1`): + +```bash +# macOS / Linux +echo $COMFYUI_DRIVESEND_ALLOW_SETUP + +# Windows (Command Prompt) +echo %COMFYUI_DRIVESEND_ALLOW_SETUP% + +# Windows (PowerShell) +echo $env:COMFYUI_DRIVESEND_ALLOW_SETUP +``` + +Now launch ComfyUI from that same shell. + +**2. Add the Setup Node to your workflow.** Configure: + +| Field | Value | +|---|---| +| `auth_method` | `oauth` | +| `folder_id` | Your Drive folder ID | +| `storage_method` | `env_file` (values written to `.env`, mode `0600`, in the plugin directory) | +| `encryption_key_method` | `save_to_env` (or `off`) | +| `reconnect` | leave off | + +**3. Click "Set credentials..."** on the node. In the modal: + +| Field | What to paste | +|---|---| +| Google OAuth Client ID | Your Client ID | +| Google OAuth Client Secret | Your Client Secret | +| Auth Code | Leave blank | + +Click **Save**. + +**4. Click Queue.** The node prints an OAuth URL. Open it, authorize at Google, and copy the auth code. + +**5. Paste the auth code.** A second modal auto-opens. Paste Client ID, Client Secret, and Auth Code. Click **Save**. + +**6. Click Queue again.** Credentials are written to `.env` in the plugin directory and loaded into the running ComfyUI process. No restart needed. + +> Once setup is complete, you can drop `COMFYUI_DRIVESEND_ALLOW_SETUP` from your environment. The AutoUploader does not need it. + +--- + +### C) Service Account (Google Workspace only) + +If you have a Google Workspace (paid) account, you can use a service account instead of OAuth. Service accounts do not work with personal Gmail. + +1. In Google Cloud Console, open **IAM & Admin** then **Service Accounts**. +2. Create a service account, then create a JSON key for it. +3. Share your Drive folder with the service account's email (e.g. `something@project-id.iam.gserviceaccount.com`). +4. Set `auth_method` on the Setup Node to `service_account`. +5. Click **Set credentials...** and paste the entire JSON body of the key file into the **Service Account JSON** textarea. +6. Click Save, then Queue. + +The service account JSON is base64-encoded into `GOOGLE_SERVICE_ACCOUNT_JSON` for you to copy into RunPod Secrets (display_only) or write to `.env` (env_file). + +--- + +## Step 3: Use the AutoUploader + +1. Add the **DriveSend AutoUploader Node** to your workflow. 2. Configure: - - `auth_method`: **oauth** - - `folder_id`: Paste your Folder ID - - `client_id`: Paste your Client ID - - `client_secret`: Paste your Client Secret - - `auth_code`: Leave blank (for now) - - `storage_method`: **env_file** - - `encryption_key_method`: **save to .env** (or off) -3. Run the node โ†’ Click URL in terminal โ†’ Authorize โ†’ Copy auth code -4. Paste code into `auth_code` field โ†’ Run again -5. Restart ComfyUI to load the `.env` file -6. Add **DriveSend AutoUploader Node** and run! +| Field | What it does | +|---|---| +| `watch_folder` | Folder to monitor for new files (defaults to ComfyUI output, see clamp note in Security section) | +| `auth_method` | `oauth` or `service_account` (must match Step 2) | +| `enable_encryption` | If true, encrypts each file with AES (Fernet) before uploading | +| `Post_Delete_Enc` | After upload, delete the local `.enc` file | +| `Subfolder_Monitor` | Recursively watch subfolders | +| `run_process` | Set to `True` to start uploading | + +3. Run a workflow that generates an image. The AutoUploader picks it up, optionally encrypts it, uploads it to your Drive folder, and prints status to the ComfyUI console. --- -## ๐Ÿ” Decryption Scripts (Local Use Only) +## Decryption (Local Use Only) -The `/scripts/` folder contains scripts to decrypt `.enc` files on your local machine. +If you uploaded with `enable_encryption=True`, files arrive in Drive as `filename.ext.enc`. Decrypt them on your local machine after downloading. -> โš ๏ธ **LOCAL USE ONLY** โ€” Run these after downloading encrypted files from Google Drive. +> **Decrypt files only after they have been moved to your local computer or external drive.** Decrypting on the cloud pod defeats the purpose of encrypting before upload. -> **ONLY DECRYPT ONCE FILES HAVE BEEN MOVED TO YOUR LOCAL COMPUTER OR EXTERNAL HD. OTHERWISE IT DEFEATS THE PURPOSE OF ENCRYPTION BEFORE SENDING TO CLOUD STORAGE.** - -### Prerequisites on local machine for decrypting +### Install the decrypt dependency ```bash pip install cryptography ``` -### Store Your Encryption Key +### Store your encryption key locally **macOS (Keychain):** -1. Open Keychain Access โ†’ File โ†’ New Password Item -2. Name: `ComfyUI_Encryption_Key`, Account: `ComfyUI`, Password: your key +1. Open Keychain Access, choose File then New Password Item. +2. Name: `ComfyUI_Encryption_Key`. Account: `ComfyUI`. Password: your key. -**Windows (Environment Variable):** -1. Win + R โ†’ `sysdm.cpl` โ†’ Advanced โ†’ Environment Variables -2. New User Variable: `COMFYUI_ENCRYPTION_KEY` = your key +**Windows:** +1. Press `Win + R`, run `sysdm.cpl`, open **Advanced**, click **Environment Variables**. +2. Under **User variables**, click **New**. Name: `COMFYUI_ENCRYPTION_KEY`. Value: your key. -**Linux (Environment Variable):** +**Linux:** ```bash echo 'export COMFYUI_ENCRYPTION_KEY="your_key_here"' >> ~/.bashrc source ~/.bashrc ``` -### Run Decryption +### Run the decrypt script -### Local Decryption (Decrypt Files on Your Computer) +1. Open the `/scripts/` folder in this repository and pick your platform: -Once your encrypted `.enc` files have been downloaded from the cloud, you can decrypt them locally using the platform-specific scripts provided in this repo. + | Platform | File | + |---|---| + | macOS | `mac/decrypt_folder_mac.sh` | + | Windows | `win/decrypt_folder_win.py` | + | Linux | `linux/decrypt_folder_linux.sh` | + | Cross-platform | `decrypt_folder.py` (Python) | -#### Step-by-Step +2. Save the script to a convenient location (your home folder works well). -1. **Go to the [`/scripts`](https://github.com/machinepainting/ComfyUI_DriveSendNode/tree/main/scripts) folder** in this repository. +3. Open a terminal in that location and run: -2. **Select your platform folder** (`mac/`, `win/`, or `linux/`) and download the decrypt script for your OS. If you want to use the cross-platform Python version, grab `decrypt_folder.py` from the scripts root instead. - - | Platform | Folder | Script | - |---|---|---| - | macOS | `mac/` | `decrypt_folder_mac.sh` | - | Windows | `win/` | `decrypt_folder_win.py` | - | Linux | `linux/` | `decrypt_folder_linux.sh` | - | Cross-platform | scripts root | `decrypt_folder.py` | - -3. **Download the script** and place it in a convenient location on your computer. Your user root/home directory is recommended for easy terminal access (e.g. `~/` on macOS/Linux or `C:\Users\YourName\` on Windows). - -4. **Open a terminal** and navigate to the folder where you saved the script: - ```bash - # macOS / Linux โ€” example if saved to home directory - cd ~/ - - # Windows (Command Prompt) โ€” example if saved to user folder - cd C:\Users\YourName\ - ``` - -5. **Run the decryption script:** ```bash # macOS ./decrypt_folder_mac.sh - # Windows - python decrypt_folder_win.py - # Linux ./decrypt_folder_linux.sh - # Cross-platform (Python) + # Windows + python decrypt_folder_win.py + + # Cross-platform Python python decrypt_folder.py ``` -6. **When prompted, drag in (or paste the path to) the folder containing your `.enc` files.** +4. When prompted, drag in (or paste the path to) the folder containing your `.enc` files. -7. **The script will decrypt your files.** Once complete, you'll be asked if you'd like to move the original `.enc` files to a separate folder for cleanup. Note: the script will **never delete your files** โ€” you can remove the `.enc` originals manually at your own discretion. +5. The script decrypts each file alongside the original. Once finished, you are asked whether to move the `.enc` originals into a separate cleanup folder. Originals are never deleted automatically. -That's it โ€” your decrypted files are now available locally on your machine, off the cloud and ready to use. - -> **๐Ÿ’ก Manual Encryption:** Each platform folder also contains an encryption script if you ever want to encrypt files manually outside of ComfyUI โ€” for example, encrypting media before uploading it yourself. These are separate from the node's built-in encryption, which handles everything automatically during upload. +> Each platform folder also contains an encryption script if you ever want to encrypt files manually outside of ComfyUI. The node itself handles encryption automatically during upload. --- -### "storageQuotaExceeded" Error -- **Cause:** Using service account with personal Gmail -- **Fix:** Use OAuth (the instructions above). Service accounts only work with Google Workspace. +## Troubleshooting -### "invalid_grant" Error -- **Cause:** Auth code expired or already used -- **Fix:** Clear the `auth_code` field, run setup again, get new auth code +### Files are not uploading +- Set `run_process` to `True`. +- Check the ComfyUI console for upload errors. +- Confirm the Drive folder is shared with your Google account (or service account email if using one). -### OAuth Credentials Not Found -- **Cause:** Environment variables not set -- **Fix:** Check RunPod Secrets AND Environment Variables in pod template +### "Encryption key not found" +- Confirm the secret name is exactly `COMFYUI_ENCRYPTION_KEY`. +- Verify the env var is set in the pod template (then restart the pod). -### "Access blocked: App not verified" -- **Cause:** OAuth consent screen not configured -- **Fix:** Complete Step 3 (Configure Google Auth Platform) +### "Setup is disabled" or "Reconnect refused" -### Encryption Key Not Found -- **Cause:** Environment variable not set -- **Fix:** Verify secret name is exactly `COMFYUI_ENCRYPTION_KEY`, restart pod +The Setup Node prints this when `COMFYUI_DRIVESEND_ALLOW_SETUP=1` is not visible to the running ComfyUI process. Stop ComfyUI, set the variable in the same terminal you launch from, verify it, and start ComfyUI again. -### Files Not Uploading -- Check console for error messages -- Verify Folder ID is correct -- Confirm `run_process` is **True** +```bash +# macOS / Linux +export COMFYUI_DRIVESEND_ALLOW_SETUP=1 +echo $COMFYUI_DRIVESEND_ALLOW_SETUP # prints: 1 +``` + +```cmd +:: Windows (Command Prompt) +set COMFYUI_DRIVESEND_ALLOW_SETUP=1 +echo %COMFYUI_DRIVESEND_ALLOW_SETUP% +``` + +```powershell +# Windows (PowerShell) +$env:COMFYUI_DRIVESEND_ALLOW_SETUP = "1" +echo $env:COMFYUI_DRIVESEND_ALLOW_SETUP +``` + +For RunPod or Docker, add the variable to the pod template and restart. + +> The variable must be set in the **same shell that launches ComfyUI**. Setting it in one terminal and starting ComfyUI from a different one (or via Pinokio, ComfyUI Desktop, a launcher script, etc.) will not work. Adding it to `~/.bashrc` or `~/.zshrc` and opening a fresh terminal also works. + +### "Browser delivery refused" + +The Setup Node printed this because your prompt was submitted without a `client_id` (typical for `curl` or SDK submissions that omit it). Submitting that way would broadcast credentials to every connected WebSocket client, so the node refuses. Submit the Setup workflow from the ComfyUI web UI instead, or switch `storage_method` to `env_file` to write to disk. + +### Browser panel did not appear after `display_only` setup +- Re-run the Setup Node. Credentials are not stored on the pod in this mode, so the panel is the only retrieval path and it regenerates on each run. +- Check the browser's popup or panel blocker. +- If it still does not appear, switch `storage_method` to `env_file` and `cat` the `.env` file from a shell. + +### "Storage quota exceeded" with a service account +Service accounts only work with Google Workspace. If you are using a personal Gmail, switch to OAuth. + +### Authorization failed +- Confirm `COMFYUI_DRIVESEND_ALLOW_SETUP=1` is set (see above). +- Re-run the Setup Node with `reconnect=True` to clear stale state, then retry. +- If Google does not show a consent screen, revoke access at https://myaccount.google.com/permissions and re-run. --- -## โš ๏ธ Security Best Practices +## Security Best Practices -1. **Never commit your `.env` file** โ€” Add to `.gitignore` -2. **Backup your encryption key** โ€” Without it, encrypted files cannot be recovered -3. **Use secure key storage** โ€” Prefer Keychain/Environment Variables over plain text +1. **Never commit your `.env` file.** It is already in `.gitignore`. +2. **Back up your encryption key.** Without it, encrypted files cannot be recovered. +3. **Prefer Keychain or environment variables** over plaintext storage of the key. + +### Threat model for network-reachable ComfyUI hosts + +If your ComfyUI is reachable over the network (RunPod, tunnels, LAN, public web UI), anyone who can submit a workflow can in principle set node inputs. To keep DriveSend safe in that setting, the nodes enforce these protections. + +- **`watch_folder` is clamped to the ComfyUI output directory.** The AutoUploader runs a recursive Watchdog observer that uploads everything it sees, so an unrestricted path would be an arbitrary-file-read primitive. By default, `watch_folder` must resolve inside the directory returned by `folder_paths.get_output_directory()`. To monitor an additional location, set `COMFYUI_DRIVESEND_ALLOWED_WATCH_PATHS` on the host (`os.pathsep`-separated absolute paths) before starting ComfyUI. Workflow inputs cannot expand this list. + +- **Setup Node secret fields are not workflow inputs.** `google_client_id`, `google_client_secret`, `auth_code`, and the service-account JSON body do not appear in the node's `INPUT_TYPES`. They are entered only via a browser-only modal and POSTed to a same-origin route (`/drivesend/setup/stash`) that is gated on a live WebSocket session. They never enter the workflow JSON, PNG metadata, ComfyUI's localStorage auto-save, copy-pasted nodes, or the unauthenticated `/history` endpoint. + +- **The Setup Node refuses to write or clear credentials unless explicitly opted in.** This is what `COMFYUI_DRIVESEND_ALLOW_SETUP=1` enforces. A remote workflow submitter can send JSON to ComfyUI but cannot set environment variables on your machine, so the gate proves a human with host access opted in. The AutoUploader never needs this flag. + +- **The stash route is hardened.** Same-origin origin check (rejects cross-origin POSTs), live-session check (the supplied `client_id` must match a connected WebSocket), per-IP rate limit (30 POSTs per 10 seconds), 32 KB body cap, JSON validation, 32-entry capacity with 60-second TTL, lock-protected access, one-shot consumption when `setup()` runs. + +- **Credentials never travel through the node's `ui` or `result` channels.** ComfyUI persists both into `PromptServer.history`, served on the unauthenticated `/history` HTTP endpoint. Instead the Setup Node delivers credentials via: + - **`env_file`** (local installs only): values are written to `.env` in the plugin directory (mode `0600`, race-free open) and the node returns only a non-secret confirmation. The OAuth `token.json` (if used) is also written with mode `0600`. + - **`display_only`** (cloud or hosted ComfyUI): values are pushed to the originating browser session via a one-shot WebSocket message and rendered in a panel inside the workflow tab. Nothing is written to disk on the pod. + +- **WebSocket deliveries are sid-targeted, never broadcast.** ComfyUI's `send_sync(event, data, sid=None)` would broadcast to every connected client when `sid` is `None`. The Setup Node refuses to send credential or refresh notifications without a `client_id` and returns a clear "Browser delivery refused" message. As defense in depth, the originating `client_id` is echoed inside each WebSocket payload and the JS handler verifies it matches `api.clientId` before rendering the credentials panel. + +- **Logging stores no secrets.** The plugin's `drivesend.log` records file paths, watcher events, and errors. No tokens, refresh tokens, client secrets, service-account JSON, or encryption keys are written to logs or stdout in plaintext. The post-setup banner uses last-4-character redaction so values are not captured by stdout aggregators. + +- **Sensitive files are gitignored.** `.env`, `token.json`, `client_secret.json`, `service_account.json`, encryption-key files, editor backups, and Claude Code session data are all listed in `.gitignore`. + +### Additional considerations + +- **`COMFYUI_DRIVESEND_ALLOW_SETUP` is process-wide, not per-action.** Once set, every workflow submission in that ComfyUI process is unblocked, including hostile submissions during the setup window. Recommended discipline: set the gate, run Setup once, restart ComfyUI without the gate. + +- **The OAuth authorization URL contains your Google `client_id`.** Google treats `client_id` as a public identifier (it appears in every OAuth URL), but the URL string itself ends up in `/history` when the Setup Node returns it during the first run. This reveals which Google project the host is paired with. It does not, on its own, allow access to your data. + +- **Plain ComfyUI runs over HTTP.** The browser-only credential delivery in `display_only` mode rides the same WebSocket the rest of ComfyUI uses, in cleartext on the wire. On a network you do not fully trust, terminate ComfyUI behind HTTPS (reverse proxy, tunnel) before relying on this path. The most conservative alternative is to run the Setup Node on a local ComfyUI install and copy the values into your cloud secrets manager directly. Credentials never touch the cloud pod's filesystem or its network. --- -## ๐Ÿงช Tested On +## Tested On +- macOS 13+ (Ventura, Sonoma) - Python 3.10 / 3.11 -- ComfyUI (Feb 2026) +- ComfyUI (May 2026) - RunPod GPU instances -- Google Drive API v3 --- -## ๐Ÿ“ Repository Structure +## Repository Structure ``` ComfyUI_DriveSendNode/ @@ -384,9 +478,13 @@ ComfyUI_DriveSendNode/ โ”œโ”€โ”€ gdrive_auth_manager.py โ”œโ”€โ”€ encrypt_file.py โ”œโ”€โ”€ monitor_output.py +โ”œโ”€โ”€ safe_paths.py +โ”œโ”€โ”€ decrypt_folder.py โ”œโ”€โ”€ requirements.txt โ”œโ”€โ”€ README.md โ”œโ”€โ”€ .gitignore +โ”œโ”€โ”€ web/ +โ”‚ โ””โ”€โ”€ drivesend_oauth.js โ””โ”€โ”€ scripts/ โ”œโ”€โ”€ decrypt_folder.py โ”œโ”€โ”€ mac/ @@ -402,24 +500,10 @@ ComfyUI_DriveSendNode/ --- -## ๐Ÿข Google Workspace Users (Service Account Method) - -> **Note:** This method is for paid Google Workspace accounts only. Personal Gmail users should use OAuth (instructions above). - -Service accounts have **0 GB storage quota** on personal Gmail and will fail with "storageQuotaExceeded" error. They only work with Google Workspace Shared Drives. - -If you have a Google Workspace account and want to use a service account: - -1. Create a service account in Google Cloud Console -2. Download the JSON key file -3. Create a Shared Drive and add the service account as Content Manager -4. Base64 encode the JSON and add as `GOOGLE_SERVICE_ACCOUNT_JSON` secret -5. Set `auth_method` to **service_account** in the nodes - -For detailed instructions, see the [Google Workspace documentation](https://cloud.google.com/iam/docs/service-accounts). - ---- - ## License MIT + +--- + +Shout-out to Adam for his contributions to this node build. diff --git a/__init__.py b/__init__.py index 2bb442a..5dda3bf 100644 --- a/__init__.py +++ b/__init__.py @@ -3,22 +3,74 @@ ComfyUI DriveSend Node Automatic Google Drive uploads with optional encryption """ +import os +import logging +from pathlib import Path +from dotenv import load_dotenv + +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- +# Attach a FileHandler to the root logger so every child logger in this +# plugin (logging.getLogger(__name__) in our modules) gets its records +# persisted to drivesend.log alongside ComfyUI's own stdout. We do NOT +# use logging.basicConfig() here: ComfyUI's main.py configures the root +# logger before our package is imported, and basicConfig is a no-op +# once any caller has already configured root. +# +# Race-free perms: pre-create the file with 0o600 via os.open() before +# FileHandler can create it under the default umask (typically 0o644). +# A follow-up chmod tightens perms if the file already existed with +# looser bits. +_LOG_PATH = os.path.join(os.path.dirname(__file__), "drivesend.log") +try: + _fd = os.open(_LOG_PATH, os.O_CREAT | os.O_WRONLY | os.O_APPEND, 0o600) + os.close(_fd) +except OSError: + pass +_root_logger = logging.getLogger() +if not any( + isinstance(h, logging.FileHandler) + and getattr(h, "baseFilename", "") == _LOG_PATH + for h in _root_logger.handlers +): + try: + _fh = logging.FileHandler(_LOG_PATH) + _fh.setFormatter( + logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s") + ) + _root_logger.addHandler(_fh) + if _root_logger.level == logging.NOTSET or _root_logger.level > logging.INFO: + _root_logger.setLevel(logging.INFO) + except OSError: + pass +try: + os.chmod(_LOG_PATH, 0o600) +except OSError: + pass + +# --------------------------------------------------------------------------- +# .env loading +# --------------------------------------------------------------------------- +# Load .env from the plugin directory at import time so GOOGLE_CLIENT_ID, +# GOOGLE_CLIENT_SECRET, GOOGLE_REFRESH_TOKEN, GOOGLE_DRIVE_FOLDER_ID, +# COMFYUI_ENCRYPTION_KEY, and the uploader's persisted settings are +# visible in os.environ before any node runs. override=False is +# intentional. Host-injected env vars (RunPod secrets, Docker env, +# systemd EnvironmentFile) take precedence over the plugin's local .env. +_PLUGIN_ENV = Path(__file__).parent / ".env" +if _PLUGIN_ENV.exists(): + load_dotenv(_PLUGIN_ENV, override=False) + from .drivesend_setup_node import NODE_CLASS_MAPPINGS as SETUP_MAPPINGS from .drivesend_setup_node import NODE_DISPLAY_NAME_MAPPINGS as SETUP_DISPLAY_MAPPINGS from .drivesend_uploader_node import NODE_CLASS_MAPPINGS as UPLOADER_MAPPINGS from .drivesend_uploader_node import NODE_DISPLAY_NAME_MAPPINGS as UPLOADER_DISPLAY_MAPPINGS -# Merge mappings NODE_CLASS_MAPPINGS = {**SETUP_MAPPINGS, **UPLOADER_MAPPINGS} NODE_DISPLAY_NAME_MAPPINGS = {**SETUP_DISPLAY_MAPPINGS, **UPLOADER_DISPLAY_MAPPINGS} -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] +# Define web directory for JavaScript extensions +WEB_DIRECTORY = "./web" -# Load environment variables from .env file if present -from pathlib import Path -from dotenv import load_dotenv - -env_path = Path(__file__).parent / '.env' -if env_path.exists(): - load_dotenv(env_path) - print(f"[DriveSend] Loaded environment from {env_path}") +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] diff --git a/drivesend_setup_node.py b/drivesend_setup_node.py index 409e92a..f197dac 100644 --- a/drivesend_setup_node.py +++ b/drivesend_setup_node.py @@ -1,277 +1,859 @@ """ DriveSend Setup Node -Handles Google Drive API authentication setup for both OAuth and Service Account methods +Handles Google Drive API authentication setup for both OAuth and Service Account methods. + +Architecture: this node has NO credential-bearing inputs. client_id, +client_secret, and auth_code are entered only via a browser-only modal +launched from the JS extension's "Set credentials..." button. The modal +POSTs values directly to the same-origin route /drivesend/setup/stash +and clears its inputs. setup() consumes the per-session stash entry on +the next Queue. This keeps every credential out of PromptServer.history +(served on the unauthenticated /history endpoint), workflow JSON, PNG +metadata, ComfyUI's localStorage auto-save, and copy-pasted nodes. """ import os +import re +import time import json import base64 +import threading +import logging from pathlib import Path from cryptography.fernet import Fernet +from dotenv import dotenv_values, load_dotenv +logger = logging.getLogger(__name__) NODE_DIR = Path(__file__).parent +# --------------------------------------------------------------------------- +# Out-of-band secret transport +# --------------------------------------------------------------------------- +# ComfyUI persists every queued prompt's input JSON in +# PromptServer.history, served on the unauthenticated /history endpoint. +# If the Setup Node accepted client_id/client_secret/auth_code as +# workflow inputs, those values would be visible there for the lifetime +# of the history entry. To close that path: the JS extension shows a +# browser-only modal, POSTs the secret values to /drivesend/setup/stash +# (this module's route), clears the inputs, then lets the queue go +# through. The prompt JSON ComfyUI stores in history therefore has no +# credential fields. setup() pulls the real values from the stash +# (one-shot, keyed by originating client_id). +# +# Entries auto-expire after _SETUP_STASH_TTL_SEC so a stash POST that's +# never consumed (user closed the tab between Save and Queue, etc.) +# does not sit in memory. +# --------------------------------------------------------------------------- + +_setup_secret_stash = {} +_setup_secret_lock = threading.Lock() +_SETUP_STASH_TTL_SEC = 60 +_SETUP_STASH_MAX_ENTRIES = 32 +_CLIENT_ID_PATTERN = re.compile(r"^[A-Za-z0-9._-]{1,128}$") + +# Per-IP rate limit on the stash route, defense in depth against +# stash-spam attempts to evict a legitimate user's entry. +_STASH_RL_WINDOW_SEC = 10 +_STASH_RL_MAX_PER_WINDOW = 30 +_stash_rl_state = {} +_stash_rl_lock = threading.Lock() + + +def _stash_rl_check(remote_ip): + if not remote_ip: + return True + now = time.time() + cutoff = now - _STASH_RL_WINDOW_SEC + with _stash_rl_lock: + if len(_stash_rl_state) > 256: + stale = [ip for ip, hits in _stash_rl_state.items() if not hits or hits[-1] < cutoff] + for ip in stale: + _stash_rl_state.pop(ip, None) + hits = _stash_rl_state.setdefault(remote_ip, []) + i = 0 + while i < len(hits) and hits[i] < cutoff: + i += 1 + if i: + del hits[:i] + if len(hits) >= _STASH_RL_MAX_PER_WINDOW: + return False + hits.append(now) + return True + + +def _stash_prune_expired_locked(): + now = time.time() + expired = [k for k, (ts, _v) in _setup_secret_stash.items() if now - ts > _SETUP_STASH_TTL_SEC] + for k in expired: + _setup_secret_stash.pop(k, None) + + +def _stash_set(client_id, payload): + with _setup_secret_lock: + _stash_prune_expired_locked() + if len(_setup_secret_stash) >= _SETUP_STASH_MAX_ENTRIES and client_id not in _setup_secret_stash: + oldest_key = min(_setup_secret_stash, key=lambda k: _setup_secret_stash[k][0]) + _setup_secret_stash.pop(oldest_key, None) + _setup_secret_stash[client_id] = (time.time(), payload) + + +def _stash_consume(client_id): + with _setup_secret_lock: + _stash_prune_expired_locked() + entry = _setup_secret_stash.pop(client_id, None) + if entry is None: + return None + _ts, payload = entry + return payload + + +def _register_stash_route_once(): + """Register the /drivesend/setup/stash POST handler on PromptServer. + + Idempotent. If PromptServer is not ready yet, registration is silently + skipped and the Setup Node will surface a clear error when the user + tries to save credentials. + """ + if getattr(_register_stash_route_once, "_registered", False): + return + try: + from server import PromptServer + from aiohttp import web + except Exception: + return + instance = getattr(PromptServer, "instance", None) + if instance is None or not hasattr(instance, "routes"): + return + + @instance.routes.post("/drivesend/setup/stash") + async def _stash_handler(request): + # Same-origin enforcement. Without this, a malicious page in + # another tab could attempt a CSRF POST against + # http://127.0.0.1:8188. + origin = request.headers.get("Origin", "") + if origin: + try: + from urllib.parse import urlparse + origin_host = urlparse(origin).netloc + except Exception: + origin_host = "" + if not origin_host or origin_host != request.host: + return web.json_response({"error": "cross-origin denied"}, status=403) + + if request.content_type and request.content_type != "application/json": + return web.json_response( + {"error": "Content-Type must be application/json"}, + status=415, + ) + + remote_ip = request.headers.get("X-Forwarded-For", request.remote or "").split(",")[0].strip() + if not _stash_rl_check(remote_ip): + return web.json_response({"error": "rate limit"}, status=429) + + if request.content_length is not None and request.content_length > 32 * 1024: + return web.json_response({"error": "payload too large"}, status=413) + try: + data = await request.json() + except Exception: + return web.json_response({"error": "invalid json"}, status=400) + client_id_field = data.get("client_id") if isinstance(data, dict) else None + if not client_id_field or not isinstance(client_id_field, str) or not _CLIENT_ID_PATTERN.match(client_id_field): + return web.json_response({"error": "client_id required"}, status=400) + + # Authenticate by liveness: only accept stash writes for a + # client_id that currently has an active WebSocket connection + # to this PromptServer instance. + try: + sockets = getattr(instance, "sockets", None) + if not sockets or client_id_field not in sockets: + return web.json_response({"error": "client_id has no active WebSocket"}, status=403) + except Exception: + return web.json_response({"error": "internal"}, status=500) + + payload = { + "google_client_id": str(data.get("google_client_id") or ""), + "google_client_secret": str(data.get("google_client_secret") or ""), + "auth_code": str(data.get("auth_code") or ""), + "service_account_json": str(data.get("service_account_json") or ""), + } + if not any(payload.values()): + return web.json_response({"ok": True, "stashed": False}) + _stash_set(client_id_field, payload) + return web.json_response({"ok": True, "stashed": True}) + + _register_stash_route_once._registered = True + print("[DriveSend Setup] Registered /drivesend/setup/stash route (out-of-band secret transport)") + + +_register_stash_route_once() + + +# --------------------------------------------------------------------------- +# Opt-in gate +# --------------------------------------------------------------------------- +# Workflow inputs to the Setup Node can come from any caller that can +# submit a workflow to ComfyUI. On a network-reachable host that means a +# remote attacker can wipe credentials or overwrite the stored Google +# Drive app to redirect uploads to their own account. We require the +# operator to opt in by setting this env var on the host before the +# Setup Node performs any state-changing action. +_SETUP_OPT_IN_VAR = "COMFYUI_DRIVESEND_ALLOW_SETUP" +_SETUP_OPT_IN_HELP = ( + f"Setup is disabled. To allow this node to write or clear credentials, " + f"set {_SETUP_OPT_IN_VAR}=1 in the host environment before starting " + f"ComfyUI, or provision GOOGLE_CLIENT_ID / GOOGLE_CLIENT_SECRET / " + f"GOOGLE_REFRESH_TOKEN directly via environment variables (e.g. RunPod " + f"secrets). This guard prevents remote workflow submitters from " + f"hijacking or wiping your Google Drive connection." +) + + +def _setup_opt_in_enabled(): + return os.getenv(_SETUP_OPT_IN_VAR, "").strip().lower() in ("1", "true", "yes") + + +# --------------------------------------------------------------------------- +# Redaction helpers for the post-setup banner +# --------------------------------------------------------------------------- +_SECRET_FIELDS = { + "GOOGLE_CLIENT_ID", + "GOOGLE_CLIENT_SECRET", + "GOOGLE_REFRESH_TOKEN", + "GOOGLE_SERVICE_ACCOUNT_JSON", + "COMFYUI_ENCRYPTION_KEY", +} + + +def _redact(value, visible_tail=4): + if not value: + return "" + if len(value) <= visible_tail: + return "*" * len(value) + return "*" * (len(value) - visible_tail) + value[-visible_tail:] + + +def _redact_kv_line(line): + name, sep, value = line.partition("=") + if sep and name in _SECRET_FIELDS: + return f"{name}={_redact(value)}" + return line + + +# --------------------------------------------------------------------------- +# The Setup Node +# --------------------------------------------------------------------------- + class DriveSendSetupNode: """Setup node for configuring Google Drive authentication.""" - + CATEGORY = "DriveSend" FUNCTION = "setup" RETURN_TYPES = ("STRING",) RETURN_NAMES = ("status",) OUTPUT_NODE = True - + @classmethod def INPUT_TYPES(cls): + # Architecture: NO credential-bearing inputs. google_client_id, + # google_client_secret, auth_code, and the service-account JSON + # body all live exclusively in the per-session stash populated + # by the browser modal. return { "required": { "auth_method": (["oauth", "service_account"], {"default": "oauth"}), "folder_id": ("STRING", { - "default": "", + "default": "", "multiline": False, - "tooltip": "Google Drive folder ID from URL (after /folders/)" + "tooltip": "Google Drive folder ID from URL (the part after /folders/)" + }), + "storage_method": (["display_only", "env_file"], { + "label": "Credential Storage Method", + "default": "display_only" + }), + "encryption_key_method": (["off", "display_only", "save_to_env"], { + "label": "Encryption Key Method", + "default": "off" }), - "storage_method": (["display_only", "env_file"], {"default": "display_only"}), - "encryption_key_method": (["off", "display_only", "save_to_env"], {"default": "off"}), }, "optional": { - "client_id": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "OAuth Client ID (from Google Cloud Console)" - }), - "client_secret": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "OAuth Client Secret (from Google Cloud Console)" - }), - "auth_code": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "Authorization code from OAuth flow (leave blank first run)" + "reconnect": ("BOOLEAN", { + "label": "Reset stored credentials", + "default": False }), } } - - def setup(self, auth_method, folder_id, storage_method, encryption_key_method, - client_id="", client_secret="", auth_code=""): - """ - Set up Google Drive authentication. - - For OAuth: - - First run: Opens browser for authorization, displays auth URL - - Second run: Exchanges auth code for refresh token - - For Service Account: - - Validates service_account.json exists - - Displays environment variable format for RunPod - """ - - results = [] - env_vars = {} - - # Validate folder_id - if not folder_id: - return ("โŒ Error: folder_id is required. Get it from your Google Drive folder URL.",) - - env_vars['GOOGLE_DRIVE_FOLDER_ID'] = folder_id - - # Handle encryption key - encryption_key = None - if encryption_key_method != "off": - # Check for existing key - # Check for existing key (multiple possible names for compatibility) - existing_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 existing_key: - encryption_key = existing_key - results.append("Using existing encryption key from environment") - else: - # Generate new key - encryption_key = Fernet.generate_key().decode('utf-8') - results.append("Generated new encryption key") - - env_vars['COMFYUI_ENCRYPTION_KEY'] = encryption_key - - # === OAuth Setup === - if auth_method == 'oauth': - if not client_id or not client_secret: - return ( - "โŒ Error: OAuth requires client_id and client_secret.\n\n" - "To get these:\n" - "1. Go to Google Cloud Console โ†’ APIs & Services โ†’ Credentials\n" - "2. Create OAuth 2.0 Client ID (Desktop app)\n" - "3. Copy Client ID and Client Secret here", + + @classmethod + def IS_CHANGED(cls, *args, **kwargs): + # Force re-execution on every Queue. setup() is intentionally + # side-effecting, and the input dict here does not reflect the + # actual work (the real driver is the out-of-band stash entry). + return float("NaN") + + def setup(self, auth_method, folder_id, storage_method, encryption_key_method, reconnect=False): + try: + print(f"[DriveSend Setup] Called: auth_method={auth_method}, " + f"reconnect={reconnect}, storage_method={storage_method}, " + f"encryption_key_method={encryption_key_method}") + + # Pull stash for this client (browser-only entry). + google_client_id = None + google_client_secret = None + auth_code = None + service_account_json = None + sid = None + try: + from server import PromptServer + sid = getattr(PromptServer.instance, "client_id", None) + if sid: + stashed = _stash_consume(sid) + if stashed: + google_client_id = stashed.get("google_client_id") or None + google_client_secret = stashed.get("google_client_secret") or None + auth_code = stashed.get("auth_code") or None + service_account_json = stashed.get("service_account_json") or None + print("[DriveSend Setup] Loaded secrets from out-of-band stash") + except Exception as e: + print(f"[DriveSend Setup] Warning: could not consume secret stash: {e}") + + # ----- Reconnect (clear all stored credentials) ----- + if reconnect: + if not _setup_opt_in_enabled(): + message = "Reconnect refused: " + _SETUP_OPT_IN_HELP + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + print("[DriveSend Setup] Reconnect requested - clearing all credentials") + + env_path = NODE_DIR / ".env" + token_path = NODE_DIR / "token.json" + + if env_path.exists(): + print(f"[DriveSend Setup] Removing .env file: {env_path}") + try: + os.remove(env_path) + except OSError as e: + print(f"[DriveSend Setup] Warning: could not remove .env: {e}") + + if token_path.exists(): + print(f"[DriveSend Setup] Removing token.json: {token_path}") + try: + os.remove(token_path) + except OSError as e: + print(f"[DriveSend Setup] Warning: could not remove token.json: {e}") + + # Clear in-process os.environ so subsequent runs do not + # short-circuit on the previously-loaded values. + for key in ( + "GOOGLE_CLIENT_ID", "GOOGLE_CLIENT_SECRET", "GOOGLE_REFRESH_TOKEN", + "GOOGLE_SERVICE_ACCOUNT_JSON", + ): + if key in os.environ: + os.environ.pop(key, None) + print(f"[DriveSend Setup] Cleared {key} from process environment") + + # WebSocket notification, sid-targeted only (refuse to + # broadcast on sid=None). + try: + from server import PromptServer + notify_sid = getattr(PromptServer.instance, "client_id", None) + if not notify_sid: + print( + "[DriveSend Setup] Skipping reconnect-complete " + "broadcast: prompt has no client_id (CLI " + "submission?). Submit from the ComfyUI web UI " + "to receive the refresh notification." + ) + else: + PromptServer.instance.send_sync( + "drivesend_reconnect_complete", + { + "type": "drivesend_reconnect_complete", + "success": True, + "client_id": notify_sid, + "message": "Credentials cleared. ComfyUI will refresh to show auth fields", + }, + notify_sid, + ) + print("[DriveSend Setup] Sent WebSocket notification for reconnect completion") + except Exception as e: + print(f"[DriveSend Setup] Warning: Could not send WebSocket notification: {e}") + + message = "Google Drive credentials cleared." + return {"ui": {"text": [message]}, "result": (message,)} + + # ----- folder_id is required ----- + if not folder_id or not str(folder_id).strip(): + message = "Error: folder_id is required. Get it from your Google Drive folder URL (the part after /folders/)." + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + from .safe_paths import validate_drive_folder_id + try: + folder_id = validate_drive_folder_id(folder_id) + except ValueError as e: + message = f"Error: {e}" + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + # ----- env_vars short-circuit (already configured) ----- + # + # If credentials are already provisioned in the host environment + # (RunPod secrets, Docker env, etc.), there's nothing to do. + if auth_method == "oauth": + env_vars_set = all([ + os.getenv("GOOGLE_CLIENT_ID"), + os.getenv("GOOGLE_CLIENT_SECRET"), + os.getenv("GOOGLE_REFRESH_TOKEN"), + ]) + if env_vars_set and not auth_code and not google_client_id: + # Persist the folder_id only (non-sensitive). Use the + # secure write helper. + _persist_env_writes([("GOOGLE_DRIVE_FOLDER_ID", folder_id)]) + message = ( + "Google Drive credentials already configured via host environment. " + "Folder ID saved. Ready to upload." + ) + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + elif auth_method == "service_account": + if os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON") and not service_account_json: + _persist_env_writes([("GOOGLE_DRIVE_FOLDER_ID", folder_id)]) + message = ( + "Service account already configured via host environment. " + "Folder ID saved. Ready to upload." + ) + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + # ----- Anything below this line writes credential state ----- + if not _setup_opt_in_enabled(): + message = _SETUP_OPT_IN_HELP + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + # ----- Generate or reuse encryption key ----- + encryption_key = None + if encryption_key_method != "off": + existing_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") ) - - from .gdrive_auth_manager import get_oauth_credentials_for_setup - - if not auth_code: - # First run - generate auth URL - result = get_oauth_credentials_for_setup(client_id, client_secret) - - if 'auth_url' in result: - auth_url = result['auth_url'] - - print("\n" + "="*60) - print("[DriveSend Setup] OAuth Authorization Required") - print("="*60) - print("\nClick the URL below to authorize DriveSend:\n") - print(auth_url) - print("\n" + "="*60) - print("After authorizing, copy the authorization code and paste") - print("it into the 'auth_code' field, then run this node again.") - print("="*60 + "\n") - - return ( - f"๐Ÿ” OAuth Step 1: Authorization Required\n\n" - f"1. Click the URL in the terminal/console\n" - f"2. Sign in to Google and authorize DriveSend\n" - f"3. Copy the authorization code shown\n" - f"4. Paste it in the 'auth_code' field\n" - f"5. Run this node again", + if existing_key: + encryption_key = existing_key + print("[DriveSend Setup] Using existing encryption key from environment") + else: + encryption_key = Fernet.generate_key().decode("utf-8") + print("[DriveSend Setup] Generated new encryption key") + + # ----- Service Account flow ----- + if auth_method == "service_account": + return self._handle_service_account( + folder_id, storage_method, encryption_key_method, + encryption_key, service_account_json, + ) + + # ----- OAuth flow ----- + return self._handle_oauth( + folder_id, storage_method, encryption_key_method, + encryption_key, google_client_id, google_client_secret, + auth_code, sid, + ) + + except Exception as e: + logger.exception("[DriveSend Setup] Unexpected error") + message = f"Setup failed with an unexpected error: {e}" + return {"ui": {"text": [message]}, "result": (message,)} + + # ----------------------------------------------------------------- + # OAuth path + # ----------------------------------------------------------------- + + def _handle_oauth(self, folder_id, storage_method, encryption_key_method, + encryption_key, google_client_id, google_client_secret, + auth_code, sid): + if not google_client_id or not google_client_secret: + # No client_id / client_secret in stash. Tell the user to + # click "Set credentials..." on the node. + message = ( + "OAuth requires google_client_id and google_client_secret.\n\n" + "Click 'Set credentials...' on this node, paste your OAuth " + "Client ID and Client Secret from Google Cloud Console, leave " + "the auth code blank, click Save, then click Queue.\n\n" + "Get a Client ID and Client Secret here:\n" + " Google Cloud Console -> APIs & Services -> Credentials\n" + " Create OAuth 2.0 Client ID (Desktop app)" + ) + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + from .gdrive_auth_manager import get_oauth_credentials_for_setup + + if not auth_code: + # First run: generate the auth URL and ask the browser to + # auto-open the entry modal so the user can paste the auth + # code on the next round. + result = get_oauth_credentials_for_setup(google_client_id, google_client_secret) + if "auth_url" not in result: + err = result.get("error", "Unknown error") + message = f"Error generating auth URL: {err}" + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + auth_url = result["auth_url"] + + # The auth URL contains your OAuth client_id, which Google + # treats as a public identifier (it's in every OAuth URL + # ever issued by your app). Safe to print. + print() + print("=" * 80) + print("DRIVESEND - GOOGLE DRIVE AUTHORIZATION REQUIRED") + print("=" * 80) + print("Open this URL in your browser to authorize:") + print() + print(f" {auth_url}") + print() + print("After authorizing, Google will display an authorization code.") + print("Paste it into the credentials modal that auto-opens in your") + print("browser, then click Save and re-queue this workflow.") + print("=" * 80) + print() + + # Auto-trigger the entry modal in the originating browser. + try: + from server import PromptServer + if sid: + PromptServer.instance.send_sync( + "drivesend_credentials_needed", + {"client_id": sid, "stage": "auth_code"}, + sid, + ) + print("[DriveSend Setup] Sent credentials-needed event to browser") + except Exception as e: + print(f"[DriveSend Setup] Warning: could not send credentials-needed event: {e}") + + message = ( + "Google OAuth Ready!\n\n" + "Click the link below (also printed in the ComfyUI terminal):\n\n" + f"{auth_url}\n\n" + "Authorize, copy the code Google shows you, then paste your " + "Client ID, Client Secret, and Auth Code in the modal that " + "auto-opens. Click Save, then click Queue." + ) + return {"ui": {"text": [message]}, "result": (message,)} + + # Second run: exchange the auth code. + result = get_oauth_credentials_for_setup( + google_client_id, google_client_secret, auth_code + ) + if "error" in result: + message = ( + f"OAuth exchange failed: {result['error']}\n\n" + "The auth code may have expired or already been used.\n" + "Click 'Set credentials...' again to enter a fresh one, " + "or rerun the node with no auth code to regenerate the URL." + ) + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + creds = result.get("credentials", {}) + refresh_token = creds.get("refresh_token") + if not refresh_token: + message = ( + "OAuth exchange returned no refresh token. This usually means " + "the user did not see the consent screen (already authorized). " + "Revoke access at https://myaccount.google.com/permissions and " + "rerun setup." + ) + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + return self._deliver_credentials( + auth_method="oauth", + folder_id=folder_id, + storage_method=storage_method, + encryption_key_method=encryption_key_method, + encryption_key=encryption_key, + sid=sid, + credentials_payload=[ + ("GOOGLE_CLIENT_ID", google_client_id), + ("GOOGLE_CLIENT_SECRET", google_client_secret), + ("GOOGLE_REFRESH_TOKEN", refresh_token), + ], + ) + + # ----------------------------------------------------------------- + # Service Account path + # ----------------------------------------------------------------- + + def _handle_service_account(self, folder_id, storage_method, encryption_key_method, + encryption_key, service_account_json): + # Three sources of service account JSON, checked in order: + # 1. Stashed by the browser modal (paste into the dedicated field) + # 2. service_account.json in the plugin directory (legacy) + # 3. nothing -> tell the user how to provide it + sa_data = None + sa_source = None + + if service_account_json: + try: + sa_data = json.loads(service_account_json) + sa_source = "modal" + except json.JSONDecodeError as e: + message = f"Service account JSON could not be parsed: {e}" + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + if sa_data is None: + sa_file = NODE_DIR / "service_account.json" + if sa_file.exists(): + try: + with open(sa_file, "r") as f: + sa_data = json.load(f) + sa_source = "file" + except (OSError, json.JSONDecodeError) as e: + message = f"Could not read service_account.json: {e}" + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + if sa_data is None: + message = ( + "No service account configured. Two options:\n\n" + "(a) Click 'Set credentials...' on this node and paste the " + "JSON contents of your service-account key file into the " + "Service Account JSON field.\n\n" + "(b) Place the file at " + "ComfyUI/custom_nodes/ComfyUI_DriveSendNode/service_account.json\n\n" + "Note: service accounts only work with Google Workspace (paid). " + "For personal Gmail accounts, use OAuth instead." + ) + print(f"[DriveSend Setup] {message}") + return {"ui": {"text": [message]}, "result": (message,)} + + sa_email = sa_data.get("client_email", "unknown") + sa_json_b64 = base64.b64encode(json.dumps(sa_data).encode()).decode() + print(f"[DriveSend Setup] Service account loaded ({sa_source}): {sa_email}") + + return self._deliver_credentials( + auth_method="service_account", + folder_id=folder_id, + storage_method=storage_method, + encryption_key_method=encryption_key_method, + encryption_key=encryption_key, + sid=None, + credentials_payload=[("GOOGLE_SERVICE_ACCOUNT_JSON", sa_json_b64)], + extra_message=( + f"Service account email: {sa_email}\n" + f"Make sure your Drive folder is shared with that email " + f"(otherwise uploads will fail with a 'file not found' error)." + ), + ) + + # ----------------------------------------------------------------- + # Credential delivery (env_file or display_only) + # ----------------------------------------------------------------- + + def _deliver_credentials(self, auth_method, folder_id, storage_method, + encryption_key_method, encryption_key, sid, + credentials_payload, extra_message=None): + # folder_id is non-sensitive and always persisted to .env so + # the AutoUploader can find it. + env_writes = [("GOOGLE_DRIVE_FOLDER_ID", folder_id)] + ws_payload = {} + + if storage_method == "env_file": + env_writes.extend(credentials_payload) + else: + for name, value in credentials_payload: + ws_payload[name] = value + + if encryption_key_method == "save_to_env" and encryption_key: + env_writes.append(("COMFYUI_ENCRYPTION_KEY", encryption_key)) + elif encryption_key_method == "display_only" and encryption_key: + ws_payload["COMFYUI_ENCRYPTION_KEY"] = encryption_key + + # Persist .env. Race-free 0o600 open. Merges with existing + # values so re-running setup does not wipe non-setup keys + # (uploader settings, etc.). + env_path = _persist_env_writes(env_writes, scrub_setup_keys=True) + + # OAuth: also write token.json for local CLI use, with secure + # permissions. Only if storage_method is env_file (otherwise + # display_only's whole point is "no secrets on disk"). + if auth_method == "oauth" and storage_method == "env_file": + creds_dict = {name: value for name, value in credentials_payload} + if all(k in creds_dict for k in ("GOOGLE_CLIENT_ID", "GOOGLE_CLIENT_SECRET", "GOOGLE_REFRESH_TOKEN")): + _write_token_json( + client_id=creds_dict["GOOGLE_CLIENT_ID"], + client_secret=creds_dict["GOOGLE_CLIENT_SECRET"], + refresh_token=creds_dict["GOOGLE_REFRESH_TOKEN"], + ) + + # WebSocket delivery for display_only (sid-targeted, never broadcast). + ws_delivered = False + ws_delivery_error = None + if ws_payload: + try: + from server import PromptServer + deliver_sid = sid or getattr(PromptServer.instance, "client_id", None) + if not deliver_sid: + ws_delivery_error = ( + "prompt has no client_id (typical for CLI/curl/SDK " + "submissions). display_only delivery requires " + "submitting from the ComfyUI web browser so a " + "client_id is attached to the prompt." + ) + print( + "[DriveSend Setup] Refusing to deliver credentials: " + "prompt has no client_id. Sending to sid=None would " + "broadcast to every connected browser. Submit from " + "the ComfyUI web UI instead." ) else: - return (f"โŒ Error generating auth URL: {result.get('error', 'Unknown error')}",) - - else: - # Second run - exchange code for tokens - result = get_oauth_credentials_for_setup(client_id, client_secret, auth_code) - - if 'error' in result: - return ( - f"โŒ Error: {result['error']}\n\n" - f"The auth code may have expired or already been used.\n" - f"To fix:\n" - f"1. Clear the 'auth_code' field (delete the code)\n" - f"2. Run this node again to get a new URL\n" - f"3. Click the URL and authorize again\n" - f"4. Paste the new auth code", + PromptServer.instance.send_sync( + "drivesend_credentials_ready", + {"credentials": dict(ws_payload), "client_id": deliver_sid}, + deliver_sid, ) - - if 'credentials' in result: - creds = result['credentials'] - - env_vars['GOOGLE_CLIENT_ID'] = creds['client_id'] - env_vars['GOOGLE_CLIENT_SECRET'] = creds['client_secret'] - env_vars['GOOGLE_REFRESH_TOKEN'] = creds['refresh_token'] - - results.append("โœ“ OAuth authorization successful!") - results.append(f"โœ“ Refresh token obtained") - - # === Service Account Setup === - elif auth_method == 'service_account': - sa_file = NODE_DIR / 'service_account.json' - - if not sa_file.exists(): - return ( - "โŒ Error: service_account.json not found.\n\n" - "To create one:\n" - "1. Go to Google Cloud Console โ†’ IAM & Admin โ†’ Service Accounts\n" - "2. Create a service account\n" - "3. Create a JSON key\n" - "4. Rename to 'service_account.json'\n" - "5. Place in: ComfyUI/custom_nodes/ComfyUI_DriveSendNode/\n\n" - "โš ๏ธ NOTE: Service accounts only work with Google Workspace (paid).\n" - "For personal Gmail accounts, use OAuth instead.", - ) - - # Read and encode service account JSON - try: - with open(sa_file, 'r') as f: - sa_data = json.load(f) - - sa_email = sa_data.get('client_email', 'unknown') - sa_json_b64 = base64.b64encode(json.dumps(sa_data).encode()).decode() - - env_vars['GOOGLE_SERVICE_ACCOUNT_JSON'] = sa_json_b64 - - results.append("โœ“ Service account loaded") - results.append(f"โœ“ Service account email: {sa_email}") - results.append("") - results.append("โš ๏ธ IMPORTANT: Share your Google Drive folder with:") - results.append(f" {sa_email}") - results.append("") - results.append("โš ๏ธ NOTE: Service accounts require Google Workspace.") - results.append(" They do NOT work with personal Gmail accounts!") - + ws_delivered = True except Exception as e: - return (f"โŒ Error reading service_account.json: {e}",) - - # === Output Results === - print("\n" + "="*60) - print("[DriveSend Setup] Configuration Complete!") - print("="*60) - print("\nCopy these values to your Environment Variables:\n") - - # Print FULL values (not truncated) - for key, value in env_vars.items(): - print(f"{key}={value}") - - print("\n" + "="*60) - print("Copy the values above to your RunPod Secrets") - print("="*60 + "\n") - - # Save token.json for local use - if auth_method == 'oauth' and 'GOOGLE_REFRESH_TOKEN' in env_vars: - token_file = NODE_DIR / 'token.json' - try: - token_data = { - "token": None, - "refresh_token": env_vars['GOOGLE_REFRESH_TOKEN'], - "token_uri": "https://oauth2.googleapis.com/token", - "client_id": env_vars['GOOGLE_CLIENT_ID'], - "client_secret": env_vars['GOOGLE_CLIENT_SECRET'], - "scopes": ["https://www.googleapis.com/auth/drive.file"] - } - with open(token_file, 'w') as f: - json.dump(token_data, f, indent=2) - results.append(f"โœ“ Saved token.json for local use") - except Exception as e: - results.append(f"โš ๏ธ Could not save token.json: {e}") - - # Save to .env file if requested - if storage_method == 'env_file': - env_file = NODE_DIR / '.env' - try: - with open(env_file, 'w') as f: - for key, value in env_vars.items(): - f.write(f"{key}={value}\n") - results.append(f"โœ“ Saved to {env_file}") - except Exception as e: - results.append(f"โš ๏ธ Could not save .env file: {e}") - - # Build status message - status_lines = [ - "โœ… DriveSend Setup Complete!", - "", - f"Auth Method: {auth_method}", - f"Folder ID: {folder_id}", - f"Encryption: {'Enabled' if encryption_key else 'Disabled'}", - "", - ] - status_lines.extend(results) - status_lines.extend([ - "", - "Check the console/terminal for credentials to copy.", - ]) - - if storage_method == 'display_only': - status_lines.extend([ - "", - "For RunPod: Copy the values from console to RunPod Secrets.", - ]) - - return ("\n".join(status_lines),) + ws_delivery_error = str(e) + print(f"[DriveSend Setup] WebSocket delivery failed: {e}") + + # Build status message. Contains NO credential values. + message_parts = ["Google Drive connected."] + if extra_message: + message_parts.append(extra_message) + + if ws_payload and ws_delivered: + message_parts.append( + "Next steps:\n" + " 1. Copy each value from the DriveSend Credentials panel " + "to a safe place (password manager, secure note).\n" + " 2. Add them to your platform's secrets configuration " + "(RunPod Secrets, Docker env, systemd EnvironmentFile).\n" + " 3. Restart your pod / container so the new secrets are " + "visible to the ComfyUI process.\n" + " 4. On the restarted pod, remove this Setup node from " + "your workflow and run the 'DriveSend - AutoUploader' node." + ) + elif ws_payload and not ws_delivered: + message_parts.append( + f"Browser delivery refused: {ws_delivery_error}\n\n" + f"What to do:\n" + f" - Cloud / RunPod / hosted ComfyUI: Submit this workflow " + f"from your web browser (not via curl / API / SDK). The " + f"browser's WebSocket connection attaches a client_id " + f"automatically. display_only never writes credentials to " + f"the pod's filesystem.\n" + f" - Local install only: switch storage_method to " + f"'env_file' and rerun if you don't mind credentials in a " + f"0600 .env file on disk." + ) + elif env_writes: + message_parts.append( + "You can remove this Setup node from your workflow and run " + "the 'DriveSend - AutoUploader' node." + ) + + message = "\n\n".join(message_parts) + + # Console banner. Confirmation only, values redacted. + print("\n" + "=" * 80) + print("DRIVESEND SETUP COMPLETE (values redacted in console)") + if env_writes: + print(f"Wrote to file: {env_path}") + for name, value in env_writes: + print(" " + _redact_kv_line(f"{name}={value}")) + if ws_payload: + label = "Sent to browser (display_only):" + print(label) + for name in ws_payload: + # Field names only for display_only. The values rode the + # WS payload, never stdout. + print(f" {name}=") + print("=" * 80 + "\n") + + return {"ui": {"text": [message]}, "result": (message,)} +# --------------------------------------------------------------------------- +# Persistence helpers (race-free, 0o600) +# --------------------------------------------------------------------------- + +# Keys the Setup Node "owns". On a setup re-run that does NOT write a +# given key, it gets scrubbed from .env so old values don't linger when +# switching env_file -> display_only. GOOGLE_DRIVE_FOLDER_ID is owned by +# the AutoUploader, so it's not in this set. +_SETUP_KEYS = { + "GOOGLE_CLIENT_ID", + "GOOGLE_CLIENT_SECRET", + "GOOGLE_REFRESH_TOKEN", + "GOOGLE_SERVICE_ACCOUNT_JSON", + "COMFYUI_ENCRYPTION_KEY", +} + + +def _persist_env_writes(env_writes, scrub_setup_keys=False): + """Write (key, value) pairs to .env with 0o600 permissions. + + Merges with existing .env contents so non-setup keys (uploader + settings, etc.) survive re-runs unchanged. If scrub_setup_keys is + True, setup-owned keys NOT being rewritten by this run are dropped + from disk (covers the env_file -> display_only switch). + """ + env_path = NODE_DIR / ".env" + + existing = {} + if env_path.exists(): + existing = dict(dotenv_values(env_path)) + + if scrub_setup_keys: + existing = {k: v for k, v in existing.items() if k not in _SETUP_KEYS} + + merged = dict(existing) + for name, value in env_writes: + merged[name] = value + + 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 freshly written values into os.environ so the + # AutoUploader can use them in the same session. + load_dotenv(env_path, override=True) + elif env_path.exists(): + os.remove(env_path) + + return env_path + + +def _write_token_json(client_id, client_secret, refresh_token): + """Write token.json with 0o600 permissions for local CLI use.""" + token_path = NODE_DIR / "token.json" + token_data = { + "token": None, + "refresh_token": refresh_token, + "token_uri": "https://oauth2.googleapis.com/token", + "client_id": client_id, + "client_secret": client_secret, + "scopes": ["https://www.googleapis.com/auth/drive.file"], + } + fd = os.open(str(token_path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as f: + json.dump(token_data, f, indent=2) + + +# --------------------------------------------------------------------------- # Node registration +# --------------------------------------------------------------------------- NODE_CLASS_MAPPINGS = { "DriveSendSetup": DriveSendSetupNode } diff --git a/drivesend_uploader_node.py b/drivesend_uploader_node.py index 347c18c..4fcf4c9 100644 --- a/drivesend_uploader_node.py +++ b/drivesend_uploader_node.py @@ -6,167 +6,227 @@ 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.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s', - handlers=[ - logging.FileHandler('drivesend.log'), - logging.StreamHandler() - ] -) +# 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__) -# Define encrypt_observer at module level +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): - # Try to get default watch folder - default_watch = os.path.join(os.getcwd(), "output") - + # 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": "oauth"}), - "enable_encryption": ("BOOLEAN", {"default": True}), - "Post_Delete_Enc": ("BOOLEAN", {"default": False}), - "Subfolder_Monitor": ("BOOLEAN", {"default": True}), - "run_process": ("BOOLEAN", {"default": True, "label": "Run Process"}), + "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, + + 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 - - This preserves original files so ComfyUI maintains sequential naming. + 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(f"[DriveSend] Starting AutoUploader: watch_folder={watch_folder}, " - f"auth_method={auth_method}, encryption={enable_encryption}, " - f"Post_Delete_Enc={Post_Delete_Enc}, Subfolder_Monitor={Subfolder_Monitor}, " - f"run_process={run_process}") - # Handle stopping the process + 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("[DriveSend] Stopping AutoUploader monitoring") - - # Stop upload watcher + logger.info("Stopping DriveSend AutoUploader monitoring") + if watcher_observer and watcher_observer.is_alive(): stop_monitoring() - logger.info("[DriveSend] Upload watcher stopped") - - # Stop encryption watcher + logger.info("Upload watcher stopped") + if encrypt_observer and encrypt_observer.is_alive(): encrypt_observer.stop() encrypt_observer.join() - logger.info("[DriveSend] Encryption watcher stopped") - - # Stop queue processors + logger.info("Encryption watcher stopped") + stop_queue_processor() stop_encrypt_queue_processor() - - stop_message = f""" -===================================================================== -๐Ÿš™๐Ÿ›‘ DriveSend - AutoUploader - 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. -===================================================================== -""" + + 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 + # ----- 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"[DriveSend] Watch folder does not exist: {watch_folder}") - return (f"โŒ Error: Watch folder does not exist: {watch_folder}",) + 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"[DriveSend] Watch folder is not a directory: {watch_folder}") - return (f"โŒ Error: Watch folder is not a directory: {watch_folder}",) + logger.error(f"Watch folder is not a directory: {watch_folder}") + return (f"Error: Watch folder is not a directory: {watch_folder}",) - # Get folder ID from environment - folder_id = os.environ.get('GOOGLE_DRIVE_FOLDER_ID') + # ----- 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" + "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", + " 1. Copy credentials to RunPod Secrets / your secrets manager\n" + " 2. Restart your pod / container so the env vars load", ) - # 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') - + # ----- 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 = Path(__file__).parent / 'token.json' + token_file = NODE_DIR / "token.json" if not token_file.exists(): return ( - "โŒ Error: OAuth credentials not found.\n\n" + "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" + " - 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' - + + 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" + "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" + " 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.", ) - # Validate encryption key if enabled + # ----- Encryption key requirement ----- if enable_encryption: enc_key = get_encryption_key() if not enc_key: - logger.error("[DriveSend] Encryption enabled but no encryption key found") + logger.error("Encryption enabled but no encryption key found") return ( - "โŒ Error: Encryption enabled but no key found.\n\n" + "Error: Encryption enabled but no key found.\n\n" "Set COMFYUI_ENCRYPTION_KEY environment variable\n" "or run DriveSend Setup with encryption enabled.", ) - # Reset queue processors before starting + # ----- 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 + # ----- Start upload monitor ----- try: start_monitoring( watch_folder=watch_folder, @@ -174,24 +234,23 @@ Set 'run_process' to True and run the node again to resume. auth_method=auth_method, enable_encryption=enable_encryption, delete_enc=Post_Delete_Enc, - subfolder_monitor=Subfolder_Monitor + subfolder_monitor=Subfolder_Monitor, ) except Exception as e: - logger.error(f"[DriveSend] Failed to start upload monitor: {e}") - return (f"โŒ Error starting upload monitor: {e}",) + logger.error(f"Failed to start upload monitor: {e}") + return (f"Error starting upload monitor: {e}",) - # Start encryption watcher if enabled (separate from upload watcher) + # ----- Start encryption watcher if enabled ----- if enable_encryption: - # Stop existing encrypt observer if running 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"[DriveSend] Starting encryption monitor for {watch_folder}") + logger.info(f"Starting encryption monitor for {watch_folder}") def keep_encrypt_alive(): try: @@ -205,35 +264,37 @@ Set 'run_process' to True and run the node again to resume. threading.Thread(target=keep_encrypt_alive, daemon=True).start() - # Build status message + # ----- Status message (no credentials) ----- status_lines = [ - "โœ… DriveSend AutoUploader started!", - "", - f"๐Ÿ“ Watching: {watch_folder}", - f"โ˜๏ธ Uploading to: Google 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'}", + "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 output folder.", "", + "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.", ]) - return ("\n".join(status_lines),) + banner = "\n".join(status_lines) + print(banner) + logger.info(banner) + return (banner,) -# Node registration NODE_CLASS_MAPPINGS = { "DriveSendAutoUploader": DriveSendAutoUploaderNode } diff --git a/encrypt_file.py b/encrypt_file.py index cf3b1bb..1a306dd 100644 --- a/encrypt_file.py +++ b/encrypt_file.py @@ -10,14 +10,8 @@ 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() - ] -) +# logging is configured centrally in __init__.py via a FileHandler on +# the root logger. logger = logging.getLogger(__name__) ENCRYPT_EXTENSIONS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.mp4', '.avi', '.mov') @@ -75,9 +69,14 @@ class FileEncryptHandler(FileSystemEventHandler): 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) + if event.is_directory or not event.src_path.lower().endswith(ENCRYPT_EXTENSIONS): + return + from .safe_paths import is_safe_event_path + if not is_safe_event_path(event.src_path): + logger.warning(f"[DriveSend] Skipping unsafe path (symlink or outside output root): {event.src_path}") + return + 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(): @@ -119,16 +118,3 @@ class FileEncryptHandler(FileSystemEventHandler): 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: - while True: - time.sleep(1) - except KeyboardInterrupt: - observer.stop() - observer.join() diff --git a/gdrive_auth_manager.py b/gdrive_auth_manager.py index b4a9762..c078f7a 100644 --- a/gdrive_auth_manager.py +++ b/gdrive_auth_manager.py @@ -106,8 +106,10 @@ def get_oauth_credentials(): credentials = Credentials.from_authorized_user_file(str(TOKEN_FILE), SCOPES) if credentials and credentials.expired and credentials.refresh_token: credentials.refresh(Request()) - # Save refreshed token - with open(TOKEN_FILE, 'w') as token: + # Save refreshed token. Race-free 0o600 open so the file + # is not world-readable in the open->chmod window. + fd = os.open(str(TOKEN_FILE), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, 'w') as token: token.write(credentials.to_json()) print("[DriveSend] OAuth credentials loaded and refreshed from token file") elif credentials and credentials.valid: @@ -128,9 +130,10 @@ def get_oauth_credentials(): str(CLIENT_SECRET_FILE), SCOPES ) credentials = flow.run_local_server(port=0) - - # Save the credentials for future use - with open(TOKEN_FILE, 'w') as token: + + # Save the credentials for future use. Race-free 0o600 open. + fd = os.open(str(TOKEN_FILE), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, 'w') as token: token.write(credentials.to_json()) print("[DriveSend] โœ“ OAuth authorization complete") diff --git a/monitor_output.py b/monitor_output.py index 07b5a6c..362a98f 100644 --- a/monitor_output.py +++ b/monitor_output.py @@ -11,14 +11,9 @@ from watchdog.events import FileSystemEventHandler from .gdrive_upload import upload_file from .encrypt_file import ENCRYPT_EXTENSIONS -logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s', - handlers=[ - logging.FileHandler('drivesend.log'), - logging.StreamHandler() - ] -) +# logging is configured centrally in __init__.py via a FileHandler on +# the root logger. logging.getLogger here returns a child logger that +# propagates records into that handler. logger = logging.getLogger(__name__) watcher_observer = None @@ -85,12 +80,17 @@ class NewFileHandler(FileSystemEventHandler): if event.is_directory: return file_path = event.src_path - + + from .safe_paths import is_safe_event_path + if not is_safe_event_path(file_path): + logger.warning(f"[DriveSend] Skipping unsafe path (symlink or outside output root): {file_path}") + return + # 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 - + # When encryption is disabled, skip .enc files and only process supported extensions if not encryption_enabled: if file_path.lower().endswith('.enc'): @@ -99,7 +99,7 @@ class NewFileHandler(FileSystemEventHandler): if not file_path.lower().endswith(ENCRYPT_EXTENSIONS): logger.debug(f"[DriveSend] Skipping unsupported file type: {file_path}") return - + logger.info(f"[DriveSend] Detected new file: {file_path}") self.file_queue.put(file_path) diff --git a/pyproject.toml b/pyproject.toml index d937349..336482f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [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.1" +version = "1.1.0" license = { file = "LICENSE" } dependencies = ["requests", "python-dotenv", "watchdog", "cryptography", "google-auth", "google-auth-oauthlib", "google-api-python-client"] diff --git a/safe_paths.py b/safe_paths.py new file mode 100644 index 0000000..c0c3c29 --- /dev/null +++ b/safe_paths.py @@ -0,0 +1,139 @@ +# safe_paths.py +# +# Path validation helpers used by DriveSend nodes to constrain user-supplied +# paths to safe locations. +# +# Threat model: a ComfyUI instance may be reachable over the network. Any +# workflow author can therefore set node inputs. Without restriction, the +# DriveSend uploader's watch_folder input becomes an arbitrary-file-read +# primitive. The watcher recursively uploads everything under the path to +# the operator's Google Drive account. We constrain watch_folder to the +# ComfyUI output directory tree, with an operator-controlled allowlist for +# advanced setups. + +import os +import logging + +logger = logging.getLogger(__name__) + + +def get_output_root(): + """Return the absolute, real path of ComfyUI's output directory. + + Falls back to /output if folder_paths is unavailable (e.g. running + unit tests outside a ComfyUI process). + """ + try: + import folder_paths # type: ignore + root = folder_paths.get_output_directory() + except Exception: + root = os.path.join(os.getcwd(), "output") + return os.path.realpath(root) + + +def get_allowed_roots(): + """Return the list of directory roots a watch_folder is allowed to be in. + + Always includes the ComfyUI output directory. Operators who need to + monitor additional locations (e.g. a custom output mount) can opt in by + setting COMFYUI_DRIVESEND_ALLOWED_WATCH_PATHS to an os.pathsep-separated + list of absolute paths. Workflow inputs cannot expand this list. + """ + roots = [get_output_root()] + extra = os.environ.get("COMFYUI_DRIVESEND_ALLOWED_WATCH_PATHS", "") + if extra: + for raw in extra.split(os.pathsep): + raw = raw.strip() + if raw: + roots.append(os.path.realpath(raw)) + return roots + + +def _is_within(child, parent): + """True iff `child` is `parent` or a descendant of it (after realpath).""" + try: + return os.path.commonpath([child, parent]) == parent + except ValueError: + return False + + +def resolve_safe_watch_folder(user_input): + """Resolve a user-supplied watch_folder to a safe absolute path. + + - Empty input -> the ComfyUI output directory. + - Relative input -> joined under the output directory. + - Absolute input -> must resolve under one of the allowed roots. + + Raises ValueError with a clear message if the input escapes the allowed + set (path traversal, symlinks pointing outside, unrelated absolute path). + """ + output_root = get_output_root() + allowed_roots = get_allowed_roots() + + if user_input is None or str(user_input).strip() == "": + return output_root + + candidate = str(user_input).strip() + + if not os.path.isabs(candidate): + candidate = os.path.join(output_root, candidate) + + resolved = os.path.realpath(candidate) + + for root in allowed_roots: + if _is_within(resolved, root): + return resolved + + raise ValueError( + "watch_folder must be inside the ComfyUI output directory " + f"({output_root}). Got: {user_input!r} (resolved: {resolved}). " + "If you need to monitor a different location, set the " + "COMFYUI_DRIVESEND_ALLOWED_WATCH_PATHS environment variable on the " + "host before starting ComfyUI." + ) + + +def is_safe_event_path(path): + """True iff a Watchdog file event path is safe to process. + + Rejects (a) symlinks and (b) paths whose realpath escapes the allowed + roots. Defends against symlinks placed inside the watched directory + that point at sensitive files outside it. + """ + try: + if os.path.islink(path): + return False + resolved = os.path.realpath(path) + except OSError: + return False + + for root in get_allowed_roots(): + if _is_within(resolved, root): + return True + return False + + +def validate_drive_folder_id(folder_id): + """Validate a Google Drive folder ID string. + + Drive folder IDs are URL-safe base64 strings, typically 28-44 chars. + We reject empty values, NUL bytes, and obvious format violations. + Defense in depth. Drive itself rejects malformed IDs server-side. + """ + if folder_id is None or str(folder_id).strip() == "": + raise ValueError("folder_id must not be empty.") + + value = str(folder_id).strip() + + if "\x00" in value: + raise ValueError("folder_id contains a NUL byte.") + if "/" in value or "\\" in value: + raise ValueError( + f"folder_id must not contain path separators. " + f"Provide just the ID after /folders/ in your Drive URL. " + f"Got: {value!r}" + ) + if len(value) > 128: + raise ValueError(f"folder_id is suspiciously long ({len(value)} chars).") + + return value diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_security.py b/tests/test_security.py new file mode 100644 index 0000000..0194e2f --- /dev/null +++ b/tests/test_security.py @@ -0,0 +1,292 @@ +""" +Security regression tests. + +These tests pin behavior that, if it regresses silently, would +re-introduce a previously-fixed vulnerability. They run without a real +ComfyUI process by stubbing the `server.PromptServer` module. + +Run from the repo root: + python -m pytest tests/ + +Or directly: + python -m unittest tests.test_security +""" + +import importlib +import os +import stat +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + + +REPO_ROOT = Path(__file__).resolve().parent.parent +PARENT_OF_REPO = REPO_ROOT.parent + + +# Install ghost modules for the google-auth / googleapiclient packages +# BEFORE the package under test imports them. This lets the test run +# without those packages installed in the environment. +def _install_lib_ghosts(): + """Install fake modules for runtime deps so the package can import + without google-auth / googleapiclient / watchdog installed.""" + for name in [ + "google", "google.oauth2", "google.auth", "google.auth.transport", + "googleapiclient", "google_auth_oauthlib", + "watchdog", "watchdog.observers", "watchdog.events", + ]: + if name not in sys.modules: + sys.modules[name] = types.ModuleType(name) + sys.modules["google.oauth2.credentials"] = types.ModuleType("google.oauth2.credentials") + sys.modules["google.oauth2.credentials"].Credentials = mock.MagicMock() + sys.modules["google.oauth2.service_account"] = types.ModuleType("google.oauth2.service_account") + sys.modules["google.oauth2.service_account"].Credentials = mock.MagicMock() + sys.modules["google.auth.transport.requests"] = types.ModuleType("google.auth.transport.requests") + sys.modules["google.auth.transport.requests"].Request = mock.MagicMock() + sys.modules["googleapiclient.discovery"] = types.ModuleType("googleapiclient.discovery") + sys.modules["googleapiclient.discovery"].build = mock.MagicMock() + sys.modules["googleapiclient.http"] = types.ModuleType("googleapiclient.http") + sys.modules["googleapiclient.http"].MediaFileUpload = mock.MagicMock() + sys.modules["googleapiclient.errors"] = types.ModuleType("googleapiclient.errors") + sys.modules["googleapiclient.errors"].HttpError = type("HttpError", (Exception,), {}) + sys.modules["google_auth_oauthlib.flow"] = types.ModuleType("google_auth_oauthlib.flow") + sys.modules["google_auth_oauthlib.flow"].Flow = mock.MagicMock() + sys.modules["google_auth_oauthlib.flow"].InstalledAppFlow = mock.MagicMock() + sys.modules["watchdog.observers"].Observer = mock.MagicMock() + sys.modules["watchdog.events"].FileSystemEventHandler = type("FileSystemEventHandler", (object,), {}) + + +_install_lib_ghosts() + + +def _install_fake_promptserver(client_id, sockets=None): + """Install a stubbed `server` module so the Setup Node can import it. + + Returns a `Mock` for `PromptServer.instance.send_sync` so tests can + assert on calls. + """ + fake_module = types.ModuleType("server") + fake_instance = mock.MagicMock() + fake_instance.client_id = client_id + fake_instance.sockets = sockets if sockets is not None else {} + fake_instance.send_sync = mock.MagicMock() + fake_module.PromptServer = mock.MagicMock() + fake_module.PromptServer.instance = fake_instance + sys.modules["server"] = fake_module + return fake_instance.send_sync + + +def _install_fake_gdrive_auth_manager(oauth_response=None): + """Install a fake gdrive_auth_manager module so tests do not require + the google-auth packages to be installed. + """ + fake = types.ModuleType("gdrive_auth_manager") + fake.get_oauth_credentials_for_setup = mock.MagicMock( + return_value=oauth_response or {"credentials": { + "client_id": "stub.apps.googleusercontent.com", + "client_secret": "GOCSPX-stub", + "refresh_token": "1//stub", + "token": "ya29.stub", + "token_uri": "https://oauth2.googleapis.com/token", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + }} + ) + fake.get_drive_service = mock.MagicMock() + fake.get_folder_id = mock.MagicMock(return_value=None) + fake.get_oauth_credentials = mock.MagicMock(return_value=None) + fake.get_service_account_credentials = mock.MagicMock(return_value=None) + fake.refresh_access_token = mock.MagicMock() + sys.modules["gdrive_auth_manager"] = fake + sys.modules["ComfyUI_DriveSendNode.gdrive_auth_manager"] = fake + return fake + + +def _import_setup_node_fresh(): + """Re-import the setup node module as part of the ComfyUI_DriveSendNode + package so relative imports (e.g. `from .gdrive_auth_manager import ...`) + resolve correctly. + """ + pkg_name = "ComfyUI_DriveSendNode" + mod_name = f"{pkg_name}.drivesend_setup_node" + # Strip cached entries so import-time route registration sees the + # current `server` mock. + for cached in list(sys.modules): + if cached == mod_name or cached.startswith(mod_name + "."): + sys.modules.pop(cached, None) + sys.modules.pop(mod_name, None) + if str(PARENT_OF_REPO) not in sys.path: + sys.path.insert(0, str(PARENT_OF_REPO)) + return importlib.import_module(mod_name) + + +class TestWebSocketBroadcastGuard(unittest.TestCase): + """Setup must NEVER call send_sync(..., sid=None) for credential events. + + ComfyUI's send_sync with sid=None broadcasts to every connected + WebSocket client. When PromptServer.instance.client_id is None + (typical for curl/SDK submissions that omit client_id), credentials + delivered via display_only would otherwise leak to any concurrent + listener on /ws. + """ + + def test_no_send_sync_when_client_id_is_none(self): + send_sync = _install_fake_promptserver(client_id=None) + setup_mod = _import_setup_node_fresh() + + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(setup_mod, "NODE_DIR", Path(tmp)): + with mock.patch.dict(os.environ, {"COMFYUI_DRIVESEND_ALLOW_SETUP": "1"}, clear=False): + node = setup_mod.DriveSendSetupNode() + # Force a state-changing path: stash an OAuth attempt + # with a usable client_id+secret+auth_code, then call + # setup() with no real client_id present. + setup_mod._stash_set("any-client", { + "google_client_id": "stub.apps.googleusercontent.com", + "google_client_secret": "GOCSPX-stub", + "auth_code": "code-stub", + "service_account_json": "", + }) + # The fake PromptServer reports client_id=None, so + # _stash_consume(None) returns nothing. The OAuth + # branch will still proceed but find no inputs and + # return a "click Set credentials..." message + # without touching send_sync. + with mock.patch( + "ComfyUI_DriveSendNode.gdrive_auth_manager.get_oauth_credentials_for_setup", + return_value={"credentials": { + "client_id": "stub.apps.googleusercontent.com", + "client_secret": "GOCSPX-stub", + "refresh_token": "1//stub", + "token": "ya29.stub", + "token_uri": "https://oauth2.googleapis.com/token", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + }}, + ): + node.setup( + auth_method="oauth", + folder_id="stub_folder_id", + storage_method="display_only", + encryption_key_method="off", + reconnect=False, + ) + + # Assert: every send_sync call was scoped to a non-None sid. + for call in send_sync.call_args_list: + args, kwargs = call + sid = args[2] if len(args) >= 3 else kwargs.get("sid") + self.assertIsNotNone( + sid, + f"send_sync was called with sid=None: {call}. " + "This would broadcast to every connected WebSocket client.", + ) + + def test_send_sync_called_with_real_sid(self): + """Sanity check: when a real client_id is present, send_sync IS called + and is scoped to that sid.""" + send_sync = _install_fake_promptserver( + client_id="real-sid", sockets={"real-sid": object()} + ) + setup_mod = _import_setup_node_fresh() + + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(setup_mod, "NODE_DIR", Path(tmp)): + with mock.patch.dict(os.environ, {"COMFYUI_DRIVESEND_ALLOW_SETUP": "1"}, clear=False): + setup_mod._stash_set("real-sid", { + "google_client_id": "stub.apps.googleusercontent.com", + "google_client_secret": "GOCSPX-stub", + "auth_code": "code-stub", + "service_account_json": "", + }) + with mock.patch( + "ComfyUI_DriveSendNode.gdrive_auth_manager.get_oauth_credentials_for_setup", + return_value={"credentials": { + "client_id": "stub.apps.googleusercontent.com", + "client_secret": "GOCSPX-stub", + "refresh_token": "1//stub", + "token": "ya29.stub", + "token_uri": "https://oauth2.googleapis.com/token", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + }}, + ): + node = setup_mod.DriveSendSetupNode() + node.setup( + auth_method="oauth", + folder_id="stub_folder_id", + storage_method="display_only", + encryption_key_method="off", + reconnect=False, + ) + + # Expect at least one send_sync delivering the credentials_ready + # event scoped to real-sid. + creds_calls = [ + c for c in send_sync.call_args_list + if c.args and c.args[0] == "drivesend_credentials_ready" + ] + self.assertTrue(creds_calls, "Expected drivesend_credentials_ready to be sent") + for call in creds_calls: + self.assertEqual(call.args[2], "real-sid") + payload = call.args[1] + self.assertEqual(payload.get("client_id"), "real-sid", + "Payload must echo the sid for JS-side defense in depth") + + +class TestEnvFilePermissions(unittest.TestCase): + """The .env file must be created with mode 0o600. Default umask + typically produces 0o644, which is world-readable on multi-tenant + hosts. The Setup and AutoUploader nodes both use os.open with + O_CREAT|0o600 to create the file race-free.""" + + def test_persist_env_writes_creates_0600(self): + _install_fake_promptserver(client_id=None) + setup_mod = _import_setup_node_fresh() + + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(setup_mod, "NODE_DIR", Path(tmp)): + env_path = setup_mod._persist_env_writes( + [("GOOGLE_DRIVE_FOLDER_ID", "stub")], + scrub_setup_keys=False, + ) + self.assertTrue(env_path.exists()) + mode = stat.S_IMODE(env_path.stat().st_mode) + self.assertEqual( + mode, 0o600, + f".env file mode is {oct(mode)}, expected 0o600. " + "Default-umask creation would expose secrets to other local users.", + ) + + +class TestStashHardening(unittest.TestCase): + """The /drivesend/setup/stash route must enforce: client_id pattern, + payload size cap, JSON content type, live-WebSocket gating. These + are the hardening primitives that prevent stash poisoning.""" + + def test_client_id_pattern_rejects_bad_input(self): + _install_fake_promptserver(client_id=None) + setup_mod = _import_setup_node_fresh() + + good = ["abc", "abc-123", "client.id_42", "X" * 128] + bad = ["", "abc 123", "abc/../etc", "X" * 129, "