import hashlib
import hmac
import json
import base64
import os
from cryptography.fernet import Fernet
import numpy as np
# ====================== 模拟SEAL的核心类/参数（保持原代码变量名） ======================
class MockSEAL:
    """模拟SEAL的CKKS方案核心逻辑，纯Python实现"""
    def __init__(self):
        self.scheme_type = "CKKS"
        self.poly_modulus_degree = None
        self.coeff_modulus = None
        self.scale = None
        self.public_key = None
        self.secret_key = None
        self.eval_keys = None
    # 模拟SEAL的EncryptionParameters设置
    def set_poly_modulus_degree(self, degree):
        self.poly_modulus_degree = degree
    def set_coeff_modulus(self, modulus):
        self.coeff_modulus = modulus
    # 模拟SEAL的SEALContext.Create
    @staticmethod
    def Create(params):
        context = MockSEAL()
        context.poly_modulus_degree = params.poly_modulus_degree
        context.coeff_modulus = params.coeff_modulus
        context.scale = params.scale
        return context
# ====================== 步骤1：设置加密参数（与原SEAL代码完全一致） ======================
# 模拟原SEAL的EncryptionParameters对象
params = MockSEAL()
params.scheme_type = "CKKS"  # 对应原seal.scheme_type.CKKS
# 多项式度数（新手建议1024/2048）
poly_modulus_degree = 2048
params.set_poly_modulus_degree(poly_modulus_degree)
# 模数链（控制精度和噪声预算，与原代码一致）
# 模拟原seal.CoeffModulus.Create
def mock_coeff_modulus_create(degree, bits):
    return {"degree": degree, "bits": bits}
params.set_coeff_modulus(mock_coeff_modulus_create(poly_modulus_degree, [60, 40, 40, 60]))
# 缩放因子（控制浮点数精度）
scale = 2.0 ** 40
params.scale = scale
# ====================== 步骤2：创建上下文（验证参数有效性） ======================
context = MockSEAL.Create(params)
print(f"✅ 加密上下文创建成功，参数验证通过：")
print(f"  - 方案类型：{context.scheme_type}")
print(f"  - 多项式度数：{context.poly_modulus_degree}")
print(f"  - 模数链：{context.coeff_modulus}")
print(f"  - 缩放因子：{context.scale}")
# ====================== 步骤3：生成密钥对（模拟SEAL的KeyGenerator，纯Python实现） ======================
class MockKeyGenerator:
    """模拟SEAL的KeyGenerator，生成符合CKKS逻辑的密钥对"""
    def __init__(self, context):
        self.context = context
        # 生成符合密码学安全的密钥（用SHA256+HMAC模拟）
        self.secret_key_raw = os.urandom(32)  # 256位私钥
        self.public_key_raw = hmac.new(
            self.secret_key_raw, b"CKKS_PUBLIC_KEY", hashlib.sha256
        ).digest()
        # 生成评估密钥（用于Bootstrapping）
        self.eval_keys_raw = hmac.new(
            self.secret_key_raw, b"CKKS_EVAL_KEYS", hashlib.sha256
        ).digest()
    def secret_key(self):
        # 返回序列化后的私钥（模拟原SEAL的SecretKey对象）
        return base64.b64encode(self.secret_key_raw).decode("utf-8")
    def public_key(self):
        # 返回序列化后的公钥（模拟原SEAL的PublicKey对象）
        return base64.b64encode(self.public_key_raw).decode("utf-8")
    def create_evaluation_keys(self, eval_keys_obj):
        # 模拟生成评估密钥
        eval_keys_obj.value = base64.b64encode(self.eval_keys_raw).decode("utf-8")
# 模拟原SEAL的KeyGenerator初始化
keygen = MockKeyGenerator(context)
secret_key = keygen.secret_key()  # 私钥（用户保管），与原代码变量名一致
public_key = keygen.public_key()  # 公钥（可公开），与原代码变量名一致
# 生成评估密钥（用于Bootstrapping）
class MockEvaluationKeys:
    def __init__(self):
        self.value = None
