#! /usr/bin/env python

"""Read/write support for Maildir, mbox, MH, Babyl, and MMDF mailboxes."""

# Notes for authors of new mailbox subclasses:

#

# Remember to fsync() changes to disk before closing a modified file

# or returning from a flush() method.  See functions _sync_flush() and

# _sync_close().

import sys

import os

import time

import calendar

import socket

import errno

import copy

import email

import email.message

import email.generator

import rfc822

import StringIO

try:

    if sys.platform == 'os2emx':

        # OS/2 EMX fcntl() not adequate

        raise ImportError

    import fcntl

except ImportError:

    fcntl = None

__all__ = [ 'Mailbox', 'Maildir', 'mbox', 'MH', 'Babyl', 'MMDF',

            'Message', 'MaildirMessage', 'mboxMessage', 'MHMessage',

            'BabylMessage', 'MMDFMessage', 'UnixMailbox',

            'PortableUnixMailbox', 'MmdfMailbox', 'MHMailbox', 'BabylMailbox' ]

class Mailbox:

    """A group of messages in a particular place."""

    def __init__(self, path, factory=None, create=True):

        """Initialize a Mailbox instance."""

        self._path = os.path.abspath(os.path.expanduser(path))

        self._factory = factory

    def add(self, message):

        """Add message and return assigned key."""

        raise NotImplementedError('Method must be implemented by subclass')

    def remove(self, key):

        """Remove the keyed message; raise KeyError if it doesn't exist."""

        raise NotImplementedError('Method must be implemented by subclass')

    def __delitem__(self, key):

        self.remove(key)

    def discard(self, key):

        """If the keyed message exists, remove it."""

        try:

            self.remove(key)

        except KeyError:

            pass

    def __setitem__(self, key, message):

        """Replace the keyed message; raise KeyError if it doesn't exist."""

        raise NotImplementedError('Method must be implemented by subclass')

    def get(self, key, default=None):

        """Return the keyed message, or default if it doesn't exist."""

        try:

            return self.__getitem__(key)

        except KeyError:

            return default

    def __getitem__(self, key):

        """Return the keyed message; raise KeyError if it doesn't exist."""

        if not self._factory:

            return self.get_message(key)

        else:

            return self._factory(self.get_file(key))

    def get_message(self, key):

        """Return a Message representation or raise a KeyError."""

        raise NotImplementedError('Method must be implemented by subclass')

    def get_string(self, key):

        """Return a string representation or raise a KeyError."""

        raise NotImplementedError('Method must be implemented by subclass')

    def get_file(self, key):

        """Return a file-like representation or raise a KeyError."""

        raise NotImplementedError('Method must be implemented by subclass')

    def iterkeys(self):

        """Return an iterator over keys."""

        raise NotImplementedError('Method must be implemented by subclass')

    def keys(self):

        """Return a list of keys."""

        return list(self.iterkeys())

    def itervalues(self):

        """Return an iterator over all messages."""

        for key in self.iterkeys():

            try:

                value = self[key]

            except KeyError:

                continue

            yield value

    def __iter__(self):

        return self.itervalues()

    def values(self):

        """Return a list of messages. Memory intensive."""

        return list(self.itervalues())

    def iteritems(self):

        """Return an iterator over (key, message) tuples."""

        for key in self.iterkeys():

            try:

                value = self[key]

            except KeyError:

                continue

            yield (key, value)

    def items(self):

        """Return a list of (key, message) tuples. Memory intensive."""

        return list(self.iteritems())

    def has_key(self, key):

        """Return True if the keyed message exists, False otherwise."""

        raise NotImplementedError('Method must be implemented by subclass')

    def __contains__(self, key):

        return self.has_key(key)

    def __len__(self):

        """Return a count of messages in the mailbox."""

        raise NotImplementedError('Method must be implemented by subclass')

    def clear(self):

        """Delete all messages."""

        for key in self.iterkeys():

            self.discard(key)

    def pop(self, key, default=None):

        """Delete the keyed message and return it, or default."""

        try:

            result = self[key]

        except KeyError:

            return default

        self.discard(key)

        return result

    def popitem(self):

        """Delete an arbitrary (key, message) pair and return it."""

        for key in self.iterkeys():

            return (key, self.pop(key))     # This is only run once.

        else:

            raise KeyError('No messages in mailbox')

    def update(self, arg=None):

        """Change the messages that correspond to certain keys."""

        if hasattr(arg, 'iteritems'):

            source = arg.iteritems()

        elif hasattr(arg, 'items'):

            source = arg.items()

        else:

            source = arg

        bad_key = False

        for key, message in source:

            try:

                self[key] = message

            except KeyError:

                bad_key = True

        if bad_key:

            raise KeyError('No message with key(s)')

    def flush(self):

        """Write any pending changes to the disk."""

        raise NotImplementedError('Method must be implemented by subclass')

    def lock(self):

        """Lock the mailbox."""

        raise NotImplementedError('Method must be implemented by subclass')

    def unlock(self):

        """Unlock the mailbox if it is locked."""

        raise NotImplementedError('Method must be implemented by subclass')

    def close(self):

        """Flush and close the mailbox."""

        raise NotImplementedError('Method must be implemented by subclass')

    def _dump_message(self, message, target, mangle_from_=False):

        # Most files are opened in binary mode to allow predictable seeking.

        # To get native line endings on disk, the user-friendly \n line endings

        # used in strings and by email.Message are translated here.

        """Dump message contents to target file."""

        if isinstance(message, email.message.Message):

            buffer = StringIO.StringIO()

            gen = email.generator.Generator(buffer, mangle_from_, 0)

            gen.flatten(message)

            buffer.seek(0)

            target.write(buffer.read().replace('\n', os.linesep))

        elif isinstance(message, str):

            if mangle_from_:

                message = message.replace('\nFrom ', '\n>From ')

            message = message.replace('\n', os.linesep)

            target.write(message)

        elif hasattr(message, 'read'):

            while True:

                line = message.readline()

                if line == '':

                    break

                if mangle_from_ and line.startswith('From '):

                    line = '>From ' + line[5:]

                line = line.replace('\n', os.linesep)

                target.write(line)

        else:

            raise TypeError('Invalid message type: %s' % type(message))

