import os
from mnemonic import Mnemonic
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Util.Padding import pad, unpad

class SentinelVault:
    def __init__(self, salt=b'Sentinel_Salt_2026'): # 固定的盐值，用于增加派生强度
        self.mnemo = Mnemonic("english")
        self.salt = salt

    def derive_key(self, mnemonic_phrase):
        """
        使用 PBKDF2 将助记词转换为 AES-256 密钥 (32 bytes)
        """
        # 种子生成遵循 BIP-39 逻辑
        seed = self.mnemo.to_seed(mnemonic_phrase, passphrase="")
        # 派生出一个 32 字节的强密钥
        key = PBKDF2(seed, self.salt, dkLen=32, count=100000)
        return key

    def encrypt_data(self, key, plaintext):
        """
        使用 AES-256 GCM 模式进行加密 (具备完整性校验)
        """
        cipher = AES.new(key, AES.MODE_GCM)
        nonce = cipher.nonce
        ciphertext, tag = cipher.encrypt_and_digest(plaintext.encode('utf-8'))
        # 返回：随机数 + 校验位 + 密文
        return nonce + tag + ciphertext

    def decrypt_data(self, key, encrypted_blob):
        """
        AES-256 GCM 解密
        """
        nonce = encrypted_blob[:16]
        tag = encrypted_blob[16:32]
        ciphertext = encrypted_blob[32:]
        
        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
        try:
            plaintext = cipher.decrypt_and_verify(ciphertext, tag)
            return plaintext.decode('utf-8')
        except ValueError:
            return "【解密失败】：密钥错误或数据被篡改"

if __name__ == "__main__":
    # --- 模拟 Sentinel 协议完整业务流 ---

    # 1. 假设这是通过之前 SSS 算法恢复出来的 12 词
    recovered_mnemonic = "apple banana cherry dog elephant fish goat horse ice jacket kangaroo lion"
    try:
        with open("words.txt", "r") as f:
            recovered_mnemonic = f.read().strip()
    except FileNotFoundError:
        print("words.txt 文件未找到，使用默认助记词进行演示。")

    print(f"Demo助记词：{recovered_mnemonic}")
    vault = SentinelVault()

    # 2. 生成加密密钥
    aes_key = vault.derive_key(recovered_mnemonic)
    aes_key_hex = aes_key.hex()
    print(f"【密钥派生完成】：len:{len(aes_key_hex)} -> {aes_key_hex[:20]}...")

    # 3. 用户生前加密资产（如：银行账户、数字遗产）
    my_legacy = "我的瑞士银行账号是：CH123456789，密码是：Sentinel2026"
    print(f"【Demo资产信息】：{my_legacy}")
    encrypted_asset = vault.encrypt_data(aes_key, my_legacy)
    encrypted_asset_hex = encrypted_asset.hex()
    print(f"【数据已加密】：len:{len(encrypted_asset_hex)} -> {encrypted_asset_hex[:40]}...")

    # 4. 模拟继承人通过分片拼凑后进行解密
    print("-" * 50)
    decrypted_content = vault.decrypt_data(aes_key, encrypted_asset)
    print(f"【继承人解密成功】：{decrypted_content}")