nodes/Suzu008-ComfyUI-CryptIO
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

remote redundant code, remove original upload widget

Browse Source
This commit is contained in:
Suzu008
2025-10-12 01:14:59 +08:00
parent e5d09afc01
commit 355cd52d8c
2 changed files with 13 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cryptio/upload_image.py
+7 -93
View File
@@ -25,16 +25,11 @@ class UploadImageCryptIO:
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))
all_files = sorted(set(encrypted_files))
return {
"required": {
"image": (all_files, {"image_upload": True}),
"encrypted": ("BOOLEAN", {"default": False}),
},
"required": {"image": (all_files, {"image_upload": True})},
}
RETURN_TYPES = ("IMAGE", "MASK")
@@ -42,14 +37,13 @@ class UploadImageCryptIO:
FUNCTION = "load_image"
CATEGORY = "CryptIO"
def load_image(self, image, encrypted):
def load_image(self, image):
"""
加载并解密图片
如果 encrypted=True，从加密文件读取并解密
"""
try:
# 检查是否是加密文件（文件名以 .encrypted 结尾）
if encrypted and isinstance(image, str) and image.endswith(".encrypted"):
if isinstance(image, str) and image.endswith(".encrypted"):
# 从加密文件读取
image_path = folder_paths.get_annotated_filepath(image)
@@ -80,9 +74,9 @@ class UploadImageCryptIO:
# 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())
private_key = serialization.load_pem_private_key(private_pem, password=None, backend=default_backend()) # type: ignore
aes_key_bytes = private_key.decrypt(
aes_key_bytes = private_key.decrypt( # type: ignore
encrypted_aes_key,
padding.OAEP(
mgf=padding.MGF1(algorithm=hashes.SHA256()),
@@ -105,88 +99,8 @@ class UploadImageCryptIO:
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)
return (None, None)
# 处理图片（与LoadImage节点相同的处理逻辑）
output_images = []
ts/cryptioUploadImage.ts
+6 -11
View File
@@ -451,10 +451,12 @@ app.registerExtension({
nodeType.prototype.onNodeCreated = function (this: any) {
const r = onNodeCreated?.apply(this, arguments);
// 移除默认的upload widget
this.removeWidgetByName("upload")
// 添加上传button widget
const uploadWidget = this.addWidget(
"button",
"choose file to upload (cryptIO)",
"cryptio-upload",
"image",
() => {
const fileInput = document.createElement("input");
@@ -478,18 +480,10 @@ app.registerExtension({
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);
app.rootGraph?.setDirtyCanvas(true, false);
} catch (error) {
console.error("Upload error:", error);
alert("Failed to upload encrypted image: " + error);
@@ -501,6 +495,7 @@ app.registerExtension({
fileInput.click();
}
);
uploadWidget.label = "choose file to upload";
// 添加预览更新方法
this.updatePreview = async function (filename: string) {
@@ -515,7 +510,7 @@ app.registerExtension({
const img = new Image();
img.onload = () => {
this.imgs = [img];
app.graph?.setDirtyCanvas(true, false);
app.rootGraph?.setDirtyCanvas(true, false);
};
img.src = imageUrl;
} catch (error) {