[PATCH v3 5/9] atom: add manage_payload()

[Muhammad Rizki <kiizuha@gnuweeb.org>]

Add manage_payload() to handle the email decoding to utf-8. This include
a non-UTF8 character and base64 decoding.

Signed-off-by: Muhammad Rizki <kiizuha@gnuweeb.org>
---
 daemon/atom/utils.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/daemon/atom/utils.py b/daemon/atom/utils.py
index f554f6f..ed5ca03 100644
--- a/daemon/atom/utils.py
+++ b/daemon/atom/utils.py
@@ -8,6 +8,7 @@ from pyrogram.types import Chat, InlineKeyboardMarkup, InlineKeyboardButton
 from email.message import Message
 from typing import Dict, Union
 from slugify import slugify
+from base64 import b64decode
 import hashlib
 import uuid
 import os
@@ -15,6 +16,7 @@ import re
 import shutil
 import httpx
 import html
+import quopri
 
 
 def get_email_msg_id(mail):
@@ -136,7 +138,7 @@ def gen_temp(name: str, platform: str):
 
 def extract_body(thread: Message, platform: str):
 	if not thread.is_multipart():
-		p = thread.get_payload(decode=True).decode(errors='replace')
+		p = manage_payload(thread)
 
 		if platform == "discord":
 			p = quote_reply(p)
@@ -253,6 +255,18 @@ def fix_utf8_char(text: str, html_escape: bool = True):
 	return t
 
 
+def manage_payload(payload: Message):
+	p = str(payload.get_payload())
+	tf_encode = payload.get("Content-Transfer-Encoding")
+
+	if tf_encode == "base64":
+		return b64decode(p).decode("utf-8")
+	if tf_encode == "quoted-printable":
+		return quopri.decodestring(p.encode()).decode()
+
+	return p.encode().decode("utf-8", errors="replace")
+
+
 EMAIL_MSG_ID_PATTERN = r"<([^\<\>]+)>"
 def extract_email_msg_id(msg_id):
 	ret = re.search(EMAIL_MSG_ID_PATTERN, msg_id)
-- 
Muhammad Rizki