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.

7 files changed, 125 insertions, 201 deletions

diff --git a/alot/commands/envelope.py b/alot/commands/envelope.py
index ad6edc3f..f7337968 100644
--- a/alot/commands/envelope.py
+++ b/alot/commands/envelope.py
@@ -57,7 +57,7 @@ class AttachCommand(Command):
 
         logging.info("attaching: %s", files)
         for path in files:
-            envelope.attach(path)
+            envelope.attach_file(path)
         ui.current_buffer.rebuild()
 
 
diff --git a/alot/commands/globals.py b/alot/commands/globals.py
index 65e8973b..6c73e6c9 100644
--- a/alot/commands/globals.py
+++ b/alot/commands/globals.py
@@ -801,7 +801,7 @@ class ComposeCommand(Command):
                 logging.debug('is file')
                 if account.signature_as_attachment:
                     name = account.signature_filename or None
-                    self.envelope.attach(sig, filename=name)
+                    self.envelope.attach_file(sig, filename = name)
                     logging.debug('attached')
                 else:
                     try:
@@ -947,7 +947,7 @@ class ComposeCommand(Command):
         if self.attach:
             for gpath in self.attach:
                 for a in glob.glob(gpath):
-                    self.envelope.attach(a)
+                    self.envelope.attach_file(a)
                     logging.debug('attaching: %s', a)
 
     async def __apply(self, ui):
diff --git a/alot/commands/thread.py b/alot/commands/thread.py
index 0a13b868..e4334a28 100644
--- a/alot/commands/thread.py
+++ b/alot/commands/thread.py
@@ -364,11 +364,9 @@ class ForwardCommand(Command):
 
         else:  # attach original mode
             # attach original msg
-            original_mail = Message()
-            original_mail.set_type('message/rfc822')
-            original_mail['Content-Disposition'] = 'attachment'
-            original_mail.set_payload(mail.as_string(policy=email.policy.SMTP))
-            envelope.attach(Attachment(original_mail))
+            data = mail.as_string()
+            a    = Attachment(data, 'message/rfc822', None, ())
+            envelope.attach(a)
 
         # copy subject
         subject = mail.get('Subject', '')
@@ -854,48 +852,73 @@ class SaveAttachmentCommand(Command):
         self.all = all
         self.path = path
 
+    def _save_attachment(self, ui, path, attachment):
+        is_dir   = os.path.isdir(path)
+        dst_path = path
+        dst      = None
+
+        try:
+            if is_dir and not attachment.filename:
+                # generate a random filename if we don't have one
+                dst  = tempfile.NamedTemporaryFile(delete = False,
+                                                   dir = dst_path)
+                dst_path = f.name
+            else:
+                if is_dir:
+                    dst_path = os.path.join(dst_path, attachment.filename)
+
+                dst = open(dst_path, 'xb')
+
+            dst.write(attachment.data)
+            dst.close()
+        except Exception as e:
+            if dst:
+                os.remove(dst_path)
+                dst.close()
+            if isinstance(e, IOError) or isinstance(e, OSError):
+                ui.notify('Error saving attachment: %s' % str(e),
+                          priority = 'error')
+            else:
+                raise
+        else:
+            ui.notify('saved %s as: %s' % (attachment, dst_path))
+
     async def apply(self, ui):
         pcomplete = PathCompleter()
-        savedir = settings.get('attachment_prefix', '~')
+        savedir   = settings.get('attachment_prefix', '~')
+        path      = self.path
+
         if self.all:
             msg = ui.current_buffer.get_selected_message()
-            if not self.path:
-                self.path = await ui.prompt('save attachments to',
-                                            text=os.path.join(savedir, ''),
-                                            completer=pcomplete)
-            if self.path:
-                if os.path.isdir(os.path.expanduser(self.path)):
-                    for a in msg.iter_attachments():
-                        dest = a.save(self.path)
-                        name = a.get_filename()
-                        if name:
-                            ui.notify('saved %s as: %s' % (name, dest))
-                        else:
-                            ui.notify('saved attachment as: %s' % dest)
-                else:
-                    ui.notify('not a directory: %s' % self.path,
-                              priority='error')
-            else:
-                raise CommandCanceled()
+            if not path:
+                path = await ui.prompt('save attachments to',
+                                       text = os.path.join(savedir, ''),
+                                       completer = pcomplete)
+                if not path:
+                    raise CommandCanceled()
+
+            path = os.path.expanduser(path)
+            if not os.path.isdir(path):
+                ui.notify('not a directory: %s' % path,
+                          priority = 'error')
+                return
+
+            for a in msg.iter_attachments():
+                self._save_attachment(ui, path, a)
         else:  # save focussed attachment
