238 lines
7.6 KiB
Python
238 lines
7.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
ComfyUI Encryption Scripts - Cross-Platform Decryption
|
|
Works with both DropSend and DriveSend nodes
|
|
Compatible with macOS, Windows, and Linux
|
|
|
|
============================================================
|
|
LOCAL USE ONLY - Run this on your local machine to decrypt
|
|
.enc files downloaded/synced from Dropbox or Google Drive.
|
|
Do not run on cloud instances (RunPod, etc.)
|
|
============================================================
|
|
|
|
Key retrieval methods (checked in order):
|
|
- COMFYUI_ENCRYPTION_KEY environment variable
|
|
- comfyui_encryption_key environment variable
|
|
- DROPSEND_ENCRYPTION_KEY environment variable
|
|
- DRIVESEND_ENCRYPTION_KEY environment variable
|
|
- macOS Keychain (ComfyUI_Encryption_Key, DropSend_Encryption_Key, or DriveSend_Encryption_Key)
|
|
- Linux Secret Service (ComfyUI, DropSend, or DriveSend)
|
|
- Manual input
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import platform
|
|
import shutil
|
|
|
|
try:
|
|
from cryptography.fernet import Fernet
|
|
except ImportError:
|
|
print("Error: cryptography library not installed.")
|
|
print("Run: pip install cryptography")
|
|
sys.exit(1)
|
|
|
|
|
|
def get_key_from_macos_keychain():
|
|
"""Try to retrieve the key from macOS Keychain."""
|
|
try:
|
|
import subprocess
|
|
# Try multiple keychain item names for compatibility
|
|
for service, account in [
|
|
('ComfyUI_Encryption_Key', 'ComfyUI'),
|
|
('DropSend_Encryption_Key', 'DropSend'),
|
|
('DriveSend_Encryption_Key', 'DriveSend')
|
|
]:
|
|
result = subprocess.run(
|
|
['security', 'find-generic-password', '-s', service, '-a', account, '-w'],
|
|
capture_output=True, text=True
|
|
)
|
|
if result.returncode == 0 and result.stdout.strip():
|
|
return result.stdout.strip()
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def get_key_from_linux_secret_service():
|
|
"""Try to retrieve the key from Linux Secret Service."""
|
|
try:
|
|
import subprocess
|
|
# Try multiple service names for compatibility
|
|
for service, username in [
|
|
('ComfyUI', 'ComfyUI'),
|
|
('DropSend', 'DropSend'),
|
|
('DriveSend', 'DriveSend')
|
|
]:
|
|
result = subprocess.run(
|
|
['secret-tool', 'lookup', 'service', service, 'username', username],
|
|
capture_output=True, text=True
|
|
)
|
|
if result.returncode == 0 and result.stdout.strip():
|
|
return result.stdout.strip()
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def get_encryption_key():
|
|
"""Get encryption key from environment, platform credential store, or user input."""
|
|
system = platform.system()
|
|
|
|
# Try multiple environment variable names for compatibility
|
|
key = (os.environ.get('COMFYUI_ENCRYPTION_KEY') or
|
|
os.environ.get('comfyui_encryption_key') or
|
|
os.environ.get('DROPSEND_ENCRYPTION_KEY') or
|
|
os.environ.get('DRIVESEND_ENCRYPTION_KEY'))
|
|
|
|
if key:
|
|
print("Using key from environment variable.")
|
|
return key
|
|
|
|
# Try platform-specific credential storage
|
|
if system == 'Darwin': # macOS
|
|
key = get_key_from_macos_keychain()
|
|
if key:
|
|
print("Using key from macOS Keychain.")
|
|
return key
|
|
elif system == 'Linux':
|
|
key = get_key_from_linux_secret_service()
|
|
if key:
|
|
print("Using key from Secret Service.")
|
|
return key
|
|
|
|
# Prompt user
|
|
print("")
|
|
print("Encryption key not found automatically.")
|
|
print("")
|
|
print("To set up automatic key retrieval:")
|
|
if system == 'Darwin':
|
|
print(" - Store in Keychain: Keychain Access > File > New Password Item")
|
|
print(" Item Name: ComfyUI_Encryption_Key, Account: ComfyUI")
|
|
elif system == 'Windows':
|
|
print(" - Set environment variable COMFYUI_ENCRYPTION_KEY")
|
|
print(" (Win+R > sysdm.cpl > Advanced > Environment Variables)")
|
|
elif system == 'Linux':
|
|
print(" - Set environment variable: export COMFYUI_ENCRYPTION_KEY=\"your_key\"")
|
|
print(" - Or use secret-tool: echo -n \"key\" | secret-tool store --label=\"ComfyUI\" service ComfyUI username ComfyUI")
|
|
print("")
|
|
|
|
key = input("Enter your encryption key: ").strip()
|
|
|
|
if not key:
|
|
print("Error: No key provided.")
|
|
sys.exit(1)
|
|
|
|
return key
|
|
|
|
|
|
def decrypt_file(encrypted_path, output_path, fernet):
|
|
"""Decrypt a single file."""
|
|
try:
|
|
with open(encrypted_path, 'rb') as enc_file:
|
|
encrypted_data = enc_file.read()
|
|
decrypted_data = fernet.decrypt(encrypted_data)
|
|
with open(output_path, 'wb') as dec_file:
|
|
dec_file.write(decrypted_data)
|
|
return True
|
|
except Exception as e:
|
|
print(f" X Error decrypting {os.path.basename(encrypted_path)}: {e}")
|
|
return False
|
|
|
|
|
|
def main():
|
|
system = platform.system()
|
|
print(f"=== ComfyUI File Decryption ({system}) ===")
|
|
print("")
|
|
print("Decrypts .enc files to restore original images/videos.")
|
|
print("Works with both DropSend and DriveSend nodes.")
|
|
print("")
|
|
|
|
# Get folder path
|
|
folder = input("Enter the folder path containing .enc files: ").strip()
|
|
folder = folder.strip('"\'')
|
|
folder = folder.replace('\\ ', ' ')
|
|
|
|
if not folder or not os.path.isdir(folder):
|
|
print(f"Error: '{folder}' is not a valid directory.")
|
|
sys.exit(1)
|
|
|
|
print(f"Processing folder: {folder}")
|
|
print("")
|
|
|
|
# Recursive option
|
|
recursive_input = input("Decrypt files recursively (including subfolders)? (Y/N): ").strip().upper()
|
|
recursive = recursive_input == 'Y'
|
|
|
|
# Get encryption key
|
|
print("")
|
|
key = get_encryption_key()
|
|
|
|
try:
|
|
fernet = Fernet(key.encode())
|
|
except Exception as e:
|
|
print(f"Error: Invalid encryption key format: {e}")
|
|
sys.exit(1)
|
|
|
|
print("")
|
|
print("Scanning for .enc files...")
|
|
|
|
# Find all .enc files
|
|
enc_files = []
|
|
if recursive:
|
|
for root, _, files in os.walk(folder):
|
|
for file in files:
|
|
if file.endswith('.enc'):
|
|
enc_files.append(os.path.join(root, file))
|
|
else:
|
|
for file in os.listdir(folder):
|
|
if file.endswith('.enc'):
|
|
enc_files.append(os.path.join(folder, file))
|
|
|
|
if not enc_files:
|
|
print("No .enc files found in the specified location.")
|
|
sys.exit(0)
|
|
|
|
print(f"Found {len(enc_files)} .enc file(s) to decrypt.")
|
|
print("")
|
|
|
|
success_count = 0
|
|
error_count = 0
|
|
|
|
for enc_path in enc_files:
|
|
out_path = enc_path[:-4] # Remove .enc
|
|
|
|
if decrypt_file(enc_path, out_path, fernet):
|
|
print(f" + Decrypted: {os.path.basename(out_path)}")
|
|
success_count += 1
|
|
else:
|
|
error_count += 1
|
|
|
|
print("")
|
|
print(f"Decryption complete: {success_count} successful, {error_count} failed")
|
|
print("")
|
|
|
|
# Offer to move .enc files
|
|
move_input = input("Move all .enc files to a separate folder? (Y/N): ").strip().upper()
|
|
if move_input == 'Y':
|
|
enc_folder = os.path.join(folder, "_encrypted_originals")
|
|
os.makedirs(enc_folder, exist_ok=True)
|
|
moved_count = 0
|
|
|
|
for enc_path in enc_files:
|
|
if os.path.exists(enc_path):
|
|
dest = os.path.join(enc_folder, os.path.basename(enc_path))
|
|
shutil.move(enc_path, dest)
|
|
moved_count += 1
|
|
|
|
print(f"Moved {moved_count} .enc file(s) to: {enc_folder}")
|
|
else:
|
|
print("Leaving .enc files in place.")
|
|
|
|
print("")
|
|
print("Done!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|