├── hash_data_structures.py ├── img ├── merkle_proof.jpeg └── merkle_tree.jpeg ├── merkle_proof.py ├── merkle_tree.py ├── node.py ├── readme.md ├── test.py ├── test_hash_data_structures.py ├── test_merkle_proof.py ├── test_merkle_tree.py ├── test_sanity.py └── utils.py /hash_data_structures.py: -------------------------------------------------------------------------------- 1 | from utils import * 2 | 3 | 4 | SECURE_HASH_FUNCTIONS = ['sha1', 'sha224', 'sha256', 'sha384', 'sha512'] 5 | 6 | class HashLeaf(object): 7 | """Attach two pieces of string data and store the hash of the concatenated 8 | strings 9 | """ 10 | def __init__(self, left, right, hash_function): 11 | assert isinstance(hash_function, str), ( 12 | "Hash function is not of type `String`") 13 | self._hash_function = hash_function 14 | self._left = left 15 | self._right = right 16 | self._data = self._evaluate() 17 | self._height = 1 18 | 19 | def _evaluate(self): 20 | """Ensure data is in the form of a string""" 21 | assert isinstance(self._left, str), "Data is not of type `String`" 22 | assert isinstance(self._right, str), "Data is not of type `String`" 23 | return hash_data(self._left + self._right, self._hash_function) 24 | 25 | @property 26 | def data(self): 27 | """str: Allow the user to query the hashed data stored in the 28 | HashLeaf 29 | """ 30 | return self._data 31 | 32 | @property 33 | def height(self): 34 | """int: Allow the user to query the height stored in the HashLeaf""" 35 | return self._height 36 | 37 | class HashNode(HashLeaf): 38 | """Attach two HashLeaf structures and store the hash of their concatenated 39 | data 40 | """ 41 | def __init__(self, left, right, hash_function): 42 | super().__init__(left, right, hash_function) 43 | assert left._hash_function == hash_function, ( 44 | "Hash functions incompatible") 45 | assert right._hash_function == hash_function, ( 46 | "Hash functions incompatible") 47 | self._height = self._left.height + 1 48 | 49 | def _evaluate(self): 50 | """Ensure data is in the form of a HashLeaf data structures and has 51 | the correct height. Separate method from `HashLeaf` as there are 52 | different requirements 53 | """ 54 | assert isinstance(self._left, HashLeaf), ( 55 | "Data is not of type `HashLeaf`") 56 | assert isinstance(self._right, HashLeaf), ( 57 | "Data is not of type `HashLeaf`") 58 | assert self._left.height == self._right.height, ( 59 | "Left and right branch not balanced") 60 | return hash_data(self._left.data + self._right.data, 61 | self._hash_function) 62 | -------------------------------------------------------------------------------- /img/merkle_proof.jpeg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Blockchain-for-Developers/merkle-tree/5b7a96d6ae317d6d583c5f85ef43a4d6c9c3808a/img/merkle_proof.jpeg -------------------------------------------------------------------------------- /img/merkle_tree.jpeg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Blockchain-for-Developers/merkle-tree/5b7a96d6ae317d6d583c5f85ef43a4d6c9c3808a/img/merkle_tree.jpeg -------------------------------------------------------------------------------- /merkle_proof.py: -------------------------------------------------------------------------------- 1 | from utils import * 2 | import math 3 | from node import Node 4 | 5 | 6 | def merkle_proof(tx, merkle_tree): 7 | """Given a tx and a Merkle tree object, retrieve its list of tx's and 8 | parse through it to arrive at the minimum amount of information required 9 | to arrive at the correct block header. This does not include the tx 10 | itself. 11 | 12 | Return this data as a list; remember that order matters! 13 | """ 14 | #### YOUR CODE HERE 15 | 16 | 17 | 18 | def verify_proof(tx, merkle_proof): 19 | """Given a Merkle proof - constructed via `merkle_proof(...)` - verify 20 | that the correct block header can be retrieved by properly hashing the tx 21 | along with every other piece of data in the proof in the correct order 22 | """ 23 | #### YOUR CODE HERE 24 | 25 | 26 | 27 | -------------------------------------------------------------------------------- /merkle_tree.py: -------------------------------------------------------------------------------- 1 | from hash_data_structures import * 2 | 3 | 4 | class MerkleTree(object): 5 | """Merkle Tree implementation, default hash function is 'sha256'. 6 | Nodes are reconstructed upon every tx addition but the list of tx 7 | persistent 8 | """ 9 | def __init__(self, tx_list, hash_function='sha256'): 10 | hash_function = hash_function.lower() 11 | assert tx_list, "No transactions to be hashed" 12 | assert hash_function in SECURE_HASH_FUNCTIONS, ( 13 | "{} is not a valid hash function".format(hash_function)) 14 | self._hash_function = hash_function 15 | self._leaves = tx_list 16 | self._nodes = [] 17 | self._root = self._evaluate() 18 | self._height = self._root.height 19 | self._block_header = self._root.data 20 | 21 | def add_tx(self, *tx): 22 | """Add an arbitrary amount of tx's to the tree. It needs to be 23 | reconstructed every time this happens and the block header 24 | changes as well 25 | """ 26 | tx_in = list(tx) 27 | if type(tx_in[0]) == list: 28 | tx_in = tx_in[0] 29 | self._leaves += tx_in 30 | self._reevaluate() 31 | 32 | def reset_tree(self, hash_function='sha256'): 33 | """Clear the tree data""" 34 | self._hash_function = hash_function 35 | self._nodes = [] 36 | self._height = 0 37 | self._block_header = None 38 | 39 | def _evaluate(self): 40 | """Used to construct the tree and arrive at the block header""" 41 | leaves = list(self._leaves) 42 | if not is_power_of_two(len(leaves)) or len(leaves) < 2: 43 | last_tx = leaves[-1] 44 | while not is_power_of_two(len(leaves)) or len(leaves) < 2: 45 | leaves.append(last_tx) 46 | for tx in range(0, len(leaves), 2): 47 | self._nodes.append(HashLeaf(leaves[tx], leaves[tx+1], 48 | self._hash_function)) 49 | nodes = list(self._nodes) 50 | while len(nodes) > 2: 51 | left = nodes.pop(0) 52 | right = nodes.pop(0) 53 | node = HashNode(left, right, self._hash_function) 54 | nodes.append(node) 55 | if len(nodes) == 1: 56 | return nodes[0] 57 | return HashNode(nodes[0], nodes[1], self._hash_function) 58 | 59 | def _reevaluate(self): 60 | """Resets the tree and makes a call to `_evaluate(...)` to reconstruct 61 | the tree given its persistent list of tx's 62 | """ 63 | self.reset_tree(self._hash_function) 64 | self._root = self._evaluate() 65 | self._height = self._root.height 66 | self._block_header = self._root.data 67 | 68 | @property 69 | def hash_function(self): 70 | """func: Allow the user to query the tree's hash function""" 71 | return self._hash_function 72 | 73 | # @hash_function.setter 74 | def hash_function(self, value): 75 | """Allows the user to change the tree's hash function. Requires that 76 | the tree be rebuilt to accomodate this change 77 | """ 78 | value = value.lower() 79 | assert value in SECURE_HASH_FUNCTIONS, ( 80 | "{} is not a valid hash function".format(value)) 81 | self._hash_function = value 82 | 83 | @property 84 | def block_header(self): 85 | """str: Allow the user to query the tree's block header""" 86 | return self._block_header 87 | 88 | @property 89 | def height(self): 90 | """int: Allow the user to query the tree's height""" 91 | return self._height 92 | 93 | @property 94 | def leaves(self): 95 | """list: Allow the user to query the tree's list of tx's""" 96 | return self._leaves 97 | -------------------------------------------------------------------------------- /node.py: -------------------------------------------------------------------------------- 1 | 2 | class Node: 3 | """Node wrapper class implementation. 4 | Useful for tracking depth of a node when 5 | constructing the merkle proof""" 6 | def __init__(self, direction, tx): 7 | self._direction = direction 8 | self._tx = tx 9 | 10 | def __eq__(self, other): 11 | """Overrides the default implementation""" 12 | if isinstance(self, other.__class__): 13 | return self.__dict__ == other.__dict__ 14 | return False 15 | 16 | def __cmp__(self, other): 17 | """Overrides the default implementation""" 18 | if isinstance(self, other.__class__): 19 | return self.__dict__ == other.__dict__ 20 | return False 21 | 22 | 23 | @property 24 | def direction(self): 25 | """int: Allow user to query node for its depth""" 26 | return self._direction 27 | 28 | @property 29 | def tx(self): 30 | """string: Allow user to query node for its tx string""" 31 | return self._tx 32 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # Merkle Proof and Verification 2 | 3 | ## Introduction 4 | 5 | Merkle trees are hash-based data structures used to prove the integrity of transaction data stored in the block. **For this exercise you may assume that all trees are binary, balanced, and that the number of transactions to be stored are some exponent of two.** 6 | 7 | ![Merkle Tree](img/merkle_tree.jpeg "Merkle Tree") 8 | _Source: [Grid+](https://blog.gridplus.io/efficiently-bridging-evm-blockchains-8421504e9ced)_ 9 | 10 | Above you can see what this tree would look like. The eight transactions in the block (A-H) are lined up in the bottom row. The second row contains four hashes (S(X) = sha3 hash) of the child transactions. The third row contains hashes of the child hashes, and the root contains a hash of the hashes of the hashes of the transactions. Generically, this is how the transaction part of an Ethereum block is laid out and the root here is what we know of as a transaction header (one of the 15 pieces of information that goes into the block header). 11 | 12 | ## The Problem 13 | 14 | The reason we use Merkle trees to store block data (i.e. transactions) is that verification is very efficient. This verification is called a Merkle proof. 15 | 16 | Suppose we want to prove that transaction C was indeed in the block that formed the header shown above. 17 | 18 | ![Merkle Proof](img/merkle_proof.jpeg "Merkle Proof") 19 | _Source: [Grid+](https://blog.gridplus.io/efficiently-bridging-evm-blockchains-8421504e9ced)_ 20 | 21 | In addition to the transaction hash C , we also need D, S(A,B), and S(S(E,F),S(G,H)) to form the proof. The verification itself performs the following steps on the proof: 22 | 23 | * Hash C and D to produce S(C,D). 24 | * Hash S(A,B) and S(C,D) to produce S(S(A,B),S(C,D)). 25 | * Hash S(S(A,B),S(C,D)) and S(S(E,F),S(G,H)) to produce the root. 26 | * Check that the root is the same as what has been stored previously. 27 | 28 | The efficiency here is that we proved a transaction belonged in a block with only 3 accompanying pieces of information (instead of the 7 other transactions that were stored in the block). This efficiency becomes exponentially more pronounced with larger trees. 29 | 30 | It will be your job to implement the Merkle proof functionality and to verify it. 31 | 32 | ### A Note on Project Layout 33 | 34 | You need only to make changes to `merkle_proof.py`, please make sure to familiarize yourselves with the layout of the entire directory however. The roles of the other files are as follows: 35 | 36 | * `hash_data_structures.py`: contains intermediary object classes that act as nodes in the tree. 37 | * `merkle_tree.py`: contains the Merkle Tree implementation; make sure you familiarize yourself with how it works. It takes in a list of transactions as inputs. Transactions are generally in the form of strings in this case, for brevity. 38 | * `test.py`: a script to help run all the other tests. 39 | * `test_*.py`: testing scripts built on the `unittest` module. Each test ensures the integrity of their respective file. 40 | * `test_sanity.py`: tests that highlight the basic use of each data structure. Take a look at this file if you find yourself lost on how they are implemented. 41 | * `utils.py`: contains multiple helpful methods. 42 | 43 | ## Running Tests 44 | 45 | Use `python3 test.py` to run all the tests. You may add your own tests and run them with `python [filename].py`. Make sure to format tests according to the `unittest` module (take a look at the other files to see examples). 46 | 47 | Writing your own tests is not required but highly recommended. All that matters is you pass our internal test cases. 48 | 49 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import unittest 3 | 4 | test_files = glob.glob('test_*.py') 5 | module_strings = [test_file[0:len(test_file)-3] for test_file in test_files] 6 | suites = [unittest.defaultTestLoader.loadTestsFromName(test_file) 7 | for test_file in module_strings] 8 | test_suite = unittest.TestSuite(suites) 9 | test_runner = unittest.TextTestRunner().run(test_suite) 10 | -------------------------------------------------------------------------------- /test_hash_data_structures.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from hash_data_structures import * 3 | 4 | 5 | class TestHashDataStructures(unittest.TestCase): 6 | 7 | def test_incompatible_hash_functions(self): 8 | """Ensure that the hash function used across the tree 9 | is the same 10 | """ 11 | tx1 = 'a' 12 | tx2 = 'b' 13 | tx3 = 'c' 14 | tx4 = 'd' 15 | 16 | hash_leaf1 = HashLeaf(tx1, tx2, 'sha1') 17 | hash_leaf2 = HashLeaf(tx3, tx4, 'sha224') 18 | 19 | self.assertRaises(AssertionError, HashNode, hash_leaf1, hash_leaf2, 20 | 'sha256') 21 | 22 | if __name__ == '__main__': 23 | unittest.main() 24 | -------------------------------------------------------------------------------- /test_merkle_proof.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from merkle_proof import * 3 | from merkle_tree import * 4 | from node import Node 5 | 6 | 7 | class TestMerkleProof(unittest.TestCase): 8 | 9 | def test_one_proof(self): 10 | """Test that the proof can handle a tree with only one transaction. 11 | No other data is necessary to arrive at the block header 12 | """ 13 | tx1 = 'a' 14 | 15 | merkle_tree = MerkleTree([tx1]) 16 | 17 | self.assertEqual([], merkle_proof(tx1, merkle_tree)) 18 | 19 | def test_small_proof(self): 20 | """Test that the proof can handle a tree with only two transactions""" 21 | tx1 = 'a' 22 | tx2 = 'b' 23 | 24 | merkle_tree = MerkleTree([tx1, tx2]) 25 | 26 | self.assertEqual([Node('r', tx2)], merkle_proof(tx1, merkle_tree)) 27 | self.assertEqual([Node('l', tx1)], merkle_proof(tx2, merkle_tree)) 28 | 29 | def test_medium_proof(self): 30 | """Test that the proof can handle a tree with up to four 31 | transactions 32 | """ 33 | tx1 = 'a' 34 | tx2 = 'b' 35 | tx3 = 'c' 36 | tx4 = 'd' 37 | 38 | merkle_tree = MerkleTree([tx1, tx2, tx3, tx4]) 39 | 40 | data = tx1 + tx2 41 | data = hash_data(data, 'sha256') 42 | 43 | self.assertEqual([Node('l', data), Node('l', tx3)], merkle_proof(tx4, merkle_tree)) 44 | self.assertEqual([Node('l', data), Node('r', tx4)], merkle_proof(tx3, merkle_tree)) 45 | 46 | data = tx3 + tx4 47 | data = hash_data(data, 'sha256') 48 | 49 | self.assertEqual([Node('r', data), Node('l', tx1)], merkle_proof(tx2, merkle_tree)) 50 | self.assertEqual([Node('r', data), Node('r', tx2)], merkle_proof(tx1, merkle_tree)) 51 | 52 | def test_large_proof(self): 53 | """Test that the proof can handle a tree with up to eight 54 | transaction""" 55 | tx1 = 'a' 56 | tx2 = 'b' 57 | tx3 = 'c' 58 | tx4 = 'd' 59 | tx5 = 'e' 60 | tx6 = 'f' 61 | tx7 = 'g' 62 | tx8 = 'h' 63 | 64 | data1 = tx1 + tx2 65 | data1 = hash_data(data1, 'sha256') 66 | data2 = tx5 + tx6 67 | data2 = hash_data(data2, 'sha256') 68 | data3 = tx7 + tx8 69 | data3 = hash_data(data3, 'sha256') 70 | data4 = data2 + data3 71 | data4 = hash_data(data4, 'sha256') 72 | 73 | merkle_tree = MerkleTree([tx1, tx2, tx3, tx4, tx5, tx6, tx7, tx8]) 74 | self.assertEqual([Node('r', data4), Node('l', data1), Node('r', tx4)], merkle_proof(tx3, merkle_tree)) 75 | 76 | def test_extra_large_proof(self): 77 | """Test that the proof can handle a tree with up to eight 78 | transaction""" 79 | tx1 = 'a' 80 | tx2 = 'b' 81 | tx3 = 'c' 82 | tx4 = 'd' 83 | tx5 = 'e' 84 | tx6 = 'f' 85 | tx7 = 'g' 86 | tx8 = 'h' 87 | tx9 = 'i' 88 | tx10 = 'j' 89 | tx11 = 'k' 90 | tx12 = 'l' 91 | tx13 = 'm' 92 | tx14 = 'n' 93 | tx15 = 'o' 94 | tx16 = 'p' 95 | 96 | data1 = hash_data(tx1 + tx2, 'sha256') 97 | data2 = hash_data(tx3 + tx4, 'sha256') 98 | data3 = hash_data(tx5 + tx6, 'sha256') 99 | data4 = hash_data(tx7 + tx8, 'sha256') 100 | data5 = hash_data(hash_data(data1 + data2, 'sha256') + hash_data(data3 + data4, 'sha256'), 'sha256') 101 | 102 | data6 = hash_data(tx15 + tx16, 'sha256') 103 | data7 = hash_data(tx13 + tx14, 'sha256') 104 | data8 = hash_data(data7 + data6, 'sha256') 105 | data9 = hash_data(tx11 + tx12, 'sha256') 106 | 107 | 108 | merkle_tree = MerkleTree([tx1, tx2, tx3, tx4, tx5, tx6, tx7, tx8, tx9, tx10, tx11, tx12, tx13, tx14, tx15, tx16]) 109 | self.assertEqual([Node('l', data5), Node('r', data8), Node('r', data9), Node('r', tx10)], merkle_proof(tx9, merkle_tree)) 110 | 111 | def test_verify_proof_small(self): 112 | """Test that the proof can be verified; the hash must be reconstructed 113 | exactly right. Issues may come up with the order in which data is 114 | hashed 115 | """ 116 | tx1 = 'a' 117 | tx2 = 'b' 118 | tx3 = 'c' 119 | tx4 = 'd' 120 | 121 | merkle_tree = MerkleTree([tx1, tx2, tx3, tx4]) 122 | proof = merkle_proof(tx1, merkle_tree) 123 | verified_hash = verify_proof(tx1, proof) 124 | 125 | self.assertEqual(verified_hash, merkle_tree.block_header) 126 | 127 | def test_verify_proof_big(self): 128 | """Test that the proof can be verified; the hash must be reconstructed 129 | exactly right. Issues may come up with the order in which data is 130 | hashed 131 | """ 132 | tx1 = 'a' 133 | tx2 = 'b' 134 | tx3 = 'c' 135 | tx4 = 'd' 136 | tx5 = 'e' 137 | tx6 = 'f' 138 | tx7 = 'g' 139 | tx8 = 'h' 140 | tx9 = 'i' 141 | tx10 = 'j' 142 | tx11 = 'k' 143 | tx12 = 'l' 144 | tx13 = 'm' 145 | tx14 = 'n' 146 | tx15 = 'o' 147 | tx16 = 'p' 148 | 149 | merkle_tree = MerkleTree([tx1, tx2, tx3, tx4, tx5, tx6, tx7, tx8, tx9, tx10, tx11, tx12, tx13, tx14, tx15, tx16]) 150 | proof = merkle_proof(tx2, merkle_tree) 151 | verified_hash = verify_proof(tx2, proof) 152 | 153 | self.assertEqual(verified_hash, merkle_tree.block_header) 154 | 155 | if __name__ == '__main__': 156 | unittest.main() 157 | -------------------------------------------------------------------------------- /test_merkle_tree.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from merkle_tree import * 3 | 4 | 5 | class TestMerkleTree(unittest.TestCase): 6 | 7 | def test_one_tx(self): 8 | """Test that the tree can construct the proper block header given 9 | only one transaction 10 | """ 11 | tx1 = 'a' 12 | 13 | data = tx1 + tx1 14 | data = hash_data(data, 'sha256') 15 | 16 | merkle_tree = MerkleTree([tx1]) 17 | 18 | self.assertEqual(merkle_tree.block_header, data) 19 | self.assertEqual(merkle_tree.height, 1) 20 | 21 | def test_two_tx(self): 22 | """Test that the tree can construct the proper block header given 23 | only two transactions 24 | """ 25 | tx1 = 'a' 26 | tx2 = 'b' 27 | 28 | data = tx1 + tx2 29 | data = hash_data(data, 'sha256') 30 | 31 | merkle_tree = MerkleTree([tx1, tx2]) 32 | 33 | self.assertEqual(merkle_tree.block_header, data) 34 | self.assertEqual(merkle_tree.height, 1) 35 | 36 | 37 | def test_less_tx(self): 38 | """Test that the tree can construct the proper block header given 39 | some non-exponent-of-two amount of transactions 40 | """ 41 | tx1 = 'a' 42 | tx2 = 'b' 43 | tx3 = 'c' 44 | 45 | data1 = tx1 + tx2 46 | data1 = hash_data(data1, 'sha1') 47 | data2 = tx3 + tx3 48 | data2 = hash_data(data2, 'sha1') 49 | data = data1 + data2 50 | data = hash_data(data, 'sha1') 51 | 52 | merkle_tree = MerkleTree([tx1, tx2, tx3], 'sha1') 53 | 54 | self.assertEqual(merkle_tree.block_header, data) 55 | self.assertEqual(merkle_tree.height, 2) 56 | 57 | def test_less_tx_again(self): 58 | """Test that the tree can construct the proper block header given 59 | a greater non-exponent-of-two amount of transactions 60 | """ 61 | tx1 = 'a' 62 | tx2 = 'b' 63 | tx3 = 'c' 64 | tx4 = 'd' 65 | tx5 = 'e' 66 | 67 | data1 = tx1 + tx2 68 | data1 = hash_data(data1, 'sha256') 69 | data2 = tx3 + tx4 70 | data2 = hash_data(data2, 'sha256') 71 | data3 = data1 + data2 72 | data3 = hash_data(data3, 'sha256') 73 | data4 = tx5 + tx5 74 | data4 = hash_data(data4, 'sha256') 75 | data5 = tx5 + tx5 76 | data5 = hash_data(data5, 'sha256') 77 | data6 = data4 + data5 78 | data6 = hash_data(data6, 'sha256') 79 | data = data3 + data6 80 | data = hash_data(data, 'sha256') 81 | 82 | merkle_tree = MerkleTree([tx1, tx2, tx3, tx4, tx5]) 83 | 84 | self.assertEqual(merkle_tree.block_header, data) 85 | self.assertEqual(merkle_tree.height, 3) 86 | 87 | def test_reset_tree(self): 88 | """Test that users can wipe the tree successfully""" 89 | tx1 = 'a' 90 | tx2 = 'b' 91 | 92 | merkle_tree = MerkleTree([tx1, tx2]) 93 | self.assertEqual(merkle_tree.height, 1) 94 | 95 | merkle_tree.hash_function('sha1') 96 | 97 | merkle_tree.reset_tree() 98 | self.assertEqual(merkle_tree.height, 0) 99 | 100 | def test_add_tx(self): 101 | """Test that users can add tx's to the tree successfully. It should 102 | be reset and reconstructed from the new list 103 | """ 104 | tx1 = 'a' 105 | tx2 = 'b' 106 | tx3 = 'c' 107 | tx4 = 'd' 108 | 109 | data1 = tx1 + tx2 110 | data1 = hash_data(data1, 'sha256') 111 | data2 = tx3 + tx3 112 | data2 = hash_data(data2, 'sha256') 113 | data = data1 + data2 114 | data = hash_data(data, 'sha256') 115 | 116 | merkle_tree = MerkleTree([tx1, tx2]) 117 | merkle_tree.add_tx(tx3) 118 | 119 | self.assertEqual(merkle_tree.block_header, data) 120 | 121 | data2 = tx3 + tx4 122 | data2 = hash_data(data2, 'sha256') 123 | data = data1 + data2 124 | data = hash_data(data, 'sha256') 125 | 126 | merkle_tree = MerkleTree([tx1, tx2]) 127 | merkle_tree.add_tx(tx3, tx4) 128 | 129 | self.assertEqual(merkle_tree.block_header, data) 130 | 131 | merkle_tree = MerkleTree([tx1, tx2]) 132 | merkle_tree.add_tx([tx3, tx4]) 133 | 134 | self.assertEqual(merkle_tree.block_header, data) 135 | 136 | if __name__ == '__main__': 137 | unittest.main() 138 | -------------------------------------------------------------------------------- /test_sanity.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from merkle_tree import * 3 | 4 | 5 | class SanityCheck(unittest.TestCase): 6 | 7 | def test_sanity_HashLeaf(self): 8 | """Ensure HashLeaf structures work as intended""" 9 | tx1 = 'a' 10 | tx2 = 'b' 11 | 12 | data = tx1 + tx2 13 | data = hash_data(data, 'sha256') 14 | 15 | hash_leaf = HashLeaf(tx1, tx2, 'sha256') 16 | 17 | self.assertEqual(hash_leaf.data, data) 18 | self.assertEqual(hash_leaf.height, 1) 19 | 20 | def test_sanity_HashNode(self): 21 | """Ensure HashNode structures work as intended""" 22 | tx1 = 'a' 23 | tx2 = 'b' 24 | tx3 = 'c' 25 | tx4 = 'd' 26 | 27 | data1 = tx1 + tx2 28 | data1 = hash_data(data1, 'sha256') 29 | data2 = tx3 + tx4 30 | data2 = hash_data(data2, 'sha256') 31 | data = data1 + data2 32 | data = hash_data(data, 'sha256') 33 | 34 | hash_leaf1 = HashLeaf(tx1, tx2, 'sha256') 35 | hash_leaf2 = HashLeaf(tx3, tx4, 'sha256') 36 | hash_node = HashNode(hash_leaf1, hash_leaf2, 'sha256') 37 | 38 | self.assertEqual(hash_node.data, data) 39 | self.assertEqual(hash_node.height, 2) 40 | 41 | def test_sanity_check_MerkleTree(self): 42 | """Ensure MerkleTree structures work as intended""" 43 | tx1 = 'a' 44 | tx2 = 'b' 45 | tx3 = 'c' 46 | tx4 = 'd' 47 | 48 | data1 = tx1 + tx2 49 | data1 = hash_data(data1, 'sha256') 50 | data2 = tx3 + tx4 51 | data2 = hash_data(data2, 'sha256') 52 | data = data1 + data2 53 | data = hash_data(data, 'sha256') 54 | 55 | merkle_tree = MerkleTree([tx1, tx2, tx3, tx4]) 56 | 57 | self.assertEqual(merkle_tree.block_header, data) 58 | 59 | if __name__ == '__main__': 60 | unittest.main() 61 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | 3 | 4 | def is_power_of_two(n): 5 | """Check whether `n` is an exponent of two 6 | 7 | >>> is_power_of_two(0) 8 | False 9 | >>> is_power_of_two(1) 10 | True 11 | >>> is_power_of_two(2) 12 | True 13 | >>> is_power_of_two(3) 14 | False 15 | >>> if_power_of_two(16) 16 | True 17 | """ 18 | return n != 0 and ((n & (n - 1)) == 0) 19 | 20 | def hash_data(data, hash_function='sha256'): 21 | """One-way function, takes various standard algorithm names as 22 | `hash_function` input and uses it to hash string `data`. The default 23 | algorithm is 'sha256'. Even small changes in `data` input cause 24 | significant changes to the output 25 | 26 | >>> example = 'hello' 27 | >>> hash_data(example) 28 | '2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824' 29 | >>> hash_data(example, 'sha1') 30 | 'aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d' 31 | >>> example = 'hello!' 32 | >>> hash_data(example) 33 | 'ce06092fb948d9ffac7d1a376e404b26b7575bcc11ee05a4615fef4fec3a308b' 34 | """ 35 | hash_function = getattr(hashlib, hash_function) 36 | data = data.encode('utf-8') 37 | return hash_function(data).hexdigest() 38 | 39 | def concat_and_hash_list(lst, hash_function='sha256'): 40 | """Helper function for quickly concatenate pairs of values and hash them. 41 | The process is repeated until one value is returned: the final hash. 42 | Assumes that the length of the `lst` is an exponent of two 43 | 44 | >>> concat_and_hash_list(['a', 'b']) 45 | 'fb8e20fc2e4c3f248c60c39bd652f3c1347298bb977b8b4d5903b85055620603' 46 | """ 47 | assert len(lst) >= 2, "No transactions to be hashed" 48 | while len(lst) > 1: 49 | a = lst.pop(0) 50 | b = lst.pop(0) 51 | lst.append(hash_data(a + b, hash_function)) 52 | return lst[0] 53 | --------------------------------------------------------------------------------