adapting mail controllers POST and PUT to work with attachements

Issue #548

6 files changed, 161 insertions, 44 deletions

diff --git a/service/pixelated/adapter/model/mail.py b/service/pixelated/adapter/model/mail.py
index bb47a434..b1111ad4 100644
--- a/service/pixelated/adapter/model/mail.py
+++ b/service/pixelated/adapter/model/mail.py
@@ -20,6 +20,7 @@ from email import message_from_file
 from email.mime.text import MIMEText
 from email.header import Header
 from email.MIMEMultipart import MIMEMultipart
+from email.mime.nonmultipart import MIMENonMultipart
 from pycryptopp.hash import sha256
 import leap.mail.walk as walk
 from pixelated.adapter.model.status import Status
@@ -68,13 +69,21 @@ class Mail(object):
         return self.fdoc.content.get('mbox', 'INBOX')
 
     def _encode_header_value_list(self, header_value_list):
-        return [self._encode_header_value(v) for v in header_value_list]
+        encoded_header_list = [self._encode_header_value(v) for v in header_value_list]
+        return ', '.join(encoded_header_list)
 
     def _encode_header_value(self, header_value):
         if isinstance(header_value, unicode):
             return str(Header(header_value, 'utf-8'))
+        return str(header_value)
+
+    def _attach_body_mimetext_to(self, mime_multipart, body_to_use=None):
+        body_to_use = body_to_use or self.body
+        if isinstance(body_to_use, list):
+            for part in body_to_use:
+                mime_multipart.attach(MIMEText(part['raw'], part['content-type']))
         else:
-            return str(header_value)
+            mime_multipart.attach(MIMEText(body_to_use, 'plain', self._charset()))
 
     @property
     def _mime_multipart(self):
@@ -83,24 +92,28 @@ class Mail(object):
         mime = MIMEMultipart()
         for key, value in self.headers.items():
             if isinstance(value, list):
-                mime[str(key)] = ', '.join(self._encode_header_value_list(value))
+                mime[str(key)] = self._encode_header_value_list(value)
             else:
                 mime[str(key)] = self._encode_header_value(value)
 
-        try:
-            body_to_use = self.body
-        except AttributeError:
-            body_to_use = self.text_plain_body
+        body_to_use = getattr(self, 'body', None) or getattr(self, 'text_plain_body', None)
+        self._attach_body_mimetext_to(mime, body_to_use)
+
+        if self._attachments:
+            for attachment in self._attachments:
+                major, sub = attachment['content-type'].split('/')
+                a = MIMENonMultipart(major, sub, Content_Disposition='attachment; filename=%s' % attachment['filename'])
+                a.set_payload(attachment['raw'])
+                mime.attach(a)
 
-        mime.attach(MIMEText(body_to_use, 'plain', self._charset()))
         self._mime = mime
         return mime
 
     def _charset(self):
-        if 'content_type' in self.headers and 'charset' in self.headers['content_type']:
-            return self._parse_charset_header(self.headers['content_type'])
-        else:
-            return 'utf-8'
+        content_type = self.headers.get('content_type', {})
+        if 'charset' in content_type:
+            return self._parse_charset_header(content_type)
+        return 'utf-8'
 
     def _parse_charset_header(self, charset_header, default_charset='utf-8'):
         try:
@@ -129,6 +142,7 @@ class InputMail(Mail):
         self.headers = {}
         self.body = ''
         self._status = []
+        self._attachments = []
 
     @property
     def ident(self):
@@ -151,11 +165,7 @@ class InputMail(Mail):
             mime_multipart['From'] = self.headers['From']
 
         mime_multipart['Date'] = self.headers['Date']
-        if type(self.body) is list:
-            for part in self.body:
-                mime_multipart.attach(MIMEText(part['raw'], part['content-type']))
-        else:
-            mime_multipart.attach(MIMEText(self.body, 'plain', 'utf-8'))
+        self._attach_body_mimetext_to(mime_multipart)
         return mime_multipart
 
     def to_smtp_format(self):
