177 lines
6.4 KiB
Python
177 lines
6.4 KiB
Python
import os
|
||
import string
|
||
import secrets
|
||
import json
|
||
import base64
|
||
from aiohttp import web
|
||
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
|
||
|
||
|
||
@PromptServer.instance.routes.post("/cryptio/exchange_keys")
|
||
async def exchange_keys(request):
|
||
"""
|
||
密钥交换端点：
|
||
1. 接收客户端公钥
|
||
2. 保存客户端公钥到服务器
|
||
3. 使用客户端公钥加密服务端密钥对
|
||
4. 返回加密后的服务端密钥
|
||
"""
|
||
try:
|
||
data = await request.json()
|
||
client_public_key_pem = data.get("client_public_key")
|
||
|
||
if not client_public_key_pem:
|
||
return web.json_response({"error": "client_public_key is required"}, status=400)
|
||
|
||
# 保存客户端公钥
|
||
client_public_key_bytes = client_public_key_pem.encode("utf-8")
|
||
_key_manager.update_client_key(client_public_key_bytes)
|
||
|
||
# 获取服务端密钥 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_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()
|
||
)
|
||
|
||
# 准备要加密的数据（服务端公钥和私钥）
|
||
keys_data = json.dumps(
|
||
{
|
||
"server_public_key": server_public_key_pem.decode("utf-8"),
|
||
"server_private_key": server_private_key_pem.decode("utf-8"),
|
||
}
|
||
)
|
||
|
||
# 由于RSA加密有大小限制，使用混合加密方案（AES-CBC + RSA-OAEP）
|
||
# 1. 生成随机AES密钥和IV
|
||
aes_key, iv = generate_aes_key_and_iv()
|
||
|
||
# 2. 使用AES加密数据
|
||
encrypted_data = encrypt_with_aes_cbc(keys_data.encode("utf-8"), aes_key, iv)
|
||
|
||
# 3. 使用RSA加密AES密钥
|
||
encrypted_aes_key = encrypt_data_with_public_key(aes_key, client_public_key)
|
||
|
||
# 4. 返回加密后的数据
|
||
return web.json_response(
|
||
{
|
||
"encrypted_aes_key": base64.b64encode(encrypted_aes_key).decode("ascii"),
|
||
"iv": base64.b64encode(iv).decode("ascii"),
|
||
"encrypted_data": base64.b64encode(encrypted_data).decode("ascii"),
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
print(f"Error in key exchange: {e}")
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
|
||
@PromptServer.instance.routes.get("/cryptio/public_key")
|
||
async def get_public_key(request):
|
||
"""兼容旧的公钥获取接口（不加密传输）"""
|
||
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)
|
||
|
||
|
||
@PromptServer.instance.routes.post("/cryptio/upload_encrypted")
|
||
async def upload_encrypted_image(request):
|
||
"""接收加密的图片数据并保存到文件"""
|
||
try:
|
||
# 读取 multipart 表单数据
|
||
reader = await request.multipart()
|
||
|
||
encrypted_data = None
|
||
original_filename = None
|
||
|
||
async for field in reader:
|
||
if field.name == "image":
|
||
# 读取加密的文件数据
|
||
encrypted_data = await field.read()
|
||
original_filename = field.filename
|
||
|
||
if not encrypted_data:
|
||
return web.json_response({"error": "No image data received"}, status=400)
|
||
|
||
# 生成唯一的加密文件名，保留原始文件名并追加随机字符串
|
||
input_dir = folder_paths.get_input_directory()
|
||
|
||
safe_original = os.path.basename(original_filename) if original_filename else "cryptio"
|
||
if safe_original.lower().endswith(".encrypted"):
|
||
safe_original = safe_original[:-len(".encrypted")]
|
||
|
||
base_name, file_ext = os.path.splitext(safe_original)
|
||
|
||
base_name = base_name.strip().replace(" ", "_") or "cryptio"
|
||
file_ext = file_ext or ".enc"
|
||
|
||
charset = string.ascii_lowercase + string.digits
|
||
|
||
def build_filename() -> str:
|
||
random_suffix = ''.join(secrets.choice(charset) for _ in range(6))
|
||
return f"{base_name}_{random_suffix}{file_ext}.encrypted"
|
||
|
||
unique_filename = build_filename()
|
||
while os.path.exists(os.path.join(input_dir, unique_filename)):
|
||
unique_filename = build_filename()
|
||
|
||
# 保存到 input 目录
|
||
file_path = os.path.join(input_dir, unique_filename)
|
||
|
||
# 写入加密数据
|
||
with open(file_path, "wb") as f:
|
||
f.write(encrypted_data)
|
||
|
||
return web.json_response({"name": unique_filename, "subfolder": "", "type": "input"})
|
||
|
||
except Exception as e:
|
||
print(f"Error uploading encrypted image: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
|
||
@PromptServer.instance.routes.get("/cryptio/view_encrypted")
|
||
async def view_encrypted_image(request):
|
||
"""获取加密图片用于前端预览（返回二进制加密数据，由前端解密）"""
|
||
try:
|
||
filename = request.query.get("filename")
|
||
if not filename:
|
||
return web.json_response({"error": "filename parameter is required"}, status=400)
|
||
|
||
# 获取文件路径
|
||
image_path = folder_paths.get_annotated_filepath(filename)
|
||
|
||
if not os.path.exists(image_path):
|
||
return web.json_response({"error": "File not found"}, status=404)
|
||
|
||
# 读取加密文件
|
||
with open(image_path, "rb") as f:
|
||
encrypted_data = f.read()
|
||
|
||
# 返回二进制数据，使用 Content-Disposition header 传递文件名
|
||
return web.Response(
|
||
body=encrypted_data, content_type="application/octet-stream", headers={"Content-Disposition": f'filename="{filename}"'}
|
||
)
|
||
|
||
except Exception as e:
|
||
print(f"Error viewing encrypted image: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|