encrypt input images during workflow execution

comfy-nodes/input_image.py

@@ -34,11 +34,12 @@ def xor_decrypt_bytes(data: bytes, key: bytes) -> bytes:
     return decrypted.tobytes()
 
 
-def safe_open_image(image_bytes):
+def safe_open_image(image_bytes, log_error=True):
     try:
         image_pil = Image.open(BytesIO(image_bytes))
     except PIL.UnidentifiedImageError as e:
-        print(e)
+        if log_error:
+            print(e)
         # Convert response content (bytes) to a NumPy array
         image_array = np.frombuffer(image_bytes, np.uint8)
         
@@ -55,6 +56,34 @@ def safe_open_image(image_bytes):
             raise ValueError("The image cannot be identified by neither PIL nor OpenCV")
     return image_pil
 
+
+def open_image_bytes(image_bytes, decrypt=False):
+    if not decrypt:
+        return safe_open_image(image_bytes)
+
+    decrypted_bytes = xor_decrypt_bytes(image_bytes, ENCRYPTION_KEY)
+    return safe_open_image(decrypted_bytes, log_error=False)
+
+
+def open_encrypted_image_file(image_path):
+    with open(image_path, 'rb') as f:
+        image_bytes = f.read()
+
+    try:
+        return open_image_bytes(image_bytes, decrypt=True)
+    except ValueError as decrypt_error:
+        try:
+            image = safe_open_image(image_bytes, log_error=False)
+        except Exception:
+            raise decrypt_error
+
+    encrypted_path = f"{image_path}.encrypted-{uuid.uuid4().hex}"
+    with open(encrypted_path, 'wb') as f:
+        f.write(xor_decrypt_bytes(image_bytes, ENCRYPTION_KEY))
+    os.replace(encrypted_path, image_path)
+    return image
+
+
 class ShellAgentPluginInputImage:
     @classmethod
     def INPUT_TYPES(s):
@@ -180,11 +209,7 @@ class ShellAgentPluginInputImage:
                 response = requests.get(image_path)
                 image_bytes = response.content
 
-                # Decrypt if encryption is enabled
-                if encrypt:
-                    image_bytes = xor_decrypt_bytes(image_bytes, ENCRYPTION_KEY)
-
-                image = safe_open_image(image_bytes)
+                image = open_image_bytes(image_bytes, decrypt=encrypt)
             elif image_path.startswith('data:image/png;base64,') or image_path.startswith('data:image/jpeg;base64,') or image_path.startswith('data:image/jpg;base64,'):
                 import base64
                 from io import BytesIO
@@ -192,22 +217,15 @@ class ShellAgentPluginInputImage:
                 base64_image = image_path[image_path.find(",")+1:]
                 decoded_image = base64.b64decode(base64_image)
 
-                # Decrypt if encryption is enabled
-                if encrypt:
-                    decoded_image = xor_decrypt_bytes(decoded_image, ENCRYPTION_KEY)
-
-                image = safe_open_image(decoded_image)
+                image = open_image_bytes(decoded_image, decrypt=encrypt)
             else:
                 if not os.path.isfile(image_path): # abs path
                     # local path
                     image_path = os.path.join(input_dir, image_path)
 
-                # For local files with encryption enabled, read as bytes and decrypt
+                # Encrypt plain local inputs before downstream nodes can preview them.
                 if encrypt:
-                    with open(image_path, 'rb') as f:
-                        image_bytes = f.read()
-                    image_bytes = xor_decrypt_bytes(image_bytes, ENCRYPTION_KEY)
-                    image = safe_open_image(image_bytes)
+                    image = open_encrypted_image_file(image_path)
                 else:
                     image = node_helpers.pillow(Image.open, image_path)
 
@@ -302,4 +320,4 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     "ShellAgentPluginInputImage": "Input Image (ShellAgent Plugin)",
     # "ShellAgentPluginInputVideo": "Input Video (ShellAgent Plugin)"
-}
+}

tests/test_input_image_encryption.py

@@ -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()