@@ -174,6 +184,7 @@ class InputMail(Mail):
         })
 
     @staticmethod
+    # def from_dict(mail_dict, attachments):
     def from_dict(mail_dict):
         input_mail = InputMail()
         input_mail.headers = {key.capitalize(): value for key, value in mail_dict.get('header', {}).items()}
@@ -184,6 +195,7 @@ class InputMail(Mail):
         input_mail.body = mail_dict.get('body', '')
         input_mail.tags = set(mail_dict.get('tags', []))
         input_mail._status = set(mail_dict.get('status', []))
+        input_mail._attachments = mail_dict.get('attachments', [])
         return input_mail
 
     @staticmethod
diff --git a/service/pixelated/adapter/services/draft_service.py b/service/pixelated/adapter/services/draft_service.py
index 98444442..504d92db 100644
--- a/service/pixelated/adapter/services/draft_service.py
+++ b/service/pixelated/adapter/services/draft_service.py
@@ -33,3 +33,8 @@ class DraftService(object):
         if removed:
             new_draft = yield self.create_draft(input_mail)
             defer.returnValue(new_draft)
+
+    def process_draft(self, ident, input_mail):
+        if ident:
+            return self.update_draft(ident, input_mail)
+        return self.create_draft(input_mail)
diff --git a/service/pixelated/resources/mails_resource.py b/service/pixelated/resources/mails_resource.py
index 613c5e6f..8874fe16 100644
--- a/service/pixelated/resources/mails_resource.py
+++ b/service/pixelated/resources/mails_resource.py
@@ -154,14 +154,6 @@ class MailsResource(Resource):
         return NOT_DONE_YET
 
     def render_POST(self, request):
-        content_dict = json.loads(request.content.read())
-
-        deferred = self._mail_service.send_mail(content_dict)
-
-        def onSuccess(sent_mail):
-            data = sent_mail.as_dict()
-            respond_json_deferred(data, request)
-
         def onError(error):
             if isinstance(error.value, SMTPDownException):
                 respond_json_deferred({'message': str(error.value)}, request, status_code=503)
@@ -169,31 +161,45 @@ class MailsResource(Resource):
                 err(error, 'something failed')
                 respond_json_deferred({'message': 'an error occurred while sending'}, request, status_code=422)
 
-        deferred.addCallback(onSuccess)
+        deferred = self._handle_post(request)
         deferred.addErrback(onError)
 
         return server.NOT_DONE_YET
 
     def render_PUT(self, request):
-        content_dict = json.loads(request.content.read())
-        _mail = InputMail.from_dict(content_dict)
-        draft_id = content_dict.get('ident')
-
         def onError(error):
                 err(error, 'error saving draft')
                 respond_json_deferred("", request, status_code=422)
 
-        def updateCallback(pixelated_mail):
-            if not pixelated_mail:
-                respond_json_deferred("", request, status_code=422)
-            else:
-                respond_json_deferred({'ident': pixelated_mail.ident}, request)
-
-        if draft_id:
-            deferred = self._draft_service.update_draft(draft_id, _mail)
-        else:
-            deferred = self._draft_service.create_draft(_mail)
-        deferred.addCallback(updateCallback)
+        deferred = self._handle_put(request)
         deferred.addErrback(onError)
 
         return server.NOT_DONE_YET
