enhancements for email filter and attachment handling

mailbot/callback.py

@@ -12,6 +12,12 @@
 class Callback(object):
     """Base class for callbacks."""
 
+    # DEFINITIONS
+    GET_ALL_ATTACHMENTS = 255 # get all attachments from eMail
+    GET_PDF_ATTACHMENTS = 1 # get only PDFs
+    GET_MSOFFICE_ATTACHMENTS = 2 # get only Word documents
+    GET_IMAGE_ATTACHMENTS = 4 # get only images
+
     def __init__(self, message, rules):
         self.matches = defaultdict(list)
         self.message = message
@@ -93,6 +99,47 @@ def get_email_body(self, message=None):
 
         return ''
 
+    def get_attachments(self, message=None, documents=GET_ALL_ATTACHMENTS):
+        """Return a list of attachments.
+
+        Return all attached files within a list of 2-tuples filename and
+        content of the attachment to be stored locally.
+
+        """
+        if message is None:
+            message = self.message
+
+        if not hasattr(message, 'walk'):  # not an email.Message instance?
+            return None
+
+        content = []
+        for part in message.walk():
+            content_type = part.get_content_type()
+            filename = part.get_filename()
+
+            # check for desired file types
+            if documents == self.GET_ALL_ATTACHMENTS:
+
+                # get any file
+                if filename is not None:
+                    content = content + [(filename, part.get_payload(decode=True))]
+
+            else:
+
+                # PDF file type
+                if documents & self.GET_PDF_ATTACHMENTS and content_type == 'application/pdf' and filename is not None:
+                    content = content + [(filename,part.get_payload(decode=True))]
+
+                # MS OFFICE file types
+                if documents & self.GET_MSOFFICE_ATTACHMENTS and content_type.find('application/ms') > -1 and filename is not None:
+                    content = content + [(filename,part.get_payload(decode=True))]
+
+                # IMAGE file types
+                if documents & self.GET_IMAGE_ATTACHMENTS and content_type.find('image') > -1 and filename is not None:
+                    content = content + [(filename,part.get_payload(decode=True))]
+
+        return content
+
     def trigger(self):
         """Called when a mail matching the registered rules is received."""
         raise NotImplementedError("Must be implemented in a child class.")

mailbot/mailbot.py

@@ -17,6 +17,11 @@ class MailBot(object):
     home_folder = 'INBOX'
     imapclient = IMAPClient
 
+    # DEFINITIONS
+    CHECK_PRESENT_EMAILS = 1 # check all not deleted eMails
+    CHECK_UNSEEN_EMAILS = 2 # check all unseen eMails
+    CHECK_UNFLAGGED_EMAILS = 4 # check all unflagged eMails
+
     def __init__(self, host, username, password, port=None, use_uid=True,
                  ssl=False, stream=False, timeout=None):
         """Create, connect and login the MailBot.
@@ -36,13 +41,23 @@ def __init__(self, host, username, password, port=None, use_uid=True,
         self.client.normalise_times = False  # deal with UTC everywhere
         self.timeout = timeout
 
-    def get_message_ids(self):
+    def get_message_ids(self, check_pref=6):
         """Return the list of IDs of messages to process."""
-        return self.client.search(['Unseen', 'Unflagged'])
 
-    def get_messages(self):
+        # choose eMails to be regarded in check
+        lstSearchPref = []
+        if check_pref & self.CHECK_PRESENT_EMAILS:
+            lstSearchPref.append('not deleted')
+        if check_pref & self.CHECK_UNSEEN_EMAILS:
+            lstSearchPref.append('Unseen')
+        if check_pref & self.CHECK_UNFLAGGED_EMAILS:
+            lstSearchPref.append('Unflagged')
+
+        return self.client.search(lstSearchPref)
+
+    def get_messages(self, check_pref=6):
         """Return the list of messages to process."""
-        ids = self.get_message_ids()
+        ids = self.get_message_ids(check_pref)
         return self.client.fetch(ids, ['RFC822'])
 
     def process_message(self, message, callback_class, rules):
@@ -51,11 +66,11 @@ def process_message(self, message, callback_class, rules):
         if callback.check_rules():
             return callback.trigger()
 
-    def process_messages(self):
+    def process_messages(self, check_pref=6):
         """Process messages: check which callbacks should be triggered."""
         from . import CALLBACKS_MAP
         self.reset_timeout_messages()
-        messages = self.get_messages()
+        messages = self.get_messages(check_pref)
 
         for uid, msg in messages.items():
             self.mark_processing(uid)