fetchallcerts.py: Store certificates.

1 files changed, 30 insertions, 25 deletions

diff --git a/tools/fetchallcerts.py b/tools/fetchallcerts.py
index dad5241..2276e68 100644
--- a/tools/fetchallcerts.py
+++ b/tools/fetchallcerts.py
@@ -18,6 +18,8 @@ from certtools import *
 parser = argparse.ArgumentParser(description='')
 parser.add_argument('baseurl', help="Base URL for CT server")
 parser.add_argument('--store', default=None, metavar="dir", help='Store certificates in directory dir')
+parser.add_argument('--start', default=0, metavar="n", type=int, help='Start at index n')
+parser.add_argument('--verify', action='store_true', help='Verify STH')
 args = parser.parse_args()
 
 def extract_original_entry(entry):
@@ -28,14 +30,15 @@ def extract_original_entry(entry):
     return [leaf_cert] + certchain
 
 def get_entries_wrapper(baseurl, start, end):
-    fetched_entries = []
-    while start + len(fetched_entries) < (end + 1):
-        print "fetching from", start + len(fetched_entries)
-        entries = get_entries(baseurl, start + len(fetched_entries), end)["entries"]
+    fetched_entries = 0
+    while start + fetched_entries < (end + 1):
+        print "fetching from", start + fetched_entries
+        entries = get_entries(baseurl, start + fetched_entries, end)["entries"]
         if len(entries) == 0:
             break
-        fetched_entries.extend(entries)
-    return fetched_entries
+        for entry in entries:
+            fetched_entries += 1
+            yield entry
 
 def print_layer(layer):
     for entry in layer:
@@ -48,28 +51,30 @@ root_hash = base64.decodestring(sth["sha256_root_hash"])
 print "tree size", tree_size
 print "root hash", base64.b16encode(root_hash)
 
-entries = get_entries_wrapper(args.baseurl, 0, tree_size - 1)
+entries = get_entries_wrapper(args.baseurl, args.start, tree_size - 1)
 
-print "fetched", len(entries), "entries"
+if args.verify:
+    layer0 = [get_leaf_hash(base64.decodestring(entry["leaf_input"])) for entry in entries]
 
-layer0 = [get_leaf_hash(base64.decodestring(entry["leaf_input"])) for entry in entries]
+    tree = build_merkle_tree(layer0)
 
-tree = build_merkle_tree(layer0)
+    calculated_root_hash = tree[-1][0]
 
-calculated_root_hash = tree[-1][0]
+    print "calculated root hash", base64.b16encode(calculated_root_hash)
 
-print "calculated root hash", base64.b16encode(calculated_root_hash)
+    if calculated_root_hash != root_hash:
+        print "fetched root hash and calculated root hash different, aborting"
+        sys.exit(1)
 
-if calculated_root_hash != root_hash:
-    print "fetched root hash and calculated root hash different, aborting"
-    sys.exit(1)
-
-if args.store:
-    for entry, i in zip(entries, range(0, len(entries))):
-        chain = extract_original_entry(entry)
-        f = open(args.store + "/" + ("%06d" % i), "w")
-        for cert in chain:
-            print >> f, "-----BEGIN CERTIFICATE-----"
-            print >> f, base64.encodestring(cert).rstrip()
-            print >> f, "-----END CERTIFICATE-----"
-            print >> f, ""
+elif args.store:
+    for entry, i in itertools.izip(entries, itertools.count(args.start)):
+        try:
+            chain = extract_original_entry(entry)
+            f = open(args.store + "/" + ("%08d" % i), "w")
+            for cert in chain:
+                print >> f, "-----BEGIN CERTIFICATE-----"
+                print >> f, base64.encodestring(cert).rstrip()
+                print >> f, "-----END CERTIFICATE-----"
+                print >> f, ""
+        except AssertionError:
+            print "error for cert", i