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
This commit is contained in:
MachinePainting
+41
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
+62
-10
@@ -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']
|
||||
|
||||
+814
-232
File diff suppressed because it is too large
Load Diff
+168
-107
@@ -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
|
||||
}
|
||||
|
||||
+10
-24
@@ -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()
|
||||
|
||||
@@ -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")
|
||||
|
||||
+11
-11
@@ -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)
|
||||
|
||||
|
||||
+1
-1
@@ -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"]
|
||||
|
||||
|
||||
+139
@@ -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 <cwd>/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
|
||||
@@ -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, "<script>"]
|
||||
|
||||
for s in good:
|
||||
self.assertIsNotNone(setup_mod._CLIENT_ID_PATTERN.match(s),
|
||||
f"client_id {s!r} should match pattern but did not")
|
||||
for s in bad:
|
||||
self.assertIsNone(setup_mod._CLIENT_ID_PATTERN.match(s),
|
||||
f"client_id {s!r} should NOT match pattern but did")
|
||||
|
||||
def test_stash_ttl_eviction(self):
|
||||
_install_fake_promptserver(client_id=None)
|
||||
setup_mod = _import_setup_node_fresh()
|
||||
|
||||
# Set a stash entry, then jump time past TTL and confirm it's gone.
|
||||
setup_mod._stash_set("test-sid", {"google_client_id": "x"})
|
||||
with mock.patch("time.time", return_value=9_999_999_999):
|
||||
self.assertIsNone(setup_mod._stash_consume("test-sid"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,571 @@
|
||||
// drivesend_oauth.js - ComfyUI extension for the DriveSend Setup Node.
|
||||
//
|
||||
// Architecture summary
|
||||
// --------------------
|
||||
// The Setup Node deliberately has NO secret-bearing inputs in its
|
||||
// Python INPUT_TYPES. google_client_id, google_client_secret, auth_code,
|
||||
// and the service-account JSON body are entered only via a browser-only
|
||||
// modal launched from the "Set credentials..." button this extension
|
||||
// installs on the node. The modal POSTs the values directly to
|
||||
// /drivesend/setup/stash (a same-origin route registered by
|
||||
// drivesend_setup_node.py) and closes. setup() consumes the stash entry
|
||||
// on the next Queue.
|
||||
//
|
||||
// What this guarantees
|
||||
// - Secrets never enter any LiteGraph widget value, so they cannot be
|
||||
// serialized to workflow JSON, ComfyUI's localStorage auto-save,
|
||||
// PNG metadata, copy-pasted nodes, or PromptServer.history /
|
||||
// /history.
|
||||
// - Browser-native masking (input type=password) covers shoulder-
|
||||
// surfing during typing.
|
||||
// - Modal closes after Save with no in-process JS cache. User pastes
|
||||
// fresh values each time the modal opens.
|
||||
|
||||
import { app } from "../../scripts/app.js";
|
||||
import { api } from "../../scripts/api.js";
|
||||
|
||||
const SETUP_NODE_TYPE = "DriveSendSetup";
|
||||
|
||||
let oauthPopup = null;
|
||||
|
||||
function findSetupNodes() {
|
||||
try {
|
||||
return app.graph.nodes.filter(n => n.type === SETUP_NODE_TYPE);
|
||||
} catch (e) {
|
||||
console.warn("[DriveSendNode] Could not enumerate graph nodes:", e);
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// "Set credentials..." button widget on the Setup Node
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function ensureCredentialButton(node) {
|
||||
if (!node || node.type !== SETUP_NODE_TYPE) return;
|
||||
if (node.__drivesendButtonAdded) return;
|
||||
if (typeof node.addWidget !== "function") {
|
||||
console.warn("[DriveSendNode] node.addWidget unavailable");
|
||||
return;
|
||||
}
|
||||
node.addWidget("button", "Set credentials...", null, () => {
|
||||
openCredentialEntryModal({ reason: "user-clicked", node });
|
||||
});
|
||||
node.__drivesendButtonAdded = true;
|
||||
try {
|
||||
if (typeof node.computeSize === "function" && typeof node.setSize === "function") {
|
||||
node.setSize(node.computeSize());
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn("[DriveSendNode] could not recompute node size:", e);
|
||||
}
|
||||
try { app.graph.setDirtyCanvas(true, true); } catch (_) {}
|
||||
console.log(`[DriveSendNode] Installed Set-credentials button on ${SETUP_NODE_TYPE}`);
|
||||
}
|
||||
|
||||
// Read the auth_method widget value from the given node so the modal
|
||||
// can show only the relevant fields (OAuth vs service account).
|
||||
function getAuthMethod(node) {
|
||||
try {
|
||||
if (!node || !node.widgets) return "oauth";
|
||||
const w = node.widgets.find(x => x.name === "auth_method");
|
||||
return (w && w.value) || "oauth";
|
||||
} catch (e) {
|
||||
return "oauth";
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Credential ENTRY modal
|
||||
// ---------------------------------------------------------------------------
|
||||
//
|
||||
// This is the modal the user types secrets INTO. Distinct from
|
||||
// showCredentialModal below, which displays the credentials delivered
|
||||
// at the end of a successful Setup run.
|
||||
|
||||
function openCredentialEntryModal({ reason, node }) {
|
||||
const existing = document.getElementById("drivesend-credentials-entry-modal");
|
||||
if (existing) {
|
||||
const firstInput = existing.querySelector("input, textarea");
|
||||
if (firstInput) firstInput.focus();
|
||||
return;
|
||||
}
|
||||
|
||||
// If no node was passed (auto-trigger from server event), pick the
|
||||
// first Setup Node in the graph so we can read its auth_method.
|
||||
if (!node) {
|
||||
const nodes = findSetupNodes();
|
||||
node = nodes && nodes.length ? nodes[0] : null;
|
||||
}
|
||||
const authMethod = getAuthMethod(node);
|
||||
|
||||
const overlay = document.createElement("div");
|
||||
overlay.id = "drivesend-credentials-entry-modal";
|
||||
overlay.style.cssText = [
|
||||
"position:fixed", "inset:0",
|
||||
"background:rgba(0,0,0,0.7)",
|
||||
"z-index:10001",
|
||||
"display:flex", "align-items:center", "justify-content:center",
|
||||
'font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif'
|
||||
].join(";");
|
||||
|
||||
const dialog = document.createElement("div");
|
||||
dialog.style.cssText = [
|
||||
"background:#1e1e1e", "color:#eaeaea",
|
||||
"border-radius:12px", "padding:24px",
|
||||
"max-width:560px", "width:90%",
|
||||
"max-height:85vh", "overflow-y:auto",
|
||||
"box-shadow:0 20px 60px rgba(0,0,0,0.5)",
|
||||
"box-sizing:border-box"
|
||||
].join(";");
|
||||
|
||||
const header = document.createElement("div");
|
||||
header.style.cssText = "display:flex;justify-content:space-between;align-items:center;margin-bottom:8px";
|
||||
const title = document.createElement("h2");
|
||||
title.textContent = "Set DriveSend Credentials";
|
||||
title.style.cssText = "margin:0;font-size:1.2rem;color:#fff";
|
||||
const closeX = document.createElement("button");
|
||||
closeX.textContent = "x";
|
||||
closeX.setAttribute("aria-label", "Close");
|
||||
closeX.style.cssText = "background:none;border:none;color:#999;font-size:1.5rem;cursor:pointer;padding:0 8px;line-height:1";
|
||||
closeX.onclick = () => overlay.remove();
|
||||
header.appendChild(title);
|
||||
header.appendChild(closeX);
|
||||
dialog.appendChild(header);
|
||||
|
||||
const intro = document.createElement("p");
|
||||
intro.style.cssText = "color:#bbb;font-size:0.85rem;line-height:1.5;margin:0 0 16px 0";
|
||||
if (authMethod === "service_account") {
|
||||
intro.textContent =
|
||||
"Paste the entire JSON body of your service account key file " +
|
||||
"below. Values are sent to the server (browser only, never " +
|
||||
"written to disk in display_only mode) and are NOT saved by " +
|
||||
"ComfyUI in any workflow JSON, PNG metadata, or browser cache.";
|
||||
} else {
|
||||
intro.textContent =
|
||||
reason === "auto-after-oauth-url"
|
||||
? "Authorization URL printed in the ComfyUI terminal. Authorize at Google, copy the auth code, then paste Client ID + Client Secret + Auth Code below and click Save."
|
||||
: "Paste your Google OAuth Client ID and Client Secret. Leave the Auth Code blank for now; you'll paste it later after authorizing at Google. Values are sent to the server (browser only, never written to disk in display_only mode) and are NOT saved by ComfyUI in any workflow JSON, PNG metadata, or browser cache.";
|
||||
}
|
||||
dialog.appendChild(intro);
|
||||
|
||||
const inputs = {};
|
||||
|
||||
if (authMethod === "service_account") {
|
||||
// One large textarea for the JSON body.
|
||||
const row = document.createElement("div");
|
||||
row.style.cssText = "margin-bottom:12px";
|
||||
const lbl = document.createElement("label");
|
||||
lbl.textContent = "Service Account JSON";
|
||||
lbl.style.cssText = "display:block;font-size:0.8rem;color:#aaa;margin-bottom:4px;letter-spacing:0.02em";
|
||||
row.appendChild(lbl);
|
||||
|
||||
const ta = document.createElement("textarea");
|
||||
ta.autocomplete = "off";
|
||||
ta.spellcheck = false;
|
||||
ta.value = "";
|
||||
ta.rows = 12;
|
||||
ta.style.cssText = "width:100%;box-sizing:border-box;padding:8px;font-family:Menlo,Monaco,monospace;font-size:0.8rem;background:#2a2a2a;color:#eaeaea;border:1px solid #444;border-radius:4px;resize:vertical";
|
||||
ta.placeholder = '{ "type": "service_account", "project_id": "...", ... }';
|
||||
row.appendChild(ta);
|
||||
dialog.appendChild(row);
|
||||
inputs.service_account_json = ta;
|
||||
} else {
|
||||
const fields = [
|
||||
{ name: "google_client_id", label: "Google OAuth Client ID" },
|
||||
{ name: "google_client_secret", label: "Google OAuth Client Secret" },
|
||||
{ name: "auth_code", label: "Auth Code (paste after authorizing at Google)" },
|
||||
];
|
||||
fields.forEach(({ name, label }) => {
|
||||
const row = document.createElement("div");
|
||||
row.style.cssText = "margin-bottom:12px";
|
||||
|
||||
const lbl = document.createElement("label");
|
||||
lbl.textContent = label;
|
||||
lbl.style.cssText = "display:block;font-size:0.8rem;color:#aaa;margin-bottom:4px;letter-spacing:0.02em";
|
||||
row.appendChild(lbl);
|
||||
|
||||
const input = document.createElement("input");
|
||||
input.type = "password";
|
||||
input.autocomplete = "off";
|
||||
input.spellcheck = false;
|
||||
input.value = "";
|
||||
input.style.cssText = "width:100%;box-sizing:border-box;padding:8px;font-family:Menlo,Monaco,monospace;font-size:0.9rem;background:#2a2a2a;color:#eaeaea;border:1px solid #444;border-radius:4px";
|
||||
row.appendChild(input);
|
||||
dialog.appendChild(row);
|
||||
inputs[name] = input;
|
||||
});
|
||||
}
|
||||
|
||||
const status = document.createElement("p");
|
||||
status.style.cssText = "color:#f87171;font-size:0.8rem;margin:6px 0 0 0;min-height:1em";
|
||||
status.textContent = "";
|
||||
dialog.appendChild(status);
|
||||
|
||||
const footer = document.createElement("div");
|
||||
footer.style.cssText = "display:flex;gap:10px;margin-top:16px;padding-top:12px;border-top:1px solid #333";
|
||||
|
||||
const cancelBtn = document.createElement("button");
|
||||
cancelBtn.textContent = "Cancel";
|
||||
cancelBtn.style.cssText = "flex:1;padding:9px;background:#444;color:#fff;border:none;border-radius:6px;cursor:pointer;font-size:0.9rem";
|
||||
cancelBtn.onclick = () => overlay.remove();
|
||||
|
||||
const saveBtn = document.createElement("button");
|
||||
saveBtn.textContent = "Save";
|
||||
saveBtn.style.cssText = "flex:1;padding:9px;background:#3b82f6;color:#fff;border:none;border-radius:6px;cursor:pointer;font-size:0.9rem";
|
||||
saveBtn.onclick = async () => {
|
||||
const payload = {
|
||||
google_client_id: (inputs.google_client_id && inputs.google_client_id.value) || "",
|
||||
google_client_secret: (inputs.google_client_secret && inputs.google_client_secret.value) || "",
|
||||
auth_code: (inputs.auth_code && inputs.auth_code.value) || "",
|
||||
service_account_json: (inputs.service_account_json && inputs.service_account_json.value) || "",
|
||||
};
|
||||
const anyValue = Object.values(payload).some(v => v && v.trim());
|
||||
if (!anyValue) {
|
||||
status.textContent = "Nothing to save. Paste values into at least one field.";
|
||||
return;
|
||||
}
|
||||
const clientId = api && api.clientId ? api.clientId : null;
|
||||
if (!clientId) {
|
||||
status.textContent = "Browser not yet connected to ComfyUI. Wait a moment and try again.";
|
||||
return;
|
||||
}
|
||||
saveBtn.disabled = true;
|
||||
cancelBtn.disabled = true;
|
||||
try {
|
||||
const resp = await fetch("/drivesend/setup/stash", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ client_id: clientId, ...payload }),
|
||||
});
|
||||
if (!resp.ok) throw new Error("server returned " + resp.status);
|
||||
} catch (e) {
|
||||
console.error("[DriveSendNode] Stash POST failed:", e);
|
||||
status.textContent = "Could not save: " + e.message;
|
||||
saveBtn.disabled = false;
|
||||
cancelBtn.disabled = false;
|
||||
return;
|
||||
}
|
||||
Object.values(inputs).forEach(i => { if (i) i.value = ""; });
|
||||
overlay.remove();
|
||||
console.log("[DriveSendNode] Stashed credentials via entry modal; click Queue to run setup");
|
||||
};
|
||||
|
||||
footer.appendChild(cancelBtn);
|
||||
footer.appendChild(saveBtn);
|
||||
dialog.appendChild(footer);
|
||||
|
||||
Object.values(inputs).forEach(input => {
|
||||
if (!input) return;
|
||||
input.addEventListener("keydown", (ev) => {
|
||||
if (ev.key === "Enter" && input.tagName !== "TEXTAREA") {
|
||||
ev.preventDefault();
|
||||
saveBtn.click();
|
||||
}
|
||||
if (ev.key === "Escape") {
|
||||
ev.preventDefault();
|
||||
cancelBtn.click();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
overlay.appendChild(dialog);
|
||||
document.body.appendChild(overlay);
|
||||
|
||||
setTimeout(() => {
|
||||
const list = Object.values(inputs).filter(Boolean);
|
||||
const firstEmpty = list.find(i => !i.value);
|
||||
if (firstEmpty) firstEmpty.focus();
|
||||
}, 0);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// OAuth popup window (delegate to Google)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
window.openDriveSendOAuth = function (url) {
|
||||
console.log("[DriveSendNode] Opening OAuth popup:", url);
|
||||
if (oauthPopup && !oauthPopup.closed) oauthPopup.close();
|
||||
const popup = window.open(
|
||||
url,
|
||||
"google_oauth",
|
||||
"width=520,height=720,scrollbars=yes,resizable=yes,status=no,location=no,toolbar=no,menubar=no"
|
||||
);
|
||||
if (popup) {
|
||||
oauthPopup = popup;
|
||||
popup.focus();
|
||||
const checkClosed = setInterval(() => {
|
||||
if (popup.closed) {
|
||||
console.log("[DriveSendNode] OAuth popup closed");
|
||||
clearInterval(checkClosed);
|
||||
oauthPopup = null;
|
||||
}
|
||||
}, 500);
|
||||
return popup;
|
||||
}
|
||||
console.error("[DriveSendNode] OAuth popup blocked by browser?");
|
||||
return null;
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Credential OUTPUT modal (display_only delivery surface)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function showCredentialModal(creds) {
|
||||
const existing = document.getElementById("drivesend-credentials-modal");
|
||||
if (existing) existing.remove();
|
||||
|
||||
const overlay = document.createElement("div");
|
||||
overlay.id = "drivesend-credentials-modal";
|
||||
overlay.style.cssText = [
|
||||
"position:fixed", "inset:0",
|
||||
"background:rgba(0,0,0,0.7)",
|
||||
"z-index:10000",
|
||||
"display:flex", "align-items:center", "justify-content:center",
|
||||
'font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif'
|
||||
].join(";");
|
||||
|
||||
const dialog = document.createElement("div");
|
||||
dialog.style.cssText = [
|
||||
"background:#1e1e1e", "color:#eaeaea",
|
||||
"border-radius:12px", "padding:24px",
|
||||
"max-width:640px", "width:90%", "max-height:85vh",
|
||||
"overflow-y:auto",
|
||||
"box-shadow:0 20px 60px rgba(0,0,0,0.5)",
|
||||
"box-sizing:border-box"
|
||||
].join(";");
|
||||
|
||||
const header = document.createElement("div");
|
||||
header.style.cssText = "display:flex;justify-content:space-between;align-items:center;margin-bottom:12px";
|
||||
const title = document.createElement("h2");
|
||||
title.textContent = "DriveSend Credentials";
|
||||
title.style.cssText = "margin:0;font-size:1.25rem;color:#fff";
|
||||
const closeX = document.createElement("button");
|
||||
closeX.textContent = "x";
|
||||
closeX.setAttribute("aria-label", "Close");
|
||||
closeX.style.cssText = "background:none;border:none;color:#999;font-size:1.5rem;cursor:pointer;padding:0 8px;line-height:1";
|
||||
closeX.onclick = () => overlay.remove();
|
||||
header.appendChild(title);
|
||||
header.appendChild(closeX);
|
||||
dialog.appendChild(header);
|
||||
|
||||
const intro = document.createElement("p");
|
||||
intro.style.cssText = "color:#bbb;font-size:0.9rem;line-height:1.5;margin:0 0 18px 0";
|
||||
intro.textContent =
|
||||
"These values are shown only in this browser and will be discarded when you close " +
|
||||
"this dialog. They are NOT written to disk in this mode. Save each value somewhere " +
|
||||
"safe now (password manager, secure note). Then configure them as your platform's " +
|
||||
"secrets (RunPod Secrets, Docker env, systemd EnvironmentFile). Most cloud platforms " +
|
||||
"require a pod or container restart for new secrets to take effect.";
|
||||
dialog.appendChild(intro);
|
||||
|
||||
const flashCopied = (btn) => {
|
||||
const original = btn.textContent;
|
||||
btn.textContent = "Copied";
|
||||
setTimeout(() => { btn.textContent = original; }, 1200);
|
||||
};
|
||||
|
||||
Object.entries(creds).forEach(([name, value]) => {
|
||||
const row = document.createElement("div");
|
||||
row.style.cssText = "margin-bottom:14px";
|
||||
|
||||
const labelRow = document.createElement("div");
|
||||
labelRow.style.cssText = "display:flex;justify-content:space-between;align-items:center;margin-bottom:4px";
|
||||
const label = document.createElement("label");
|
||||
label.textContent = name;
|
||||
label.style.cssText = "font-weight:600;font-size:0.85rem;color:#aaa;letter-spacing:0.02em";
|
||||
|
||||
const copyBtn = document.createElement("button");
|
||||
copyBtn.textContent = "Copy";
|
||||
copyBtn.style.cssText = "padding:4px 12px;font-size:0.8rem;background:#3b82f6;color:#fff;border:none;border-radius:4px;cursor:pointer";
|
||||
copyBtn.onclick = async () => {
|
||||
try {
|
||||
await navigator.clipboard.writeText(value);
|
||||
flashCopied(copyBtn);
|
||||
} catch (e) {
|
||||
input.select();
|
||||
document.execCommand("copy");
|
||||
flashCopied(copyBtn);
|
||||
}
|
||||
};
|
||||
labelRow.appendChild(label);
|
||||
labelRow.appendChild(copyBtn);
|
||||
|
||||
// Service-account JSON is long; render it as a textarea.
|
||||
const isLong = (value && value.length > 120) || /JSON|SERVICE_ACCOUNT/i.test(name);
|
||||
const input = document.createElement(isLong ? "textarea" : "input");
|
||||
if (!isLong) input.type = "text";
|
||||
input.value = value;
|
||||
input.readOnly = true;
|
||||
if (isLong) input.rows = 5;
|
||||
input.style.cssText = "width:100%;box-sizing:border-box;padding:8px;font-family:Menlo,Monaco,monospace;font-size:0.8rem;background:#2a2a2a;color:#eaeaea;border:1px solid #444;border-radius:4px;resize:vertical";
|
||||
input.onclick = () => input.select();
|
||||
input.onfocus = () => input.select();
|
||||
|
||||
row.appendChild(labelRow);
|
||||
row.appendChild(input);
|
||||
dialog.appendChild(row);
|
||||
});
|
||||
|
||||
const footer = document.createElement("div");
|
||||
footer.style.cssText = "display:flex;gap:10px;margin-top:18px;padding-top:14px;border-top:1px solid #333";
|
||||
|
||||
const copyAllBtn = document.createElement("button");
|
||||
copyAllBtn.textContent = "Copy all as NAME=value";
|
||||
copyAllBtn.style.cssText = "flex:1;padding:10px;background:#3b82f6;color:#fff;border:none;border-radius:6px;cursor:pointer;font-size:0.9rem";
|
||||
copyAllBtn.onclick = async () => {
|
||||
const text = Object.entries(creds).map(([k, v]) => `${k}=${v}`).join("\n");
|
||||
try {
|
||||
await navigator.clipboard.writeText(text);
|
||||
flashCopied(copyAllBtn);
|
||||
} catch (e) {
|
||||
console.error("[DriveSendNode] Clipboard write failed:", e);
|
||||
alert("Clipboard write failed; select and copy values manually.");
|
||||
}
|
||||
};
|
||||
|
||||
const doneBtn = document.createElement("button");
|
||||
doneBtn.textContent = "I have copied them";
|
||||
doneBtn.style.cssText = "flex:1;padding:10px;background:#444;color:#fff;border:none;border-radius:6px;cursor:pointer;font-size:0.9rem";
|
||||
doneBtn.onclick = () => overlay.remove();
|
||||
|
||||
footer.appendChild(copyAllBtn);
|
||||
footer.appendChild(doneBtn);
|
||||
dialog.appendChild(footer);
|
||||
|
||||
overlay.appendChild(dialog);
|
||||
document.body.appendChild(overlay);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension registration
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
app.registerExtension({
|
||||
name: "DriveSendNode.OAuthHandler",
|
||||
|
||||
async setup() {
|
||||
console.log("[DriveSendNode] OAuth extension loaded");
|
||||
|
||||
try {
|
||||
findSetupNodes().forEach(n => ensureCredentialButton(n));
|
||||
} catch (e) {
|
||||
console.warn("[DriveSendNode] setup-time sweep failed:", e);
|
||||
}
|
||||
|
||||
api.addEventListener("drivesend_reconnect_complete", (event) => {
|
||||
console.log("[DriveSendNode] drivesend_reconnect_complete:", event.detail);
|
||||
const data = event.detail;
|
||||
if (data && data.success) {
|
||||
this.handleReconnectSuccess();
|
||||
}
|
||||
});
|
||||
|
||||
api.addEventListener("drivesend_credentials_needed", (event) => {
|
||||
console.log("[DriveSendNode] drivesend_credentials_needed:", event.detail);
|
||||
const data = event.detail || {};
|
||||
const expectedClientId = data.client_id;
|
||||
const ourClientId = api && typeof api.clientId !== "undefined" ? api.clientId : null;
|
||||
if (expectedClientId && ourClientId && expectedClientId !== ourClientId) {
|
||||
console.warn(
|
||||
"[DriveSendNode] Ignoring credentials-needed event: client_id mismatch"
|
||||
);
|
||||
return;
|
||||
}
|
||||
openCredentialEntryModal({ reason: "auto-after-oauth-url" });
|
||||
});
|
||||
|
||||
api.addEventListener("drivesend_credentials_ready", (event) => {
|
||||
console.log("[DriveSendNode] drivesend_credentials_ready");
|
||||
const data = event.detail;
|
||||
if (!data || !data.credentials || typeof data.credentials !== "object") {
|
||||
console.error("[DriveSendNode] credentials_ready payload missing/malformed", event);
|
||||
return;
|
||||
}
|
||||
const expectedClientId = data.client_id;
|
||||
const ourClientId = api && typeof api.clientId !== "undefined" ? api.clientId : null;
|
||||
if (ourClientId == null) {
|
||||
console.warn(
|
||||
"[DriveSendNode] api.clientId not available, cannot defense-in-depth verify; " +
|
||||
"trusting the server's sid-targeted delivery."
|
||||
);
|
||||
} else if (expectedClientId && expectedClientId !== ourClientId) {
|
||||
console.warn(
|
||||
"[DriveSendNode] Ignoring credentials event: client_id mismatch"
|
||||
);
|
||||
return;
|
||||
}
|
||||
const entry = document.getElementById("drivesend-credentials-entry-modal");
|
||||
if (entry) entry.remove();
|
||||
showCredentialModal(data.credentials);
|
||||
});
|
||||
},
|
||||
|
||||
resetReconnectFields() {
|
||||
try {
|
||||
findSetupNodes().forEach(node => {
|
||||
if (!node.widgets) return;
|
||||
const reconnectWidget = node.widgets.find(w => w.name === "reconnect");
|
||||
if (reconnectWidget && reconnectWidget.value === true) {
|
||||
reconnectWidget.value = false;
|
||||
if (node.onWidgetChanged) node.onWidgetChanged("reconnect", false);
|
||||
}
|
||||
});
|
||||
} catch (e) {
|
||||
console.log("[DriveSendNode] Could not reset reconnect fields:", e);
|
||||
}
|
||||
},
|
||||
|
||||
handleReconnectSuccess() {
|
||||
console.log("[DriveSendNode] Reconnect success - refreshing UI to show auth fields");
|
||||
this.resetReconnectFields();
|
||||
const entry = document.getElementById("drivesend-credentials-entry-modal");
|
||||
if (entry) entry.remove();
|
||||
setTimeout(() => {
|
||||
console.log("[DriveSendNode] Reloading after reconnect");
|
||||
window.location.reload();
|
||||
}, 1000);
|
||||
},
|
||||
|
||||
nodeCreated(node) {
|
||||
ensureCredentialButton(node);
|
||||
},
|
||||
|
||||
loadedGraphNode(node, _app) {
|
||||
ensureCredentialButton(node);
|
||||
},
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, _app) {
|
||||
if (nodeData.name !== SETUP_NODE_TYPE) return;
|
||||
console.log(`[DriveSendNode] Registered ${SETUP_NODE_TYPE} enhancement`);
|
||||
|
||||
const origOnNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
nodeType.prototype.onNodeCreated = function () {
|
||||
const r = origOnNodeCreated ? origOnNodeCreated.apply(this, arguments) : undefined;
|
||||
try {
|
||||
ensureCredentialButton(this);
|
||||
} catch (e) {
|
||||
console.warn("[DriveSendNode] ensureCredentialButton via prototype hook failed:", e);
|
||||
}
|
||||
return r;
|
||||
};
|
||||
|
||||
const originalExecuted = nodeType.prototype.onExecuted;
|
||||
nodeType.prototype.onExecuted = function (message) {
|
||||
try {
|
||||
if (message && message.text && message.text[0]) {
|
||||
const text = message.text[0];
|
||||
if (text.includes("Google OAuth Ready!")) {
|
||||
const m = text.match(/(https:\/\/accounts\.google\.com\/o\/oauth2\/auth[^\s]+)/);
|
||||
if (m) {
|
||||
setTimeout(() => window.openDriveSendOAuth(m[1]), 500);
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn("[DriveSendNode] onExecuted hook error:", e);
|
||||
}
|
||||
if (originalExecuted) return originalExecuted.apply(this, arguments);
|
||||
};
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user