comfyui_node

2025-09-13 16:12:47 +08:00
parent d277bacc4e
commit fa4e9b88fa
4 changed files with 304 additions and 0 deletions
@@ -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"]
@@ -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,)
@@ -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,)
@@ -0,0 +1,3 @@
+torch
+numpy
+pycryptodome