From 1fcf166c00a236bb3aa767cdfbe764f81ef889fe Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Thu, 15 Oct 2020 16:24:04 +0200
Subject: [PATCH] small refactor: add should_add_dkim_signature

---
 app/email_utils.py | 13 ++++++++++++-
 email_handler.py   |  8 ++------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/app/email_utils.py b/app/email_utils.py
index 17b115c9..adbcf566 100644
--- a/app/email_utils.py
+++ b/app/email_utils.py
@@ -37,7 +37,7 @@ from app.config import (
 from app.dns_utils import get_mx_domains
 from app.extensions import db
 from app.log import LOG
-from app.models import Mailbox, User, SentAlert, CustomDomain
+from app.models import Mailbox, User, SentAlert, CustomDomain, PublicDomain
 
 
 def render(template_name, **kwargs) -> str:
@@ -627,3 +627,14 @@ def to_bytes(msg: Message):
         except UnicodeEncodeError:
             LOG.warning("as_bytes fails with SMTP policy, try SMTPUTF8 policy")
             return msg.as_bytes(policy=email.policy.SMTPUTF8)
+
+
+def should_add_dkim_signature(domain: str) -> bool:
+    if PublicDomain.get_by(domain=domain):
+        return True
+
+    custom_domain: CustomDomain = CustomDomain.get_by(domain=domain)
+    if custom_domain.dkim_verified:
+        return True
+
+    return False
diff --git a/email_handler.py b/email_handler.py
index ff41fcb3..1c07a9cb 100644
--- a/email_handler.py
+++ b/email_handler.py
@@ -98,6 +98,7 @@ from app.email_utils import (
     get_header_from_bounce,
     send_email_at_most_times,
     is_valid_alias_address_domain,
+    should_add_dkim_signature,
 )
 from app.extensions import db
 from app.greylisting import greylisting_needed
@@ -874,13 +875,8 @@ def handle_reply(envelope, msg: Message, rcpt_to: str) -> (bool, str):
         else:
             msg = replace_str_in_msg(msg, reply_email, contact.website_email)
 
-    if alias_domain in ALIAS_DOMAINS or alias_domain in PREMIUM_ALIAS_DOMAINS:
+    if should_add_dkim_signature(alias_domain):
         add_dkim_signature(msg, alias_domain)
-    # add DKIM-Signature for custom-domain alias
-    else:
-        custom_domain: CustomDomain = CustomDomain.get_by(domain=alias_domain)
-        if custom_domain.dkim_verified:
-            add_dkim_signature(msg, alias_domain)
 
     # create PGP email if needed
     if contact.pgp_finger_print and user.is_premium():