"""Two Merkle trees, side by side: the naive textbook one and RFC 6962's.

Both expose the same interface — Tree(blocks), .root, .proof(i), and a
classmethod verify(root, block, proof) -> bool. Proofs are lists of
(sibling_digest, side) pairs, bottom-up, where `side` says which side the
sibling sits on. There is no clock and no randomness anywhere: SHA-256 over
fixed bytes makes every digest on this page byte-identical on every run.

They differ in exactly two places, and both are the point:

  NaiveTree   leaf = H(block)          node = H(left || right)
  TaggedTree  leaf = H(0x00 || block)  node = H(0x01 || left || right)

  NaiveTree   an odd level duplicates its last node to make a pair
  TaggedTree  an n-leaf tree splits at the largest power of two below n

verify() is written once, on NaiveTree, and inherited unchanged. The two
schemes run the identical folding loop; only leaf_hash and node_hash differ.
"""

import hashlib

LEFT = "L"   # the sibling is on the left; fold it as H(sibling || running)
RIGHT = "R"  # the sibling is on the right; fold it as H(running || sibling)


class NaiveTree:
    """A Merkle tree the way it is usually first explained: hash the blocks,
    hash pairs of hashes, repeat. A level with an odd count duplicates its
    last node to make a pair.
    """

    def __init__(self, blocks):
        if not blocks:
            raise ValueError("a Merkle tree needs at least one block")
        level = [self.leaf_hash(b) for b in blocks]
        self.levels = []
        while len(level) > 1:
            if len(level) % 2:
                level = level + [level[-1]]   # the CVE-2012-2459 line
            self.levels.append(level)
            level = [self.node_hash(level[i], level[i + 1])
                     for i in range(0, len(level), 2)]
        self.levels.append(level)

    @staticmethod
    def leaf_hash(block):
        return hashlib.sha256(block).digest()

    @staticmethod
    def node_hash(left, right):
        return hashlib.sha256(left + right).digest()

    @property
    def root(self):
        return self.levels[-1][0]

    def proof(self, index):
        """The sibling digests needed to recompute the root from block
        `index`, bottom level first.
        """
        path = []
        for level in self.levels[:-1]:
            sibling = index ^ 1
            path.append((level[sibling], RIGHT if index % 2 == 0 else LEFT))
            index //= 2
        return path

    @classmethod
    def verify(cls, root, block, proof):
        """Does `block` sit under `root`, given `proof`?

        Note what is NOT an argument: the number of leaves, or the expected
        proof length. The verifier folds whatever it is handed and compares.
        See commentary section 5.4 — this omission is the attack surface.
        """
        running = cls.leaf_hash(block)
        for sibling, side in proof:
            if side == LEFT:
                running = cls.node_hash(sibling, running)
            else:
                running = cls.node_hash(running, sibling)
        return running == root


class TaggedTree(NaiveTree):
    """RFC 6962's Merkle tree. Two changes from NaiveTree, both about making
    the hash commit to the tree's *shape* and not just its contents:

    1. Domain separation. A leaf hash is H(0x00 || block) and an internal
       hash is H(0x01 || left || right), so no leaf hash can ever collide
       with an internal hash.
    2. No duplication. An n-leaf tree splits at the largest power of two
       below n, so [a, b, c] and [a, b, c, c] are structurally different
       trees rather than the same one.
    """

    def __init__(self, blocks):
        if not blocks:
            raise ValueError("a Merkle tree needs at least one block")
        self.leaves = [self.leaf_hash(b) for b in blocks]
        self._cache = {}
        self._subtree(0, len(self.leaves))

    @staticmethod
    def leaf_hash(block):
        return hashlib.sha256(b"\x00" + block).digest()

    @staticmethod
    def node_hash(left, right):
        return hashlib.sha256(b"\x01" + left + right).digest()

    @staticmethod
    def _split(n):
        """RFC 6962 section 2.1: k is the largest power of two strictly less
        than n. So a 5-leaf tree is a perfect 4-leaf tree beside a 1-leaf
        tree, never a padded 8-leaf one.
        """
        return 1 << ((n - 1).bit_length() - 1)

    def _subtree(self, lo, hi):
        if (lo, hi) not in self._cache:
            if hi - lo == 1:
                self._cache[lo, hi] = self.leaves[lo]
            else:
                k = lo + self._split(hi - lo)
                self._cache[lo, hi] = self.node_hash(self._subtree(lo, k),
                                                     self._subtree(k, hi))
        return self._cache[lo, hi]

    @property
    def root(self):
        return self._subtree(0, len(self.leaves))

    def proof(self, index):
        """Same walk as NaiveTree's, but down the recursive split instead of
        along stored levels. Collected top-down, so it is reversed to hand
        back the same bottom-up order.
        """
        path = []
        lo, hi = 0, len(self.leaves)
        while hi - lo > 1:
            k = lo + self._split(hi - lo)
            if index < k:
                path.append((self._subtree(k, hi), RIGHT))
                hi = k
            else:
                path.append((self._subtree(lo, k), LEFT))
                lo = k
        path.reverse()
        return path


def hexes(proof):
    """A proof rendered for printing: 8 hex characters per sibling."""
    return [f"{side}:{sibling.hex()[:8]}" for sibling, side in proof]