# compact_block_relay.py — 模拟 Compact Block 的短哈希生成与重构过程
import hashlib
import struct
from typing import List, Dict

TXID_SIZE = 32        # 完整 txid 为 32 字节
SHORT_ID_SIZE = 6     # BIP-152 短哈希为 6 字节
SHORT_ID_MASK = (1 << 48) - 1


def short_txid(txid: bytes, nonce: int) -> int:
    """根据 BIP-152 生成交易的 6 字节短标识符。"""
    h = hashlib.sha256(txid + struct.pack("<Q", nonce)).digest()
    return struct.unpack("<Q", h[:8])[0] & SHORT_ID_MASK


class Mempool:
    """本地交易池，用于演示短哈希查找。"""
    def __init__(self):
        self.txns: Dict[bytes, dict] = {}

    def add(self, tx: dict):
        txid = bytes.fromhex(tx["txid"])
        self.txns[txid] = tx

    def lookup_by_short_id(self, sid: int, nonce: int) -> dict | None:
        for txid, tx in self.txns.items():
            if short_txid(txid, nonce) == sid:
                return tx
        return None


def test_compact_block_reconstruction():
    """演示 compact block 的短哈希匹配与区块重构。"""
    nonce = 0xDEADBEEF
    pool = Mempool()

    # 模拟本地已收到的 4 笔交易（tx1~tx3 在矿工新区块中）
    txs = [
        {"txid": "aa" * 32, "data": "tx1_data"},
        {"txid": "bb" * 32, "data": "tx2_data"},
        {"txid": "cc" * 32, "data": "tx3_data"},
        {"txid": "dd" * 32, "data": "tx4_data"},  # 不在新区块中
    ]
    for tx in txs:
        pool.add(tx)

    # 矿工发送：新区块包含 tx1, tx2, tx3，以及一笔新交易 tx_new
    block_txids = [
        bytes.fromhex("aa" * 32),
        bytes.fromhex("bb" * 32),
        bytes.fromhex("cc" * 32),
        bytes.fromhex("ee" * 32),  # 新交易，接收方未知
    ]

    compact_msg = []
    missing_txns = []
    for txid in block_txids:
        sid = short_txid(txid, nonce)
        found = pool.lookup_by_short_id(sid, nonce)
        if found:
            compact_msg.append(("matched", sid, found))
        else:
            compact_msg.append(("missing", txid.hex()))
            missing_txns.append(txid.hex())

    print("=== Compact Block 重构结果 ===")
    for item in compact_msg:
        if item[0] == "matched":
            print(f"  [短哈希匹配] sid={item[1]:016x} -> tx: {item[2]}")
        else:
            print(f"  [缺失交易]   需要请求完整 txid: {item[3]}")

    assert len(missing_txns) == 1, "只有 tx_new 应该缺失"
    print("\n✅ 区块重构成功：3/4 交易由短哈希从 Mempool 重建。")


if __name__ == "__main__":
    test_compact_block_reconstruction()

# mempool_priority_queue.py — 模拟手续费率优先级的 Mempool
import heapq
from dataclasses import dataclass, field
from typing import List


@dataclass(order=True)
class TxEntry:
    fee_rate: float  # sat/vB（实际存储负值以便 Python 最小堆模拟最大堆）
    txid: str = field(compare=False)
    size_vbytes: int = field(compare=False)
    fee: int = field(compare=False)


class MempoolSim:
    """模拟手续费率优先级的 Mempool（最小堆实现）。"""
    MAX_SIZE_BYTES = 10_000  # 模拟限制

    def __init__(self):
        self._heap: List[TxEntry] = []
        self._total_vsize = 0

    def add_tx(self, txid: str, fee_sat: int, size_vb: int):
        fee_rate = fee_sat / size_vb
        entry = TxEntry(fee_rate=-fee_rate, txid=txid,
                        size_vbytes=size_vb, fee=fee_sat)
        # Python heapq 为最小堆，用负值实现最大堆
        heapq.heappush(self._heap, entry)
        self._total_vsize += size_vb
        self._evict_if_needed()

    def _evict_if_needed(self):
        while self._total_vsize > self.MAX_SIZE_BYTES and self._heap:
            worst = heapq.heappop(self._heap)
            self._total_vsize -= worst.size_vbytes
            print(f"  [驱逐] tx={worst.txid} fee_rate={-worst.fee_rate:.1f} sat/vB")

    def pop_best_block(self, block_max_vsize: int = 4000) -> List[TxEntry]:
        """模拟矿工打包：从堆顶取最高手续费率交易，凑满一个区块。"""
        selected = []
        used_vsize = 0
        temp_heap = list(self._heap)
        heapq.heapify(temp_heap)

        while temp_heap and used_vsize < block_max_vsize:
            entry = heapq.heappop(temp_heap)
            if used_vsize + entry.size_vbytes <= block_max_vsize:
                selected.append(entry)
                used_vsize += entry.size_vbytes
        return selected


