A simple implementation of a Merkle Tree

simple-merkle-tree.py

import hashlib


class MerkleNode(object):
    def __init__(self, u=None, v=None, parent=None):
        self.u = u
        self.v = v
        self.parent = parent

    def compute_hash(self) -> bytes:
        if self.u is None and self.v is None:
            raise Exception("No nodes")
        if self.v is None:
            res = self.u.compute_hash()
        else:
            res = hashlib.sha256(self.u.compute_hash() + self.v.compute_hash()).hexdigest().encode()
        return res

    def find_depth(self):
        d = 0
        n = self.parent
        while n is not None:
            d+=1
            n = n.parent
        return d

    def is_free(self):
        return self.u is None or self.v is None

    def __repr__(self):
        return self.compute_hash().decode()


class MerkleLeaf(MerkleNode):
    def __init__(self, el:bytes):
        super().__init__()
        self.el = el

    def compute_hash(self):
        # print("leaf: "+ self.el.decode() + " " + self.el.decode())
        return hashlib.sha256(self.el).hexdigest().encode()


class MerkleTree(object):

    def __init__(self, elements=None):
        self._max_depth = 0

        if elements is not None:
            self.leaves = [MerkleLeaf(e) for e in elements]
            temp = self.leaves
            temp2 = []
            while len(temp) != 1:
                self._max_depth+=1
                for x in range(0, len(temp), 2):
                    try:
                        cur_u, cur_v = temp[x:x + 2]
                        new_node = MerkleNode(u=cur_u, v=cur_v)
                        cur_u.parent = new_node
                        cur_v.parent = new_node
                        temp2.append(new_node)
                    except ValueError:
                        temp2.append(temp[x])

                temp = temp2
                temp2 = []
            self.root = temp[0]
        else:
            self.root = None
            self.leaves = []

    def add(self, el):
        new_leaf = MerkleLeaf(el)
        if self.root is None:
            self.root = MerkleNode(u=new_leaf)
            self.leaves.append(new_leaf)
            new_leaf.parent=self.root
            self._max_depth+=1
        else:
            last_leaf = self.leaves[-1]
            free_parent = self._find_free_node(last_leaf, 0)
            if free_parent.u is None:
                free_parent.u = new_leaf
            else:
                free_parent.v = new_leaf
            new_leaf.parent = free_parent
            self.leaves.append(new_leaf)


    def _find_free_node(self, node:MerkleNode, depth):
        if node.parent is None:
            # node is the root and is not free
            if self._max_depth == depth:
                new_root = MerkleNode(u=node)
                node.parent = new_root
                self._max_depth+=1
                self.root = new_root
                return new_root
            else:
                new_node = MerkleNode(u=node.v, parent=node)
                node.v = new_node
                return new_node
        elif node.parent.is_free():
            return node.parent
        else:
            return self._find_free_node(node.parent, depth+1)

        pass

    def compute_hash(self):
        if self.root is None:
            raise Exception("Tree is empty")
        return self.root.compute_hash()

if __name__ == '__main__':
    

    m = MerkleTree([b"a", b"b", b"c", b"d", b"e", b"f"])
    
    print(m.compute_hash())
    
    dm = MerkleTree()
    dm.add(b"a")
    dm.add(b"b")
    dm.add(b"c")
    dm.add(b"d")
    dm.add(b"e")
    dm.add(b"f")
    print(dm.compute_hash())