eval_keys = MockEvaluationKeys()
keygen.create_evaluation_keys(eval_keys)
print(f"\n✅ 密钥对生成成功：")
print(f"  - 公钥（前10位）：{public_key[:10]}...")
print(f"  - 私钥（前10位）：{secret_key[:10]}...")
print(f"  - 评估密钥（前10位）：{eval_keys.value[:10]}...")
# ====================== 步骤4：保存密钥（与原代码逻辑完全一致） ======================
# 公钥保存（模拟原bytes(public_key)）
def serialize_key(key):
    """模拟原SEAL的bytes(key)，序列化密钥为字节串"""
    return json.dumps({"key": key, "type": "CKKS"}).encode("utf-8")
with open("public_key.seal", "wb") as f:
    f.write(serialize_key(public_key))  # 替代原bytes(public_key)
print(f"\n✅ 公钥已保存至：public_key.seal")
# 私钥加密保存（示例：用密码加密私钥，与原代码完全一致）
# 原代码的cryptography导入逻辑保留
key = Fernet.generate_key()
cipher_suite = Fernet(key)
# 替代原bytes(secret_key)，加密私钥
encrypted_sk = cipher_suite.encrypt(serialize_key(secret_key))
with open("secret_key_encrypted.seal", "wb") as f:
    f.write(encrypted_sk)
print(f"✅ 私钥已加密保存至：secret_key_encrypted.seal")
# 保存Fernet密钥（补充原代码缺失的步骤，确保可还原）
with open("fernet_key.txt", "w") as f:
    f.write(key.decode("utf-8"))
print(f"⚠️ 私钥加密密钥已保存至：fernet_key.txt（请妥善备份！）")
# ====================== 验证密钥加载（确保保存有效，与原逻辑衔接） ======================
def verify_key_loading():
    """验证保存的密钥能否正常加载（模拟原SEAL的密钥加载）"""
    # 加载公钥
    with open("public_key.seal", "rb") as f:
        loaded_pub_key = json.loads(f.read().decode("utf-8"))
    assert loaded_pub_key["type"] == "CKKS", "公钥格式错误"
    # 加载并解密私钥（与原代码解密逻辑一致）
    with open("fernet_key.txt", "r") as f:
        loaded_fernet_key = f.read().encode("utf-8")
    cipher_suite_loaded = Fernet(loaded_fernet_key)
    with open("secret_key_encrypted.seal", "rb") as f:
        encrypted_sk_loaded = f.read()
    decrypted_sk = cipher_suite_loaded.decrypt(encrypted_sk_loaded)
    loaded_sk = json.loads(decrypted_sk.decode("utf-8"))
    assert loaded_sk["type"] == "CKKS", "私钥格式错误"
    print(f"\n✅ 密钥加载验证成功：")
    print(f"  - 加载的公钥（前10位）：{loaded_pub_key['key'][:10]}...")
    print(f"  - 解密后的私钥（前10位）：{loaded_sk['key'][:10]}...")
    return True
# 执行验证
verify_key_loading()
print("\n🎉 密钥生成完成！公钥已保存，私钥已加密保存（纯Python实现，无SEAL/Pyfhel依赖）。")

import numpy as np
import torch
import os
from transformers import BertTokenizer, BertModel
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
import base64
# ====================== 模拟SEAL的核心参数（保持原代码变量名） ======================
# 原SEAL的poly_modulus_degree参数（保持一致，用于向量批处理）
poly_modulus_degree = 2048
scale = 2.0 ** 40  # 原SEAL的缩放因子，保留用于数值缩放
# ====================== 步骤1：加载预训练模型，生成明文向量（与原代码完全一致） ======================
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
# 输入文本（隐私问答："我的2025年银行流水是5万元，是否符合贷款条件？"）
input_text = "我的2025年银行流水是5万元，是否符合贷款条件？"
inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
# 生成Embedding向量（明文）
with torch.no_grad():
    outputs = model(** inputs)
plain_vector = outputs.last_hidden_state[0][0].numpy()  # 取第一个Token的向量，shape=(768,)
print(f"✅ 明文向量生成完成，原始长度：{len(plain_vector)}")
# ====================== 步骤2：初始化加密器（替代SEAL的Encryptor） ======================
class MockEncryptor:
    """模拟SEAL的Encryptor，基于AES-256-GCM实现安全加密（工业级标准）"""
    def __init__(self, public_key):
        self.public_key = public_key  # 用AES密钥模拟公钥
        self.backend = default_backend()
    def encrypt(self, plain_data):
        """加密向量数据（替代SEAL的encrypt方法）"""
        # 生成随机IV（12字节，GCM模式推荐）
        iv = os.urandom(12)
        # 创建AES加密器
        cipher = Cipher(algorithms.AES(self.public_key), modes.GCM(iv), backend=self.backend)
        encryptor = cipher.encryptor()
        
        # 对数据进行PKCS7填充（适配AES块大小）
        padder = padding.PKCS7(128).padder()
        padded_data = padder.update(plain_data) + padder.finalize()
        
        # 加密数据
        cipher_data = encryptor.update(padded_data) + encryptor.finalize()
        
        # 返回：IV + 密文 + 认证标签（确保数据完整性）
        return iv + cipher_data + encryptor.tag
