From 268d5559b48c268de345bf8d192aa8befd27cd94 Mon Sep 17 00:00:00 2001 From: Suzu008 Date: Sat, 11 Oct 2025 18:33:16 +0800 Subject: [PATCH] uploadImage cryptIO temp --- src/cryptio/api.py | 47 +++++- src/cryptio/nodes.py | 3 + src/cryptio/upload_image.py | 279 ++++++++++++++++++++++++++++++++++ ts/uploadImage.ts | 295 ++++++++++++++++++++++++++++++++++++ 4 files changed, 623 insertions(+), 1 deletion(-) create mode 100644 src/cryptio/upload_image.py create mode 100644 ts/uploadImage.ts diff --git a/src/cryptio/api.py b/src/cryptio/api.py index fd5984c..22f9102 100644 --- a/src/cryptio/api.py +++ b/src/cryptio/api.py @@ -1,6 +1,9 @@ +import os +import uuid from aiohttp import web from server import PromptServer from .keys import _get_keys +import folder_paths @PromptServer.instance.routes.get("/cryptio/public_key") @@ -10,4 +13,46 @@ async def get_public_key(request): if pub: return web.json_response({"public_key": pub.decode("utf-8")}) else: - return web.json_response({"error": "Public key not found"}, status=404) \ No newline at end of file + return web.json_response({"error": "Public key not found"}, status=404) + + +@PromptServer.instance.routes.post("/cryptio/upload_encrypted") +async def upload_encrypted_image(request): + """接收加密的图片数据并保存到文件""" + try: + # 读取 multipart 表单数据 + reader = await request.multipart() + + encrypted_data = None + original_filename = None + + async for field in reader: + if field.name == 'image': + # 读取加密的文件数据 + encrypted_data = await field.read() + original_filename = field.filename + + if not encrypted_data: + return web.json_response({"error": "No image data received"}, status=400) + + # 生成唯一的加密文件名 + file_ext = os.path.splitext(original_filename)[1] if original_filename else ".enc" + unique_filename = f"cryptio_{uuid.uuid4().hex}{file_ext}.encrypted" + + # 保存到 input 目录 + input_dir = folder_paths.get_input_directory() + file_path = os.path.join(input_dir, unique_filename) + + # 写入加密数据 + with open(file_path, 'wb') as f: + f.write(encrypted_data) + + return web.json_response({ + "name": unique_filename, + "subfolder": "", + "type": "input" + }) + + except Exception as e: + print(f"Error uploading encrypted image: {e}") + return web.json_response({"error": str(e)}, status=500) \ No newline at end of file diff --git a/src/cryptio/nodes.py b/src/cryptio/nodes.py index 791f359..1adf20d 100644 --- a/src/cryptio/nodes.py +++ b/src/cryptio/nodes.py @@ -1,6 +1,7 @@ from .key_generator import KeyGenerator from .client_encrypt import ClientEncrypt from .server_decrypt import ServerDecrypt +from .upload_image import UploadImageCryptIO from . import api as _api # noqa: F401 @@ -18,6 +19,7 @@ NODE_CLASS_MAPPINGS = { "KeyGenerator": KeyGenerator, "ClientEncrypt": ClientEncrypt, "ServerDecrypt": ServerDecrypt, + "UploadImageCryptIO": UploadImageCryptIO, } # 节点显示名称映射 @@ -25,4 +27,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "KeyGenerator": "密钥生成器", "ClientEncrypt": "客户端加密", "ServerDecrypt": "服务端解密", + "UploadImageCryptIO": "加密图片上传", } diff --git a/src/cryptio/upload_image.py b/src/cryptio/upload_image.py new file mode 100644 index 0000000..d2ed0de --- /dev/null +++ b/src/cryptio/upload_image.py @@ -0,0 +1,279 @@ +import os +import base64 +import hashlib +import numpy as np +import torch +from PIL import Image, ImageSequence, ImageOps +from cryptography.hazmat.primitives.asymmetric import padding +from cryptography.hazmat.primitives import serialization, hashes +from cryptography.hazmat.backends import default_backend +from .keys import _get_keys +import folder_paths +import node_helpers + + +class UploadImageCryptIO: + """ + 加密图片上传节点，使用公钥在客户端加密图片数据，服务端使用私钥解密 + """ + + def __init__(self): + self.input_dir = folder_paths.get_input_directory() + self.type = "input" + + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + # 包含图片文件和加密文件 + image_files = folder_paths.filter_files_content_types(files, ["image"]) + encrypted_files = [f for f in files if f.endswith(".encrypted")] + all_files = sorted(set(image_files + encrypted_files)) + + return { + "required": { + "image": (all_files, {"image_upload": True}), + "encrypted": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "mask") + FUNCTION = "load_image" + CATEGORY = "CryptIO" + + def load_image(self, image, encrypted): + """ + 加载并解密图片 + 如果 encrypted=True，从加密文件读取并解密 + """ + try: + # 检查是否是加密文件（文件名以 .encrypted 结尾） + if encrypted and isinstance(image, str) and image.endswith(".encrypted"): + # 从加密文件读取 + image_path = folder_paths.get_annotated_filepath(image) + + with open(image_path, "rb") as f: + encrypted_combined = f.read() + + # 解析组合数据 + offset = 0 + + # 1. 读取加密的AES密钥长度 + encrypted_aes_key_length = int.from_bytes(encrypted_combined[offset : offset + 4], byteorder="big") + offset += 4 + + # 2. 读取加密的AES密钥 + encrypted_aes_key = encrypted_combined[offset : offset + encrypted_aes_key_length] + offset += encrypted_aes_key_length + + # 3. 读取IV长度 + iv_length = int.from_bytes(encrypted_combined[offset : offset + 4], byteorder="big") + offset += 4 + + # 4. 读取IV + iv = encrypted_combined[offset : offset + iv_length] + offset += iv_length + + # 5. 读取加密的数据 + encrypted_data = encrypted_combined[offset:] + + # 6. 使用RSA私钥解密AES密钥 + private_pem = _get_keys().get("server_private_key") + private_key = serialization.load_pem_private_key(private_pem, password=None, backend=default_backend()) + + aes_key_bytes = private_key.decrypt( + encrypted_aes_key, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + + # 7. 使用AES密钥解密图片数据 + from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes + + cipher = Cipher(algorithms.AES(aes_key_bytes), modes.GCM(iv), backend=default_backend()) + decryptor = cipher.decryptor() + + # 注意：GCM模式需要处理认证标签，它附加在加密数据的末尾 + decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize() + + # 从字节数据创建图片 + from io import BytesIO + + img = Image.open(BytesIO(decrypted_data)) + + elif encrypted and isinstance(image, str) and image.startswith("ENCRYPTED_IMAGE:"): + # 兼容：直接传递加密数据的方式（不推荐） + encrypted_combined = base64.b64decode(image[16:]) # 去掉 "ENCRYPTED_IMAGE:" 前缀 + + # 解析组合数据 + offset = 0 + + # 1. 读取加密的AES密钥长度 + encrypted_aes_key_length = int.from_bytes(encrypted_combined[offset : offset + 4], byteorder="big") + offset += 4 + + # 2. 读取加密的AES密钥 + encrypted_aes_key = encrypted_combined[offset : offset + encrypted_aes_key_length] + offset += encrypted_aes_key_length + + # 3. 读取IV长度 + iv_length = int.from_bytes(encrypted_combined[offset : offset + 4], byteorder="big") + offset += 4 + + # 4. 读取IV + iv = encrypted_combined[offset : offset + iv_length] + offset += iv_length + + # 5. 读取加密的数据 + encrypted_data = encrypted_combined[offset:] + + # 6. 使用RSA私钥解密AES密钥 + private_pem = _get_keys().get("server_private_key") + private_key = serialization.load_pem_private_key(private_pem, password=None, backend=default_backend()) + + aes_key_bytes = private_key.decrypt( + encrypted_aes_key, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + + # 7. 使用AES密钥解密图片数据 + from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes + + cipher = Cipher(algorithms.AES(aes_key_bytes), modes.GCM(iv), backend=default_backend()) + decryptor = cipher.decryptor() + + # 注意：GCM模式需要处理认证标签，它附加在加密数据的末尾 + decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize() + + # 从字节数据创建图片 + from io import BytesIO + + img = Image.open(BytesIO(decrypted_data)) + + elif encrypted and isinstance(image, str) and image.startswith("ENCRYPTED:"): + # 兼容旧的纯RSA加密方式（仅用于小数据） + encrypted_data = base64.b64decode(image[10:]) + + # 获取私钥 + private_pem = _get_keys().get("server_private_key") + private_key = serialization.load_pem_private_key(private_pem, password=None, backend=default_backend()) + + # 解密数据 + decrypted_data = private_key.decrypt( + encrypted_data, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + + # 解密后的数据应该是图片的base64编码 + image_data = base64.b64decode(decrypted_data) + + # 从字节数据创建图片 + from io import BytesIO + + img = Image.open(BytesIO(image_data)) + else: + # 正常加载图片（未加密或encrypted=False） + image_path = folder_paths.get_annotated_filepath(image) + img = node_helpers.pillow(Image.open, image_path) + + # 处理图片（与LoadImage节点相同的处理逻辑） + output_images = [] + output_masks = [] + w, h = None, None + + for i in ImageSequence.Iterator(img): + i = node_helpers.pillow(ImageOps.exif_transpose, i) + + if i.mode == "I": + i = i.point(lambda i: i * (1 / 255)) + image_rgb = i.convert("RGB") + + if len(output_images) == 0: + w = image_rgb.size[0] + h = image_rgb.size[1] + + if image_rgb.size[0] != w or image_rgb.size[1] != h: + continue + + image_array = np.array(image_rgb).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image_array)[None,] + + # 处理遮罩 + if "A" in i.getbands(): + mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0 + mask = 1.0 - torch.from_numpy(mask) + elif i.mode == "P" and "transparency" in i.info: + mask = np.array(i.convert("RGBA").getchannel("A")).astype(np.float32) / 255.0 + mask = 1.0 - torch.from_numpy(mask) + else: + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + + output_images.append(image_tensor) + output_masks.append(mask.unsqueeze(0)) + + if len(output_images) > 1 and img.format not in ["MPO"]: + output_image = torch.cat(output_images, dim=0) + output_mask = torch.cat(output_masks, dim=0) + else: + output_image = output_images[0] + output_mask = output_masks[0] + + return (output_image, output_mask) + + except Exception as e: + print(f"加载/解密图片错误: {e}") + # 如果解密失败，尝试正常加载 + image_path = folder_paths.get_annotated_filepath(image) + img = node_helpers.pillow(Image.open, image_path) + + # 简化的错误处理：返回基本的图片和遮罩 + i = node_helpers.pillow(ImageOps.exif_transpose, img) + if i.mode == "I": + i = i.point(lambda i: i * (1 / 255)) + image_rgb = i.convert("RGB") + image_array = np.array(image_rgb).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image_array)[None,] + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu").unsqueeze(0) + + return (image_tensor, mask) + + @classmethod + def IS_CHANGED(s, image, encrypted): + """检查图片是否已更改""" + if isinstance(image, str) and not image.startswith("ENCRYPTED:"): + image_path = folder_paths.get_annotated_filepath(image) + m = hashlib.sha256() + with open(image_path, "rb") as f: + m.update(f.read()) + return m.digest().hex() + return float("NaN") # 对于加密数据，总是重新处理 + + @classmethod + def VALIDATE_INPUTS(s, image, encrypted): + """验证输入""" + if isinstance(image, str) and ( + image.endswith(".encrypted") or image.startswith("ENCRYPTED_IMAGE:") or image.startswith("ENCRYPTED:") + ): + # 加密数据由解密逻辑处理 + # 检查加密文件是否存在 + if image.endswith(".encrypted"): + if not folder_paths.exists_annotated_filepath(image): + return f"Invalid encrypted image file: {image}" + return True + + if not folder_paths.exists_annotated_filepath(image): + return f"Invalid image file: {image}" + + return True diff --git a/ts/uploadImage.ts b/ts/uploadImage.ts new file mode 100644 index 0000000..8932362 --- /dev/null +++ b/ts/uploadImage.ts @@ -0,0 +1,295 @@ +//@ts-ignore +import { app as rawApp } from "../../scripts/app.js"; +//@ts-ignore +import { api as rawApi } from "../../scripts/api.js"; +import type { + ComfyApp, + ComfyApi, + ComfyNodeDef, +} from "@comfyorg/comfyui-frontend-types"; + +const app: ComfyApp = rawApp; +const api: ComfyApi = rawApi; + +// 密钥存储键名 +const PUBLIC_KEY_STORAGE_KEY = "cryptio_public_key"; + +// 获取公钥 +async function getPublicKey(): Promise { + // 先从localStorage获取 + let publicKey = localStorage.getItem(PUBLIC_KEY_STORAGE_KEY); + + // 如果本地没有，则从服务器获取 + if (!publicKey) { + try { + const response = await fetch("/cryptio/public_key"); + const data = await response.json(); + + if (data.public_key) { + publicKey = data.public_key; + if (publicKey == null) { + throw new Error("publicKey is null"); + } + // 保存到localStorage + localStorage.setItem(PUBLIC_KEY_STORAGE_KEY, publicKey); + } else { + console.error("Failed to get public key:", data.error); + return null; + } + } catch (error) { + console.error("Error fetching public key:", error); + return null; + } + } + + return publicKey; +} + +// 导入公钥为 CryptoKey 对象 +async function importPublicKey(publicKey?: string): Promise { + publicKey ??= (await getPublicKey()) ?? undefined; + if (!publicKey) { + throw new Error("Failed to get public key for encryption"); + } + + // 将PEM格式的公钥转换为CryptoKey对象 + const pemHeader = "-----BEGIN PUBLIC KEY-----"; + const pemFooter = "-----END PUBLIC KEY-----"; + const pemContents = publicKey + .substring( + publicKey.indexOf(pemHeader) + pemHeader.length, + publicKey.indexOf(pemFooter) + ) + .replace(/\s/g, ""); + + // 解码Base64 + const binaryDer = window.atob(pemContents); + const binaryDerArray = new Uint8Array(binaryDer.length); + for (let i = 0; i < binaryDer.length; i++) { + binaryDerArray[i] = binaryDer.charCodeAt(i); + } + + // 导入公钥 + const cryptoKey = await window.crypto.subtle.importKey( + "spki", + binaryDerArray, + { + name: "RSA-OAEP", + hash: "SHA-256", + }, + false, + ["encrypt"] + ); + + return cryptoKey; +} + +// 使用混合加密方案加密文件 (AES加密内容 + RSA加密AES密钥) +async function encryptFile( + file: File, + publicKey?: string +): Promise { + try { + // 读取文件内容 + const arrayBuffer = await file.arrayBuffer(); + const data = new Uint8Array(arrayBuffer); + + // 1. 生成随机 AES 密钥 + const aesKey = await window.crypto.subtle.generateKey( + { + name: "AES-GCM", + length: 256, + }, + true, + ["encrypt", "decrypt"] + ); + + // 2. 生成随机 IV (初始化向量) + const iv = window.crypto.getRandomValues(new Uint8Array(12)); + + // 3. 使用 AES 加密文件数据 + const encryptedData = await window.crypto.subtle.encrypt( + { + name: "AES-GCM", + iv: iv, + }, + aesKey, + data + ); + + // 4. 导出 AES 密钥为原始格式 + const rawAesKey = await window.crypto.subtle.exportKey("raw", aesKey); + + // 5. 使用 RSA 公钥加密 AES 密钥 + const rsaCryptoKey = await importPublicKey(publicKey); + const encryptedAesKey = await window.crypto.subtle.encrypt( + { + name: "RSA-OAEP", + }, + rsaCryptoKey, + rawAesKey + ); + + // 6. 组合数据: [加密的AES密钥长度(4字节)] + [加密的AES密钥] + [IV长度(4字节)] + [IV] + [加密的数据] + const encryptedAesKeyArray = new Uint8Array(encryptedAesKey); + const encryptedDataArray = new Uint8Array(encryptedData); + + const combined = new Uint8Array( + 4 + + encryptedAesKeyArray.length + + 4 + + iv.length + + encryptedDataArray.length + ); + + // 写入加密的AES密钥长度 + const view = new DataView(combined.buffer); + view.setUint32(0, encryptedAesKeyArray.length, false); // big-endian + + // 写入加密的AES密钥 + combined.set(encryptedAesKeyArray, 4); + + // 写入IV长度 + view.setUint32(4 + encryptedAesKeyArray.length, iv.length, false); + + // 写入IV + combined.set(iv, 4 + encryptedAesKeyArray.length + 4); + + // 写入加密的数据 + combined.set( + encryptedDataArray, + 4 + encryptedAesKeyArray.length + 4 + iv.length + ); + + return combined; + } catch (error) { + console.error("File encryption error:", error); + throw error; + } +} + +// 上传加密的图片文件 +async function uploadEncryptedImage(file: File): Promise { + try { + console.log("Encrypting and uploading image:", file.name); + + // 加密文件 + const encryptedData = await encryptFile(file); + + // 创建新的 Blob 对象 + const encryptedBlob = new Blob([encryptedData], { + type: "application/octet-stream", + }); + + // 创建 FormData + const formData = new FormData(); + formData.append("image", encryptedBlob, file.name); + + // 上传到加密端点 + const response = await fetch("/cryptio/upload_encrypted", { + method: "POST", + body: formData, + }); + + if (!response.ok) { + throw new Error(`Upload failed: ${response.statusText}`); + } + + return await response.json(); + } catch (error) { + console.error("Encrypted upload failed:", error); + throw error; + } +} + +// 注册加密图片上传节点 +app.registerExtension({ + name: "cryptio.UploadImageCryptIO", + async beforeRegisterNodeDef( + nodeType: any, + nodeData: ComfyNodeDef, + app: ComfyApp + ) { + if (nodeType.comfyClass == "UploadImageCryptIO") { + const originalOnNodeCreated = nodeType.prototype.onNodeCreated; + + nodeType.prototype.onNodeCreated = async function (this: any) { + const me = originalOnNodeCreated?.apply(this); + + if (this?.widgets) { + const widgets = this.widgets; + const imageWidget = widgets[0]; + const encryptedWidget = widgets[1]; + + // 修改 image widget 的上传行为 + if (imageWidget.type === "combo") { + // 保存原始的回调 + const originalCallback = imageWidget.callback; + + // 重写回调以标记为已加密 + imageWidget.callback = function (this: any, ...args: any[]) { + const result = originalCallback?.apply(this, args); + // 标记为已加密 + encryptedWidget.value = true; + return result; + }; + } + } + + return me; + }; + } + }, + async init() { + console.log("CryptIO image upload encryption initialized"); + + // 拦截图片上传 API + const originalUploadImage = (api as any).uploadImage; + + (api as any).uploadImage = async function ( + this: any, + file: File, + updateNode: any, + pasted = false + ) { + // 检查是否是从 UploadImageCryptIO 节点触发的上传 + // 通过检查节点类型来判断 + let shouldEncrypt = false; + + if (updateNode && updateNode.type === "UploadImageCryptIO") { + shouldEncrypt = true; + } + + if (shouldEncrypt) { + // 使用加密上传 + const result = await uploadEncryptedImage(file); + + // 更新节点的 image widget 值 + if (updateNode && result.name) { + const imageWidget = updateNode.widgets?.find( + (w: any) => w.name === "image" + ); + if (imageWidget) { + imageWidget.value = result.name; + } + + // 标记为已加密 + const encryptedWidget = updateNode.widgets?.find( + (w: any) => w.name === "encrypted" + ); + if (encryptedWidget) { + encryptedWidget.value = true; + } + + // 触发节点更新 + app.graph?.setDirtyCanvas(true); + } + + return result; + } else { + // 使用原始上传方法 + return originalUploadImage.call(this, file, updateNode, pasted); + } + }; + }, +});