def test_mempool_sim():
    pool = MempoolSim()
    pool.add_tx("tx_a", 300_000, 300)    # 1000 sat/vB
    pool.add_tx("tx_b", 100_000, 400)    # 250 sat/vB
    pool.add_tx("tx_c", 60_000, 200)     # 300 sat/vB
    pool.add_tx("tx_d", 10_000, 500)     # 20  sat/vB
    pool.add_tx("tx_e", 500_000, 1000)   # 500 sat/vB

    block = pool.pop_best_block()
    print("\n=== 矿工打包的最佳区块 ===")
    for tx in block:
        print(f"  {tx.txid}: {-tx.fee_rate:.1f} sat/vB  fee={tx.fee}  size={tx.size_vbytes}vB")

    assert len(block) <= 5
    assert all(tx.fee_rate <= 0 for tx in block)  # 负值表示排序正确
    print("\n✅ Mempool 手续费率排序与区块选择演示完成。")


if __name__ == "__main__":
    test_mempool_sim()

# tx_validation_pipeline.py — 模拟交易的4层验证流水线
from dataclasses import dataclass
from typing import Set


@dataclass
class UTXO:
    txid: str
    vout: int
    amount: int
    script_pubkey: str


class TxValidationPipeline:
    """模拟交易的4层验证流水线。"""
    def __init__(self):
        self.utxo_set: Set[str] = set()          # 格式: "txid:vout"
        self.mempool_spent: Set[str] = set()     # 已在 Mempool 中花掉的 UTXO
        self.utxo_details: dict = {}

    def _level1_syntax(self, raw_tx: dict) -> bool:
        """第1层：语法检查"""
        required = {"version", "vin", "vout"}
        if not required.issubset(raw_tx.keys()):
            return False
        if not isinstance(raw_tx["vin"], list) or not raw_tx["vin"]:
            return False
        if not isinstance(raw_tx["vout"], list) or not raw_tx["vout"]:
            return False
        for vin in raw_tx["vin"]:
            if not {"txid", "vout"}.issubset(vin.keys()):
                return False
        return True

    def _level2_script(self, raw_tx: dict) -> bool:
        """第2层：脚本检查（简化模拟，真实实现涉及 ECDSA/secp256k1 验证）"""
        return True

    def _level3_utxo(self, raw_tx: dict) -> bool:
        """第3层：UTXO 存在性"""
        for vin in raw_tx["vin"]:
            key = f"{vin['txid']}:{vin['vout']}"
            if key not in self.utxo_set:
                print(f"    [失败] 输入 {key} 不在 UTXO 集中")
                return False
        return True

    def _level4_doublespend(self, raw_tx: dict) -> bool:
        """第4层：双花检测（比对 Mempool）"""
        for vin in raw_tx["vin"]:
            key = f"{vin['txid']}:{vin['vout']}"
            if key in self.mempool_spent:
                print(f"    [失败] 输入 {key} 已在 Mempool 中被花费")
                return False
        return True

    def validate(self, raw_tx: dict) -> bool:
        print("\n--- 交易验证流水线 ---")
        print(f"  txid: {raw_tx.get('txid', 'unknown')}")

        if not self._level1_syntax(raw_tx):
            print("  ❌ 第1层失败：语法检查")
            return False
        print("  ✅ 第1层通过：语法检查")

        if not self._level2_script(raw_tx):
            print("  ❌ 第2层失败：脚本/签名检查")
            return False
        print("  ✅ 第2层通过：脚本/签名检查")

        if not self._level3_utxo(raw_tx):
            print("  ❌ 第3层失败：UTXO 存在性")
            return False
        print("  ✅ 第3层通过：UTXO 存在性")

        if not self._level4_doublespend(raw_tx):
            print("  ❌ 第4层失败：双花检测")
            return False
        print("  ✅ 第4层通过：双花检测")

        for vin in raw_tx["vin"]:
            self.mempool_spent.add(f"{vin['txid']}:{vin['vout']}")

        print("  🎯 交易验证通过，加入 Mempool！")
        return True


