aio multi backend smtp proxy.

smtpd-multi-proxy.py

#!/usr/bin/python3

import asyncio
from aiosmtpd.controller import Controller
from aiosmtpd.smtp import SMTP
from pprint import pprint
import email
import smtplib
import string
import random
import socket
import signal
import syslog
import sys

#bind = ('127.0.0.1', 2525)
bind = ('193.7.176.63', 25)
host_primary = (('193.7.176.64', 25), 'mailstore.bawue.net')
host_secondaries = [('193.7.176.38', 25),]
__version__ = '0.1'

debug = True
details = True

REQ = {}

class CustomController(Controller):
     def factory(self):
         server = SMTP(self.handler, enable_SMTPUTF8=self.enable_SMTPUTF8, hostname=host_primary[1])
         server.__ident__ = 'ESMTP MP %s' % (__version__,)
         return server

class CustomHandler():
    def __init__(self):
        self.conn = {}
        self.lid = ''

    def rand_str(self):
        return ''.join([random.choice(string.ascii_letters + string.digits) for n in range(10)])

    def log(self, message, lid=None):
        if lid is None:
            lid = self.lid
        if not debug:
            syslog.syslog('%s: %s' % (lid, message))
        else:
            print('DEBUG:', lid, message)

    def send_ehlo(self, host, port, remote_addr, remote_ehlo):
        try:
          dnsname = socket.gethostbyaddr(remote_addr)[0]
        except socket.herror:
          dnsname = 'unknown'

        host_port = '%s_%s' % (host, port)

        if not self.conn:
            self.conn = {}
        self.conn[host_port] = smtplib.SMTP()
        self.conn[host_port].connect(host, port)
        r = self.conn[host_port].docmd('XCLIENT', 'NAME=%s ADDR=%s' % (dnsname, remote_addr))
        if r[0] not in range(200, 299):
            self.log('Backend %s does not accept our XCLIENT command: %s' % (host, r[1]))
            return (450, b'Backend misconfigured. Please retry later.')
        self.conn[host_port].ehlo(name=remote_ehlo)
        if r[0] not in range(200, 299):
            self.log('Backend %s does not accept our EHLO command: %s' % (host, r[1]))
            return r
        return r

    def send_mail(self, host, port, address):
        host_port = '%s_%s' % (host, port)
        r = self.conn[host_port].mail(address)
        return r

    def send_rcpt(self, host, port, recip):
        host_port = '%s_%s' % (host, port)
        r = self.conn[host_port].rcpt(recip)
        return r

    def send_data(self, host, port, msg):
        host_port = '%s_%s' % (host, port)
        r = self.conn[host_port].data(msg)
        return r

    def send_quit(self, host, port):
        host_port = '%s_%s' % (host, port)
        r = self.conn[host_port].quit()
        return r

    def send_rset(self, host, port):
        host_port = '%s_%s' % (host, port)
        r = self.conn[host_port].rset()
        return r

    def fanout(self, function, *args):
        try:
            host = host_primary[0][0]
            port = host_primary[0][1]
            pr = getattr(self, function)(host, port, *args)
        except Exception as e:
            self.log("Exception from backend: %s:%s: %s" % (host, port, str(sys.exc_info())))
            return ('450 Backend error. Please retry later.')
        if pr[0] not in range(200, 299):
            self.log("Error from backend %s:%i: %s" % (host, port, pr))

        for host, port in host_secondaries:
            try:
                sr = getattr(self, function)(host, port, *args)
                if sr[0] not in range(200, 299):
                    self.log("Error from backend %s:%i: %s" % (host, port, sr))
            except Exception as e:
                self.log("Exception from backend: %s:%s: %s" % (host, port, str(sys.exc_info())))

        try:
            sr
        except UnboundLocalError:
            sr = (0, b'undefined')

        if function in ['send_data'] and details:
            self.log("%s: %s" % (function, pr))
            self.log("%s: %s" % (function, sr))
        elif details:
            self.log("%s: %s: %s" % (function, args, pr))
            self.log("%s: %s: %s" % (function, args, sr))

        if pr[0] != sr[0] and function not in ['send_ehlo']:
            self.log('Unexpected diff between primary and secondary during %s: %i/%s vs. %i/%s' % (function, pr[0], pr[1].decode("utf-8"), sr[0], sr[1].decode("utf-8")))

        return '%i %s' % (pr[0], pr[1].decode("utf-8"))

    @asyncio.coroutine
    def handle_QUIT(self, server, session, envelope):
        self.fanout('send_quit',)
        return '250 Bye'

    @asyncio.coroutine
    def handle_RSET(self, server, session, envelope):
        return self.fanout('send_rset',)

    @asyncio.coroutine
    def handle_DATA(self, server, session, envelope):
        msg = envelope.content
        r = self.fanout('send_data', msg)

        self.msgid = email.message_from_bytes(msg).get('message-id', 'unknown')
        self.log('Delivered mail %s' % (self.msgid,))
        return r

    @asyncio.coroutine
    def handle_MAIL(self, server, session, envelope, address, mail_options):
        r = self.fanout('send_mail', address)
        if r.startswith('2'):
            envelope.mail_from = address
            return r
        else:
          return r

    @asyncio.coroutine
    def handle_RCPT(self, server, session, envelope, address, rcpt_options):
        r = self.fanout('send_rcpt', address)
        if r.startswith('2'):
            envelope.rcpt_tos.append(address)
            return r
        else:
          return r

    @asyncio.coroutine
    def handle_EHLO(self, server, session, envelope, hostname):
        self.lid = self.rand_str()
        self.log('Connection from %s claiming to be %s' % (session.peer[0], hostname))
        r = self.fanout('send_ehlo', session.peer[0], hostname)
        if r.startswith('2'):
          session.host_name = hostname
          return '250 HELP'
        else:
          return r

    @asyncio.coroutine
    def handle_HELO(self, server, session, envelope, hostname):
        self.lid = self.rand_str()
        self.log('Connection from %s claiming to be %s' % (session.peer[0], hostname))
        r = self.fanout('send_ehlo', session.peer[0], hostname)
        if r.startswith('2'):
          session.host_name = hostname
          return '250 {}'.format(server.hostname)
        else:
          return r

    @asyncio.coroutine
    def handle_exception(self, error):
        print(error)
        self.log("Unexpected exception: %s" % (str(sys.exc_info())))
        return '452 Internal server error'

if __name__ == '__main__':
    if not debug:
      syslog.openlog(ident='smtp-multi-proxy', logoption=syslog.LOG_PID, facility=syslog.LOG_MAIL)
    CustomHandler.log(CustomHandler, 'Started proxy with backends %s:%i and %s' % (host_primary[0][0], host_primary[0][1], host_secondaries), '')
    handler = CustomHandler()
    controller = CustomController(handler, hostname=bind[0], port=bind[1])
    controller.start()
    sig = signal.sigwait([signal.SIGINT, signal.SIGQUIT])
    controller.stop()