From 4ddc3838c0039fbf009600b611257d8506c9824f Mon Sep 17 00:00:00 2001
From: Magnus Ahltorp <map@kth.se>
Date: Fri, 7 Apr 2017 15:41:15 +0200
Subject: Submit certificates on the same http session.

Retry submission when server replies with 429 Too Many Requests.
---
 tools/certtools.py  | 44 ++++++++++++++++++++++++++------------------
 tools/submitcert.py | 20 ++++++++++++++++----
 2 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/tools/certtools.py b/tools/certtools.py
index 84d0bd9..485a8ac 100644
--- a/tools/certtools.py
+++ b/tools/certtools.py
@@ -180,14 +180,18 @@ def unpack_tls_array(packed_data, length_len):
 
 def add_chain(baseurl, submission, session=None):
     try:
-        result = urlpost(baseurl + "ct/v1/add-chain", json.dumps(submission), session=session)
-        if result.status_code == requests.codes.ok:
-            return result.json()
-        else:
-            print "ERROR:", result.status_code, result.text
-            if result.status_code == 400:
-                return None
-            sys.exit(1)
+        while True:
+            result = urlpost(baseurl + "ct/v1/add-chain", json.dumps(submission), session=session)
+            if result.status_code == requests.codes.ok:
+                return result.json()
+            else:
+                print "ERROR:", result.status_code, result.text
+                if result.status_code == 400:
+                    return None
+                if result.status_code == 429:
+                    sleep(1)
+                    continue
+                sys.exit(1)
     except ValueError, e:
         print "==== FAILED REQUEST ===="
         print submission
@@ -198,16 +202,20 @@ def add_chain(baseurl, submission, session=None):
 
 def add_prechain(baseurl, submission, session=None):
     try:
-        result = urlpost(baseurl + "ct/v1/add-pre-chain",
-                         json.dumps(submission), session=session)
-
-        if result.status_code == requests.codes.ok:
-            return result.json()
-        else:
-            print "ERROR:", result.status_code, result.text
-            if result.status_code == 400:
-                return None
-            sys.exit(1)
+        while True:
+            result = urlpost(baseurl + "ct/v1/add-pre-chain",
+                            json.dumps(submission), session=session)
+
+            if result.status_code == requests.codes.ok:
+                return result.json()
+            else:
+                print "ERROR:", result.status_code, result.text
+                if result.status_code == 400:
+                    return None
+                if result.status_code == 429:
+                    sleep(1)
+                    continue
+                sys.exit(1)
     except ValueError, e:
         print "==== FAILED REQUEST ===="
         print submission
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
 
-- 
cgit v1.1