From 73e610e31aa2116d7d9b8371e8dbcf074565f84f Mon Sep 17 00:00:00 2001 From: Chengwei Ouyang Date: Mon, 15 Dec 2025 17:27:23 +0800 Subject: [PATCH] feat: add encryption support for saved images and videos - Implemented XOR encryption for image files in output_image.py, allowing users to save encrypted images. - Introduced a new module output_video_encrypt.py for combining images into videos with optional encryption. - Added encryption options in the save_images method and video combine functionality, ensuring output files cannot be viewed directly without decryption. --- comfy-nodes/output_image.py | 45 ++- comfy-nodes/output_video_encrypt.py | 408 ++++++++++++++++++++++++++++ 2 files changed, 451 insertions(+), 2 deletions(-) create mode 100644 comfy-nodes/output_video_encrypt.py diff --git a/comfy-nodes/output_image.py b/comfy-nodes/output_image.py index 7a884b6..68ef388 100644 --- a/comfy-nodes/output_image.py +++ b/comfy-nodes/output_image.py @@ -1,6 +1,36 @@ import folder_paths from nodes import SaveImage import os +import numpy as np + +# Fixed encryption key - use the same key for decryption +ENCRYPTION_KEY = b"ShellAgentSecretKey2024!" + +def xor_encrypt_file(filepath, key): + """ + Fast XOR encryption using NumPy vectorization. + Completely corrupts the file so it cannot be viewed directly. + To decrypt, simply call this function again with the same key. + """ + with open(filepath, 'rb') as f: + data = np.frombuffer(f.read(), dtype=np.uint8) + + # Create repeated key array matching data length + key_array = np.frombuffer(key * ((len(data) // len(key)) + 1), dtype=np.uint8)[:len(data)] + + # Vectorized XOR - 50-100x faster than Python loop + encrypted = np.bitwise_xor(data, key_array) + + with open(filepath, 'wb') as f: + f.write(encrypted.tobytes()) + +def xor_decrypt_file(filepath, key): + """ + Decrypt a file encrypted with xor_encrypt_file. + XOR encryption is symmetric, so decryption is the same as encryption. + """ + xor_encrypt_file(filepath, key) + class ShellAgentSaveImages(SaveImage): @classmethod @@ -9,7 +39,8 @@ class ShellAgentSaveImages(SaveImage): "required": { "images": ("IMAGE", {"tooltip": "The images to save."}), "output_name": ("STRING", {"multiline": False, "default": "output_image"},), - "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."}) + "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."}), + "encrypt": ("BOOLEAN", {"default": False, "tooltip": "If enabled, the saved image will be encrypted and cannot be viewed directly. Use the same key to decrypt elsewhere."}) }, "hidden": { "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" @@ -30,8 +61,18 @@ class ShellAgentSaveImages(SaveImage): } return schema - def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, **extra_kwargs): + def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, encrypt=False, **extra_kwargs): results = super().save_images(images, filename_prefix, prompt, extra_pnginfo) + + # Encrypt saved files if encrypt option is enabled + if encrypt: + for img_info in results.get("ui", {}).get("images", []): + filename = img_info["filename"] + subfolder = img_info.get("subfolder", "") + full_path = os.path.join(self.output_dir, subfolder, filename) + if os.path.exists(full_path): + xor_encrypt_file(full_path, ENCRYPTION_KEY) + results["shellagent_kwargs"] = extra_kwargs return results diff --git a/comfy-nodes/output_video_encrypt.py b/comfy-nodes/output_video_encrypt.py new file mode 100644 index 0000000..913c857 --- /dev/null +++ b/comfy-nodes/output_video_encrypt.py @@ -0,0 +1,408 @@ +import folder_paths +import os +import json +import subprocess +import numpy as np +import re +import shutil +import datetime +from PIL import Image +from PIL.PngImagePlugin import PngInfo +from comfy.utils import ProgressBar + +# Fixed encryption key - use the same key for decryption +ENCRYPTION_KEY = b"ShellAgentSecretKey2024!" + + +def xor_encrypt_file(filepath, key): + """ + Fast XOR encryption using NumPy vectorization. + Completely corrupts the file so it cannot be viewed directly. + To decrypt, simply call this function again with the same key. + """ + with open(filepath, 'rb') as f: + data = np.frombuffer(f.read(), dtype=np.uint8) + + # Create repeated key array matching data length + key_array = np.frombuffer(key * ((len(data) // len(key)) + 1), dtype=np.uint8)[:len(data)] + + # Vectorized XOR - 50-100x faster than Python loop + encrypted = np.bitwise_xor(data, key_array) + + with open(filepath, 'wb') as f: + f.write(encrypted.tobytes()) + + +def xor_decrypt_file(filepath, key): + """ + Decrypt a file encrypted with xor_encrypt_file. + XOR encryption is symmetric, so decryption is the same as encryption. + """ + xor_encrypt_file(filepath, key) + + +def tensor_to_bytes(tensor): + """Convert tensor to uint8 bytes""" + tensor = tensor.cpu().numpy() * 255 + return np.clip(tensor, 0, 255).astype(np.uint8) + + +def find_ffmpeg(): + """Find ffmpeg executable path.""" + # Try to find ffmpeg in PATH + ffmpeg = shutil.which("ffmpeg") + if ffmpeg: + return ffmpeg + + # Try common locations + common_paths = [ + os.path.join(os.path.dirname(__file__), "ffmpeg"), + os.path.join(os.path.dirname(__file__), "ffmpeg.exe"), + os.path.join(folder_paths.base_path, "ffmpeg"), + os.path.join(folder_paths.base_path, "ffmpeg.exe"), + "C:/ffmpeg/bin/ffmpeg.exe", + "/usr/bin/ffmpeg", + "/usr/local/bin/ffmpeg", + ] + + for path in common_paths: + if os.path.isfile(path): + return path + + # Try imageio-ffmpeg as fallback + try: + import imageio_ffmpeg + return imageio_ffmpeg.get_ffmpeg_exe() + except ImportError: + pass + + return None + + +# Find ffmpeg at module load +FFMPEG_PATH = find_ffmpeg() + + +# Video format configurations +VIDEO_FORMATS = { + "h264-mp4": { + "extension": "mp4", + "main_pass": ["-c:v", "libx264", "-preset", "medium", "-crf", "19", "-pix_fmt", "yuv420p"], + "dim_alignment": 2, + "description": "H.264 MP4 - Best compatibility", + }, + "h265-mp4": { + "extension": "mp4", + "main_pass": ["-c:v", "libx265", "-preset", "medium", "-crf", "23", "-pix_fmt", "yuv420p", "-tag:v", "hvc1"], + "dim_alignment": 2, + "description": "H.265/HEVC MP4 - Better compression", + }, + "vp9-webm": { + "extension": "webm", + "main_pass": ["-c:v", "libvpx-vp9", "-crf", "30", "-b:v", "0", "-pix_fmt", "yuv420p"], + "dim_alignment": 2, + "description": "VP9 WebM - Web friendly", + }, + "avi": { + "extension": "avi", + "main_pass": ["-c:v", "mjpeg", "-q:v", "3", "-pix_fmt", "yuvj420p"], + "dim_alignment": 2, + "description": "Motion JPEG AVI", + }, + "mov": { + "extension": "mov", + "main_pass": ["-c:v", "libx264", "-preset", "medium", "-crf", "19", "-pix_fmt", "yuv420p"], + "dim_alignment": 2, + "description": "QuickTime MOV", + }, +} + + +def get_format_list(): + """Get list of available formats for the dropdown.""" + formats = ["image/gif", "image/webp"] + for fmt in VIDEO_FORMATS.keys(): + formats.append(f"video/{fmt}") + return formats + + +class ShellAgentVideoCombineEncrypt: + """ + Video combine node with encryption support. + Combines images into video and encrypts the output files. + The encrypted files cannot be played or viewed directly. + + Supports: GIF, WebP, MP4 (H.264/H.265), WebM (VP9), AVI, MOV + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "frame_rate": ("FLOAT", {"default": 24, "min": 1, "max": 120, "step": 1}), + "loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1, "tooltip": "Number of loops. 0 = infinite for GIF/WebP"}), + "filename_prefix": ("STRING", {"default": "ShellAgent_Encrypted"}), + "format": (get_format_list(),), + "quality": ("INT", {"default": 85, "min": 1, "max": 100, "step": 1, "tooltip": "Quality level (higher = better quality, larger file)"}), + "pingpong": ("BOOLEAN", {"default": False, "tooltip": "Reverse and append frames for seamless loop"}), + "encrypt": ("BOOLEAN", {"default": True, "tooltip": "If enabled, output files will be encrypted and cannot be viewed directly."}), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + RETURN_TYPES = ("VHS_FILENAMES",) + RETURN_NAMES = ("Filenames",) + OUTPUT_NODE = True + CATEGORY = "shellagent" + FUNCTION = "combine_video" + DESCRIPTION = "Combines images into video with optional encryption. Supports GIF, WebP, MP4, WebM, AVI, MOV. Encrypted files cannot be viewed directly." + + @classmethod + def VALIDATE_INPUTS(cls, format, **kwargs): + return True + + def combine_video( + self, + images, + frame_rate: float, + loop_count: int, + filename_prefix="ShellAgent_Encrypted", + format="image/gif", + quality=85, + pingpong=False, + encrypt=True, + prompt=None, + extra_pnginfo=None, + ): + if images is None or len(images) == 0: + return ((True, []),) + + num_frames = len(images) + pbar = ProgressBar(num_frames) + first_image = images[0] + + # Get output directory + output_dir = folder_paths.get_output_directory() + ( + full_output_folder, + filename, + _, + subfolder, + _, + ) = folder_paths.get_save_image_path(filename_prefix, output_dir) + + output_files = [] + + # Setup metadata + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + metadata.add_text("CreationTime", datetime.datetime.now().isoformat(" ")[:19]) + + # Find next counter + max_counter = 0 + matcher = re.compile(f"{re.escape(filename)}_(\\d+)\\D*\\..+", re.IGNORECASE) + for existing_file in os.listdir(full_output_folder): + match = matcher.fullmatch(existing_file) + if match: + file_counter = int(match.group(1)) + if file_counter > max_counter: + max_counter = file_counter + counter = max_counter + 1 + + # Save first frame as png to keep metadata + png_file = f"{filename}_{counter:05}.png" + png_path = os.path.join(full_output_folder, png_file) + Image.fromarray(tensor_to_bytes(first_image)).save( + png_path, + pnginfo=metadata, + compress_level=4, + ) + output_files.append(png_path) + + # Handle pingpong - create reversed frames + if pingpong: + images_list = list(images) + # Add reversed frames (excluding first and last to avoid duplicates) + images = images_list + images_list[-2:0:-1] + num_frames = len(images) + + format_type, format_ext = format.split("/") + + if format_type == "image": + # Use PIL for gif/webp + self._create_animated_image( + images, full_output_folder, filename, counter, + format_ext, frame_rate, loop_count, quality, + output_files, pbar, num_frames + ) + else: + # Use ffmpeg for video formats + self._create_video( + images, full_output_folder, filename, counter, + format_ext, frame_rate, loop_count, quality, + output_files, pbar, first_image + ) + + # Encrypt all output files if encryption is enabled + if encrypt: + for filepath in output_files: + if os.path.exists(filepath): + xor_encrypt_file(filepath, ENCRYPTION_KEY) + + # Return filenames in VHS format + previews = [ + { + "filename": os.path.basename(output_files[-1]), + "subfolder": subfolder, + "type": "output", + "format": format, + "frame_rate": frame_rate, + } + ] + + return {"ui": {"gifs": previews}, "result": ((True, output_files),)} + + def _create_animated_image(self, images, output_folder, filename, counter, + format_ext, frame_rate, loop_count, quality, + output_files, pbar, num_frames): + """Create animated GIF or WebP using PIL.""" + image_kwargs = {} + + if format_ext == "gif": + image_kwargs['disposal'] = 2 + image_kwargs['optimize'] = False + elif format_ext == "webp": + image_kwargs['quality'] = quality + image_kwargs['method'] = 4 + + file = f"{filename}_{counter:05}.{format_ext}" + file_path = os.path.join(output_folder, file) + + frames = [] + for img in images: + frames.append(Image.fromarray(tensor_to_bytes(img))) + pbar.update(1) + + frames[0].save( + file_path, + format=format_ext.upper(), + save_all=True, + append_images=frames[1:], + duration=round(1000 / frame_rate), + loop=loop_count, + **image_kwargs + ) + output_files.append(file_path) + + def _create_video(self, images, output_folder, filename, counter, + format_ext, frame_rate, loop_count, quality, + output_files, pbar, first_image): + """Create video using ffmpeg.""" + if FFMPEG_PATH is None: + raise ProcessLookupError( + "ffmpeg is required for video outputs and could not be found.\n" + "Please install ffmpeg:\n" + " - Windows: Download from https://ffmpeg.org/download.html\n" + " - Linux: sudo apt install ffmpeg\n" + " - Or install imageio-ffmpeg: pip install imageio-ffmpeg" + ) + + # Get format configuration + video_format = VIDEO_FORMATS.get(format_ext, VIDEO_FORMATS["h264-mp4"]) + + # Calculate dimensions with alignment + height, width = first_image.shape[0], first_image.shape[1] + dim_alignment = video_format.get("dim_alignment", 2) + + # Ensure dimensions are divisible by alignment + aligned_width = ((width + dim_alignment - 1) // dim_alignment) * dim_alignment + aligned_height = ((height + dim_alignment - 1) // dim_alignment) * dim_alignment + + dimensions = f"{aligned_width}x{aligned_height}" + + # Output file path + extension = video_format.get("extension", "mp4") + file = f"{filename}_{counter:05}.{extension}" + file_path = os.path.join(output_folder, file) + + # Adjust quality based on format + main_pass = video_format["main_pass"].copy() + + # Map quality (1-100) to CRF (51-0) for x264/x265, or to appropriate scale + if "-crf" in main_pass: + crf_index = main_pass.index("-crf") + 1 + # Quality 100 -> CRF 0, Quality 1 -> CRF 51 + crf_value = int(51 - (quality / 100 * 51)) + main_pass[crf_index] = str(crf_value) + elif "-q:v" in main_pass: + q_index = main_pass.index("-q:v") + 1 + # For MJPEG: quality 100 -> 1, quality 1 -> 31 + q_value = int(1 + ((100 - quality) / 100 * 30)) + main_pass[q_index] = str(q_value) + + # Build ffmpeg command + args = [ + FFMPEG_PATH, + "-v", "error", + "-y", # Overwrite output + "-f", "rawvideo", + "-pix_fmt", "rgb24", + "-s", dimensions, + "-r", str(frame_rate), + "-i", "-", # Read from stdin + ] + main_pass + [file_path] + + # Prepare frame data with padding if needed + frame_data_list = [] + for img in images: + img_bytes = tensor_to_bytes(img) + h, w = img_bytes.shape[:2] + + # Pad to aligned dimensions if necessary + if h != aligned_height or w != aligned_width: + padded = np.zeros((aligned_height, aligned_width, 3), dtype=np.uint8) + padded[:h, :w] = img_bytes + frame_data_list.append(padded.tobytes()) + else: + frame_data_list.append(img_bytes.tobytes()) + pbar.update(1) + + frame_data = b''.join(frame_data_list) + + # Run ffmpeg + try: + result = subprocess.run( + args, + input=frame_data, + capture_output=True, + check=True + ) + except subprocess.CalledProcessError as e: + error_msg = e.stderr.decode('utf-8') if e.stderr else "Unknown error" + raise Exception( + f"FFmpeg error:\n{error_msg}\n" + f"Command: {' '.join(args)}" + ) + except FileNotFoundError: + raise ProcessLookupError( + f"ffmpeg not found at: {FFMPEG_PATH}\n" + "Please ensure ffmpeg is properly installed." + ) + + output_files.append(file_path) + + +NODE_CLASS_MAPPINGS = { + "ShellAgentPluginVideoCombineEncrypt": ShellAgentVideoCombineEncrypt, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "ShellAgentPluginVideoCombineEncrypt": "Video Combine Encrypt (ShellAgent Plugin)", +}