def test_validation_pipeline():
    pipeline = TxValidationPipeline()
    pipeline.utxo_set.add("prev_tx:0")
    pipeline.utxo_details["prev_tx:0"] = UTXO("prev_tx", 0, 50000, "script")

    # 有效交易
    tx_valid = {
        "txid": "valid_tx", "version": 2,
        "vin": [{"txid": "prev_tx", "vout": 0, "scriptSig": "..."}],
        "vout": [{"value": 49000, "scriptPubKey": "..."}],
    }
    assert pipeline.validate(tx_valid) is True

    # 双花尝试（同一 UTXO）
    tx_double = {
        "txid": "double_tx", "version": 2,
        "vin": [{"txid": "prev_tx", "vout": 0, "scriptSig": "..."}],
        "vout": [{"value": 49000, "scriptPubKey": "..."}],
    }
    assert pipeline.validate(tx_double) is False

    # 无效语法
    tx_bad_syntax = {"txid": "bad_tx", "version": 1}
    assert pipeline.validate(tx_bad_syntax) is False

    print("\n✅ 交易验证流水线全部测试通过！")


if __name__ == "__main__":
    test_validation_pipeline()

# ddos_rate_limiter.py — 速率限制与 Ban Score DDoS 防御模拟
import time
from collections import defaultdict
from dataclasses import dataclass, field


@dataclass
class RateLimiter:
    """按连接和消息类型做速率限制。"""
    max_per_second: int = 10
    _window: dict = field(default_factory=lambda: defaultdict(list))

    def allow(self, peer_id: str, msg_type: str) -> bool:
        now = time.monotonic()
        key = (peer_id, msg_type)
        timestamps = self._window[key]
        while timestamps and timestamps[0] < now - 1.0:
            timestamps.pop(0)
        if len(timestamps) >= self.max_per_second:
            return False
        timestamps.append(now)
        return True


class BanScore:
    """基于行为的惩罚计数器。"""
    BAN_THRESHOLD = 100

    def __init__(self):
        self._scores: dict = defaultdict(int)

    def add(self, peer_id: str, points: int) -> bool:
        self._scores[peer_id] += points
        return self._scores[peer_id] >= self.BAN_THRESHOLD

    def get_score(self, peer_id: str) -> int:
        return self._scores.get(peer_id, 0)

    def reset(self, peer_id: str):
        self._scores.pop(peer_id, None)


def test_ddos_defense():
    limiter = RateLimiter(max_per_second=3)
    ban = BanScore()
    peer = "192.168.1.100:8333"

    print("=== 模拟 DDoS 防御 ===")

    # 模拟快速发送 5 条 inv 消息（应触发限制）
    for i in range(5):
        allowed = limiter.allow(peer, "inv")
        status = "✅ 允许" if allowed else "❌ 限制"
        print(f"  消息 {i+1}: {status}")
        if not allowed:
            ban.add(peer, 30)

    print(f"\n  Ban Score: {ban.get_score(peer)}")
    assert ban.get_score(peer) > 0, "应有惩罚分"

    for _ in range(4):
        ban.add(peer, 30)
    assert ban.get_score(peer) >= BanScore.BAN_THRESHOLD
    print(f"  ⚠️ 节点 {peer} 已超过阈值得分 {ban.get_score(peer)}，断开连接")

    print("\n✅ DDoS 防御机制演示完成。")


if __name__ == "__main__":
    test_ddos_defense()

# merkle_proof_verification.py — Merkle 路径验证演示
import hashlib
from typing import List


def dhash(data: bytes) -> bytes:
    """双重 SHA256：SHA256(SHA256(data))"""
    return hashlib.sha256(hashlib.sha256(data).digest()).digest()


