Refactor error handling in MicromailClient to raise exceptions instead of exiting; improve connection logic and command responses
This commit is contained in:
parent
f592335c8f
commit
36dc32fef3
GPG key ID:
9833D3C5A25BD227
1 changed files with 68 additions and 75 deletions
143
micromail.py
143
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",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue