Adding new settings allowing one to specify a list of mimetypes that are permitted to be stored in the message body of the e-mail message model.

django_mailbox/models.py

@@ -1,4 +1,5 @@
 import email
+from email.message import Message as EmailMessage
 from email.utils import formatdate
 import rfc822
 import urllib
@@ -20,6 +21,24 @@ SKIPPED_EXTENSIONS = getattr(
     'DJANGO_MAILBOX_SKIPPED_EXTENSIONS', 
     ['.p7s']
 )
+STRIP_UNALLOWED_MIMETYPES = getattr(
+    settings,
+    'DJANGO_MAILBOX_STRIP_UNALLOWED_MIMETYPES',
+    False
+)
+ALLOWED_MIMETYPES = getattr(
+    settings,
+    'DJANGO_MAILBOX_ALLOWED_MIMETYPES',
+    [
+        'text/plain',
+        'text/html'
+    ]
+)
+ALTERED_MESSAGE_HEADER = getattr(
+    settings,
+    'DJANGO_MAILBOX_ALTERED_MESSAGE_HEADER',
+    'X-Django-Mailbox-Altered-Message'
+)
 
 class ActiveMailboxManager(models.Manager):
     def get_query_set(self):
@@ -155,6 +174,29 @@ class Mailbox(models.Model):
         msg.save()
         return msg
 
+    def _filter_message_body(self, message):
+        if not message.is_multipart() or not STRIP_UNALLOWED_MIMETYPES:
+            return message
+        stripped_content = {}
+        new = EmailMessage()
+        for header, value in message.items():
+            new[header] = value
+        for part in message.walk():
+            content_type = part.get_content_type()
+            print content_type
+            if not content_type in ALLOWED_MIMETYPES:
+                if content_type not in stripped_content:
+                    stripped_content[content_type] = 0
+                stripped_content[content_type] = (
+                    stripped_content[content_type] + 1
+                )
+                continue
+            new.attach(part)
+        new[ALTERED_MESSAGE_HEADER] = 'Stripped ' + ', '.join(
+            ['%s*%s' % (key, value) for key, value in stripped_content.items()]
+        )
+        return new
+
     def _process_message(self, message):
         msg = Message()
         msg.mailbox = self
@@ -162,14 +204,8 @@ class Mailbox(models.Model):
         msg.message_id = message['message-id'][0:255]
         msg.from_header = message['from']
         msg.to_header = message['to']
-        # if message is multipart want to only store the actual message body in database (not attachments)
-        if message.is_multipart():
-            for part in message.walk():
-                if part.get_content_type() == 'text/plain' or part.get_content_type() == 'text/html':
-                    msg.body = part.get_payload()
-        # else store the whole message in database
-        else:
-            msg.body = message.as_string()
+        message = self._filter_message_body(message)
+        msg.body = message.as_string()
         if message['in-reply-to']:
             try:
                 msg.in_reply_to = Message.objects.filter(message_id=message['in-reply-to'])[0]

django_mailbox/tests/__init__.py

@@ -5,9 +5,10 @@ import shutil
 from django.db import models
 from django.test import TestCase
 
+import django_mailbox
 from django_mailbox.models import Mailbox, Message
 
-class TestProcessMessage(TestCase):
+class EmailMessageTestCase(TestCase):
     def _get_email_object(self, name):
         with open(os.path.join(os.path.dirname(__file__), name), 'r') as f:
             return email.message_from_string(
@@ -17,7 +18,8 @@ class TestProcessMessage(TestCase):
     def tearDown(self):
         for message in Message.objects.all():
             message.delete()
-    
+
+class TestProcessEmail(EmailMessageTestCase):
     def test_message_without_attachments(self):
         message = self._get_email_object('generic_message.eml')
 
@@ -75,3 +77,49 @@ class TestProcessMessage(TestCase):
             expected_results,
             actual_results,
         )
+
+class TestFilterMessageBody(EmailMessageTestCase):
+    def setUp(self):
+        django_mailbox.models.STRIP_UNALLOWED_MIMETYPES = True
+        super(TestFilterMessageBody, self).setUp()
+
+    def tearDown(self):
+        django_mailbox.models.STRIP_UNALLOWED_MIMETYPES = False
+        super(TestFilterMessageBody, self).tearDown()
+
+    def test_filter_message_does_not_filter_message_if_disabled(self):
+        django_mailbox.models.STRIP_UNALLOWED_MIMETYPES = False
+        message = self._get_email_object('message_with_attachment.eml')
+        mailbox = Mailbox.objects.create()
+
+        self.assertEquals(
+            message.as_string(),
+            mailbox._filter_message_body(message).as_string()
+        )
+
+    def test_filter_message_removes_unknown_content_if_disabled(self):
+        # The below is the _same_ as message_with_attachment.eml, but missing
+        # its attached png image, and adding the expected message altered header.
+        message_without_non_plaintext = (
+            "MIME-Version: 1.0\n"
+            "Received: by 10.221.0.211 with HTTP; Sun, 20 Jan 2013 12:07:07 -0800 (PST)\n"
+            "X-Originating-IP: [24.22.122.177]\n"
+            "Date: Sun, 20 Jan 2013 12:07:07 -0800\n"
+            "Delivered-To: test@adamcoddington.net\n"
+            "Message-ID: <CAMdmm+jYCgrxrekAxszmDnBjAytcBym-Ec+uM-+HEtzuKy=M_g@mail.gmail.com>\n"
+            "Subject: Message With Attachment\n"
+            "From: Adam Coddington <test@adamcoddington.net>\n"
+            "To: Adam Coddington <test@adamcoddington.net>\n"
+            "Content-Type: multipart/mixed; boundary=047d7b33dd729737fe04d3bde348\n"
+            "X-Django-Mailbox-Altered-Message: Stripped image/png*1, multipart/mixed*1\n"
+            "\n--047d7b33dd729737fe04d3bde348\n"
+            "Content-Type: text/plain; charset=UTF-8\n\n"
+            "This message has an attachment.\n\n--047d7b33dd729737fe04d3bde348--"
+        )
+        message = self._get_email_object('message_with_attachment.eml')
+        mailbox = Mailbox.objects.create()
+
+        self.assertEquals(
+            message_without_non_plaintext,
+            mailbox._filter_message_body(message).as_string()
+        )