diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..a3ef7e5 --- /dev/null +++ b/__init__.py @@ -0,0 +1,27 @@ +from .comfyui_chaos_encrypt_node import ChaosEncrypt, ChaosDecrypt +from .comfyui_crypto_embed_nodes import ( + ChaCha20EncryptEmbed, + ChaCha20DecryptEmbed, + AESCTREncryptEmbed, + AESCTRDecryptEmbed, +) + +NODE_CLASS_MAPPINGS = { + "ChaosEncrypt": ChaosEncrypt, + "ChaosDecrypt": ChaosDecrypt, + "ChaCha20EncryptEmbed": ChaCha20EncryptEmbed, + "ChaCha20DecryptEmbed": ChaCha20DecryptEmbed, + "AESCTREncryptEmbed": AESCTREncryptEmbed, + "AESCTRDecryptEmbed": AESCTRDecryptEmbed, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ChaosEncrypt": "Chaos Encrypt", + "ChaosDecrypt": "Chaos Decrypt", + "ChaCha20EncryptEmbed": "ChaCha20 Encrypt", + "ChaCha20DecryptEmbed": "ChaCha20 Decrypt", + "AESCTREncryptEmbed": "AES-CTR Encrypt", + "AESCTRDecryptEmbed": "AES-CTR Decrypt", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/comfyui_chaos_encrypt_node.py b/comfyui_chaos_encrypt_node.py new file mode 100644 index 0000000..cb52fe5 --- /dev/null +++ b/comfyui_chaos_encrypt_node.py @@ -0,0 +1,99 @@ +""" +ComfyUI nodes: Chaos-based Image Encrypt / Decrypt (IMAGE in / IMAGE out) + +Usage: +- ChaosEncrypt: input a ComfyUI IMAGE (tensor), set `key` and select `mode`. +- ChaosDecrypt: input encrypted IMAGE and same `key`/`mode`, get decrypted IMAGE. + +WARNING: This is a pedagogical / experimental chaotic encryption demo. +It is NOT cryptographically secure. +""" + +import torch + +# ------------------- Chaotic Keystream Generator ------------------- +def logistic_map_keystream(seed: float, size: int, r: float = 3.99, skip: int = 100) -> torch.ByteTensor: + if isinstance(seed, int): + seed = (seed % 2**31) / float(2**31 - 1) + seed = float(seed) + if seed <= 0 or seed >= 1: + seed = (seed % 1.0) + if seed == 0: + seed = 0.3141592653 + x = seed + out = torch.empty(size, dtype=torch.uint8) + for _ in range(skip): + x = r * x * (1.0 - x) + i = 0 + while i < size: + x = r * x * (1.0 - x) + val = int((x * 2**32) % 256) + out[i] = val + i += 1 + return out + +# ------------------- Encryption / Decryption ------------------- +def chaotic_encrypt_tensor(img: torch.Tensor, seed: float, mode: str = 'xor') -> torch.Tensor: + b, h, w, c = img.shape + flat = (img * 255.0).round().to(torch.uint8).view(-1) + ks = logistic_map_keystream(seed, flat.numel()) + if mode == 'xor': + enc = torch.bitwise_xor(flat, ks) + else: # addmod + enc = ((flat.to(torch.int16) + ks.to(torch.int16)) % 256).to(torch.uint8) + out = enc.view(b, h, w, c).to(torch.float32) / 255.0 + return out + +def chaotic_decrypt_tensor(img: torch.Tensor, seed: float, mode: str = 'xor') -> torch.Tensor: + b, h, w, c = img.shape + flat = (img * 255.0).round().to(torch.uint8).view(-1) + ks = logistic_map_keystream(seed, flat.numel()) + if mode == 'xor': + dec = torch.bitwise_xor(flat, ks) + else: # addmod + dec = ((flat.to(torch.int16) - ks.to(torch.int16)) % 256).to(torch.uint8) + out = dec.view(b, h, w, c).to(torch.float32) / 255.0 + return out + +# ------------------- ComfyUI Node Classes ------------------- +class ChaosEncrypt: + CATEGORY = "Crypto/Chaos" + RETURN_TYPES = ("IMAGE",) + FUNCTION = "encrypt" + DISPLAY_NAME = "Chaos Encrypt" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "key": ("INT", {"default": 123456, "min": 0}), + "mode": (["xor", "addmod"], {"default": "xor"}), # 改为下拉选项 + } + } + + def encrypt(self, image, key=123456, mode="xor"): + seed = ((int(key) & 0x7fffffff) / float(2**31 - 1)) * 0.999999 + out = chaotic_encrypt_tensor(image, seed, mode=mode) + return (out,) + +class ChaosDecrypt: + CATEGORY = "Crypto/Chaos" + RETURN_TYPES = ("IMAGE",) + FUNCTION = "decrypt" + DISPLAY_NAME = "Chaos Decrypt" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "key": ("INT", {"default": 123456, "min": 0}), + "mode": (["xor", "addmod"], {"default": "xor"}), # 改为下拉选项 + } + } + + def decrypt(self, image, key=123456, mode="xor"): + seed = ((int(key) & 0x7fffffff) / float(2**31 - 1)) * 0.999999 + out = chaotic_decrypt_tensor(image, seed, mode=mode) + return (out,) diff --git a/comfyui_crypto_embed_nodes.py b/comfyui_crypto_embed_nodes.py new file mode 100644 index 0000000..7098e75 --- /dev/null +++ b/comfyui_crypto_embed_nodes.py @@ -0,0 +1,175 @@ +""" +ComfyUI nodes: ChaCha20 / AES-CTR Image Encrypt-Decrypt (RGB encrypt + Alpha metadata) + +加密逻辑: +- RGB 通道存放 ciphertext(和原图同样大小)。 +- Alpha 通道嵌入 metadata(nonce、算法、shape)。 +- 解密时先从 Alpha 取 metadata,再解密 RGB。 + +""" + +import torch +import numpy as np +import json, zlib +from Crypto.Cipher import ChaCha20, AES +from Crypto.Protocol.KDF import PBKDF2 +from Crypto.Random import get_random_bytes + +# --------- 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 + + 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 + +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 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) + 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"]) + cipher = ChaCha20.new(key=key, nonce=nonce) + plain = cipher.decrypt(cipher_bytes) + return np.frombuffer(plain, dtype=np.uint8).reshape(shape) + +# --------- 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) + 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"]) + cipher = AES.new(key, AES.MODE_CTR, nonce=nonce) + plain = cipher.decrypt(cipher_bytes) + return np.frombuffer(plain, dtype=np.uint8).reshape(shape) + +# --------- ComfyUI Nodes --------- +class ChaCha20EncryptEmbed: + CATEGORY = "Crypto/Embed" + RETURN_TYPES = ("IMAGE",) + FUNCTION = "encrypt" + DISPLAY_NAME = "ChaCha20 Encrypt (Embed)" + + @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)" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "passphrase": ("STRING", {"default": "changeme"}), + "nonce_bytes": ("INT", {"default": 8, "min": 4, "max": 16}), + } + } + + def encrypt(self, image, passphrase, nonce_bytes=8): + 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) + return (out,) + +class AESCTRDecryptEmbed: + CATEGORY = "Crypto/Embed" + RETURN_TYPES = ("IMAGE",) + FUNCTION = "decrypt" + DISPLAY_NAME = "AES-CTR 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"] == "AES-CTR": + plain = decrypt_aes_ctr(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,) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..4a26500 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +torch +numpy +pycryptodome \ No newline at end of file