arborist/tests/test_merkle.py

"""Merkle round-trip and tamper-detection tests.

These exercise the proxy.unturf.com Go-merkle conventions ported into Python:
- non-commutative HashCombine (0x03 prefix)
- self-duplicate odd elements
- explicit IsLeft per sibling (no lexical ordering)
"""

from __future__ import annotations

import pytest

from aborist.merkle import (
    HASH_LEN,
    MerkleTree,
    ZERO_HASH,
    hash_combine,
    hash_leaf,
    proof_from_dict,
    proof_to_dict,
    verify_proof,
)


def _leaf(s: str) -> bytes:
    return hash_leaf(s.encode("utf-8"))


def test_empty_tree_root_is_zero_hash():
    tree = MerkleTree.build([])
    assert tree.root == ZERO_HASH


def test_single_leaf_root_equals_leaf():
    leaves = [_leaf("only-chunk")]
    tree = MerkleTree.build(leaves)
    # No interior layer needed; root is the single leaf itself.
    assert tree.root == leaves[0]


def test_combine_is_non_commutative():
    a = _leaf("a")
    b = _leaf("b")
    assert hash_combine(a, b) != hash_combine(b, a)


def test_two_leaves_round_trip():
    leaves = [_leaf("alpha"), _leaf("beta")]
    tree = MerkleTree.build(leaves)
    expected_root = hash_combine(leaves[0], leaves[1])
    assert tree.root == expected_root

    for i in range(len(leaves)):
        proof = tree.proof(i)
        assert verify_proof(proof)
        assert proof.root == tree.root


def test_three_leaves_self_duplicate_odd():
    leaves = [_leaf("a"), _leaf("b"), _leaf("c")]
    tree = MerkleTree.build(leaves)
    # Layer 1: combine(a,b), combine(c,c)
    n01 = hash_combine(leaves[0], leaves[1])
    n22 = hash_combine(leaves[2], leaves[2])
    expected_root = hash_combine(n01, n22)
    assert tree.root == expected_root

    for i in range(3):
        proof = tree.proof(i)
        assert verify_proof(proof), f"proof for index {i} should verify"


def test_four_leaves_full_round_trip():
    leaves = [_leaf(f"chunk-{i}") for i in range(4)]
    tree = MerkleTree.build(leaves)
    for i in range(4):
        proof = tree.proof(i)
        assert verify_proof(proof)
        assert len(proof.siblings) == 2  # log2(4) = 2 levels


def test_seven_leaves_full_round_trip():
    """Odd intermediate layers self-duplicate; every proof must still verify."""
    leaves = [_leaf(f"chunk-{i}") for i in range(7)]
    tree = MerkleTree.build(leaves)
    for i in range(7):
        proof = tree.proof(i)
        assert verify_proof(proof), f"proof for index {i} should verify"


def test_proof_serialization_round_trip():
    leaves = [_leaf(f"x-{i}") for i in range(5)]
    tree = MerkleTree.build(leaves)
    proof = tree.proof(2)
    d = proof_to_dict(proof)
    restored = proof_from_dict(d)
    assert restored == proof
    assert verify_proof(restored)


def test_tampered_leaf_fails_verification():
    leaves = [_leaf(f"chunk-{i}") for i in range(4)]
    tree = MerkleTree.build(leaves)
    proof = tree.proof(1)
    # Tamper: replace the leaf bytes.
    bad_leaf = _leaf("not-the-real-chunk")
    bad_proof = type(proof)(
        leaf=bad_leaf,
        leaf_index=proof.leaf_index,
        siblings=proof.siblings,
        root=proof.root,
    )
    assert not verify_proof(bad_proof)


def test_tampered_sibling_fails_verification():
    leaves = [_leaf(f"chunk-{i}") for i in range(4)]
    tree = MerkleTree.build(leaves)
    proof = tree.proof(0)
    from aborist.merkle import ProofNode

    bad_siblings = list(proof.siblings)
    s0 = bad_siblings[0]
    bad_siblings[0] = ProofNode(hash=b"\xff" * HASH_LEN, is_left=s0.is_left)
    bad_proof = type(proof)(
        leaf=proof.leaf,
        leaf_index=proof.leaf_index,
        siblings=tuple(bad_siblings),
        root=proof.root,
    )
    assert not verify_proof(bad_proof)


def test_swapped_is_left_flag_fails():
    """Order must be preserved — flipping the IsLeft flag must reject."""
    leaves = [_leaf(f"chunk-{i}") for i in range(4)]
    tree = MerkleTree.build(leaves)
    proof = tree.proof(1)
    from aborist.merkle import MerkleProof, ProofNode

    flipped = MerkleProof(
        leaf=proof.leaf,
        leaf_index=proof.leaf_index,
        siblings=tuple(
            ProofNode(hash=s.hash, is_left=not s.is_left) for s in proof.siblings
        ),
        root=proof.root,
    )
    assert not verify_proof(flipped)


def test_out_of_range_leaf_index():
    tree = MerkleTree.build([_leaf("x"), _leaf("y")])
    with pytest.raises(IndexError):
        tree.proof(5)