132 lines
3.7 KiB
Python
132 lines
3.7 KiB
Python
"""
|
|
File Encryption Module
|
|
Handles Fernet (AES-128) encryption for files
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from cryptography.fernet import Fernet
|
|
|
|
|
|
def generate_key():
|
|
"""Generate a new Fernet encryption key."""
|
|
return Fernet.generate_key().decode('utf-8')
|
|
|
|
|
|
def get_encryption_key():
|
|
"""
|
|
Get the encryption key from environment variable.
|
|
|
|
Returns:
|
|
str: The encryption key, or None if not set
|
|
"""
|
|
return os.environ.get('comfyui_encryption_key')
|
|
|
|
|
|
def encrypt_file(input_path, output_path=None, key=None):
|
|
"""
|
|
Encrypt a file using Fernet (AES-128).
|
|
|
|
Args:
|
|
input_path: Path to the file to encrypt
|
|
output_path: Path for the encrypted file (default: input_path + '.enc')
|
|
key: Encryption key (default: from environment variable)
|
|
|
|
Returns:
|
|
dict with encryption result
|
|
"""
|
|
input_path = Path(input_path)
|
|
|
|
if not input_path.exists():
|
|
return {'success': False, 'error': f'File not found: {input_path}'}
|
|
|
|
if key is None:
|
|
key = get_encryption_key()
|
|
|
|
if not key:
|
|
return {'success': False, 'error': 'No encryption key provided'}
|
|
|
|
if output_path is None:
|
|
output_path = input_path.with_suffix(input_path.suffix + '.enc')
|
|
else:
|
|
output_path = Path(output_path)
|
|
|
|
try:
|
|
fernet = Fernet(key.encode() if isinstance(key, str) else key)
|
|
|
|
with open(input_path, 'rb') as f:
|
|
data = f.read()
|
|
|
|
encrypted_data = fernet.encrypt(data)
|
|
|
|
with open(output_path, 'wb') as f:
|
|
f.write(encrypted_data)
|
|
|
|
print(f"[DriveSend] Encrypted: {input_path.name} → {output_path.name}")
|
|
return {
|
|
'success': True,
|
|
'input_path': str(input_path),
|
|
'output_path': str(output_path),
|
|
'original_size': len(data),
|
|
'encrypted_size': len(encrypted_data)
|
|
}
|
|
|
|
except Exception as e:
|
|
return {'success': False, 'error': f'Encryption failed: {e}'}
|
|
|
|
|
|
def decrypt_file(input_path, output_path=None, key=None):
|
|
"""
|
|
Decrypt a file using Fernet (AES-128).
|
|
|
|
Args:
|
|
input_path: Path to the encrypted file
|
|
output_path: Path for the decrypted file (default: removes .enc extension)
|
|
key: Encryption key (default: from environment variable)
|
|
|
|
Returns:
|
|
dict with decryption result
|
|
"""
|
|
input_path = Path(input_path)
|
|
|
|
if not input_path.exists():
|
|
return {'success': False, 'error': f'File not found: {input_path}'}
|
|
|
|
if key is None:
|
|
key = get_encryption_key()
|
|
|
|
if not key:
|
|
return {'success': False, 'error': 'No encryption key provided'}
|
|
|
|
if output_path is None:
|
|
# Remove .enc extension
|
|
if input_path.suffix == '.enc':
|
|
output_path = input_path.with_suffix('')
|
|
else:
|
|
output_path = input_path.with_suffix('.decrypted')
|
|
else:
|
|
output_path = Path(output_path)
|
|
|
|
try:
|
|
fernet = Fernet(key.encode() if isinstance(key, str) else key)
|
|
|
|
with open(input_path, 'rb') as f:
|
|
encrypted_data = f.read()
|
|
|
|
decrypted_data = fernet.decrypt(encrypted_data)
|
|
|
|
with open(output_path, 'wb') as f:
|
|
f.write(decrypted_data)
|
|
|
|
print(f"[DriveSend] Decrypted: {input_path.name} → {output_path.name}")
|
|
return {
|
|
'success': True,
|
|
'input_path': str(input_path),
|
|
'output_path': str(output_path),
|
|
'encrypted_size': len(encrypted_data),
|
|
'decrypted_size': len(decrypted_data)
|
|
}
|
|
|
|
except Exception as e:
|
|
return {'success': False, 'error': f'Decryption failed: {e}'}
|