1 files changed, 39 insertions, 30 deletions

diff --git a/src/jwebmail/render_mail.py b/src/jwebmail/render_mail.py
index f76c40a..3f72d19 100644
--- a/src/jwebmail/render_mail.py
+++ b/src/jwebmail/render_mail.py
@@ -40,18 +40,20 @@ def render_multipart_alternative(_subtype, content, path):
     T += f"<li class=is-active data=0><a>{to_mime_type(init['head'])}</a></li>"
 
     C += "<div class=jwm-mail-body-multipart-alternative-body>\n"
-    C += mime_render(
-        *to_mime_types(init["head"]), init["body"], path + (len(parts) - 1,)
-    )
+    if init["head"].get("content_disposition", "").lower() == "attachment":
+        C += _as_download(init, path + (len(parts) - 1,))
+    else:
+        C += mime_render(init["head"], init.get("body"), path + (len(parts) - 1,))
     C += "</div>\n"
 
     for i, r in enumerate(rest, 1):
         T += f"<li data={i}><a>{to_mime_type(r['head'])}</a></li>\n"
 
         C += '<div class="jwm-mail-body-multipart-alternative-body is-hidden">\n'
-        C += mime_render(
-            *to_mime_types(r["head"]), r["body"], path + (len(parts) - 1 - i,)
-        )
+        if r["head"].get("content_disposition", "").lower() == "attachment":
+            C += _as_download(r, path + (len(parts) - 1 - i,))
+        else:
+            C += mime_render(r["head"], r.get("body"), path + (len(parts) - 1 - i,))
         C += "</div>\n"
 
     C += "</div>"
@@ -70,24 +72,9 @@ def render_multipart(_subtype, content, path):
             or p["head"]["content_disposition"].lower() == "none"
             or p["head"]["content_disposition"].lower() == "inline"
         ):
-            R += mime_render(*to_mime_types(p["head"]), p["body"], path + (i,))
+            R += mime_render(p["head"], p["body"], path + (i,))
         elif p["head"]["content_disposition"].lower() == "attachment":
-            link_text = gettext("Attachment {filename} of type {filetype}").format(
-                filename=p["head"]["filename"], filetype=to_mime_type(p["head"])
-            )
-
-            ref_url = url_for(
-                "read",
-                folder=request.view_args.get("folder"),
-                msgid=request.view_args["msgid"],
-                path=".".join(map(str, [*path, i])),
-                format="raw",
-            )
-
-            R += "<p>"
-            R += f'<a href="{ref_url}" download="{escape(p["head"]["filename"])}">\n'
-            R += f"{escape(link_text)}</a>\n"
-            R += "</p>\n"
+            R += _as_download(p, [*path, i])
         else:
             current_app.log.warning(
                 "unknown Content-Disposition %s", p["head"]["content_disposition"]
@@ -99,6 +86,26 @@ def render_multipart(_subtype, content, path):
     return R + "</div>\n"
 
 
+def _as_download(part, path):
+    link_text = gettext("Attachment {filename} of type {filetype}").format(
+        filename=part["head"]["filename"], filetype=to_mime_type(part["head"])
+    )
+
+    ref_url = url_for(
+        "read",
+        folder=request.view_args.get("folder"),
+        msgid=request.view_args["msgid"],
+        path=".".join(map(str, path)),
+        format="raw",
+    )
+
+    R = "<p>"
+    R += f'<a href="{ref_url}" download="{escape(part["head"]["filename"])}">\n'
+    R += f"{escape(link_text)}</a>\n"
+    R += "</p>\n"
+    return R
+
+
 def _format_header(category, value):
     R = ""
 
@@ -135,7 +142,7 @@ def render_message(subtype, msg, path):
     R += f"<dd>{to_mime_type(msg['head']['mime'])}</dd>\n"
     R += "</dl>\n"
 
-    R += mime_render(*to_mime_types(msg["head"]["mime"]), msg["body"], path + (0,))
+    R += mime_render(msg["head"]["mime"], msg["body"], path + (0,))
 
     return R + "</div>\n"
 
@@ -150,7 +157,9 @@ MIMERenderSubs = {
 }
 
 
-def mime_render(maintype, subtype, content, path):
+def mime_render(mime_header, content, path):
+    maintype = escape(mime_header["content_maintype"])
+    subtype = escape(mime_header["content_subtype"])
     renderer = MIMERenderSubs.get((maintype, subtype)) or MIMERenderSubs.get(maintype)
 
     if not renderer:
@@ -165,9 +174,9 @@ def to_mime_type(mime):
     return escape(f"{mime['content_maintype']}/{mime['content_subtype']}".lower())
 
 
-def to_mime_types(mime):
-    return escape(mime["content_maintype"]), escape(mime["content_subtype"])
-
-
 def format_mail(mail):
-    return Markup(mime_render("message", "rfc822", mail, tuple()))
+    return Markup(
+        mime_render(
+            {"content_maintype": "message", "content_subtype": "rfc822"}, mail, tuple()
+        )
+    )