refactor(cryptio): 重构密钥管理模块并优化加密功能
将密钥管理逻辑从 keys.py 迁移到 utils 模块，实现 KeyManager 单例类 新增 crypto_utils 模块集中管理加密/解密功能 优化混合加密方案实现，提高代码复用性和可维护性
This commit is contained in:
Suzu008
+27
-42
@@ -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)
|
||||
|
||||
|
||||
@@ -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,)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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,)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -15,7 +15,7 @@ import { loadEncryptedImageFromFilename } from "./imageLoader.js";
|
||||
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 encryptFileWithServerKey(file);
|
||||
|
||||
Reference in New Issue
Block a user