From d0238a1adb4002703b664980200637157375e177 2024-05-09 16:57:52
From: x <xbr@c3l.lu>
Date: 2024-05-09 16:57:52
Subject: [PATCH] style: refactor mail checking and add TLSDetails class

---

diff --git a/backend/check_domains.py b/backend/check_domains.py
index 0c86db38c8f33791c7507d27b557fb3756a7b6bf..ac915dc2b38ee3e45ea65c740d85be2c8bd7b79b 100644
--- a/backend/check_domains.py
+++ b/backend/check_domains.py
@@ -1,134 +1,70 @@
 #!/usr/bin/env python3
 import json
 import ssl
-import smtplib
-import imaplib
 import socket
+import os
 from rich.console import Console
 from cryptography import x509
-import datetime
-import math
 
-def get_expiry_timestamps(expiry_timestamp: int, now_timestamp: int = datetime.datetime.now().timestamp()) -> (bool, int):
-    seconds_left = expiry_timestamp - now_timestamp
-    days_left = math.floor(seconds_left / 86400)
-    return (seconds_left >= 0,days_left)
-
-def get_validity_days(cert) -> (bool, int):
-    # Get expiry date
-    notAfter = cert['notAfter']
-    notAfter_date = datetime.datetime.strptime(notAfter, '%b %d %H:%M:%S %Y %Z')
-
-    # datetime to UNIX time
-    notAfter_timestamp = notAfter_date.timestamp()
-    expiry = get_expiry_timestamps(notAfter_timestamp)
-    return (expiry[0], abs(expiry[1]))
-
-console = Console()
-
-# Parse the input file
-with open('input.json') as raw_data:
-    input = json.load(raw_data)
-
-console.log("[white]Checking web domains...")
-
-context = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
-
-def web_noconn_expiry_days(web_domain: str) -> int | None:
-    try:
-        pem_cert = ssl.get_server_certificate((web_domain, 443), timeout=5)
-        cert = x509.load_pem_x509_certificate(pem_cert.encode())
-    except Exception as e:
-        console = Console()
-        console.log("Could not grab server cert for", "[orange bold underline]"+web_domain, ":", e, style="orange")
-        return None
-    
-    not_after = cert.not_valid_after.timestamp()
-    return get_expiry_timestamps(not_after)
-
-for web_domain in input["domains"]["web"]:
-    # Initiate TLS connection
-    with context.wrap_socket(socket.socket(), server_hostname=web_domain) as s:
-        try:
-            s.connect((web_domain, 443))
-            cert = s.getpeercert()
-        except ssl.SSLCertVerificationError as e:
-            saved = e
-            if e.verify_code == 10:
-                expiry = web_noconn_expiry_days(web_domain)[1]
-                if(expiry != None):
-                    # TODO: add the TLS expiry stuff here
-                    # possibly a list of domains that have expired
-                    # if its already in here, dont add it again
-                    console.log("[red bold underline]" + web_domain, "expired", abs(expiry), "days ago.", style="red")
-            elif e.verify_code == 23:
-                console.log("[red bold underline]" + web_domain, "was revoked.", style="red")
-            elif e.verify_code == 18:
-                console.log("[red bold underline]" + web_domain, "is self-signed.", style="red")
-            elif e.verify_code == 19:
-                console.log("[red bold underline]" + web_domain, "invalid: root not trusted.", style="red")
-            else:
-                console.log("[red bold underline]" + web_domain, "failed verification:", e.verify_message + ".", style="red")
-            continue
-        except ssl.SSLError as e:
-            console.log("[orange bold underline]" + web_domain, "could not establish a secure connection:", e.reason, style="orange")
-            continue
-        except Exception as e:
-            print(e)
-            continue
-
-    validity = get_validity_days(cert)[1]
-    # Print expiry date
-    console.log("[green bold underline]" + web_domain, "expires in", validity, "days", style="green")
-    # TODO: remove known expired certs
-    # If the cert was expired before, we know that it is now valid
-    # -> remove it from the list of expirjuded certs
-
-def __mail_connection(host, port, verification: bool, initializer, starttls_args, closer) -> (bool, int):
-    connection = initializer(host, port)
-    if verification:
-        connection.starttls(**starttls_args)
-    else:
-        connection.starttls()
-    cert = connection.sock.getpeercert()
-    closer(connection)
-    return get_validity_days(cert)
-
-def __smtp_closing(connection):
-    connection.quit()
-    
-def __smtp_connection(domain, port, verification: bool) -> (bool, int):
-    initializer = smtplib.SMTP
-    starttls_args = {"context": context}
-    return __mail_connection(domain, port, verification, initializer, starttls_args, __smtp_closing)
-
-def __imap_closing(connection):
-    connection.logout()
-
-def __imap_connection(domain, port, verification: bool) -> (bool, int):
-    initializer = imaplib.IMAP4
-    starttls_args = {"ssl_context": context}
-    return __mail_connection(domain, port, verification, initializer, starttls_args, __imap_closing)
-
-def __mail_connect(domain, port, protocol_func):
-    try:
-        expiry = protocol_func(domain, port, True)[1]
-        console.log("[green bold underline]" + domain, "expires in", expiry, "days", style="green")
-    except ssl.SSLCertVerificationError as e:
-        if (e.verify_code == 10):
-            expiry = protocol_func(domain, port, False)[1]
-            console.log("[red bold underline]" + web_domain, "expired", abs(expiry), "days ago.", style="red")
-        else:
-            console.log("[red bold underline]" + domain, "failed verification:", e.verify_message + ".", style="red")
-
-def smtp_connect(domain, port):
-    __mail_connect(domain, port, __smtp_connection)
-
-def imap_connect(domain, port):
-    __mail_connect(domain, port, __imap_connection)
-
-for smtp_entry in input["domains"]["smtp"]:
-    smtp_connect(smtp_entry["host"], smtp_entry["port"])
-
-for imap_entry in input["domains"]["imap"]:
-    imap_connect(imap_entry["host"], imap_entry["port"])
+import web
+from mail import *
+import tls_utils
+
+if __name__ == "__main__":
+    console = Console()
+
+    # Parse the input file
+    path = os.path.split(__file__)[0] + "/"
+    with open(path + 'input.json') as raw_data:
+        input = json.load(raw_data)
+
+    context = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
+
+    console.log("[white]Checking web domains...")
+
+    for web_domain in input["domains"]["web"]:
+        # Initiate TLS connection
+        with context.wrap_socket(socket.socket(), server_hostname=web_domain) as s:
+            try:
+                s.connect((web_domain, 443))
+                cert = s.getpeercert()
+            except ssl.SSLCertVerificationError as e:
+                saved = e
+                if e.verify_code == 10:
+                    expiry = web.web_noconn_expiry_days(web_domain)[1]
+                    if(expiry != None):
+                        # TODO: add the TLS expiry stuff here
+                        # possibly a list of domains that have expired
+                        # if its already in here, dont add it again
+                        console.log("[red bold underline]" + web_domain, "expired", abs(expiry), "days ago.", style="red")
+                elif e.verify_code == 23:
+                    console.log("[red bold underline]" + web_domain, "was revoked.", style="red")
+                elif e.verify_code == 18:
+                    console.log("[red bold underline]" + web_domain, "is self-signed.", style="red")
+                elif e.verify_code == 19:
+                    console.log("[red bold underline]" + web_domain, "invalid: root not trusted.", style="red")
+                else:
+                    console.log("[red bold underline]" + web_domain, "failed verification:", e.verify_message + ".", style="red")
+                continue
+            except ssl.SSLError as e:
+                console.log("[orange bold underline]" + web_domain, "could not establish a secure connection:", e.reason, style="orange")
+                continue
+            except Exception as e:
+                print(e)
+                continue
+
+        validity = tls_utils.get_validity_days(cert)[1]
+        # Print expiry date
+        console.log("[green bold underline]" + web_domain, "expires in", validity, "days", style="green")
+        # TODO: remove known expired certs
+        # If the cert was expired before, we know that it is now valid
+        # -> remove it from the list of expirjuded certs
+
+    mail = MailVerificator(context)
+    for smtp_entry in input["domains"]["smtp"]:
+        result = mail.connect(smtp_entry["host"], smtp_entry["port"], "smtp")
+        result.print(console)
+
+    for imap_entry in input["domains"]["imap"]:
+        result = mail.connect(imap_entry["host"], imap_entry["port"], "imap")
+        result.print(console)
diff --git a/backend/mail.py b/backend/mail.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8ecd52b4f00555c4eaba8ef93b43630b4755b55
--- /dev/null
+++ b/backend/mail.py
@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+import ssl
+import smtplib
+import imaplib
+from rich.console import Console
+from cryptography import x509
+import tls_utils
+from tls_utils import TLSDetails
+from abc import ABC, abstractmethod
+
+class MailHandler(ABC):
+    def __init__(self, host: str, port: int, context: ssl.SSLContext):
+        self.host = host
+        self.port = port
+        self.context = context
+    
+    def connect(self, verification: bool) -> tuple[bool, int]:
+        connection = self.protocol_init(self.host, self.port)
+        if verification:
+            connection.starttls(**self.protocol_starttls_args())
+        else:
+            connection.starttls()
+        cert = connection.sock.getpeercert()
+        self.protocol_close(connection)
+        return tls_utils.get_validity_days(cert)
+
+    @abstractmethod
+    def protocol_init(self, host, port):
+        raise NotImplementedError()
+    @abstractmethod
+    def protocol_close(self, connection):
+        raise NotImplementedError()
+    @abstractmethod
+    def protocol_starttls_args(self):
+        raise NotImplementedError()
+    
+    @staticmethod
+    def create_handler(protocol: str):
+        if protocol == "smtp":
+            return SMTPHandler
+        elif protocol == "imap":
+            return IMAPHandler
+        else:
+            raise ValueError("Invalid protocol")
+
+class IMAPHandler(MailHandler):
+    def protocol_init(self, host, port):
+        return imaplib.IMAP4(host, port)
+    def protocol_close(self, connection):
+        connection.logout()
+    def protocol_starttls_args(self):
+        return {"ssl_context": self.context}
+
+class SMTPHandler(MailHandler):
+    def protocol_init(self, host, port):
+        return smtplib.SMTP(host, port)
+    def protocol_close(self, connection):
+        connection.quit()
+    def protocol_starttls_args(self):
+        return {"context": self.context}
+
+class MailVerificator:
+    def __init__(self, context: ssl.SSLContext):
+        self.context = context
+
+    def connect(self, domain: str, port: int, protocol: str) -> TLSDetails:
+        mail = MailHandler.create_handler(protocol)(domain, port, self.context)
+        try:
+            expiry = mail.connect(True)[1]
+            return TLSDetails(domain_name=domain, expires_in_days=expiry)
+        except ssl.SSLCertVerificationError as e:
+            if (e.verify_code == 10):
+                expiry = mail.connect(False)[1]
+                return TLSDetails(domain_name=domain, expires_in_days=expiry)
+            else:
+                error = "failed verification:", e.verify_message + "."
+                return TLSDetails(domain_name=domain, error_message=error)
\ No newline at end of file
diff --git a/backend/tls_utils.py b/backend/tls_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e6ae78ef00284f356c46dac20096e38151e58d4
--- /dev/null
+++ b/backend/tls_utils.py
@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+from rich.console import Console
+import datetime
+import math
+
+class TLSDetails:
+    domain_name = None
+    expires_in_days = None
+    error_message = None
+
+    def __init__(self, domain_name : str = None, expires_in_days : str = None, error_message : str = None):
+        self.domain_name = domain_name
+        self.expires_in_days = expires_in_days
+        self.error_message = error_message
+    
+    def print(self, console: Console):
+        if self.error_message != None:
+            console.log("[red bold underline]" + self.domain_name, self.error_message, style="red")
+        elif self.expires_in_days < 0:
+            console.log("[red bold underline]" + self.domain_name, "expired", abs(self.expires_in_days), "days ago.", style="red")
+        else:
+            console.log("[green bold underline]" + self.domain_name, "expires in", self.expires_in_days, "days", style="green")
+
+def get_expiry_timestamps(expiry_timestamp: int, now_timestamp: int = datetime.datetime.now().timestamp()) -> tuple[bool, int]:
+    seconds_left = expiry_timestamp - now_timestamp
+    days_left = math.floor(seconds_left / 86400)
+    return (seconds_left >= 0, days_left)
+
+def get_validity_days(cert) -> tuple[bool, int]:
+    # Get expiry date
+    notAfter = cert['notAfter']
+    notAfter_date = datetime.datetime.strptime(notAfter, '%b %d %H:%M:%S %Y %Z')
+
+    # datetime to UNIX time
+    notAfter_timestamp = notAfter_date.timestamp()
+    expiry = get_expiry_timestamps(notAfter_timestamp)
+    return (expiry[0], abs(expiry[1]))
\ No newline at end of file
diff --git a/backend/web.py b/backend/web.py
new file mode 100644
index 0000000000000000000000000000000000000000..60b3c3cce11a587cb1a3b0b44158a0e7e161cdca
--- /dev/null
+++ b/backend/web.py
@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+import ssl
+from rich.console import Console
+from cryptography import x509
+
+import tls_utils as tls_utils
+
+def web_noconn_expiry_days(web_domain: str) -> int | None:
+    try:
+        pem_cert = ssl.get_server_certificate((web_domain, 443), timeout=5)
+        cert = x509.load_pem_x509_certificate(pem_cert.encode())
+    except Exception as e:
+        console = Console()
+        console.log("Could not grab server cert for", "[orange bold underline]"+web_domain, ":", e, style="orange")
+        return None
+    
+    not_after = cert.not_valid_after.timestamp()
+    return tls_utils.get_expiry_timestamps(not_after)
\ No newline at end of file