#!/usr/bin/env python
##
## Name:     eyepopper.py
## Purpose:  Simple mailbox-to-POP3 gateway.
## Author:   M. J. Fromberger <http://www.dartmouth.edu/~sting/>
## Info:     $Id: eyepopper.py 664 2008-09-19 03:21:03Z sting $
##

from __future__ import with_statement

from SocketServer import TCPServer, StreamRequestHandler
import getopt, hashlib, os, re, socket, sys

import pdb
__version__ = "1.1"

# {{ class MessageBase 

class MessageBase (object):
    """Base class representing a message stored in a Unix mailbox file.
    Subclasses must override:

    ._data    -- the complete mailbox data (string).
    ._path    -- the pathname of the mailbox file (string).

    Instances provide the following attributes:
    .start    -- position of first character of message
    .end      -- position after last character of message
    .poplen   -- length of message including CRLF conversion
    .bodypos  -- position of first character of mesage body
    .uid      -- unique identifier for message
    .content  -- the complete content of the message (computed)
    .head     -- the headers of the message (computed)
    .body     -- the body of the message (computed)
    """
    def __init__(self, start, end):
        self.start  = start
        self.end    = end

        text = self._data[start:end]
        self.poplen = (end - start) + len(re.findall(r'\n', text))
        bmark = re.search(r'\n\n', text)
        if bmark is None:
            pdb.set_trace()
        self.bodypos = start + bmark.end()

        hash = hashlib.md5()
        hash.update(text)
        hash.update(self._path)

        self.uid    = hash.hexdigest()

    def get_content(self):
        return self._data[self.start : self.end]
    def get_head(self):
        return self._data[self.start : self.bodypos]
    def get_body(self):
        return self._data[self.bodypos : self.end]
    def get_top(self, n):
        """As .get_content(), but return at most n lines from the
        beginning of the message body, n >= 0.
        """
        assert n >= 0

        lines = self.body.split('\n')
        return self.head + '\n'.join(lines[:n])

    content = property(get_content, None, None,
                       "Return the complete content of the message.")
    head    = property(get_head, None, None,
                       "Return the headers of the message.")
    body    = property(get_body, None, None,
                       "Return the body of the message.")

# }}

# {{ class MailboxFile 

class MailboxFile (object):
    """Abstracts a Unix mailbox file stored on disk."""
    def __init__(self, path):
        """Construct a new mailbox abstraction around the given file."""
        self._path = os.path.realpath(path)
        self.build_index()

    def __len__(self):
        """Return the number of messages in this mailbox."""
        return len(self._msgs)

    def __getitem__(self, itm):
        """Index into the mailbox by message."""
        return self._msgs[itm]
    
    def __iter__(self):
        """Iterate over the messages in this mailbox."""
        return iter(self._msgs)

    def mailbox_path(self):
        """Return the pathname of the mailbox."""
        return self._path
    
    def mailbox_size(self):
        """Return the total size of the mailbox in octets."""
        return sum(m.poplen for m in self._msgs)
        
    def build_index(self):
        """Construct an index of the contents of the associated disk file."""
        ex = re.compile(r'^From \w+.*\n\w+:', re.MULTILINE | re.UNICODE)
        with file(self._path, 'rt') as fp:
            data = fp.read()

            class Message (MessageBase):
                _data = data
                _path = self._path

            msgs = [{}] # sentinel
            
            for msg in ex.finditer(data):
                msgs[-1]['end'] = msg.start() - 1
                msgs.append({'start': msg.start()})

            msgs[-1]['end'] = len(data)
            msgs.pop(0)

            self._msgs = list(Message(m['start'], m['end'])
                              for m in msgs)
            self._data = data

# }}

# {{ class POP3Server 

