|
2 | 2 | import struct |
3 | 3 | from abc import abstractmethod, ABC |
4 | 4 | from enum import Enum |
| 5 | +from smtplib import SMTP, SMTPException |
5 | 6 | from typing import ClassVar, Optional |
6 | 7 |
|
7 | 8 |
|
@@ -61,20 +62,41 @@ def prepare_socket_for_tls_handshake(self, sock: socket.socket) -> None: |
61 | 62 | class _SmtpHelper(_OpportunisticTlsHelper): |
62 | 63 | """Perform an SMTP StartTLS negotiation.""" |
63 | 64 |
|
| 65 | + def __init__(self, smtp_ehlo_hostname: str): |
| 66 | + self._smtp_ehlo_hostname = smtp_ehlo_hostname |
| 67 | + |
64 | 68 | def prepare_socket_for_tls_handshake(self, sock: socket.socket) -> None: |
65 | | - # Get the SMTP banner |
66 | | - sock.recv(2048) |
| 69 | + # SMTP parsing has some complicated areas and some unusual but legal |
| 70 | + # server behavior - this code uses Python's smtplib to handle the protocol. |
| 71 | + smtp = SMTP(local_hostname=self._smtp_ehlo_hostname) |
| 72 | + smtp.sock = sock |
67 | 73 |
|
68 | | - # Send a EHLO and wait for the 250 status |
69 | | - sock.send(b"EHLO sslyze.scan\r\n") |
70 | | - data = sock.recv(2048) |
71 | | - if b"250 " not in data: |
72 | | - raise OpportunisticTlsError(f"SMTP EHLO was rejected: {repr(data)}") |
| 74 | + try: |
| 75 | + code, server_reply = smtp.getreply() |
| 76 | + message = server_reply.decode() |
| 77 | + except SMTPException as exc: |
| 78 | + code, message = -1, str(exc) |
| 79 | + if code != 220: |
| 80 | + raise OpportunisticTlsError(f"Unable to find 220 service ready response: {message}") |
73 | 81 |
|
74 | | - # Send a STARTTLS |
75 | | - sock.send(b"STARTTLS\r\n") |
76 | | - if b"220" not in sock.recv(2048): |
77 | | - raise OpportunisticTlsError("SMTP STARTTLS not supported") |
| 82 | + try: |
| 83 | + code, server_reply = smtp.ehlo() |
| 84 | + message = server_reply.decode() |
| 85 | + except SMTPException as exc: |
| 86 | + code, message = -1, str(exc) |
| 87 | + if code != 250: |
| 88 | + raise OpportunisticTlsError(f"SMTP EHLO was rejected: {message}") |
| 89 | + |
| 90 | + if not smtp.has_extn("starttls"): |
| 91 | + raise OpportunisticTlsError(f"Server does not support STARTTLS: {message}") |
| 92 | + |
| 93 | + try: |
| 94 | + code, server_reply = smtp.docmd("STARTTLS") |
| 95 | + message = server_reply.decode() |
| 96 | + except SMTPException as exc: |
| 97 | + code, message = -1, str(exc) |
| 98 | + if code != 220: |
| 99 | + raise OpportunisticTlsError(f"SMTP STARTTLS rejected: {message}") |
78 | 100 |
|
79 | 101 |
|
80 | 102 | class _XmppHelper(_OpportunisticTlsHelper): |
@@ -220,14 +242,16 @@ class _PostgresHelper(_GenericOpportunisticTlsHelper): |
220 | 242 |
|
221 | 243 |
|
222 | 244 | def get_opportunistic_tls_helper( |
223 | | - protocol: ProtocolWithOpportunisticTlsEnum, xmpp_to_hostname: Optional[str] |
| 245 | + protocol: ProtocolWithOpportunisticTlsEnum, xmpp_to_hostname: Optional[str], smtp_ehlo_hostname: str |
224 | 246 | ) -> _OpportunisticTlsHelper: |
225 | 247 | helper_cls = _START_TLS_HELPER_CLASSES[protocol] |
226 | | - if protocol not in [ProtocolWithOpportunisticTlsEnum.XMPP, ProtocolWithOpportunisticTlsEnum.XMPP_SERVER]: |
227 | | - opportunistic_tls_helper = helper_cls() |
228 | | - else: |
| 248 | + if protocol in [ProtocolWithOpportunisticTlsEnum.XMPP, ProtocolWithOpportunisticTlsEnum.XMPP_SERVER]: |
229 | 249 | if xmpp_to_hostname is None: |
230 | 250 | raise ValueError("Received None for xmpp_to_hostname") |
231 | 251 | opportunistic_tls_helper = helper_cls(xmpp_to=xmpp_to_hostname) |
| 252 | + elif protocol == ProtocolWithOpportunisticTlsEnum.SMTP: |
| 253 | + opportunistic_tls_helper = helper_cls(smtp_ehlo_hostname) |
| 254 | + else: |
| 255 | + opportunistic_tls_helper = helper_cls() |
232 | 256 |
|
233 | 257 | return opportunistic_tls_helper |
0 commit comments