uploadImage cryptIO temp
This commit is contained in:
Suzu008
+46
-1
@@ -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)
|
||||
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)
|
||||
@@ -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": "加密图片上传",
|
||||
}
|
||||
|
||||
@@ -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
|
||||
@@ -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<string | null> {
|
||||
// 先从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<CryptoKey> {
|
||||
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<Uint8Array> {
|
||||
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<any> {
|
||||
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);
|
||||
}
|
||||
};
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user