refactor email-handler: move handle_forward, handle_reply to outside of MailHandler class

email_handler.py

@@ -176,286 +176,284 @@ def try_auto_create(alias: str) -> Optional[GenEmail]:
     return gen_email
 
 
-class MailHandler:
-    async def handle_DATA(self, server, session, envelope):
-        LOG.debug(">>> New message <<<")
+def handle_forward(envelope, smtp: SMTP, msg: Message, rcpt_to: str) -> str:
+    """return *status_code message*"""
+    alias = rcpt_to.lower()  # alias@SL
+
+    gen_email = GenEmail.get_by(email=alias)
+    if not gen_email:
+        LOG.d("alias %s not exist. Try to see if it can be created on the fly", alias)
+        gen_email = try_auto_create(alias)
+        if not gen_email:
+            LOG.d("alias %s cannot be created on-the-fly, return 510", alias)
+            return "510 Email not exist"
 
-        LOG.debug("Mail from %s", envelope.mail_from)
-        LOG.debug("Rcpt to %s", envelope.rcpt_tos)
-        message_data = envelope.content.decode("utf8", errors="replace")
+    if gen_email.mailbox_id:
+        mailbox_email = gen_email.mailbox.email
+    else:
+        mailbox_email = gen_email.user.email
 
-        smtp = SMTP(POSTFIX_SERVER, 25)
-        msg = Parser(policy=SMTPUTF8).parsestr(message_data)
+    website_email = get_email_part(msg["From"])
 
-        for rcpt_to in envelope.rcpt_tos:
-            # Reply case
-            # recipient starts with "reply+" or "ra+" (ra=reverse-alias) prefix
-            if rcpt_to.startswith("reply+") or rcpt_to.startswith("ra+"):
-                LOG.debug("Reply phase")
-                app = new_app()
+    forward_email = ForwardEmail.get_by(
+        gen_email_id=gen_email.id, website_email=website_email
+    )
+    if forward_email:
+        # update the From header if needed
+        if forward_email.website_from != msg["From"]:
+            LOG.d("Update From header for %s", forward_email)
+            forward_email.website_from = msg["From"]
+            db.session.commit()
+    else:
+        LOG.debug(
+            "create forward email for alias %s and website email %s",
+            alias,
+            website_email,
+        )
 
-                with app.app_context():
-                    return self.handle_reply(envelope, smtp, msg, rcpt_to)
-            else:  # Forward case
-                LOG.debug("Forward phase")
-                app = new_app()
+        # generate a reply_email, make sure it is unique
+        # not use while to avoid infinite loop
+        for _ in range(1000):
+            reply_email = f"reply+{random_string(30)}@{EMAIL_DOMAIN}"
+            if not ForwardEmail.get_by(reply_email=reply_email):
+                break
+
+        forward_email = ForwardEmail.create(
+            gen_email_id=gen_email.id,
+            website_email=website_email,
+            website_from=msg["From"],
+            reply_email=reply_email,
+        )
+        db.session.commit()
 
-                with app.app_context():
-                    return self.handle_forward(envelope, smtp, msg, rcpt_to)
+    forward_log = ForwardEmailLog.create(forward_id=forward_email.id)
 
-    def handle_forward(self, envelope, smtp: SMTP, msg: Message, rcpt_to: str) -> str:
-        """return *status_code message*"""
-        alias = rcpt_to.lower()  # alias@SL
+    if gen_email.enabled:
+        # add custom header
+        add_or_replace_header(msg, "X-SimpleLogin-Type", "Forward")
 
-        gen_email = GenEmail.get_by(email=alias)
-        if not gen_email:
-            LOG.d(
-                "alias %s not exist. Try to see if it can be created on the fly", alias
-            )
-            gen_email = try_auto_create(alias)
-            if not gen_email:
-                LOG.d("alias %s cannot be created on-the-fly, return 510", alias)
-                return "510 Email not exist"
+        # remove reply-to header if present
+        delete_header(msg, "Reply-To")
 
