diff --git a/__init__.py b/__init__.py index a3ef7e5..9879c42 100644 --- a/__init__.py +++ b/__init__.py @@ -1,27 +1,20 @@ from .comfyui_chaos_encrypt_node import ChaosEncrypt, ChaosDecrypt from .comfyui_crypto_embed_nodes import ( - ChaCha20EncryptEmbed, - ChaCha20DecryptEmbed, - AESCTREncryptEmbed, - AESCTRDecryptEmbed, + PNGEncryptNode,PNGDecryptNode ) NODE_CLASS_MAPPINGS = { "ChaosEncrypt": ChaosEncrypt, "ChaosDecrypt": ChaosDecrypt, - "ChaCha20EncryptEmbed": ChaCha20EncryptEmbed, - "ChaCha20DecryptEmbed": ChaCha20DecryptEmbed, - "AESCTREncryptEmbed": AESCTREncryptEmbed, - "AESCTRDecryptEmbed": AESCTRDecryptEmbed, + "PNGEncryptNode": PNGEncryptNode, + "PNGDecryptNode": PNGDecryptNode, } NODE_DISPLAY_NAME_MAPPINGS = { "ChaosEncrypt": "Chaos Encrypt", "ChaosDecrypt": "Chaos Decrypt", - "ChaCha20EncryptEmbed": "ChaCha20 Encrypt", - "ChaCha20DecryptEmbed": "ChaCha20 Decrypt", - "AESCTREncryptEmbed": "AES-CTR Encrypt", - "AESCTRDecryptEmbed": "AES-CTR Decrypt", + "PNGEncryptNode": "Encrypt Image", + "PNGDecryptNode": "Decrypt PNG", } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/comfyui_crypto_embed_nodes.py b/comfyui_crypto_embed_nodes.py index 7098e75..6023359 100644 --- a/comfyui_crypto_embed_nodes.py +++ b/comfyui_crypto_embed_nodes.py @@ -1,62 +1,49 @@ -""" -ComfyUI nodes: ChaCha20 / AES-CTR Image Encrypt-Decrypt (RGB encrypt + Alpha metadata) - -加密逻辑： -- RGB 通道存放 ciphertext（和原图同样大小）。 -- Alpha 通道嵌入 metadata（nonce、算法、shape）。 -- 解密时先从 Alpha 取 metadata，再解密 RGB。 - -""" - -import torch +import os +import json +import zlib import numpy as np -import json, zlib +import torch +from PIL import Image, PngImagePlugin from Crypto.Cipher import ChaCha20, AES from Crypto.Protocol.KDF import PBKDF2 from Crypto.Random import get_random_bytes +import folder_paths -# --------- Metadata helpers --------- -def embed_metadata_in_alpha(rgb_arr: np.ndarray, meta: dict) -> np.ndarray: - """把 metadata 压缩后写到 Alpha 通道""" - h, w, _ = rgb_arr.shape - rgba = np.zeros((h, w, 4), dtype=np.uint8) - rgba[:, :, :3] = rgb_arr - +# ========= PNG metadata helpers ========= +def save_png_with_metadata(rgb_arr: np.ndarray, meta: dict, path: str): + img = Image.fromarray(rgb_arr, mode="RGB") meta_json = json.dumps(meta).encode("utf-8") meta_compressed = zlib.compress(meta_json) - meta_bytes = np.frombuffer(meta_compressed, dtype=np.uint8) - capacity = h * w - if len(meta_bytes) > capacity: - raise ValueError( - f"Metadata ({len(meta_bytes)} bytes) too large for alpha capacity {capacity}" - ) - alpha = np.zeros(capacity, dtype=np.uint8) - alpha[: len(meta_bytes)] = meta_bytes - rgba[:, :, 3] = alpha.reshape(h, w) - return rgba + pnginfo = PngImagePlugin.PngInfo() + pnginfo.add_text("crypto_meta", meta_compressed.hex()) + img.save(path, "PNG", pnginfo=pnginfo) -def extract_metadata_from_alpha(rgba: np.ndarray) -> dict: - """从 Alpha 通道解出 metadata""" - h, w, _ = rgba.shape - alpha = rgba[:, :, 3].flatten() - # 去掉末尾多余的 0 - raw = bytes(alpha).rstrip(b"\x00") - meta_json = zlib.decompress(raw) - return json.loads(meta_json.decode("utf-8")) -# --------- ChaCha20 --------- +def load_png_with_metadata(path: str): + img = Image.open(path).convert("RGB") + rgb_arr = np.array(img) + if "crypto_meta" not in img.info: + raise ValueError("No crypto_meta in PNG metadata") + meta_hex = img.info["crypto_meta"] + meta_compressed = bytes.fromhex(meta_hex) + meta_json = zlib.decompress(meta_compressed) + meta = json.loads(meta_json.decode("utf-8")) + return rgb_arr, meta + + +# ========= ChaCha20 ========= def encrypt_chacha20(arr: np.ndarray, passphrase: str, nonce_bytes: int = 12): h, w, c = arr.shape key = PBKDF2(passphrase, b"comfyui-salt", dkLen=32) nonce = get_random_bytes(nonce_bytes) cipher = ChaCha20.new(key=key, nonce=nonce) - data = arr.tobytes() - enc = cipher.encrypt(data) + enc = cipher.encrypt(arr.tobytes()) rgb = np.frombuffer(enc, dtype=np.uint8).reshape(h, w, c) meta = {"algo": "ChaCha20", "nonce": nonce.hex(), "shape": [h, w, c]} return rgb, meta + def decrypt_chacha20(cipher_bytes: bytes, passphrase: str, meta: dict, shape): key = PBKDF2(passphrase, b"comfyui-salt", dkLen=32) nonce = bytes.fromhex(meta["nonce"]) @@ -64,18 +51,19 @@ def decrypt_chacha20(cipher_bytes: bytes, passphrase: str, meta: dict, shape): plain = cipher.decrypt(cipher_bytes) return np.frombuffer(plain, dtype=np.uint8).reshape(shape) -# --------- AES-CTR --------- + +# ========= AES-CTR ========= def encrypt_aes_ctr(arr: np.ndarray, passphrase: str, nonce_bytes: int = 8): h, w, c = arr.shape key = PBKDF2(passphrase, b"comfyui-salt", dkLen=32) nonce = get_random_bytes(nonce_bytes) cipher = AES.new(key, AES.MODE_CTR, nonce=nonce) - data = arr.tobytes() - enc = cipher.encrypt(data) + enc = cipher.encrypt(arr.tobytes()) rgb = np.frombuffer(enc, dtype=np.uint8).reshape(h, w, c) meta = {"algo": "AES-CTR", "nonce": nonce.hex(), "shape": [h, w, c]} return rgb, meta + def decrypt_aes_ctr(cipher_bytes: bytes, passphrase: str, meta: dict, shape): key = PBKDF2(passphrase, b"comfyui-salt", dkLen=32) nonce = bytes.fromhex(meta["nonce"]) @@ -83,93 +71,105 @@ def decrypt_aes_ctr(cipher_bytes: bytes, passphrase: str, meta: dict, shape): plain = cipher.decrypt(cipher_bytes) return np.frombuffer(plain, dtype=np.uint8).reshape(shape) -# --------- ComfyUI Nodes --------- -class ChaCha20EncryptEmbed: - CATEGORY = "Crypto/Embed" + +# ========= ComfyUI Nodes ========= +import os +import numpy as np +import torch +from PIL import Image +import folder_paths + +class PNGEncryptNode: + CATEGORY = "Crypto/PNG" RETURN_TYPES = ("IMAGE",) FUNCTION = "encrypt" - DISPLAY_NAME = "ChaCha20 Encrypt (Embed)" + DISPLAY_NAME = "Encrypt Image" - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "passphrase": ("STRING", {"default": "changeme", "multiline": False}), - "nonce_bytes": ("INT", {"default": 12, "min": 8, "max": 16}), - } - } - - def encrypt(self, image, passphrase, nonce_bytes=12): - arr = (image[0].cpu().numpy() * 255).astype(np.uint8) - rgb, meta = encrypt_chacha20(arr, passphrase, nonce_bytes) - rgba = embed_metadata_in_alpha(rgb, meta) - out = torch.from_numpy(rgba.astype(np.float32) / 255.0).unsqueeze(0) - return (out,) - -class ChaCha20DecryptEmbed: - CATEGORY = "Crypto/Embed" - RETURN_TYPES = ("IMAGE",) - FUNCTION = "decrypt" - DISPLAY_NAME = "ChaCha20 Decrypt (Embed)" - - @classmethod - def INPUT_TYPES(cls): - return {"required": {"image": ("IMAGE",), "passphrase": ("STRING", {"default": "changeme"})}} - - def decrypt(self, image, passphrase): - arr = (image[0].cpu().numpy() * 255).astype(np.uint8) - meta = extract_metadata_from_alpha(arr) - h, w, c = meta["shape"] - cipher_bytes = arr[:, :, :3].tobytes() - if meta["algo"] == "ChaCha20": - plain = decrypt_chacha20(cipher_bytes, passphrase, meta, (h, w, c)) - else: - raise ValueError(f"Unsupported algo {meta['algo']}") - out = torch.from_numpy(plain.astype(np.float32) / 255.0).unsqueeze(0) - return (out,) - -class AESCTREncryptEmbed: - CATEGORY = "Crypto/Embed" - RETURN_TYPES = ("IMAGE",) - FUNCTION = "encrypt" - DISPLAY_NAME = "AES-CTR Encrypt (Embed)" + counter = 1 # 静态计数器 @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), + "filename_prefix": ("STRING", {"default": "EncImage_"}), "passphrase": ("STRING", {"default": "changeme"}), - "nonce_bytes": ("INT", {"default": 8, "min": 4, "max": 16}), + "algo": (["ChaCha20", "AES-CTR"],), } } - def encrypt(self, image, passphrase, nonce_bytes=8): + def encrypt(self, image, filename_prefix, passphrase, algo): + workdir = folder_paths.get_output_directory() + #os.makedirs(workdir, exist_ok=True) + + filename = f"{filename_prefix}_{PNGEncryptNode.counter:03}.png" + PNGEncryptNode.counter += 1 + out_path = os.path.join(workdir, filename) + arr = (image[0].cpu().numpy() * 255).astype(np.uint8) - rgb, meta = encrypt_aes_ctr(arr, passphrase, nonce_bytes) - rgba = embed_metadata_in_alpha(rgb, meta) - out = torch.from_numpy(rgba.astype(np.float32) / 255.0).unsqueeze(0) + + if algo == "ChaCha20": + rgb, meta = encrypt_chacha20(arr, passphrase) + else: # AES-CTR + rgb, meta = encrypt_aes_ctr(arr, passphrase) + + save_png_with_metadata(rgb, meta, out_path) + + saved_img = Image.open(out_path).convert("RGB") + saved_np = np.array(saved_img).astype(np.float32) / 255.0 + out = torch.from_numpy(saved_np).unsqueeze(0) + return (out,) -class AESCTRDecryptEmbed: - CATEGORY = "Crypto/Embed" - RETURN_TYPES = ("IMAGE",) +class PNGDecryptNode: + CATEGORY = "Crypto/PNG" FUNCTION = "decrypt" - DISPLAY_NAME = "AES-CTR Decrypt (Embed)" + RETURN_TYPES = ("IMAGE",) + DISPLAY_NAME = "Decrypt PNG" @classmethod def INPUT_TYPES(cls): - return {"required": {"image": ("IMAGE",), "passphrase": ("STRING", {"default": "changeme"})}} + output_dir = folder_paths.get_output_directory() + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(output_dir) if os.path.isfile(os.path.join(output_dir, f))] + files = folder_paths.filter_files_content_types(files, ["image"]) + return { + "required": { + "passphrase": ("STRING", {"default": "changeme"}), + "algo": (["ChaCha20", "AES-CTR"],), + "image": (sorted(files), {"image_upload": True}) + } + } + + def decrypt(self, image, passphrase, algo): + path = folder_paths.get_annotated_filepath(image) + rgb_enc, meta = load_png_with_metadata(path) + cipher_bytes = rgb_enc.tobytes() + shape = tuple(meta["shape"]) - def decrypt(self, image, passphrase): - arr = (image[0].cpu().numpy() * 255).astype(np.uint8) - meta = extract_metadata_from_alpha(arr) - h, w, c = meta["shape"] - cipher_bytes = arr[:, :, :3].tobytes() - if meta["algo"] == "AES-CTR": - plain = decrypt_aes_ctr(cipher_bytes, passphrase, meta, (h, w, c)) + if meta["algo"] == "ChaCha20": + if algo == "ChaCha20": + plain = decrypt_chacha20(cipher_bytes, passphrase, meta, shape) + else: + raise ValueError("Wrong algorithm selection") + elif meta["algo"] == "AES-CTR": + if algo == "AES-CTR": + plain = decrypt_aes_ctr(cipher_bytes, passphrase, meta, shape) + else: + raise ValueError("Wrong algorithm selection") else: - raise ValueError(f"Unsupported algo {meta['algo']}") + raise ValueError(f"Unsupported algorithm: {meta['algo']}") + out = torch.from_numpy(plain.astype(np.float32) / 255.0).unsqueeze(0) return (out,) + +# 注册节点 +NODE_CLASS_MAPPINGS = { + "PNGEncryptNode": PNGEncryptNode, + "PNGDecryptNode": PNGDecryptNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PNGEncryptNode": "Encrypt Image", + "PNGDecryptNode": "Decrypt PNG", +} diff --git a/requirements.txt b/requirements.txt index 4a26500..903c43f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ torch numpy -pycryptodome \ No newline at end of file +pycryptodome +pillow \ No newline at end of file