#!/usr/bin/env python

"""SMTP server base class.

Note: the class in this module doesn't implement a fully functional SMTP
server; it simply provides a framework that subclasses can use to build
their own special-purpose SMTP servers.

Contents:

- SMTPServer: SMTP server class
- ForkingSMTPServer: forking SMTP server class
- ThreadingSMTPServer: threaded SMTP server class
- BaseSMTPRequestHandler: SMTP request handler base class

TODO:

- support mailbox delivery
- ESMTP support
- STARTTLS and SMTPS
- SMTP AUTH
"""

__version__ = '0.3'

__all__ = ['BaseSMTPRequestHandler', 'SMTPServer',
           'ForkingSMTPServer', 'ThreadingSMTPServer']

import logging
import SocketServer


class SMTPServer(SocketServer.TCPServer):
    allow_reuse_address = 1

class ForkingSMTPServer(SocketServer.ForkingMixIn, SMTPServer): pass
class ThreadingSMTPServer(SocketServer.ThreadingMixIn, SMTPServer): pass


class BaseSMTPRequestHandler(SocketServer.StreamRequestHandler):

    """SMTP request handler base class.

    The following methods are provided:

    - handle(), which is called by the SMTPServer object to handle an
    incoming connection;

    - read_command(), which is called by handle() to read
    an SMTP command from the client;

    - send_response(), which is called by handle() to send
    an SMTP response code to the client;

    - reset_state(), which resets the state of the SMTP session;

    To handle each SMTP command, the handle() method will look for a
    method with the name "smtp_", plus the upper-case name of the SMTP
    command.  So, for example, the smtp_HELO() method is a handler for
    the HELO command.  Support for additional SMTP commands can be added
    by simply defining the smtp_*() method for each new command.

    SMTP command handlers take one argument, which is a list containing
    the parsed command line returned by read_command().  They return a
    tuple containing an SMTP response code and a one-line human-readable
    message.  In most cases, this is something like (250, 'ok').

    Message delivery is done via the process_message() method, which
    must be provided by subclasses.

    Access control is provided via the allow_host(), allow_sender(), and
    allow_recipient() methods.  The implementation of those methods in
    this class will allow all hosts, senders, and recipients by default.
    Subclasses may override these methods to provide whatever access
    control policies are appropriate.

    The following instance variables are supported:

    - close_flag is a boolean variable that can be set to true by an
    SMTP command handler to indicate to the handle() method that the
    SMTP connection should be closed;

    - helo is a string containing the argument sent by the client to the
    HELO command;

    - allow_host is a boolean variable containing the result of the
    allow_host() method;

    - state is set to of the class variables STATE_INIT, STATE_MAIL,
    or STAT_RCPT, depending on the current phase of the SMTP transction;

    - mail_from is a string value containing the sender address (set
    to a non-None value when in STATE_MAIL or STATE_RCPT);

    - rcpt_to is a list containing one or more recipient addresses
    (non-empty when in STATE_RCPT);

    In addition, all of the instance variables documented for the
    SocketServer.StreamRequestHandler base class are also available.  In
    particular, this includes rfile, wfile, server, and client_address.
    """

    # SMTP state values
    STATE_INIT = 0       # no transaction currently in progress
    STATE_MAIL = 1       # after MAIL command
    STATE_RCPT = 2       # after first RCPT command, before DATA command

    def handle(self):
        """Main loop for the SMTP connection."""
        logging.debug('new SMTP connection from: %s' % (self.client_address,))
        self.close_flag = False
        self.helo = None
        self.reset_state()
        self.allow_host = self.allow_host(self.client_address)
        try:
            self.send_response(220, 'hi there')
            while not self.close_flag:
                args = self.read_command()
                if args is None:
                    # EOF
                    break
                cmd = args.pop(0).upper()
                method_name = 'smtp_' + cmd
                if hasattr(self, method_name):
                    method = getattr(self, method_name)
                    code, msg = method(args)
                else:
                    code = 501
                    msg = 'unknown command "%s"' % cmd
                self.send_response(code, msg)
        except socket.error, e:
            logging.error('socket error: %s' % e)
        logging.debug('SMTP connection closed')

    #
    # utility methods
    #

    def reset_state(self):
        """Utility function to reset the state of the SMTP session.
        Should be called before a new SMTP transaction is started.
        """
        self.state = self.STATE_INIT
        self.mail_from = None
        self.rcpt_to = [ ]

    def read_command(self):
        """Read a command from the client, parse it into arguments,
        and return the argument list.
        """
        line = self.rfile.readline()
        if not line:
            # return None on EOF
            return None
        line = line.strip()
        logging.debug('>>> ' + line)
        return line.split(' ')

    def send_response(self, code, msg):
        """Sends an SMTP response to the client."""
        response = '%d %s' % (code, msg)
        logging.debug('<<< ' + response)
        self.wfile.write(response + '\n')

    #
    # SMTP command handlers
    #

    def smtp_HELO(self, args):
        if len(args) != 1:
            return 501, 'usage: HELO hostname'
        if self.helo:
            return 503, 'already said HELO'
        self.helo = args[0]
        return 250, 'ok'

    def smtp_QUIT(self, args):
        # Setting this flag tells the main loop in the handler() method
        # that the connection should be closed after sending this
        # response.
        self.close_flag = True
        return 221, 'closing connection'

    def smtp_RSET(self, args):
        if len(args) != 1:
            return 501, 'usage: RSET'
        self.reset_state()
        return 250, 'ok'

    def smtp_NOOP(self, args):
        if len(args) != 1:
            return 501, 'usage: NOOP'
        return 250, 'ok'

    def smtp_MAIL(self, args):
        # handle "MAIL FROM:<address>" or "MAIL FROM <address>"
        if len(args) == 1 and args[0].upper().startswith('FROM:') and \
           len(args[0]) > 5:
            args = [ 'FROM', args[0][5:] ]
        elif args[0].endswith(':'):
            args[0] = args[0][:len(args[0]) - 1]
        if len(args) != 2 or args[0].upper() != 'FROM':
            return 501, 'usage: MAIL FROM address'

        if self.allow_host:
            return self.allow_host
        if self.state != self.STATE_INIT:
            return 503, 'transaction already in progress - use RSET to abort'
        response = self.allow_sender(args[1])
        if response:
            return response

        self.state = self.STATE_MAIL
        self.mail_from = args[1]
        return 250, 'sender ok'

    def smtp_RCPT(self, args):
        # handle "RCPT TO:<address>" or "RCPT TO <address>"
        if len(args) == 1 and args[0].upper().startswith('TO:') and \
           len(args[0]) > 3:
            args = [ 'TO', args[0][3:] ]
        elif args[0].endswith(':'):
            args[0] = args[0][:len(args[0]) - 1]
        if len(args) != 2 or args[0].upper() != 'TO':
            return 501, 'usage: RCPT TO address'

        if self.state not in (self.STATE_MAIL, self.STATE_RCPT):
            return 503, 'send MAIL command first'
        response = self.allow_recipient(args[1])
        if response:
            return response

        self.state = self.STATE_RCPT
        self.rcpt_to.append(args[1])
        return 250, 'recipient ok'

    def smtp_DATA(self, args):
        if args:
            return 501, 'usage: DATA'
        if self.state != self.STATE_RCPT:
            return 503, 'send RCPT command first'
        self.send_response(354, 'end DATA with <CR><LF>.<CR><LF>')

        data = [ ]
        while True:
            line = self.rfile.readline().rstrip('\n\r')
            if line.startswith('.'):
                if len(line) == 1:
                    break
                data.append(line[1:])
                continue
            data.append(line)

        result = self.process_message(self.mail_from, self.rcpt_to,
                                      '\n'.join(data))
        self.reset_state()
        if result:
            return result
        return 250, 'message accepted'

    #
    # methods to be implemented in subclasses
    #

    def process_message(self, sender, recipients, msg):
        """Override this abstract method to handle messages from the client.

        sender is the raw address the client claims the message is coming
        from.

        recipients is a list of raw addresses the client wishes to deliver
        the message to.

        msg is a string containing the entire full text of the message,
        headers (if supplied) and all.  It has been `de-transparencied'
        according to RFC 821, Section 4.5.2.  In other words, a line
        containing a `.' followed by other text has had the leading dot
        removed.

        This method should return None if the message was processed
        successfully.  On failure, it returns a tuple containing the SMTP
        response code and a one-line human-readable message; for example,
        to defer a message, it might return (450, 'temporary failure').
        """
        raise NotImplementedError

    def allow_sender(self, sender):
        """Access control function to allow or deny the specified sender.
        This implementation will allow all senders.  Subclasses may
        override this method.

        This method should return None if the sender should be allowed.
        Otherwise, it returns a tuple containing the SMTP response code
        and a one-line human-readable message; for example, to reject a
        sender, it might return (553, 'access denied').
        """
        return None

    def allow_recipient(self, recipient):
        """Access control function to allow or deny the specified recipient.
        This implementation will allow all recipients.  Subclasses may
        override this method.

        This method should return None if the recipient should be allowed.
        Otherwise, it returns a tuple containing the SMTP response code
        and a one-line human-readable message; for example, to reject a
        recipient, it might return (553, 'relaying denied').
        """
        return None

    def allow_host(self, host):
        """Access control function to allow or deny the specified client
        host.  This implementation will allow all hosts.  Subclasses may
        override this method.

        This method should return None if the client host should be allowed.
        Otherwise, it returns a tuple containing the SMTP response code
        and a one-line human-readable message (to be returned to the
        client when it issues the MAIL command); for example, to reject a
        client, it might return (550, 'access denied').
        """
        return None


class TestSMTPRequestHandler(BaseSMTPRequestHandler):

    """A minimal subclass that handles incoming messages by printing
    them out.  Not practically useful; just a demonstraction.
    """

    def process_message(self, sender, recipients, msg):
        print 'sender: ' + sender
        for rcpt in recipients:
            print 'recipient: ' + rcpt
        print 'msg:'
        print msg
        return None


if __name__ == '__main__':
    server = ThreadingSMTPServer(('', 2025), TestSMTPRequestHandler)
    server.serve_forever()