From 2c49196561ef38060845d18785adcdf920e07f02 Mon Sep 17 00:00:00 2001 From: MachinePainting Date: Sun, 1 Feb 2026 02:34:43 -0800 Subject: [PATCH] Major update: OAuth support for personal Gmail accounts - Added OAuth 2.0 authentication with auto-refresh tokens - OAuth tokens refresh automatically on pod startup (works like DropSend Node) - Service Account method now clearly marked as Google Workspace only - Updated README with clear auth requirements - Added wait_for_complete_write() for reliable file uploads - Removed broken ownership transfer code - Default auth method changed to OAuth - Better error messages for quota and auth issues --- README.md | 437 +++++++++++++++++-------------------- drivesend_setup_node.py | 380 +++++++++++++++++--------------- drivesend_uploader_node.py | 261 ++++++++++++---------- gdrive_auth_manager.py | 429 ++++++++++++++++++++++-------------- gdrive_upload.py | 138 ++++-------- 5 files changed, 859 insertions(+), 786 deletions(-) diff --git a/README.md b/README.md index deddbca..ace34b2 100644 --- a/README.md +++ b/README.md @@ -1,48 +1,41 @@ # ComfyUI DriveSend Node -A ComfyUI custom node for seamless Google Drive uploads with **optional** encryption capabilities. Automatically upload your ComfyUI output files (images and videos) to Google Drive cloud storage. - -## โš ๏ธ Authentication Methods - READ FIRST - -| Method | Lifespan | Best For | Effort | -|--------|----------|----------|--------| -| **Service Account** | โœ… **Permanent** (never expires) | Cloud/RunPod | Medium setup | -| **OAuth 2.0** | โŒ **7 days** (must re-auth weekly) | Not recommended | Easy setup | - -**We strongly recommend Service Account** for any persistent use. OAuth tokens expire every 7 days in Google's "testing mode" and there is no workaround without paying for Google Workspace or going through Google's app verification process. +A ComfyUI custom node for uploading output files to Google Drive with optional encryption. Automatically monitors your output folder and uploads images/videos to your Google Drive. --- -## ๐Ÿ”„ How It Works +## โš ๏ธ Important: Authentication Requirements -``` -โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” -โ”‚ CLOUD (RunPod, etc.) โ”‚ -โ”‚ โ”‚ -โ”‚ ComfyUI generates files โ”€โ”€โ†’ DriveSend Node โ”€โ”€โ†’ Uploads to Google Drive โ”‚ -โ”‚ (png, mp4, etc.) โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ -โ”‚ โ–ผ โ”‚ -โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ -โ”‚ โ”‚ Encryption OPTIONAL โ”‚ โ”‚ -โ”‚ โ”‚ โ˜ OFF: file.png โ”‚ โ”‚ -โ”‚ โ”‚ โ˜‘ ON: file.png.enc โ”‚ โ”‚ -โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ -โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ - โ”‚ - โ–ผ - ๐Ÿ“ GOOGLE DRIVE - โ”‚ - โ–ผ -โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” -โ”‚ YOUR LOCAL MACHINE โ”‚ -โ”‚ โ”‚ -โ”‚ Google Drive syncs/downloads โ”€โ”€โ†’ If encrypted: Run decrypt script (local) โ”‚ -โ”‚ โ”€โ”€โ†’ file.png (viewable!) โ”‚ -โ”‚ โ”‚ -โ”‚ If not encrypted: Ready to use! โ”‚ -โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ -``` +DriveSend supports two authentication methods, each with different requirements: + +| Method | Account Type | Cost | Token Expiry | Recommended For | +|--------|-------------|------|--------------|-----------------| +| **Service Account** | Google Workspace only | ~$7+/month | Never expires | Business/paid users | +| **OAuth 2.0** | Personal Gmail | Free | Auto-refreshes* | Personal/free users | + +**\* OAuth tokens auto-refresh on each pod startup, effectively working like permanent credentials.** + +### Why These Limitations? + +- **Service Accounts** have 0 GB storage quota on personal Gmail accounts. Google blocks uploads entirely. +- **Google Workspace** ($7+/month) provides Shared Drives where service accounts can upload. +- **OAuth 2.0** works with personal Gmail but requires initial browser authorization. + +--- + +## ๐Ÿ“ค๐Ÿ“ฆ Features + +- **DriveSend AutoUploader Node** - Automatically uploads new files to Google Drive + - Monitors output folder in real-time + - Supports: `.png`, `.jpg`, `.jpeg`, `.webp`, `.gif`, `.mp4`, `.avi`, `.mov` + - Optional AES-128 encryption before upload + - SHA256 checksum verification + - Queue system for reliable uploads + +- **DriveSend Setup Node** - Configures authentication and credentials + - Generates OAuth refresh tokens + - Outputs credentials for RunPod environment variables + - Optional encryption key generation --- @@ -51,243 +44,225 @@ A ComfyUI custom node for seamless Google Drive uploads with **optional** encryp ```bash cd ComfyUI/custom_nodes/ git clone https://github.com/machinepainting/ComfyUI_DriveSendNode.git -cd ComfyUI_DriveSendNode -pip install -r requirements.txt +pip install -r ComfyUI_DriveSendNode/requirements.txt ``` +Restart ComfyUI after installation. + --- -## ๐Ÿ”ง Google Cloud Setup (Service Account - Recommended) +## ๐Ÿ”ง Google Cloud Setup (Required for Both Methods) ### Step 1: Create Google Cloud Project 1. Go to [Google Cloud Console](https://console.cloud.google.com/) -2. Sign in with a **personal Gmail account** (not work/school - they often have restrictions) -3. Click **Select a project** โ†’ **New Project** -4. Name it (e.g., `ComfyUI-DriveSend`) โ†’ **Create** +2. Sign in with your Google account +3. Click **Select a project** (top, next to "Google Cloud" logo) โ†’ **New Project** +4. Name it (e.g., `ComfyUI-DriveSend`) โ†’ Click **Create** +5. Make sure your new project is selected (check dropdown at top) ### Step 2: Enable Google Drive API -1. Go to **APIs & Services** โ†’ **Library** +1. Go to **APIs & Services** โ†’ **Library** (left sidebar) 2. Search for **Google Drive API** -3. Click **Enable** +3. Click on it โ†’ Click **Enable** -### Step 3: Create Service Account +### Step 3: Create OAuth Credentials + +1. Go to **APIs & Services** โ†’ **Credentials** (left sidebar) +2. Click **+ Create Credentials** โ†’ **OAuth client ID** +3. If prompted, configure the OAuth consent screen: + - User Type: **External** โ†’ Click **Create** + - App name: `DriveSend` + - User support email: Your email + - Developer contact: Your email + - Click **Save and Continue** through all steps +4. Back in Credentials, click **+ Create Credentials** โ†’ **OAuth client ID** +5. Application type: **Desktop app** +6. Name: `DriveSend` +7. Click **Create** +8. **Download JSON** and save as `client_secret.json` in the DriveSend node folder: + ``` + ComfyUI/custom_nodes/ComfyUI_DriveSendNode/client_secret.json + ``` + +### Step 4: Add Yourself as Test User + +1. Go to **APIs & Services** โ†’ **OAuth consent screen** +2. Under **Test users**, click **+ Add Users** +3. Add your Gmail address +4. Click **Save** + +--- + +## ๐Ÿ” Option A: OAuth 2.0 Setup (Personal Gmail - FREE) + +**Best for:** Personal Gmail accounts, RunPod users who want free Google Drive uploads. + +### How It Works + +OAuth uses a refresh token that auto-refreshes on every pod startup. You authorize once, and it works indefinitely (as long as you don't revoke access). + +### Local Setup + +1. Add the **DriveSend Setup** node to your workflow +2. Configure: + - `auth_method`: **oauth** + - `storage_method`: **display_only** (for cloud) or **env_file** (for local) + - `encryption_key_method`: Your preference +3. Run the node +4. A browser window opens - sign in and authorize the app +5. Copy the credentials from the console output + +### RunPod Setup + +1. Complete local setup first to get your refresh token +2. In RunPod, create these **Secrets**: + + | Secret Name | Value | + |-------------|-------| + | `GOOGLE_CLIENT_ID` | Your OAuth client ID | + | `GOOGLE_CLIENT_SECRET` | Your OAuth client secret | + | `GOOGLE_REFRESH_TOKEN` | The refresh token from setup | + | `GOOGLE_DRIVE_FOLDER_ID` | Your Drive folder ID | + | `comfyui_encryption_key` | (Optional) Your encryption key | + +3. Add **Environment Variables** to your pod template: + - `GOOGLE_CLIENT_ID` โ†’ `{{ RUNPOD_SECRET_GOOGLE_CLIENT_ID }}` + - `GOOGLE_CLIENT_SECRET` โ†’ `{{ RUNPOD_SECRET_GOOGLE_CLIENT_SECRET }}` + - `GOOGLE_REFRESH_TOKEN` โ†’ `{{ RUNPOD_SECRET_GOOGLE_REFRESH_TOKEN }}` + - `GOOGLE_DRIVE_FOLDER_ID` โ†’ `{{ RUNPOD_SECRET_GOOGLE_DRIVE_FOLDER_ID }}` + - `comfyui_encryption_key` โ†’ `{{ RUNPOD_SECRET_comfyui_encryption_key }}` + +4. Deploy your pod and add the **DriveSend AutoUploader** node with: + - `auth_method`: **oauth** + - `run_process`: **True** + +--- + +## ๐Ÿข Option B: Service Account Setup (Google Workspace - PAID) + +**Best for:** Google Workspace subscribers who want simpler, permanent authentication. + +**Requirements:** +- Google Workspace account (~$7+/month) +- Custom domain (e.g., yourbusiness.com) + +### Why Service Account Requires Google Workspace + +Service accounts have 0 GB storage quota. On personal Gmail, uploads fail immediately with "storageQuotaExceeded". Google Workspace provides Shared Drives where service accounts CAN upload. + +### Setup Steps 1. Go to **IAM & Admin** โ†’ **Service Accounts** 2. Click **+ Create Service Account** -3. Name it (e.g., `comfyui-uploader`) -4. Click **Create and Continue** -5. Skip optional steps โ†’ **Done** +3. Name it (e.g., `comfyui-uploader`) โ†’ Click **Create and Continue** โ†’ **Done** +4. Click on the service account โ†’ **Keys** tab โ†’ **Add Key** โ†’ **Create new key** โ†’ **JSON** +5. Rename downloaded file to `service_account.json` +6. Place in: `ComfyUI/custom_nodes/ComfyUI_DriveSendNode/service_account.json` -### Step 4: Fix Organization Policy (IMPORTANT - New Google Accounts) +### Create a Shared Drive -Google now blocks service account key creation by default. You MUST disable this policy first. +1. In Google Drive, click **Shared drives** (left sidebar) +2. Click **+ New** to create a Shared Drive +3. Add the service account email as a **Content Manager** +4. Get the Shared Drive ID from the URL -#### 4a. Grant Yourself Policy Admin Role +### RunPod Configuration -1. Go to **IAM & Admin** โ†’ **IAM** -2. Click the dropdown at top-left and select your **organization** (your email domain), not the project -3. Click **+ Grant Access** -4. Principal: **your email** -5. Role: search for **Organization Policy Administrator** -6. Click **Save** +Create secrets and environment variables similar to OAuth setup, but use: +- `GOOGLE_SERVICE_ACCOUNT_JSON` (base64-encoded JSON) +- `GOOGLE_DRIVE_FOLDER_ID` (Shared Drive ID) -#### 4b. Enable the Organization Policy API +--- -Open **Cloud Shell** (terminal icon at top-right) and run: +## ๐Ÿ“ Getting Your Folder ID +1. Open Google Drive +2. Navigate to your target folder (or Shared Drive) +3. Look at the URL: + ``` + https://drive.google.com/drive/folders/XXXXXXXXXXXXXXXXX + ``` +4. Copy the long string after `/folders/` - that's your **Folder ID** + +--- + +## ๐Ÿ”’ Encryption (Optional) + +DriveSend supports AES-128 encryption for files before upload. + +### Enable Encryption + +1. In **DriveSend Setup** node, set `encryption_key_method` to **Display Only** or **save to .env** +2. Run the setup node to generate a key +3. Save the key securely (you'll need it to decrypt files) +4. In **DriveSend AutoUploader**, set `enable_encryption` to **True** + +### Decrypt Files Locally + +Use the scripts in the `/scripts/` folder: + +**macOS:** ```bash -gcloud services enable orgpolicy.googleapis.com --project=YOUR_PROJECT_ID +./scripts/decrypt_folder_mac.sh ``` -Replace `YOUR_PROJECT_ID` with your actual project ID. - -#### 4c. Disable the Key Creation Restrictions - -Run these commands in Cloud Shell: - +**Windows:** ```bash -gcloud org-policies reset iam.disableServiceAccountKeyCreation --project=YOUR_PROJECT_ID - -gcloud org-policies reset iam.managed.disableServiceAccountKeyCreation --project=YOUR_PROJECT_ID +python scripts/decrypt_folder_win.py ``` -### Step 5: Create and Download the Key - -1. Switch back to your **project** (click dropdown top-left, select your project) -2. Go to **IAM & Admin** โ†’ **Service Accounts** -3. Click on your service account -4. Go to **Keys** tab -5. Click **Add Key** โ†’ **Create new key** โ†’ **JSON** โ†’ **Create** -6. **IMPORTANT:** Rename the downloaded file to exactly: `service_account.json` - -### Step 6: Share Your Google Drive Folder - -1. Go to [Google Drive](https://drive.google.com) -2. Create a new folder (e.g., `ComfyUI_Uploads`) -3. Right-click the folder โ†’ **Share** -4. Open your `service_account.json` file and find the `client_email` value (looks like `name@project-id.iam.gserviceaccount.com`) -5. Paste that email in the Share dialog -6. Set permission to **Editor** -7. Click **Share** (uncheck "Notify people" if prompted) - -### Step 7: Get Your Folder ID - -1. Open the folder in Google Drive -2. Look at the URL: `https://drive.google.com/drive/folders/XXXXXXXXXXXXXXXXX` -3. Copy the long string after `/folders/` โ€” that's your **Folder ID** - ---- - -## ๐Ÿš€ Local Setup (Quick Test) - -1. Place `service_account.json` in the `ComfyUI_DriveSendNode` folder -2. Add the **DriveSend Setup** node to ComfyUI -3. Configure: - - `auth_method`: **service_account** - - `folder_id`: your folder ID - - `owner_email`: **your Gmail address** (required!) - - `storage_method`: **env_file** - - `encryption_key_method`: **off** (for testing) -4. Run the node -5. Add **DriveSend AutoUploader** node and run a workflow - ---- - -## โ˜๏ธ RunPod / Cloud Setup (Persistent) - -### Option A: Custom Template (Recommended - One-Time Setup) - -1. Run the **DriveSend Setup** node locally with `storage_method`: **display_only** -2. Copy the output values from the console -3. In RunPod, go to **Secrets** and create: - -| Secret Name | Value | -|-------------|-------| -| `GOOGLE_DRIVE_FOLDER_ID` | Your folder ID | -| `GOOGLE_SERVICE_ACCOUNT_JSON` | The base64 string from setup | -| `GOOGLE_OWNER_EMAIL` | Your Gmail address | -| `comfyui_encryption_key` | (only if using encryption) | - -4. Create or edit a **Pod Template**: - - Click **Edit Template** โ†’ **Environment Variables** - - Add each variable, linking to secrets: - - Key: `GOOGLE_DRIVE_FOLDER_ID` โ†’ Value: `{{ RUNPOD_SECRET_GOOGLE_DRIVE_FOLDER_ID }}` - - Key: `GOOGLE_SERVICE_ACCOUNT_JSON` โ†’ Value: `{{ RUNPOD_SECRET_GOOGLE_SERVICE_ACCOUNT_JSON }}` - - Key: `GOOGLE_OWNER_EMAIL` โ†’ Value: `{{ RUNPOD_SECRET_GOOGLE_OWNER_EMAIL }}` - - Key: `comfyui_encryption_key` โ†’ Value: `{{ RUNPOD_SECRET_comfyui_encryption_key }}` - - Click **Set Overrides** - -5. Deploy pods using this template โ€” credentials persist automatically! - -### Option B: Manual Environment Variables (Each Pod) - -If you don't want to create a template, manually add the environment variables each time you create a pod. This is more tedious but works the same way. - ---- - -## ๐Ÿ“‹ Node Settings Reference - -### DriveSend Setup Node - -| Field | Description | -|-------|-------------| -| `auth_method` | `service_account` (recommended) or `oauth` | -| `folder_id` | Google Drive folder ID from URL | -| `owner_email` | **Your Gmail** - required for service account! | -| `storage_method` | `display_only` (cloud) or `env_file` (local) | -| `encryption_key_method` | `off`, `Display Only`, or `save to .env` | -| `service_account_path` | Default: `service_account.json` | - -### DriveSend AutoUploader Node - -| Field | Description | -|-------|-------------| -| `watch_directory` | Folder to monitor (default: ComfyUI output) | -| `auth_method` | Must match setup node | -| `folder_id` | Override folder ID (or leave blank to use env var) | -| `owner_email` | Override owner email (or leave blank to use env var) | -| `enable_encryption` | Encrypt files before upload | -| `Post_Delete_Enc` | Delete .enc files after upload | -| `Subfolder_Monitor` | Watch subfolders too | -| `run_process` | Start/stop the monitor | - ---- - -## โ“ Why is `owner_email` Required? - -**Service accounts have 0 GB storage quota.** When a service account uploads a file, it owns that file โ€” but it has no storage space! - -The `owner_email` setting transfers ownership to your personal Gmail account after upload, so the file uses YOUR storage quota (15 GB free). - -Without this, uploads will fail with: `403 storageQuotaExceeded` - ---- - -## ๐Ÿ” Encryption (Optional) - -Enable encryption to protect files in cloud storage: - -1. In Setup node: set `encryption_key_method` to **Display Only** or **save to .env** -2. In AutoUploader: set `enable_encryption` to **True** -3. Save your encryption key securely โ€” you need it to decrypt files! - -See the `/scripts/` folder for decryption scripts (run on your local machine after downloading). +**Linux:** +```bash +./scripts/decrypt_folder_linux.sh +``` --- ## ๐Ÿ› ๏ธ Troubleshooting -### "service_account.json NOT FOUND" -- Rename your downloaded key file to exactly `service_account.json` -- Place it in the `ComfyUI_DriveSendNode` folder +### "storageQuotaExceeded" Error +- **Cause:** Using service account with personal Gmail +- **Fix:** Use OAuth instead, or upgrade to Google Workspace -### "403 storageQuotaExceeded" -- Set `owner_email` to your Gmail address -- Make sure the folder is shared with your service account email +### "invalid_grant" Error +- **Cause:** Refresh token expired or revoked +- **Fix:** Run setup node again to get new refresh token -### "Organization Policy blocks key creation" -- Follow Step 4 in the setup guide to disable the policy restrictions -- Run the gcloud commands in Cloud Shell +### "Access blocked: App not verified" +- **Cause:** OAuth consent screen in testing mode +- **Fix:** Add your email as a test user (Step 4 in Google Cloud Setup) -### "Permission denied" on upload -- Make sure you shared the Google Drive folder with the service account email -- The service account needs **Editor** access +### Files not uploading +- **Check:** Console output for error messages +- **Check:** Folder ID is correct +- **Check:** Auth credentials are set in environment variables + +### Browser doesn't open for OAuth +- **Cause:** Running on headless server (RunPod) +- **Fix:** Complete OAuth setup locally first, then copy credentials to RunPod --- -## ๐Ÿ“ Repository Structure +## ๐Ÿ“Š Comparison: DriveSend vs DropSend -``` -ComfyUI_DriveSendNode/ -โ”œโ”€โ”€ __init__.py -โ”œโ”€โ”€ drivesend_uploader_node.py -โ”œโ”€โ”€ drivesend_setup_node.py -โ”œโ”€โ”€ gdrive_upload.py -โ”œโ”€โ”€ gdrive_auth_manager.py -โ”œโ”€โ”€ encrypt_file.py -โ”œโ”€โ”€ monitor_output.py -โ”œโ”€โ”€ requirements.txt -โ”œโ”€โ”€ README.md -โ”œโ”€โ”€ .gitignore -โ””โ”€โ”€ scripts/ - โ”œโ”€โ”€ decrypt_folder.py - โ”œโ”€โ”€ decrypt_folder_mac.sh - โ”œโ”€โ”€ encrypt_folder_mac.sh - โ”œโ”€โ”€ decrypt_folder_win.py - โ”œโ”€โ”€ encrypt_folder_win.py - โ”œโ”€โ”€ decrypt_folder_linux.sh - โ””โ”€โ”€ encrypt_folder_linux.sh -``` +| Feature | DriveSend (Google Drive) | DropSend (Dropbox) | +|---------|-------------------------|-------------------| +| Free with personal account | โœ… OAuth only | โœ… Yes | +| Service account support | โš ๏ธ Workspace only | N/A | +| Storage (free tier) | 15 GB | 2 GB | +| Token management | Auto-refresh | Auto-refresh | +| Setup complexity | Medium | Easy | --- -## ๐Ÿงช Status +## ๐Ÿงช Tested On -**Testing in Progress** โ€” Based on the working [DropSend Node](https://github.com/machinepainting/ComfyUI_DropSendNode). - -Please report issues on GitHub! +- Python 3.10 / 3.11 +- ComfyUI (Feb 2026) +- RunPod GPU instances +- Google Drive API v3 --- diff --git a/drivesend_setup_node.py b/drivesend_setup_node.py index 3f21ecc..82ef854 100644 --- a/drivesend_setup_node.py +++ b/drivesend_setup_node.py @@ -1,235 +1,257 @@ """ DriveSend Setup Node -Handles Google Drive API authentication setup and encryption key generation +Handles Google Drive API authentication setup for both OAuth and Service Account methods """ import os import json import base64 +import webbrowser from pathlib import Path - -from .encrypt_file import generate_key +from cryptography.fernet import Fernet NODE_DIR = Path(__file__).parent class DriveSendSetupNode: - """ - ComfyUI node for setting up Google Drive API access and encryption. - """ + """Setup node for configuring Google Drive authentication.""" + + CATEGORY = "DriveSend" + FUNCTION = "setup" + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("status",) + OUTPUT_NODE = True @classmethod def INPUT_TYPES(cls): return { "required": { - "auth_method": (["service_account", "oauth"],), + "auth_method": (["oauth", "service_account"], {"default": "oauth"}), "folder_id": ("STRING", { - "default": "", + "default": "", "multiline": False, - "placeholder": "Google Drive Folder ID (from URL)" + "tooltip": "Google Drive folder ID from URL (after /folders/)" }), - "owner_email": ("STRING", { - "default": "", - "multiline": False, - "placeholder": "Your Gmail address (for ownership transfer)" - }), - "storage_method": (["display_only", "env_file"],), - "encryption_key_method": (["off", "Display Only", "save to .env"],), + "storage_method": (["display_only", "env_file"], {"default": "display_only"}), + "encryption_key_method": (["off", "display_only", "save_to_env"], {"default": "off"}), }, "optional": { - "service_account_path": ("STRING", { - "default": "service_account.json", + "client_id": ("STRING", { + "default": "", "multiline": False, - "placeholder": "Path to service_account.json" + "tooltip": "OAuth Client ID (from Google Cloud Console)" + }), + "client_secret": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "OAuth Client Secret (from Google Cloud Console)" + }), + "auth_code": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "Authorization code from OAuth flow (leave blank first run)" }), - "reconnect": ("BOOLEAN", {"default": False}), } } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("status",) - FUNCTION = "setup" - CATEGORY = "DriveSend" - OUTPUT_NODE = True - - def setup( - self, - auth_method, - folder_id, - owner_email, - storage_method, - encryption_key_method, - service_account_path="service_account.json", - reconnect=False - ): - output_lines = [] + def setup(self, auth_method, folder_id, storage_method, encryption_key_method, + client_id="", client_secret="", auth_code=""): + """ + Set up Google Drive authentication. + + For OAuth: + - First run: Opens browser for authorization, displays auth URL + - Second run: Exchanges auth code for refresh token + + For Service Account: + - Validates service_account.json exists + - Displays environment variable format for RunPod + """ + + results = [] env_vars = {} - print("\n" + "="*60) - print("[DriveSend Setup] Starting configuration...") - print("="*60) + # Validate folder_id + if not folder_id: + return ("โŒ Error: folder_id is required. Get it from your Google Drive folder URL.",) - # Handle folder ID - if folder_id: - env_vars['GOOGLE_DRIVE_FOLDER_ID'] = folder_id - output_lines.append(f"โœ“ Folder ID configured") - else: - output_lines.append("โš  No Folder ID provided - will upload to Drive root") - - # Handle owner email (required for service account) - if owner_email: - env_vars['GOOGLE_OWNER_EMAIL'] = owner_email - output_lines.append(f"โœ“ Owner Email: {owner_email}") - elif auth_method == "service_account": - output_lines.append("โš  WARNING: No owner_email set!") - output_lines.append(" Service accounts have 0 GB storage quota.") - output_lines.append(" Files will fail to upload without ownership transfer.") - - # Handle authentication based on method - if auth_method == "service_account": - # Try to find service account file - sa_path = None - - # Check provided path first - if service_account_path: - check_path = Path(service_account_path) - if check_path.exists(): - sa_path = check_path - elif (NODE_DIR / service_account_path).exists(): - sa_path = NODE_DIR / service_account_path - - # Fall back to default name in node directory - if not sa_path and (NODE_DIR / 'service_account.json').exists(): - sa_path = NODE_DIR / 'service_account.json' - - if sa_path: - try: - with open(sa_path, 'r') as f: - sa_json = f.read() - - # Base64 encode for environment variable storage - sa_base64 = base64.b64encode(sa_json.encode()).decode() - env_vars['GOOGLE_SERVICE_ACCOUNT_JSON'] = sa_base64 - - # Parse to get client email for display - sa_data = json.loads(sa_json) - client_email = sa_data.get('client_email', 'unknown') - - output_lines.append(f"โœ“ Service Account loaded: {client_email}") - output_lines.append(f" Make sure your Drive folder is shared with this email!") - - print(f"\n[DriveSend Setup] Service Account: {client_email}") - print(f"[DriveSend Setup] Share your Google Drive folder with: {client_email}") - - except Exception as e: - output_lines.append(f"โœ— Error loading service account: {e}") - else: - output_lines.append("โœ— service_account.json NOT FOUND") - output_lines.append("") - output_lines.append("Please rename your downloaded JSON key file to:") - output_lines.append(" service_account.json") - output_lines.append("") - output_lines.append("And place it in:") - output_lines.append(f" {NODE_DIR}/") - return ("\n".join(output_lines),) - - elif auth_method == "oauth": - output_lines.append("") - output_lines.append("โš  OAuth WARNING:") - output_lines.append(" OAuth tokens expire every 7 DAYS in testing mode.") - output_lines.append(" You will need to re-authenticate weekly.") - output_lines.append(" Service Account is recommended for persistent use.") - output_lines.append("") - - creds_path = NODE_DIR / 'credentials.json' - token_path = NODE_DIR / 'token.json' - - if creds_path.exists(): - output_lines.append("โœ“ OAuth credentials.json found") - - if token_path.exists() and not reconnect: - output_lines.append("โœ“ Existing token.json found - already authenticated") - else: - output_lines.append("โ†’ Run the AutoUploader node to complete OAuth authentication") - if reconnect: - try: - token_path.unlink() - output_lines.append(" Deleted existing token for reconnection") - except: - pass - else: - output_lines.append("โœ— credentials.json NOT FOUND") - output_lines.append(" Download OAuth credentials from Google Cloud Console") + env_vars['GOOGLE_DRIVE_FOLDER_ID'] = folder_id # Handle encryption key + encryption_key = None if encryption_key_method != "off": + # Check for existing key existing_key = os.environ.get('comfyui_encryption_key') - - if existing_key and not reconnect: - output_lines.append("โœ“ Encryption key already configured") - env_vars['comfyui_encryption_key'] = existing_key + if existing_key: + encryption_key = existing_key + results.append("Using existing encryption key from environment") else: - new_key = generate_key() - env_vars['comfyui_encryption_key'] = new_key - output_lines.append("โœ“ New encryption key generated") + # Generate new key + encryption_key = Fernet.generate_key().decode('utf-8') + results.append("Generated new encryption key") + + env_vars['comfyui_encryption_key'] = encryption_key - # Output based on storage method - if storage_method == "env_file" and env_vars: - env_path = NODE_DIR / '.env' + # === OAuth Setup === + if auth_method == 'oauth': + if not client_id or not client_secret: + return ( + "โŒ Error: OAuth requires client_id and client_secret.\n\n" + "To get these:\n" + "1. Go to Google Cloud Console โ†’ APIs & Services โ†’ Credentials\n" + "2. Create OAuth 2.0 Client ID (Desktop app)\n" + "3. Copy Client ID and Client Secret here", + ) - # Load existing .env if present - existing_vars = {} - if env_path.exists(): - with open(env_path, 'r') as f: - for line in f: - line = line.strip() - if '=' in line and not line.startswith('#'): - key, value = line.split('=', 1) - existing_vars[key] = value + from .gdrive_auth_manager import get_oauth_credentials_for_setup - # Merge with new vars - existing_vars.update(env_vars) + if not auth_code: + # First run - generate auth URL + result = get_oauth_credentials_for_setup(client_id, client_secret) + + if 'auth_url' in result: + auth_url = result['auth_url'] + + # Try to open browser + try: + webbrowser.open(auth_url) + browser_msg = "Browser opened!" + except: + browser_msg = "Could not open browser automatically." + + print("\n" + "="*60) + print("[DriveSend Setup] OAuth Authorization Required") + print("="*60) + print(f"\n{browser_msg}\n") + print("If browser didn't open, visit this URL:") + print(f"\n{auth_url}\n") + print("After authorizing, copy the code and paste it in the") + print("'auth_code' field, then run this node again.") + print("="*60 + "\n") + + return ( + f"๐Ÿ” OAuth Step 1: Authorization Required\n\n" + f"{browser_msg}\n\n" + f"1. Sign in to Google and authorize the app\n" + f"2. Copy the authorization code\n" + f"3. Paste it in the 'auth_code' field\n" + f"4. Run this node again", + ) + else: + return (f"โŒ Error generating auth URL: {result.get('error', 'Unknown error')}",) - # Write .env file - with open(env_path, 'w') as f: - for key, value in existing_vars.items(): - f.write(f"{key}={value}\n") - - output_lines.append(f"\nโœ“ Credentials saved to: {env_path}") + else: + # Second run - exchange code for tokens + result = get_oauth_credentials_for_setup(client_id, client_secret, auth_code) + + if 'error' in result: + return (f"โŒ Error: {result['error']}\n\nTry generating a new auth code.",) + + if 'credentials' in result: + creds = result['credentials'] + + env_vars['GOOGLE_CLIENT_ID'] = creds['client_id'] + env_vars['GOOGLE_CLIENT_SECRET'] = creds['client_secret'] + env_vars['GOOGLE_REFRESH_TOKEN'] = creds['refresh_token'] + + results.append("โœ“ OAuth authorization successful!") + results.append(f"โœ“ Refresh token obtained") - elif storage_method == "display_only": - output_lines.append("") - output_lines.append("="*50) - output_lines.append("COPY THESE TO RUNPOD SECRETS / ENV VARIABLES:") - output_lines.append("="*50) + # === Service Account Setup === + elif auth_method == 'service_account': + sa_file = NODE_DIR / 'service_account.json' - # Print to console for copying - print("\n" + "="*60) - print("COPY THESE VALUES TO YOUR ENVIRONMENT VARIABLES:") - print("="*60) + if not sa_file.exists(): + return ( + "โŒ Error: service_account.json not found.\n\n" + "To create one:\n" + "1. Go to Google Cloud Console โ†’ IAM & Admin โ†’ Service Accounts\n" + "2. Create a service account\n" + "3. Create a JSON key\n" + "4. Rename to 'service_account.json'\n" + "5. Place in: ComfyUI/custom_nodes/ComfyUI_DriveSendNode/\n\n" + "โš ๏ธ NOTE: Service accounts only work with Google Workspace (paid).\n" + "For personal Gmail accounts, use OAuth instead.", + ) - for key, value in env_vars.items(): - print(f"\n{key}={value}") - - print("\n" + "="*60) - print("See console output above for full values to copy") - print("="*60) + # Read and encode service account JSON + try: + with open(sa_file, 'r') as f: + sa_data = json.load(f) + + sa_email = sa_data.get('client_email', 'unknown') + sa_json_b64 = base64.b64encode(json.dumps(sa_data).encode()).decode() + + env_vars['GOOGLE_SERVICE_ACCOUNT_JSON'] = sa_json_b64 + + results.append("โœ“ Service account loaded") + results.append(f"โœ“ Service account email: {sa_email}") + results.append("") + results.append("โš ๏ธ IMPORTANT: Share your Google Drive folder with:") + results.append(f" {sa_email}") + results.append("") + results.append("โš ๏ธ NOTE: Service accounts require Google Workspace.") + results.append(" They do NOT work with personal Gmail accounts!") + + except Exception as e: + return (f"โŒ Error reading service_account.json: {e}",) - # Print summary + # === Output Results === print("\n" + "="*60) - print("[DriveSend Setup] Configuration complete!") + print("[DriveSend Setup] Configuration Complete!") + print("="*60) + print("\nCopy these values to your Environment Variables:\n") + + for key, value in env_vars.items(): + # Truncate long values for display + display_value = value if len(value) < 50 else value[:47] + "..." + print(f"{key}={display_value}") + + print("\n" + "="*60) + print("See console output above for full values to copy") print("="*60 + "\n") - status = "\n".join(output_lines) - return (status,) + # Save to .env file if requested + if storage_method == 'env_file': + env_file = NODE_DIR / '.env' + try: + with open(env_file, 'w') as f: + for key, value in env_vars.items(): + f.write(f"{key}={value}\n") + results.append(f"โœ“ Saved to {env_file}") + except Exception as e: + results.append(f"โš ๏ธ Could not save .env file: {e}") + + # Build status message + status_lines = [ + "โœ… DriveSend Setup Complete!", + "", + f"Auth Method: {auth_method}", + f"Folder ID: {folder_id}", + f"Encryption: {'Enabled' if encryption_key else 'Disabled'}", + "", + ] + status_lines.extend(results) + status_lines.extend([ + "", + "Check the console/terminal for credentials to copy.", + ]) + + if storage_method == 'display_only': + status_lines.extend([ + "", + "For RunPod: Copy the values from console to RunPod Secrets.", + ]) + + return ("\n".join(status_lines),) # Node registration NODE_CLASS_MAPPINGS = { - "DriveSendSetupNode": DriveSendSetupNode + "DriveSendSetup": DriveSendSetupNode } NODE_DISPLAY_NAME_MAPPINGS = { - "DriveSendSetupNode": "DriveSend Setup" + "DriveSendSetup": "DriveSend Setup" } diff --git a/drivesend_uploader_node.py b/drivesend_uploader_node.py index 84d9464..f1e9f0e 100644 --- a/drivesend_uploader_node.py +++ b/drivesend_uploader_node.py @@ -1,6 +1,6 @@ """ DriveSend AutoUploader Node -Monitors ComfyUI output directory and uploads files to Google Drive +Monitors output folder and uploads new files to Google Drive """ import os @@ -9,147 +9,176 @@ from pathlib import Path from .monitor_output import start_monitor, stop_monitor, get_monitor -# Default output directory -DEFAULT_OUTPUT_DIR = Path(__file__).parent.parent.parent / "output" - - class DriveSendAutoUploaderNode: - """ - ComfyUI node for automatic Google Drive uploads with optional encryption. - """ + """ComfyUI node for automatic Google Drive uploads.""" + + CATEGORY = "DriveSend" + FUNCTION = "run" + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("status",) + OUTPUT_NODE = True @classmethod def INPUT_TYPES(cls): return { "required": { - "watch_directory": ("STRING", { - "default": str(DEFAULT_OUTPUT_DIR), - "multiline": False, - "placeholder": "Directory to monitor for new files" - }), - "auth_method": (["service_account", "oauth"],), + "auth_method": (["oauth", "service_account"], {"default": "oauth"}), "run_process": ("BOOLEAN", {"default": True}), + "enable_encryption": ("BOOLEAN", {"default": False}), + "Subfolder_Monitor": ("BOOLEAN", {"default": True}), + "Post_Delete_Enc": ("BOOLEAN", {"default": False}), }, "optional": { + "watch_folder": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "Folder to monitor (leave blank for ComfyUI output folder)" + }), "folder_id": ("STRING", { "default": "", "multiline": False, - "placeholder": "Google Drive Folder ID (uses env var if blank)" + "tooltip": "Google Drive folder ID (uses env var if blank)" }), - "owner_email": ("STRING", { - "default": "", - "multiline": False, - "placeholder": "Your Gmail address (required for service account)" - }), - "enable_encryption": ("BOOLEAN", {"default": False}), - "Post_Delete_Enc": ("BOOLEAN", {"default": False}), - "Subfolder_Monitor": ("BOOLEAN", {"default": True}), - "any_input": ("*", {}), # Passthrough for workflow chaining } } - RETURN_TYPES = ("STRING", "*") - RETURN_NAMES = ("status", "passthrough") - FUNCTION = "run" - CATEGORY = "DriveSend" - - def run( - self, - watch_directory, - auth_method, - run_process, - folder_id="", - owner_email="", - enable_encryption=False, - Post_Delete_Enc=False, - Subfolder_Monitor=True, - any_input=None - ): - status_lines = [] + def run(self, auth_method, run_process, enable_encryption, Subfolder_Monitor, + Post_Delete_Enc, watch_folder="", folder_id=""): + """ + Start or stop the Google Drive upload monitor. - # Get current monitor state - current_monitor = get_monitor() + Args: + auth_method: 'oauth' or 'service_account' + run_process: True to start, False to stop + enable_encryption: Encrypt files before upload + Subfolder_Monitor: Monitor subfolders recursively + Post_Delete_Enc: Delete .enc files after upload + watch_folder: Folder to monitor (default: ComfyUI output) + folder_id: Google Drive folder ID (default: from env var) + """ - if run_process: - # Start or restart monitor - if current_monitor and current_monitor.is_running(): + # Stop if requested + if not run_process: + monitor = get_monitor() + if monitor and monitor.is_running(): stop_monitor() - status_lines.append("Stopped existing monitor") - - # Validate watch directory - watch_path = Path(watch_directory) - if not watch_path.exists(): - try: - watch_path.mkdir(parents=True, exist_ok=True) - status_lines.append(f"Created watch directory: {watch_path}") - except Exception as e: - status_lines.append(f"Error: Could not create watch directory: {e}") - return ("\n".join(status_lines), any_input) - - # Get owner email from field or environment - effective_owner_email = owner_email if owner_email else os.environ.get('GOOGLE_OWNER_EMAIL', '') - - # Warn if using service account without owner email - if auth_method == 'service_account' and not effective_owner_email: - status_lines.append("โš  WARNING: No owner_email set!") - status_lines.append(" Service accounts have 0 GB storage quota.") - status_lines.append(" Set owner_email to your Gmail or set GOOGLE_OWNER_EMAIL env var.") - status_lines.append("") - - # Set owner email in environment for upload module - if effective_owner_email: - os.environ['GOOGLE_OWNER_EMAIL'] = effective_owner_email - - # Start monitor - try: - monitor = start_monitor( - watch_dir=str(watch_path), - folder_id=folder_id if folder_id else None, - recursive=Subfolder_Monitor, - enable_encryption=enable_encryption, - post_delete_enc=Post_Delete_Enc, - auth_method=auth_method, - owner_email=effective_owner_email - ) - - status_lines.append(f"โœ“ Monitor started") - status_lines.append(f" Directory: {watch_path}") - status_lines.append(f" Auth: {auth_method}") - status_lines.append(f" Recursive: {Subfolder_Monitor}") - status_lines.append(f" Encryption: {enable_encryption}") - - if effective_owner_email: - status_lines.append(f" Owner Email: {effective_owner_email}") - - if folder_id: - status_lines.append(f" Folder ID: {folder_id}") - else: - env_folder = os.environ.get('GOOGLE_DRIVE_FOLDER_ID') - if env_folder: - status_lines.append(f" Folder ID: {env_folder} (from env)") - else: - status_lines.append(" Folder ID: Root (no folder specified)") - - except Exception as e: - status_lines.append(f"Error starting monitor: {e}") - - else: - # Stop monitor - if current_monitor and current_monitor.is_running(): - stop_monitor() - status_lines.append("โœ“ Monitor stopped") + return ("โน๏ธ DriveSend monitor stopped.",) else: - status_lines.append("Monitor not running") + return ("โ„น๏ธ Monitor was not running.",) - status = "\n".join(status_lines) - return (status, any_input) + # Determine watch folder + if watch_folder: + watch_path = Path(watch_folder) + else: + # Default to ComfyUI output folder + # Try common locations + possible_paths = [ + Path("/workspace/ComfyUI/output"), # RunPod + Path("./output"), # Relative + Path(__file__).parent.parent.parent / "output", # Relative to node + ] + + watch_path = None + for p in possible_paths: + if p.exists(): + watch_path = p + break + + if not watch_path: + return ("โŒ Error: Could not find ComfyUI output folder. Please specify watch_folder.",) + + if not watch_path.exists(): + return (f"โŒ Error: Watch folder does not exist: {watch_path}",) + + # Get folder ID from parameter or environment + effective_folder_id = folder_id or os.environ.get('GOOGLE_DRIVE_FOLDER_ID') + if not effective_folder_id: + return ( + "โŒ Error: No folder_id provided.\n\n" + "Either:\n" + "1. Set the folder_id parameter in this node\n" + "2. Set GOOGLE_DRIVE_FOLDER_ID environment variable\n" + "3. Run the DriveSend Setup node first", + ) + + # Check for encryption key if encryption is enabled + if enable_encryption: + enc_key = os.environ.get('comfyui_encryption_key') + if not enc_key: + return ( + "โŒ Error: Encryption enabled but no key found.\n\n" + "Set comfyui_encryption_key environment variable\n" + "or run DriveSend Setup with encryption enabled.", + ) + + # 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') + + if not all([client_id, client_secret, refresh_token]): + # Check for token.json file as fallback + token_file = Path(__file__).parent / 'token.json' + if not token_file.exists(): + return ( + "โŒ Error: OAuth credentials not found.\n\n" + "Set these environment variables:\n" + "- GOOGLE_CLIENT_ID\n" + "- GOOGLE_CLIENT_SECRET\n" + "- GOOGLE_REFRESH_TOKEN\n\n" + "Or run DriveSend Setup to authorize.", + ) + + elif auth_method == 'service_account': + sa_json = os.environ.get('GOOGLE_SERVICE_ACCOUNT_JSON') + sa_file = Path(__file__).parent / 'service_account.json' + + if not sa_json and not sa_file.exists(): + return ( + "โŒ Error: Service account credentials not found.\n\n" + "Either:\n" + "1. Set GOOGLE_SERVICE_ACCOUNT_JSON environment variable (base64)\n" + "2. Place service_account.json in the node folder\n\n" + "โš ๏ธ NOTE: Service accounts only work with Google Workspace (paid).\n" + "For personal Gmail, use OAuth instead.", + ) + + # Start the monitor + try: + monitor = start_monitor( + watch_dir=str(watch_path), + folder_id=effective_folder_id, + recursive=Subfolder_Monitor, + enable_encryption=enable_encryption, + post_delete_enc=Post_Delete_Enc, + auth_method=auth_method + ) + + status_lines = [ + "โœ… DriveSend monitor started!", + "", + f"๐Ÿ“ Watching: {watch_path}", + f"โ˜๏ธ Uploading to: Google Drive folder {effective_folder_id[:20]}...", + f"๐Ÿ” Auth method: {auth_method}", + f"๐Ÿ”’ Encryption: {'Enabled' if enable_encryption else 'Disabled'}", + f"๐Ÿ“‚ Subfolder monitoring: {'Enabled' if Subfolder_Monitor else 'Disabled'}", + "", + "New files will be uploaded automatically.", + "Set run_process to False to stop.", + ] + + return ("\n".join(status_lines),) + + except Exception as e: + return (f"โŒ Error starting monitor: {e}",) # Node registration NODE_CLASS_MAPPINGS = { - "DriveSendAutoUploaderNode": DriveSendAutoUploaderNode + "DriveSendAutoUploader": DriveSendAutoUploaderNode } NODE_DISPLAY_NAME_MAPPINGS = { - "DriveSendAutoUploaderNode": "DriveSend AutoUploader" + "DriveSendAutoUploader": "DriveSend AutoUploader" } diff --git a/gdrive_auth_manager.py b/gdrive_auth_manager.py index 875c5f1..b4a9762 100644 --- a/gdrive_auth_manager.py +++ b/gdrive_auth_manager.py @@ -1,6 +1,6 @@ """ Google Drive Authentication Manager -Handles both Service Account and OAuth 2.0 authentication methods +Handles OAuth 2.0 and Service Account authentication with auto-refresh support """ import os @@ -10,176 +10,277 @@ from pathlib import Path from google.oauth2 import service_account from google.oauth2.credentials import Credentials -from google_auth_oauthlib.flow import InstalledAppFlow from google.auth.transport.requests import Request from googleapiclient.discovery import build -# If modifying scopes, delete token.json +# Scopes required for Google Drive uploads SCOPES = ['https://www.googleapis.com/auth/drive.file'] +# Path to credential files NODE_DIR = Path(__file__).parent - - -def get_service_from_service_account(service_account_path=None, service_account_json=None): - """ - Authenticate using a Service Account. - - Args: - service_account_path: Path to service_account.json file - service_account_json: Base64-encoded service account JSON (for env var storage) - - Returns: - Google Drive API service object - """ - credentials = None - - # Try environment variable first (base64 encoded) - env_json = os.environ.get('GOOGLE_SERVICE_ACCOUNT_JSON') - if env_json: - try: - # Check if it's base64 encoded - try: - decoded = base64.b64decode(env_json).decode('utf-8') - service_account_info = json.loads(decoded) - except: - # Maybe it's just raw JSON - service_account_info = json.loads(env_json) - - credentials = service_account.Credentials.from_service_account_info( - service_account_info, scopes=SCOPES - ) - print("[DriveSend] Authenticated using service account from environment variable") - except Exception as e: - print(f"[DriveSend] Failed to parse service account from env: {e}") - - # Try provided JSON string - if not credentials and service_account_json: - try: - try: - decoded = base64.b64decode(service_account_json).decode('utf-8') - service_account_info = json.loads(decoded) - except: - service_account_info = json.loads(service_account_json) - - credentials = service_account.Credentials.from_service_account_info( - service_account_info, scopes=SCOPES - ) - print("[DriveSend] Authenticated using provided service account JSON") - except Exception as e: - print(f"[DriveSend] Failed to parse provided service account JSON: {e}") - - # Try file path - if not credentials and service_account_path: - path = Path(service_account_path) - if path.exists(): - try: - credentials = service_account.Credentials.from_service_account_file( - str(path), scopes=SCOPES - ) - print(f"[DriveSend] Authenticated using service account file: {path}") - except Exception as e: - print(f"[DriveSend] Failed to load service account file: {e}") - - # Try default location - if not credentials: - default_path = NODE_DIR / 'service_account.json' - if default_path.exists(): - try: - credentials = service_account.Credentials.from_service_account_file( - str(default_path), scopes=SCOPES - ) - print(f"[DriveSend] Authenticated using default service account file") - except Exception as e: - print(f"[DriveSend] Failed to load default service account file: {e}") - - if not credentials: - raise ValueError( - "No valid service account credentials found. Please provide:\n" - " - GOOGLE_SERVICE_ACCOUNT_JSON environment variable (base64 encoded), or\n" - " - service_account.json file in the node directory, or\n" - " - Path to service account JSON file" - ) - - return build('drive', 'v3', credentials=credentials) - - -def get_service_from_oauth(credentials_path=None): - """ - Authenticate using OAuth 2.0 (for local/personal use). - - Args: - credentials_path: Path to credentials.json file (OAuth client ID) - - Returns: - Google Drive API service object - """ - creds = None - token_path = NODE_DIR / 'token.json' - - # Try to load existing token - if token_path.exists(): - try: - creds = Credentials.from_authorized_user_file(str(token_path), SCOPES) - print("[DriveSend] Loaded existing OAuth token") - except Exception as e: - print(f"[DriveSend] Failed to load token: {e}") - - # If no valid credentials, authenticate - if not creds or not creds.valid: - if creds and creds.expired and creds.refresh_token: - try: - creds.refresh(Request()) - print("[DriveSend] Refreshed OAuth token") - except Exception as e: - print(f"[DriveSend] Failed to refresh token: {e}") - creds = None - - if not creds: - # Need to do full OAuth flow - creds_file = credentials_path or (NODE_DIR / 'credentials.json') - - if not Path(creds_file).exists(): - raise ValueError( - f"OAuth credentials file not found: {creds_file}\n" - "Please download credentials.json from Google Cloud Console." - ) - - flow = InstalledAppFlow.from_client_secrets_file(str(creds_file), SCOPES) - creds = flow.run_local_server(port=0) - print("[DriveSend] Completed OAuth authentication") - - # Save the credentials for next time - with open(token_path, 'w') as token: - token.write(creds.to_json()) - print(f"[DriveSend] Saved OAuth token to {token_path}") - - return build('drive', 'v3', credentials=creds) - - -def get_drive_service(auth_method='service_account', **kwargs): - """ - Get an authenticated Google Drive API service. - - Args: - auth_method: 'service_account' or 'oauth' - **kwargs: Additional arguments for the specific auth method - - Returns: - Google Drive API service object - """ - if auth_method == 'service_account': - return get_service_from_service_account( - service_account_path=kwargs.get('service_account_path'), - service_account_json=kwargs.get('service_account_json') - ) - elif auth_method == 'oauth': - return get_service_from_oauth( - credentials_path=kwargs.get('credentials_path') - ) - else: - raise ValueError(f"Unknown auth method: {auth_method}. Use 'service_account' or 'oauth'") +SERVICE_ACCOUNT_FILE = NODE_DIR / 'service_account.json' +CLIENT_SECRET_FILE = NODE_DIR / 'client_secret.json' +TOKEN_FILE = NODE_DIR / 'token.json' def get_folder_id(): - """Get the target folder ID from environment or return None (root).""" + """Get the Google Drive folder ID from environment or return None.""" return os.environ.get('GOOGLE_DRIVE_FOLDER_ID') + + +def get_service_account_credentials(): + """ + Get credentials from service account JSON file or environment variable. + + Returns: + google.oauth2.service_account.Credentials or None + """ + # Try environment variable first (base64 encoded JSON) + sa_json_b64 = os.environ.get('GOOGLE_SERVICE_ACCOUNT_JSON') + if sa_json_b64: + try: + sa_json = base64.b64decode(sa_json_b64).decode('utf-8') + sa_info = json.loads(sa_json) + credentials = service_account.Credentials.from_service_account_info( + sa_info, scopes=SCOPES + ) + print("[DriveSend] Using service account from environment variable") + return credentials + except Exception as e: + print(f"[DriveSend] Failed to load service account from env: {e}") + + # Try file + if SERVICE_ACCOUNT_FILE.exists(): + try: + credentials = service_account.Credentials.from_service_account_file( + str(SERVICE_ACCOUNT_FILE), scopes=SCOPES + ) + print("[DriveSend] Using service account from file") + return credentials + except Exception as e: + print(f"[DriveSend] Failed to load service account file: {e}") + + return None + + +def get_oauth_credentials(): + """ + Get OAuth 2.0 credentials with automatic refresh. + + Checks for credentials in this order: + 1. Environment variables (GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET, GOOGLE_REFRESH_TOKEN) + 2. Saved token.json file + 3. Interactive OAuth flow (requires browser) + + Returns: + google.oauth2.credentials.Credentials or None + """ + credentials = None + + # Option 1: Try environment variables (for RunPod) + 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 client_id and client_secret and refresh_token: + try: + credentials = Credentials( + token=None, + refresh_token=refresh_token, + token_uri='https://oauth2.googleapis.com/token', + client_id=client_id, + client_secret=client_secret, + scopes=SCOPES + ) + # Force refresh to get a valid access token + credentials.refresh(Request()) + print("[DriveSend] OAuth credentials loaded from environment variables") + print("[DriveSend] โœ“ Access token refreshed successfully") + return credentials + except Exception as e: + print(f"[DriveSend] Failed to refresh OAuth token from env: {e}") + + # Option 2: Try saved token file + if TOKEN_FILE.exists(): + try: + 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: + token.write(credentials.to_json()) + print("[DriveSend] OAuth credentials loaded and refreshed from token file") + elif credentials and credentials.valid: + print("[DriveSend] OAuth credentials loaded from token file") + return credentials + except Exception as e: + print(f"[DriveSend] Failed to load token file: {e}") + + # Option 3: Interactive OAuth flow (only works with browser) + if CLIENT_SECRET_FILE.exists(): + try: + from google_auth_oauthlib.flow import InstalledAppFlow + + print("[DriveSend] Starting OAuth authorization flow...") + print("[DriveSend] A browser window will open for authorization.") + + flow = InstalledAppFlow.from_client_secrets_file( + 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: + token.write(credentials.to_json()) + + print("[DriveSend] โœ“ OAuth authorization complete") + print(f"[DriveSend] Token saved to: {TOKEN_FILE}") + return credentials + + except Exception as e: + print(f"[DriveSend] OAuth flow failed: {e}") + print("[DriveSend] Note: OAuth flow requires a browser. For headless servers,") + print("[DriveSend] complete setup locally first and copy credentials to environment variables.") + + return None + + +def get_drive_service(auth_method='oauth', **kwargs): + """ + Get an authenticated Google Drive service. + + Args: + auth_method: 'oauth' or 'service_account' + **kwargs: Additional arguments (unused, for compatibility) + + Returns: + Google Drive API service object + + Raises: + Exception if authentication fails + """ + credentials = None + + if auth_method == 'service_account': + credentials = get_service_account_credentials() + if not credentials: + raise Exception( + "Service account credentials not found. " + "Place service_account.json in the node folder or set GOOGLE_SERVICE_ACCOUNT_JSON environment variable.\n" + "NOTE: Service accounts only work with Google Workspace accounts (paid). " + "For personal Gmail, use OAuth instead." + ) + + elif auth_method == 'oauth': + credentials = get_oauth_credentials() + if not credentials: + raise Exception( + "OAuth credentials not found. Options:\n" + "1. Set GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET, and GOOGLE_REFRESH_TOKEN environment variables\n" + "2. Place client_secret.json in the node folder and run setup to authorize\n" + "3. Run setup locally first if using a headless server" + ) + + else: + raise Exception(f"Unknown auth method: {auth_method}. Use 'oauth' or 'service_account'") + + # Build and return the Drive service + service = build('drive', 'v3', credentials=credentials) + return service + + +def get_oauth_credentials_for_setup(client_id, client_secret, auth_code=None): + """ + Get OAuth credentials during setup process. + + If auth_code is provided, exchanges it for tokens. + Otherwise, returns the authorization URL. + + Args: + client_id: OAuth client ID + client_secret: OAuth client secret + auth_code: Authorization code from OAuth flow (optional) + + Returns: + dict with either 'auth_url' or 'credentials' containing token info + """ + from google_auth_oauthlib.flow import Flow + + # Create flow from client config + client_config = { + "installed": { + "client_id": client_id, + "client_secret": client_secret, + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + "redirect_uris": ["urn:ietf:wg:oauth:2.0:oob", "http://localhost"] + } + } + + flow = Flow.from_client_config(client_config, scopes=SCOPES) + flow.redirect_uri = "urn:ietf:wg:oauth:2.0:oob" + + if not auth_code: + # Generate authorization URL + auth_url, _ = flow.authorization_url( + access_type='offline', + include_granted_scopes='true', + prompt='consent' # Force consent to get refresh token + ) + return {'auth_url': auth_url} + + else: + # Exchange auth code for tokens + try: + flow.fetch_token(code=auth_code) + credentials = flow.credentials + + return { + 'credentials': { + 'client_id': client_id, + 'client_secret': client_secret, + 'refresh_token': credentials.refresh_token, + 'token': credentials.token, + 'token_uri': credentials.token_uri, + 'scopes': list(credentials.scopes) if credentials.scopes else SCOPES + } + } + except Exception as e: + return {'error': str(e)} + + +def refresh_access_token(client_id, client_secret, refresh_token): + """ + Refresh an access token using the refresh token. + + This is called automatically but can also be used to verify credentials. + + Args: + client_id: OAuth client ID + client_secret: OAuth client secret + refresh_token: OAuth refresh token + + Returns: + dict with 'access_token' or 'error' + """ + import requests + + try: + response = requests.post( + 'https://oauth2.googleapis.com/token', + data={ + 'client_id': client_id, + 'client_secret': client_secret, + 'refresh_token': refresh_token, + 'grant_type': 'refresh_token' + } + ) + + if response.status_code == 200: + data = response.json() + return {'access_token': data.get('access_token')} + else: + return {'error': f"Token refresh failed: {response.text}"} + + except Exception as e: + return {'error': str(e)} diff --git a/gdrive_upload.py b/gdrive_upload.py index 9803258..ca11038 100644 --- a/gdrive_upload.py +++ b/gdrive_upload.py @@ -1,6 +1,6 @@ """ Google Drive Upload Module -Handles file uploads to Google Drive with integrity verification and ownership transfer +Handles file uploads to Google Drive with integrity verification """ import os @@ -42,52 +42,7 @@ def calculate_sha256(file_path): return sha256_hash.hexdigest() -def transfer_ownership(service, file_id, owner_email): - """ - Transfer file ownership to a personal Google account. - - This is required because service accounts have 0 GB storage quota. - By transferring ownership, the file counts against the personal account's quota. - - Args: - service: Google Drive API service object - file_id: ID of the uploaded file - owner_email: Email address to transfer ownership to - - Returns: - bool: True if transfer successful, False otherwise - """ - if not owner_email: - return False - - try: - # Create permission with ownership transfer - permission = { - 'type': 'user', - 'role': 'owner', - 'emailAddress': owner_email - } - - service.permissions().create( - fileId=file_id, - body=permission, - transferOwnership=True, - supportsAllDrives=True - ).execute() - - print(f"[DriveSend] โœ“ Ownership transferred to: {owner_email}") - return True - - except HttpError as e: - print(f"[DriveSend] โš  Could not transfer ownership: {e}") - return False - except Exception as e: - print(f"[DriveSend] โš  Ownership transfer error: {e}") - return False - - -def upload_file(file_path, folder_id=None, service=None, auth_method='service_account', - owner_email=None, **auth_kwargs): +def upload_file(file_path, folder_id=None, service=None, auth_method='oauth', **kwargs): """ Upload a file to Google Drive. @@ -95,9 +50,8 @@ def upload_file(file_path, folder_id=None, service=None, auth_method='service_ac file_path: Path to the file to upload folder_id: Google Drive folder ID (optional, uses env var or root if not provided) service: Existing Drive service object (optional) - auth_method: Authentication method ('service_account' or 'oauth') - owner_email: Email to transfer ownership to (required for service account uploads) - **auth_kwargs: Additional auth arguments + auth_method: Authentication method ('oauth' or 'service_account') + **kwargs: Additional arguments (for compatibility) Returns: dict with upload result including file ID and verification status @@ -110,7 +64,7 @@ def upload_file(file_path, folder_id=None, service=None, auth_method='service_ac # Get or create service if service is None: try: - service = get_drive_service(auth_method, **auth_kwargs) + service = get_drive_service(auth_method) except Exception as e: return {'success': False, 'error': f'Authentication failed: {e}'} @@ -118,10 +72,6 @@ def upload_file(file_path, folder_id=None, service=None, auth_method='service_ac if folder_id is None: folder_id = get_folder_id() - # Get owner email from parameter or environment - if owner_email is None: - owner_email = os.environ.get('GOOGLE_OWNER_EMAIL') - # Calculate local hash for verification local_hash = calculate_sha256(file_path) @@ -157,18 +107,6 @@ def upload_file(file_path, folder_id=None, service=None, auth_method='service_ac remote_size = int(file.get('size', 0)) local_size = file_path.stat().st_size - # Transfer ownership if using service account - ownership_transferred = False - if auth_method == 'service_account' and owner_email: - ownership_transferred = transfer_ownership(service, file_id, owner_email) - if not ownership_transferred: - print(f"[DriveSend] โš  Warning: File uploaded but ownership not transferred.") - print(f"[DriveSend] The file may not persist if service account quota is exceeded.") - print(f"[DriveSend] Make sure GOOGLE_OWNER_EMAIL is set to your personal Gmail.") - elif auth_method == 'service_account' and not owner_email: - print(f"[DriveSend] โš  Warning: No owner_email set. Service accounts have 0 GB quota.") - print(f"[DriveSend] Set GOOGLE_OWNER_EMAIL environment variable to your Gmail address.") - # Verify upload size_match = remote_size == local_size @@ -180,8 +118,7 @@ def upload_file(file_path, folder_id=None, service=None, auth_method='service_ac 'file_name': file.get('name'), 'size': remote_size, 'local_hash': local_hash, - 'verified': True, - 'ownership_transferred': ownership_transferred + 'verified': True } else: print(f"[DriveSend] โš  Upload size mismatch: local={local_size}, remote={remote_size}") @@ -192,14 +129,25 @@ def upload_file(file_path, folder_id=None, service=None, auth_method='service_ac 'size': remote_size, 'local_hash': local_hash, 'verified': False, - 'warning': 'Size mismatch - file may be corrupted', - 'ownership_transferred': ownership_transferred + 'warning': 'Size mismatch - file may be corrupted' } except HttpError as e: - error_msg = f'Google Drive API error: {e}' - print(f"[DriveSend] โœ— {error_msg}") - return {'success': False, 'error': error_msg} + error_msg = str(e) + + # Provide helpful error messages + if 'storageQuotaExceeded' in error_msg: + print(f"[DriveSend] โœ— Storage quota exceeded!") + print(f"[DriveSend] If using service account with personal Gmail, this won't work.") + print(f"[DriveSend] Service accounts require Google Workspace (paid).") + print(f"[DriveSend] For personal Gmail, use OAuth authentication instead.") + return { + 'success': False, + 'error': 'Storage quota exceeded. Service accounts do not work with personal Gmail. Use OAuth instead.' + } + + print(f"[DriveSend] โœ— Google Drive API error: {error_msg}") + return {'success': False, 'error': f'Google Drive API error: {error_msg}'} except Exception as e: error_msg = f'Upload failed: {e}' @@ -207,7 +155,7 @@ def upload_file(file_path, folder_id=None, service=None, auth_method='service_ac return {'success': False, 'error': error_msg} -def create_folder(folder_name, parent_folder_id=None, service=None, auth_method='service_account', **auth_kwargs): +def create_folder(folder_name, parent_folder_id=None, service=None, auth_method='oauth'): """ Create a folder in Google Drive. @@ -216,14 +164,13 @@ def create_folder(folder_name, parent_folder_id=None, service=None, auth_method= parent_folder_id: Parent folder ID (optional) service: Existing Drive service object (optional) auth_method: Authentication method - **auth_kwargs: Additional auth arguments Returns: dict with folder creation result including folder ID """ if service is None: try: - service = get_drive_service(auth_method, **auth_kwargs) + service = get_drive_service(auth_method) except Exception as e: return {'success': False, 'error': f'Authentication failed: {e}'} @@ -253,46 +200,45 @@ def create_folder(folder_name, parent_folder_id=None, service=None, auth_method= return {'success': False, 'error': f'Failed to create folder: {e}'} -def find_folder(folder_name, parent_folder_id=None, service=None, auth_method='service_account', **auth_kwargs): +def list_files(folder_id=None, service=None, auth_method='oauth', max_results=100): """ - Find a folder by name in Google Drive. + List files in a Google Drive folder. Args: - folder_name: Name of the folder to find - parent_folder_id: Parent folder ID to search in (optional) + folder_id: Folder ID to list (optional, lists root if not provided) service: Existing Drive service object (optional) auth_method: Authentication method - **auth_kwargs: Additional auth arguments + max_results: Maximum number of files to return Returns: - Folder ID if found, None otherwise + dict with list of files or error """ if service is None: try: - service = get_drive_service(auth_method, **auth_kwargs) + service = get_drive_service(auth_method) except Exception as e: - print(f"[DriveSend] Authentication failed: {e}") - return None + return {'success': False, 'error': f'Authentication failed: {e}'} - query = f"name='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false" + if folder_id is None: + folder_id = get_folder_id() - if parent_folder_id: - query += f" and '{parent_folder_id}' in parents" + query = f"'{folder_id}' in parents and trashed=false" if folder_id else "trashed=false" try: results = service.files().list( q=query, - fields='files(id, name)', - pageSize=1, + fields='files(id, name, size, mimeType, createdTime)', + pageSize=max_results, supportsAllDrives=True, includeItemsFromAllDrives=True ).execute() files = results.get('files', []) - if files: - return files[0].get('id') - return None + return { + 'success': True, + 'files': files, + 'count': len(files) + } except HttpError as e: - print(f"[DriveSend] Error finding folder: {e}") - return None + return {'success': False, 'error': f'Failed to list files: {e}'}