refactor(crypto): 重构加密模块以支持不同密钥场景

- 将通用加密解密逻辑提取为内部函数
- 添加针对不同场景的加密解密方法（服务端/客户端密钥）
- 更新相关调用代码以使用新的加密方法
- 改进注释和类型定义

src/cryptio/save_image.py

@@ -166,7 +166,7 @@ class PreviewImageCryptIO:
 
 def encrypt_data(data: bytes) -> bytes:
     """
-    使用混合加密方案加密数据
+    使用混合加密方案加密数据（使用客户端公钥加密）
     """
     # 1. 生成随机 AES 密钥和 IV
     aes_key = os.urandom(32)  # 256-bit key
@@ -176,9 +176,11 @@ def encrypt_data(data: bytes) -> bytes:
     encryptor = cipher.encryptor()
     encrypted_data = encryptor.update(data) + encryptor.finalize()
     tag = encryptor.tag  # type: ignore # 获取认证标签
-    # 3. 使用服务端公钥加密 AES 密钥
-    public_pem = _get_keys().get("server_public_key")
-    public_key = serialization.load_pem_public_key(public_pem, backend=default_backend())  # type: ignore
+    # 3. 使用客户端公钥加密 AES 密钥（这样只有客户端能解密）
+    client_public_pem = _get_keys().get("client_public_key")
+    if not client_public_pem:
+        raise ValueError("Client public key not found. Please upload an image first to exchange keys.")
+    public_key = serialization.load_pem_public_key(client_public_pem, backend=default_backend())  # type: ignore
     encrypted_aes_key = public_key.encrypt(  # type: ignore
         aes_key,
         padding.OAEP(

ts/utils/cryptoUtils.ts

@@ -3,20 +3,15 @@
  * 提供文件和数据的加密解密功能
  */
 
-import { getServerKeysWithExchange, getServerKeys, importPublicKeyFromPem, importPrivateKeyFromPem } from "./cryptoKeys.js";
+import { getServerKeysWithExchange, getClientKeyPair, importPublicKeyFromPem, importPrivateKeyFromPem, type KeyPair } from "./cryptoKeys.js";
 
 /**
- * 使用混合加密方案加密文件，使用服务端公钥加密 AES 密钥
- * @param file 要加密的文件
+ * 使用混合加密方案加密数据
+ * @param data 要加密的数据
+ * @param publicKeyPem 用于加密AES密钥的公钥（PEM格式）
  * @returns 加密后的数据
  */
-export async function encryptFile(file: File): Promise<Uint8Array> {
-    const serverKeys = await getServerKeysWithExchange();
-
-    // 读取文件内容
-    const arrayBuffer = await file.arrayBuffer();
-    const data = new Uint8Array(arrayBuffer);
-
+async function encryptData(data: Uint8Array, publicKeyPem: string): Promise<Uint8Array> {
     // 1. 生成随机 AES 密钥
     const aesKey = await window.crypto.subtle.generateKey(
         {
@@ -43,15 +38,13 @@ export async function encryptFile(file: File): Promise<Uint8Array> {
     // 4. 导出 AES 密钥
     const rawAesKey = await window.crypto.subtle.exportKey("raw", aesKey);
 
-    // 5. 使用服务端公钥加密 AES 密钥
-    const serverPublicKey = await importPublicKeyFromPem(
-        serverKeys.publicKey
-    );
+    // 5. 使用指定的公钥加密 AES 密钥
+    const publicKey = await importPublicKeyFromPem(publicKeyPem);
     const encryptedAesKey = await window.crypto.subtle.encrypt(
         {
             name: "RSA-OAEP",
         },
-        serverPublicKey,
+        publicKey,
         rawAesKey
     );
 
@@ -73,13 +66,41 @@ export async function encryptFile(file: File): Promise<Uint8Array> {
 }
 
 /**
- * 解密文件数据，使用服务端私钥解密 AES 密钥
+ * 使用混合加密方案加密文件（UploadImage场景，使用服务端公钥）
+ * @param file 要加密的文件
+ * @returns 加密后的数据
+ */
+export async function encryptFileWithServerKey(file: File): Promise<Uint8Array> {
+    const serverKeys = await getServerKeysWithExchange();
+
+    // 读取文件内容
+    const arrayBuffer = await file.arrayBuffer();
+    const data = new Uint8Array(arrayBuffer);
+
+    return await encryptData(data, serverKeys.publicKey);
+}
+
+/**
+ * 使用混合加密方案加密文件（支持自定义密钥对）
+ * @param file 要加密的文件
+ * @param keyPair 密钥对（使用其中的公钥加密）
+ * @returns 加密后的数据
+ */
+export async function encryptFileWithKeyPair(file: File, keyPair: KeyPair): Promise<Uint8Array> {
+    // 读取文件内容
+    const arrayBuffer = await file.arrayBuffer();
+    const data = new Uint8Array(arrayBuffer);
+
+    return await encryptData(data, keyPair.publicKey);
+}
+
+/**
+ * 解密文件数据（使用指定的私钥）
  * @param encryptedData 加密的数据
+ * @param privateKeyPem 用于解密AES密钥的私钥（PEM格式）
  * @returns 解密后的数据
  */
-export async function decryptFile(encryptedData: Uint8Array): Promise<Uint8Array> {
-    const serverKeys = await getServerKeys();
-
+async function decryptData(encryptedData: Uint8Array, privateKeyPem: string): Promise<Uint8Array> {
     // 解析组合数据
     let offset = 0;
 
@@ -106,15 +127,13 @@ export async function decryptFile(encryptedData: Uint8Array): Promise<Uint8Array
     // 5. 读取加密的数据
     const encrypted = encryptedData.slice(offset);
 
-    // 6. 使用服务端私钥解密AES密钥
-    const serverPrivateKey = await importPrivateKeyFromPem(
-        serverKeys.privateKey
-    );
+    // 6. 使用指定的私钥解密AES密钥
+    const privateKey = await importPrivateKeyFromPem(privateKeyPem);
     const aesKeyBuffer = await window.crypto.subtle.decrypt(
         {
             name: "RSA-OAEP",
         },
-        serverPrivateKey,
+        privateKey,
         encryptedAesKey
     );
 
@@ -142,3 +161,33 @@ export async function decryptFile(encryptedData: Uint8Array): Promise<Uint8Array
 
     return new Uint8Array(decryptedBuffer);
 }
+
+/**
+ * 解密文件数据（UploadImage预览场景，使用服务端私钥）
+ * @param encryptedData 加密的数据
+ * @returns 解密后的数据
+ */
+export async function decryptFileWithServerKey(encryptedData: Uint8Array): Promise<Uint8Array> {
+    const serverKeys = await getServerKeysWithExchange();
+    return await decryptData(encryptedData, serverKeys.privateKey);
+}
+
+/**
+ * 解密文件数据（SaveImage/PreviewImage场景，使用客户端私钥）
+ * @param encryptedData 加密的数据
+ * @returns 解密后的数据
+ */
+export async function decryptFileWithClientKey(encryptedData: Uint8Array): Promise<Uint8Array> {
+    const clientKeyPair = await getClientKeyPair();
+    return await decryptData(encryptedData, clientKeyPair.privateKey);
+}
+
+/**
+ * 解密文件数据（支持自定义密钥对）
+ * @param encryptedData 加密的数据
+ * @param keyPair 密钥对（使用其中的私钥解密）
+ * @returns 解密后的数据
+ */
+export async function decryptFileWithKeyPair(encryptedData: Uint8Array, keyPair: KeyPair): Promise<Uint8Array> {
+    return await decryptData(encryptedData, keyPair.privateKey);
+}

ts/utils/imageLoader.ts

@@ -4,7 +4,7 @@
  */
 
 import type { ComfyApi } from "@comfyorg/comfyui-frontend-types";
-import { decryptFile } from "./cryptoUtils.js";
+import { decryptFileWithServerKey, decryptFileWithClientKey } from "./cryptoUtils.js";
 import { getMimeTypeFromFilename } from "./fileUtils.js";
 
 export interface ImageLoadParams {
@@ -14,8 +14,7 @@ export interface ImageLoadParams {
 }
 
 /**
- * 从简单文件名加载加密图片（用于UploadImage场景）
- * @param api ComfyUI API实例
+ * 从简单文件名加载加密图片（用于UploadImage场景，使用服务端私钥解密）
  * @param filename 文件名
  * @returns Blob URL
  */
@@ -33,8 +32,8 @@ export async function loadEncryptedImageFromFilename(filename: string): Promise<
     const encryptedArrayBuffer = await response.arrayBuffer();
     const encryptedData = new Uint8Array(encryptedArrayBuffer);
 
-    // 解密数据
-    const decryptedData = await decryptFile(encryptedData);
+    // 使用服务端私钥解密（UploadImage场景）
+    const decryptedData = await decryptFileWithServerKey(encryptedData);
 
     // 转换为Blob URL
     const blob = new Blob([new Uint8Array(decryptedData)], {
@@ -44,7 +43,7 @@ export async function loadEncryptedImageFromFilename(filename: string): Promise<
 }
 
 /**
- * 从参数加载加密图片（用于SaveImage/PreviewImage场景）
+ * 从参数加载加密图片（用于SaveImage/PreviewImage场景，使用客户端私钥解密）
  * @param api ComfyUI API实例
  * @param params 图片参数（包含filename, subfolder, type）
  * @returns Blob URL
@@ -68,8 +67,8 @@ export async function loadEncryptedImageFromParams(api: ComfyApi, params: ImageL
     const encryptedArrayBuffer = await response.arrayBuffer();
     const encryptedData = new Uint8Array(encryptedArrayBuffer);
 
-    // 解密数据
-    const decryptedData = await decryptFile(encryptedData);
+    // 使用客户端私钥解密（SaveImage/PreviewImage场景）
+    const decryptedData = await decryptFileWithClientKey(encryptedData);
 
     // 转换为Blob URL
     const blob = new Blob([new Uint8Array(decryptedData)], {

ts/utils/uploadUtils.ts

@@ -4,21 +4,21 @@
  */
 
 import type { ComfyApp } from "@comfyorg/comfyui-frontend-types";
-import { encryptFile } from "./cryptoUtils.js";
+import { encryptFileWithServerKey } from "./cryptoUtils.js";
 import { loadEncryptedImageFromFilename } from "./imageLoader.js";
 
 /**
- * 上传加密图片到服务器
+ * 上传加密图片到服务器（使用服务端公钥加密）
  * @param file 要上传的文件
  * @returns 服务器响应
  */
 export async function uploadEncryptedImage(file: File): Promise<any> {
     console.log("Encrypting and uploading image:", file.name);
 
-    // 加密文件（如果文件已经加密则直接读取）
+    // 加密文件（使用服务端公钥，如果文件已经加密则直接读取）
     const encryptedData = file?.name?.endsWith(".encrypted")
         ? new Uint8Array(await file.arrayBuffer())
-        : await encryptFile(file);
+        : await encryptFileWithServerKey(file);
 
     // 创建FormData
     const formData = new FormData();