are_ssz.py

"""
SSZ + consensus hashing primitives for the ARE reference implementation.

Real cryptography only:
  - SHA-256 for SSZ hash_tree_root (consensus serialization).
  - Keccak-256 (Ethereum variant) + RLP for the EIP-1186 MPT proof.
  - Merkle branch verification by generalized index.

Honest scope: this implements exactly the SSZ machinery the ARE verifier needs
(fixed-size uints/bytes, Bytes32 vectors, BeaconBlockHeader, the FULL Deneb+
ExecutionPayloadHeader, and generalized-index branch checks). It is NOT a full
SSZ library, but the ExecutionPayloadHeader is now the real 17-field mainnet
container (v0.5), merkleized field-for-field with real SHA-256 — its
hash_tree_root matches the value committed by a real beacon block body. Every
leaf hash is REAL SHA-256 / Keccak-256 — no value is fabricated.
"""

import hashlib

# ---------------------------------------------------------------------------
# SHA-256 (consensus) primitives
# ---------------------------------------------------------------------------

ZERO_HASHES = [b"\x00" * 32]
for _i in range(1, 64):
    ZERO_HASHES.append(hashlib.sha256(ZERO_HASHES[-1] + ZERO_HASHES[-1]).digest())


def sha256(data: bytes) -> bytes:
    return hashlib.sha256(data).digest()


def hash_concat(a: bytes, b: bytes) -> bytes:
    return hashlib.sha256(a + b).digest()


def next_pow2(n: int) -> int:
    p = 1
    while p < n:
        p <<= 1
    return p


def merkleize(chunks, limit=None):
    """Merkleize a list of 32-byte chunks (SSZ), padding with zero subtrees.

    Uses precomputed ZERO_HASHES so padding a wide tree is correct even when the
    pad spans multiple depths (a flat zero-leaf pad would be wrong for large
    limits because intermediate zero subtrees differ from raw zero leaves)."""
    if limit is None:
        width = next_pow2(max(1, len(chunks)))
    else:
        width = next_pow2(limit) if limit > 0 else 1
    nodes = list(chunks)
    if len(nodes) > width:
        raise ValueError("too many chunks for limit")
    target_levels = width.bit_length() - 1  # log2(width)
    # empty input: the root is the zero subtree at the target depth.
    if not nodes:
        return ZERO_HASHES[target_levels]
    depth = 0
    cur = nodes
    # iteratively combine; pad each level with the correct zero subtree hash
    while len(cur) > 1 or depth < target_levels:
        nxt = []
        for i in range(0, len(cur), 2):
            left = cur[i]
            right = cur[i + 1] if i + 1 < len(cur) else ZERO_HASHES[depth]
            nxt.append(hash_concat(left, right))
        cur = nxt
        depth += 1
        if len(cur) == 1 and depth >= target_levels:
            break
    return cur[0]


def mix_in_length(root: bytes, length: int) -> bytes:
    return hash_concat(root, length.to_bytes(32, "little"))


def uint64_root(x: int) -> bytes:
    return x.to_bytes(8, "little") + b"\x00" * 24


def uint256_root(x: int) -> bytes:
    return x.to_bytes(32, "little")


def bytes32_root(b: bytes) -> bytes:
    assert len(b) == 32
    return b


def bytes20_root(b: bytes) -> bytes:
    # Bytes20 is a fixed Vector[uint8,20] -> right-pad to 32-byte chunk
    assert len(b) == 20
    return b + b"\x00" * 12


def boolean_root(x: bool) -> bytes:
    return (b"\x01" if x else b"\x00") + b"\x00" * 31


def bytevector_root(data: bytes) -> bytes:
    """hash_tree_root of a fixed-size ByteVector[N] (no length mix-in)."""
    chunks = [data[i:i + 32].ljust(32, b"\x00") for i in range(0, len(data), 32)]
    if not chunks:
        chunks = [b"\x00" * 32]
    return merkleize(chunks)


