Submit certificates on the same http session.submission-fix

Retry submission when server replies with 429 Too Many Requests.

1 files changed, 16 insertions, 4 deletions

diff --git a/tools/submitcert.py b/tools/submitcert.py
index 4e4f3c1..cc483f1 100755
--- a/tools/submitcert.py
+++ b/tools/submitcert.py
@@ -19,12 +19,14 @@ import os
 import signal
 import select
 import zipfile
+import itertools
 
 parser = argparse.ArgumentParser(description='')
 parser.add_argument('baseurl', help="Base URL for CT server")
 parser.add_argument('--store', default=None, metavar="dir", help='Get certificates from directory dir')
 parser.add_argument('--sct-file', default=None, metavar="file", help='Store SCT:s in file')
 parser.add_argument('--parallel', type=int, default=16, metavar="n", help="Number of parallel submits")
+parser.add_argument('--maxcerts', type=int, metavar="n", help="Maximum number of certificates to submit")
 parser.add_argument('--check-sct', action='store_true', help="Check SCT signature")
 parser.add_argument('--pre-warm', action='store_true', help="Wait 3 seconds after first submit")
 parser.add_argument('--publickey', default=None, metavar="file", help='Public key for the CT log')
@@ -49,6 +51,8 @@ else:
 
 sth = get_sth(baseurl)
 
+session = None
+
 def submitcert((certfile, cert)):
     timing = timing_point()
     certchain = get_certs_from_string(cert)
@@ -68,12 +72,12 @@ def submitcert((certfile, cert)):
             cleanedcert = cleanprecert(precert, issuer=issuer)
             signed_entry = pack_precert(cleanedcert, issuer_key_hash)
             leafcert = cleanedcert
-            result = add_prechain(baseurl, {"chain":map(base64.b64encode, [precert] + certchain)})
+            result = add_prechain(baseurl, {"chain":map(base64.b64encode, [precert] + certchain)}, session=session)
         else:
             signed_entry = pack_cert(certchain[0])
             leafcert = certchain[0]
             issuer_key_hash = None
-            result = add_chain(baseurl, {"chain":map(base64.b64encode, certchain)})
+            result = add_chain(baseurl, {"chain":map(base64.b64encode, certchain)}, session=session)
     except SystemExit:
         print "EXIT:", certfile
         select.select([], [], [], 1.0)
@@ -171,7 +175,12 @@ def save_sct(sct, sth, leafcert, issuer_key_hash):
     sctlog.write("\n")
     sctlog.close()
 
-p = Pool(args.parallel, lambda: signal.signal(signal.SIGINT, signal.SIG_IGN))
+def worker_init():
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+    global session
+    session = requests.sessions.Session()
+
+p = Pool(args.parallel, worker_init)
 
 nsubmitted = 0
 lastprinted = 0
@@ -179,9 +188,12 @@ lastprinted = 0
 print "listing certs"
 ncerts = get_ncerts(certfiles)
 
+if args.maxcerts:
+    ncerts = min(ncerts, args.maxcerts)
+
 print ncerts, "certs"
 
-certs = get_all_certificates(certfiles)
+certs = itertools.islice(get_all_certificates(certfiles), ncerts)
 
 errors = 0