book_rag_server/test/test_asymmetric_encryption.py

from utils.asymmetric_encryption import AsymmetricEncryption
import os


def test_asymmetric_encryption():
    """测试非对称加密功能"""
    print("开始测试非对称加密功能...")
    
    # 测试1：生成密钥对
    print("\n1. 生成密钥对测试...")
    private_pem, public_pem = AsymmetricEncryption.generate_key_pair()
    print(f"私钥长度: {len(private_pem)} 字节")
    print(f"公钥长度: {len(public_pem)} 字节")
    print(f"私钥前100字节: {private_pem[:100]!r}")
    print(f"公钥前100字节: {public_pem[:100]!r}")
    
    # 测试2：加密解密测试
    print("\n2. 加密解密测试...")
    test_message = "这是一个测试消息，用于测试非对称加密功能！"
    print(f"原始消息: {test_message}")
    
    # 加密
    encrypted = AsymmetricEncryption.encrypt(test_message, public_pem)
    print(f"加密后: {encrypted}")
    
    # 解密
    decrypted = AsymmetricEncryption.decrypt(encrypted, private_pem)
    print(f"解密后: {decrypted}")
    
    # 验证
    assert decrypted == test_message, "解密失败，结果与原始消息不符！"
    print("✓ 加密解密测试通过！")
    
    # 测试3：密钥文件保存和加载测试
    print("\n3. 密钥文件保存和加载测试...")
    private_key_path = "private_key.pem"
    public_key_path = "public_key.pem"
    
    # 保存密钥
    AsymmetricEncryption.save_key_to_file(private_pem, private_key_path)
    AsymmetricEncryption.save_key_to_file(public_pem, public_key_path)
    print(f"✓ 密钥已保存到文件: {private_key_path}, {public_key_path}")
    
    # # 加载密钥
    # loaded_private_pem = AsymmetricEncryption.load_key_from_file(private_key_path)
    # loaded_public_pem = AsymmetricEncryption.load_key_from_file(public_key_path)
    # print("✓ 密钥已从文件加载")
    
    # # 验证加载的密钥是否正确
    # encrypted2 = AsymmetricEncryption.encrypt(test_message, loaded_public_pem)
    # decrypted2 = AsymmetricEncryption.decrypt(encrypted2, loaded_private_pem)
    # assert decrypted2 == test_message, "使用加载的密钥解密失败！"
    # print("✓ 使用加载的密钥加密解密测试通过！")
    
    # # 测试4：生成密钥对并保存测试
    # print("\n4. 生成密钥对并保存测试...")
    # private_key_path2 = "test_private_key2.pem"
    # public_key_path2 = "test_public_key2.pem"
    
    # AsymmetricEncryption.generate_key_pair_and_save(private_key_path2, public_key_path2)
    # print(f"✓ 密钥对已生成并保存到文件: {private_key_path2}, {public_key_path2}")
    
    # # 验证生成并保存的密钥
    # loaded_private_pem2 = AsymmetricEncryption.load_key_from_file(private_key_path2)
    # loaded_public_pem2 = AsymmetricEncryption.load_key_from_file(public_key_path2)
    # encrypted3 = AsymmetricEncryption.encrypt(test_message, loaded_public_pem2)
    # decrypted3 = AsymmetricEncryption.decrypt(encrypted3, loaded_private_pem2)
    # assert decrypted3 == test_message, "使用生成并保存的密钥解密失败！"
    # print("✓ 使用生成并保存的密钥加密解密测试通过！")
    
    # # 清理临时文件
    # print("\n5. 清理临时文件...")
    # for file_path in [private_key_path, public_key_path, private_key_path2, public_key_path2]:
    #     if os.path.exists(file_path):
    #         os.remove(file_path)
    #         print(f"✓ 删除临时文件: {file_path}")
    
    # print("\n所有测试通过！非对称加密功能正常工作。")


if __name__ == "__main__":
    test_asymmetric_encryption()