From 5c50628d36a9d6a37da96e052f05ffc5f381d088 Mon Sep 17 00:00:00 2001
From: Sibren Vasse <github@sibrenvasse.nl>
Date: Mon, 25 May 2020 16:08:10 +0200
Subject: [PATCH] Fix reverse alias replacement multipart message

---
 email_handler.py | 52 +++++++++++++++++++++++++++++++++++-------------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/email_handler.py b/email_handler.py
index 137a3942..1b7d75c7 100644
--- a/email_handler.py
+++ b/email_handler.py
@@ -30,9 +30,13 @@ It should contain the following info:
 
 
 """
+import arrow
 import email
+import spf
 import time
 import uuid
+from aiosmtpd.controller import Controller
+from aiosmtpd.smtp import Envelope
 from email import encoders
 from email.message import Message
 from email.mime.application import MIMEApplication
@@ -42,11 +46,6 @@ from io import BytesIO
 from smtplib import SMTP
 from typing import List, Tuple
 
-import arrow
-import spf
-from aiosmtpd.controller import Controller
-from aiosmtpd.smtp import Envelope
-
 from app import pgp_utils, s3
 from app.alias_utils import try_auto_create
 from app.config import (
@@ -95,7 +94,6 @@ from app.utils import random_string
 from init_app import load_pgp_public_keys
 from server import create_app
 
-
 _IP_HEADER = "X-SimpleLogin-Client-IP"
 _MAILBOX_ID_HEADER = "X-SimpleLogin-Mailbox-ID"
 
@@ -264,6 +262,30 @@ def replace_header_when_reply(msg: Message, alias: Alias, header: str):
     add_or_replace_header(msg, header, new_header)
 
 
+def replace_str_in_msg(msg: Message, fr: str, to: str):
+    if msg.get_content_maintype() != "text":
+        return msg
+    new_body = msg.get_payload(decode=True).replace(fr.encode(), to.encode())
+
+    # If utf-8 decoding fails, do not touch message part
+    try:
+        new_body = new_body.decode("utf-8")
+    except:
+        return msg
+
+    cte = (
+        msg["Content-Transfer-Encoding"].lower()
+        if msg["Content-Transfer-Encoding"]
+        else None
+    )
+    subtype = msg.get_content_subtype()
+    delete_header(msg, "Content-Transfer-Encoding")
+    delete_header(msg, "Content-Type")
+
+    email.contentmanager.set_text_content(msg, new_body, subtype=subtype, cte=cte)
+    return msg
+
+
 def generate_reply_email():
     # generate a reply_email, make sure it is unique
     # not use while loop to avoid infinite loop
@@ -556,15 +578,17 @@ def handle_reply(envelope, smtp: SMTP, msg: Message, rcpt_to: str) -> (bool, str
         envelope.rcpt_options,
     )
 
-    # replace the "ra+string@simplelogin.co" by the alias in the email body
-    # as this is usually included in when replying
+    # replace "ra+string@simplelogin.co" by the contact email in the email body
+    # as this is usually included when replying
     if user.replace_reverse_alias:
-        payload = (
-            msg.get_payload()
-            .encode()
-            .replace(reply_email.encode(), contact.website_email.encode())
-        )
-        msg.set_payload(payload)
+        if msg.is_multipart():
+            for part in msg.walk():
+                if part.get_content_maintype() != "text":
+                    continue
+                part = replace_str_in_msg(part, reply_email, contact.website_email)
+
+        else:
+            msg = replace_str_in_msg(msg, reply_email, contact.website_email)
 
     if alias_domain in ALIAS_DOMAINS:
         add_dkim_signature(msg, alias_domain)