diff --git a/CRYPTO_SYSTEM.md b/CRYPTO_SYSTEM.md new file mode 100644 index 0000000..75df1c0 --- /dev/null +++ b/CRYPTO_SYSTEM.md @@ -0,0 +1,226 @@ +# CryptIO 加密图片上传系统 + +## 系统架构 + +本系统实现了端到端的加密图片上传和预览功能，确保图片在客户端加密，在服务器以加密形式存储，只有在使用时才在服务端解密。 + +## 密钥管理方案 + +### 密钥分布 +- **客户端持有**： + - 客户端公钥/私钥对 + - 服务端公钥/私钥对（通过加密传输获得） + +- **服务端持有**： + - 服务端公钥/私钥对 + - 客户端公钥 + +### 密钥交换流程 + +1. **客户端初始化**： + - 检查 `localStorage`，如果没有客户端密钥对则生成新的 RSA-2048 密钥对 + - 密钥对保存在 `localStorage` 中（键：`cryptio_client_keypair`） + +2. **密钥交换**（`POST /cryptio/exchange_keys`）： + ``` + 客户端 -> 服务器: 发送客户端公钥 + 服务器: 1. 保存客户端公钥到 keys.json + 2. 生成随机 AES-256 密钥 + 3. 用 AES 加密服务端密钥对（JSON格式） + 4. 用客户端公钥加密 AES 密钥 + 5. 返回: {encrypted_aes_key, iv, encrypted_data} + 客户端: 1. 用客户端私钥解密 AES 密钥 + 2. 用 AES 密钥解密服务端密钥对 + 3. 保存到 localStorage（键：`cryptio_server_keys`） + ``` + +3. **密钥缓存**： + - 客户端密钥对和服务端密钥都缓存在 `localStorage` + - 下次访问时直接使用缓存，无需重新交换 + +## 图片上传流程 + +### 1. 客户端加密上传 + +```typescript +文件选择 -> 读取文件 -> 混合加密 -> 上传加密数据 -> 服务器存储 +``` + +**混合加密方案**（AES + RSA）： + +1. 生成随机 AES-256-GCM 密钥 +2. 生成随机 IV（12字节） +3. 用 AES-GCM 加密图片数据 +4. 用**服务端公钥**加密 AES 密钥 +5. 组合数据格式： + ``` + [4字节: AES密钥长度] + + [加密的AES密钥] + + [4字节: IV长度] + + [IV] + + [加密的图片数据] + ``` + +### 2. 服务器存储 + +- 端点：`POST /cryptio/upload_encrypted` +- 接收加密数据，直接保存到 `input/` 目录 +- 文件名格式：`cryptio_{uuid}.{ext}.encrypted` +- **文件以完全加密的形式存储在磁盘上** + +### 3. 服务端解密（节点执行时） + +当 workflow 执行到 `UploadImageCryptIO` 节点时： + +1. 从 `input/` 目录读取 `.encrypted` 文件 +2. 解析组合数据（AES密钥长度、AES密钥、IV长度、IV、加密数据） +3. 用**服务端私钥**解密 AES 密钥 +4. 用 AES-GCM 解密图片数据 +5. 转换为 PIL Image 对象 +6. 处理并返回 IMAGE tensor 和 MASK tensor + +代码位置：`src/cryptio/upload_image.py:48` + +## 图片预览流程 + +### 客户端预览加密图片 + +``` +选择加密图片 -> 请求加密数据 -> 客户端解密 -> 显示预览 +``` + +1. **获取加密数据**： + - 端点：`GET /cryptio/view_encrypted?filename=xxx.encrypted` + - 服务器返回 base64 编码的加密文件内容 + +2. **客户端解密**： + - 用**服务端私钥**解密 AES 密钥 + - 用 AES 密钥解密图片数据 + - 创建 Blob URL 显示图片 + +代码位置：`ts/cryptioUploadImage.ts:350` + +## 文件结构 + +### 服务端 + +``` +src/cryptio/ +├── keys.py # 密钥管理 +├── api.py # API端点 +│ ├── POST /cryptio/exchange_keys # 密钥交换 +│ ├── POST /cryptio/upload_encrypted # 上传加密图片 +│ ├── GET /cryptio/view_encrypted # 获取加密图片 +│ └── GET /cryptio/public_key # 兼容旧接口 +├── upload_image.py # UploadImageCryptIO节点 +└── nodes.py # 节点注册 +``` + +### 客户端 + +``` +ts/ +├── cryptioUploadImage.ts # 完整的加密上传实现 +│ ├── 密钥生成和管理 +│ ├── 密钥交换 +│ ├── 文件加密/解密 +│ ├── 自定义上传widget +│ └── 预览功能 +└── clientEncrypt.ts # 文本加密节点（旧） +``` + +## 安全特性 + +1. **端到端加密**： + - 图片在客户端加密后才上传 + - 服务器存储加密文件 + - 只在节点执行时临时解密 + +2. **混合加密**： + - 结合 RSA（非对称）和 AES（对称）的优点 + - RSA 用于密钥交换，AES 用于数据加密 + - 每次上传使用不同的随机 AES 密钥和 IV + +3. **密钥安全传输**： + - 服务端密钥通过客户端公钥加密传输 + - 客户端密钥永不离开客户端 + +4. **文件级加密**： + - 加密文件直接存储在磁盘 + - 服务器重启不影响加密数据 + - 无明文缓存 + +## 使用方法 + +### 1. 基础使用 + +1. 在 ComfyUI 中添加 `UploadImageCryptIO` 节点 +2. 点击 "choose file to upload" 按钮 +3. 选择图片文件 +4. 图片自动加密、上传，并显示预览 +5. 将节点的 `encrypted` 参数设为 `True` +6. 连接到其他图片处理节点 + +### 2. 密钥管理 + +- **重新生成客户端密钥**：清空 localStorage 中的 `cryptio_client_keypair` +- **重新获取服务端密钥**：清空 localStorage 中的 `cryptio_server_keys` +- **服务端重新生成密钥**：删除 `keys/keys.json` 文件，重启 ComfyUI + +### 3. 兼容性 + +- 支持所有标准图片格式（JPEG, PNG, WebP, etc.） +- 支持透明通道和多帧图片 +- 输出与标准 `LoadImage` 节点完全兼容 + +## 技术栈 + +- **加密算法**： + - RSA-2048 (OAEP, SHA-256) + - AES-256-GCM (图片加密) + - AES-256-CBC (密钥传输) + +- **前端**： + - TypeScript + - Web Crypto API + - ComfyUI Extension API + +- **后端**： + - Python 3 + - cryptography library + - aiohttp + +## 性能考虑 + +- **客户端加密**：现代浏览器的 Web Crypto API 性能优秀 +- **文件大小**：加密后文件略大（约增加 300-400 字节的元数据） +- **内存使用**：服务端解密时需要完整加载文件到内存 +- **缓存策略**：密钥缓存在 localStorage，避免重复交换 + +## 故障排查 + +### 问题：上传失败 + +- 检查浏览器控制台错误信息 +- 确认 `/cryptio/exchange_keys` 端点可访问 +- 验证客户端密钥对已生成 + +### 问题：预览不显示 + +- 确认文件名以 `.encrypted` 结尾 +- 检查服务端私钥是否正确获取 +- 查看控制台解密错误信息 + +### 问题：节点执行失败 + +- 确认 `encrypted` 参数设为 `True` +- 验证加密文件存在于 `input/` 目录 +- 检查服务端私钥是否可用 + +## 未来改进 + +- [ ] 支持批量加密上传 +- [ ] 添加加密进度显示 +- [ ] 支持流式加密（大文件） +- [ ] 密钥轮换机制 +- [ ] 加密文件完整性校验 diff --git a/src/cryptio/api.py b/src/cryptio/api.py index 22f9102..a2d998b 100644 --- a/src/cryptio/api.py +++ b/src/cryptio/api.py @@ -1,13 +1,101 @@ import os import uuid +import json +import base64 from aiohttp import web from server import PromptServer -from .keys import _get_keys +from .keys import _get_keys, update_client_public_key import folder_paths +from cryptography.hazmat.primitives.asymmetric import padding +from cryptography.hazmat.primitives import serialization, hashes +from cryptography.hazmat.backends import default_backend + + +@PromptServer.instance.routes.post("/cryptio/exchange_keys") +async def exchange_keys(request): + """ + 密钥交换端点： + 1. 接收客户端公钥 + 2. 保存客户端公钥到服务器 + 3. 使用客户端公钥加密服务端密钥对 + 4. 返回加密后的服务端密钥 + """ + try: + data = await request.json() + client_public_key_pem = data.get("client_public_key") + + if not client_public_key_pem: + return web.json_response({"error": "client_public_key is required"}, status=400) + + # 保存客户端公钥 + client_public_key_bytes = client_public_key_pem.encode("utf-8") + update_client_public_key(client_public_key_bytes) + + # 获取服务端密钥 + keys = _get_keys() + server_public_key = keys.get("server_public_key") + server_private_key = keys.get("server_private_key") + + if not server_public_key or not server_private_key: + return web.json_response({"error": "Server keys not found"}, status=500) + + # 使用客户端公钥加密服务端密钥 + client_public_key = serialization.load_pem_public_key( + client_public_key_bytes, backend=default_backend() + ) + + # 准备要加密的数据（服务端公钥和私钥） + keys_data = json.dumps({ + "server_public_key": server_public_key.decode("utf-8"), + "server_private_key": server_private_key.decode("utf-8"), + }) + + # 由于RSA加密有大小限制，使用混合加密方案 + # 1. 生成随机AES密钥 + from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes + from cryptography.hazmat.primitives import padding as sym_padding + import secrets + + aes_key = secrets.token_bytes(32) # 256-bit AES key + iv = secrets.token_bytes(16) # 128-bit IV + + # 2. 使用AES加密数据 + cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend()) + encryptor = cipher.encryptor() + + # PKCS7 padding + padder = sym_padding.PKCS7(128).padder() + padded_data = padder.update(keys_data.encode("utf-8")) + padder.finalize() + + encrypted_data = encryptor.update(padded_data) + encryptor.finalize() + + # 3. 使用RSA加密AES密钥 + encrypted_aes_key = client_public_key.encrypt( + aes_key, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + + # 4. 返回加密后的数据 + return web.json_response({ + "encrypted_aes_key": base64.b64encode(encrypted_aes_key).decode("ascii"), + "iv": base64.b64encode(iv).decode("ascii"), + "encrypted_data": base64.b64encode(encrypted_data).decode("ascii"), + }) + + except Exception as e: + print(f"Error in key exchange: {e}") + import traceback + traceback.print_exc() + return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.get("/cryptio/public_key") async def get_public_key(request): + """兼容旧的公钥获取接口（不加密传输）""" keys = _get_keys() pub = keys.get("server_public_key") if pub: @@ -55,4 +143,33 @@ async def upload_encrypted_image(request): except Exception as e: print(f"Error uploading encrypted image: {e}") + return web.json_response({"error": str(e)}, status=500) + + +@PromptServer.instance.routes.get("/cryptio/view_encrypted") +async def view_encrypted_image(request): + """获取加密图片用于前端预览（返回加密数据，由前端解密）""" + try: + filename = request.query.get("filename") + if not filename: + return web.json_response({"error": "filename parameter is required"}, status=400) + + # 获取文件路径 + image_path = folder_paths.get_annotated_filepath(filename) + + if not os.path.exists(image_path): + return web.json_response({"error": "File not found"}, status=404) + + # 读取加密文件 + with open(image_path, 'rb') as f: + encrypted_data = f.read() + + # 返回base64编码的加密数据 + return web.json_response({ + "encrypted_data": base64.b64encode(encrypted_data).decode("ascii"), + "filename": filename + }) + + except Exception as e: + print(f"Error viewing encrypted image: {e}") return web.json_response({"error": str(e)}, status=500) \ No newline at end of file diff --git a/src/cryptio/keys.py b/src/cryptio/keys.py index 7cb01b0..6e04d94 100644 --- a/src/cryptio/keys.py +++ b/src/cryptio/keys.py @@ -111,4 +111,24 @@ def _get_keys() -> dict: else: _KEY_CACHE = {} _KEY_CACHE_MTIME = None - return _KEY_CACHE \ No newline at end of file + return _KEY_CACHE + + +def update_client_public_key(client_public_pem: bytes): + """更新客户端公钥到 keys.json""" + global _KEY_CACHE, _KEY_CACHE_MTIME + + # 确保服务端密钥存在 + _ensure_keys_exist() + keys = _load_keys_json() + + # 更新客户端公钥 + _save_keys_json( + keys["server_public_key"], + keys["server_private_key"], + client_public_pem + ) + + # 清空缓存，强制重新加载 + _KEY_CACHE = None + _KEY_CACHE_MTIME = None \ No newline at end of file diff --git a/ts/cryptioUploadImage.ts b/ts/cryptioUploadImage.ts new file mode 100644 index 0000000..94da90b --- /dev/null +++ b/ts/cryptioUploadImage.ts @@ -0,0 +1,546 @@ +//@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 CLIENT_KEYPAIR_STORAGE_KEY = "cryptio_client_keypair"; +const SERVER_KEYS_STORAGE_KEY = "cryptio_server_keys"; + +// 密钥接口 +interface KeyPair { + publicKey: string; + privateKey: string; +} + +interface ServerKeys { + publicKey: string; + privateKey: string; +} + +// 生成客户端密钥对 +async function generateClientKeyPair(): Promise { + // 生成RSA密钥对 + const keyPair = await window.crypto.subtle.generateKey( + { + name: "RSA-OAEP", + modulusLength: 2048, + publicExponent: new Uint8Array([1, 0, 1]), + hash: "SHA-256", + }, + true, + ["encrypt", "decrypt"] + ); + + // 导出公钥为PEM格式 + const publicKeyBuffer = await window.crypto.subtle.exportKey( + "spki", + keyPair.publicKey + ); + const publicKeyBase64 = btoa( + String.fromCharCode(...new Uint8Array(publicKeyBuffer)) + ); + const publicKeyPem = + "-----BEGIN PUBLIC KEY-----\n" + + publicKeyBase64.match(/.{1,64}/g)!.join("\n") + + "\n-----END PUBLIC KEY-----"; + + // 导出私钥为PEM格式 + const privateKeyBuffer = await window.crypto.subtle.exportKey( + "pkcs8", + keyPair.privateKey + ); + const privateKeyBase64 = btoa( + String.fromCharCode(...new Uint8Array(privateKeyBuffer)) + ); + const privateKeyPem = + "-----BEGIN PRIVATE KEY-----\n" + + privateKeyBase64.match(/.{1,64}/g)!.join("\n") + + "\n-----END PRIVATE KEY-----"; + + return { + publicKey: publicKeyPem, + privateKey: privateKeyPem, + }; +} + +// 获取或生成客户端密钥对 +async function getClientKeyPair(): Promise { + // 从localStorage获取 + const stored = localStorage.getItem(CLIENT_KEYPAIR_STORAGE_KEY); + if (stored) { + try { + return JSON.parse(stored); + } catch (e) { + console.error("Failed to parse stored client keypair:", e); + } + } + + // 生成新的密钥对 + console.log("Generating new client keypair..."); + const keyPair = await generateClientKeyPair(); + + // 保存到localStorage + localStorage.setItem(CLIENT_KEYPAIR_STORAGE_KEY, JSON.stringify(keyPair)); + + return keyPair; +} + +// 从PEM导入公钥 +async function importPublicKeyFromPem(pem: string): Promise { + const pemHeader = "-----BEGIN PUBLIC KEY-----"; + const pemFooter = "-----END PUBLIC KEY-----"; + const pemContents = pem + .substring( + pem.indexOf(pemHeader) + pemHeader.length, + pem.indexOf(pemFooter) + ) + .replace(/\s/g, ""); + + const binaryDer = window.atob(pemContents); + const binaryDerArray = new Uint8Array(binaryDer.length); + for (let i = 0; i < binaryDer.length; i++) { + binaryDerArray[i] = binaryDer.charCodeAt(i); + } + + return await window.crypto.subtle.importKey( + "spki", + binaryDerArray, + { + name: "RSA-OAEP", + hash: "SHA-256", + }, + false, + ["encrypt"] + ); +} + +// 从PEM导入私钥 +async function importPrivateKeyFromPem(pem: string): Promise { + const pemHeader = "-----BEGIN PRIVATE KEY-----"; + const pemFooter = "-----END PRIVATE KEY-----"; + const pemContents = pem + .substring( + pem.indexOf(pemHeader) + pemHeader.length, + pem.indexOf(pemFooter) + ) + .replace(/\s/g, ""); + + const binaryDer = window.atob(pemContents); + const binaryDerArray = new Uint8Array(binaryDer.length); + for (let i = 0; i < binaryDer.length; i++) { + binaryDerArray[i] = binaryDer.charCodeAt(i); + } + + return await window.crypto.subtle.importKey( + "pkcs8", + binaryDerArray, + { + name: "RSA-OAEP", + hash: "SHA-256", + }, + false, + ["decrypt"] + ); +} + +// 与服务器交换密钥 +async function exchangeKeys(): Promise { + // 检查缓存 + const cached = localStorage.getItem(SERVER_KEYS_STORAGE_KEY); + if (cached) { + try { + return JSON.parse(cached); + } catch (e) { + console.error("Failed to parse stored server keys:", e); + } + } + + // 获取客户端密钥对 + const clientKeyPair = await getClientKeyPair(); + + // 发送客户端公钥到服务器 + const response = await fetch("/cryptio/exchange_keys", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + client_public_key: clientKeyPair.publicKey, + }), + }); + + if (!response.ok) { + throw new Error(`Key exchange failed: ${response.statusText}`); + } + + const data = await response.json(); + + // 使用客户端私钥解密服务端密钥 + const clientPrivateKey = await importPrivateKeyFromPem( + clientKeyPair.privateKey + ); + + // 1. 解密AES密钥 + const encryptedAesKey = Uint8Array.from( + atob(data.encrypted_aes_key), + (c) => c.charCodeAt(0) + ); + const aesKeyBuffer = await window.crypto.subtle.decrypt( + { + name: "RSA-OAEP", + }, + clientPrivateKey, + encryptedAesKey + ); + + // 2. 导入AES密钥 + const aesKey = await window.crypto.subtle.importKey( + "raw", + aesKeyBuffer, + { + name: "AES-CBC", + length: 256, + }, + false, + ["decrypt"] + ); + + // 3. 解密服务端密钥数据 + const iv = Uint8Array.from(atob(data.iv), (c) => c.charCodeAt(0)); + const encryptedData = Uint8Array.from(atob(data.encrypted_data), (c) => + c.charCodeAt(0) + ); + + const decryptedBuffer = await window.crypto.subtle.decrypt( + { + name: "AES-CBC", + iv: iv, + }, + aesKey, + encryptedData + ); + + // 4. 去除PKCS7 padding + const decryptedArray = new Uint8Array(decryptedBuffer); + const paddingLength = decryptedArray[decryptedArray.length - 1]; + const unpaddedData = decryptedArray.slice(0, -paddingLength); + + // 5. 解析JSON + const decryptedText = new TextDecoder().decode(unpaddedData); + const serverKeys = JSON.parse(decryptedText); + + // 保存到localStorage + localStorage.setItem(SERVER_KEYS_STORAGE_KEY, JSON.stringify(serverKeys)); + + return serverKeys; +} + +// 获取服务端密钥 +async function getServerKeys(): Promise { + return await exchangeKeys(); +} + +// 使用混合加密方案加密文件 +async function encryptFile(file: File): Promise { + const serverKeys = await getServerKeys(); + + // 读取文件内容 + 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"] + ); + + // 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. 使用服务端公钥加密 AES 密钥 + const serverPublicKey = await importPublicKeyFromPem( + serverKeys.publicKey + ); + const encryptedAesKey = await window.crypto.subtle.encrypt( + { + name: "RSA-OAEP", + }, + serverPublicKey, + rawAesKey + ); + + // 6. 组合数据 + const encryptedAesKeyArray = new Uint8Array(encryptedAesKey); + const encryptedDataArray = new Uint8Array(encryptedData); + const combined = new Uint8Array( + 4 + encryptedAesKeyArray.length + 4 + iv.length + encryptedDataArray.length + ); + + const view = new DataView(combined.buffer); + view.setUint32(0, encryptedAesKeyArray.length, false); + combined.set(encryptedAesKeyArray, 4); + view.setUint32(4 + encryptedAesKeyArray.length, iv.length, false); + combined.set(iv, 4 + encryptedAesKeyArray.length + 4); + combined.set(encryptedDataArray, 4 + encryptedAesKeyArray.length + 4 + iv.length); + + return combined; +} + +// 解密文件（用于预览） +async function decryptFile(encryptedData: Uint8Array): Promise { + const serverKeys = await getServerKeys(); + + // 解析组合数据 + let offset = 0; + + // 1. 读取加密的AES密钥长度 + const view = new DataView(encryptedData.buffer); + const encryptedAesKeyLength = view.getUint32(offset, false); + offset += 4; + + // 2. 读取加密的AES密钥 + const encryptedAesKey = encryptedData.slice( + offset, + offset + encryptedAesKeyLength + ); + offset += encryptedAesKeyLength; + + // 3. 读取IV长度 + const ivLength = view.getUint32(offset, false); + offset += 4; + + // 4. 读取IV + const iv = encryptedData.slice(offset, offset + ivLength); + offset += ivLength; + + // 5. 读取加密的数据 + const encrypted = encryptedData.slice(offset); + + // 6. 使用服务端私钥解密AES密钥 + const serverPrivateKey = await importPrivateKeyFromPem( + serverKeys.privateKey + ); + const aesKeyBuffer = await window.crypto.subtle.decrypt( + { + name: "RSA-OAEP", + }, + serverPrivateKey, + encryptedAesKey + ); + + // 7. 导入AES密钥 + const aesKey = await window.crypto.subtle.importKey( + "raw", + aesKeyBuffer, + { + name: "AES-GCM", + length: 256, + }, + false, + ["decrypt"] + ); + + // 8. 使用AES密钥解密数据 + const decryptedBuffer = await window.crypto.subtle.decrypt( + { + name: "AES-GCM", + iv: iv, + }, + aesKey, + encrypted + ); + + return new Uint8Array(decryptedBuffer); +} + +// 自定义上传函数 +async function uploadEncryptedImage(file: File): Promise { + console.log("Encrypting and uploading image:", file.name); + + // 加密文件 + const encryptedData = await encryptFile(file); + + // 创建FormData + const formData = new FormData(); + const encryptedBlob = new Blob([encryptedData], { + type: "application/octet-stream", + }); + 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(); +} + +// 加载并解密图片用于预览 +async function loadEncryptedImageForPreview(filename: string): Promise { + // 从服务器获取加密数据 + const response = await fetch( + `/cryptio/view_encrypted?filename=${encodeURIComponent(filename)}` + ); + + if (!response.ok) { + throw new Error(`Failed to load encrypted image: ${response.statusText}`); + } + + const data = await response.json(); + + // 解密数据 + const encryptedData = Uint8Array.from(atob(data.encrypted_data), (c) => + c.charCodeAt(0) + ); + const decryptedData = await decryptFile(encryptedData); + + // 转换为Blob URL + const blob = new Blob([decryptedData]); + return URL.createObjectURL(blob); +} + +// 注册扩展 +app.registerExtension({ + name: "cryptio.UploadImageCryptIO", + async setup() { + // 初始化时交换密钥 + try { + await exchangeKeys(); + console.log("CryptIO: Keys exchanged successfully"); + } catch (error) { + console.error("CryptIO: Failed to exchange keys:", error); + } + }, + async beforeRegisterNodeDef( + nodeType: any, + nodeData: ComfyNodeDef, + app: ComfyApp + ) { + if (nodeType.comfyClass === "UploadImageCryptIO") { + // 添加自定义的upload widget + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function (this: any) { + const r = onNodeCreated?.apply(this, arguments); + + // 添加上传button widget + const uploadWidget = this.addWidget( + "button", + "choose file to upload (cryptIO)", + "image", + () => { + const fileInput = document.createElement("input"); + fileInput.type = "file"; + fileInput.accept = "image/*"; + fileInput.style.display = "none"; + document.body.appendChild(fileInput); + + fileInput.onchange = async () => { + if (fileInput.files && fileInput.files.length > 0) { + const file = fileInput.files[0]; + try { + // 上传加密图片 + const result = await uploadEncryptedImage(file); + + // 更新image widget的值 + const imageWidget = this.widgets.find( + (w: any) => w.name === "image" + ); + if (imageWidget) { + imageWidget.value = result.name; + } + + // 更新encrypted标志 + const encryptedWidget = this.widgets.find( + (w: any) => w.name === "encrypted" + ); + if (encryptedWidget) { + encryptedWidget.value = true; + } + + // 更新预览 + await this.updatePreview(result.name); + + app.graph?.setDirtyCanvas(true, false); + } catch (error) { + console.error("Upload error:", error); + alert("Failed to upload encrypted image: " + error); + } + } + document.body.removeChild(fileInput); + }; + + fileInput.click(); + } + ); + + // 添加预览更新方法 + this.updatePreview = async function (filename: string) { + if (!filename || !filename.endsWith(".encrypted")) { + return; + } + + try { + const imageUrl = await loadEncryptedImageForPreview(filename); + + // 更新节点的图片显示 + const img = new Image(); + img.onload = () => { + this.imgs = [img]; + app.graph?.setDirtyCanvas(true, false); + }; + img.src = imageUrl; + } catch (error) { + console.error("Preview error:", error); + } + }; + + // 当image值改变时更新预览 + const imageWidget = this.widgets.find((w: any) => w.name === "image"); + if (imageWidget) { + const originalCallback = imageWidget.callback; + imageWidget.callback = async function (this: any, value: any) { + if (originalCallback) { + originalCallback.call(this, value); + } + if (value && value.endsWith(".encrypted")) { + await this.node.updatePreview(value); + } + }.bind({ node: this }); + } + + return r; + }; + } + }, +});