encrypt input images during workflow execution

2026-05-12 18:34:08 +08:00
parent 312d88dc21
commit 6889f22fcf
2 changed files with 126 additions and 18 deletions
@@ -0,0 +1,90 @@
+import importlib.util
+import sys
+import tempfile
+import types
+import unittest
+from io import BytesIO
+from pathlib import Path
+
+from PIL import Image, UnidentifiedImageError
+
+
+ROOT = Path(__file__).resolve().parents[1]
+INPUT_IMAGE_PATH = ROOT / "comfy-nodes" / "input_image.py"
+
+
+def load_input_image_module():
+    sys.modules.setdefault("folder_paths", types.SimpleNamespace(get_input_directory=lambda: ""))
+    sys.modules.setdefault(
+        "node_helpers",
+        types.SimpleNamespace(pillow=lambda fn, *args, **kwargs: fn(*args, **kwargs)),
+    )
+    sys.modules.setdefault(
+        "torch",
+        types.SimpleNamespace(
+            from_numpy=lambda value: value,
+            zeros=lambda *args, **kwargs: None,
+            cat=lambda values, dim=0: values,
+            float32="float32",
+        ),
+    )
+    sys.modules.setdefault(
+        "cv2",
+        types.SimpleNamespace(
+            IMREAD_COLOR=1,
+            COLOR_BGR2RGB=1,
+            imdecode=lambda *args, **kwargs: None,
+            cvtColor=lambda image, code: image,
+        ),
+    )
+    sys.modules.setdefault(
+        "pillow_heif",
+        types.SimpleNamespace(register_heif_opener=lambda: None),
+    )
+
+    spec = importlib.util.spec_from_file_location("shellagent_input_image", INPUT_IMAGE_PATH)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def make_png_bytes():
+    buffer = BytesIO()
+    Image.new("RGB", (2, 2), (255, 0, 0)).save(buffer, format="PNG")
+    return buffer.getvalue()
+
+
+class InputImageEncryptionTests(unittest.TestCase):
+    def setUp(self):
+        self.module = load_input_image_module()
+
+    def test_encrypt_flag_encrypts_plain_image_file_for_future_runs(self):
+        plain_bytes = make_png_bytes()
+        with tempfile.TemporaryDirectory() as temp_dir:
+            image_path = Path(temp_dir) / "input.png"
+            image_path.write_bytes(plain_bytes)
+
+            image = self.module.open_encrypted_image_file(str(image_path))
+
+            self.assertEqual(image.size, (2, 2))
+            self.assertNotEqual(image_path.read_bytes(), plain_bytes)
+            with self.assertRaises(UnidentifiedImageError):
+                Image.open(image_path).verify()
+
+            decrypted = self.module.xor_decrypt_bytes(image_path.read_bytes(), self.module.ENCRYPTION_KEY)
+            Image.open(BytesIO(decrypted)).verify()
+
+    def test_encrypt_flag_accepts_previously_encrypted_image_file(self):
+        encrypted = self.module.xor_decrypt_bytes(make_png_bytes(), self.module.ENCRYPTION_KEY)
+        with tempfile.TemporaryDirectory() as temp_dir:
+            image_path = Path(temp_dir) / "input.png"
+            image_path.write_bytes(encrypted)
+
+            image = self.module.open_encrypted_image_file(str(image_path))
+
+            self.assertEqual(image.size, (2, 2))
+            self.assertEqual(image_path.read_bytes(), encrypted)
+
+
+if __name__ == "__main__":
+    unittest.main()