playground/aws_ses.py

import concurrent.futures
import functools
import logging
import pprint
import time
import typing

import boto3
import botocore.client

DnsRecord = typing.TypedDict(
    "DnsRecord",
    {
        "purpose": str,
        "type": str,
        "name": str,
        "value": str,
        "priority": typing.Optional[str],
    },
)
VerificationStatus = typing.Literal["Pending", "Success", "Failed", "TemporaryFailure", "NotStarted"]

logger = logging.getLogger(__name__)


def retry(tries: int = 5, delay: int = 1.5):
    def wrapper(fn):
        @functools.wraps(fn)
        def wrapped(*args, **kwargs):
            nonlocal tries
            while tries:
                try:
                    return fn(*args, **kwargs)
                except Exception:
                    logger.exception(f'{fn} returned error, retrying {tries} more times')
                time.sleep(delay)
                tries -= 1

        return wrapped

    return wrapper


class AwsSesService:
    def __init__(self, client: typing.Optional[botocore.client.BaseClient] = None, timeout: int = 5):
        self.ses = client or boto3.client("ses")
        self.region = self.ses._client_config.region_name
        self.timeout = timeout

    # @retry()
    def configure_domain_identity(
        self, domain: str, mail_subdomain: typing.Optional[str] = None
    ) -> typing.List[DnsRecord]:
        """
        Creates up mail identities for sending email through a domain.

        DKIM-signed messages help receiving mail servers validate that a message was not forged or altered in transit.

        :param domain: domain name to set up email identity
        :return: DNS records that need to be set
        """

        res_id = self.ses.verify_domain_identity(Domain=domain)
        id_token = res_id["VerificationToken"]

        res_dkim = self.ses.verify_domain_dkim(Domain=domain)
        dns_records = [
            {
                "purpose": "identity",
                "type": "TXT",
                "name": domain,
                "value": id_token,
            },
            *[
                {
                    "purpose": "dkim",
                    "type": "CNAME",
                    "name": f"{dkim_token}._domainkey.{domain}",
                    "value": f"{dkim_token}.dkim.amazonses.com",
                }
                for dkim_token in res_dkim["DkimTokens"]
            ],
        ]
        if mail_subdomain:
            more_dns = self.configure_from_domain(domain, mail_subdomain)
            dns_records.extend(more_dns)
        return dns_records

    @retry()
    def check_dkim_verification_status(self, domain: str) -> VerificationStatus:
        """
        Returns the DKIM verification status for an email identity.

        :param domain: domain used to set up the mail identity
        :return: one of {Pending, Success, Failed, TemporaryFailure, NotStarted}
        """
        res = self.ses.get_identity_dkim_attributes(Identities=[domain])
        return res["DkimAttributes"][domain]["DkimVerificationStatus"]

    @retry()
    def check_id_verification_status(self, domain: str) -> VerificationStatus:
        """
        Returns the id verification status of an email identity.

        :param domain: domain used to set up the mail identity
        :return: one of {Pending, Success, Failed, TemporaryFailure, NotStarted}
        """
        res = self.ses.get_identity_verification_attributes(Identities=[domain])
        return res["VerificationAttributes"][domain]["VerificationStatus"]

    def configure_from_domain(self, domain: str, mail_subdomain: str) -> typing.List[DnsRecord]:
        """
        Allows using `mail_subdomain` as FROM address when sending emails.
        Messages sent through Amazon SES will be marked as originating from your domain instead of a subdomain of amazon.com.

        :param domain: mail identity
        :param mail_subdomain: a subdomain
        :return: DNS records that need to be set
        """
        _ = self.ses.set_identity_mail_from_domain(
            Identity=domain,
            MailFromDomain=mail_subdomain,
            BehaviorOnMXFailure="UseDefaultValue",  # uses $region.amazonses.com if MX records not present on the mail subdomain
        )
        return [
            {
                "purpose": "mail_from_domain",
                "type": "MX",
                "name": mail_subdomain,
                "value": f"feedback-smtp.{self.region}.amazonses.com",
                "priority": "10",
            },
            {
                "purpose": "mail_from_domain",
                "type": "TXT",
                "name": mail_subdomain,
                "value": f'"v=spf1 include:amazonses.com ~all"',
            },
        ]

    @retry()
    def check_custom_from_domain_status(self, domain: str) -> VerificationStatus:
        """
        Returns the status of custom email FROM request verification.

        :param domain: domain used to set up email identity
        :return: one of {Pending, Success, Failed, TemporaryFailure}
        """
        res = self.ses.get_identity_mail_from_domain_attributes(Identities=[domain])
        return res["MailFromDomainAttributes"][domain].get("MailFromDomainStatus", 'NotStarted')

    def bulk_check_domain_verification(self, identity_domain: str, timeout: int = 5) -> dict:
        tasks = {
            "status_id": functools.partial(self.check_id_verification_status, identity_domain),
            "status_dkim": functools.partial(self.check_dkim_verification_status, identity_domain),
            "status_custom": functools.partial(self.check_custom_from_domain_status, identity_domain),
        }
        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as pool:
            futures = [pool.submit(t) for _, t in tasks.items()]
            results = dict(zip(tasks.keys(), [f.result(timeout=timeout) for f in futures]))

        return {
            **results,
            "identity_domain": identity_domain,
        }

    @retry()
    def bulk_configure_email_domains(
        self,
        identity_domain: str,
        mail_subdomain: typing.Optional[str] = None,
        timeout: int = 10,
    ) -> dict:
        tasks = {
            "dns": functools.partial(self.configure_domain_identity, identity_domain, mail_subdomain),
        }

        # send all requests in parallel
        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as pool:
            futures = [pool.submit(t) for _, t in tasks.items()]
            results = dict(zip(tasks.keys(), [f.result(timeout=timeout) for f in futures]))
        return {
            **results,
            "identity_domain": identity_domain,
            "mail_subdomain": mail_subdomain,
        }


if __name__ == "__main__":
    ses = AwsSesService()
    logging.basicConfig(level=logging.INFO)
    # print(ses.bulk_configure_email_domains('abdus100.dev', 'mail.abdus100.dev'))
    # print(ses.check_custom_from_domain_status('abdus101.dev'))

    tasks = [
        functools.partial(ses.bulk_check_domain_verification, d) for d in ['abdus.dev', 'abdus1.dev', 'abdus2.dev']
    ] * 3

    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as pool:
        futures = [pool.submit(t) for t in tasks]
        results = [f.result(timeout=5) for f in futures]
    pprint.pp(results)