fix some bug

2025-09-14 12:28:45 +08:00
parent fa4e9b88fa
commit b1ee082e86
3 changed files with 116 additions and 122 deletions
@@ -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"]
@@ -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",
+}
@@ -1,3 +1,4 @@
 torch
 numpy
-pycryptodome
+pycryptodome
+pillow