223 lines
7.8 KiB
Python
223 lines
7.8 KiB
Python
# image_cipher.py
|
|
# EncryptMaster — Image <-> Noise (AES-256-GCM + scrypt)
|
|
# - Image_CipherToNoise : chiffre une image en "bruit" (PNG/TIFF conseillé)
|
|
# - Image_DecipherFromNoise : reconstruit l'image originale depuis le "bruit"
|
|
#
|
|
# Remarques:
|
|
# - La sortie peut être +haute que l'entrée (quelques lignes) pour stocker l'overhead (sel+nonce+header+tag).
|
|
# - Sauvegarder en PNG/TIFF (sans perte). Le JPEG casserait les octets.
|
|
|
|
import os
|
|
import math
|
|
import struct
|
|
from typing import Tuple
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
|
|
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
|
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
|
|
CATEGORY = "EncryptMaster"
|
|
|
|
# ======== Crypto helpers ========
|
|
def _derive_key(passphrase: str, salt: bytes, length: int = 32) -> bytes:
|
|
if not passphrase:
|
|
raise ValueError("passphrase cannot be empty")
|
|
kdf = Scrypt(salt=salt, length=length, n=2**14, r=8, p=1)
|
|
return kdf.derive(passphrase.encode("utf-8"))
|
|
|
|
def _encrypt(plaintext: bytes, passphrase: str, aad: str = "") -> Tuple[bytes, bytes, bytes]:
|
|
salt = os.urandom(16)
|
|
key = _derive_key(passphrase, salt)
|
|
aes = AESGCM(key)
|
|
nonce = os.urandom(12)
|
|
ad = aad.encode("utf-8") if aad else None
|
|
cipher = aes.encrypt(nonce, plaintext, ad) # = plaintext + 16-byte tag
|
|
return salt, nonce, cipher
|
|
|
|
def _decrypt(salt: bytes, nonce: bytes, cipher: bytes, passphrase: str, aad: str = "") -> bytes:
|
|
key = _derive_key(passphrase, salt)
|
|
aes = AESGCM(key)
|
|
ad = aad.encode("utf-8") if aad else None
|
|
return aes.decrypt(nonce, cipher, ad)
|
|
|
|
# ======== Tensor <-> PIL helpers ========
|
|
def tensor2rgb_u8(img_tensor: torch.Tensor) -> np.ndarray:
|
|
"""
|
|
torch.Tensor [B,H,W,C] float32 in [0,1] -> np.uint8 [H,W,3]
|
|
"""
|
|
t = img_tensor[0].detach().cpu().numpy() # [H,W,C]
|
|
t = np.clip(t, 0.0, 1.0)
|
|
arr = (t * 255.0).round().astype(np.uint8)
|
|
if arr.ndim == 2:
|
|
arr = np.stack([arr, arr, arr], axis=-1)
|
|
if arr.shape[2] < 3:
|
|
arr = np.tile(arr[:, :, :1], (1, 1, 3))
|
|
if arr.shape[2] > 3:
|
|
arr = arr[:, :, :3]
|
|
return arr
|
|
|
|
def rgb_u8_to_tensor(arr: np.ndarray):
|
|
"""
|
|
np.uint8 [H,W,3] -> (torch.Tensor [B,H,W,C] float32 in [0,1],)
|
|
"""
|
|
arr = arr.astype(np.float32) / 255.0
|
|
t = torch.from_numpy(arr).unsqueeze(0)
|
|
return (t,)
|
|
|
|
# ======== Payload header ========
|
|
# MAGIC(6) | VER(1) | H(4) | W(4) | C(1) | CIPHER_LEN(8) | RSRV(4)
|
|
# = 6 + 1 + 4 + 4 + 1 + 8 + 4 = 28 bytes
|
|
MAGIC = b"EMIMG1"
|
|
HEADER_FMT = ">6sBIIBQ I" # Espace ignoré par struct; équiv. ">6sBII BQI"
|
|
# Pour éviter toute ambiguïté, on calcule la taille explicitement:
|
|
HEADER_SIZE = 6 + 1 + 4 + 4 + 1 + 8 + 4 # 28
|
|
|
|
def _build_header(h: int, w: int, c: int, cipher_len: int, ver: int = 1) -> bytes:
|
|
# RSRV = 0 (réservé pour futures versions)
|
|
# Comme struct ne gère pas bien les espaces dans le fmt, pack manuellement:
|
|
return (
|
|
MAGIC
|
|
+ bytes([ver & 0xFF])
|
|
+ struct.pack(">I", h)
|
|
+ struct.pack(">I", w)
|
|
+ bytes([c & 0xFF])
|
|
+ struct.pack(">Q", cipher_len)
|
|
+ struct.pack(">I", 0)
|
|
)
|
|
|
|
def _parse_header(buf: bytes):
|
|
if len(buf) < HEADER_SIZE:
|
|
raise ValueError("Header too small")
|
|
if buf[:6] != MAGIC:
|
|
raise ValueError("Invalid magic")
|
|
ver = buf[6]
|
|
h = struct.unpack(">I", buf[7:11])[0]
|
|
w = struct.unpack(">I", buf[11:15])[0]
|
|
c = buf[15]
|
|
cipher_len = struct.unpack(">Q", buf[16:24])[0]
|
|
# rsrv = struct.unpack(">I", buf[24:28])[0]
|
|
return ver, h, w, c, cipher_len
|
|
|
|
# ======== Nodes ========
|
|
class Image_CipherToNoise:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"passphrase": ("STRING", {"multiline": False, "default": "", "password": True}),
|
|
},
|
|
"optional": {
|
|
"associated_data": ("STRING", {"multiline": False, "default": ""}),
|
|
"preserve_width": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT")
|
|
RETURN_NAMES = ("image", "out_width", "out_height")
|
|
FUNCTION = "run"
|
|
CATEGORY = CATEGORY
|
|
|
|
def run(self, image, passphrase, associated_data="", preserve_width=True):
|
|
# 1) Image -> bytes
|
|
rgb = tensor2rgb_u8(image) # [H,W,3]
|
|
h, w, c = rgb.shape
|
|
plain = rgb.tobytes(order="C") # len = h*w*3
|
|
|
|
# 2) Encrypt
|
|
salt, nonce, cipher = _encrypt(plain, passphrase, associated_data)
|
|
|
|
# 3) Build payload: header + salt + nonce + cipher
|
|
header = _build_header(h, w, 3, len(cipher), ver=1)
|
|
payload = header + salt + nonce + cipher
|
|
total = len(payload)
|
|
|
|
# 4) Decide output size (capacity = H*W*3 bytes). We enlarge height if needed.
|
|
if preserve_width and w > 0:
|
|
out_w = w
|
|
pixels = math.ceil(total / 3)
|
|
out_h = math.ceil(pixels / out_w)
|
|
else:
|
|
# make it near-square if free width
|
|
pixels = math.ceil(total / 3)
|
|
side = int(math.sqrt(pixels))
|
|
out_w = max(1, side)
|
|
out_h = math.ceil(pixels / out_w)
|
|
|
|
# 5) Fill noise buffer
|
|
out = np.frombuffer(os.urandom(out_w * out_h * 3), dtype=np.uint8).copy()
|
|
out[:total] = np.frombuffer(payload, dtype=np.uint8)
|
|
out = out.reshape(out_h, out_w, 3)
|
|
|
|
return rgb_u8_to_tensor(out) + (int(out_w), int(out_h))
|
|
|
|
|
|
class Image_DecipherFromNoise:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"passphrase": ("STRING", {"multiline": False, "default": "", "password": True}),
|
|
},
|
|
"optional": {
|
|
"associated_data": ("STRING", {"multiline": False, "default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "report")
|
|
FUNCTION = "run"
|
|
CATEGORY = CATEGORY
|
|
|
|
def run(self, image, passphrase, associated_data=""):
|
|
# 1) Read raw bytes from noise image
|
|
arr = tensor2rgb_u8(image) # [H,W,3]
|
|
raw = arr.reshape(-1).tobytes(order="C")
|
|
|
|
# 2) Parse header
|
|
if len(raw) < HEADER_SIZE + 16 + 12:
|
|
return (rgb_u8_to_tensor(np.zeros((1,1,3), dtype=np.uint8))[0],
|
|
"[ImageCipher ERROR] Not enough data for header")
|
|
ver, orig_h, orig_w, orig_c, cipher_len = _parse_header(raw[:HEADER_SIZE])
|
|
|
|
# 3) Read salt, nonce, cipher
|
|
ofs = HEADER_SIZE
|
|
salt = raw[ofs:ofs+16]; ofs += 16
|
|
nonce = raw[ofs:ofs+12]; ofs += 12
|
|
end = ofs + cipher_len
|
|
if end > len(raw):
|
|
return (rgb_u8_to_tensor(np.zeros((1,1,3), dtype=np.uint8))[0],
|
|
"[ImageCipher ERROR] Truncated payload")
|
|
cipher = raw[ofs:end]
|
|
|
|
# 4) Decrypt
|
|
try:
|
|
plain = _decrypt(salt, nonce, cipher, passphrase, associated_data)
|
|
except Exception as e:
|
|
return (rgb_u8_to_tensor(np.zeros((1,1,3), dtype=np.uint8))[0],
|
|
f"[ImageCipher ERROR] {type(e).__name__}: {e}")
|
|
|
|
# 5) Rebuild original image
|
|
needed = orig_h * orig_w * orig_c
|
|
if len(plain) != needed or orig_c != 3:
|
|
return (rgb_u8_to_tensor(np.zeros((1,1,3), dtype=np.uint8))[0],
|
|
"[ImageCipher ERROR] Invalid plaintext size/channels")
|
|
rgb = np.frombuffer(plain, dtype=np.uint8).reshape(orig_h, orig_w, 3)
|
|
return rgb_u8_to_tensor(rgb)[0], f"Recovered {orig_w}x{orig_h} RGB image (v{ver})."
|
|
|
|
|
|
# ======== ComfyUI mappings ========
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Image Cipher To Noise": Image_CipherToNoise,
|
|
"Image Decipher From Noise": Image_DecipherFromNoise,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Image Cipher To Noise": "EncryptMaster — Image Cipher → Noise (AES-GCM)",
|
|
"Image Decipher From Noise": "EncryptMaster — Image Decipher ← Noise (AES-GCM)",
|
|
}
|