From 309fb25e9b089618c37f1a741fa6009cce54ac9e Mon Sep 17 00:00:00 2001 From: Anton Khirnov Date: Sat, 30 Jan 2021 16:33:56 +0100 Subject: 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. --- alot/commands/envelope.py | 2 +- alot/commands/globals.py | 4 +- alot/commands/thread.py | 113 ++++++++++++++++++++++++++++------------------ alot/db/attachment.py | 91 +++++++++---------------------------- alot/db/message.py | 4 +- alot/mail/envelope.py | 109 ++++++++++++-------------------------------- alot/widgets/globals.py | 3 -- 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') -- cgit v1.2.3