class POP3Server (TCPServer):
    """Implements a simple read-only POP3 server."""
    allow_reuse_address = True
    
    def __init__(self, port, mailboxes, debug = False, users = ()):
        """Create a new server instance.

        port       -- TCP port to listen on.
        mailboxes  -- MailboxFile objects to serve.
        debug      -- enable debugging output?
        """
        self._debug = debug
        self._host  = 'localhost'
        self._port  = port
        self._boxes = list(mailboxes)
        self._dels  = set()
        self._users = dict(users)
        
        TCPServer.__init__(self, (self._host, port),
                           POP3Handler)

    def _diag(self, msg, *args):
        if self._debug:
            print >> sys.stderr, msg % args

    def run(self):
        """Run the server until it is killed or closed."""
        self.server_activate()
        self.running = True
        try:
            self._diag('* SERVER STARTING at %s %s',
                       self._host, self._port)

            while self.running:
                self.handle_request()

            self._diag('* SERVER SHUTTING DOWN')
        except KeyboardInterrupt:
            print >> sys.stderr, "\n>> INTERRUPT <<"
        finally:
            self.server_close()
    
    def mail_total_count(self):
        """[client] Return the total number of messages in all
        known mailboxes.
        """
        return sum(len(b) for b in self._boxes)

    def mail_total_size(self):
        """[client] Return the total size of the messages in all
        known mailboxes.
        """
        return sum(b.mailbox_size() for b in self._boxes)

    def mail_get_index(self, pos):
        """[client] Fetch a message object by its global index (0-based)."""
        if pos < 0:
            raise IndexError("Out of range")
        
        for mbox in self._boxes:
            if len(mbox) > pos:
                break

            pos -= len(mbox)
        else:
            raise IndexError("Out of range")

        return mbox[pos]

    def mail_get_all(self):
        """[client] Iterate over all the messages in all known mailboxes."""
        for mbox in self._boxes:
            for msg in mbox:
                yield msg

    def mail_mark_deleted(self, pos):
        """[client] Mark the specified message as "deleted"."""
        self.mail_get_index(pos) # range test
        self._dels.add(pos)

    def mail_get_deleted(self):
        """[client] Return a set of indices of deleted messages."""
        return self._dels
    
    def mail_clear_deleted(self):
        """[client] Clear the list of "deleted" messages."""
        self._dels.clear()

    def mail_add_mailbox(self, path):
        """[client] Add a new mailbox to the list of known mailboxes."""
        for box in self._boxes:
            if box.mailbox_path() == path:
                return

        box = MailboxFile(path) # may fail
        self._boxes.append(box)
    
    def is_mail_deleted(self, pos):
        """[client] Check whether the message specified is "deleted"."""
        return pos in self._dels
    
    def is_user_ok(self, user):
        """[client] Check whether a user ID is acceptable."""
        return len(self._users) == 0 or user in self._users

    def is_auth_ok(self, user, pw):
        """[client] Check user authentication."""
        return len(self._users) == 0 or (
            user in self._users and self._users[user] == pw)

# }}

# {{ class StateError 

class StateError (Exception):
    """An exception used internally by POP3Handler."""
    pass

# }}

# {{ class POP3Handler 

