diff --git a/src/cryptio/api.py b/src/cryptio/api.py index 62ea3eb..cb7a420 100644 --- a/src/cryptio/api.py +++ b/src/cryptio/api.py @@ -3,11 +3,15 @@ import uuid import json import base64 from aiohttp import web -from server import PromptServer -from .keys import _get_keys, update_client_public_key -import folder_paths -from cryptography.hazmat.primitives.asymmetric import padding -from cryptography.hazmat.primitives import serialization, hashes +from server import PromptServer # pyright: ignore[reportMissingImports] +import folder_paths # pyright: ignore[reportMissingImports] +from .utils import _key_manager +from .utils.crypto_utils import ( + encrypt_data_with_public_key, + generate_aes_key_and_iv, + encrypt_with_aes_cbc, +) +from cryptography.hazmat.primitives import serialization from cryptography.hazmat.backends import default_backend @@ -29,55 +33,37 @@ async def exchange_keys(request): # 保存客户端公钥 client_public_key_bytes = client_public_key_pem.encode("utf-8") - update_client_public_key(client_public_key_bytes) + _key_manager.update_client_key(client_public_key_bytes) - # 获取服务端密钥 - keys = _get_keys() - server_public_key = keys.get("server_public_key") - server_private_key = keys.get("server_private_key") + # 获取服务端密钥 PEM 格式 + server_public_key_pem = _key_manager.server_public_key_pem + server_private_key_pem = _key_manager.server_private_key_pem - if not server_public_key or not server_private_key: + if not server_public_key_pem or not server_private_key_pem: return web.json_response({"error": "Server keys not found"}, status=500) # 使用客户端公钥加密服务端密钥 - client_public_key = serialization.load_pem_public_key(client_public_key_bytes, backend=default_backend()) + client_public_key = serialization.load_pem_public_key( + client_public_key_bytes, backend=default_backend() + ) # 准备要加密的数据（服务端公钥和私钥） keys_data = json.dumps( { - "server_public_key": server_public_key.decode("utf-8"), - "server_private_key": server_private_key.decode("utf-8"), + "server_public_key": server_public_key_pem.decode("utf-8"), + "server_private_key": server_private_key_pem.decode("utf-8"), } ) - # 由于RSA加密有大小限制，使用混合加密方案 - # 1. 生成随机AES密钥 - from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes - from cryptography.hazmat.primitives import padding as sym_padding - import secrets - - aes_key = secrets.token_bytes(32) # 256-bit AES key - iv = secrets.token_bytes(16) # 128-bit IV + # 由于RSA加密有大小限制，使用混合加密方案（AES-CBC + RSA-OAEP） + # 1. 生成随机AES密钥和IV + aes_key, iv = generate_aes_key_and_iv() # 2. 使用AES加密数据 - cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend()) - encryptor = cipher.encryptor() - - # PKCS7 padding - padder = sym_padding.PKCS7(128).padder() - padded_data = padder.update(keys_data.encode("utf-8")) + padder.finalize() - - encrypted_data = encryptor.update(padded_data) + encryptor.finalize() + encrypted_data = encrypt_with_aes_cbc(keys_data.encode("utf-8"), aes_key, iv) # 3. 使用RSA加密AES密钥 - encrypted_aes_key = client_public_key.encrypt( - aes_key, - padding.OAEP( - mgf=padding.MGF1(algorithm=hashes.SHA256()), - algorithm=hashes.SHA256(), - label=None, - ), - ) + encrypted_aes_key = encrypt_data_with_public_key(aes_key, client_public_key) # 4. 返回加密后的数据 return web.json_response( @@ -99,10 +85,9 @@ async def exchange_keys(request): @PromptServer.instance.routes.get("/cryptio/public_key") async def get_public_key(request): """兼容旧的公钥获取接口（不加密传输）""" - keys = _get_keys() - pub = keys.get("server_public_key") - if pub: - return web.json_response({"public_key": pub.decode("utf-8")}) + pub_pem = _key_manager.server_public_key_pem + if pub_pem: + return web.json_response({"public_key": pub_pem.decode("utf-8")}) else: return web.json_response({"error": "Public key not found"}, status=404) diff --git a/src/cryptio/key_generator.py b/src/cryptio/key_generator.py index 2e0468d..e390d6b 100644 --- a/src/cryptio/key_generator.py +++ b/src/cryptio/key_generator.py @@ -1,4 +1,4 @@ -from .keys import _ensure_keys_exist, _get_keys +from .utils import _key_manager class KeyGenerator: @@ -7,7 +7,8 @@ class KeyGenerator: """ def __init__(self): - _ensure_keys_exist() + # 密钥会在首次访问 _key_manager 时自动生成 + pass @classmethod def INPUT_TYPES(s): @@ -18,11 +19,7 @@ class KeyGenerator: FUNCTION = "generate" CATEGORY = "CryptIO" - def generate_keys_if_not_exist(self): - """兼容旧接口：确保密钥存在（现使用 JSON 存储）。""" - _ensure_keys_exist() - def generate(self): """返回公钥（使用缓存避免频繁 IO）""" - public_key = _get_keys()["server_public_key"].decode("utf-8") + public_key = _key_manager.server_public_key_pem.decode("utf-8") return (public_key,) \ No newline at end of file diff --git a/src/cryptio/keys.py b/src/cryptio/keys.py deleted file mode 100644 index 6e04d94..0000000 --- a/src/cryptio/keys.py +++ /dev/null @@ -1,134 +0,0 @@ -import os -import json -import base64 -from cryptography.hazmat.primitives.asymmetric import rsa -from cryptography.hazmat.primitives import serialization -from cryptography.hazmat.backends import default_backend - -# 密钥存储路径（JSON + Base64） -KEY_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "keys") -KEY_JSON_PATH = os.path.join(KEY_DIR, "keys.json") - -# 兼容旧版 PEM 文件路径（用于迁移） -SERVER_PUBLIC_KEY_PATH = os.path.join(KEY_DIR, "server_public_key.pem") -SERVER_PRIVATE_KEY_PATH = os.path.join(KEY_DIR, "server_private_key.pem") -CLIENT_PUBLIC_KEY_PATH = os.path.join(KEY_DIR, "client_public_key.pem") - -os.makedirs(KEY_DIR, exist_ok=True) - -# 内存缓存，避免频繁 IO -_KEY_CACHE = None -_KEY_CACHE_MTIME = None - - -def _read_file_bytes(path: str) -> bytes: - with open(path, "rb") as f: - return f.read() - - -def _save_keys_json(server_public_pem: bytes, server_private_pem: bytes, client_public_pem: bytes | None = None): - data = { - "server_public_key": base64.b64encode(server_public_pem).decode("ascii"), - "server_private_key": base64.b64encode(server_private_pem).decode("ascii"), - } - if client_public_pem is not None: - data["client_public_key"] = base64.b64encode(client_public_pem).decode("ascii") - with open(KEY_JSON_PATH, "w", encoding="utf-8") as f: - json.dump(data, f, ensure_ascii=False, indent=2) - - -def _load_keys_json() -> dict: - with open(KEY_JSON_PATH, "r", encoding="utf-8") as f: - raw = json.load(f) - keys = { - "server_public_key": base64.b64decode(raw["server_public_key"]), - "server_private_key": base64.b64decode(raw["server_private_key"]), - } - if "client_public_key" in raw: - keys["client_public_key"] = base64.b64decode(raw["client_public_key"]) - return keys - - -def _ensure_keys_exist(): - """确保 keys.json 存在；若不存在则生成或从旧文件迁移。""" - if os.path.exists(KEY_JSON_PATH): - return - - server_pub = None - server_pri = None - client_pub = None - - # 迁移旧 PEM 文件到 JSON - if os.path.exists(SERVER_PUBLIC_KEY_PATH) and os.path.exists(SERVER_PRIVATE_KEY_PATH): - try: - server_pub = _read_file_bytes(SERVER_PUBLIC_KEY_PATH) - server_pri = _read_file_bytes(SERVER_PRIVATE_KEY_PATH) - if os.path.exists(CLIENT_PUBLIC_KEY_PATH): - client_pub = _read_file_bytes(CLIENT_PUBLIC_KEY_PATH) - except Exception: - server_pub = None - server_pri = None - - # 如果没有旧文件，生成新的密钥对 - if server_pub is None or server_pri is None: - private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048, backend=default_backend()) - public_key = private_key.public_key() - - server_pri = private_key.private_bytes( - encoding=serialization.Encoding.PEM, - format=serialization.PrivateFormat.PKCS8, - encryption_algorithm=serialization.NoEncryption(), - ) - server_pub = public_key.public_bytes( - encoding=serialization.Encoding.PEM, - format=serialization.PublicFormat.SubjectPublicKeyInfo, - ) - - _save_keys_json(server_pub, server_pri, client_pub) - - -def _get_keys() -> dict: - """获取密钥（带内存缓存），当文件变化时自动刷新。""" - global _KEY_CACHE, _KEY_CACHE_MTIME - try: - mtime = os.path.getmtime(KEY_JSON_PATH) if os.path.exists(KEY_JSON_PATH) else None - except Exception: - mtime = None - - if _KEY_CACHE is not None and _KEY_CACHE_MTIME == mtime: - return _KEY_CACHE - - # 初次或文件已更新，重新加载 - _ensure_keys_exist() - if os.path.exists(KEY_JSON_PATH): - keys = _load_keys_json() - _KEY_CACHE = keys - try: - _KEY_CACHE_MTIME = os.path.getmtime(KEY_JSON_PATH) - except Exception: - _KEY_CACHE_MTIME = None - return keys - else: - _KEY_CACHE = {} - _KEY_CACHE_MTIME = None - return _KEY_CACHE - - -def update_client_public_key(client_public_pem: bytes): - """更新客户端公钥到 keys.json""" - global _KEY_CACHE, _KEY_CACHE_MTIME - - # 确保服务端密钥存在 - _ensure_keys_exist() - keys = _load_keys_json() - - # 更新客户端公钥 - _save_keys_json( - keys["server_public_key"], - keys["server_private_key"], - client_public_pem - ) - - # 清空缓存，强制重新加载 - _KEY_CACHE = None - _KEY_CACHE_MTIME = None \ No newline at end of file diff --git a/src/cryptio/save_image.py b/src/cryptio/save_image.py index 2067c20..25087a7 100644 --- a/src/cryptio/save_image.py +++ b/src/cryptio/save_image.py @@ -5,12 +5,10 @@ import numpy as np import torch from PIL import Image from PIL.PngImagePlugin import PngInfo -import folder_paths -from .keys import _get_keys -from cryptography.hazmat.primitives.asymmetric import padding -from cryptography.hazmat.primitives import serialization, hashes -from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes -from cryptography.hazmat.backends import default_backend + +import folder_paths # pyright: ignore[reportMissingImports] +from .utils import _key_manager +from .utils.crypto_utils import encrypt_data_hybrid class SaveImageCryptIO: @@ -168,36 +166,8 @@ def encrypt_data(data: bytes) -> bytes: """ 使用混合加密方案加密数据（使用客户端公钥加密） """ - # 1. 生成随机 AES 密钥和 IV - aes_key = os.urandom(32) # 256-bit key - iv = os.urandom(12) # 96-bit IV for GCM - # 2. 使用 AES-GCM 加密数据 - cipher = Cipher(algorithms.AES(aes_key), modes.GCM(iv), backend=default_backend()) - encryptor = cipher.encryptor() - encrypted_data = encryptor.update(data) + encryptor.finalize() - tag = encryptor.tag # type: ignore # 获取认证标签 - # 3. 使用客户端公钥加密 AES 密钥（这样只有客户端能解密） - client_public_pem = _get_keys().get("client_public_key") - if not client_public_pem: + # 使用客户端公钥加密（这样只有客户端能解密） + if not _key_manager.client_public_key: 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( - mgf=padding.MGF1(algorithm=hashes.SHA256()), - algorithm=hashes.SHA256(), - label=None, - ), - ) - # 4. 组合数据：[加密的AES密钥长度(4字节)] + [加密的AES密钥] + [IV长度(4字节)] + [IV] + [加密的数据] + [认证标签] - encrypted_aes_key_length = len(encrypted_aes_key).to_bytes(4, byteorder="big") - iv_length = len(iv).to_bytes(4, byteorder="big") - combined = ( - encrypted_aes_key_length - + encrypted_aes_key - + iv_length - + iv - + encrypted_data - + tag # GCM认证标签附加在末尾 - ) - return combined + + return encrypt_data_hybrid(data, _key_manager.client_public_key) diff --git a/src/cryptio/server_decrypt.py b/src/cryptio/server_decrypt.py index 4bbdf27..47b0dca 100644 --- a/src/cryptio/server_decrypt.py +++ b/src/cryptio/server_decrypt.py @@ -1,8 +1,6 @@ import base64 -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 +from .utils import _key_manager +from .utils.crypto_utils import decrypt_data_hybrid class ServerDecrypt: @@ -27,24 +25,14 @@ class ServerDecrypt: CATEGORY = "CryptIO" def decrypt(self, encrypted_text): - """使用私钥解密文本（私钥从缓存加载）""" + """使用私钥解密文本（支持混合加密格式）""" try: if encrypted_text.startswith("ENCRYPTED:"): encrypted_bytes = base64.b64decode(encrypted_text[10:]) - private_pem = _get_keys().get("server_private_key") - private_key = serialization.load_pem_private_key(private_pem, password=None, backend=default_backend()) - - decrypted_bytes = private_key.decrypt( - encrypted_bytes, - padding.OAEP( - mgf=padding.MGF1(algorithm=hashes.SHA256()), - algorithm=hashes.SHA256(), - label=None, - ), - ) - - decrypted_text = decrypted_bytes.decode("utf-8") + # 使用混合解密 + decrypted_data = decrypt_data_hybrid(encrypted_bytes, _key_manager.server_private_key) + decrypted_text = decrypted_data.decode("utf-8") return (decrypted_text,) else: return (encrypted_text,) diff --git a/src/cryptio/upload_image.py b/src/cryptio/upload_image.py index b2948ac..622f9a1 100644 --- a/src/cryptio/upload_image.py +++ b/src/cryptio/upload_image.py @@ -1,13 +1,10 @@ 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 +from .utils import _key_manager +from .utils.crypto_utils import decrypt_data_hybrid import folder_paths import node_helpers @@ -50,55 +47,8 @@ class UploadImageCryptIO: 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_with_tag = 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()) # type: ignore - - aes_key_bytes = private_key.decrypt( # type: ignore - 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 - - # GCM模式：Web Crypto API将16字节的认证标签附加在密文末尾 - # 需要分离出认证标签和密文 - tag_length = 16 - encrypted_data = encrypted_data_with_tag[:-tag_length] - tag = encrypted_data_with_tag[-tag_length:] - - cipher = Cipher(algorithms.AES(aes_key_bytes), modes.GCM(iv, tag), backend=default_backend()) - decryptor = cipher.decryptor() - - # 解密数据（认证标签已在cipher初始化时提供） - decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize() + # 使用服务端私钥解密（混合加密） + decrypted_data = decrypt_data_hybrid(encrypted_combined, _key_manager.server_private_key) # 从字节数据创建图片 from io import BytesIO diff --git a/src/cryptio/utils/__init__.py b/src/cryptio/utils/__init__.py new file mode 100644 index 0000000..98028eb --- /dev/null +++ b/src/cryptio/utils/__init__.py @@ -0,0 +1,30 @@ +""" +CryptIO Utils Module +提供加密、解密、密钥管理等通用功能 +""" + +from .crypto_utils import ( + encrypt_data_with_public_key, + decrypt_data_with_private_key, + encrypt_data_hybrid, + decrypt_data_hybrid, + generate_aes_key_and_iv, + encrypt_with_aes_cbc, + decrypt_with_aes_cbc, +) + +from .key_utils import KeyManager, _key_manager + +__all__ = [ + # 加密/解密 + "encrypt_data_with_public_key", + "decrypt_data_with_private_key", + "encrypt_data_hybrid", + "decrypt_data_hybrid", + "generate_aes_key_and_iv", + "encrypt_with_aes_cbc", + "decrypt_with_aes_cbc", + # 密钥管理 + "KeyManager", + "_key_manager", +] diff --git a/src/cryptio/utils/crypto_utils.py b/src/cryptio/utils/crypto_utils.py new file mode 100644 index 0000000..3782985 --- /dev/null +++ b/src/cryptio/utils/crypto_utils.py @@ -0,0 +1,216 @@ +""" +加密/解密工具模块 +提供数据加密和解密的核心功能 +""" + +import os +import secrets +from typing import Tuple +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes +from cryptography.hazmat.primitives.asymmetric import padding +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.backends import default_backend + + +def encrypt_data_with_public_key(data: bytes, public_key) -> bytes: + """ + 使用RSA公钥直接加密数据（适用于小数据，最大约190字节） + + Args: + data: 要加密的数据 + public_key: RSA公钥对象 + + Returns: + 加密后的数据 + """ + encrypted = public_key.encrypt( + data, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + return encrypted + + +def decrypt_data_with_private_key(encrypted_data: bytes, private_key) -> bytes: + """ + 使用RSA私钥直接解密数据 + + Args: + encrypted_data: 加密的数据 + private_key: RSA私钥对象 + + Returns: + 解密后的数据 + """ + decrypted = private_key.decrypt( + encrypted_data, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + return decrypted + + +def encrypt_data_hybrid(data: bytes, public_key) -> bytes: + """ + 使用混合加密方案加密数据（AES-GCM + RSA-OAEP） + 适用于任意大小的数据 + + 数据格式： + [加密的AES密钥长度(4字节)] + [加密的AES密钥] + [IV长度(4字节)] + [IV] + [加密的数据] + + Args: + data: 要加密的数据 + public_key: RSA公钥对象（用于加密AES密钥） + + Returns: + 组合的加密数据 + """ + # 1. 生成随机 AES 密钥和 IV + aes_key = os.urandom(32) # 256-bit key + iv = os.urandom(12) # 96-bit IV for GCM + + # 2. 使用 AES-GCM 加密数据 + cipher = Cipher(algorithms.AES(aes_key), modes.GCM(iv), backend=default_backend()) + encryptor = cipher.encryptor() + encrypted_data = encryptor.update(data) + encryptor.finalize() + tag = encryptor.tag # 获取认证标签 + + # 3. 使用RSA公钥加密 AES 密钥 + encrypted_aes_key = encrypt_data_with_public_key(aes_key, public_key) + + # 4. 组合数据 + encrypted_aes_key_length = len(encrypted_aes_key).to_bytes(4, byteorder="big") + iv_length = len(iv).to_bytes(4, byteorder="big") + + combined = ( + encrypted_aes_key_length + + encrypted_aes_key + + iv_length + + iv + + encrypted_data + + tag # GCM认证标签附加在末尾 + ) + + return combined + + +def decrypt_data_hybrid(encrypted_combined: bytes, private_key) -> bytes: + """ + 解密混合加密的数据（AES-GCM + RSA-OAEP） + + Args: + encrypted_combined: 组合的加密数据 + private_key: RSA私钥对象（用于解密AES密钥） + + Returns: + 解密后的原始数据 + """ + 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_with_tag = encrypted_combined[offset:] + + # 6. 使用RSA私钥解密AES密钥 + aes_key = decrypt_data_with_private_key(encrypted_aes_key, private_key) + + # 7. 使用AES密钥解密数据 + # GCM模式：认证标签附加在密文末尾 + tag_length = 16 + encrypted_data = encrypted_data_with_tag[:-tag_length] + tag = encrypted_data_with_tag[-tag_length:] + + cipher = Cipher( + algorithms.AES(aes_key), modes.GCM(iv, tag), backend=default_backend() + ) + decryptor = cipher.decryptor() + decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize() + + return decrypted_data + + +def generate_aes_key_and_iv() -> Tuple[bytes, bytes]: + """ + 生成随机的AES密钥和IV + + Returns: + (aes_key, iv) 元组 + """ + aes_key = secrets.token_bytes(32) # 256-bit key + iv = secrets.token_bytes(16) # 128-bit IV for CBC + return aes_key, iv + + +def encrypt_with_aes_cbc(data: bytes, aes_key: bytes, iv: bytes) -> bytes: + """ + 使用AES-CBC模式加密数据（带PKCS7填充） + + Args: + data: 要加密的数据 + aes_key: AES密钥 + iv: 初始化向量 + + Returns: + 加密后的数据 + """ + from cryptography.hazmat.primitives import padding as sym_padding + + # PKCS7 padding + padder = sym_padding.PKCS7(128).padder() + padded_data = padder.update(data) + padder.finalize() + + # AES-CBC 加密 + cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend()) + encryptor = cipher.encryptor() + encrypted_data = encryptor.update(padded_data) + encryptor.finalize() + + return encrypted_data + + +def decrypt_with_aes_cbc(encrypted_data: bytes, aes_key: bytes, iv: bytes) -> bytes: + """ + 使用AES-CBC模式解密数据（带PKCS7填充） + + Args: + encrypted_data: 加密的数据 + aes_key: AES密钥 + iv: 初始化向量 + + Returns: + 解密后的数据 + """ + from cryptography.hazmat.primitives import padding as sym_padding + + # AES-CBC 解密 + cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend()) + decryptor = cipher.decryptor() + padded_data = decryptor.update(encrypted_data) + decryptor.finalize() + + # 移除 PKCS7 padding + unpadder = sym_padding.PKCS7(128).unpadder() + data = unpadder.update(padded_data) + unpadder.finalize() + + return data diff --git a/src/cryptio/utils/key_utils.py b/src/cryptio/utils/key_utils.py new file mode 100644 index 0000000..0b6bbfe --- /dev/null +++ b/src/cryptio/utils/key_utils.py @@ -0,0 +1,183 @@ +""" +密钥工具模块 +提供密钥导入、导出、加载、存储管理等功能 +""" + +import os +import json +import base64 +from typing import Optional, cast +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey, RSAPrivateKey +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.backends import default_backend + + +class KeyManager: + """ + 密钥管理器单例类 + 负责密钥的加载、缓存、更新和持久化 + """ + + _instance: Optional["KeyManager"] = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self): + """初始化密钥管理器（仅首次调用时执行）""" + if self._initialized: + return + + # 密钥存储路径 + base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + self._key_dir = os.path.join(base_dir, "keys") + self._key_file = os.path.join(self._key_dir, "keys.json") + os.makedirs(self._key_dir, exist_ok=True) + + # 缓存状态 + self._cache_mtime: Optional[float] = None + self._server_public_key: Optional[RSAPublicKey] = None + self._server_private_key: Optional[RSAPrivateKey] = None + self._client_public_key: Optional[RSAPublicKey] = None + self._server_public_key_pem: Optional[bytes] = None + self._server_private_key_pem: Optional[bytes] = None + self._client_public_key_pem: Optional[bytes] = None + + # 初始化加载 + self._load() + self._initialized = True + + def _load(self): + """加载密钥到缓存""" + # 检查文件修改时间 + try: + current_mtime = os.path.getmtime(self._key_file) if os.path.exists(self._key_file) else None + except Exception: + current_mtime = None + + # 缓存命中 + if self._cache_mtime == current_mtime and self._server_private_key is not None: + return + + # 确保密钥文件存在 + if not os.path.exists(self._key_file): + self._generate_server_keys() + + # 从文件加载 + with open(self._key_file, "r", encoding="utf-8") as f: + data = json.load(f) + + # 解码 PEM 数据 + self._server_public_key_pem = base64.b64decode(data["server_public_key"]) + self._server_private_key_pem = base64.b64decode(data["server_private_key"]) + + if "client_public_key" in data: + self._client_public_key_pem = base64.b64decode(data["client_public_key"]) + + # 加载密钥对象 + self._server_public_key = cast( + RSAPublicKey, serialization.load_pem_public_key(self._server_public_key_pem, backend=default_backend()) + ) + self._server_private_key = cast( + RSAPrivateKey, serialization.load_pem_private_key(self._server_private_key_pem, password=None, backend=default_backend()) + ) + + if self._client_public_key_pem: + self._client_public_key = cast( + RSAPublicKey, serialization.load_pem_public_key(self._client_public_key_pem, backend=default_backend()) + ) + + # 更新缓存时间 + self._cache_mtime = current_mtime + + def _generate_server_keys(self): + """生成服务端密钥对""" + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048, backend=default_backend()) + public_key = private_key.public_key() + + self._server_private_key_pem = private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + self._server_public_key_pem = public_key.public_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PublicFormat.SubjectPublicKeyInfo, + ) + + self._save() + + def _save(self): + """保存密钥到文件""" + data: dict[str, str] = {} + if self._server_public_key_pem and self._server_private_key_pem: + data["server_public_key"] = base64.b64encode(self._server_public_key_pem).decode("ascii") + data["server_private_key"] = base64.b64encode(self._server_private_key_pem).decode("ascii") + + if self._client_public_key_pem: + data["client_public_key"] = base64.b64encode(self._client_public_key_pem).decode("ascii") + + with open(self._key_file, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=2) + + def reload(self): + """重新加载密钥（当外部修改文件时调用）""" + self._cache_mtime = None + self._load() + + def update_client_key(self, client_public_pem: bytes): + """ + 更新客户端公钥 + + Args: + client_public_pem: 客户端公钥（PEM格式） + """ + self._client_public_key_pem = client_public_pem + self._save() + self.reload() + + # ===== 属性访问器 ===== + + @property + def server_public_key(self) -> Optional[RSAPublicKey]: + """服务端公钥对象""" + self._load() + return self._server_public_key + + @property + def server_private_key(self) -> Optional[RSAPrivateKey]: + """服务端私钥对象""" + self._load() + return self._server_private_key + + @property + def client_public_key(self) -> Optional[RSAPublicKey]: + """客户端公钥对象""" + self._load() + return self._client_public_key + + @property + def server_public_key_pem(self) -> Optional[bytes]: + """服务端公钥（PEM格式）""" + self._load() + return self._server_public_key_pem + + @property + def server_private_key_pem(self) -> Optional[bytes]: + """服务端私钥（PEM格式）""" + self._load() + return self._server_private_key_pem + + @property + def client_public_key_pem(self) -> Optional[bytes]: + """客户端公钥（PEM格式）""" + self._load() + return self._client_public_key_pem + + +# ===== 全局单例实例 ===== +_key_manager = KeyManager() diff --git a/ts/utils/uploadUtils.ts b/ts/utils/uploadUtils.ts index e278289..43a71b7 100644 --- a/ts/utils/uploadUtils.ts +++ b/ts/utils/uploadUtils.ts @@ -15,7 +15,7 @@ import { loadEncryptedImageFromFilename } from "./imageLoader.js"; 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 encryptFileWithServerKey(file);