Add Merkle Tree construction and verification algorithm

blockchain/merkle_tree.py

@@ -0,0 +1,107 @@
+"""
+Merkle Tree Construction and Verification
+
+This module implements the construction of a Merkle Tree and
+verification of inclusion proofs for blockchain data integrity.
+
+Each leaf is a SHA-256 hash of a transaction, and internal nodes are
+computed by hashing the concatenation of their child nodes.
+
+References:
+https://en.wikipedia.org/wiki/Merkle_tree
+"""
+
+import hashlib
+
+
+def sha256(data: str) -> str:
+    """
+    Compute the SHA-256 hash of the given string.
+
+    Args:
+        data (str): Input string.
+
+    Returns:
+        str: Hexadecimal SHA-256 hash of the input.
+
+    Example:
+        >>> sha256("abc")
+        'ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad'
+    """
+    return hashlib.sha256(data.encode()).hexdigest()
+
+
+def build_merkle_tree(leaves: list[str]) -> list[list[str]]:
+    """
+    Build a Merkle Tree from the given leaf nodes.
+
+    Args:
+        leaves: List of data strings (transactions).
+
+    Returns:
+        A list of lists representing tree levels,
+        with the last level containing the Merkle root.
+
+    >>> len(build_merkle_tree(["a", "b", "c", "d"])[-1][0])
+    64
+    """
+    if not leaves:
+        raise ValueError("Leaf list cannot be empty.")
+
+    current_level = [sha256(x) for x in leaves]
+    tree = [current_level]
+
+    while len(current_level) > 1:
+        next_level = []
+        for i in range(0, len(current_level), 2):
+            left = current_level[i]
+            right = current_level[i + 1] if i + 1 < len(current_level) else left
+            next_level.append(sha256(left + right))
+        current_level = next_level
+        tree.append(current_level)
+
+    return tree
+
+
+def merkle_root(leaves: list[str]) -> str:
+    """
+    Return the Merkle root hash for a given list of data.
+
+    >>> r = merkle_root(["tx1", "tx2", "tx3"])
+    >>> isinstance(r, str)
+    True
+    """
+    return build_merkle_tree(leaves)[-1][0]
+
+
+def verify_proof(leaf: str, proof: list[str], root: str) -> bool:
+    """
+    Verify inclusion of a leaf using a Merkle proof.
+
+    Args:
+        leaf: Original data string.
+        proof: List of sibling hashes up the path.
+        root: Expected Merkle root hash.
+
+    Returns:
+        True if proof is valid, else False.
+
+    >>> data = ["a", "b", "c", "d"]
+    >>> tree = build_merkle_tree(data)
+    >>> root = tree[-1][0]
+    >>> leaf = "a"
+    >>> proof = [sha256("b"), sha256(sha256("c") + sha256("d"))]
+    >>> verify_proof(leaf, proof, root)
+    True
+    """
+    computed_hash = sha256(leaf)
+    for sibling in proof:
+        combined = sha256(computed_hash + sibling)
+        computed_hash = combined
+    return computed_hash == root
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()