move get_spam_score_async(), get_spam_score() to email/spam.py

app/email/spam.py

@@ -0,0 +1,63 @@
+import asyncio
+import time
+from email.message import Message
+
+import aiospamc
+
+from app.config import SPAMASSASSIN_HOST
+from app.email_utils import to_bytes
+from app.log import LOG
+from app.models import EmailLog
+from app.spamassassin_utils import SpamAssassin
+
+
+async def get_spam_score_async(message: Message) -> float:
+    sa_input = to_bytes(message)
+
+    # Spamassassin requires to have an ending linebreak
+    if not sa_input.endswith(b"\n"):
+        LOG.d("add linebreak to spamassassin input")
+        sa_input += b"\n"
+
+    try:
+        # wait for at max 300s which is the default spamd timeout-child
+        response = await asyncio.wait_for(
+            aiospamc.check(sa_input, host=SPAMASSASSIN_HOST), timeout=300
+        )
+        return response.headers["Spam"].score
+    except asyncio.TimeoutError:
+        LOG.e("SpamAssassin timeout")
+        # return a negative score so the message is always considered as ham
+        return -999
+    except Exception:
+        LOG.e("SpamAssassin exception")
+        return -999
+
+
+def get_spam_score(
+    message: Message, email_log: EmailLog, can_retry=True
+) -> (float, dict):
+    """
+    Return the spam score and spam report
+    """
+    LOG.d("get spam score for %s", email_log)
+    sa_input = to_bytes(message)
+
+    # Spamassassin requires to have an ending linebreak
+    if not sa_input.endswith(b"\n"):
+        LOG.d("add linebreak to spamassassin input")
+        sa_input += b"\n"
+
+    try:
+        # wait for at max 300s which is the default spamd timeout-child
+        sa = SpamAssassin(sa_input, host=SPAMASSASSIN_HOST, timeout=300)
+        return sa.get_score(), sa.get_report_json()
+    except Exception:
+        if can_retry:
+            LOG.w("SpamAssassin exception, retry")
+            time.sleep(3)
+            return get_spam_score(message, email_log, can_retry=False)
+        else:
+            # return a negative score so the message is always considered as ham
+            LOG.exception("SpamAssassin exception, ignore spam check")
+            return -999, None

email_handler.py

@@ -31,7 +31,6 @@ It should contain the following info:
 
 """
 import argparse
-import asyncio
 import email
 import time
 import uuid
@@ -45,7 +44,6 @@ from io import BytesIO
 from smtplib import SMTP, SMTPRecipientsRefused, SMTPServerDisconnected
 from typing import List, Tuple, Optional
 
-import aiospamc
 import arrow
 import spf
 from aiosmtpd.controller import Controller
@@ -83,6 +81,7 @@ from app.config import (
     POSTFIX_PORT_FORWARD,
     NOT_SEND_EMAIL,
 )
+from app.email.spam import get_spam_score
 from app.email_utils import (
     send_email,
     add_dkim_signature,
@@ -125,7 +124,6 @@ from app.models import (
     TransactionalEmail,
 )
 from app.pgp_utils import PGPException, sign_data_with_pgpy, sign_data
-from app.spamassassin_utils import SpamAssassin
 from app.utils import sanitize_email
 from init_app import load_pgp_public_keys
 from server import create_app, create_light_app
@@ -1672,58 +1670,6 @@ def handle_bounce(envelope, rcpt_to) -> str:
         return "550 SL E26 Email cannot be forwarded to mailbox"
 
 
-async def get_spam_score_async(message: Message) -> float:
-    sa_input = to_bytes(message)
-
-    # Spamassassin requires to have an ending linebreak
-    if not sa_input.endswith(b"\n"):
-        LOG.d("add linebreak to spamassassin input")
-        sa_input += b"\n"
-
-    try:
-        # wait for at max 300s which is the default spamd timeout-child
-        response = await asyncio.wait_for(
-            aiospamc.check(sa_input, host=SPAMASSASSIN_HOST), timeout=300
-        )
-        return response.headers["Spam"].score
-    except asyncio.TimeoutError:
-        LOG.e("SpamAssassin timeout")
-        # return a negative score so the message is always considered as ham
-        return -999
-    except Exception:
-        LOG.e("SpamAssassin exception")
-        return -999
-
-
-def get_spam_score(
-    message: Message, email_log: EmailLog, can_retry=True
-) -> (float, dict):
-    """
-    Return the spam score and spam report
-    """
-    LOG.d("get spam score for %s", email_log)
-    sa_input = to_bytes(message)
-
-    # Spamassassin requires to have an ending linebreak
-    if not sa_input.endswith(b"\n"):
-        LOG.d("add linebreak to spamassassin input")
-        sa_input += b"\n"
-
-    try:
-        # wait for at max 300s which is the default spamd timeout-child
-        sa = SpamAssassin(sa_input, host=SPAMASSASSIN_HOST, timeout=300)
-        return sa.get_score(), sa.get_report_json()
-    except Exception:
-        if can_retry:
-            LOG.w("SpamAssassin exception, retry")
-            time.sleep(3)
-            return get_spam_score(message, email_log, can_retry=False)
-        else:
-            # return a negative score so the message is always considered as ham
-            LOG.exception("SpamAssassin exception, ignore spam check")
-            return -999, None
-
-
 def sl_sendmail(
     from_addr,
     to_addr,