merkle.py

# Merkle Tree implementation in Python

from collections import deque
from hashlib import sha3_224


class MerkleNode:
    def __init__(self, children: list = None, value: str = None):
        assert not (children and value is not None) and (children or value is not None), 'requires either children or value'

        sha = sha3_224()

        self.children = None
        if children:
            child_nodes = list(map(self.hash_node, children))
            self.children = child_nodes

            for node in child_nodes:
                sha.update(node.val)

        if value is not None:
            sha.update(bytes(value, 'utf-8'))

        self.val = bytes(sha.hexdigest(), 'utf-8')

    @staticmethod
    def hash_node(val):
        if type(val) is list:
            # if list
            return MerkleNode(children=val)
        else:
            # if a value (also str)
            return MerkleNode(value=val)


class MerkleTree:
    def __init__(self, tree_list: list):
        self.tree = tree_list
        self.root = self.hash(tree_list)

    @staticmethod
    def hash(tree_list: list):
        return MerkleNode(children=tree_list)

    def __eq__(self, other):
        if type(other) is MerkleTree:
            return self.root.val == other.root.val

        return False

    def __str__(self):
        nodes = deque([self.root])

        to_print = []

        while nodes:
            next_nodes = deque()
            while nodes:
                n = nodes.popleft()
                to_print.append(n.val.decode('utf-8'))
                to_print.append(' ')

                if n.children is not None:
                    next_nodes.extend(n.children)

            to_print.append('\n')
            nodes = next_nodes

        return ''.join(to_print)



tree_1 = ['hello', '', ['goodbye']]
tree_2 = ['hello', ['goodbye']]

m1 = MerkleTree(tree_1)
m2 = MerkleTree(tree_2)

# test equality
print(m1 != m2)
print(m1 ==  m2)

# print hashes
print(m1)