-        if gen_email.mailbox_id:
-            mailbox_email = gen_email.mailbox.email
-        else:
-            mailbox_email = gen_email.user.email
+        # change the from header so the sender comes from @SL
+        # so it can pass DMARC check
+        # replace the email part in from: header
+        from_header = (
+            get_email_name(msg["From"])
+            + ("" if get_email_name(msg["From"]) == "" else " - ")
+            + website_email.replace("@", " at ")
+            + f" <{forward_email.reply_email}>"
+        )
+        msg.replace_header("From", from_header)
+        LOG.d("new from header:%s", from_header)
 
-        website_email = get_email_part(msg["From"])
+        # add List-Unsubscribe header
+        unsubscribe_link = f"{URL}/dashboard/unsubscribe/{gen_email.id}"
+        add_or_replace_header(msg, "List-Unsubscribe", f"<{unsubscribe_link}>")
+        add_or_replace_header(
+            msg, "List-Unsubscribe-Post", "List-Unsubscribe=One-Click"
+        )
+
+        add_dkim_signature(msg, EMAIL_DOMAIN)
 
-        forward_email = ForwardEmail.get_by(
-            gen_email_id=gen_email.id, website_email=website_email
+        LOG.d(
+            "Forward mail from %s to %s, mail_options %s, rcpt_options %s ",
+            website_email,
+            mailbox_email,
+            envelope.mail_options,
+            envelope.rcpt_options,
         )
-        if forward_email:
-            # update the From header if needed
-            if forward_email.website_from != msg["From"]:
-                LOG.d("Update From header for %s", forward_email)
-                forward_email.website_from = msg["From"]
-                db.session.commit()
-        else:
-            LOG.debug(
-                "create forward email for alias %s and website email %s",
-                alias,
-                website_email,
-            )
 
-            # generate a reply_email, make sure it is unique
-            # not use while to avoid infinite loop
-            for _ in range(1000):
-                reply_email = f"reply+{random_string(30)}@{EMAIL_DOMAIN}"
-                if not ForwardEmail.get_by(reply_email=reply_email):
-                    break
-
-            forward_email = ForwardEmail.create(
-                gen_email_id=gen_email.id,
-                website_email=website_email,
-                website_from=msg["From"],
-                reply_email=reply_email,
-            )
-            db.session.commit()
+        # smtp.send_message has UnicodeEncodeErroremail issue
+        # encode message raw directly instead
+        msg_raw = msg.as_string().encode()
+        smtp.sendmail(
+            forward_email.reply_email,
+            mailbox_email,
+            msg_raw,
+            envelope.mail_options,
+            envelope.rcpt_options,
+        )
+    else:
+        LOG.d("%s is disabled, do not forward", gen_email)
+        forward_log.blocked = True
 
-        forward_log = ForwardEmailLog.create(forward_id=forward_email.id)
+    db.session.commit()
+    return "250 Message accepted for delivery"
 
-        if gen_email.enabled:
-            # add custom header
-            add_or_replace_header(msg, "X-SimpleLogin-Type", "Forward")
 
-            # remove reply-to header if present
-            delete_header(msg, "Reply-To")
+def handle_reply(envelope, smtp: SMTP, msg: Message, rcpt_to: str) -> str:
+    reply_email = rcpt_to.lower()
 
-            # change the from header so the sender comes from @SL
-            # so it can pass DMARC check
-            # replace the email part in from: header
-            from_header = (
-                get_email_name(msg["From"])
-                + ("" if get_email_name(msg["From"]) == "" else " - ")
-                + website_email.replace("@", " at ")
-                + f" <{forward_email.reply_email}>"
-            )
-            msg.replace_header("From", from_header)
-            LOG.d("new from header:%s", from_header)
-
-            # add List-Unsubscribe header
-            unsubscribe_link = f"{URL}/dashboard/unsubscribe/{gen_email.id}"
-            add_or_replace_header(msg, "List-Unsubscribe", f"<{unsubscribe_link}>")
-            add_or_replace_header(
-                msg, "List-Unsubscribe-Post", "List-Unsubscribe=One-Click"
-            )
+    # reply_email must end with EMAIL_DOMAIN
+    if not reply_email.endswith(EMAIL_DOMAIN):
+        LOG.warning(f"Reply email {reply_email} has wrong domain")
+        return "550 wrong reply email"
 
-            add_dkim_signature(msg, EMAIL_DOMAIN)
+    forward_email = ForwardEmail.get_by(reply_email=reply_email)
+    if not forward_email:
+        LOG.warning(f"No such forward-email with {reply_email} as reply-email")
+        return "550 wrong reply email"
 
-            LOG.d(
-                "Forward mail from %s to %s, mail_options %s, rcpt_options %s ",
-                website_email,
-                mailbox_email,
-                envelope.mail_options,
-                envelope.rcpt_options,
-            )
+    alias: str = forward_email.gen_email.email
+    alias_domain = alias[alias.find("@") + 1 :]
 
-            # smtp.send_message has UnicodeEncodeErroremail issue
-            # encode message raw directly instead
-            msg_raw = msg.as_string().encode()
-            smtp.sendmail(
-                forward_email.reply_email,
-                mailbox_email,
-                msg_raw,
-                envelope.mail_options,
-                envelope.rcpt_options,
-            )
-        else:
-            LOG.d("%s is disabled, do not forward", gen_email)
-            forward_log.blocked = True
+    # alias must end with one of the ALIAS_DOMAINS or custom-domain
+    if not email_belongs_to_alias_domains(alias):
+        if not CustomDomain.get_by(domain=alias_domain):
+            return "550 alias unknown by SimpleLogin"
 
-        db.session.commit()
-        return "250 Message accepted for delivery"
+    gen_email = forward_email.gen_email
+    if gen_email.mailbox_id:
+        mailbox_email = gen_email.mailbox.email
+    else:
+        mailbox_email = gen_email.user.email
+
+    # bounce email initiated by Postfix
+    # can happen in case emails cannot be delivered to user-email
+    # in this case Postfix will try to send a bounce report to original sender, which is
+    # the "reply email"
+    if envelope.mail_from == "<>":
+        LOG.error(
+            "Bounce when sending to alias %s, user %s, from header: %s",
+            alias,
+            gen_email.user,
+            msg["From"],
+        )
+        # send the bounce email payload to admin
+        msg.replace_header("From", SUPPORT_EMAIL)
+        msg.replace_header("To", ADMIN_EMAIL)
+        add_dkim_signature(msg, get_email_domain_part(SUPPORT_EMAIL))
 
-    def handle_reply(self, envelope, smtp: SMTP, msg: Message, rcpt_to: str) -> str:
-        reply_email = rcpt_to.lower()
+        smtp.sendmail(
+            SUPPORT_EMAIL,
+            ADMIN_EMAIL,
+            msg.as_string().encode(),
+            envelope.mail_options,
+            envelope.rcpt_options,
+        )
+        return "550 ignored"
+
+    # only mailbox can send email to the reply-email
+    if envelope.mail_from.lower() != mailbox_email.lower():
+        LOG.warning(
+            f"Reply email can only be used by user email. Actual mail_from: %s. msg from header: %s, User email %s. reply_email %s",
+            envelope.mail_from,
+            msg["From"],
+            mailbox_email,
+            reply_email,
+        )
 
-        # reply_email must end with EMAIL_DOMAIN
-        if not reply_email.endswith(EMAIL_DOMAIN):
-            LOG.warning(f"Reply email {reply_email} has wrong domain")
-            return "550 wrong reply email"
+        user = gen_email.user
+        send_email(
+            mailbox_email,
+            f"Reply from your alias {alias} only works from your mailbox",
+            render(
+                "transactional/reply-must-use-personal-email.txt",
+                name=user.name,
+                alias=alias,
+                sender=envelope.mail_from,
+                mailbox_email=mailbox_email,
+            ),
+            render(
+                "transactional/reply-must-use-personal-email.html",
+                name=user.name,
+                alias=alias,
+                sender=envelope.mail_from,
+                mailbox_email=mailbox_email,
+            ),
+        )
 
-        forward_email = ForwardEmail.get_by(reply_email=reply_email)
-        if not forward_email:
-            LOG.warning(f"No such forward-email with {reply_email} as reply-email")
-            return "550 wrong reply email"
+        # Notify sender that they cannot send emails to this address
+        send_email(
+            envelope.mail_from,
+            f"Your email ({envelope.mail_from}) is not allowed to send emails to {reply_email}",
+            render(
+                "transactional/send-from-alias-from-unknown-sender.txt",
+                sender=envelope.mail_from,
+                reply_email=reply_email,
+            ),
+            "",
+        )
 
-        alias: str = forward_email.gen_email.email
-        alias_domain = alias[alias.find("@") + 1 :]
+        return "550 ignored"
 
-        # alias must end with one of the ALIAS_DOMAINS or custom-domain
-        if not email_belongs_to_alias_domains(alias):
-            if not CustomDomain.get_by(domain=alias_domain):
-                return "550 alias unknown by SimpleLogin"
+    delete_header(msg, "DKIM-Signature")
 
-        gen_email = forward_email.gen_email
-        if gen_email.mailbox_id:
-            mailbox_email = gen_email.mailbox.email
-        else:
-            mailbox_email = gen_email.user.email
+    # the email comes from alias
+    msg.replace_header("From", alias)
 
-        # bounce email initiated by Postfix
-        # can happen in case emails cannot be delivered to user-email
-        # in this case Postfix will try to send a bounce report to original sender, which is
-        # the "reply email"
-        if envelope.mail_from == "<>":
-            LOG.error(
-                "Bounce when sending to alias %s, user %s, from header: %s",
-                alias,
-                gen_email.user,
-                msg["From"],
-            )
-            # send the bounce email payload to admin
-            msg.replace_header("From", SUPPORT_EMAIL)
-            msg.replace_header("To", ADMIN_EMAIL)
-            add_dkim_signature(msg, get_email_domain_part(SUPPORT_EMAIL))
-
-            smtp.sendmail(
-                SUPPORT_EMAIL,
-                ADMIN_EMAIL,
-                msg.as_string().encode(),
-                envelope.mail_options,
-                envelope.rcpt_options,
-            )
-            return "550 ignored"
-
-        # only mailbox can send email to the reply-email
-        if envelope.mail_from.lower() != mailbox_email.lower():
-            LOG.warning(
-                f"Reply email can only be used by user email. Actual mail_from: %s. msg from header: %s, User email %s. reply_email %s",
-                envelope.mail_from,
-                msg["From"],
-                mailbox_email,
-                reply_email,
-            )
+    # some email providers like ProtonMail adds automatically the Reply-To field
+    # make sure to delete it
+    delete_header(msg, "Reply-To")
 
-            user = gen_email.user
-            send_email(
-                mailbox_email,
-                f"Reply from your alias {alias} only works from your mailbox",
-                render(
-                    "transactional/reply-must-use-personal-email.txt",
-                    name=user.name,
-                    alias=alias,
-                    sender=envelope.mail_from,
-                    mailbox_email=mailbox_email,
-                ),
-                render(
-                    "transactional/reply-must-use-personal-email.html",
-                    name=user.name,
-                    alias=alias,
-                    sender=envelope.mail_from,
-                    mailbox_email=mailbox_email,
-                ),
-            )
+    msg.replace_header("To", forward_email.website_email)
 
-            # Notify sender that they cannot send emails to this address
-            send_email(
-                envelope.mail_from,
-                f"Your email ({envelope.mail_from}) is not allowed to send emails to {reply_email}",
-                render(
-                    "transactional/send-from-alias-from-unknown-sender.txt",
-                    sender=envelope.mail_from,
-                    reply_email=reply_email,
-                ),
-                "",
-            )
+    # add List-Unsubscribe header
+    unsubscribe_link = f"{URL}/dashboard/unsubscribe/{forward_email.gen_email_id}"
+    add_or_replace_header(msg, "List-Unsubscribe", f"<{unsubscribe_link}>")
+    add_or_replace_header(msg, "List-Unsubscribe-Post", "List-Unsubscribe=One-Click")
 
-            return "550 ignored"
+    # Received-SPF is injected by postfix-policyd-spf-python can reveal user original email
+    delete_header(msg, "Received-SPF")
 
-        delete_header(msg, "DKIM-Signature")
+    LOG.d(
+        "send email from %s to %s, mail_options:%s,rcpt_options:%s",
+        alias,
+        forward_email.website_email,
+        envelope.mail_options,
+        envelope.rcpt_options,
+    )
 
-        # the email comes from alias
-        msg.replace_header("From", alias)
+    if alias_domain in ALIAS_DOMAINS:
+        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)
 
