Card #30 - adds method to retrieve attachments from soledad

2 files changed, 42 insertions, 1 deletions

diff --git a/service/pixelated/adapter/soledad_querier.py b/service/pixelated/adapter/soledad_querier.py
index 4123c71d..aa49d3b4 100644
--- a/service/pixelated/adapter/soledad_querier.py
+++ b/service/pixelated/adapter/soledad_querier.py
@@ -15,6 +15,8 @@
 # along with Pixelated. If not, see <http://www.gnu.org/licenses/>.
 from pixelated.adapter.mail import PixelatedMail
 import re
+import base64
+import quopri
 
 
 class SoledadQuerier:
@@ -110,6 +112,20 @@ class SoledadQuerier:
 
         return PixelatedMail.from_soledad(fdoc, hdoc, bdoc, soledad_querier=self, parts=parts)
 
+    def attachment(self, ident, encoding):
+        bdoc = self.soledad.get_from_index('by-type-and-payloadhash', 'cnt', ident)[0]
+        return {'content': self._try_decode(bdoc.content['raw'], encoding),
+                'content-type': bdoc.content['content-type']}
+
+    def _try_decode(self, raw, encoding):
+        encoding = encoding.lower()
+        if encoding == 'base64':
+            return base64.decodestring(raw)
+        elif encoding == 'quoted-printable':
+            return quopri.decodestring(raw)
+        else:
+            return str(raw)
+
     def mails(self, idents):
         fdocs_chash = [(self.soledad.get_from_index('by-type-and-contenthash', 'flags', ident), ident) for ident in idents]
         fdocs_chash = [(result[0], ident) for result, ident in fdocs_chash if result]
@@ -137,7 +153,8 @@ class SoledadQuerier:
 
     def _extract_attachment(self, hdoc, headers_dict):
         content_disposition = headers_dict['Content-Disposition']
-        match = re.compile('.*name=\"(.*)\".*').match(content_disposition)
+        match = re.compile('.*name=\"(.*)\".*').search(content_disposition)
+        filename = ''
         if match:
             filename = match.group(1)
         return {'headers': headers_dict, 'ident': hdoc['phash'], 'name': filename}
diff --git a/service/test/unit/adapter/soledad_querier_test.py b/service/test/unit/adapter/soledad_querier_test.py
index c63a3962..84f2bd5e 100644
--- a/service/test/unit/adapter/soledad_querier_test.py
+++ b/service/test/unit/adapter/soledad_querier_test.py
@@ -18,6 +18,8 @@ from pixelated.adapter.soledad_querier import SoledadQuerier
 from mockito import mock, when, any
 import json
 import os
+import base64
+import quopri
 
 
 class SoledadQuerierTest(unittest.TestCase):
@@ -64,3 +66,25 @@ class SoledadQuerierTest(unittest.TestCase):
         parts = querier._extract_parts(hdoc)
 
         self.assertEquals(bdoc.content['raw'], parts['alternatives'][0]['content'])
+
+    def test_attachment_base64(self):
+        soledad = mock()
+        bdoc = mock()
+        bdoc.content = {'raw': base64.encodestring('esse papo seu ta qualquer coisa'), 'content-type': 'text/plain'}
+        when(soledad).get_from_index('by-type-and-payloadhash', 'cnt', any(unicode)).thenReturn([bdoc])
+        querier = SoledadQuerier(soledad)
+
+        attachment = querier.attachment(u'0400BEBACAFE', 'base64')
+
+        self.assertEquals('esse papo seu ta qualquer coisa', attachment['content'])
+
+    def test_attachment_quoted_printable(self):
+        soledad = mock()
+        bdoc = mock()
+        bdoc.content = {'raw': quopri.encodestring('esse papo seu ta qualquer coisa'), 'content-type': 'text/plain'}
+        when(soledad).get_from_index('by-type-and-payloadhash', 'cnt', any(unicode)).thenReturn([bdoc])
+        querier = SoledadQuerier(soledad)
+
+        attachment = querier.attachment(u'0400BEBACAFE', 'quoted-printable')
+
+        self.assertEquals('esse papo seu ta qualquer coisa', attachment['content'])