diff --git a/django_mailbox/tests/base.py b/django_mailbox/tests/base.py index a559b9e..72c3249 100644 --- a/django_mailbox/tests/base.py +++ b/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 = [ diff --git a/django_mailbox/tests/test_transports.py b/django_mailbox/tests/test_transports.py index 7f044d4..b701eae 100644 --- a/django_mailbox/tests/test_transports.py +++ b/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 - super(TestImapTransport, self).setUp() + self.transport.server = self.imap_server def test_get_email_message(self): - with mock.patch.object(self.transport, 'server') as server: - 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') - ], - ')', - ] - ) - actual_messages = list(self.transport.get_message()) - - self.assertEqual(len(actual_messages), 1) - + actual_messages = list(self.transport.get_message()) + self.assertEqual(len(actual_messages), 27) actual_message = actual_messages[0] expected_message = self._get_email_object('generic_message.eml') - self.assertEqual(expected_message, actual_message) - -class TestImapArchivedTransport(EmailMessageTestCase): +class TestImapArchivedTransport(TestImapTransport): def setUp(self): - self.arbitrary_hostname = 'one.two.three' - self.arbitrary_port = 100 - self.ssl = False + super(TestImapArchivedTransport, self).setUp() self.archive = 'Archive' self.transport = ImapTransport( self.arbitrary_hostname, @@ -58,50 +64,34 @@ class TestImapArchivedTransport(EmailMessageTestCase): self.ssl, self.archive ) - self.transport.server = None - super(TestImapArchivedTransport, self).setUp() + self.transport.server = self.imap_server + +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: - 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()) - + 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) + diff --git a/django_mailbox/transports/imap.py b/django_mailbox/transports/imap.py index 078392e..b6255b1 100644 --- a/django_mailbox/transports/imap.py +++ b/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): - typ, inbox = self.server.search(None, 'ALL') + def _get_all_message_ids(self): + # 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 diff --git a/docs/topics/settings.rst b/docs/topics/settings.rst index 6023046..290e884 100644 --- a/docs/topics/settings.rst +++ b/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.