Move get_merkle_hash_64k to certtools.py

1 files changed, 70 insertions, 0 deletions

diff --git a/tools/certtools.py b/tools/certtools.py
index 428d623..6a144c9 100644
--- a/tools/certtools.py
+++ b/tools/certtools.py
@@ -12,6 +12,7 @@ import hashlib
 import ecdsa
 import datetime
 import cStringIO
+import zipfile
 
 publickeys = {
     "https://ct.googleapis.com/pilot/":
@@ -301,3 +302,72 @@ def build_merkle_tree(layer0):
         current_layer = next_merkle_layer(current_layer)
         layers.append(current_layer)
     return layers
+
+def print_inclusion_proof(proof):
+    audit_path = proof[u'audit_path']
+    n = proof[u'leaf_index']
+    level = 0
+    for s in audit_path:
+        entry = base64.b16encode(base64.b64decode(s))
+        n ^= 1
+        print level, n, entry
+        n >>= 1
+        level += 1
+
+def get_one_cert(store, i):
+    filename = i / 10000
+    zf = zipfile.ZipFile("%s/%04d.zip" % (store, i / 10000))
+    cert = zf.read("%08d" % i)
+    zf.close()
+    return cert
+
+def get_hash_from_certfile(cert):
+    for line in cert.split("\n"):
+        if line.startswith("-----"):
+            return None
+        if line.startswith("Leafhash: "):
+            return base64.b16decode(line[len("Leafhash: "):])
+    return None
+
+def get_proof(store, tree_size, n):
+    hash = get_hash_from_certfile(get_one_cert(store, n))
+    return get_proof_by_hash(args.baseurl, hash, tree_size)
+
+def get_certs_from_zipfiles(zipfiles, firstleaf, lastleaf):
+    for i in range(firstleaf, lastleaf + 1):
+        try:
+            yield zipfiles[i / 10000].read("%08d" % i)
+        except KeyError:
+            return
+
+def get_merkle_hash_64k(store, blocknumber, write_to_cache=False):
+    hashfilename = "%s/%04x.64khash" % (store, blocknumber)
+    try:
+        hash = base64.b16decode(open(hashfilename).read())
+        assert len(hash) == 32
+        return ("hash", hash)
+    except IOError:
+        pass
+    firstleaf = blocknumber * 65536
+    lastleaf = firstleaf + 65535
+    firstfile = firstleaf / 10000
+    lastfile = lastleaf / 10000
+    zipfiles = {}
+    for i in range(firstfile, lastfile + 1):
+        try:
+            zipfiles[i] = zipfile.ZipFile("%s/%04d.zip" % (store, i))
+        except IOError:
+            break
+    certs = get_certs_from_zipfiles(zipfiles, firstleaf, lastleaf)
+    layer0 = [get_hash_from_certfile(cert) for cert in certs]
+    tree = build_merkle_tree(layer0)
+    calculated_hash = tree[-1][0]
+    for zf in zipfiles.values():
+        zf.close()
+    if len(layer0) != 65536:
+        return ("incomplete", (len(layer0), calculated_hash))
+    if write_to_cache:
+        f = open(hashfilename, "w")
+        f.write(base64.b16encode(calculated_hash))
+        f.close()
+    return ("hash", calculated_hash)