# Neil Schemenauer <nas@arctrix.com>
# 2004/05/24

import sys
import os
import time
import re
import socket
import signal
import email
import tempfile
import SocketServer


# control
ME = "mail.example.com"
IPME = "192.168.1.1"
GREETING = ME
QMAIL_HOME = "/var/qmail"
RCPTHOSTS = open(QMAIL_HOME + "/control/rcpthosts").read().strip().split()
USERS = open(QMAIL_HOME + "/users.txt").read().strip().split()
RELAYCLIENTS = ["127."]
MAXHOPS = 100
USER = 65534
GROUP = 101

# spambayes (with CDB database)
SPAM_CUTOFF = 0.7
DB_FILE = "/var/qmail/spambayes/wordprobs.cdb"
OPTION_FILE = "/var/qmail/spambayes/bayescustomize.ini"
FILTER_ADDRESSES = ['joe', 'bob']

# exe content filter
EXE_PAT = re.compile(r'\.(bat|chm|cmd|com|cpl|exe|hlp|hta|js|jse|pif|scr|sct|'
                      'shs|shb|vb|vbe|vbs|wsc|wsf|wsh|lnk)($|")', re.I)


SOURCE_ROUTE_PAT = re.compile(r"@[^:]*:")

_wrote_newline = 1

def _log_timestamp():
    sys.stdout.write('[%s] %d ' % (time.strftime('%Y-%m-%d:%H:%M:%S'),
                                   os.getpid()))

def log(*args):
    global _wrote_newline
    date = time.strftime('%Y-%m-%d:%H:%M:%S')
    msg = ' '.join(map(str, args))
    sys.stdout.write('[%s] %d %s\n' % (date, os.getpid(), msg))
    sys.stdout.flush()
    _wrote_newline = 1

def logwrite(s):
    global _wrote_newline
    if _wrote_newline:
        _log_timestamp()
        _wrote_newline = 0
    sys.stdout.write(s)
    if s[:-1] == '\n':
        _wrote_newline = 1
        sys.stdout.flush()

#### email parsing #################################################
def parse_msg(msgfile):
    if msgfile.tell() > 2000000:
        return None # too big, let it through
    msgfile.seek(0)
    try:
        return email.message_from_file(msgfile)
    except:
        log('message_from_file failed')
    return None
        
#### spambayes #####################################################
os.environ['BAYESCUSTOMIZE'] = OPTION_FILE
from spambayes.cdb_classifier import CdbClassifier
from spambayes.tokenizer import tokenize

_bayes = CdbClassifier(open(DB_FILE, 'rb'))

def is_spam(msg, rcptto):
    for addr in rcptto:
        user, host = addr.split('@', 1)
        if user not in FILTER_ADDRESSES:
            return 0
    prob = _bayes.spamprob(tokenize(msg))
    log('spambayes score', ' '.join(rcptto), prob)
    return prob > SPAM_CUTOFF

####################################################################
# Executable content
def has_exe(msg):
    for part in msg.walk():
        attachname = (part.get_param('name') or
                      part.get_param('filename',
                                     header='content-disposition'))
        if attachname and EXE_PAT.search(attachname):
            log('rejecting attachment %r' % attachname)
            return 1
    return 0
    
####################################################################

class AddressSyntaxError(Exception):
    pass

def addrparse(arg):
    terminator = ">"
    i = arg.find("<")
    if i >= 0:
        arg = arg[i+1:]
    else:
        # invalid but try to parse anyhow
        terminator = " "
        i = arg.find(':')
        if i >= 0:
            arg = arg[i+1:]
        arg = arg.lstrip()

    # strip source route
    m = SOURCE_ROUTE_PAT.match(arg)
    if m:
        arg = arg[m.end():]

    flagesc = 0
    flagquoted = 0
    addr = ""
    for i in range(len(arg)):
        ch = arg[i]
        if flagesc:
            addr += ch
        else:
            if not flagquoted and ch == terminator:
                break
            if ch == "\\":
                flagesc = 1
            elif ch == '"':
                flagquoted = not flagquoted
            else:
                addr += ch

    if len(addr) > 900:
        raise AddressSyntaxError
    return addr

