significantly reduced memory usage while building trees

1 files changed, 53 insertions, 23 deletions

diff --git a/tools/josef_auditor.py b/tools/josef_auditor.py
index 4cb0d04..4e60f3c 100755
--- a/tools/josef_auditor.py
+++ b/tools/josef_auditor.py
@@ -31,13 +31,46 @@ parser.add_argument('--audit', action='store_true', help="run lightweight audito
 parser.add_argument('--build-sth', action='store_true', help="get all entries and construct STH")
 
 
-def reduce_leafs_to_root(layer0):
-    if len(layer0) == 0:
-        return [[hashlib.sha256().digest()]]
-    current_layer = layer0
-    while len(current_layer) > 1:
-        current_layer = next_merkle_layer(current_layer)
-    return current_layer
+# def reduce_leafs_to_root(layer0):
+#     if len(layer0) == 0:
+#         return [[hashlib.sha256().digest()]]
+#     current_layer = layer0
+#     while len(current_layer) > 1:
+#         current_layer = next_merkle_layer(current_layer)
+#     return current_layer
+
+def reduce_layer(layer):
+	new_layer = []
+	while len(layer) > 1:
+		e1 = layer.pop(0)
+		e2 = layer.pop(0)
+		new_layer.append(internal_hash((e1,e2)))
+	return new_layer
+
+def reduce_tree(entries, layers):
+	if len(entries) == 0 and layers is []:
+		return [[hashlib.sha256().digest()]]
+  
+	layer_idx = 0
+	layers[layer_idx] += entries
+
+	while len(layers[layer_idx]) > 1:
+		if len(layers) == layer_idx + 1:
+			layers.append([])
+
+		layers[layer_idx + 1] += reduce_layer(layers[layer_idx]) 
+		layer_idx += 1
+	return layers
+
+def reduce_subtree_to_root(layers):
+	while len(layers) > 1:
+		layers[1] += next_merkle_layer(layers[0])
+		del layers[0]
+
+	if len(layers[0]) > 1:
+		return next_merkle_layer(layers[0])
+	return layers[0]
+
 
 # Get STH and verify signature
 def fetch_all_sth():
@@ -81,26 +114,23 @@ def verify_consistency(old, new):
 
 
 def fetch_and_build_tree(old_sth, base_url):
-	print "Getting all entries from " + base_url
-
 	sth = old_sth[base_url]
+	subtree = [[]]
+	idx = 0
 
-	entries = []
-	leafs = []
-
-	while len(leafs) < sth["tree_size"]:
-		pre_size = len(leafs)
-		entries = get_entries(base_url, len(leafs), sth["tree_size"])["entries"]
+	print "Getting all entries from " + base_url
+	while idx < sth["tree_size"]:
+		pre_size = idx
+		entries = get_entries(base_url, idx, sth["tree_size"])["entries"]
 		
+		new_leafs = []
 		for item in entries:
-			leafs.append(get_leaf_hash(base64.b64decode(item["leaf_input"])))
-		print "Got entries " + str(pre_size) + " to " + str(len(leafs)) + " (total leaf size: " + str(asizeof(leafs)/1024/1024) + "MB)"
-
+			new_leafs.append(get_leaf_hash(base64.b64decode(item["leaf_input"])))
+		idx += len(new_leafs)
+		print "Got entries " + str(pre_size) + " to " + str(idx) + " (tree size: " + str(asizeof(subtree)) + " B)"
+		subtree = reduce_tree(new_leafs, subtree)
 
-	# tree = build_merkle_tree(leafs)
-	root = base64.b64encode(reduce_leafs_to_root(leafs)[0])
-	# print "Tree size: " + str(asizeof(tree)/1024/1024) + "MB"
-	# print len(leafs)
+	root = base64.b64encode(reduce_subtree_to_root(subtree)[0])
 
 	if root == sth["sha256_root_hash"]:
 		print "Verifying root hashes...OK."
@@ -115,6 +145,7 @@ def main(args):
 	old_sth = fetch_all_sth()
 
 	if args.build_sth:
+		print "Building trees from entries. This may take a while, go get coffee or something..."
 		fetch_and_build_tree(old_sth, base_urls[2])
 
 	if args.audit:
@@ -122,7 +153,6 @@ def main(args):
 
 		while True:
 			time.sleep(1*60-4)
-			# print time.strftime("%H:%M:%S", time.gmtime()) + " checking for updates..."
 			new_sth = fetch_all_sth()
 			verify_consistency(old_sth, new_sth)
 			old_sth = new_sth