diff --git a/app/email_utils.py b/app/email_utils.py index 0e519b38..feb940fb 100644 --- a/app/email_utils.py +++ b/app/email_utils.py @@ -1,5 +1,6 @@ import base64 import email +import enum import os import quopri import random @@ -673,7 +674,13 @@ def is_valid_email(email_address: str) -> bool: ) -def get_encoding(msg: Message) -> str: +class EmailEncoding(enum.Enum): + BASE64 = "base64" + QUOTED = "quoted-printable" + NO = "no-encoding" + + +def get_encoding(msg: Message) -> EmailEncoding: """ Return the message encoding, possible values: - quoted-printable @@ -682,21 +689,24 @@ def get_encoding(msg: Message) -> str: """ cte = str(msg.get("content-transfer-encoding", "")).lower() if cte in ("", "7bit"): - return "7bit" + return EmailEncoding.NO - if cte in ("quoted-printable", "base64"): - return cte + if cte == "base64": + return EmailEncoding.BASE64 + + if cte == "quoted-printable": + return EmailEncoding.QUOTED LOG.exception("Unknown encoding %s", cte) - return "7bit" + return EmailEncoding.NO -def encode_text(text: str, encoding: str = "7bit") -> str: - if encoding == "quoted-printable": +def encode_text(text: str, encoding: EmailEncoding = EmailEncoding.NO) -> str: + if encoding == EmailEncoding.QUOTED: encoded = quopri.encodestring(text.encode("utf-8")) return str(encoded, "utf-8") - elif encoding == "base64": + elif encoding == EmailEncoding.BASE64: encoded = base64.b64encode(text.encode("utf-8")) return str(encoded, "utf-8") else: # 7bit - no encoding diff --git a/tests/test_email_utils.py b/tests/test_email_utils.py index 6aeb571b..b786b538 100644 --- a/tests/test_email_utils.py +++ b/tests/test_email_utils.py @@ -20,6 +20,7 @@ from app.email_utils import ( normalize_reply_email, get_encoding, encode_text, + EmailEncoding, ) from app.extensions import db from app.models import User, CustomDomain @@ -422,24 +423,24 @@ def test_normalize_reply_email(flask_client): def test_get_encoding(): msg = email.message_from_string("") - assert get_encoding(msg) == "7bit" + assert get_encoding(msg) == EmailEncoding.NO msg = email.message_from_string("Content-TRANSFER-encoding: Invalid") - assert get_encoding(msg) == "7bit" + assert get_encoding(msg) == EmailEncoding.NO msg = email.message_from_string("Content-TRANSFER-encoding: quoted-printable") - assert get_encoding(msg) == "quoted-printable" + assert get_encoding(msg) == EmailEncoding.QUOTED msg = email.message_from_string("Content-TRANSFER-encoding: base64") - assert get_encoding(msg) == "base64" + assert get_encoding(msg) == EmailEncoding.BASE64 def test_encode_text(): assert encode_text("") == "" assert encode_text("ascii") == "ascii" - assert encode_text("ascii", "base64") == "YXNjaWk=" - assert encode_text("ascii", "quoted-printable") == "ascii" + assert encode_text("ascii", EmailEncoding.BASE64) == "YXNjaWk=" + assert encode_text("ascii", EmailEncoding.QUOTED) == "ascii" assert encode_text("mèo méo") == "mèo méo" - assert encode_text("mèo méo", "base64") == "bcOobyBtw6lv" - assert encode_text("mèo méo", "quoted-printable") == "m=C3=A8o m=C3=A9o" + assert encode_text("mèo méo", EmailEncoding.BASE64) == "bcOobyBtw6lv" + assert encode_text("mèo méo", EmailEncoding.QUOTED) == "m=C3=A8o m=C3=A9o"