def userallowed(addr):
    i = addr.rfind('@')
    if i == -1:
        user = addr
    else:
        user = addr[:i]
    user = user.lower()
    for pat in USERS:
        pat = pat.lower()
        if user == pat or user.startswith(pat + "-"):
            return 1
    return 0
    
def hostallowed(addr):
    i = addr.rfind('@')
    if i == -1:
        return 1
    host = addr[i+1:]
    for pat in RCPTHOSTS:
        if host.endswith(pat):
            return 1
    return 0

def safe(s, pat=re.compile("[^.@%+/=:a-zA-Z0-9-]")):
    return pat.sub("?", s)

RECEIVED_PAT = re.compile("(received|delivered)", re.I)

#try:
#    import thread
#    ServerBase = SocketServer.ThreadingTCPServer
#except ImportError:
#    ServerBase = SocketServer.ForkingTCPServer
ServerBase = SocketServer.ForkingTCPServer


class QmailQueue:

    ERRORS = {
        115: "Denvelope address too long for qq (#5.1.3)",
        11: "Denvelope address too long for qq (#5.1.3)",
        31: "Dmail server permanently rejected message (#5.3.0)",
        51: "Zqq out of memory (#4.3.0)",
        52: "Zqq timeout (#4.3.0)",
        53: "Zqq write error or disk full (#4.3.0)",
        54: "Zqq read error (#4.3.0)",
        55: "Zqq unable to read configuration (#4.3.0)",
        56: "Zqq trouble making network connection (#4.3.0)",
        61: "Zqq trouble in home directory (#4.3.0)",
        62: "Zqq trouble creating files in queue (#4.3.0)",
        63: "Zqq trouble creating files in queue (#4.3.0)",
        64: "Zqq trouble creating files in queue (#4.3.0)",
        65: "Zqq trouble creating files in queue (#4.3.0)",
        66: "Zqq trouble creating files in queue (#4.3.0)",
        71: "Zmail server temporarily rejected message (#4.3.0)",
        72: "Zconnection to mail server timed out (#4.4.1)",
        73: "Zconnection to mail server rejected (#4.4.1)",
        74: "Zcommunication with mail server failed (#4.4.2)",
        91: "Zqq internal bug (#4.3.0)",
        81: "Zqq internal bug (#4.3.0)",
        120: "Zunable to exec qq (#4.3.0)",
    }

    ARGS = ("bin/qmail-queue",)

    def __init__(self):
        pim = os.pipe()
        try:
            pie = os.pipe()
        except:
            os.close(pim[0])
            os.close(pim[1])
            raise
        self.pid = os.fork()
        if self.pid == 0:
            try:
                os.close(pim[1])
                os.close(pie[1])
                os.dup2(pim[0], 0)
                os.close(pim[0])
                os.dup2(pie[0], 1)
                os.close(pie[0])
                try:
                    os.chdir(QMAIL_HOME)
                except OSError:
                    os._exit(61)
                os.execv(self.ARGS[0], self.ARGS)
                raise IOError
            except OSError:
                os._exit(120)
        else:
            self.fdm = pim[1]
            os.close(pim[0])
            self.fde = pie[1]
            os.close(pie[0])
            self.ss = os.fdopen(self.fdm, "w")
            self.flagerr = 0

    def get_pid(self):
        return self.pid

    def fail(self):
        self.flagerr = 1

    def write(self, s):
        if not self.flagerr:
            self.ss.write(s)

    def write_from(self, s):
        self.ss.flush()
        self.ss.close()
        self.ss = os.fdopen(self.fde, "w")
        self.write("F" + s + "\0")

    def write_to(self, s):
        self.write("T" + s + "\0")

    def close(self):
        self.write("\0")
        self.ss.flush()
        self.ss.close()
        pid, rv = os.waitpid(self.pid, 0)
        exitcode = os.WEXITSTATUS(rv)
        if exitcode in self.ERRORS:
            return self.ERRORS[exitcode]
        elif exitcode == 0 and not self.flagerr:
            return ""
        elif exitcode >= 11 and exitcode <= 40:
            return "Dqq permanent problem (#5.3.0)";
        else:
            return "Zqq temporary problem (#4.3.0)";