-        # some email providers like ProtonMail adds automatically the Reply-To field
-        # make sure to delete it
-        delete_header(msg, "Reply-To")
+    msg_raw = msg.as_string().encode()
+    smtp.sendmail(
+        alias,
+        forward_email.website_email,
+        msg_raw,
+        envelope.mail_options,
+        envelope.rcpt_options,
+    )
 
-        msg.replace_header("To", forward_email.website_email)
+    ForwardEmailLog.create(forward_id=forward_email.id, is_reply=True)
+    db.session.commit()
 
-        # add List-Unsubscribe header
-        unsubscribe_link = f"{URL}/dashboard/unsubscribe/{forward_email.gen_email_id}"
-        add_or_replace_header(msg, "List-Unsubscribe", f"<{unsubscribe_link}>")
-        add_or_replace_header(
-            msg, "List-Unsubscribe-Post", "List-Unsubscribe=One-Click"
-        )
+    return "250 Message accepted for delivery"
 
-        # Received-SPF is injected by postfix-policyd-spf-python can reveal user original email
-        delete_header(msg, "Received-SPF")
 
-        LOG.d(
-            "send email from %s to %s, mail_options:%s,rcpt_options:%s",
-            alias,
-            forward_email.website_email,
-            envelope.mail_options,
-            envelope.rcpt_options,
-        )
+class MailHandler:
+    async def handle_DATA(self, server, session, envelope):
+        LOG.debug(">>> New message <<<")
 
