From 10ff26b17f2463e11f7ed42ad6b93400ac830367 Mon Sep 17 00:00:00 2001
From: Guilhem MARION <gmarion@netc.fr>
Date: Tue, 11 Oct 2022 06:55:34 +1030
Subject: [PATCH] Perform handshakes asynchronously

---
 certo/__main__.py        | 21 +++++++++++++++------
 certo/checks/hostname.py | 30 ++++++++++++++++++++++++++----
 certo/report.py          |  6 +++---
 pyproject.toml           |  2 +-
 4 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/certo/__main__.py b/certo/__main__.py
index e25814c..e0ff646 100644
--- a/certo/__main__.py
+++ b/certo/__main__.py
@@ -1,6 +1,6 @@
 """
 Usage:
-    certo [-vj] <hostnames>... [-d DAYS|--days-to-expiration=DAYS] [-t SECONDS|--timeout=SECONDS]
+    certo [-vj] [-d DAYS|--days-to-expiration=DAYS] [-t SECONDS|--timeout=SECONDS] <hostnames>...
     certo -h | --help
 
 Options:
@@ -10,6 +10,7 @@ Options:
    -d DAYS --days-to-expiration=DAYS  Warn about near expiration if within DAYS of the cert's notAfter [default: 5].
    -t SECONDS --timeout=SECONDS       Timeout for SSL Handshake [default: 5].
 """
+import asyncio
 import logging
 
 from docopt import docopt
@@ -17,7 +18,8 @@ from docopt import docopt
 from certo.checks.hostname import check_host_certificate_expiration
 from certo.report import JSONReporter, DefaultReporter
 
-if __name__ == "__main__":
+
+async def main():
     args = docopt(__doc__)
 
     output_as_json = args.get("-j")
@@ -32,12 +34,19 @@ if __name__ == "__main__":
     else:
         reporter = DefaultReporter()
 
-    # @todo async
-    for hs in hostnames:
-        logging.info(f"Getting CERT from {hs}")
-        reporter.add_check(check_host_certificate_expiration(hs, days_to_expiration))
+    jobs = {
+        check_host_certificate_expiration(hs, days_to_expiration) for hs in hostnames
+    }
+    checks = await asyncio.gather(*jobs)
+
+    for check in checks:
+        reporter.append(check)
 
     if log := reporter.report():
         print(log)
 
     exit(reporter.num_failed())
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
diff --git a/certo/checks/hostname.py b/certo/checks/hostname.py
index 7f59bd4..025b432 100644
--- a/certo/checks/hostname.py
+++ b/certo/checks/hostname.py
@@ -1,4 +1,6 @@
+import asyncio
 import datetime
+import logging
 from collections import namedtuple
 
 from dateutil.parser import parse as dtparse
@@ -11,18 +13,38 @@ CertCheckResult = namedtuple(
 )
 
 
-def get_cert(hostname, timeout):
+# Unit of time slept asynchronously to simulate async socket handling
+AWAIT_IOTA = 0.001
+
+
+async def get_cert(hostname, timeout):
     ctx = ssl.create_default_context()
-    with ctx.wrap_socket(socket.socket(), server_hostname=hostname) as s:
+    with ctx.wrap_socket(
+        socket.socket(), server_hostname=hostname, do_handshake_on_connect=False
+    ) as s:
         s.settimeout(timeout)
 
+        # @todo simulate async connect
         s.connect((hostname, 443))
+
+        s.setblocking(False)
+        # Cannot await the handshake: simulate it with asyncio sleep
+        while "Handshake not finished":
+            try:
+                s.do_handshake()
+                break
+            except ssl.SSLWantReadError:
+                await asyncio.sleep(AWAIT_IOTA)
+            except ssl.SSLWantWriteError:
+                await asyncio.sleep(AWAIT_IOTA)
+
         return s.getpeercert()
 
 
-def check_host_certificate_expiration(hostname, days_to_expiration, timeout=5):
+async def check_host_certificate_expiration(hostname, days_to_expiration, timeout=5):
+    logging.info(f"Getting CERT from {hostname}")
     try:
-        cert = get_cert(hostname, timeout)
+        cert = await get_cert(hostname, timeout)
     except ssl.SSLCertVerificationError as e:
         return CertCheckResult(hostname, False, None, e.strerror)
 
diff --git a/certo/report.py b/certo/report.py
index 51387ee..f91b1cf 100644
--- a/certo/report.py
+++ b/certo/report.py
@@ -5,7 +5,7 @@ class CheckReporter:
     def __init__(self):
         self.checks = list()
 
-    def add_check(self, check):
+    def append(self, check):
         self.checks.append(check)
 
     def failed(self):
@@ -43,8 +43,8 @@ class JSONReporter(CheckReporter):
 
 
 class DefaultReporter(CheckReporter):
-    def add_check(self, check):
-        super().add_check(check)
+    def append(self, check):
+        super().append(check)
         result = f"[{'PASS' if check.check_successful else 'FAIL'}] Check host {check.hostname}"
         if check.debug:
             result += f" - {check.debug}"
diff --git a/pyproject.toml b/pyproject.toml
index 87d579a..a35d94c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "certo"
 version = "0.1.0"
-description = ""
+description = "A certificate expiration checker and reminder"
 authors = ["Guilhem MARION <gmarion@netc.fr>"]
 
 [tool.poetry.dependencies]