class POP3Handler (StreamRequestHandler):
    """Implements a very simple POP3 service.

    In addition to the usual POP3 suite, this handler implements the
    following client commands:

    SHUTDOWN       -- close down the server immediately.
    ADDBOX <path>  -- load a new mailbox into the running server.
    """
    welcome_banner = 'POP3 server ready'

    def cmd_unknown(self, cmd, data):
        """Handle unknown commands."""
        self.wfile.write('-ERR Command not understood\r\n')

    def _diag(self, msg, *args):
        self.server._diag(msg, *args)
    
    def check_state(self, cmd, data, state):
        if self.state <> state:
            self.wfile.write('-ERR Invalid command\r\n')
            raise StateError

    def check_args(self, cmd, data, min, max):
        args = re.split(r'\s+', data) if data else []

        if min <= len(args) and (max < 0 or len(args) <= max):
            return args
        else:
            self.wfile.write('-ERR Wrong number of arguments (%s)\r\n' % cmd)
            raise StateError

    def parse_args(self, args):
        try:
            return list(int(x) for x in args)
        except ValueError:
            self.wfile.write('-ERR Invalid argument\r\n')
            raise StateError

    def check_message(self, pos):
        try:
            msg = self.server.mail_get_index(pos)
            if self.server.is_mail_deleted(pos):
                self.wfile.write('-ERR Message is deleted\r\n')
                raise StateError
            return msg
        except IndexError:
            self.wfile.write('-ERR Message out of range\r\n')
            raise StateError
    
    def send_data(self, data):
        esc  = re.compile(r'^\.', re.MULTILINE | re.UNICODE)
        text = esc.sub('..', data.replace('\n', '\r\n'))
        self.wfile.write('+OK %d octets\r\n' % len(text))
        self.wfile.write(text)
        if not text.endswith('\r\n'):
            self.wfile.write('\r\n')
        self.wfile.write('.\r\n')
    
    def cmd_NOOP(self, cmd, data):
        self.check_args(cmd, data, 0, 0)
        self.wfile.write('+OK Nothing accomplished\r\n')

    def cmd_QUIT(self, cmd, data):
        self.check_args(cmd, data, 0, 0)
        self.state = 'UPDATE' # triggers exit from handler loop
        self._diag('- Received QUIT command, entering UPDATE state.')
    
    def cmd_USER(self, cmd, data):
        try:
            self.check_state(cmd, data, 'AUTH')
            self.check_args(cmd, data, 1, 1)
            if self.server.is_user_ok(data):
                self.wfile.write('+OK %s\r\n' % data)
                self.state  = 'USER'
                self.userid = data
            else:
                self.wfile.write('-ERR User invalid\r\n')
        except StateError:
            if self.state == 'USER':
                self.state = 'AUTH'

    def cmd_PASS(self, cmd, data):
        self.check_state(cmd, data, 'USER')
        self.check_args(cmd, data, 1, 1)
        if self.server.is_auth_ok(self.userid, data):
            self.wfile.write('+OK Ready\r\n')
            self.state = 'TRANS'
            self._diag('- Authenticated "%s", entering TRANSACTION state.',
                       self.userid)
        else:
            self.wfile.write('-ERR Access denied\r\n')
            self.state = 'AUTH'

    def cmd_STAT(self, cmd, data):
        self.check_state(cmd, data, 'TRANS')
        self.check_args(cmd, data, 0, 0)
        self.wfile.write('+OK %d %d\r\n' %
                         (self.server.mail_total_count(),
                          self.server.mail_total_size()))

    def list_cmd(extract):
        def do_command(self, cmd, data):
            self.check_state(cmd, data, 'TRANS')
            args = self.parse_args(self.check_args(cmd, data, 0, 1))
            if args:
                msg = self.check_message(args[0] - 1)
                elt = extract(msg)
                self.wfile.write('+OK %d %s\r\n' % (args[0], elt))
            else:
                self.wfile.write('+OK %d messages (%d octets)\r\n' % (
                    self.server.mail_total_count(),
                    self.server.mail_total_size()))
                for pos, msg in enumerate(self.server.mail_get_all()):
                    if not self.server.is_mail_deleted(pos):
                        elt = extract(msg)
                        self.wfile.write('%d %s\r\n' % (pos + 1, elt))
                self.wfile.write('.\r\n')
        return do_command

    cmd_LIST = list_cmd(lambda m: m.poplen)
    cmd_UIDL = list_cmd(lambda m: m.uid)
    
    def cmd_RETR(self, cmd, data):
        self.check_state(cmd, data, 'TRANS')
        args = self.parse_args(self.check_args(cmd, data, 1, 1))
        msg  = self.check_message(args[0] - 1)
        self.send_data(msg.content)
    
    def cmd_DELE(self, cmd, data):
        self.check_state(cmd, data, 'TRANS')
        args = self.parse_args(self.check_args(cmd, data, 1, 1))
        msg  = self.check_message(args[0] - 1)
        self.server.mail_mark_deleted(args[0] - 1)
        self.wfile.write('+OK Message %d deleted\r\n' % args[0])
        self._diag('- Marked message %d for deletion.', args[0])

    def cmd_RSET(self, cmd, data):
        self.check_state(cmd, data, 'TRANS')
        self.check_args(cmd, data, 0, 0)
        self.server.mail_clear_deleted()
        self.wfile.write('+OK Reset\r\n')
        self._diag('- Reset deleted messages list.')

    def cmd_TOP(self, cmd, data):
        self.check_state(cmd, data, 'TRANS')
        msgid, n = self.parse_args(self.check_args(cmd, data, 2, 2))

        if n < 0:
            self.wfile.write('-ERR Invalid argument\r\n')
            return
        
        msg  = self.check_message(msgid - 1)
        self.send_data(msg.get_top(n))

    def cmd_SHUTDOWN(self, cmd, data):
        self.check_state(cmd, data, 'TRANS')
        self.check_args(cmd, data, 0, 0)
        self.server.running = False
        self.state = 'UPDATE' # triggers exit from handler loop
        self._diag('- Client requested server SHUTDOWN.')

    def cmd_ADDBOX(self, cmd, data):
        self.check_state(cmd, data, 'TRANS')
        if not data or data.isspace():
            self.wfile.write('-ERR Invalid argument\r\n')
            return
        
        try:
            self.server.mail_add_mailbox(data)
            self.wfile.write('+OK Mailbox added %s\r\n' % data)
            self._diag('- Client requested mailbox "%s".', data)
        except (OSError, IOError), e:
            self.wfile.write('-ERR Mailbox not added (%s)\r\n' % e.strerror)
    
    def handle(self):
        """Required entry point for use with SocketServer.  Dispatches
        received commands to .cmd_XXXX() methods based on the first
        word of the command line received from the client.

        Command names are case-insensitive, but the method names to
        implement them must be capitalized.  If .cmd_unknown() is
        defined, it is given all commands that are not otherwise
        recognized.

        Note: This implementation is not thread-safe; in particular,
        it uses state on the server object without locks.
        """
        try:
            self.wfile.write('+OK %s\r\n' % self.welcome_banner)
            self.state = 'AUTH'
            self._diag('- Client connected, entering AUTH state.')

            while self.state <> 'UPDATE':
                line      = self.rfile.readline()
                args      = line.rstrip().split(' ', 1)
                cmd       = args[0]
                data      = ''
                if len(args) > 1:
                    data = args[1]
                hname     = 'cmd_%s' % cmd.upper()

                try:
                    getattr(self, hname, self.cmd_unknown)(cmd, data)
                except StateError:
                    pass

            # When control reaches here, we are in the "UPDATE" state
            # of the POP3 protocol.  Since this is read-only, deleted
            # messages will not actually be removed; the POP3 RFC
            # wants this to be an error.
            if len(self.server.mail_get_deleted()) == 0:
                self.wfile.write('+OK Goodnight sweet prince\r\n')
            else:
                self.wfile.write('-ERR Unable to delete messages\r\n')

            self.server.mail_clear_deleted()
            self._diag('* REQUEST COMPLETE\n')
        
        except socket.error, e:
            pass

        self.wfile.close()
        self.rfile.close()