+
+    @defer.inlineCallbacks
+    def _fetch_attachment_contents(self, attachments):
+        for attachment in attachments:
+            retrieved_attachment = yield self._mail_service.attachment(attachment['id'])
+            attachment['raw'] = retrieved_attachment['content']
+
+    @defer.inlineCallbacks
+    def _handle_post(self, request):
+        content_dict = json.loads(request.content.read())
+        self._fetch_attachment_contents(content_dict.get('attachments', []))
+
+        sent_mail = yield self._mail_service.send_mail(content_dict)
+        respond_json_deferred(sent_mail.as_dict(), request)
+
+    @defer.inlineCallbacks
+    def _handle_put(self, request):
+        content_dict = json.loads(request.content.read())
+        self._fetch_attachment_contents(content_dict.get('attachments', []))
+
+        _mail = InputMail.from_dict(content_dict)
+        draft_id = content_dict.get('ident')
+        pixelated_mail = yield self._draft_service.process_draft(draft_id, _mail)
+
+        if not pixelated_mail:
+            respond_json_deferred("", request, status_code=422)
+        else:
+            respond_json_deferred({'ident': pixelated_mail.ident}, request)
diff --git a/service/test/integration/test_draft_service.py b/service/test/integration/test_draft_service.py
index 00b1fcfe..4f4a9181 100644
--- a/service/test/integration/test_draft_service.py
+++ b/service/test/integration/test_draft_service.py
@@ -24,9 +24,24 @@ class DraftServiceTest(SoledadTestBase):
     @defer.inlineCallbacks
     def test_store_and_load_draft(self):
         input_mail = MailBuilder().with_body('some test text').build_input_mail()
-
-        stored_draft = yield self.draft_service.create_draft(input_mail)
+        draft_id = None
+        stored_draft = yield self.draft_service.process_draft(draft_id, input_mail)
 
         draft = yield self.mail_store.get_mail(stored_draft.ident, include_body=True)
 
         self.assertEqual('some test text', draft.body)
+
+    @defer.inlineCallbacks
+    def test_update_draft(self):
+        input_mail = MailBuilder().with_body('some test text').build_input_mail()
+        saved_mail = yield self.mail_store.add_mail('DRAFTS', input_mail.raw)
+        draft_id = saved_mail.ident
+        new_email = MailBuilder().with_body('other test text').with_ident(draft_id).build_input_mail()
+
+        stored_draft = yield self.draft_service.process_draft(draft_id, new_email)
+
+        old_draft = yield self.mail_store.get_mail(draft_id, include_body=True)
+        draft = yield self.mail_store.get_mail(stored_draft.ident, include_body=True)
+
+        self.assertIsNone(old_draft)
+        self.assertEqual('other test text', draft.body)
diff --git a/service/test/unit/adapter/test_mail.py b/service/test/unit/adapter/test_mail.py
index 5c5d4c59..006bde77 100644
--- a/service/test/unit/adapter/test_mail.py
+++ b/service/test/unit/adapter/test_mail.py
@@ -14,6 +14,9 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with Pixelated. If not, see <http://www.gnu.org/licenses/>.
+from email.mime.nonmultipart import MIMENonMultipart
+from email.mime.multipart import MIMEMultipart
+
 from twisted.trial import unittest
 
 import pixelated.support.date
@@ -50,6 +53,23 @@ def multipart_mail_dict():
     }
 
 
+def with_attachment_mail_dict():
+    return {
+        'attachments': [{'content-type': 'text/plain', 'filename': 'ayoyo.txt', 'raw': 'Hamburg Ayoyoyooooo!!!', 'id': 'some_attachment_id'},
+                        {'content-type': 'text/html', 'filename': 'hello.html', 'raw': '<p>Hello html Hamburg!</p>', 'id': 'other_attachment_id'}],
+        'body': [{'content-type': 'plain', 'raw': 'Hello world!'},
+                 {'content-type': 'html', 'raw': '<p>Hello html world!</p>'}],
+        'header': {
+            'cc': ['cc@pixelated.org', 'anothercc@pixelated.org'],
+            'to': ['to@pixelated.org', 'anotherto@pixelated.org'],
+            'bcc': ['bcc@pixelated.org', 'anotherbcc@pixelated.org'],
+            'subject': 'Oi',
+        },
+        'ident': '',
+        'tags': ['sent']
+    }
+
+
 class InputMailTest(unittest.TestCase):
 
     def test_to_mime_multipart_should_add_blank_fields(self):
