all in one key file
This commit is contained in:
Suzu008
+116
-46
@@ -8,21 +8,114 @@ from cryptography.hazmat.backends import default_backend
|
||||
from server import PromptServer
|
||||
from aiohttp import web
|
||||
|
||||
# 密钥存储路径
|
||||
# 密钥存储路径（整合为单一 JSON 文件，使用 Base64 编码）
|
||||
KEY_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "keys")
|
||||
PUBLIC_KEY_PATH = os.path.join(KEY_DIR, "server_public_key.pem")
|
||||
PRIVATE_KEY_PATH = os.path.join(KEY_DIR, "server_private_key.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") # 兼容旧存储（迁移用）
|
||||
KEY_JSON_PATH = os.path.join(KEY_DIR, "keys.json")
|
||||
os.makedirs(KEY_DIR, exist_ok=True)
|
||||
|
||||
# 读取密钥缓存，避免频繁读取文件
|
||||
_KEY_CACHE = None
|
||||
_KEY_CACHE_MTIME = None
|
||||
|
||||
def _read_file_bytes(path: str):
|
||||
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():
|
||||
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():
|
||||
"""获取密钥（带内存缓存），当文件变化时自动刷新。"""
|
||||
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
|
||||
|
||||
|
||||
@PromptServer.instance.routes.get("/cryptio/public_key")
|
||||
async def get_public_key(request):
|
||||
if os.path.exists(PUBLIC_KEY_PATH):
|
||||
with open(PUBLIC_KEY_PATH, "rb") as f:
|
||||
public_key = f.read().decode("utf-8")
|
||||
return web.json_response({"public_key": public_key})
|
||||
keys = _get_keys()
|
||||
pub = keys.get("server_public_key")
|
||||
if pub:
|
||||
return web.json_response({"public_key": pub.decode("utf-8")})
|
||||
else:
|
||||
return web.json_response({"error": "Public key not found"}, status=404)
|
||||
|
||||
@@ -33,7 +126,7 @@ class KeyGenerator:
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.generate_keys_if_not_exist()
|
||||
_ensure_keys_exist()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -45,34 +138,12 @@ class KeyGenerator:
|
||||
CATEGORY = "CryptIO"
|
||||
|
||||
def generate_keys_if_not_exist(self):
|
||||
"""如果密钥不存在，则生成新的密钥对"""
|
||||
if not (os.path.exists(PUBLIC_KEY_PATH) and os.path.exists(PRIVATE_KEY_PATH)):
|
||||
# 生成RSA密钥对
|
||||
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048, backend=default_backend())
|
||||
|
||||
# 获取公钥
|
||||
public_key = private_key.public_key()
|
||||
|
||||
# 将私钥保存到文件
|
||||
with open(PRIVATE_KEY_PATH, "wb") as f:
|
||||
f.write(
|
||||
private_key.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.NoEncryption(),
|
||||
)
|
||||
)
|
||||
|
||||
# 将公钥保存到文件
|
||||
with open(PUBLIC_KEY_PATH, "wb") as f:
|
||||
f.write(
|
||||
public_key.public_bytes(encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo)
|
||||
)
|
||||
"""兼容旧接口：确保密钥存在（现使用 JSON 存储）。"""
|
||||
_ensure_keys_exist()
|
||||
|
||||
def generate(self):
|
||||
"""返回公钥"""
|
||||
with open(PUBLIC_KEY_PATH, "rb") as f:
|
||||
public_key = f.read().decode("utf-8")
|
||||
"""返回公钥（使用缓存避免频繁 IO）"""
|
||||
public_key = _get_keys()["server_public_key"].decode("utf-8")
|
||||
return (public_key,)
|
||||
|
||||
|
||||
@@ -127,27 +198,26 @@ class ServerDecrypt:
|
||||
CATEGORY = "CryptIO"
|
||||
|
||||
def decrypt(self, encrypted_text):
|
||||
"""使用私钥解密文本"""
|
||||
"""使用私钥解密文本（私钥从缓存加载）"""
|
||||
try:
|
||||
# 如果前端已经加密，则解密
|
||||
if encrypted_text.startswith("ENCRYPTED:"):
|
||||
# 移除前缀并解码Base64
|
||||
encrypted_bytes = base64.b64decode(encrypted_text[10:])
|
||||
|
||||
# 加载私钥
|
||||
with open(PRIVATE_KEY_PATH, "rb") as key_file:
|
||||
private_key = serialization.load_pem_private_key(key_file.read(), password=None, backend=default_backend())
|
||||
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)
|
||||
encrypted_bytes,
|
||||
padding.OAEP(
|
||||
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
||||
algorithm=hashes.SHA256(),
|
||||
label=None,
|
||||
),
|
||||
)
|
||||
|
||||
# 解码为字符串
|
||||
decrypted_text = decrypted_bytes.decode("utf-8")
|
||||
return (decrypted_text,)
|
||||
else:
|
||||
# 如果未加密，直接返回原始文本
|
||||
return (encrypted_text,)
|
||||
except Exception as e:
|
||||
print(f"解密错误: {e}")
|
||||
|
||||
Reference in New Issue
Block a user