Merge remote-tracking branch 'alexlovelltroy/size_limiting' into issue_19

django_mailbox/tests/base.py

@@ -8,6 +8,16 @@ from django.test import TestCase
 from django_mailbox import models
 from django_mailbox.models import Mailbox, Message
 
+def get_email_as_text(name):
+    with open(
+        os.path.join(
+            os.path.dirname(__file__),
+            'messages',
+            name,
+        ),
+        'rb'
+    ) as f:
+        return f.read()
 
 class EmailMessageTestCase(TestCase):
     ALLOWED_EXTRA_HEADERS = [

django_mailbox/tests/test_transports.py

@@ -1,12 +1,42 @@
 import mock
-import six
 
-from django_mailbox.tests.base import EmailMessageTestCase
+from django.test.utils import override_settings
+
+from django_mailbox.tests.base import EmailMessageTestCase, get_email_as_text
 from django_mailbox.transports import ImapTransport, Pop3Transport
 
+FAKE_UID_SEARCH_ANSWER = ('OK', ['18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44'])
+FAKE_UID_FETCH_SIZES = ('OK', ['1 (UID 18 RFC822.SIZE 58070000000)', '2 (UID 19 RFC822.SIZE 2593)'])
+FAKE_UID_FETCH_MSG = ('OK', [('1 (UID 18 RFC822 {5807}', get_email_as_text('generic_message.eml') ),])
+FAKE_UID_COPY_MSG = ('OK', ['[COPYUID 1 2 2] (Success)'])
+FAKE_LIST_ARCHIVE_FOLDERS_ANSWERS = ('OK', ['(\\HasNoChildren \\All) "/" "[Gmail]/All Mail"'])
 
-class TestImapTransport(EmailMessageTestCase):
+class IMAPTestCase(EmailMessageTestCase):
     def setUp(self):
+        def imap_server_uid_method(*args):
+            cmd = args[0]
+            arg2 = args[2]
+            if cmd == 'search':
+                return FAKE_UID_SEARCH_ANSWER
+            if cmd == 'copy':
+                return FAKE_UID_COPY_MSG
+            if cmd == 'fetch':
+                if arg2 == '(RFC822.SIZE)':
+                    return FAKE_UID_FETCH_SIZES
+                if arg2 == '(RFC822)':
+                    return FAKE_UID_FETCH_MSG
+        def imap_server_list_method(pattern=None):
+            return FAKE_LIST_ARCHIVE_FOLDERS_ANSWERS
+
+        self.imap_server = mock.Mock()
+        self.imap_server.uid = imap_server_uid_method
+        self.imap_server.list = imap_server_list_method
+        super(IMAPTestCase, self).setUp()
+
+
+class TestImapTransport(IMAPTestCase):
+    def setUp(self):
+        super(TestImapTransport, self).setUp()
         self.arbitrary_hostname = 'one.two.three'
         self.arbitrary_port = 100
         self.ssl = False
@@ -15,42 +45,18 @@ class TestImapTransport(EmailMessageTestCase):
             self.arbitrary_port,
             self.ssl
         )
-        self.transport.server = None
+        self.transport.server = self.imap_server
-        super(TestImapTransport, self).setUp()
 
     def test_get_email_message(self):
-        with mock.patch.object(self.transport, 'server') as server:
+        actual_messages = list(self.transport.get_message())
-            server.search.return_value = (
+        self.assertEqual(len(actual_messages), 27)
-                'OK',
-                [
-                    'One',  # This is totally arbitrary
-                ]
-            )
-            server.fetch.return_value = (
-                'OK',
-                [
-                    [
-                        '1 (RFC822 {8190}',  # Wat?
-                        self._get_email_as_text('generic_message.eml')
-                    ],
-                    ')',
-                ]
-            )
-            actual_messages = list(self.transport.get_message())
-
-        self.assertEqual(len(actual_messages), 1)
-
         actual_message = actual_messages[0]
         expected_message = self._get_email_object('generic_message.eml')
-
         self.assertEqual(expected_message, actual_message)
 
-
+class TestImapArchivedTransport(TestImapTransport):
-class TestImapArchivedTransport(EmailMessageTestCase):
     def setUp(self):
-        self.arbitrary_hostname = 'one.two.three'
+        super(TestImapArchivedTransport, self).setUp()
-        self.arbitrary_port = 100
-        self.ssl = False
         self.archive = 'Archive'
         self.transport = ImapTransport(
             self.arbitrary_hostname,
@@ -58,50 +64,34 @@ class TestImapArchivedTransport(EmailMessageTestCase):
             self.ssl,
             self.archive
         )
-        self.transport.server = None
+        self.transport.server = self.imap_server
-        super(TestImapArchivedTransport, self).setUp()
+
+class TestMaxSizeImapTransport(TestImapTransport):
+
+    @override_settings(DJANGO_MAILBOX_MAX_MESSAGE_SIZE=5807)
+    def setUp(self):
+        super(TestMaxSizeImapTransport, self).setUp()
+
+        self.transport = ImapTransport(
+            self.arbitrary_hostname,
+            self.arbitrary_port,
+            self.ssl,
+        )
+        self.transport.server = self.imap_server
+
+    def test_size_limit(self):
+        all_message_ids = self.transport._get_all_message_ids()
+        small_message_ids = self.transport._get_small_message_ids(all_message_ids)
+        self.assertEqual(len(small_message_ids), 1)
 
     def test_get_email_message(self):
