upstream-zulip/zproject/email_backends.py

# https://zulip.readthedocs.io/en/latest/subsystems/email.html#testing-in-a-real-email-client
import configparser
import logging
import smtplib
from collections.abc import Sequence
from email.message import Message
from typing import Any

import tenacity
from django.conf import settings
from django.core.mail import EmailMultiAlternatives
from django.core.mail.backends.smtp import EmailBackend
from django.core.mail.message import EmailAlternative, EmailMessage
from django.template import loader
from django.utils.timezone import now as timezone_now
from typing_extensions import override

MAX_CONNECTION_TRIES = 3

# SMTPException is a subclass of OSError, so retry on connection-level
# errors (ConnectionError, TimeoutError) and SMTPServerDisconnected
# (dropped connection), but not on SMTP protocol errors like
# SMTPAuthenticationError, SMTPRecipientsRefused, or SMTPDataError.
smtp_connection_backoff = tenacity.retry(
    wait=tenacity.wait_exponential_jitter(),
    retry=tenacity.retry_if_exception(
        lambda exc: (
            isinstance(exc, OSError)
            and (
                not isinstance(exc, smtplib.SMTPException)
                or isinstance(exc, smtplib.SMTPServerDisconnected)
            )
        )
    ),
    stop=tenacity.stop_after_attempt(MAX_CONNECTION_TRIES),
    reraise=True,
)


def get_forward_address() -> str:
    config = configparser.ConfigParser()
    config.read(settings.FORWARD_ADDRESS_CONFIG_FILE)
    try:
        return config.get("DEV_EMAIL", "forward_address")
    except (configparser.NoSectionError, configparser.NoOptionError):
        return ""


def set_forward_address(forward_address: str) -> None:
    config = configparser.ConfigParser()
    config.read(settings.FORWARD_ADDRESS_CONFIG_FILE)

    if not config.has_section("DEV_EMAIL"):
        config.add_section("DEV_EMAIL")
    config.set("DEV_EMAIL", "forward_address", forward_address)

    with open(settings.FORWARD_ADDRESS_CONFIG_FILE, "w") as cfgfile:
        config.write(cfgfile)


class EmailLogBackEnd(EmailBackend):
    @staticmethod
    def log_email(email: EmailMessage) -> None:
        """Used in development to record sent emails in a nice HTML log"""
        html_message: bytes | EmailMessage | Message | str = "Missing HTML message"
        assert isinstance(email, EmailMultiAlternatives)
        if len(email.alternatives) > 0:
            html_message = email.alternatives[0][0]

        context = {
            "subject": email.subject,
            "envelope_from": email.from_email,
            "from_email": email.extra_headers.get("From", email.from_email),
            "reply_to": email.reply_to,
            "recipients": email.to,
            "body": email.body,
            "date": email.extra_headers.get("Date", "?"),
            "html_message": html_message,
        }

        new_email = loader.render_to_string("zerver/email.html", context)

        # Read in the pre-existing log, so that we can add the new entry
        # at the top.
        try:
            with open(settings.EMAIL_CONTENT_LOG_PATH) as f:
                previous_emails = f.read()
        except FileNotFoundError:
            previous_emails = ""

        with open(settings.EMAIL_CONTENT_LOG_PATH, "w+") as f:
            f.write(new_email + previous_emails)

    @staticmethod
    def prepare_email_messages_for_forwarding(email_messages: Sequence[EmailMessage]) -> None:
        localhost_email_images_base_url = settings.ROOT_DOMAIN_URI + "/static/images/emails"
        czo_email_images_base_url = "https://chat.zulip.org/static/images/emails"

        for email_message in email_messages:
            assert isinstance(email_message, EmailMultiAlternatives)
            # Here, we replace the image URLs used in development with
            # chat.zulip.org URLs, so that web email providers like Gmail
            # will be able to fetch the illustrations used in the emails.
            assert isinstance(email_message.alternatives[0], EmailAlternative)
            original_content = email_message.alternatives[0].content
            original_mimetype = email_message.alternatives[0].mimetype
            assert isinstance(original_content, str)
            email_message.alternatives[0] = EmailAlternative(
                content=original_content.replace(
                    localhost_email_images_base_url, czo_email_images_base_url
                ),
                mimetype=original_mimetype,
            )

            email_message.to = [get_forward_address()]

    # This wrapper function exists to allow tests easily to mock the
    # step of trying to send the emails. Previously, we had mocked
    # Django's connection.send_messages(), which caused unexplained
    # test failures when running test-backend at very high
    # concurrency.
    def _do_send_messages(self, email_messages: Sequence[EmailMessage]) -> int:
        return super().send_messages(email_messages)  # nocoverage

    @override
    def send_messages(self, email_messages: Sequence[EmailMessage]) -> int:
        num_sent = len(email_messages)
        if get_forward_address():
            self.prepare_email_messages_for_forwarding(email_messages)
            num_sent = self._do_send_messages(email_messages)

        if settings.DEVELOPMENT_LOG_EMAILS:
            for email in email_messages:
                self.log_email(email)
                email_log_url = settings.ROOT_DOMAIN_URI + "/emails"
                logging.info("Emails sent in development are available at %s", email_log_url)
        return num_sent


class PersistentSMTPEmailBackend(EmailBackend):
    @override
    def open(self, **kwargs: Any) -> bool:
        if super().open(**kwargs):
            self.opened_at = timezone_now()
        # Always return False so that Django's send_messages does not
        # auto-close the connection after sending.
        return False

    @smtp_connection_backoff
    def ensure_connected(self) -> None:
        """Open the connection if needed, and validate that it is
        still alive, reconnecting if necessary."""
        self.open()
        self._validate_or_reconnect()

    def _validate_or_reconnect(self) -> None:
        """Validate that the existing SMTP connection is still alive,
        reconnecting if necessary. Called with backoff protection,
        not from open()."""
        status = None
        time_elapsed = (timezone_now() - self.opened_at).total_seconds() / 60
        if (
            settings.EMAIL_MAX_CONNECTION_LIFETIME_IN_MINUTES is None
            or time_elapsed <= settings.EMAIL_MAX_CONNECTION_LIFETIME_IN_MINUTES
        ):
            # No-op to ensure that we don't return a connection that has been
            # closed by the mail server.
            try:
                assert self.connection is not None
                status = self.connection.noop()[0]
            except Exception:
                pass
        if status is None or status != 250:
            # Close and connect again.
            super().close()
            self.open()