diff --git a/src/cryptio/save_image.py b/src/cryptio/save_image.py index 950d691..2067c20 100644 --- a/src/cryptio/save_image.py +++ b/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( diff --git a/ts/utils/cryptoUtils.ts b/ts/utils/cryptoUtils.ts index 17edf30..7d8b9a8 100644 --- a/ts/utils/cryptoUtils.ts +++ b/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 { - const serverKeys = await getServerKeysWithExchange(); - - // 读取文件内容 - const arrayBuffer = await file.arrayBuffer(); - const data = new Uint8Array(arrayBuffer); - +async function encryptData(data: Uint8Array, publicKeyPem: string): Promise { // 1. 生成随机 AES 密钥 const aesKey = await window.crypto.subtle.generateKey( { @@ -43,15 +38,13 @@ export async function encryptFile(file: File): Promise { // 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 { } /** - * 解密文件数据，使用服务端私钥解密 AES 密钥 + * 使用混合加密方案加密文件（UploadImage场景，使用服务端公钥） + * @param file 要加密的文件 + * @returns 加密后的数据 + */ +export async function encryptFileWithServerKey(file: File): Promise { + 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 { + // 读取文件内容 + 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 { - const serverKeys = await getServerKeys(); - +async function decryptData(encryptedData: Uint8Array, privateKeyPem: string): Promise { // 解析组合数据 let offset = 0; @@ -106,15 +127,13 @@ export async function decryptFile(encryptedData: Uint8Array): Promise { + const serverKeys = await getServerKeysWithExchange(); + return await decryptData(encryptedData, serverKeys.privateKey); +} + +/** + * 解密文件数据（SaveImage/PreviewImage场景，使用客户端私钥） + * @param encryptedData 加密的数据 + * @returns 解密后的数据 + */ +export async function decryptFileWithClientKey(encryptedData: Uint8Array): Promise { + const clientKeyPair = await getClientKeyPair(); + return await decryptData(encryptedData, clientKeyPair.privateKey); +} + +/** + * 解密文件数据（支持自定义密钥对） + * @param encryptedData 加密的数据 + * @param keyPair 密钥对（使用其中的私钥解密） + * @returns 解密后的数据 + */ +export async function decryptFileWithKeyPair(encryptedData: Uint8Array, keyPair: KeyPair): Promise { + return await decryptData(encryptedData, keyPair.privateKey); +} diff --git a/ts/utils/imageLoader.ts b/ts/utils/imageLoader.ts index dc322da..c43a7c4 100644 --- a/ts/utils/imageLoader.ts +++ b/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)], { diff --git a/ts/utils/uploadUtils.ts b/ts/utils/uploadUtils.ts index c405aa5..e278289 100644 --- a/ts/utils/uploadUtils.ts +++ b/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 { 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();