-            attachment = ui.current_buffer.get_selected_attachment()
-            if attachment is not None:
-                filename = attachment.get_filename()
-                if not self.path:
-                    msg = 'save attachment (%s) to ' % filename
-                    initialtext = os.path.join(savedir, filename)
-                    self.path = await ui.prompt(msg,
-                                                completer=pcomplete,
-                                                text=initialtext)
-                if self.path:
-                    try:
-                        dest = attachment.save(self.path)
-                        ui.notify('saved attachment as: %s' % dest)
-                    except (IOError, OSError) as e:
-                        ui.notify(str(e), priority='error')
-                else:
+            a = ui.current_buffer.get_selected_attachment()
+            if not a:
+                return
+
+            if not path:
+                msg = 'save attachment (%s) to ' % a.filename
+                initialtext = os.path.join(savedir, '')
+                path = await ui.prompt(msg, completer = pcomplete,
+                                       text = savedir)
+                if not path:
                     raise CommandCanceled()
 
+            self._save_attachment(ui, path, a)
 
 class OpenAttachmentCommand(Command):
 
@@ -910,12 +933,12 @@ class OpenAttachmentCommand(Command):
 
     async def apply(self, ui):
         logging.info('open attachment')
-        data     = self.attachment.get_data()
-        mimetype = self.attachment.get_content_type()
-        part     = self.attachment.get_mime_representation()
-        fname    = self.attachment.get_filename()
+        data     = self.attachment.data
+        mimetype = self.attachment.content_type
+        params   = self.attachment.params
+        fname    = self.attachment.filename
 
-        h = MailcapHandler(data, mimetype, part.get_params(), fname, 'view')
+        h = MailcapHandler(data, mimetype, params, fname, 'view')
         if not h:
             ui.notify('No handler for: %s' % mimetype)
             return
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
diff --git a/alot/mail/envelope.py b/alot/mail/envelope.py
index 5a798f19..540c0342 100644
--- a/alot/mail/envelope.py
+++ b/alot/mail/envelope.py
@@ -7,12 +7,7 @@ import os
 import re
 import email
 import email.policy
-from email.encoders import encode_7or8bit
 from email.message import MIMEPart
-from email.mime.audio import MIMEAudio
-from email.mime.base import MIMEBase
-from email.mime.image import MIMEImage
-from email.mime.text import MIMEText
 import email.charset as charset
 from urllib.parse import unquote
 
@@ -85,57 +80,6 @@ def _guess_encoding(blob):
     else:
         raise Exception('Unknown magic API')
 
