From 8b61b34f1ed71c04afbeeb45f08a65d35a18423d Mon Sep 17 00:00:00 2001
From: mnandri <mnandri@eunglick.corporate.thoughtworks.com>
Date: Fri, 18 Dec 2015 18:24:52 +0100
Subject: adapting mail controllers POST and PUT to work with attachements
 Issue #548

---
 service/pixelated/adapter/model/mail.py | 46 +++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 17 deletions(-)

(limited to 'service/pixelated/adapter/model/mail.py')

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
-- 
cgit v1.2.3