From 79d0ef8906066df5d493eda1a53b58341af29848 Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Fri, 4 Jun 2021 17:15:59 +0200
Subject: [PATCH] Use Postfix queue-id as log message-id

---
 app/email_utils.py        | 19 ++++++++++++++++++-
 app/log.py                |  2 +-
 email_handler.py          |  4 ++++
 tests/test_email_utils.py |  9 +++++++++
 4 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/app/email_utils.py b/app/email_utils.py
index df91aad7..595c4d07 100644
--- a/app/email_utils.py
+++ b/app/email_utils.py
@@ -13,7 +13,7 @@ from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
 from email.utils import make_msgid, formatdate, parseaddr
 from smtplib import SMTP, SMTPServerDisconnected
-from typing import Tuple, List
+from typing import Tuple, List, Optional
 
 import arrow
 import dkim
@@ -1198,3 +1198,20 @@ def sl_sendmail(
             )
         else:
             raise
+
+
+def get_queue_id(msg: Message) -> Optional[str]:
+    """Get the Postfix queue-id from a message"""
+    received_header = msg["Received"]
+    if not received_header:
+        return
+
+    # received_header looks like 'from mail-wr1-x434.google.com (mail-wr1-x434.google.com [IPv6:2a00:1450:4864:20::434])\r\n\t(using TLSv1.3 with cipher TLS_AES_128_GCM_SHA256 (128/128 bits))\r\n\t(No client certificate requested)\r\n\tby mx1.simplelogin.co (Postfix) with ESMTPS id 4FxQmw1DXdz2vK2\r\n\tfor <jglfdjgld@alias.com>; Fri,  4 Jun 2021 14:55:43 +0000 (UTC)'
+    search_result = re.search("with ESMTPS id [0-9a-zA-Z]{1,}", received_header)
+    if not search_result:
+        return
+
+    # the "with ESMTPS id 4FxQmw1DXdz2vK2" part
+    with_esmtps = received_header[search_result.start() : search_result.end()]
+
+    return with_esmtps[len("with ESMTPS id ") :]
diff --git a/app/log.py b/app/log.py
index f3124a6d..09065457 100644
--- a/app/log.py
+++ b/app/log.py
@@ -18,7 +18,7 @@ _MESSAGE_ID = ""
 
 def set_message_id(message_id):
     global _MESSAGE_ID
-    print("set message_id", message_id)
+    LOG.d("set message_id %s", message_id)
     _MESSAGE_ID = message_id
 
 
diff --git a/email_handler.py b/email_handler.py
index c882b7d1..04c059f6 100644
--- a/email_handler.py
+++ b/email_handler.py
@@ -106,6 +106,7 @@ from app.email_utils import (
     spf_pass,
     sl_sendmail,
     sanitize_header,
+    get_queue_id,
 )
 from app.extensions import db
 from app.greylisting import greylisting_needed
@@ -1535,6 +1536,9 @@ def handle(envelope: Envelope) -> str:
     envelope.rcpt_tos = rcpt_tos
 
     msg = email.message_from_bytes(envelope.original_content)
+    postfix_queue_id = get_queue_id(msg)
+    if postfix_queue_id:
+        set_message_id(postfix_queue_id)
 
     # sanitize email headers
     sanitize_header(msg, "from")
diff --git a/tests/test_email_utils.py b/tests/test_email_utils.py
index fcc848c8..1cea140a 100644
--- a/tests/test_email_utils.py
+++ b/tests/test_email_utils.py
@@ -27,6 +27,7 @@ from app.email_utils import (
     should_disable,
     decode_text,
     parse_id_from_bounce,
+    get_queue_id,
 )
 from app.extensions import db
 from app.models import User, CustomDomain, Alias, Contact, EmailLog
@@ -720,3 +721,11 @@ def test_parse_id_from_bounce():
     assert parse_id_from_bounce("bounces+1234+@local") == 1234
     assert parse_id_from_bounce("anything+1234+@local") == 1234
     assert parse_id_from_bounce(BOUNCE_EMAIL.format(1234)) == 1234
+
+
+def test_get_queue_id():
+    msg = email.message_from_string(
+        "Received: from mail-wr1-x434.google.com (mail-wr1-x434.google.com [IPv6:2a00:1450:4864:20::434])\r\n\t(using TLSv1.3 with cipher TLS_AES_128_GCM_SHA256 (128/128 bits))\r\n\t(No client certificate requested)\r\n\tby mx1.simplelogin.co (Postfix) with ESMTPS id 4FxQmw1DXdz2vK2\r\n\tfor <jglfdjgld@alias.com>; Fri,  4 Jun 2021 14:55:43 +0000 (UTC)"
+    )
+
+    assert get_queue_id(msg) == "4FxQmw1DXdz2vK2"