Basic IMAP implementation.

startupstorm · Jul 15, 2013 · 7d56c96 · 7d56c96
1 parent 4c4ff34
commit 7d56c96
Show file tree

Hide file tree

Showing 3 changed files with 185 additions and 8 deletions.
diff --git a/mailpile/commands.py b/mailpile/commands.py
@@ -7,6 +7,7 @@
 # commands have been defined and what their names and command-line flags are.
 #
 import datetime
+import logging
 import os
 import os.path
 import re
@@ -23,6 +24,7 @@
 except ImportError:
   GnuPG = None
 
+logger = logging.getLogger(__name__)
 
 class Command:
   """Generic command object all others inherit from"""
@@ -956,13 +958,16 @@ def command(self):
     if fn in config.get('mailbox', {}).values():
       session.ui.warning('Already in the pile: %s' % fn)
     else:
-      if os.path.exists(fn):
-        arg = os.path.abspath(fn)
-        if config.parse_set(session,
-                            'mailbox:%s=%s' % (config.nid('mailbox'), fn)):
-          self._serialize('Save config', lambda: config.save())
+      if fn.startswith("imap://"):
+        arg = fn
       else:
-        return self._error('No such file/directory: %s' % raw_fn)
+        if os.path.exists(fn):
+          arg = os.path.abspath(fn)
+        else:
+          return self._error('No such file/directory: %s' % raw_fn)
+      if config.parse_set(session,
+                          'mailbox:%s=%s' % (config.nid('mailbox'), fn)):
+        self._serialize('Save config', lambda: config.save())
     return True
 
 

