it-infra/m-labs-intl/smtp_router.py

279 lines
9.2 KiB
Python

import asyncio
import email
import logging
from smtplib import SMTP as SMTPClient
from typing import Dict
from aiosmtpd.proxy_protocol import ProxyData
from python_socks.sync import Proxy
import os
import ssl
from aiosmtpd.controller import Controller
from aiosmtpd.smtp import Envelope, Session, SMTP
from email.message import Message
import dkim
import spf
from functools import lru_cache
import dns.resolver
PROXY = Proxy.from_url('socks5://5.78.86.156:2025')
END_HOST = "localhost"
END_PORT = 25
PROXY_SMTP_HOST = "*"
PROXY_SMTP_PORT = 2025
END_SMTP_HOST = "*"
END_SMTP_PORT = 3025
GENERIC_SMTP_HOST = "*"
GENERIC_SMTP_PORT = 25
log = logging.getLogger("smtphandler")
VALID_PROXY_ADDRS = {"5.78.86.156", "2a01:4ff:1f0:83de::1", "m-labs-intl.com", "mail.m-labs-intl.com"}
VALID_END_ADDRS = {"localhost", "127.0.0.1"}
VALID_RECEPIENTS = {"m-labs.hk", "m-labs.ph", "m-labs-intl.com", "193thz.com", "malloctech.fr"}
@lru_cache(maxsize=256)
def get_mx(domain):
records = dns.resolver.resolve(domain, "MX")
if not records:
return None
result = max(records, key=lambda r: r.preference)
return str(result.exchange)
class ProxiedSMTP(SMTPClient):
def _get_socket(self, host, port, timeout):
return PROXY.connect(dest_host=host, dest_port=port, timeout=timeout)
class GenericMailServer:
"""
Accepts mail from everywhere, and passes it to the real SMTP server, after proper SPF and DKIM checks.
"""
async def handle_MAIL(self,
server: SMTP,
session: Session,
envelope: Envelope,
address: str,
mail_options: list):
ip = session.peer[0]
result, description = spf.check2(ip, address, session.host_name)
valid_spf = result == 'pass'
envelope.spf = valid_spf
log.info("SPF: %s, %s", result, description)
if not valid_spf:
return '550 SPF validation failed'
envelope.mail_from = address
envelope.mail_options.extend(mail_options)
return '250 OK'
async def handle_RCPT(self,
server: SMTP,
session: Session,
envelope: Envelope,
address: str,
rcpt_options: list):
if address.split("@")[1] not in VALID_RECEPIENTS:
return '550 not relaying to that domain'
log.debug("Handle RCPT for %s", address)
envelope.rcpt_tos.append(address)
return '250 OK'
async def handle_DATA(self,
server: SMTP,
session: Session,
envelope: Envelope):
valid_dkim = dkim.verify(envelope.content)
envelope.dkim = valid_dkim
log.info("DKIM: %s", valid_dkim)
message: Message = email.message_from_bytes(envelope.content)
if not valid_dkim:
return '550 DKIM validation failed'
log.info('Message: %s', message)
try:
with SMTPClient(END_HOST, END_PORT) as client:
client.sendmail(
from_addr=envelope.mail_from,
to_addrs=envelope.rcpt_tos,
msg=envelope.original_content
)
except BaseException as e:
print(e)
return '500 Could not process your message'
return '250 Message accepted for delivery'
class EndMailServer:
"""
Accepts mail only from end server. Checks if mail is signed to be from proxy, and sends to proxy if needed.
SPF and DKIM needs to be already included.
"""
async def handle_DATA(self,
server: SMTP,
session: Session,
envelope: Envelope):
ip = session.peer[0]
if ip not in VALID_END_ADDRS:
return "521 Server doesn't accept mail"
message: Message = email.message_from_bytes(envelope.content)
log.info('Message: %s', message)
mx_rcpt: Dict[str, list[str]] = {}
for rcpt in envelope.rcpt_tos:
_, _, domain = rcpt.partition("@")
mx = get_mx(domain)
if mx is None:
continue
mx_rcpt.setdefault(mx, []).append(rcpt)
try:
for mx, rcpts in mx_rcpt.items():
if envelope.mail_from in VALID_PROXY_ADDRS:
with ProxiedSMTP(mx, 25) as client:
client.sendmail(
from_addr=envelope.mail_from,
to_addrs=rcpts,
msg=envelope.original_content
)
else:
with SMTPClient(mx, 25) as client:
client.sendmail(
from_addr=envelope.mail_from,
to_addrs=rcpts,
msg=envelope.original_content
)
except BaseException as e:
print(e)
return '500 Could not process your message'
return '250 Message accepted for delivery'
class ProxyMailServer:
"""
Accepts mail only from proxy server, and passes it to the real SMTP server, after proper SPF and DKIM checks.
"""
async def handle_PROXY(self, server: SMTP, session: Session, envelope: Envelope, proxy_data: ProxyData):
ip = session.peer[0]
envelope.proxy_data = proxy_data
return ip in VALID_PROXY_ADDRS
async def handle_MAIL(self,
server: SMTP,
session: Session,
envelope: Envelope,
address: str,
mail_options: list):
ip = envelope.proxy_data.src_addr
result, description = spf.check2(ip, address, session.host_name)
valid_spf = result == 'pass'
envelope.spf = valid_spf
log.info("SPF: %s, %s", result, description)
if not valid_spf:
return '550 SPF validation failed'
envelope.mail_from = address
envelope.mail_options.extend(mail_options)
return '250 OK'
async def handle_RCPT(self,
server: SMTP,
session: Session,
envelope: Envelope,
address: str,
rcpt_options: list):
if address.split("@")[1] not in VALID_RECEPIENTS:
return '550 not relaying to that domain'
log.debug("Handle RCPT for %s", address)
envelope.rcpt_tos.append(address)
return '250 OK'
async def handle_DATA(self,
server: SMTP,
session: Session,
envelope: Envelope):
valid_dkim = dkim.verify(envelope.content)
envelope.dkim = valid_dkim
log.info("DKIM: %s", valid_dkim)
message: Message = email.message_from_bytes(envelope.content)
if not valid_dkim:
return '550 DKIM validation failed'
log.info('Message: %s', message)
try:
with SMTPClient(END_HOST, END_PORT) as client:
client.sendmail(
from_addr=envelope.mail_from,
to_addrs=envelope.rcpt_tos,
msg=envelope.original_content
)
except BaseException as e:
print(e)
return '500 Could not process your message'
return '250 Message accepted for delivery'
if __name__ == '__main__':
host = os.getenv('SMTP_HOST', '*')
port = int(os.getenv('SMTP_PORT', '25'))
accept_host = os.getenv('ACCEPT_HOST')
ssl_keys = os.getenv('SSL_KEYS')
loop = asyncio.get_event_loop()
end_handler = EndMailServer()
proxy_handler = ProxyMailServer()
generic_handler = GenericMailServer()
ssl_context = None
if ssl_keys:
ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
ssl_context.load_cert_chain(ssl_keys + '.crt', ssl_keys + '.key')
generic_controller = Controller(GenericMailServer, hostname=GENERIC_SMTP_HOST, port=GENERIC_SMTP_PORT)
generic_controller.factory = lambda: SMTP(GenericMailServer, enable_SMTPUTF8=True, tls_context=ssl_context)
generic_controller.start()
log.info("Generic SMTP server started on %s:%s", GENERIC_SMTP_HOST, GENERIC_SMTP_PORT)
end_controller = Controller(EndMailServer, hostname=END_SMTP_HOST, port=END_SMTP_PORT)
end_controller.factory = lambda: SMTP(EndMailServer, enable_SMTPUTF8=True)
end_controller.start()
log.info("End SMTP server started on %s:%s", END_SMTP_HOST, END_SMTP_PORT)
proxy_controller = Controller(ProxyMailServer, hostname=PROXY_SMTP_HOST, port=PROXY_SMTP_PORT)
proxy_controller.factory = lambda: SMTP(ProxyMailServer, enable_SMTPUTF8=True, tls_context=ssl_context)
proxy_controller.start()
log.info("Proxy SMTP server started on %s:%s", PROXY_SMTP_HOST, PROXY_SMTP_PORT)
try:
loop.run_forever()
except KeyboardInterrupt:
print("Shutting down")
finally:
generic_controller.stop()
end_controller.stop()
proxy_controller.stop()
loop.stop()
loop.close()