merkleTree.py

# This is a basic implementation of merkle tree, which is extensively used in Blockchain based
# technologies for verification and storage of transactions in a block of blockchain

import hashlib
import math

# Node class that represents a node of a binary tree
class Node:
    #function to initialise the node with data, left child and right child, representing a binary
    #tree
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None
    
    #utility function to check if a given node has 2 children
    def isFull(self):
        return self.left and self.right
    
    #utility function to print the contents of the node
    def __str__(self):
        return self.data
    
    #utility function to check if a given node has 0 children (is a leaf node)
    def isLeaf(self):
        return ((self.left == None) and (self.right == None))

# Merkle tree class for actual implementation of the tree
class merkleTree:

    # function to initialise the tree with "None" as root and merkleRoot set to blank
    def __init__(self):
        self.root = None
        self._merkleRoot = ''
        
    # private utility function to return a hash of an input string, using the sha256 algorithm    
    def __returnHash(self, x):
            return (hashlib.sha256(x.encode()).hexdigest())
        
    # function to generate a tree from a list. The list is expected to contain
    # the list of transactions(as strings) we wish to seal in a block    
    def makeTreeFromArray(self, arr):
        
        # all the elements of the list must be leaf nodes, hence total nodes in the tree
        # is [2*(noOfLeafNodes) - 2]
        def __noOfNodesReqd(arr):
            x = len(arr)
            return (2*x - 1)
        
        # Private function to build tree with a list generated by a sequence corresponding to the
        # number of nodes required for the tree.
        # For e.g. if we have 4 transactions, i.e. 4 leafNodes, we will have a total of 7 nodes
        # in the tree. Hence this function takes in a list that looks like [1,2,3,4,5,6,7] and 
        # creates a binary tree by recursively inserting elements in inorder fasion.
        def __buildTree(arr, root, i, n): 
      
            # Base case for recursion  
            if i < n: 
                temp = Node(str(arr[i]))  
                root = temp  

                # insert left child  
                root.left = __buildTree(arr, root.left,2 * i + 1, n)  

                # insert right child  
                root.right = __buildTree(arr, root.right,2 * i + 2, n) 
                
            return root
        
        # This function adds our transactions (leaf Node data) to the Leaf Nodes in the tree
        # hence now the tree looks like :
                #                 ""
                #                 /\
                #         ""              ""
                #         /\              /\
                #    (txn1) (txn2)  (txn3) (txn4)      
        def __addLeafData(arr, node):
            
            if not node:
                return
            
            __addLeafData(arr, node.left)
            if node.isLeaf():
                node.data = self.__returnHash(arr.pop())
            else:
                node.data = ''
            __addLeafData(arr, node.right)
        
        # Driver function calls of the main function makeTreeFromArray
        nodesReqd = __noOfNodesReqd(arr)
        nodeArr = [num for num in range(1,nodesReqd+1)]
        self.root = __buildTree(nodeArr)
        __addLeafData(arr,self.root)

    # utility function to traverse the tree using inorder Traversal algorithm    
    def inorderTraversal(self, node):
        if not node:
            return
        
        self.inorderTraversal(node.left)
        print(node)
        self.inorderTraversal(node.right)
    
    # function to calculate merkle root of the tree. This is a recursive algorithm.
    # Essentially, the data of parent node P, is the hash of concatenation of data of its two 
    # child nodes, A and B. If H(x) is the hash function,
    # P = H( H(A) + H(B) )
    # If R is the root, then data of R is the merkle hash or merkle root.  
    def calculateMerkleRoot(self):
         
        def __merkleHash(node):
            if node.isLeaf():
                return node
            
            left = __merkleHash(node.left).data
            right = __merkleHash(node.right).data
            node.data = self.__returnHash(left+right)
            return node
        
        merkleRoot = __merkleHash(self.root)
        self._merkleRoot = merkleRoot.data
        
        return self._merkleRoot 

    #utility function to get the merkle Root of that tree
    def getMerkleRoot(self):
        return self._merkleRoot
    
    # utility function to verify if the transactions are intact with respect to this tree.
    # say we know transaction list A is intact and original. We create a merkle Tree with that.
    # Now we want to know whether the original list of transactions has been tampered with,
    # this function takes in a new list, creates a new merkle tree, calculate its merkle root,
    # and returns if merkle roots match. If they match, transaction list is intact and verified.
    # If transactions are tampered or changed, this function returns False.
    def verifyUtil(self, arr1):
            hash1 = self.getMerkleRoot()
            new_tree = merkleTree()
            new_tree.makeTreeFromArray(arr1)
            new_tree.calculateMerkleRoot()
            hash2 = new_tree.getMerkleRoot()
            if hash1 == hash2 :
                print("Transactions verified Successfully")
                return True
            else:
                print("Transactions have been tampered")
                return False