236 lines
8.9 KiB
Python
236 lines
8.9 KiB
Python
"""
|
|
DriveSend Setup Node
|
|
Handles Google Drive API authentication setup and encryption key generation
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import base64
|
|
from pathlib import Path
|
|
|
|
from .encrypt_file import generate_key
|
|
|
|
|
|
NODE_DIR = Path(__file__).parent
|
|
|
|
|
|
class DriveSendSetupNode:
|
|
"""
|
|
ComfyUI node for setting up Google Drive API access and encryption.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"auth_method": (["service_account", "oauth"],),
|
|
"folder_id": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"placeholder": "Google Drive Folder ID (from URL)"
|
|
}),
|
|
"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"],),
|
|
},
|
|
"optional": {
|
|
"service_account_path": ("STRING", {
|
|
"default": "service_account.json",
|
|
"multiline": False,
|
|
"placeholder": "Path to service_account.json"
|
|
}),
|
|
"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 = []
|
|
env_vars = {}
|
|
|
|
print("\n" + "="*60)
|
|
print("[DriveSend Setup] Starting configuration...")
|
|
print("="*60)
|
|
|
|
# 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")
|
|
|
|
# Handle encryption key
|
|
if encryption_key_method != "off":
|
|
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
|
|
else:
|
|
new_key = generate_key()
|
|
env_vars['comfyui_encryption_key'] = new_key
|
|
output_lines.append("✓ New encryption key generated")
|
|
|
|
# Output based on storage method
|
|
if storage_method == "env_file" and env_vars:
|
|
env_path = NODE_DIR / '.env'
|
|
|
|
# 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
|
|
|
|
# Merge with new vars
|
|
existing_vars.update(env_vars)
|
|
|
|
# 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}")
|
|
|
|
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)
|
|
|
|
# Print to console for copying
|
|
print("\n" + "="*60)
|
|
print("COPY THESE VALUES TO YOUR ENVIRONMENT VARIABLES:")
|
|
print("="*60)
|
|
|
|
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)
|
|
|
|
# Print summary
|
|
print("\n" + "="*60)
|
|
print("[DriveSend Setup] Configuration complete!")
|
|
print("="*60 + "\n")
|
|
|
|
status = "\n".join(output_lines)
|
|
return (status,)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DriveSendSetupNode": DriveSendSetupNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DriveSendSetupNode": "DriveSend Setup"
|
|
}
|