diff --git a/comfy-nodes/input_image.py b/comfy-nodes/input_image.py index 101512e..f0b3fd6 100755 --- a/comfy-nodes/input_image.py +++ b/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)" -} \ No newline at end of file +} diff --git a/tests/test_input_image_encryption.py b/tests/test_input_image_encryption.py new file mode 100644 index 0000000..37955e6 --- /dev/null +++ b/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()