-        if alias_domain in ALIAS_DOMAINS:
-            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)
+        LOG.debug("Mail from %s", envelope.mail_from)
+        LOG.debug("Rcpt to %s", envelope.rcpt_tos)
+        message_data = envelope.content.decode("utf8", errors="replace")
 
-        msg_raw = msg.as_string().encode()
-        smtp.sendmail(
-            alias,
-            forward_email.website_email,
-            msg_raw,
-            envelope.mail_options,
-            envelope.rcpt_options,
-        )
+        smtp = SMTP(POSTFIX_SERVER, 25)
+        msg = Parser(policy=SMTPUTF8).parsestr(message_data)
 
-        ForwardEmailLog.create(forward_id=forward_email.id, is_reply=True)
-        db.session.commit()
+        for rcpt_to in envelope.rcpt_tos:
+            # Reply case
+            # recipient starts with "reply+" or "ra+" (ra=reverse-alias) prefix
+            if rcpt_to.startswith("reply+") or rcpt_to.startswith("ra+"):
+                LOG.debug("Reply phase")
+                app = new_app()
 
-        return "250 Message accepted for delivery"
+                with app.app_context():
+                    return handle_reply(envelope, smtp, msg, rcpt_to)
+            else:  # Forward case
+                LOG.debug("Forward phase")
+                app = new_app()
+
+                with app.app_context():
+                    return handle_forward(envelope, smtp, msg, rcpt_to)
 
 
 if __name__ == "__main__":