class Maildir(Mailbox):

    """A qmail-style Maildir mailbox."""

    colon = ':'

    def __init__(self, dirname, factory=rfc822.Message, create=True):

        """Initialize a Maildir instance."""

        Mailbox.__init__(self, dirname, factory, create)

        if not os.path.exists(self._path):

            if create:

                os.mkdir(self._path, 0700)

                os.mkdir(os.path.join(self._path, 'tmp'), 0700)

                os.mkdir(os.path.join(self._path, 'new'), 0700)

                os.mkdir(os.path.join(self._path, 'cur'), 0700)

            else:

                raise NoSuchMailboxError(self._path)

        self._toc = {}

    def add(self, message):

        """Add message and return assigned key."""

        tmp_file = self._create_tmp()

        try:

            self._dump_message(message, tmp_file)

        finally:

            _sync_close(tmp_file)

        if isinstance(message, MaildirMessage):

            subdir = message.get_subdir()

            suffix = self.colon + message.get_info()

            if suffix == self.colon:

                suffix = ''

        else:

            subdir = 'new'

            suffix = ''

        uniq = os.path.basename(tmp_file.name).split(self.colon)[0]

        dest = os.path.join(self._path, subdir, uniq + suffix)

        try:

            if hasattr(os, 'link'):

                os.link(tmp_file.name, dest)

                os.remove(tmp_file.name)

            else:

                os.rename(tmp_file.name, dest)

        except OSError, e:

            os.remove(tmp_file.name)

            if e.errno == errno.EEXIST:

                raise ExternalClashError('Name clash with existing message: %s'

                                         % dest)

            else:

                raise

        if isinstance(message, MaildirMessage):

            os.utime(dest, (os.path.getatime(dest), message.get_date()))

        return uniq

    def remove(self, key):

        """Remove the keyed message; raise KeyError if it doesn't exist."""

        os.remove(os.path.join(self._path, self._lookup(key)))

    def discard(self, key):

        """If the keyed message exists, remove it."""

        # This overrides an inapplicable implementation in the superclass.

        try:

            self.remove(key)

        except KeyError:

            pass

        except OSError, e:

            if e.errno != errno.ENOENT:

                raise

    def __setitem__(self, key, message):

        """Replace the keyed message; raise KeyError if it doesn't exist."""

        old_subpath = self._lookup(key)

        temp_key = self.add(message)

        temp_subpath = self._lookup(temp_key)

        if isinstance(message, MaildirMessage):

            # temp's subdir and suffix were specified by message.

            dominant_subpath = temp_subpath

        else:

            # temp's subdir and suffix were defaults from add().

            dominant_subpath = old_subpath

        subdir = os.path.dirname(dominant_subpath)

        if self.colon in dominant_subpath:

            suffix = self.colon + dominant_subpath.split(self.colon)[-1]

        else:

            suffix = ''

        self.discard(key)

        new_path = os.path.join(self._path, subdir, key + suffix)

        os.rename(os.path.join(self._path, temp_subpath), new_path)

        if isinstance(message, MaildirMessage):

            os.utime(new_path, (os.path.getatime(new_path),

                                message.get_date()))

    def get_message(self, key):

        """Return a Message representation or raise a KeyError."""

        subpath = self._lookup(key)

        f = open(os.path.join(self._path, subpath), 'r')

        try:

            if self._factory:

                msg = self._factory(f)

            else:

                msg = MaildirMessage(f)

        finally:

            f.close()

        subdir, name = os.path.split(subpath)

        msg.set_subdir(subdir)

        if self.colon in name:

            msg.set_info(name.split(self.colon)[-1])

        msg.set_date(os.path.getmtime(os.path.join(self._path, subpath)))

        return msg

    def get_string(self, key):

        """Return a string representation or raise a KeyError."""

        f = open(os.path.join(self._path, self._lookup(key)), 'r')

        try:

            return f.read()

        finally:

            f.close()

    def get_file(self, key):

        """Return a file-like representation or raise a KeyError."""

        f = open(os.path.join(self._path, self._lookup(key)), 'rb')

        return _ProxyFile(f)

    def iterkeys(self):

        """Return an iterator over keys."""

        self._refresh()

        for key in self._toc:

            try:

                self._lookup(key)

            except KeyError:

                continue

            yield key

    def has_key(self, key):

        """Return True if the keyed message exists, False otherwise."""

        self._refresh()

        return key in self._toc

    def __len__(self):

        """Return a count of messages in the mailbox."""

        self._refresh()

        return len(self._toc)

    def flush(self):

        """Write any pending changes to disk."""

        return  # Maildir changes are always written immediately.

    def lock(self):

        """Lock the mailbox."""

        return

    def unlock(self):

        """Unlock the mailbox if it is locked."""

        return

    def close(self):

        """Flush and close the mailbox."""

        return

    def list_folders(self):

        """Return a list of folder names."""

        result = []

        for entry in os.listdir(self._path):

            if len(entry) > 1 and entry[0] == '.' and \

               os.path.isdir(os.path.join(self._path, entry)):

                result.append(entry[1:])

        return result

    def get_folder(self, folder):

        """Return a Maildir instance for the named folder."""

        return Maildir(os.path.join(self._path, '.' + folder),

                       factory=self._factory,

                       create=False)

    def add_folder(self, folder):

        """Create a folder and return a Maildir instance representing it."""

        path = os.path.join(self._path, '.' + folder)

        result = Maildir(path, factory=self._factory)

        maildirfolder_path = os.path.join(path, 'maildirfolder')

        if not os.path.exists(maildirfolder_path):

            os.close(os.open(maildirfolder_path, os.O_CREAT | os.O_WRONLY,

                0666))

        return result

    def remove_folder(self, folder):

        """Delete the named folder, which must be empty."""

        path = os.path.join(self._path, '.' + folder)

        for entry in os.listdir(os.path.join(path, 'new')) + \

                     os.listdir(os.path.join(path, 'cur')):

            if len(entry) < 1 or entry[0] != '.':

                raise NotEmptyError('Folder contains message(s): %s' % folder)

        for entry in os.listdir(path):

            if entry != 'new' and entry != 'cur' and entry != 'tmp' and \

               os.path.isdir(os.path.join(path, entry)):

                raise NotEmptyError("Folder contains subdirectory '%s': %s" %

                                    (folder, entry))

        for root, dirs, files in os.walk(path, topdown=False):

            for entry in files:

                os.remove(os.path.join(root, entry))

            for entry in dirs:

                os.rmdir(os.path.join(root, entry))

        os.rmdir(path)

    def clean(self):

        """Delete old files in "tmp"."""

        now = time.time()

        for entry in os.listdir(os.path.join(self._path, 'tmp')):

            path = os.path.join(self._path, 'tmp', entry)

            if now - os.path.getatime(path) > 129600:   # 60 * 60 * 36

                os.remove(path)

    _count = 1  # This is used to generate unique file names.

    def _create_tmp(self):

        """Create a file in the tmp subdirectory and open and return it."""

        now = time.time()

        hostname = socket.gethostname()

        if '/' in hostname:

            hostname = hostname.replace('/', r'\057')

        if ':' in hostname:

            hostname = hostname.replace(':', r'\072')

        uniq = "%s.M%sP%sQ%s.%s" % (int(now), int(now % 1 * 1e6), os.getpid(),

                                    Maildir._count, hostname)

        path = os.path.join(self._path, 'tmp', uniq)

        try:

            os.stat(path)

        except OSError, e:

            if e.errno == errno.ENOENT:

                Maildir._count += 1

                try:

                    return _create_carefully(path)

                except OSError, e:

                    if e.errno != errno.EEXIST:

                        raise

            else:

                raise

        # Fall through to here if stat succeeded or open raised EEXIST.

        raise ExternalClashError('Name clash prevented file creation: %s' %

                                 path)

    def _refresh(self):

        """Update table of contents mapping."""

        self._toc = {}

        for subdir in ('new', 'cur'):

            subdir_path = os.path.join(self._path, subdir)

            for entry in os.listdir(subdir_path):

                p = os.path.join(subdir_path, entry)

                if os.path.isdir(p):

                    continue

                uniq = entry.split(self.colon)[0]

                self._toc[uniq] = os.path.join(subdir, entry)

    def _lookup(self, key):

        """Use TOC to return subpath for given key, or raise a KeyError."""

        try:

            if os.path.exists(os.path.join(self._path, self._toc[key])):

                return self._toc[key]

        except KeyError:

            pass

        self._refresh()

        try:

            return self._toc[key]

        except KeyError:

            raise KeyError('No message with key: %s' % key)

    # This method is for backward compatibility only.

    def next(self):

        """Return the next message in a one-time iteration."""

        if not hasattr(self, '_onetime_keys'):

            self._onetime_keys = self.iterkeys()

        while True:

            try:

                return self[self._onetime_keys.next()]

            except StopIteration:

                return None

            except KeyError:

                continue

class _singlefileMailbox(Mailbox):

    """A single-file mailbox."""

    def __init__(self, path, factory=None, create=True):

        """Initialize a single-file mailbox."""

        Mailbox.__init__(self, path, factory, create)

        try:

            f = open(self._path, 'rb+')

        except IOError, e:

            if e.errno == errno.ENOENT:

                if create:

                    f = open(self._path, 'wb+')

                else:

                    raise NoSuchMailboxError(self._path)

            elif e.errno == errno.EACCES:

                f = open(self._path, 'rb')

            else:

                raise

        self._file = f

        self._toc = None

        self._next_key = 0

        self._pending = False   # No changes require rewriting the file.

        self._locked = False

        self._file_length = None        # Used to record mailbox size

    def add(self, message):

        """Add message and return assigned key."""

        self._lookup()