def verify_merkle_proof(txid: bytes, merkle_path: List[bytes],
                        index: int, merkle_root: bytes) -> bool:
    """验证 Merkle 路径证明。

    Args:
        txid: 交易哈希 (32 字节)
        merkle_path: 路径中每层的兄弟节点哈希列表
        index: 交易在区块交易列表中的位置（从 0 开始）
        merkle_root: 已知的区块头中的 Merkle Root

    Returns:
        True 如果路径哈希运算后等于 merkle_root
    """
    current = txid
    pos = index

    for sibling in merkle_path:
        if pos % 2 == 0:
            current = dhash(current + sibling)   # 当前在左侧
        else:
            current = dhash(sibling + current)   # 当前在右侧
        pos //= 2

    return current == merkle_root


def test_merkle_proof():
    txids = [
        bytes.fromhex("aa" * 32),
        bytes.fromhex("bb" * 32),
        bytes.fromhex("cc" * 32),
        bytes.fromhex("dd" * 32),
    ]

    # 构建 Merkle 树
    l1_0 = dhash(txids[0] + txids[1])
    l1_1 = dhash(txids[2] + txids[3])
    merkle_root = dhash(l1_0 + l1_1)

    # 验证 txids[1] (index=1)
    txid_to_prove = txids[1]
    merkle_path = [txids[0], l1_1]
    index = 1

    result = verify_merkle_proof(txid_to_prove, merkle_path, index, merkle_root)
    print(f"Merkle Root: {merkle_root.hex()}")
    print(f"验证交易 index={index}，路径节点数={len(merkle_path)}")
    print(f"验证结果: {'✅ 通过' if result else '❌ 失败'}")

    # 虚假交易应失败
    fake_txid = bytes.fromhex("ee" * 32)
    bad_result = verify_merkle_proof(fake_txid, merkle_path, index, merkle_root)
    print(f"错误 txid 验证结果 (应失败): {'✅ 通过' if bad_result else '❌ 失败'}")

    assert result is True
    assert bad_result is False
    print("\n✅ Merkle 路径验证演示完成。")


if __name__ == "__main__":
    test_merkle_proof()

# sync_mode_comparison.py — 以太坊同步模式对比与推荐
from dataclasses import dataclass
from typing import Optional


@dataclass
class SyncMode:
    name: str
    storage_gb: int
    sync_time_hours: float
    full_validation: bool
    description: str


SYNC_MODES = [
    SyncMode("Light Sync (轻同步)", 1, 0.3, False,
             "仅下载区块头，依赖全节点提供状态，适合移动钱包"),
    SyncMode("Fast Sync (快速同步)", 700, 4, False,
             "下载全部区块+枢轴前的状态快照，旧版 Geth 默认"),
    SyncMode("Snap Sync (快照同步)", 400, 2, True,
             "下载部分状态 + 快照修复并行验证，当前默认模式"),
    SyncMode("Full Archive (全归档)", 15000, 168, True,
             "从创世完整执行每笔交易，精确但最慢最贵"),
]


def recommend_sync(use_case: str) -> Optional[SyncMode]:
    """根据用途推荐同步模式。"""
    case_map = {
        "validator": "Snap Sync (快照同步)",
        "wallet": "Light Sync (轻同步)",
        "explorer": "Full Archive (全归档)",
        "quick_node": "Fast Sync (快速同步)",
        "default": "Snap Sync (快照同步)",
    }
    name = case_map.get(use_case, case_map["default"])
    for m in SYNC_MODES:
        if m.name == name:
            return m
    return None


def test_sync_recommendation():
    print("=== 以太坊同步模式推荐 ===")
    cases = ["validator", "wallet", "explorer", "quick_node", "unknown"]
    for case in cases:
        mode = recommend_sync(case)
        print(f"\n  用途: [{case}]")
        if mode:
            print(f"    推荐模式: {mode.name}")
            print(f"    存储需求: ~{mode.storage_gb:,} GB")
            print(f"    同步时间: ~{mode.sync_time_hours} 小时")
            print(f"    完全验证: {'是' if mode.full_validation else '否'}")
            print(f"    说明: {mode.description}")

    assert recommend_sync("validator").name == "Snap Sync (快照同步)"
    assert recommend_sync("wallet").name == "Light Sync (轻同步)"
    print("\n✅ 同步模式对比与推荐演示完成。")


if __name__ == "__main__":
    test_sync_recommendation()