db/attachment: simplify the Attachment class

Make it a plain container around raw data and a few bits of metadata,
rather than around a whole MIME part.

2 files changed, 24 insertions, 71 deletions

diff --git a/alot/db/attachment.py b/alot/db/attachment.py
index 0ebe6fc9..5c993a68 100644
--- a/alot/db/attachment.py
+++ b/alot/db/attachment.py
@@ -2,14 +2,8 @@
 # Copyright © 2018 Dylan Baker
 # This file is released under the GNU GPL, version 3 or a later revision.
 # For further details see the COPYING file
-import os
-import tempfile
-import email.charset as charset
-from copy import deepcopy
 
-from ..helper import guess_mimetype
-
-charset.add_charset('utf-8', charset.QP, charset.QP, 'utf-8')
+import os.path
 
 def _humanize_size(size):
     """Create a nice human readable representation of the given number
@@ -35,71 +29,28 @@ def _humanize_size(size):
 class Attachment:
     """represents a mail attachment"""
 
-    def __init__(self, emailpart):
-        """
-        :param emailpart: a non-multipart email that is the attachment
-        :type emailpart: :class:`email.message.Message`
-        """
-        self.part = emailpart
-
-    def __str__(self):
-        return '%s:%s (%s)' % (self.get_content_type(),
-                               self.get_filename(),
-                               _humanize_size(self.get_size()))
-
-    def get_filename(self):
-        """
-        return name of attached file.
-        If the content-disposition header contains no file name,
-        this returns `None`
-        """
-        fname = self.part.get_filename()
-        if fname:
-            return os.path.basename(fname)
-        return None
-
-    def get_content_type(self):
-        """mime type of the attachment part"""
-        ctype = self.part.get_content_type()
-        # replace underspecified mime description by a better guess
-        if ctype in ['octet/stream', 'application/octet-stream',
-                     'application/octetstream']:
-            ctype = guess_mimetype(self.get_data())
-        return ctype
+    data         = None
+    content_type = None
+    filename     = None
+    params       = None
 
-    def get_size(self):
-        """returns attachments size in bytes"""
-        return len(self.part.get_payload())
+    def __init__(self, data, ctype, filename, params):
+        self.data         = data
+        self.content_type = ctype
+        self.params       = params
 
-    def save(self, path):
-        """
-        save the attachment to disk. Uses :meth:`~get_filename` in case path
-        is a directory
-        """
-        filename = self.get_filename()
-        path = os.path.expanduser(path)
-        if os.path.isdir(path):
-            if filename:
-                basename = os.path.basename(filename)
-                file_ = open(os.path.join(path, basename), "wb")
-            else:
-                file_ = tempfile.NamedTemporaryFile(delete=False, dir=path)
-        else:
-            file_ = open(path, "wb")  # this throws IOErrors for invalid path
-        self.write(file_)
-        file_.close()
-        return file_.name
+        # make sure the filename is a relative path
+        # that does not go upwards
+        filename = os.path.normpath(filename)
+        if filename.startswith('/') or filename.startswith('..'):
+            raise ValueError('Dangerous attachment filename: %s' % filename)
 
-    def write(self, fhandle):
-        """writes content to a given filehandle"""
-        fhandle.write(self.get_data())
+        self.filename = filename
 
-    def get_data(self):
-        """return data blob from wrapped file"""
-        return self.part.get_payload(decode=True)
+    def __str__(self):
+        ret = self.content_type
+        if self.filename:
+            ret += ':' + self.filename
+        ret += ' (%s)' % _humanize_size(len(self.data))
 
-    def get_mime_representation(self):
-        """returns mime part that constitutes this attachment"""
-        part = deepcopy(self.part)
-        part.set_param('maxlinelen', '78', header='Content-Disposition')
-        return part
+        return ret
diff --git a/alot/db/message.py b/alot/db/message.py
index 8ca8f2b3..5db42c0a 100644
--- a/alot/db/message.py
+++ b/alot/db/message.py
@@ -121,7 +121,9 @@ class _MimeTree:
             cd = part.get_content_disposition()
             fn = part.get_filename()
             if cd == 'attachment' or fn is not None:
-                self.attachment = Attachment(part)
+                data = part.get_content()
+                self.attachment = Attachment(data, self.content_type,
+                                             fn, part.get_params())
 
     def __str__(self):
         return 'MimePart(%s)' % self.content_type