change content transfer encoding before setting payload

gpgmymail

@@ -58,124 +58,6 @@ def is_message_encrypted(message: email.message.Message) -> bool:
 
     return message.get_content_subtype() == "encrypted"
 
-def decode_email(message: email.message.Message) -> email.message.Message:
-    """Turn a quoted-printable or base64 encoded email into a 7or8bit encoded
-    email
-
-    :param message: email.message.Message to be decoded
-    :return: decoded email.message.Message"""
-    def decoded_bytes_to_return_value(decoded_bytes: bytes) -> email.message.Message:
-        """
-        if at any point you want to return, return a call of this function and
-        pass decoded_bytes
-
-        :param decoded_bytes: an email as an ASCII byte array; this should be
-                              stored in decoded_bytes
-        :return: the expected return value of decode_email
-        """
-        # if i do message_from_bytes it bizarrely changes it back to base64?
-        # utf-8 has encoding issues so do latin1
-        return email.message_from_string(decoded_bytes.decode("latin1"))
-
-    # this is a kinda hacky way to do this by manipulating the message as a
-    # string but i couldn't get it to work any other way
-
-    # sometimes this raises an exception and idk why
-    try:
-        decoded_bytes = message.as_bytes()
-    except UnicodeEncodeError:
-        decoded_bytes = message.as_string().encode()
-    # if email doesn't need decoding
-    has_quopri = b'Content-Transfer-Encoding: quoted-printable' in decoded_bytes
-    has_base64 = b'Content-Transfer-Encoding: base64' in decoded_bytes
-    if not (has_quopri or has_base64):
-        return message
-    decoded_bytes = quopri.decodestring(decoded_bytes)
-
-    # replace any instances of the Content-Transfer-Encoding header
-    # quopri version, we do base64 version down there
-    decoded_bytes = decoded_bytes.replace(
-        b'Content-Transfer-Encoding: quoted-printable',
-        b'Content-Transfer-Encoding: 7bit'
-    )
-
-    # now exit if there's no base64 as i think that's the most fucky
-    if not has_base64:
-        return decoded_bytes_to_return_value(decoded_bytes)
-
-    # REALLY hacky but i had issues with the more sensible ways to do this.
-    # iterates through a Message object to find CTEs of base64
-    # gets the b64 payload and the decoded payload
-    # then find and replaces in decoded_bytes the b64 payload
-    # with the decoded payload
-    # lol
-
-    def decode_b64_part(
-        part: email.message.Message,
-        decoded_bytes: bytes,
-        most_recent_boundary: str = None
-    ) -> bytes:
-        """
-        change decoded_bytes such that part is decoded if base64 (unchanged if
-        not)
-
-        see usage below for examples
-
-        :param part: email.message.Message to be decoded
-        :param decoded_bytes: the email as a bytes object (ie not a string with
-                              encoding), will have modified version returned
-        :param most_recent_boundary: str of the most recent boundary so we
-                                     don't overwrite this
-        :return: bytes of decoded_bytes with part decoded if base64
-        """
-        if part.get("Content-Transfer-Encoding") == "base64":
-            b64_str = part.get_payload()
-            # remove the boundary as we don't want to change this
-            if most_recent_boundary:
-                b64_str = b64_str.replace(most_recent_boundary, "")
-            # sometimes we have leftover hyphens from a boundary, so strip:
-            # hyphens not in base64 so we know not to use them
-            # strip whitespace first
-            b64_str = b64_str.strip()
-            b64_str = b64_str.strip('-')
-            b64_str = b64_str.encode()  # turn into bytes-like object
-            # this will also decode the boundary so there'll be some nonsese
-            # chars at end of email but it's nbd
-            decoded_b64_str = part.get_payload(decode=True)
-            return decoded_bytes.replace(
-                b64_str,
-                decoded_b64_str
-            )
-        
-        return decoded_bytes
-
-    quopri_decoded_message = email.message_from_bytes(decoded_bytes)
-    if quopri_decoded_message.is_multipart():
-        most_recent_boundary = None
-        for part in quopri_decoded_message.walk():
-            # multipart and has boundary (not None)
-            if part.is_multipart() and part.get_boundary():
-                most_recent_boundary = part.get_boundary()
-            else:
-                decoded_bytes = decode_b64_part(
-                    part,
-                    decoded_bytes,
-                    most_recent_boundary
-                )
-    else:
-        decoded_bytes = decode_b64_part(
-            quopri_decoded_message,
-            decoded_bytes,
-            None
-        )
-
-    decoded_bytes = decoded_bytes.replace(
-        b'Content-Transfer-Encoding: base64',
-        b'Content-Transfer-Encoding: 7bit'
-    )
-    
-    return decoded_bytes_to_return_value(decoded_bytes)
-
 def set_email_header(
     message: email.message.Message,
     name: str,
@@ -197,6 +79,24 @@ def set_email_header(
     else:
         message.add_header(name, value)
 
+def decode_email(message: email.message.Message) -> email.message.Message:
+    """Turn a quoted-printable or base64 encoded email into a 7or8bit encoded
+    email
+
+    :param message: email.message.Message to be decoded
+    :return: decoded email.message.Message"""
+    if message.is_multipart():
+        for part in message.walk():
+            if part.get('Content-Transfer-Encoding') in ('quoted-printable', 'base64'):
+                set_email_header(part, 'Content-Transfer-Encoding', '7bit')
+                part.set_payload(part.get_payload(decode=True))
+    else:
+        if message.get('Content-Transfer-Encoding') in ('quoted-printable', 'base64'):
+            set_email_header(message, 'Content-Transfer-Encoding', '7bit')
+            message.set_payload(message.get_payload(decode=True))
+
+    return message
+
 def encrypt(
         message: email.message.Message,
         recipients: typing.List[str],