From 36dc32fef38959a6f83ba795d2271028f031a57d Mon Sep 17 00:00:00 2001 From: "Edgar P. Burkhart" Date: Mon, 19 May 2025 12:33:37 +0200 Subject: [PATCH] Refactor error handling in MicromailClient to raise exceptions instead of exiting; improve connection logic and command responses --- micromail.py | 143 ++++++++++++++++++++++++--------------------------- 1 file changed, 68 insertions(+), 75 deletions(-) diff --git a/micromail.py b/micromail.py index 7bd3198..17cec66 100644 --- a/micromail.py +++ b/micromail.py @@ -1,6 +1,6 @@ +import errno import socket import ssl -import sys from binascii import b2a_base64 from datetime import datetime, timedelta, timezone @@ -11,8 +11,7 @@ class MicromailClient: self.port = port if ssl and starttls: - print("Cannot use both SSL and STARTTLS at the same time.") - sys.exit(1) + raise Error("Cannot use both SSL and STARTTLS at the same time.") self.ssl = ssl self.starttls = starttls @@ -21,22 +20,28 @@ class MicromailClient: self.username = None def connect(self): - for addr_info in socket.getaddrinfo( - self.host, self.port, 0, socket.SOCK_STREAM, socket.IPPROTO_TCP - ): + con_errors = [] + try: + ai = socket.getaddrinfo( + self.host, self.port, 0, socket.SOCK_STREAM, socket.IPPROTO_TCP + ) + except OSError as e: + ai = [] + con_errors.append(e) + + for addr_info in ai: af, socktype, proto, _, sa = addr_info - print(f"Trying {sa}") try: self.socket = socket.socket(af, socktype, proto) - except OSError: - print(f"Could not create socket: {sa}") + except OSError as e: + con_errors.append(e) self.socket = None continue try: self.socket.connect(sa) - except OSError: - print(f"Could not connect to {sa}") + except OSError as e: + con_errors.append(e) self.socket.close() self.socket = None continue @@ -44,19 +49,16 @@ class MicromailClient: break if self.socket is None: - print("Could not open socket.") - sys.exit(1) + err = ", ".join([errno.errorcode.get(e.errno, str(e)) for e in con_errors]) + raise Error( + f"Could not connect to server: {err}. Are you connected to the network?" + ) + print("Connected to server.") if self.ssl: self.socket = ssl.wrap_socket(self.socket) - code = self.socket.read(3) - self.socket.readline() - - if code != b"220": - print(f"Error: got code {code} on initial connection") - sys.exit(1) - + self.read_res(code=220) self.features = self.ehlo() if self.starttls and b"STARTTLS" in self.features: @@ -70,72 +72,65 @@ class MicromailClient: if feature.startswith(b"AUTH"): auth_methods = feature[4:].upper().split() break + if auth_methods is None: - print("No authentication methods available") - sys.exit(1) + raise Error( + f"No authentication methods available. Server features: {self.features}" + ) if b"PLAIN" in auth_methods: encoded_auth = b2a_base64(f"\0{username}\0{password}", newline=False) - code, res = self.send_command(b"AUTH PLAIN ", encoded_auth) + self.send_command(b"AUTH PLAIN ", encoded_auth, code=235) elif b"LOGIN" in auth_methods: encoded_username = b2a_base64(username, newline=False) - code, res = self.send_command(b"AUTH LOGIN ", encoded_username) - if code != b"334": - print(f"Error: got code {code} on AUTH LOGIN") - sys.exit(1) - encoded_password = b2a_base64(password, newline=False) - code, res = self.send_command(encoded_password) - if code != b"235": - print(f"Error: got code {code} on AUTH LOGIN") - sys.exit(1) - else: - print(f"Unsupported authentication method: {', '.join(auth_methods)}") - sys.exit(1) + self.send_command(b"AUTH LOGIN ", encoded_username, code=334) - print(f"Authenticated as {username}") + encoded_password = b2a_base64(password, newline=False) + self.send_command(encoded_password, code=235) + else: + raise Error( + f"Unsupported authentication methods: {', '.join(auth_methods)}. Micromail only supports PLAIN and LOGIN." + ) + + print(f"Authenticated as {username}.") def write(self, data): if self.socket is None: - print("No socket to write to") - sys.exit(1) + raise Error("No socket to write to.") n = self.socket.write(data) if n != len(data): print(f"Failure writing data <{data}>: not all bytes written") - def send_command(self, *cmd): - self.write(b" ".join(cmd) + b"\r\n") + def send_command(self, *cmd, code=[]): + cmd = b" ".join(cmd) + self.write(cmd + b"\r\n") + return self.read_res(cmd, code=code) + def read_res(self, cmd="", code=[]): response = [] next = True while next: - code = self.socket.read(3) + r_code = int(self.socket.read(3)) next = self.socket.read(1) == b"-" response.append(self.socket.readline().strip()) - return code, response + if isinstance(code, int): + code = [code] + if code != [] and r_code not in code: + raise SMTPError(cmd, r_code, b" ".join(response).decode()) + return r_code, response def ehlo(self): - code, res = self.send_command(b"EHLO _") - if code != b"250": - print(f"Error: got code {code} on EHLO") - sys.exit(1) + _, res = self.send_command(b"EHLO _", code=250) return res def start_starttls(self): - code, res = self.send_command(b"STARTTLS") - if code != b"220": - print(f"Error: got code {code} on STARTTLS") - sys.exit(1) - - print(res) + self.send_command(b"STARTTLS", code=220) self.socket = ssl.wrap_socket(self.socket) def new_message(self, to, sender=None): - code, _ = self.send_command(b"RSET") - if code != b"250": - print(f"Error: got code {code} on RSET") - sys.exit(1) + self.send_command(b"RSET", code=250) if sender is None: sender = self.username @@ -145,20 +140,12 @@ class MicromailClient: self.sender = sender self.to = to - code, _ = self.send_command(f"MAIL FROM:<{sender}>".encode()) - if code != b"250": - print(f"Error: got code {code} on MAIL FROM:<{sender}>") - sys.exit(1) + self.send_command(f"MAIL FROM:<{sender}>".encode(), code=250) for recipient in to: - code, _ = self.send_command(f"RCPT TO:<{recipient}>".encode()) - if code not in [b"250", b"251"]: - print(f"Error: got code {code} on RCPT TO:<{recipient}>") + self.send_command(f"RCPT TO:<{recipient}>".encode(), code=[250, 251]) - code, _ = self.send_command(b"DATA") - if code != b"354": - print(f"Error: got code {code} on DATA") - sys.exit(1) + self.send_command(b"DATA", code=354) def headers(self, headers={}): date = headers.get("date", datetime.now(tz=timezone(timedelta()))) @@ -188,21 +175,27 @@ class MicromailClient: self.write(content + b"\r\n") def send(self): - code, res = self.send_command(b"\r\n.") - if code != b"250": - print(f"Error: got code {code} on send") - print(res) - sys.exit(1) + self.send_command(b"\r\n.", code=250) print("Message sent successfully") def quit(self): - code, _ = self.send_command(b"QUIT") - if code != b"221": - print(f"Error: got code {code} on QUIT") + self.send_command(b"QUIT", code=221) print("Disconnected from server") self.socket.close() +class Error(Exception): + pass + + +class SMTPError(Error): + def __init__(self, cmd, code, message): + super().__init__(f"{code} {message}") + + self.cmd = cmd + self.code = code + + MONTHS = [ "Jan", "Feb",