class SMTPHandler(SocketServer.StreamRequestHandler):

    COMMANDS = {
        # name : flush
        "rcpt": 0,
        "mail": 0,
        "data": 1,
        "quit": 1,
        "helo": 1,
        "ehlo": 1,
        "rset": 0,
        "help": 1,
        "noop": 1,
        "vrfy": 1,
        ""    : 1,
        }

    def setup(self):
        SocketServer.StreamRequestHandler.setup(self)
        self.seenmail = 0
        self.rcptto = []
        self.mailfrom = ""
        self.fakehelo = ""
        self.remoteip = self.client_address[0]
        try:
            self.remotehost = socket.gethostbyaddr(self.remoteip)[0].lower()
        except socket.error:
            self.remotehost = "unknown"
        log('connection from %s (%s)' % (self.remoteip, self.remotehost))
        signal.signal(signal.SIGALRM, self.die_alarm)
        signal.alarm(3600)
        self.dohelo(self.remotehost)

    def out(self, s):
        logwrite(s)
        self.wfile.write(s)

    def flush(self):
        self.wfile.flush()

    def msg_die(self, msg):
        self.out(msg + "\r\n")
        self.flush()
        os._exit(1)

    def die_read(self):
        os._exit(1)

    def die_alarm(self, *args):
        self.msg_die("451 timeout (#4.4.2)")

    def die_nomem(self):
        self.msg_die("421 out of memory (#4.3.0)")

    def die_control(self):
        self.msg_die("421 unable to read controls (#4.3.0)")

    def die_ipme(self):
        self.msg_die("421 unable to figure out my IP addresses (#4.3.0)")

    def straynewline(self):
        self.msg_die("451 See http://pobox.com/~djb/docs/smtplf.html.")

    def die_spam(self):
        self.msg_die("553 rejected, message looks like spam.")

    def die_exe(self):
        self.msg_die("553 rejected, executable attachment.")

    def err_nogateway(self):
        self.out("553 sorry, that domain isn't in my list of allowed rcpthosts "
                 "(#5.7.1)\r\n")

    def err_invaliduser(self):
        self.out("553 sorry, no user by that name is here (#5.7.1)\r\n")

    def err_unimpl(self):
        self.out("502 unimplemented (#5.5.1)\r\n")
        self.flush()

    def err_syntax(self):
        self.out("555 syntax error (#5.5.4)\r\n")

    def err_wantmail(self):
        self.out("503 MAIL first (#5.5.1)\r\n")

    def err_wantrcpt(self):
        self.out("503 RCPT first (#5.5.1)\r\n")

    def err_qqt(self):
        self.out("451 qqt failure (#4.3.0)\r\n")

    def greet(self, code):
        self.out(code)
        self.out(GREETING)

    def smtp_help(self, arg):
        self.out("214 qmail home page: http://pobox.com/~djb/qmail.html\r\n")

    def smtp_helo(self, arg):
        self.greet("250 ")
        self.out("\r\n")
        self.seenmail = 0
        self.dohelo(arg)
        
    def smtp_ehlo(self, arg):
        self.greet("250-")
        self.out("\r\n250-PIPELINING\r\n250 8BITMIME\r\n")
        self.seenmail = 0
        self.dohelo(arg)

    def smtp_rset(self, arg):
        self.seenmail = 0
        self.out("250 flushed\r\n")

    def smtp_noop(self, arg):
        self.out("250 ok\r\n")

    def smtp_vrfy(self, arg):
        self.out("252 send some mail, i'll try my best\r\n")

    def smtp_mail(self, arg):
        try:
            addr = addrparse(arg)
        except AddressSyntaxError:
            self.err_syntax()
            return
        self.seenmail = 1
        self.rcptto = []
        self.mailfrom = addr
        self.out("250 ok\r\n")

    def smtp_rcpt(self, arg):
        if not self.seenmail:
            self.err_wantmail()
            return
        try:
            addr = addrparse(arg)
        except AddressSyntaxError:
            self.err_syntax()
            return
        for pat in RELAYCLIENTS:
            if self.remoteip.startswith(pat):
                break
        else:
            if not hostallowed(addr):
                self.err_nogateway()
                return
            if not userallowed(addr):
                self.err_invaliduser()
                return
        self.rcptto.append(addr)
        self.out("250 ok\r\n")

    def _received(self, qq):
        qq.write("Received: from %s" % safe(self.remotehost))
        if self.fakehelo:
            qq.write(" (HELO %s)" % safe(self.fakehelo))
        qq.write(" (%s)\n" % safe(self.remoteip))
        qq.write("  by %s with SMTP; " % safe(IPME))
        qq.write(time.strftime("%d %b %Y %H:%M:%S %z"))
        qq.write("\n")
        
    def _blast(self, qq, msgfile):
        header = 1
        dot = 0
        hops = 0
        while 1:
            line = self.rfile.readline()
            if not line.endswith('\r\n'):
                self.straynewline()
            if header:
                if line == '\r\n':
                    header = 0
                elif RECEIVED_PAT.match(line):
                    hops += 1
            if line == '.\r\n':
                break
            qq.write(line[:-2])
            qq.write("\n")
            msgfile.write(line[:-2])
            msgfile.write("\n")
        return hops

    def _acceptmessage(self, pid):
        self.out("250 ok %d qp %d\r\n" % (time.time(), pid))
        log('accepted mail from %s to %s' % (self.remotehost, self.rcptto))

    def smtp_data(self, arg):
        if not self.seenmail:
            self.err_wantmail()
            return
        if not self.rcptto:
            self.err_wantrcpt()
            return
        self.seenmail = 0
        qq = QmailQueue()
        msgfile = tempfile.TemporaryFile(prefix='qmail-smtp-tmp')
        self.out("354 go ahead\r\n")
        self._received(qq)
        hops = self._blast(qq, msgfile)
        hops = hops >= MAXHOPS
        if hops:
            qq.fail()
        msg = parse_msg(msgfile)
        msgfile.close()
        spam = exe = False
        if msg is not None:
            exe = has_exe(msg)
            if exe:
                qq.fail()
            else:
                spam = is_spam(msg, self.rcptto)
                if spam:
                    qq.fail()
        qq.write_from(self.mailfrom)
        for addr in self.rcptto:
            qq.write_to(addr)
        status = qq.close()
        if not status:
            self._acceptmessage(qq.get_pid())
            return
        if hops:
            self.out("554 too many hops, this message is looping "
                     "(#5.4.6)\r\n")
            return
        if exe:
            self.die_exe()
        if spam:
            self.die_spam()
        if status[0] == "D":
            self.out("554 ")
        else:
            self.out("451 ")
        self.out(status[1:])
        self.out("\r\n")
        
    def smtp_quit(self, arg):
        self.greet("221 ")
        self.out("\r\n")
        self.flush()
        os._exit(0)

    def dohelo(self, arg):
        helohost = arg
        if helohost.lower() != self.remotehost.lower():
            self.fakehelo = helohost
        else:
            self.fakehelo = ""

    def commands(self):
        while 1:
            cmd = ""
            while 1:
                c = self.rfile.read(1)
                if len(c) != 1:
                    self.die_read()
                if c == "\n":
                    break
                cmd += c
            if len(cmd) > 0 and cmd[-1] == "\r":
                cmd = cmd[:-1]
            if cmd.find(" ") >= 0:
                cmd, arg = cmd.split(" ", 1)
                arg = arg.lstrip()
            else:
                arg = ""
            log('cmd', cmd + ' ' + arg)
            cmd = cmd.lower()
            doflush = self.COMMANDS.get(cmd)
            if doflush is None:
                self.err_unimpl()
            else:
                func = getattr(self, "smtp_%s" % cmd)
                func(arg)
                if doflush:
                    self.flush()

    def handle(self):
        try:
            self.greet("220 ")
            self.out(" ESMTP\r\n")
            self.commands()
        except MemoryError:
            self.die_nomem()


class SMTPServer(ServerBase):

    allow_reuse_address = True

    def __init__(self, port):
        ServerBase.__init__(self, ("0", port), SMTPHandler)

    def server_bind(self):
        ServerBase.server_bind(self)
        if os.geteuid() == 0:
            os.setgid(GROUP)
            os.setuid(USER)

def main():
    if len(sys.argv) < 2:
        port = 25
    else:
        port = int(sys.argv[1])
    log('qmail-smtpd.py listening on port', port)
    s = SMTPServer(port)
    s.serve_forever()

main()