import unittest
import sys
import os

# 确保可以导入上级目录的 core 包
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from core.sp_trust_sharding import SentinelKeyEngine
from core.sp_vault_aes import SentinelVault
from core.sp_gateway_rsa import SentinelSystemProvider

class TestSentinelCore(unittest.TestCase):

    def setUp(self):
        self.key_engine = SentinelKeyEngine()
        self.vault = SentinelVault()
        self.sys_provider = SentinelSystemProvider()

    def test_01_sharding_recovery(self):
        """测试 SSS 密钥分片与恢复 (3选2)"""
        print("\n[Test] Running Sharding & Recovery...")
        
        # 1. 生成
        words, entropy = self.key_engine.generate_vault_keys()
        self.assertEqual(len(words.split()), 12, "助记词应为12个单词")
        
        # 2. 分片
        shares = self.key_engine.split_to_shares(entropy)
        self.assertEqual(len(shares), 3, "应生成3个分片")
        
        # 3. 验证所有组合 (3选2)
        # 组合 A+B
        rec_ab = self.key_engine.recover_from_shares(shares[0], shares[1])
        self.assertEqual(rec_ab, words, "分片 A+B 恢复失败")
        
        # 组合 B+C
        rec_bc = self.key_engine.recover_from_shares(shares[1], shares[2])
        self.assertEqual(rec_bc, words, "分片 B+C 恢复失败")
        
        # 组合 A+C
        rec_ac = self.key_engine.recover_from_shares(shares[0], shares[2])
        self.assertEqual(rec_ac, words, "分片 A+C 恢复失败")

    def test_02_vault_encryption(self):
        """测试 AES 金库加密与解密"""
        print("[Test] Running AES Vault...")
        words = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"
        key = self.vault.derive_key(words)
        data = "Sentinel Top Secret Data"
        
        # 加密
        encrypted = self.vault.encrypt_data(key, data)
        self.assertNotEqual(encrypted, data.encode(), "密文不应与明文相同")
        
        # 解密
        decrypted = self.vault.decrypt_data(key, encrypted)
        self.assertEqual(decrypted, data, "解密后数据应与原始数据一致")
        
        # 错误密钥测试
        wrong_key = self.vault.derive_key("zoo zoo zoo zoo zoo zoo zoo zoo zoo zoo zoo zoo")
        result = self.vault.decrypt_data(wrong_key, encrypted)
        self.assertIn("解密失败", result, "使用错误密钥应返回失败信息")

    def test_03_gateway_rsa(self):
        """测试 RSA 系统网关加壳流程"""
        print("[Test] Running RSA Gateway...")
        priv_pem, pub_pem = self.sys_provider.generate_system_keys()
        payload = b"User Encrypted Blob Data"
        
        # 加密
        cipher = self.sys_provider.encrypt_with_system_public(pub_pem, payload)
        self.assertNotEqual(cipher, payload)
        
        # 解密
        restored = self.sys_provider.decrypt_with_system_private(priv_pem, cipher)
        self.assertEqual(restored, payload, "RSA 解密还原失败")

if __name__ == '__main__':
    unittest.main()