@@ -120,3 +140,17 @@ class InputMailTest(unittest.TestCase):
 
         self.assertRegexpMatches(mime_multipart.as_string(), part_one)
         self.assertRegexpMatches(mime_multipart.as_string(), part_two)
+
+    def test_raw_with_attachment_data(self):
+        input_mail = InputMail.from_dict(with_attachment_mail_dict())
+
+        attachment = MIMENonMultipart('text', 'plain', Content_Disposition='attachment; filename=ayoyo.txt')
+        attachment.set_payload('Hello World')
+        mail = MIMEMultipart()
+        mail.attach(attachment)
+
+        part_one = 'Content-Type: text/plain; Content-Disposition="attachment; filename=ayoyo.txt"\nMIME-Version: 1.0\n\nHamburg Ayoyoyooooo!!!'
+        part_two = 'Content-Type: text/html; Content-Disposition="attachment; filename=hello.html"\nMIME-Version: 1.0\n\n<p>Hello html Hamburg!</p>'
+
+        self.assertRegexpMatches(input_mail.raw, part_one)
+        self.assertRegexpMatches(input_mail.raw, part_two)
diff --git a/service/test/unit/resources/test_mails_resource.py b/service/test/unit/resources/test_mails_resource.py
index 02b17bf1..44aebd4d 100644
--- a/service/test/unit/resources/test_mails_resource.py
+++ b/service/test/unit/resources/test_mails_resource.py
@@ -24,7 +24,7 @@ from twisted.internet import defer
 from mock import patch
 
 
-class TestArchiveResource(unittest.TestCase):
+class TestMailsResource(unittest.TestCase):
     def setUp(self):
         self.mail_service = mock()
 
@@ -49,3 +49,48 @@ class TestArchiveResource(unittest.TestCase):
 
         d.addCallback(assert_response)
         return d
+
+    @patch('leap.common.events.register')
+    def test_render_PUT_should_store_draft_with_attachments(self, mock_register):
+        request = DummyRequest(['/mails'])
+        request.method = 'PUT'
+        content = mock()
+        when(content).read().thenReturn('{"attachments": [{"id": "some fake attachment id"}]}')
+        when(self.mail_service).attachment('some fake attachment id').thenReturn(defer.Deferred())
+        request.content = content
+
+        mails_resource = MailsResource(self.mail_service, mock())
+        mails_resource.isLeaf = True
+        web = DummySite(mails_resource)
+        d = web.get(request)
+
+        def assert_response(_):
+            verify(self.mail_service).attachment('some fake attachment id')
+
+        d.addCallback(assert_response)
+        return d
+
+    @patch('leap.common.events.register')
+    def test_render_POST_should_send_email_with_attachments(self, mock_register):
+        request = DummyRequest(['/mails'])
+        request.method = 'POST'
+        content = mock()
+        when(content).read().thenReturn('{"attachments": [{"id": "some fake attachment id"}]}')
+        when(self.mail_service).attachment('some fake attachment id').thenReturn(defer.succeed({"content": "some content"}))
+        as_dictable = mock()
+        when(as_dictable).as_dict().thenReturn({})
+        when(self.mail_service).send_mail({"attachments": [{"id": "some fake attachment id", "raw": "some content"}]})\
+            .thenReturn(defer.succeed(as_dictable))
+        request.content = content
+
+        mails_resource = MailsResource(self.mail_service, mock())
+        mails_resource.isLeaf = True
+        web = DummySite(mails_resource)
+        d = web.get(request)
+
+        def assert_response(_):
+            verify(self.mail_service).attachment('some fake attachment id')
+            verify(self.mail_service).send_mail({"attachments": [{"id": "some fake attachment id", "raw": "some content"}]})
+
+        d.addCallback(assert_response)
+        return d