ERC-7908: 金库管理中的 HD 钱包

entity_hash = sha256(f"ENTITY:{entity}".encode()).digest()  
entity_index = int.from_bytes(entity_hash[:4], "big") | 0x80000000   # 2^31 ≤ index < 2^32

dept_hash = sha256(f"DEPT:{entity_hash}:{department}".encode()).digest()  
dept_index = int.from_bytes(dept_hash[:4], "big") | 0x80000000   # 2^31 ≤ index < 2^32

"""
Secure Treasury Management System
企业级分层确定性钱包实现，符合 BIP-44
"""

import hashlib
import logging
from typing import Tuple, Dict
from bip32utils import BIP32Key
from eth_account import Account
from mnemonic import Mnemonic  # 添加 BIP39 支持

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("TreasurySystem")

class TreasurySystem:
    def __init__(self, mnemonic: str):
        """
        初始化金库系统
        :param mnemonic: BIP-39 助记词（12/24 个单词）
        """
        if not Mnemonic("english").check(mnemonic):
            raise ValueError("Invalid BIP-39 mnemonic")
        
        # 使用标准 BIP-39 生成种子
        self.seed = Mnemonic.to_seed(mnemonic, passphrase="")
        self.root_key = BIP32Key.fromEntropy(self.seed)
        
        logger.info("Treasury system initialized. Master key fingerprint: %s", 
                  self.root_key.Fingerprint().hex())

    @staticmethod
    def _hierarchical_hash(entity: str, department: str) -> Tuple[int, int]:
        """
        分层哈希计算（符合提案规范）
        Returns: (entity_index, department_index)
        """
        # 实体哈希
        entity_hash = hashlib.sha256(f"ENTITY:{entity}".encode()).digest()
        entity_index = int.from_bytes(entity_hash[:4], 'big') | 0x80000000
        
        # 部门哈希（链式）
        dept_input = f"DEPT:{entity_hash.hex()}:{department}".encode()
        dept_hash = hashlib.sha256(dept_input).digest()
        dept_index = int.from_bytes(dept_hash[:4], 'big') | 0x80000000
        
        return entity_index, dept_index

    def _derive_key(self, path: list) -> BIP32Key:
        """通用密钥派生方法"""
        current_key = self.root_key
        for index in path:
            if not isinstance(index, int):
                raise TypeError(f"Invalid derivation index type: {type(index)}")
            current_key = current_key.ChildKey(index)
        return current_key

    def generate_account(self, entity: str, department: str, 
                        account_idx: int = 0) -> Dict[str, str]:
        """
        生成部门账户（BIP44 5 层结构）
        Path: m/44'/60'/entity'/dept'/account_idx
        """
        e_idx, d_idx = self._hierarchical_hash(entity, department)
        
        # BIP-44 标准路径
        derivation_path = [
            0x8000002C,  # 44' (强化)
            0x8000003C,  # 60' (以太坊)
            e_idx,       # entity_index (强化)
            d_idx,       # department_index (强化)
            account_idx  # 地址索引
        ]
        
        key = self._derive_key(derivation_path)
        priv_key = key.PrivateKey().hex()
        
        return {
            'path': f"m/44'/60'/{e_idx}'/{d_idx}'/{account_idx}",
            'private_key': priv_key,  # 警告：切勿在生产环境中暴露此信息
            'address': Account.from_key(priv_key).address
        }

    def get_audit_xpub(self, entity: str, department: str) -> str:
        """
        检索部门级扩展公钥（用于审计）
        Path: m/44'/60'/entity'/dept'
        """
        e_idx, d_idx = self._hierarchical_hash(entity, department)
        path = [
            0x8000002C,  # 44'
            0x8000003C,  # 60'
            e_idx,       # entity'
            d_idx        # dept'
        ]
        return self._derive_key(path).ExtendedKey()

    def get_dept_xprv(self, entity: str, department: str) -> str:
        """
        获取部门级扩展私钥（严格控制）
        Path: m/44'/60'/entity'/dept'
        """
        e_idx, d_idx = self._hierarchical_hash(entity, department)
        path = [
            0x8000002C,  # 44'
            0x8000003C,  # 60'
            e_idx,       # entity'
            d_idx        # dept'
        ]
        return self._derive_key(path).ExtendedKey()


    @staticmethod
    def derive_addresses_from_xpub(xpub: str, count: int = 20) -> list:
        """从扩展公钥派生地址（审计用途）"""
        audit_key = BIP32Key.fromExtendedKey(xpub)
        return [
            Account.from_key(
                audit_key
                        .ChildKey(i)   # 地址索引
                        .PrivateKey()
            ).address
            for i in range(count)
        ]


if __name__ == "__main__":
    # 示例用法（删除生产环境中的私钥打印）
    try:
        # 使用标准助记词
        mnemo = Mnemonic("english")
        mnemonic = mnemo.generate(strength=256)
        treasury = TreasurySystem(mnemonic)
        print(f"mnemonic: {mnemonic}")
        
        print("\n=== 财务部门账户生成 ===")
        finance_acc1 = treasury.generate_account("GroupA", "Finance", 0)
        finance_acc2 = treasury.generate_account("GroupA", "Finance", 1)
        print(f"Account1 path: {finance_acc1['path']}")
        print(f"Account1 address: {finance_acc1['address']}")
        print(f"Account1 private key: {finance_acc1['private_key']}")
        print(f"Account2 path: {finance_acc2['path']}")
        print(f"Account2 address: {finance_acc2['address']}")
        print(f"Account2 private key: {finance_acc2['private_key']}")

        print("\n=== 审计验证测试 ===")
        audit_xpub = treasury.get_audit_xpub("GroupA", "Finance")
        print(f"Audit xpub: {audit_xpub}")
        audit_addresses = TreasurySystem.derive_addresses_from_xpub(audit_xpub, 2)
        print(f"Audit-derived addresses: {audit_addresses}")
        
        assert finance_acc1['address'] in audit_addresses, "Audit verification failed"
        assert finance_acc2['address'] in audit_addresses, "Audit verification failed"
        print("✅ Audit verification successful")

        print("\n=== 部门隔离测试 ===")
        other_dept_acc = treasury.generate_account("GroupA", "Audit", 0)
        print(f"Account3 path: {other_dept_acc['path']}")
        print(f"Account3 address: {other_dept_acc['address']}")
        assert other_dept_acc['address'] not in audit_addresses, "Isolation breach"
        print("✅ Department isolation effective")


        print("\n=== 部门私钥共享测试 ===")
        # 获取部门级扩展私钥
        dept_xprv = treasury.get_audit_xpub("GroupA", "Finance").replace('xpub', 'xprv')  # 实际应通过专用方法获取
        print(f"Fiance xprv: {dept_xprv}")
        # 从扩展私钥派生账户私钥
        dept_key = BIP32Key.fromExtendedKey(dept_xprv)
        derived_acc0_key = dept_key.ChildKey(0).PrivateKey().hex()
        derived_acc1_key = dept_key.ChildKey(1).PrivateKey().hex()
        print(f"Fiance derived_acc0_key: {derived_acc0_key}")
        print(f"Fiance derived_acc1_key: {derived_acc1_key}")
        # 验证私钥派生能力
        assert derived_acc0_key == finance_acc1['private_key'], \
            "Account 0 private key derivation failed"
        assert derived_acc1_key == finance_acc2['private_key'], \
            "Account 1 private key derivation failed"
        print("✅ Private key derivation from department xprv successful")

    except Exception as e:
        logger.error("System error: %s", e, exc_info=True)