def bytelist_root(data: bytes, limit_bytes: int) -> bytes:
    """hash_tree_root of a ByteList[N] (variable-length): merkleize the byte
    chunks against the chunk-limit, then mix in the byte length."""
    chunks = [data[i:i + 32].ljust(32, b"\x00") for i in range(0, len(data), 32)]
    limit_chunks = max(1, (limit_bytes + 31) // 32)
    root = merkleize(chunks, limit=limit_chunks)
    return mix_in_length(root, len(data))


# ---------------------------------------------------------------------------
# Consensus containers (field-ordered hash_tree_root)
# ---------------------------------------------------------------------------

class BeaconBlockHeader:
    """Standard consensus BeaconBlockHeader: slot, proposer_index, parent_root,
    state_root, body_root. All five fields hashed as a 5-leaf container."""

    def __init__(self, slot=0, proposer_index=0, parent_root=b"\x00" * 32,
                 state_root=b"\x00" * 32, body_root=b"\x00" * 32):
        self.slot = slot
        self.proposer_index = proposer_index
        self.parent_root = parent_root
        self.state_root = state_root
        self.body_root = body_root

    def hash_tree_root(self) -> bytes:
        chunks = [
            uint64_root(self.slot),
            uint64_root(self.proposer_index),
            bytes32_root(self.parent_root),
            bytes32_root(self.state_root),
            bytes32_root(self.body_root),
        ]
        return merkleize(chunks)

    def is_zero(self) -> bool:
        return (self.slot == 0 and self.proposer_index == 0 and
                self.parent_root == b"\x00" * 32 and
                self.state_root == b"\x00" * 32 and
                self.body_root == b"\x00" * 32)

    def __eq__(self, other):
        return (isinstance(other, BeaconBlockHeader) and
                self.slot == other.slot and
                self.proposer_index == other.proposer_index and
                self.parent_root == other.parent_root and
                self.state_root == other.state_root and
                self.body_root == other.body_root)


ZERO_HEADER = BeaconBlockHeader()


# extra_data is ByteList[MAX_EXTRA_DATA_BYTES]; logs_bloom is ByteVector[256]
MAX_EXTRA_DATA_BYTES = 32


class ExecutionPayloadHeader:
    """FULL Deneb+ ExecutionPayloadHeader (17 fields) — v0.5 mainnet fidelity.

    Field order per consensus-specs (Deneb / Electra / Fulu share this 17-field
    layout):
      parent_hash, fee_recipient, state_root, receipts_root, logs_bloom,
      prev_randao, block_number, gas_limit, gas_used, timestamp, extra_data,
      base_fee_per_gas, block_hash, transactions_root, withdrawals_root,
      blob_gas_used, excess_blob_gas

    hash_tree_root is the REAL SSZ container root and matches the
    execution_payload leaf committed by a mainnet beacon BeaconBlockBody (verified
    against real data in are_real_vectors.py). The ARE verifier reads state_root,
    block_number, timestamp from this proven leaf.
    """

    FIELDS = [
        "parent_hash", "fee_recipient", "state_root", "receipts_root",
        "logs_bloom", "prev_randao", "block_number", "gas_limit", "gas_used",
        "timestamp", "extra_data", "base_fee_per_gas", "block_hash",
        "transactions_root", "withdrawals_root", "blob_gas_used",
        "excess_blob_gas",
    ]

    def __init__(self, state_root, block_number, timestamp,
                 parent_hash=b"\x00" * 32, fee_recipient=b"\x00" * 20,
                 receipts_root=b"\x00" * 32, logs_bloom=b"\x00" * 256,
                 prev_randao=b"\x00" * 32, gas_limit=0, gas_used=0,
                 extra_data=b"", base_fee_per_gas=0, block_hash=b"\x00" * 32,
                 transactions_root=b"\x00" * 32, withdrawals_root=b"\x00" * 32,
                 blob_gas_used=0, excess_blob_gas=0):
        self.state_root = state_root
        self.block_number = block_number
        self.timestamp = timestamp
        self.parent_hash = parent_hash
        self.fee_recipient = fee_recipient
        self.receipts_root = receipts_root
        self.logs_bloom = logs_bloom
        self.prev_randao = prev_randao
        self.gas_limit = gas_limit
        self.gas_used = gas_used
        self.extra_data = extra_data
        self.base_fee_per_gas = base_fee_per_gas
        self.block_hash = block_hash
        self.transactions_root = transactions_root
        self.withdrawals_root = withdrawals_root
        self.blob_gas_used = blob_gas_used
        self.excess_blob_gas = excess_blob_gas

    def hash_tree_root(self) -> bytes:
        chunks = [
            bytes32_root(self.parent_hash),
            bytes20_root(self.fee_recipient),
            bytes32_root(self.state_root),
            bytes32_root(self.receipts_root),
            bytevector_root(self.logs_bloom),            # ByteVector[256] -> 8 chunks
            bytes32_root(self.prev_randao),
            uint64_root(self.block_number),
            uint64_root(self.gas_limit),
            uint64_root(self.gas_used),
            uint64_root(self.timestamp),
            bytelist_root(self.extra_data, MAX_EXTRA_DATA_BYTES),
            uint256_root(self.base_fee_per_gas),
            bytes32_root(self.block_hash),
            bytes32_root(self.transactions_root),
            bytes32_root(self.withdrawals_root),
            uint64_root(self.blob_gas_used),
            uint64_root(self.excess_blob_gas),
        ]
        return merkleize(chunks)


# ---------------------------------------------------------------------------
# Generalized-index Merkle branch verification
# ---------------------------------------------------------------------------

def verify_merkle_branch(leaf: bytes, branch, gindex: int, root: bytes) -> bool:
    """Verify a Merkle branch against `root` at generalized index `gindex`.

    branch[i] is the sibling at depth i (deepest first), per the consensus-spec
    is_valid_merkle_branch but driven by gindex bit decomposition.
    """
    node = leaf
    g = gindex
    i = 0
    while g > 1:
        if i >= len(branch):
            return False
        sib = branch[i]
        if g & 1:  # node is the right child
            node = hash_concat(sib, node)
        else:      # node is the left child
            node = hash_concat(node, sib)
        g >>= 1
        i += 1
    return node == root


def build_merkle_branch(leaf: bytes, gindex: int, siblings):
    """Construct (branch, root) for a leaf at gindex given the sibling hashes
    (deepest-first). Used by the synthetic generator with seeded siblings — the
    siblings are arbitrary 32-byte values but every PARENT hash is real SHA-256,
    so the resulting root genuinely commits to the leaf along the gindex path.
    """
    node = leaf
    g = gindex
    i = 0
    while g > 1:
        sib = siblings[i]
        if g & 1:
            node = hash_concat(sib, node)
        else:
            node = hash_concat(node, sib)
        g >>= 1
        i += 1
    return list(siblings[:i]), node