diff --git a/ts/cryptioSaveImage.ts b/ts/cryptioSaveImage.ts index 045e56f..1dde129 100644 --- a/ts/cryptioSaveImage.ts +++ b/ts/cryptioSaveImage.ts @@ -8,190 +8,11 @@ import type { ComfyNodeDef, } from "@comfyorg/comfyui-frontend-types"; +import { loadEncryptedImageFromParams, downloadDecryptedImage } from "./utils/imageLoader.js"; + const app: ComfyApp = rawApp; const api: ComfyApi = rawApi; -// 密钥存储键名 -const SERVER_KEYS_STORAGE_KEY = "cryptio_server_keys"; - -// 密钥接口 -interface ServerKeys { - publicKey: string; - privateKey: string; -} - -// 从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 getServerKeys(): 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); - } - } - throw new Error("Server keys not found. Please upload an image first to exchange keys."); -} - -// 解密文件（用于预览） -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); -} - -// 从文件名中提取MIME类型 -function getMimeTypeFromFilename(filename: string): string { - const baseFilename = filename.replace(/\.encrypted$/, ''); - // 提取原始扩展名 - const lastDotIndex = baseFilename.lastIndexOf('.'); - if (lastDotIndex === -1) { - // 如果没有扩展名，默认返回image/png - return 'image/png'; - } - - const extension = baseFilename.substring(lastDotIndex + 1).toLowerCase(); - - // 扩展名到MIME类型的映射 - const mimeTypes: { [key: string]: string } = { - // 图片格式 - 'jpg': 'image/jpeg', - 'jpeg': 'image/jpeg', - 'png': 'image/png', - 'gif': 'image/gif', - 'webp': 'image/webp', - 'bmp': 'image/bmp', - 'svg': 'image/svg+xml', - 'ico': 'image/x-icon', - 'tiff': 'image/tiff', - 'tif': 'image/tiff', - }; - - return mimeTypes[extension] || 'image/png'; // 默认返回image/png -} - -// 加载并解密图片用于预览 -async function loadEncryptedImageForPreview(params: { - filename: string; - subfolder?: string; - type: string; -}): Promise { - // 构建URL - const url = api.apiURL( - `/view?filename=${encodeURIComponent(params.filename)}&type=${params.type}${params.subfolder ? `&subfolder=${encodeURIComponent(params.subfolder)}` : "" - }` - ); - - // 从服务器获取加密数据（二进制格式） - const response = await fetch(url); - - if (!response.ok) { - throw new Error(`Failed to load encrypted image: ${response.statusText}`); - } - - // 直接读取二进制数据 - const encryptedArrayBuffer = await response.arrayBuffer(); - const encryptedData = new Uint8Array(encryptedArrayBuffer); - - // 解密数据 - const decryptedData = await decryptFile(encryptedData); - - // 转换为Blob URL - const blob = new Blob([new Uint8Array(decryptedData)], { - type: getMimeTypeFromFilename(params.filename), - }); - return URL.createObjectURL(blob); -} - // 注册扩展 app.registerExtension({ name: "cryptio.SaveImageCryptIO", @@ -211,13 +32,13 @@ app.registerExtension({ // 处理加密图片 if (message?.cryptio_images) { - const autoDownload = this.widgets.find((n: any) => n.name === "auto_download").value + const autoDownload = this.widgets.find((n: any) => n.name === "auto_download")?.value; for (let i = 0; i < message.cryptio_images.length; i++) { const imageInfo = message.cryptio_images[i]; if (imageInfo.filename && imageInfo.filename.endsWith(".encrypted")) { // 异步加载并解密图片 - loadEncryptedImageForPreview(imageInfo) + loadEncryptedImageFromParams(api, imageInfo) .then((imageUrl) => { // 更新节点的图片显示 const img = new Image(); @@ -230,16 +51,8 @@ app.registerExtension({ // 如果开启了自动下载 if (autoDownload) { - // 创建下载链接 - const a = document.createElement("a"); - a.href = imageUrl; - // 移除 .encrypted 后缀 - const originalFilename = imageInfo.filename.replace(/\.encrypted$/, ""); - a.download = originalFilename; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - console.log(`Auto-downloaded decrypted image: ${originalFilename}`); + downloadDecryptedImage(imageUrl, imageInfo.filename); + console.log(`Auto-downloaded decrypted image: ${imageInfo.filename.replace(/\.encrypted$/, "")}`); } }; img.src = imageUrl; @@ -267,15 +80,9 @@ app.registerExtension({ for (let i = 0; i < this.imgs.length; i++) { const img = this.imgs[i]; if (img && img.src) { - // 创建下载链接 - const a = document.createElement("a"); - a.href = img.src; - // 移除 .encrypted 后缀 - const originalFilename = this.images?.[i]?.filename?.replace(/\.encrypted$/, "") || `image_${i}.png`; - a.download = originalFilename; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); + // 获取原始文件名 + const filename = this.images?.[i]?.filename || `image_${i}.png.encrypted`; + downloadDecryptedImage(img.src, filename); } } }, diff --git a/ts/cryptioUploadImage.ts b/ts/cryptioUploadImage.ts index 15a169a..e3c0fe3 100644 --- a/ts/cryptioUploadImage.ts +++ b/ts/cryptioUploadImage.ts @@ -8,486 +8,12 @@ import type { ComfyNodeDef, } from "@comfyorg/comfyui-frontend-types"; +import { exchangeKeys } from "./utils/cryptoKeys.js"; +import { handleFileUpload, createPreviewUpdateFunction } from "./utils/uploadUtils.js"; + 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. 解析JSON - const decryptedText = new TextDecoder().decode(decryptedBuffer); - const rawServerKeys = JSON.parse(decryptedText); - const serverKeys = { - publicKey: rawServerKeys.server_public_key, - privateKey: rawServerKeys.server_private_key, - }; - - // 5. 保存到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 = file?.name?.endsWith(".encrypted") ? new Uint8Array(await file.arrayBuffer()) : await encryptFile(file); - - // 创建FormData - const formData = new FormData(); - const encryptedBlob = new Blob([new Uint8Array(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(); -} - -// 从文件名中提取MIME类型 -function getMimeTypeFromFilename(filename: string): string { - const baseFilename = filename.replace(/\.encrypted$/, ''); - // 提取原始扩展名 - const lastDotIndex = baseFilename.lastIndexOf('.'); - if (lastDotIndex === -1) { - // 如果没有扩展名，默认返回image/png - return 'image/png'; - } - - const extension = baseFilename.substring(lastDotIndex + 1).toLowerCase(); - - // 扩展名到MIME类型的映射 - const mimeTypes: { [key: string]: string } = { - // 图片格式 - 'jpg': 'image/jpeg', - 'jpeg': 'image/jpeg', - 'png': 'image/png', - 'gif': 'image/gif', - 'webp': 'image/webp', - 'bmp': 'image/bmp', - 'svg': 'image/svg+xml', - 'ico': 'image/x-icon', - 'tiff': 'image/tiff', - 'tif': 'image/tiff', - // 视频格式 - 'mp4': 'video/mp4', - 'webm': 'video/webm', - 'mkv': 'video/x-matroska' - }; - - return mimeTypes[extension] || 'image/png'; // 默认返回image/png -} - -// 加载并解密图片用于预览 -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 encryptedArrayBuffer = await response.arrayBuffer(); - const encryptedData = new Uint8Array(encryptedArrayBuffer); - - // 解密数据 - const decryptedData = await decryptFile(encryptedData); - - // 转换为Blob URL - const blob = new Blob([new Uint8Array(decryptedData)], { - type: getMimeTypeFromFilename(filename), - }); - return URL.createObjectURL(blob); -} - -// 处理文件上传的通用方法 -async function handleFileUpload(node: any, file: File) { - try { - // 上传加密图片 - const result = await uploadEncryptedImage(file); - - // 更新image widget的值 - const imageWidget = node.widgets.find( - (w: any) => w.name === "image" - ); - if (imageWidget) { - imageWidget.value = result.name; - } - - // 更新预览 - await node.updatePreview(result.name); - - app.rootGraph?.setDirtyCanvas(true, false); - } catch (error) { - console.error("Upload error:", error); - alert("Failed to upload encrypted image: " + error); - } -} - // 注册扩展 app.registerExtension({ name: "cryptio.UploadImageCryptIO", @@ -528,7 +54,7 @@ app.registerExtension({ fileInput.onchange = async () => { if (fileInput.files && fileInput.files.length > 0) { const file = fileInput.files[0]; - await handleFileUpload(this, file); + await handleFileUpload(this, file, app); } document.body.removeChild(fileInput); }; @@ -539,25 +65,7 @@ app.registerExtension({ uploadWidget.label = "choose file to upload"; // 添加预览更新方法 - 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.rootGraph?.setDirtyCanvas(true, false); - }; - img.src = imageUrl; - } catch (error) { - console.error("Preview error:", error); - } - }; + this.updatePreview = createPreviewUpdateFunction(app); // 当image值改变时更新预览 const imageWidget = this.widgets.find((w: any) => w.name === "image"); @@ -607,7 +115,7 @@ app.registerExtension({ } // 上传文件 - await handleFileUpload(this, file); + await handleFileUpload(this, file, app); return true; }; @@ -620,7 +128,7 @@ app.registerExtension({ continue } // 上传文件 - await handleFileUpload(this, file); + await handleFileUpload(this, file, app); return true; } diff --git a/ts/utils/cryptoKeys.ts b/ts/utils/cryptoKeys.ts new file mode 100644 index 0000000..dca19cd --- /dev/null +++ b/ts/utils/cryptoKeys.ts @@ -0,0 +1,272 @@ +/** + * 密钥管理模块 + * 负责客户端和服务端密钥的生成、导入、交换和存储 + */ + +// 密钥存储键名 +export const CLIENT_KEYPAIR_STORAGE_KEY = "cryptio_client_keypair"; +export const SERVER_KEYS_STORAGE_KEY = "cryptio_server_keys"; + +// 密钥接口 +export interface KeyPair { + publicKey: string; + privateKey: string; +} + +export interface ServerKeys { + publicKey: string; + privateKey: string; +} + +/** + * 生成客户端RSA密钥对 + */ +export 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, + }; +} + +/** + * 获取或生成客户端密钥对 + */ +export 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格式导入公钥 + */ +export 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格式导入私钥 + */ +export 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"] + ); +} + +/** + * 与服务器交换密钥，使用客户端公钥加密 AES 密钥 + */ +export 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. 解析JSON + const decryptedText = new TextDecoder().decode(decryptedBuffer); + const rawServerKeys = JSON.parse(decryptedText); + const serverKeys = { + publicKey: rawServerKeys.server_public_key, + privateKey: rawServerKeys.server_private_key, + }; + + // 5. 保存到localStorage + localStorage.setItem(SERVER_KEYS_STORAGE_KEY, JSON.stringify(serverKeys)); + + return serverKeys; +} + +/** + * 获取服务端密钥 + */ +export async function getServerKeys(): 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); + } + } + + // 如果缓存不存在，抛出错误（SaveImage场景） + // 或者尝试交换密钥（UploadImage场景） + throw new Error("Server keys not found. Please upload an image first to exchange keys."); +} + +/** + * 获取服务端密钥（带自动交换） + */ +export async function getServerKeysWithExchange(): Promise { + return await exchangeKeys(); +} diff --git a/ts/utils/cryptoUtils.ts b/ts/utils/cryptoUtils.ts new file mode 100644 index 0000000..17edf30 --- /dev/null +++ b/ts/utils/cryptoUtils.ts @@ -0,0 +1,144 @@ +/** + * 加密/解密工具模块 + * 提供文件和数据的加密解密功能 + */ + +import { getServerKeysWithExchange, getServerKeys, importPublicKeyFromPem, importPrivateKeyFromPem } from "./cryptoKeys.js"; + +/** + * 使用混合加密方案加密文件，使用服务端公钥加密 AES 密钥 + * @param file 要加密的文件 + * @returns 加密后的数据 + */ +export async function encryptFile(file: File): Promise { + const serverKeys = await getServerKeysWithExchange(); + + // 读取文件内容 + 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; +} + +/** + * 解密文件数据，使用服务端私钥解密 AES 密钥 + * @param encryptedData 加密的数据 + * @returns 解密后的数据 + */ +export 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); +} diff --git a/ts/utils/fileUtils.ts b/ts/utils/fileUtils.ts new file mode 100644 index 0000000..ad1022c --- /dev/null +++ b/ts/utils/fileUtils.ts @@ -0,0 +1,42 @@ +/** + * 文件工具模块 + * 提供文件相关的实用函数 + */ + +/** + * 从文件名中提取MIME类型 + * @param filename 文件名 + * @returns MIME类型字符串 + */ +export function getMimeTypeFromFilename(filename: string): string { + const baseFilename = filename.replace(/\.encrypted$/, ''); + // 提取原始扩展名 + const lastDotIndex = baseFilename.lastIndexOf('.'); + if (lastDotIndex === -1) { + // 如果没有扩展名，默认返回image/png + return 'image/png'; + } + + const extension = baseFilename.substring(lastDotIndex + 1).toLowerCase(); + + // 扩展名到MIME类型的映射 + const mimeTypes: { [key: string]: string } = { + // 图片格式 + 'jpg': 'image/jpeg', + 'jpeg': 'image/jpeg', + 'png': 'image/png', + 'gif': 'image/gif', + 'webp': 'image/webp', + 'bmp': 'image/bmp', + 'svg': 'image/svg+xml', + 'ico': 'image/x-icon', + 'tiff': 'image/tiff', + 'tif': 'image/tiff', + // 视频格式 + 'mp4': 'video/mp4', + 'webm': 'video/webm', + 'mkv': 'video/x-matroska' + }; + + return mimeTypes[extension] || 'image/png'; // 默认返回image/png +} diff --git a/ts/utils/imageLoader.ts b/ts/utils/imageLoader.ts new file mode 100644 index 0000000..dc322da --- /dev/null +++ b/ts/utils/imageLoader.ts @@ -0,0 +1,95 @@ +/** + * 图片加载模块 + * 负责从服务器加载并解密图片用于预览 + */ + +import type { ComfyApi } from "@comfyorg/comfyui-frontend-types"; +import { decryptFile } from "./cryptoUtils.js"; +import { getMimeTypeFromFilename } from "./fileUtils.js"; + +export interface ImageLoadParams { + filename: string; + subfolder?: string; + type: string; +} + +/** + * 从简单文件名加载加密图片（用于UploadImage场景） + * @param api ComfyUI API实例 + * @param filename 文件名 + * @returns Blob URL + */ +export async function loadEncryptedImageFromFilename(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 encryptedArrayBuffer = await response.arrayBuffer(); + const encryptedData = new Uint8Array(encryptedArrayBuffer); + + // 解密数据 + const decryptedData = await decryptFile(encryptedData); + + // 转换为Blob URL + const blob = new Blob([new Uint8Array(decryptedData)], { + type: getMimeTypeFromFilename(filename), + }); + return URL.createObjectURL(blob); +} + +/** + * 从参数加载加密图片（用于SaveImage/PreviewImage场景） + * @param api ComfyUI API实例 + * @param params 图片参数（包含filename, subfolder, type） + * @returns Blob URL + */ +export async function loadEncryptedImageFromParams(api: ComfyApi, params: ImageLoadParams): Promise { + // 构建URL + const url = api.apiURL( + `/view?filename=${encodeURIComponent(params.filename)}&type=${params.type}${ + params.subfolder ? `&subfolder=${encodeURIComponent(params.subfolder)}` : "" + }` + ); + + // 从服务器获取加密数据（二进制格式） + const response = await fetch(url); + + if (!response.ok) { + throw new Error(`Failed to load encrypted image: ${response.statusText}`); + } + + // 直接读取二进制数据 + const encryptedArrayBuffer = await response.arrayBuffer(); + const encryptedData = new Uint8Array(encryptedArrayBuffer); + + // 解密数据 + const decryptedData = await decryptFile(encryptedData); + + // 转换为Blob URL + const blob = new Blob([new Uint8Array(decryptedData)], { + type: getMimeTypeFromFilename(params.filename), + }); + return URL.createObjectURL(blob); +} + +/** + * 下载解密后的图片 + * @param blobUrl Blob URL + * @param filename 原始文件名（带 .encrypted 后缀） + */ +export function downloadDecryptedImage(blobUrl: string, filename: string): void { + const a = document.createElement("a"); + a.href = blobUrl; + // 移除 .encrypted 后缀 + const originalFilename = filename.replace(/\.encrypted$/, ""); + a.download = originalFilename; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); +} diff --git a/ts/utils/uploadUtils.ts b/ts/utils/uploadUtils.ts new file mode 100644 index 0000000..c405aa5 --- /dev/null +++ b/ts/utils/uploadUtils.ts @@ -0,0 +1,97 @@ +/** + * 上传工具模块 + * 提供文件上传相关功能 + */ + +import type { ComfyApp } from "@comfyorg/comfyui-frontend-types"; +import { encryptFile } 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); + + // 创建FormData + const formData = new FormData(); + const encryptedBlob = new Blob([new Uint8Array(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(); +} + +/** + * 处理文件上传到节点的通用方法 + * @param node 节点实例 + * @param file 要上传的文件 + * @param app ComfyUI App实例 + */ +export async function handleFileUpload(node: any, file: File, app: ComfyApp) { + try { + // 上传加密图片 + const result = await uploadEncryptedImage(file); + + // 更新image widget的值 + const imageWidget = node.widgets.find( + (w: any) => w.name === "image" + ); + if (imageWidget) { + imageWidget.value = result.name; + } + + // 更新预览 + await node.updatePreview(result.name); + + app.rootGraph?.setDirtyCanvas(true, false); + } catch (error) { + console.error("Upload error:", error); + alert("Failed to upload encrypted image: " + error); + } +} + +/** + * 创建节点预览更新函数 + * @param app ComfyUI App实例 + * @returns 预览更新函数 + */ +export function createPreviewUpdateFunction(app: ComfyApp) { + return async function (this: any, filename: string) { + if (!filename || !filename.endsWith(".encrypted")) { + return; + } + + try { + const imageUrl = await loadEncryptedImageFromFilename(filename); + + // 更新节点的图片显示 + const img = new Image(); + img.onload = () => { + this.imgs = [img]; + app.rootGraph?.setDirtyCanvas(true, false); + }; + img.src = imageUrl; + } catch (error) { + console.error("Preview error:", error); + } + }; +}