-        with mock.patch.object(self.transport, 'server') as server:
+        actual_messages = list(self.transport.get_message())
-            server.search.return_value = (
-                'OK',
-                [
-                    'One',  # This is totally arbitrary
-                ]
-            )
-            server.fetch.return_value = (
-                'OK',
-                [
-                    [
-                        '1 (RFC822 {8190}',  # Wat?
-                        self._get_email_as_text('generic_message.eml')
-                    ],
-                    ')',
-                ]
-            )
-            server.list.return_value = (
-                'OK',
-                [
-                    '(\\HasNoChildren) "/" "Archive"'
-                ]
-            )
-            server.copy.return_value = (
-                'OK',
-                [
-                    '[COPYUID 1 2 2] (Success)'
-                ]
-            )
-            actual_messages = list(self.transport.get_message())
-
         self.assertEqual(len(actual_messages), 1)
-
         actual_message = actual_messages[0]
         expected_message = self._get_email_object('generic_message.eml')
-
         self.assertEqual(expected_message, actual_message)
 
 
-
 class TestPop3Transport(EmailMessageTestCase):
     def setUp(self):
         self.arbitrary_hostname = 'one.two.three'
@@ -139,3 +129,4 @@ class TestPop3Transport(EmailMessageTestCase):
         expected_message = self._get_email_object('generic_message.eml')
 
         self.assertEqual(expected_message, actual_message)
+

django_mailbox/transports/imap.py

@@ -1,13 +1,25 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
 from imaplib import IMAP4, IMAP4_SSL
 
 from .base import EmailTransport, MessageParseError
 
+from django.conf import settings
+
 
 class ImapTransport(EmailTransport):
     def __init__(self, hostname, port=None, ssl=False, archive=''):
+        MAX_MESSAGE_SIZE = getattr(
+            settings,
+            'DJANGO_MAILBOX_MAX_MESSAGE_SIZE',
+            False
+        )
         self.hostname = hostname
         self.port = port
         self.archive = archive
+        self.MAX_MSG_SIZE = MAX_MESSAGE_SIZE
         if ssl:
             self.transport = IMAP4_SSL
             if not self.port:
@@ -22,10 +34,42 @@ class ImapTransport(EmailTransport):
         typ, msg = self.server.login(username, password)
         self.server.select()
 
-    def get_message(self):
+    def _get_all_message_ids(self):
-        typ, inbox = self.server.search(None, 'ALL')
+        # Fetch all the message uids
+        response, message_ids = self.server.uid('search', None, 'ALL')
+        return message_ids[0].split(' ')
 
-        if not inbox[0]:
+    def _get_small_message_ids(self, message_ids):
+        # Using existing message uids, get the sizes and
+        # return only those that are under the size
+        # limit
+        safe_message_ids = []
+
+        status, data = self.server.uid(
+            'fetch',
+            ','.join(message_ids),
+            '(RFC822.SIZE)'
+        )
+
+        for each_msg in data:
+            try:
+                uid = each_msg.split(' ')[2]
+                size = each_msg.split(' ')[4].rstrip(')')
+                if int(size) <= int(self.MAX_MSG_SIZE):
+                    safe_message_ids.append(uid)
+            except ValueError as e:
+                logger.warning("ValueError: %s working on %s" % (e, each_msg[0]))
+                pass
+        return safe_message_ids
+
+    def get_message(self):
+        message_ids = self._get_all_message_ids()
+
+        # Limit the uids to the small ones if we care about that
+        if self.MAX_MSG_SIZE:
+            message_ids = self._get_small_message_ids(message_ids)
+
+        if not message_ids:
             return
 
         if self.archive:
@@ -34,17 +78,17 @@ class ImapTransport(EmailTransport):
                 # If the archive folder does not exist, create it
                 self.server.create(self.archive)
 
-        for key in inbox[0].split():
+        for uid in message_ids:
             try:
-                typ, msg_contents = self.server.fetch(key, '(RFC822)')
+                typ, msg_contents = self.server.uid('fetch', uid, '(RFC822)')
                 message = self.get_email_from_bytes(msg_contents[0][1])
                 yield message
             except MessageParseError:
                 continue
 
             if self.archive:
-                self.server.copy(key, self.archive)
+                self.server.uid('copy', uid, self.archive)
 
-            self.server.store(key, "+FLAGS", "\\Deleted")
+            self.server.uid('store', uid, "+FLAGS", "\\Deleted")
         self.server.expunge()
         return

docs/topics/settings.rst

@@ -58,3 +58,11 @@ Settings
     a message instance's ``get_email_object`` method being called.  Value of
     this field is the primary key of the ``django_mailbox.MessageAttachment``
     instance currently storing this payload component's contents.
+
+* ``DJANGO_MAILBOX_MAX_MESSAGE_SIZE``
+
+  * Default: ``False``
+  * Type: ``integer``
+  * If this is set, it will be read as a number of
+    bytes.  Any messages above that size will not be
+    downloaded.  ``2000000`` is 2 Megabytes.