"""Stdlib-only tests (no pytest): plain asserts in functions called from a
__main__ block. Run: `python3 test_merkle.py`.

The headline tests pin the two failures the commentary is about — the forged
inclusion proof (section 6.2) and the duplicated-last-block root collision
(section 6.3) — and pin that RFC 6962's tagged tree refuses both. The rest
are round-trip checks so a "fix" that breaks honest proofs can't pass.
"""

from merkle import NaiveTree, TaggedTree

BLOCKS = [b"blk%d" % i + b"." * 59 for i in range(4)]  # 64 bytes each
THREE = [b"a" * 64, b"b" * 64, b"c" * 64]


def test_honest_proofs_round_trip():
    """Both schemes verify every block at every tree size from 1 to 9."""
    for cls in (NaiveTree, TaggedTree):
        for n in range(1, 10):
            blocks = [b"%02d" % i * 32 for i in range(n)]
            tree = cls(blocks)
            for i in range(n):
                assert cls.verify(tree.root, blocks[i], tree.proof(i)), (cls, n, i)


def test_honest_proof_rejects_the_wrong_block():
    """Sanity: the naive tree is not simply accepting everything."""
    tree = NaiveTree(BLOCKS)
    assert NaiveTree.verify(tree.root, BLOCKS[0], tree.proof(0)) is True
    assert NaiveTree.verify(tree.root, BLOCKS[1], tree.proof(0)) is False
    assert NaiveTree.verify(tree.root, b"z" * 64, tree.proof(0)) is False


def test_forged_leaf_verifies_under_naive():
    """THE AHA. H(blk0) || H(blk1) is 64 bytes, is not any block in the
    file, and its leaf hash IS the internal node above blk0 and blk1 — so
    the honest proof minus its first step verifies it against the real root.
    """
    tree = NaiveTree(BLOCKS)
    proof0 = tree.proof(0)
    forged = NaiveTree.leaf_hash(BLOCKS[0]) + proof0[0][0]

    assert len(forged) == 64            # a legal block, size-wise
    assert forged not in BLOCKS         # and never in the file
    assert len(proof0) == 2
    assert NaiveTree.verify(tree.root, forged, proof0[1:]) is True


def test_forged_leaf_fails_under_tagged():
    """Same recipe, adapted to the tagged leaf hash, against the tagged
    tree. The 0x00/0x01 prefixes put leaf and internal hashes in disjoint
    domains, so the forged block's leaf hash can never be a node hash.
    """
    tree = TaggedTree(BLOCKS)
    proof0 = tree.proof(0)
    forged = TaggedTree.leaf_hash(BLOCKS[0]) + proof0[0][0]

    assert len(forged) == 64
    assert TaggedTree.verify(tree.root, forged, proof0[1:]) is False
    # ...and the honest proof still works, i.e. the fix isn't "reject all".
    assert TaggedTree.verify(tree.root, BLOCKS[0], proof0) is True


def test_duplicating_the_last_block_keeps_the_naive_root():
    """CVE-2012-2459: odd levels pad by repeating their last node, so a list
    and that list with its tail repeated build the identical tree. Holds for
    3-vs-4 leaves, 5-vs-6, and (repeating the last *pair*) 6-vs-8.
    """
    six = [b"%02d" % i * 32 for i in range(6)]
    assert NaiveTree(THREE).root == NaiveTree(THREE + THREE[-1:]).root
    assert NaiveTree(six[:5]).root == NaiveTree(six[:5] + six[4:5]).root
    assert NaiveTree(six).root == NaiveTree(six + six[4:6]).root


def test_tagged_split_separates_those_lists():
    """RFC 6962 never pads, so those pairs are different trees."""
    six = [b"%02d" % i * 32 for i in range(6)]
    assert TaggedTree(THREE).root != TaggedTree(THREE + THREE[-1:]).root
    assert TaggedTree(six[:5]).root != TaggedTree(six[:5] + six[4:5]).root
    assert TaggedTree(six).root != TaggedTree(six + six[4:6]).root


def test_proof_length_would_have_caught_the_forgery_but_not_the_collision():
    """The cheap defence (section 7.2) and its limit. A 4-leaf tree always
    yields length-2 proofs, so a length check rejects the forged length-1
    proof. It does nothing about the collision: 3 leaves and 4 leaves both
    produce length-2 proofs, so the check cannot tell the two trees apart.
    """
    tree = NaiveTree(BLOCKS)
    assert len(tree.proof(0)) == 2
    assert len(NaiveTree(BLOCKS).proof(0)[1:]) == 1      # forged, catchable
    assert len(NaiveTree(THREE).proof(0)) == 2           # but 3 leaves...
    assert len(NaiveTree(THREE + THREE[-1:]).proof(0)) == 2   # ...matches 4


def test_empty_block_list_is_rejected():
    for cls in (NaiveTree, TaggedTree):
        try:
            cls([])
        except ValueError:
            pass
        else:
            raise AssertionError(f"{cls.__name__} accepted an empty list")


TESTS = [
    test_honest_proofs_round_trip,
    test_honest_proof_rejects_the_wrong_block,
    test_forged_leaf_verifies_under_naive,
    test_forged_leaf_fails_under_tagged,
    test_duplicating_the_last_block_keeps_the_naive_root,
    test_tagged_split_separates_those_lists,
    test_proof_length_would_have_caught_the_forgery_but_not_the_collision,
    test_empty_block_list_is_rejected,
]


if __name__ == "__main__":
    failed = 0
    for test in TESTS:
        try:
            test()
        except AssertionError as exc:
            failed += 1
            print(f"FAIL  {test.__name__}: {exc}")
        else:
            print(f"PASS  {test.__name__}")
    if failed:
        print(f"\n{failed} of {len(TESTS)} tests FAILED")
        raise SystemExit(1)
    print(f"\nAll {len(TESTS)} tests PASSED")