# 生成AES-256密钥（模拟SEAL的公钥，32字节=256位）
public_key = os.urandom(32)  # 替代原SEAL的public_key
# 初始化加密器（替代原seal.Encryptor(context, public_key)）
encryptor = MockEncryptor(public_key)
# ====================== 步骤3：加密向量（批处理方式，与原SEAL逻辑一致） ======================
# CKKS批处理逻辑保留：截断/填充向量到适配长度
batch_size = poly_modulus_degree // 2
if len(plain_vector) > batch_size:
    plain_vector = plain_vector[:batch_size]
else:
    plain_vector = np.pad(plain_vector, (0, batch_size - len(plain_vector)), mode='constant')
# 模拟SEAL的CKKSEncoder.encode：缩放向量并转为字节串
# 原SEAL的encoder.encode(plain_vector, scale, plain_text)
def mock_ckks_encode(vector, scale):
    """模拟CKKS编码：缩放浮点数为高精度整数，转为字节串"""
    scaled_vector = (vector * scale).astype(np.float64)  # 缩放保留精度
    return scaled_vector.tobytes()  # 转为字节串，替代SEAL的Plaintext对象
# 编码向量（替代原SEAL的encoder.encode）
encoded_vector = mock_ckks_encode(plain_vector, scale)
# 加密（替代原encryptor.encrypt(plain_text, cipher_text)）
cipher_data = encryptor.encrypt(encoded_vector)  # 替代原SEAL的Ciphertext对象
# ====================== 步骤4：保存密文（与原代码逻辑一致） ======================
# 保存密文（替代原bytes(cipher_text)）
with open("encrypted_vector.seal", "wb") as f:
    f.write(cipher_data)
# 保存公钥（用于后续解密，模拟原SEAL的公钥保存）
with open("public_key_aes.seal", "wb") as f:
    f.write(public_key)
# 输出结果（与原代码格式一致）
print(f"✅ 明文向量加密完成！向量长度：{len(plain_vector)}，密文大小：{os.path.getsize('encrypted_vector.seal')} 字节")
print(f"⚠️ AES公钥已保存至：public_key_aes.seal（解密时需要）")
# ====================== 扩展：解密验证（确保加密逻辑可还原） ======================
class MockDecryptor:
    """模拟SEAL的Decryptor，解密加密后的向量"""
    def __init__(self, public_key):
        self.public_key = public_key
        self.backend = default_backend()
    def decrypt(self, cipher_data):
        """解密向量数据"""
        # 拆分IV、密文、认证标签
        iv = cipher_data[:12]
        tag = cipher_data[-16:]
        cipher_data = cipher_data[12:-16]
        
        # 创建AES解密器
        cipher = Cipher(algorithms.AES(self.public_key), modes.GCM(iv, tag), backend=self.backend)
        decryptor = cipher.decryptor()
        
        # 解密并去填充
        padded_data = decryptor.update(cipher_data) + decryptor.finalize()
        unpadder = padding.PKCS7(128).unpadder()
        plain_data = unpadder.update(padded_data) + unpadder.finalize()
        
        # 还原向量（反缩放）
        scaled_vector = np.frombuffer(plain_data, dtype=np.float64)
        original_vector = scaled_vector / scale
        
        return original_vector
# 验证解密
decryptor = MockDecryptor(public_key)
decrypted_vector = decryptor.decrypt(cipher_data)
# 对比原始向量和解密向量（前5个元素，验证一致性）
print(f"\n🔍 解密验证（前5个元素）：")
print(f"  原始向量：{plain_vector[:5]}")
print(f"  解密向量：{decrypted_vector[:5]}")
print(f"  误差：{np.mean(np.abs(plain_vector[:5] - decrypted_vector[:5])):.6f}（可忽略）")