# }}

# {{ main(argv) 

def main(argv):
    def usage(short = True):
        """Print a human-readable usage message."""
        print >> sys.stderr, "Usage: eyepopper.py [options] mailbox*"
        if short:
            print >> sys.stderr, " [use -h or --help for options]\n"
        else:
            print >> sys.stderr, """
This program implements a local read-only POP3 server that serves up
messages stored in plain text mailbox files.  It can be used to import
such messages into clients that do not fully grok Unix mailbox format.

Options:
  -h/--help             : display this help message.
  -p/--port <port>      : listen on the specified port (default: %s).
  -q/--quiet            : disable diagnostic output (run quietly)
  -u/--user <name>:<pw> : allow this username/password combination.

The POP_PORT environment variable will be used to specify the port, if
it is defined and the --port option is not provided.  Only one client
will be served at a time, and only clients connecting from localhost
are accepted.

By default, any username and password are accepted.  If the --user
option is used one or more times, only the user/password combinations
specified are granted access.  The name and password are separated by
a colon.  Username and password combinations can also be supplied via
the POP_USERS environment variable, where multiple entries are
delimited by carriage returns.
""" % listen_port

    # Process command-line options
    try:
        opts, args = getopt.gnu_getopt(argv, 'hp:qu:',
                                       ('help', 'port=', 'quiet', 'user='))
    except getopt.GetoptError, e:
        print >> sys.stderr, "Error: %s" % e
        usage()
        return 1

    listen_port  = int(os.getenv('POP_PORT', 1110))
    debugging    = True
    legal_users  = {}
    for opt, arg in opts:
        if opt in ('-h', '--help'):
            usage(short = False)
            return 0
        elif opt in ('-q', '--quiet'):
            debugging = False
        elif opt in ('-p', '--port'):
            listen_port = int(arg)
        elif opt in ('-u', '--user'):
            try:
                name, pw = arg.strip().split(':', 1)
                legal_users[name] = pw
            except ValueError:
                print >> sys.stderr, "Error: Invalid user specification: %s" % arg
                usage()
                return 1
        else:
            raise NotImplementedError("unimplemented option")

    # Load users from the environment, if any are defined
    for line in os.getenv('POP_USERS', '').split('\n'):
        try:
            name, pw = line.strip().split(':', 1)
            legal_users[name] = pw
        except ValueError:
            pass # skip quietly
    
    # Load mailbox files
    try:
        boxes = [MailboxFile(p) for p in args]
    except (OSError, IOError), e:
        print >> sys.stderr, "Error: %s" % e
        return 2
    
    # Start up server
    if debugging:
        print >> sys.stderr, "EyePopper v. %s by M. J. Fromberger" % __version__
        if boxes:
            print >> sys.stderr, "Mailboxes:\n -",
            print >> sys.stderr, '\n - '.join('%s (%d msg)' %
                                              (b.mailbox_path(), len(b))
                                              for b in boxes)
            print >> sys.stderr
        if legal_users:
            print >> sys.stderr, "Legal users: %s\n" % ', '.join(legal_users)
    
    pop = POP3Server(listen_port, boxes, debug = debugging, users = legal_users)
    pop.run()

    return 0

# }}

if __name__ == "__main__":
    res = main(sys.argv[1:])
    sys.exit(res)

# Here there be dragons