-# TODO: make this work on blobs, not paths
-def _mimewrap(path, filename, ctype):
-    """Take the contents of the given path and wrap them into an email MIME
-    part according to the content type.  The content type is auto detected from
-    the actual file contents and the file name if it is not given.
-
-    :param path: the path to the file contents
-    :type path: str
-    :param filename: the file name to use in the generated MIME part
-    :type filename: str or None
-    :param ctype: the content type of the file contents in path
-    :type ctype: str or None
-    :returns: the message MIME part storing the data from path
-    :rtype: subclasses of email.mime.base.MIMEBase
-    """
-
-    with open(path, 'rb') as f:
-        content = f.read()
-    if not ctype:
-        ctype = helper.guess_mimetype(content)
-        # libmagic < 5.12 incorrectly detects excel/powerpoint files as
-        # 'application/msword' (see #179 and #186 in libmagic bugtracker)
-        # This is a workaround, based on file extension, useful as long
-        # as distributions still ship libmagic 5.11.
-        if (ctype == 'application/msword' and
-            not _libmagic_version_at_least(513)):
-            mimetype, _ = mimetypes.guess_type(path)
-            if mimetype:
-                ctype = mimetype
-
-    maintype, subtype = ctype.split('/', 1)
-    if maintype == 'text':
-        part = MIMEText(content.decode(_guess_encoding(content), 'replace'),
-                        _subtype=subtype,
-                        _charset='utf-8')
-    elif maintype == 'image':
-        part = MIMEImage(content, _subtype=subtype)
-    elif maintype == 'audio':
-        part = MIMEAudio(content, _subtype=subtype)
-    else:
-        part = MIMEBase(maintype, subtype)
-        part.set_payload(content)
-        # Encode the payload using Base64
-        email.encoders.encode_base64(part)
-    # Set the filename parameter
-    if not filename:
-        filename = os.path.basename(path)
-    part.add_header('Content-Disposition', 'attachment',
-                    filename=filename)
-    return part
-
 class Envelope:
     """
     a message that is not yet sent and still editable.
@@ -263,28 +207,37 @@ class Envelope:
         if self.sent_time:
             self.modified_since_sent = True
 
-    def attach(self, attachment, filename=None, ctype=None):
+    def attach_file(self, path, filename = None):
+        with open(path, 'rb') as f:
+            data = f.read()
+
+        ctype = helper.guess_mimetype(data)
+        # libmagic < 5.12 incorrectly detects excel/powerpoint files as
+        # 'application/msword' (see #179 and #186 in libmagic bugtracker)
+        # This is a workaround, based on file extension, useful as long
+        # as distributions still ship libmagic 5.11.
+        if (ctype == 'application/msword' and
+            not _libmagic_version_at_least(513)):
+            mimetype, _ = mimetypes.guess_type(path)
+            if mimetype:
+                ctype = mimetype
+
+        # Set the filename parameter
+        if not filename:
+            filename = os.path.basename(path)
+
+        attachment = Attachment(data, ctype, filename, ())
+        self.attach(attachment)
+
+    def attach(self, attachment):
         """
         attach a file
 
         :param attachment: File to attach, given as
-            :class:`~alot.db.attachment.Attachment` object or path to a file.
-        :type attachment: :class:`~alot.db.attachment.Attachment` or str
-        :param filename: filename to use in content-disposition.
-            Will be ignored if `path` matches multiple files
-        :param ctype: force content-type to be used for this attachment
-        :type ctype: str
+            :class:`~alot.db.attachment.Attachment` object
+        :type attachment: :class:`~alot.db.attachment.Attachment`
         """
-
-        if isinstance(attachment, Attachment):
-            self.attachments.append(attachment)
-        elif isinstance(attachment, str):
-            path = os.path.expanduser(attachment)
-            part = _mimewrap(path, filename, ctype)
-            self.attachments.append(Attachment(part))
-        else:
-            raise TypeError('attach accepts an Attachment or str')
-
+        self.attachments.append(attachment)
         if self.sent_time:
             self.modified_since_sent = True
 
@@ -305,11 +258,9 @@ class Envelope:
 
         # add attachments
         for a in self.attachments:
-            maintype, _, subtype = a.get_content_type().partition('/')
-            fname                = a.get_filename()
-            data                 = a.get_data()
-            mail.add_attachment(data, filename = fname,
-                                maintype = maintype, subtype = subtype)
+            mail.add_attachment(a.data, filename = a.filename,
+                                maintype = a.content_maintype,
+                                subtype  = a.content_subtype)
 
         if self.sign:
             to_sign = mail
@@ -456,7 +407,7 @@ class Envelope:
                                   if os.path.isfile(g)]
                 logging.debug('Attaching: %s', to_attach)
                 for path in to_attach:
-                    self.attach(path)
+                    self.attach_file(path)
                 del self['Attach']
 
         self.body = raw[headerEndPos:].strip()
diff --git a/alot/widgets/globals.py b/alot/widgets/globals.py
index 8b9b3065..7f2a9dcd 100644
--- a/alot/widgets/globals.py
+++ b/alot/widgets/globals.py
@@ -10,7 +10,6 @@ import operator
 import urwid
 
 from ..settings.const import settings
-from ..db.attachment import Attachment
 from ..errors import CompletionError
 
 
@@ -22,8 +21,6 @@ class AttachmentWidget(urwid.WidgetWrap):
     def __init__(self, attachment, selectable=True):
         self._selectable = selectable
         self.attachment = attachment
-        if not isinstance(attachment, Attachment):
-            self.attachment = Attachment(self.attachment)
         att = settings.get_theming_attribute('thread', 'attachment')
         focus_att = settings.get_theming_attribute('thread',
                                                    'attachment_focus')