From 9296ddb8fe53eadffff920d529ebdfa572d82e0a Mon Sep 17 00:00:00 2001 From: bilt Date: Mon, 7 Jul 2025 13:26:57 +0800 Subject: [PATCH] encrypt video node --- comfy_extras/nodes_video.py | 148 +++++++++++++++++++++++++++++++++++- 1 file changed, 144 insertions(+), 4 deletions(-) diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py index 61f7171b2..c4f091191 100644 --- a/comfy_extras/nodes_video.py +++ b/comfy_extras/nodes_video.py @@ -93,6 +93,8 @@ class SaveVideo(ComfyNodeABC): "filename_prefix": ("STRING", {"default": "video/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."}), "format": (VideoContainer.as_input(), {"default": "auto", "tooltip": "The format to save the video as."}), "codec": (VideoCodec.as_input(), {"default": "auto", "tooltip": "The codec to use for the video."}), + "encrypt": ("BOOLEAN", {"default": False, "tooltip": "Enable simple XOR encryption"}), + "encryption_key": ("STRING", {"default": "mysecretkey", "tooltip": "Encryption key for XOR encryption"}), }, "hidden": { "prompt": "PROMPT", @@ -108,7 +110,18 @@ class SaveVideo(ComfyNodeABC): CATEGORY = "image/video" DESCRIPTION = "Saves the input images to your ComfyUI output directory." - def save_video(self, video: VideoInput, filename_prefix, format, codec, prompt=None, extra_pnginfo=None): + def simple_xor_encrypt(self, data: bytes, key: str) -> bytes: + """简单的XOR加密""" + key_bytes = key.encode('utf-8') + key_length = len(key_bytes) + encrypted = bytearray() + + for i, byte in enumerate(data): + encrypted.append(byte ^ key_bytes[i % key_length]) + + return bytes(encrypted) + + def save_video(self, video: VideoInput, filename_prefix, format, codec, encrypt, encryption_key, prompt=None, extra_pnginfo=None): filename_prefix += self.prefix_append width, height = video.get_dimensions() full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( @@ -127,16 +140,41 @@ class SaveVideo(ComfyNodeABC): metadata["prompt"] = prompt if len(metadata) > 0: saved_metadata = metadata - file = f"{filename}_{counter:05}_.{VideoContainer.get_extension(format)}" + + # 生成文件路径 + file_extension = VideoContainer.get_extension(format) + temp_file_path = os.path.join(full_output_folder, f"temp_{filename}_{counter:05}_.{file_extension}") + final_file_path = os.path.join(full_output_folder, f"{filename}_{counter:05}_.{file_extension}") + + # 先保存到临时文件 video.save_to( - os.path.join(full_output_folder, file), + temp_file_path, format=format, codec=codec, metadata=saved_metadata ) + + # 如果启用加密，对文件进行XOR加密 + if encrypt: + # 读取文件内容 + with open(temp_file_path, 'rb') as f: + video_data = f.read() + + # 删除临时文件 + os.remove(temp_file_path) + + # 加密数据 + encrypted_data = self.simple_xor_encrypt(video_data, encryption_key) + + # 保存加密后的文件 + with open(final_file_path, 'wb') as f: + f.write(encrypted_data) + else: + # 不加密，直接重命名 + os.rename(temp_file_path, final_file_path) results.append({ - "filename": file, + "filename": os.path.basename(final_file_path), "subfolder": subfolder, "type": self.type }) @@ -225,12 +263,113 @@ class LoadVideo(ComfyNodeABC): return True +class LoadVideoEncrypted(ComfyNodeABC): + """支持解密加密视频的加载节点""" + + @classmethod + def INPUT_TYPES(cls): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + + # 使用PIL来过滤视频和动画文件，包括WebP动画 + def is_video_or_animated_file(file_path): + """使用PIL判断文件是否为视频或动画文件""" + try: + from PIL import Image + with Image.open(file_path) as img: + # 检查是否为动画文件（包括WebP动画、GIF等） + if hasattr(img, 'is_animated') and img.is_animated: + return True + + # 检查帧数，如果大于1帧，可能是视频 + if hasattr(img, 'n_frames') and img.n_frames > 1: + return True + + # 检查持续时间，如果有持续时间信息，可能是视频 + if hasattr(img, 'duration') and img.duration > 0: + return True + + return False + except Exception: + # 如果PIL无法打开，使用备用方案 + import mimetypes + mime_type, _ = mimetypes.guess_type(file_path) + if mime_type: + return mime_type.startswith('video/') + + video_extensions = {'.mp4', '.avi', '.mov', '.wmv', '.flv', '.webm', '.mkv', '.m4v'} + file_ext = os.path.splitext(file_path)[1].lower() + return file_ext in video_extensions + + # 过滤出视频和动画文件 + video_files = [] + for file in files: + file_path = os.path.join(input_dir, file) + if is_video_or_animated_file(file_path): + video_files.append(file) + + return { + "required": { + "file": (sorted(video_files), {"video_upload": True}), + "is_encrypted": ("BOOLEAN", {"default": False, "tooltip": "Whether the file is XOR encrypted"}), + "encryption_key": ("STRING", {"default": "mysecretkey", "tooltip": "Encryption key for XOR decryption"}), + } + } + + CATEGORY = "image/video" + RETURN_TYPES = (IO.VIDEO,) + FUNCTION = "load_video_encrypted" + + def simple_xor_decrypt(self, data: bytes, key: str) -> bytes: + """简单的XOR解密（XOR加密是对称的）""" + key_bytes = key.encode('utf-8') + key_length = len(key_bytes) + decrypted = bytearray() + + for i, byte in enumerate(data): + decrypted.append(byte ^ key_bytes[i % key_length]) + + return bytes(decrypted) + + def load_video_encrypted(self, file, is_encrypted, encryption_key): + video_path = folder_paths.get_annotated_filepath(file) + + if not is_encrypted: + # 不加密，直接加载 + return (VideoFromFile(video_path),) + else: + # 读取加密文件 + with open(video_path, 'rb') as f: + encrypted_data = f.read() + + # 解密数据 + decrypted_data = self.simple_xor_decrypt(encrypted_data, encryption_key) + + # 直接使用BytesIO，无需创建临时文件 + import io + video_bytes_io = io.BytesIO(decrypted_data) + + return (VideoFromFile(video_bytes_io),) + + @classmethod + def IS_CHANGED(cls, file, is_encrypted, encryption_key): + video_path = folder_paths.get_annotated_filepath(file) + mod_time = os.path.getmtime(video_path) + return mod_time + + @classmethod + def VALIDATE_INPUTS(cls, file, is_encrypted, encryption_key): + if not folder_paths.exists_annotated_filepath(file): + return "Invalid video file: {}".format(file) + return True + NODE_CLASS_MAPPINGS = { "SaveWEBM": SaveWEBM, "SaveVideo": SaveVideo, "CreateVideo": CreateVideo, "GetVideoComponents": GetVideoComponents, "LoadVideo": LoadVideo, + "LoadVideoEncrypted": LoadVideoEncrypted, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -238,4 +377,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CreateVideo": "Create Video", "GetVideoComponents": "Get Video Components", "LoadVideo": "Load Video", + "LoadVideoEncrypted": "Load Video (Encrypted)", }