diff --git a/mailpile/imap_mailbox.py b/mailpile/imap_mailbox.py
@@ -0,0 +1,105 @@
+import logging
+try:
+    import cStringIO as StringIO
+except ImportError:
+    import StringIO
+
+from imaplib import IMAP4, IMAP4_SSL
+from mailbox import Mailbox, Message
+
+logger = logging.getLogger(__name__)
+
+class IMAPMailbox(Mailbox):
+    """
+    Basic implementation of IMAP Mailbox. Needs a lot of work.
+
+    As of now only get_* is implemented.
+    """
+    def __init__(self, host, port=993, user=None, password=None, mailbox=None, use_ssl=True, factory=None):
+        """Initialize a Mailbox instance."""
+        logger.debug("Opening IMAP mailbox %s:%d. SSL:%s" % (host, port, use_ssl))
+        if use_ssl:
+            self._mailbox = IMAP4_SSL(host, port)
+        else:
+            self._mailbox = IMAP4(host, port)
+        self._mailbox.login(user, password)
+        if not mailbox:
+            mailbox = "INBOX"
+        self.mailbox = mailbox
+        self._mailbox.select(mailbox)
+        self._factory = factory
+
+    def add(self, message):
+        """Add message and return assigned key."""
+        # TODO(halldor): not tested...
+        self._mailbox.append(self.mailbox, message=message)
+
+    def remove(self, key):
+        """Remove the keyed message; raise KeyError if it doesn't exist."""
+        # TODO(halldor): not tested...
+        self._mailbox.store(key, "+FLAGS", r"\Deleted")
+
+    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):
+        logger.debug("Fetching %s" % key)
+        typ, data = self._mailbox.fetch(key, '(RFC822)')
+        response = data[0]
+        if typ != "OK" or response is None:
+            raise KeyError
+        return response[1]
+
+    def get_message(self, key):
+        """Return a Message representation or raise a KeyError."""
+        return Message(self._get(key))
+
+    def get_bytes(self, key):
+        """Return a byte 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."""
+        message = self._get(key)
+        fd = StringIO.StringIO()
+        fd.write(message)
+        fd.seek(0)
+        return fd
+
+    def iterkeys(self):
+        """Return an iterator over keys."""
+        typ, data = self._mailbox.search(None, "ALL")
+        return data[0].split()
+
+    def __contains__(self, key):
+        """Return True if the keyed message exists, False otherwise."""
+        typ, data = self._mailbox.fetch(key, '(RFC822)')
+        response = data[0]
+        if response is None:
+            return False
+        return True
+
+    def __len__(self):
+        """Return a count of messages in the mailbox."""
+        return len(self.iterkeys())
+
+    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."""
+        self._mailbox.close()
+        self._mailbox.logout()
+
+    # Whether each message must end in a newline
+    _append_newline = False
diff --git a/mailpile/mailutils.py b/mailpile/mailutils.py
@@ -16,6 +16,7 @@
 import email.parser
 import email.utils
 import errno
+import logging
 import mailbox
 import mimetypes
 import os
@@ -32,6 +33,7 @@
 
 from smtplib import SMTP, SMTP_SSL
 
+from mailpile import imap_mailbox
 from mailpile.util import *
 from lxml.html.clean import Cleaner
 
@@ -41,6 +43,7 @@
 except ImportError:
   GnuPG = PGPMimeParser = None
 
+logger = logging.getLogger(__name__)
 
 class NotEditableError(ValueError):
   pass
@@ -60,6 +63,11 @@ class NoSuchMailboxError(OSError):
 
 def ParseMessage(fd, pgpmime=True):
   pos = fd.tell()
+  if logger.isEnabledFor(logging.DEBUG):
+    fd.seek(0, os.SEEK_END)
+    _len = fd.tell()
+    fd.seek(pos)
+    logger.debug("Parsing message (fd: %s, len: %d, pos: %d)" % (fd, _len, pos))
   header = [fd.readline()]
   while header[-1] not in ('', '\n', '\r\n'):
     line = fd.readline()
@@ -193,7 +201,7 @@ def closer():
 DULL_HEADERS = ('in-reply-to', 'references')
 def HeaderPrint(message):
   """Generate a fingerprint from message headers which identifies the MUA."""
-  headers = [x.split(':', 1)[0] for x in message.raw_header]
+  headers = message.keys() #[x.split(':', 1)[0] for x in message.raw_header]
 
   while headers and headers[0].lower() not in MUA_HEADERS:
     headers.pop(0)
@@ -204,14 +212,73 @@ def HeaderPrint(message):
   return b64w(sha1b64('\n'.join(headers))).lower()
 
 
-def OpenMailbox(fn):
+def OpenMailbox(fn): 
+  logger.debug("Opening mailbox %s" % fn)
+  if fn.startswith("imap://"):
+    # FIXME(halldor): waaaayy too naive - expects imap://username:password@server/mailbox
+    url = fn[7:]
+    try:
+      serverpart, mailbox = url.split("/")
+    except ValueError:
+      serverpart = url
+      mailbox = None
+    try:
+      userpart, server = serverpart.split("@")
+      user, password = userpart.split(":")
+    except ValueError:
+      raise
+
+    return IncrementalIMAPMailbox(server, user=user, password=password, mailbox=mailbox)
   if os.path.isdir(fn) and os.path.exists(os.path.join(fn, 'cur')):
     return IncrementalMaildir(fn)
   elif os.path.isdir(fn) and os.path.exists(os.path.join(fn, 'db')):
     return IncrementalGmvault(fn)
   else:
     return IncrementalMbox(fn)
 
+class IncrementalIMAPMailbox(imap_mailbox.IMAPMailbox):
+  editable = True
+  save_to = None
+  parsed = {}
+
+  def __setstate__(self, dict):
+    self.__dict__.update(dict)
+    self.update_toc()
+
+  def __getstate__(self):
+    odict = self.__dict__.copy()
+    # Pickle can't handle file objects.
+    del odict['_mailbox']
+    return odict
+
+  def save(self, session=None, to=None):
+    if to:
+      self.save_to = to
+    if self.save_to and len(self) > 0:
+      if session: session.ui.mark('Saving state to %s' % self.save_to)
+      fd = open(self.save_to, 'w')
+      cPickle.dump(self, fd)
+      fd.close()
+
+  def unparsed(self):
+    return [i for i in self.keys() if i not in self.parsed]
+
+  def mark_parsed(self, i):
+    self.parsed[i] = True
+
+  def update_toc(self):
+    self._refresh()
+
+  def get_msg_size(self, toc_id):
+    fd = self.get_file(toc_id)
+    fd.seek(0, 2)
+    return fd.tell()
+
+  def get_msg_ptr(self, idx, toc_id):
+    return '%s%s' % (idx, toc_id)
+
+  def get_file_by_ptr(self, msg_ptr):
+    return self.get_file(msg_ptr[3:])
 
 class IncrementalMaildir(mailbox.Maildir):
   """A Maildir class that supports pickling and a few mailpile specifics."""