extracted a leap attachment store, handling all attachment responsibilities, including saving attachments.

Issue #548

2 files changed, 61 insertions, 22 deletions

diff --git a/service/pixelated/adapter/mailstore/leap_attachment_store.py b/service/pixelated/adapter/mailstore/leap_attachment_store.py
new file mode 100644
index 00000000..86121db9
--- /dev/null
+++ b/service/pixelated/adapter/mailstore/leap_attachment_store.py
@@ -0,0 +1,61 @@
+
+import quopri
+import base64
+from email import encoders
+from leap.mail.adaptors.soledad import SoledadMailAdaptor, ContentDocWrapper
+from twisted.internet import defer
+from email.mime.nonmultipart import MIMENonMultipart
+from email.mime.multipart import MIMEMultipart
+from leap.mail.mail import Message
+
+
+class LeapAttachmentStore(object):
+
+    def __init__(self, soledad):
+        self.soledad = soledad
+
+    @defer.inlineCallbacks
+    def get_mail_attachment(self, attachment_id):
+        results = yield self.soledad.get_from_index('by-type-and-payloadhash', 'cnt', attachment_id) if attachment_id else []
+        if len(results):
+            content = ContentDocWrapper(**results[0].content)
+            defer.returnValue({'content-type': content.content_type, 'content': self._try_decode(
+                content.raw, content.content_transfer_encoding)})
+        else:
+            raise ValueError('No attachment with id %s found!' % attachment_id)
+
+    @defer.inlineCallbacks
+    def add_attachment(self, content, content_type):
+        cdoc = self._attachment_to_cdoc(content, content_type)
+        yield self.soledad.create_doc(cdoc.serialize(), doc_id=cdoc.phash)
+        defer.returnValue(cdoc.phash)
+
+    def _try_decode(self, raw, encoding):
+        encoding = encoding.lower()
+        if encoding == 'base64':
+            data = base64.decodestring(raw)
+        elif encoding == 'quoted-printable':
+            data = quopri.decodestring(raw)
+        else:
+            data = str(raw)
+
+        return bytearray(data)
+
+    def _attachment_to_cdoc(self, content, content_type, encoder=encoders.encode_base64):
+        major, sub = content_type.split('/')
+        attachment = MIMENonMultipart(major, sub)
+        attachment.set_payload(content)
+        encoder(attachment)
+        attachment.add_header('Content-Disposition', 'attachment', filename='does_not_matter.txt')
+
+        pseudo_mail = MIMEMultipart()
+        pseudo_mail.attach(attachment)
+
+        tmp_mail = SoledadMailAdaptor().get_msg_from_string(MessageClass=Message, raw_msg=pseudo_mail.as_string())
+
+        cdoc = tmp_mail.get_wrapper().cdocs[1]
+        return cdoc
+
+    def _calc_attachment_id_(self, content, content_type, encoder=encoders.encode_base64):
+        cdoc = self._attachment_to_cdoc(content, content_type, encoder)
+        return cdoc.phash
diff --git a/service/pixelated/adapter/mailstore/leap_mailstore.py b/service/pixelated/adapter/mailstore/leap_mailstore.py
index 14b0e417..2660a208 100644
--- a/service/pixelated/adapter/mailstore/leap_mailstore.py
+++ b/service/pixelated/adapter/mailstore/leap_mailstore.py
@@ -16,7 +16,6 @@
 import base64
 from email.header import decode_header
 from email.utils import parseaddr
-import quopri
 from uuid import uuid4
 
 import re
@@ -222,27 +221,6 @@ class LeapMailStore(MailStore):
         return defer.gatherResults(deferreds, consumeErrors=True)
 
     @defer.inlineCallbacks
-    def get_mail_attachment(self, attachment_id):
-        results = yield self.soledad.get_from_index('by-type-and-payloadhash', 'cnt', attachment_id) if attachment_id else []
-        if len(results):
-            content = ContentDocWrapper(**results[0].content)
-            defer.returnValue({'content-type': content.content_type, 'content': self._try_decode(
-                content.raw, content.content_transfer_encoding)})
-        else:
-            raise ValueError('No attachment with id %s found!' % attachment_id)
-
-    def _try_decode(self, raw, encoding):
-        encoding = encoding.lower()
-        if encoding == 'base64':
-            data = base64.decodestring(raw)
-        elif encoding == 'quoted-printable':
-            data = quopri.decodestring(raw)
-        else:
-            data = str(raw)
-
-        return bytearray(data)
-
-    @defer.inlineCallbacks
     def update_mail(self, mail):
         message = yield self._fetch_msg_from_soledad(mail.mail_id)
         message.get_wrapper().set_tags(tuple(mail.tags))