-
Notifications
You must be signed in to change notification settings - Fork 0
/
prog.py
132 lines (108 loc) · 4.55 KB
/
prog.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
from hashlib import sha256
import pickle
from MiMC5Sponge import MiMC5Sponge
class IncrementalMerkleTree:
def __init__(self, depth, load_from_file=False, file_path='files/merkle_tree.pkl'):
if load_from_file:
try:
with open(file_path, 'rb') as f:
loaded_tree = pickle.load(f)
self.depth = loaded_tree.depth
self.zero_roots = loaded_tree.zero_roots
self.__tree = loaded_tree.__tree
self.__current_insert_idx = loaded_tree.__current_insert_idx
except FileNotFoundError:
print(f"No such file: '{file_path}'")
return
except Exception as e:
print(f"Error loading tree from file: {e}")
return
else:
self.depth = depth
self.zero_roots = list()
self.__set_zero_roots()
self.__tree = [[] for i in range(depth)]
self.__current_insert_idx = 0
def hash_data(self, data=None):
return int(sha256(data.encode() if data else bytes(0)).hexdigest(), 16)
def hash_lr(self, left, right):
hasher = MiMC5Sponge()
return hasher.hash([left, right])
def __set_zero_roots(self):
t_hash = self.hash_data()
self.zero_roots.append(t_hash)
for i in range(1, self.depth):
t_hash = self.hash_lr(t_hash, t_hash)
self.zero_roots.append(t_hash)
def is_tree_full(self):
return True if self.__current_insert_idx >= 2 ** (self.depth - 1) else False
def insert_leaf(self, data):
if self.is_tree_full():
print("Merkle Tree is Full.")
return
curr_idx = self.__current_insert_idx
curr_hash = self.hash_data(data)
self.__tree[0].append(curr_hash)
for i in range(1, self.depth):
next_idx = curr_idx // 2
if curr_idx % 2 == 0: # even -> left child
l = curr_hash
r = self.zero_roots[i - 1]
else: # odd -> right child
l = self.__tree[i - 1][curr_idx - 1]
r = curr_hash
curr_idx = next_idx
curr_hash = self.hash_lr(l, r)
if len(self.__tree[i]) <= next_idx: self.__tree[i].append(curr_hash)
else: self.__tree[i][curr_idx] = curr_hash
self.__current_insert_idx += 1
def get_root(self):
return self.__tree[-1][0] if len(self.__tree[-1]) == 1 else self.zero_roots[-1]
def get_leaf_index(self, data):
leaf_hash = self.hash_data(data)
return self.__tree[0].index(leaf_hash) if leaf_hash in self.__tree[0] else -1
def get_node(self, level, idx):
if not 0 <= level < self.depth:
print("Invalid Level")
return
if not 0 <= idx < 2 ** (self.depth - level - 1):
print("Invalid Index")
return
node_value = self.__tree[level][idx] if len(self.__tree[level]) > idx else self.zero_roots[level]
return node_value
def get_leaf(self, leaf_idx):
if not 0 <= leaf_idx < 2 ** (self.depth - 1):
print("Invalid Index")
return
return self.__tree[0][leaf_idx] if len(self.__tree[0]) > leaf_idx else self.zero_roots[0]
def get_path(self, leaf_idx):
if not 0 <= leaf_idx < 2 ** (self.depth - 1):
print("Invalid Index")
return
path = []
curr_idx = leaf_idx
for i in range(self.depth - 1):
neigh_side = 0 # 0 -> left, 1 -> right
if curr_idx % 2 == 0: # even -> left child
neigh_idx = curr_idx + 1
neigh_side = 1
else: # odd -> right child
neigh_idx = curr_idx - 1
neigh = self.__tree[i][neigh_idx] if len(self.__tree[i]) > neigh_idx else self.zero_roots[i]
path.append((neigh_side, neigh))
curr_idx = curr_idx // 2 # parent index
return path
def is_tree_member(self, leaf, path):
if len(path) != self.depth - 1:
print("Invalid Path")
return
root = leaf
for side, hs in path:
if side == 1:
root = self.hash_lr(root, hs)
else:
root = self.hash_lr(hs, root)
return root == self.get_root()
def store_tree(self, loc='files/merkle_tree.pkl'):
with open(loc, 'wb') as f:
pickle.dump(self, f)