4.53.2

app/.github/workflows/main.yml

@@ -163,7 +163,7 @@ jobs:
         uses: docker/build-push-action@v3
         with:
           context: .
-          platforms: linux/amd64,linux/arm64
+          platforms: linux/amd64
           push: true
           tags: ${{ steps.meta.outputs.tags }}
 

app/app/admin_model.py

@@ -34,6 +34,7 @@ from app.models import (
     DeletedAlias,
     DomainDeletedAlias,
     PartnerUser,
+    AliasMailbox,
 )
 from app.newsletter_utils import send_newsletter_to_user, send_newsletter_to_address
 
@@ -785,6 +786,25 @@ class EmailSearchHelpers:
     def mailbox_count(user: User) -> int:
         return Mailbox.filter_by(user_id=user.id).order_by(Mailbox.id.desc()).count()
 
+    @staticmethod
+    def alias_mailboxes(alias: Alias) -> list[Mailbox]:
+        return (
+            Session.query(Mailbox)
+            .filter(Mailbox.id == Alias.mailbox_id, Alias.id == alias.id)
+            .union(
+                Session.query(Mailbox)
+                .join(AliasMailbox, Mailbox.id == AliasMailbox.mailbox_id)
+                .filter(AliasMailbox.alias_id == alias.id)
+            )
+            .order_by(Mailbox.id)
+            .limit(10)
+            .all()
+        )
+
+    @staticmethod
+    def alias_mailbox_count(alias: Alias) -> int:
+        return len(alias.mailboxes)
+
     @staticmethod
     def alias_list(user: User) -> list[Alias]:
         return (

app/app/alias_utils.py

@@ -1,6 +1,7 @@
 import csv
 from io import StringIO
 import re
+from dataclasses import dataclass
 from typing import Optional, Tuple
 
 from email_validator import validate_email, EmailNotValidError
@@ -23,6 +24,7 @@ from app.email_utils import (
     send_cannot_create_domain_alias,
     send_email,
     render,
+    sl_formataddr,
 )
 from app.errors import AliasInTrashError
 from app.events.event_dispatcher import EventDispatcher
@@ -30,6 +32,7 @@ from app.events.generated.event_pb2 import (
     AliasDeleted,
     AliasStatusChanged,
     EventContent,
+    AliasCreated,
 )
 from app.log import LOG
 from app.models import (
@@ -501,6 +504,28 @@ def transfer_alias(alias, new_user, new_mailboxes: [Mailbox]):
     alias.disable_pgp = False
     alias.pinned = False
 
+    EventDispatcher.send_event(
+        old_user,
+        EventContent(
+            alias_deleted=AliasDeleted(
+                id=alias.id,
+                email=alias.email,
+            )
+        ),
+    )
+    EventDispatcher.send_event(
+        new_user,
+        EventContent(
+            alias_created=AliasCreated(
+                id=alias.id,
+                email=alias.email,
+                note=alias.note,
+                enabled=alias.enabled,
+                created_at=int(alias.created_at.timestamp),
+            )
+        ),
+    )
+
     Session.commit()
 
 
@@ -518,3 +543,30 @@ def change_alias_status(alias: Alias, enabled: bool, commit: bool = False):
 
     if commit:
         Session.commit()
+
+
+@dataclass
+class AliasRecipientName:
+    name: str
+    message: Optional[str] = None
+
+
+def get_alias_recipient_name(alias: Alias) -> AliasRecipientName:
+    """
+    Logic:
+    1. If alias has name, use it
+    2. If alias has custom domain, and custom domain has name, use it
+    3. Otherwise, use the alias email as the recipient
+    """
+    if alias.name:
+        return AliasRecipientName(
+            name=sl_formataddr((alias.name, alias.email)),
+            message=f"Put alias name {alias.name} in from header",
+        )
+    elif alias.custom_domain:
+        if alias.custom_domain.name:
+            return AliasRecipientName(
+                name=sl_formataddr((alias.custom_domain.name, alias.email)),
+                message=f"Put domain default alias name {alias.custom_domain.name} in from header",
+            )
+    return AliasRecipientName(name=alias.email)

app/app/api/views/alias.py

@@ -424,7 +424,7 @@ def create_contact_route(alias_id):
     contact_address = data.get("contact")
 
     try:
-        contact = create_contact(g.user, alias, contact_address)
+        contact = create_contact(alias, contact_address)
     except ErrContactErrorUpgradeNeeded as err:
         return jsonify(error=err.error_for_user()), 403
     except (ErrAddressInvalid, CannotCreateContactForReverseAlias) as err:

app/app/api/views/auth.py

@@ -52,8 +52,12 @@ def auth_login():
     password = data.get("password")
     device = data.get("device")
 
-    email = sanitize_email(data.get("email"))
+    email = data.get("email")
-    canonical_email = canonicalize_email(data.get("email"))
+    if not email:
+        LoginEvent(LoginEvent.ActionType.failed, LoginEvent.Source.api).send()
+        return jsonify(error="Email or password incorrect"), 400
+    email = sanitize_email(email)
+    canonical_email = canonicalize_email(email)
 
     user = User.get_by(email=email) or User.get_by(email=canonical_email)
 

app/app/api/views/user_info.py

@@ -87,7 +87,7 @@ def update_user_info():
                 File.delete(file.id)
                 s3.delete(file.path)
                 Session.flush()
-        else:
+        if data["profile_picture"] is not None:
             raw_data = base64.decodebytes(data["profile_picture"].encode())
             if detect_image_format(raw_data) == ImageFormat.Unknown:
                 return jsonify(error="Unsupported image format"), 400

app/app/config.py

@@ -653,7 +653,13 @@ def read_partner_dict(var: str) -> dict[int, str]:
     return res
 
 
-PARTNER_DOMAINS: dict[int, str] = read_partner_dict("PARTNER_DOMAINS")
+PARTNER_DNS_CUSTOM_DOMAINS: dict[int, str] = read_partner_dict(
-PARTNER_DOMAIN_VALIDATION_PREFIXES: dict[int, str] = read_partner_dict(
+    "PARTNER_DNS_CUSTOM_DOMAINS"
-    "PARTNER_DOMAIN_VALIDATION_PREFIXES"
+)
+PARTNER_CUSTOM_DOMAIN_VALIDATION_PREFIXES: dict[int, str] = read_partner_dict(
+    "PARTNER_CUSTOM_DOMAIN_VALIDATION_PREFIXES"
+)
+
+MAILBOX_VERIFICATION_OVERRIDE_CODE: Optional[str] = os.environ.get(
+    "MAILBOX_VERIFICATION_OVERRIDE_CODE", None
 )

app/app/contact_utils.py

@@ -0,0 +1,113 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+from sqlalchemy.exc import IntegrityError
+
+from app.db import Session
+from app.email_utils import generate_reply_email, parse_full_address
+from app.email_validation import is_valid_email
+from app.log import LOG
+from app.models import Contact, Alias
+from app.utils import sanitize_email
+
+
+class ContactCreateError(Enum):
+    InvalidEmail = "Invalid email"
+    NotAllowed = "Your plan does not allow to create contacts"
+
+
+@dataclass
+class ContactCreateResult:
+    contact: Optional[Contact]
+    created: bool
+    error: Optional[ContactCreateError]
+
+
+def __update_contact_if_needed(
+    contact: Contact, name: Optional[str], mail_from: Optional[str]
+) -> ContactCreateResult:
+    if name and contact.name != name:
+        LOG.d(f"Setting {contact} name to {name}")
+        contact.name = name
+        Session.commit()
+    if mail_from and contact.mail_from is None:
+        LOG.d(f"Setting {contact} mail_from to {mail_from}")
+        contact.mail_from = mail_from
+        Session.commit()
+    return ContactCreateResult(contact, created=False, error=None)
+
+
+def create_contact(
+    email: str,
+    alias: Alias,
+    name: Optional[str] = None,
+    mail_from: Optional[str] = None,
+    allow_empty_email: bool = False,
+    automatic_created: bool = False,
+    from_partner: bool = False,
+) -> ContactCreateResult:
+    # If user cannot create contacts, they still need to be created when receiving an email for an alias
+    if not automatic_created and not alias.user.can_create_contacts():
+        return ContactCreateResult(
+            None, created=False, error=ContactCreateError.NotAllowed
+        )
+    # Parse emails with form 'name <email>'
+    try:
+        email_name, email = parse_full_address(email)
+    except ValueError:
+        email = ""
+        email_name = ""
+    # If no name is explicitly given try to get it from the parsed email
+    if name is None:
+        name = email_name[: Contact.MAX_NAME_LENGTH]
+    else:
+        name = name[: Contact.MAX_NAME_LENGTH]
+    # If still no name is there, make sure the name is None instead of empty string
+    if not name:
+        name = None
+    if name is not None and "\x00" in name:
+        LOG.w("Cannot use contact name because has \\x00")
+        name = ""
+    # Sanitize email and if it's not valid only allow to create a contact if it's explicitly allowed. Otherwise fail
+    email = sanitize_email(email, not_lower=True)
+    if not is_valid_email(email):
+        LOG.w(f"invalid contact email {email}")
+        if not allow_empty_email:
+            return ContactCreateResult(
+                None, created=False, error=ContactCreateError.InvalidEmail
+            )
+        LOG.d("Create a contact with invalid email for %s", alias)
+        # either reuse a contact with empty email or create a new contact with empty email
+        email = ""
+    # If contact exists, update name and mail_from if needed
+    contact = Contact.get_by(alias_id=alias.id, website_email=email)
+    if contact is not None:
+        return __update_contact_if_needed(contact, name, mail_from)
+    # Create the contact
+    reply_email = generate_reply_email(email, alias)
+    try:
+        flags = Contact.FLAG_PARTNER_CREATED if from_partner else 0
+        contact = Contact.create(
+            user_id=alias.user_id,
+            alias_id=alias.id,
+            website_email=email,
+            name=name,
+            reply_email=reply_email,
+            mail_from=mail_from,
+            automatic_created=automatic_created,
+            flags=flags,
+            invalid_email=email == "",
+            commit=True,
+        )
+        LOG.d(
+            f"Created contact {contact} for alias {alias} with email {email} invalid_email={contact.invalid_email}"
+        )
+    except IntegrityError:
+        Session.rollback()
+        LOG.info(
+            f"Contact with email {email} for alias_id {alias.id} already existed, fetching from DB"
+        )
+        contact = Contact.get_by(alias_id=alias.id, website_email=email)
+        return __update_contact_if_needed(contact, name, mail_from)
+    return ContactCreateResult(contact, created=True, error=None)

app/app/custom_domain_utils.py

@@ -1,15 +1,18 @@
+import arrow
 import re
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import Optional
+from typing import List, Optional
 
+from app.config import JOB_DELETE_DOMAIN
 from app.db import Session
 from app.email_utils import get_email_domain_part
 from app.log import LOG
-from app.models import User, CustomDomain, SLDomain, Mailbox
+from app.models import User, CustomDomain, SLDomain, Mailbox, Job, DomainMailbox
 
 _ALLOWED_DOMAIN_REGEX = re.compile(r"^(?!-)[A-Za-z0-9-]{1,63}(?<!-)$")
+_MAX_MAILBOXES_PER_DOMAIN = 20
 
 
 @dataclass
@@ -43,6 +46,20 @@ class CannotUseDomainReason(Enum):
             raise Exception("Invalid CannotUseDomainReason")
 
 
+class CannotSetCustomDomainMailboxesCause(Enum):
+    InvalidMailbox = "Something went wrong, please retry"
+    NoMailboxes = "You must select at least 1 mailbox"
+    TooManyMailboxes = (
+        f"You can only set up to {_MAX_MAILBOXES_PER_DOMAIN} mailboxes per domain"
+    )
+
+
+@dataclass
+class SetCustomDomainMailboxesResult:
+    success: bool
+    reason: Optional[CannotSetCustomDomainMailboxesCause] = None
+
+
 def is_valid_domain(domain: str) -> bool:
     """
     Checks that a domain is valid according to RFC 1035
@@ -126,3 +143,52 @@ def create_custom_domain(
         success=True,
         instance=new_custom_domain,
     )
+
+
+def delete_custom_domain(domain: CustomDomain):
+    # Schedule delete domain job
+    LOG.w("schedule delete domain job for %s", domain)
+    domain.pending_deletion = True
+    Job.create(
+        name=JOB_DELETE_DOMAIN,
+        payload={"custom_domain_id": domain.id},
+        run_at=arrow.now(),
+        commit=True,
+    )
+
+
+def set_custom_domain_mailboxes(
+    user_id: int, custom_domain: CustomDomain, mailbox_ids: List[int]
+) -> SetCustomDomainMailboxesResult:
+    if len(mailbox_ids) == 0:
+        return SetCustomDomainMailboxesResult(
+            success=False, reason=CannotSetCustomDomainMailboxesCause.NoMailboxes
+        )
+    elif len(mailbox_ids) > _MAX_MAILBOXES_PER_DOMAIN:
+        return SetCustomDomainMailboxesResult(
+            success=False, reason=CannotSetCustomDomainMailboxesCause.TooManyMailboxes
+        )
+
+    mailboxes = (
+        Session.query(Mailbox)
+        .filter(
+            Mailbox.id.in_(mailbox_ids),
+            Mailbox.user_id == user_id,
+            Mailbox.verified == True,  # noqa: E712
+        )
+        .all()
+    )
+    if len(mailboxes) != len(mailbox_ids):
+        return SetCustomDomainMailboxesResult(
+            success=False, reason=CannotSetCustomDomainMailboxesCause.InvalidMailbox
+        )
+
+    # first remove all existing domain-mailboxes links
+    DomainMailbox.filter_by(domain_id=custom_domain.id).delete()
+    Session.flush()
+
+    for mailbox in mailboxes:
+        DomainMailbox.create(domain_id=custom_domain.id, mailbox_id=mailbox.id)
+
+    Session.commit()
+    return SetCustomDomainMailboxesResult(success=True)

app/app/custom_domain_validation.py

@@ -1,15 +1,17 @@
 from dataclasses import dataclass
-from typing import Optional
+from typing import List, Optional
 
 from app import config
 from app.constants import DMARC_RECORD
 from app.db import Session
 from app.dns_utils import (
+    MxRecord,
     DNSClient,
     is_mx_equivalent,
     get_network_dns_client,
 )
 from app.models import CustomDomain
+from app.utils import random_string
 
 
 @dataclass
@@ -28,10 +30,10 @@ class CustomDomainValidation:
     ):
         self.dkim_domain = dkim_domain
         self._dns_client = dns_client
-        self._partner_domains = partner_domains or config.PARTNER_DOMAINS
+        self._partner_domains = partner_domains or config.PARTNER_DNS_CUSTOM_DOMAINS
         self._partner_domain_validation_prefixes = (
             partner_domains_validation_prefixes
-            or config.PARTNER_DOMAIN_VALIDATION_PREFIXES
+            or config.PARTNER_CUSTOM_DOMAIN_VALIDATION_PREFIXES
         )
 
     def get_ownership_verification_record(self, domain: CustomDomain) -> str:
@@ -41,8 +43,36 @@ class CustomDomainValidation:
             and domain.partner_id in self._partner_domain_validation_prefixes
         ):
             prefix = self._partner_domain_validation_prefixes[domain.partner_id]
+
+        if not domain.ownership_txt_token:
+            domain.ownership_txt_token = random_string(30)
+            Session.commit()
+
         return f"{prefix}-verification={domain.ownership_txt_token}"
 
+    def get_expected_mx_records(self, domain: CustomDomain) -> list[MxRecord]:
+        records = []
+        if domain.partner_id is not None and domain.partner_id in self._partner_domains:
+            domain = self._partner_domains[domain.partner_id]
+            records.append(MxRecord(10, f"mx1.{domain}."))
+            records.append(MxRecord(20, f"mx2.{domain}."))
+        else:
+            # Default ones
+            for priority, domain in config.EMAIL_SERVERS_WITH_PRIORITY:
+                records.append(MxRecord(priority, domain))
+
+        return records
+
+    def get_expected_spf_domain(self, domain: CustomDomain) -> str:
+        if domain.partner_id is not None and domain.partner_id in self._partner_domains:
+            return self._partner_domains[domain.partner_id]
+        else:
+            return config.EMAIL_DOMAIN
+
+    def get_expected_spf_record(self, domain: CustomDomain) -> str:
+        spf_domain = self.get_expected_spf_domain(domain)
+        return f"v=spf1 include:{spf_domain} ~all"
+
     def get_dkim_records(self, domain: CustomDomain) -> {str: str}:
         """
         Get a list of dkim records to set up. Depending on the custom_domain, whether if it's from a partner or not,
@@ -116,11 +146,12 @@ class CustomDomainValidation:
         self, custom_domain: CustomDomain
     ) -> DomainValidationResult:
         mx_domains = self._dns_client.get_mx_domains(custom_domain.domain)
+        expected_mx_records = self.get_expected_mx_records(custom_domain)
 
-        if not is_mx_equivalent(mx_domains, config.EMAIL_SERVERS_WITH_PRIORITY):
+        if not is_mx_equivalent(mx_domains, expected_mx_records):
             return DomainValidationResult(
                 success=False,
-                errors=[f"{priority} {domain}" for (priority, domain) in mx_domains],
+                errors=[f"{record.priority} {record.domain}" for record in mx_domains],
             )
         else:
             custom_domain.verified = True
@@ -131,16 +162,19 @@ class CustomDomainValidation:
         self, custom_domain: CustomDomain
     ) -> DomainValidationResult:
         spf_domains = self._dns_client.get_spf_domain(custom_domain.domain)
-        if config.EMAIL_DOMAIN in spf_domains:
+        expected_spf_domain = self.get_expected_spf_domain(custom_domain)
+        if expected_spf_domain in spf_domains:
             custom_domain.spf_verified = True
             Session.commit()
             return DomainValidationResult(success=True, errors=[])
         else:
             custom_domain.spf_verified = False
             Session.commit()
+            txt_records = self._dns_client.get_txt_record(custom_domain.domain)
+            cleaned_records = self.__clean_spf_records(txt_records, custom_domain)
             return DomainValidationResult(
                 success=False,
-                errors=self._dns_client.get_txt_record(custom_domain.domain),
+                errors=cleaned_records,
             )
 
     def validate_dmarc_records(
@@ -155,3 +189,13 @@ class CustomDomainValidation:
             custom_domain.dmarc_verified = False
             Session.commit()
             return DomainValidationResult(success=False, errors=txt_records)
+
+    def __clean_spf_records(
+        self, txt_records: List[str], custom_domain: CustomDomain
+    ) -> List[str]:
+        final_records = []
+        verification_record = self.get_ownership_verification_record(custom_domain)
+        for record in txt_records:
+            if record != verification_record:
+                final_records.append(record)
+        return final_records

app/app/dashboard/views/alias_contact_manager.py

@@ -9,13 +9,10 @@ from sqlalchemy import and_, func, case
 from wtforms import StringField, validators, ValidationError
 
 # Need to import directly from config to allow modification from the tests
-from app import config, parallel_limiter
+from app import config, parallel_limiter, contact_utils
+from app.contact_utils import ContactCreateError
 from app.dashboard.base import dashboard_bp
 from app.db import Session
-from app.email_utils import (
-    generate_reply_email,
-    parse_full_address,
-)
 from app.email_validation import is_valid_email
 from app.errors import (
     CannotCreateContactForReverseAlias,
@@ -24,8 +21,8 @@ from app.errors import (
     ErrContactAlreadyExists,
 )
 from app.log import LOG
-from app.models import Alias, Contact, EmailLog, User
+from app.models import Alias, Contact, EmailLog
-from app.utils import sanitize_email, CSRFValidationForm
+from app.utils import CSRFValidationForm
 
 
 def email_validator():
@@ -51,7 +48,7 @@ def email_validator():
     return _check
 
 
-def create_contact(user: User, alias: Alias, contact_address: str) -> Contact:
+def create_contact(alias: Alias, contact_address: str) -> Contact:
     """
     Create a contact for a user. Can be restricted for new free users by enabling DISABLE_CREATE_CONTACTS_FOR_FREE_USERS.
     Can throw exceptions:
@@ -61,37 +58,23 @@ def create_contact(user: User, alias: Alias, contact_address: str) -> Contact:
     """
     if not contact_address:
         raise ErrAddressInvalid("Empty address")
-    try:
+    output = contact_utils.create_contact(email=contact_address, alias=alias)
-        contact_name, contact_email = parse_full_address(contact_address)
+    if output.error == ContactCreateError.InvalidEmail:
-    except ValueError:
         raise ErrAddressInvalid(contact_address)
-
+    elif output.error == ContactCreateError.NotAllowed:
-    contact_email = sanitize_email(contact_email)
-    if not is_valid_email(contact_email):
-        raise ErrAddressInvalid(contact_email)
-
-    contact = Contact.get_by(alias_id=alias.id, website_email=contact_email)
-    if contact:
-        raise ErrContactAlreadyExists(contact)
-
-    if not user.can_create_contacts():
         raise ErrContactErrorUpgradeNeeded()
+    elif output.error is not None:
+        raise ErrAddressInvalid("Invalid address")
+    elif not output.created:
+        raise ErrContactAlreadyExists(output.contact)
 
-    contact = Contact.create(
+    contact = output.contact
-        user_id=alias.user_id,
-        alias_id=alias.id,
-        website_email=contact_email,
-        name=contact_name,
-        reply_email=generate_reply_email(contact_email, alias),
-    )
-
     LOG.d(
         "create reverse-alias for %s %s, reverse alias:%s",
         contact_address,
         alias,
         contact.reply_email,
     )
-    Session.commit()
 
     return contact
 
@@ -261,7 +244,7 @@ def alias_contact_manager(alias_id):
             if new_contact_form.validate():
                 contact_address = new_contact_form.email.data.strip()
                 try:
-                    contact = create_contact(current_user, alias, contact_address)
+                    contact = create_contact(alias, contact_address)
                 except (
                     ErrContactErrorUpgradeNeeded,
                     ErrAddressInvalid,

app/app/dashboard/views/custom_domain.py

@@ -21,7 +21,9 @@ class NewCustomDomainForm(FlaskForm):
 @parallel_limiter.lock(only_when=lambda: request.method == "POST")
 def custom_domain():
     custom_domains = CustomDomain.filter_by(
-        user_id=current_user.id, is_sl_subdomain=False
+        user_id=current_user.id,
+        is_sl_subdomain=False,
+        pending_deletion=False,
     ).all()
     new_custom_domain_form = NewCustomDomainForm()
 

app/app/dashboard/views/domain_detail.py

@@ -1,26 +1,23 @@
 import re
 
-import arrow
 from flask import render_template, request, redirect, url_for, flash
 from flask_login import login_required, current_user
 from flask_wtf import FlaskForm
 from wtforms import StringField, validators, IntegerField
 
 from app.constants import DMARC_RECORD
-from app.config import EMAIL_SERVERS_WITH_PRIORITY, EMAIL_DOMAIN, JOB_DELETE_DOMAIN
+from app.config import EMAIL_SERVERS_WITH_PRIORITY, EMAIL_DOMAIN
+from app.custom_domain_utils import delete_custom_domain, set_custom_domain_mailboxes
 from app.custom_domain_validation import CustomDomainValidation
 from app.dashboard.base import dashboard_bp
 from app.db import Session
-from app.log import LOG
 from app.models import (
     CustomDomain,
     Alias,
     DomainDeletedAlias,
     Mailbox,
-    DomainMailbox,
     AutoCreateRule,
     AutoCreateRuleMailbox,
-    Job,
 )
 from app.regex_utils import regex_match
 from app.utils import random_string, CSRFValidationForm
@@ -39,8 +36,6 @@ def domain_detail_dns(custom_domain_id):
         custom_domain.ownership_txt_token = random_string(30)
         Session.commit()
 
-    spf_record = f"v=spf1 include:{EMAIL_DOMAIN} ~all"
-
     domain_validator = CustomDomainValidation(EMAIL_DOMAIN)
     csrf_form = CSRFValidationForm()
 
@@ -144,7 +139,9 @@ def domain_detail_dns(custom_domain_id):
         ownership_record=domain_validator.get_ownership_verification_record(
             custom_domain
         ),
+        expected_mx_records=domain_validator.get_expected_mx_records(custom_domain),
         dkim_records=domain_validator.get_dkim_records(custom_domain),
+        spf_record=domain_validator.get_expected_spf_record(custom_domain),
         dmarc_record=DMARC_RECORD,
         **locals(),
     )
@@ -222,40 +219,16 @@ def domain_detail(custom_domain_id):
             )
         elif request.form.get("form-name") == "update":
             mailbox_ids = request.form.getlist("mailbox_ids")
-            # check if mailbox is not tempered with
+            result = set_custom_domain_mailboxes(
-            mailboxes = []
+                user_id=current_user.id,
-            for mailbox_id in mailbox_ids:
+                custom_domain=custom_domain,
-                mailbox = Mailbox.get(mailbox_id)
+                mailbox_ids=mailbox_ids,
-                if (
-                    not mailbox
-                    or mailbox.user_id != current_user.id
-                    or not mailbox.verified
-                ):
-                    flash("Something went wrong, please retry", "warning")
-                    return redirect(
-                        url_for(
-                            "dashboard.domain_detail", custom_domain_id=custom_domain.id
-                        )
-                    )
-                mailboxes.append(mailbox)
-
-            if not mailboxes:
-                flash("You must select at least 1 mailbox", "warning")
-                return redirect(
-                    url_for(
-                        "dashboard.domain_detail", custom_domain_id=custom_domain.id
-                    )
             )
 
-            # first remove all existing domain-mailboxes links
+            if result.success:
-            DomainMailbox.filter_by(domain_id=custom_domain.id).delete()
-            Session.flush()
-
-            for mailbox in mailboxes:
-                DomainMailbox.create(domain_id=custom_domain.id, mailbox_id=mailbox.id)
-
-            Session.commit()
                 flash(f"{custom_domain.domain} mailboxes has been updated", "success")
+            else:
+                flash(result.reason.value, "warning")
 
             return redirect(
                 url_for("dashboard.domain_detail", custom_domain_id=custom_domain.id)
@@ -263,16 +236,8 @@ def domain_detail(custom_domain_id):
 
         elif request.form.get("form-name") == "delete":
             name = custom_domain.domain
-            LOG.d("Schedule deleting %s", custom_domain)
 
-            # Schedule delete domain job
+            delete_custom_domain(custom_domain)
-            LOG.w("schedule delete domain job for %s", custom_domain)
-            Job.create(
-                name=JOB_DELETE_DOMAIN,
-                payload={"custom_domain_id": custom_domain.id},
-                run_at=arrow.now(),
-                commit=True,
-            )
 
             flash(
                 f"{name} scheduled for deletion."

app/app/dns_utils.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
-from typing import List, Tuple, Optional
+from dataclasses import dataclass
+from typing import List, Optional
 
 import dns.resolver
 
@@ -8,8 +9,14 @@ from app.config import NAMESERVERS
 _include_spf = "include:"
 
 
+@dataclass
+class MxRecord:
+    priority: int
+    domain: str
+
+
 def is_mx_equivalent(
-    mx_domains: List[Tuple[int, str]], ref_mx_domains: List[Tuple[int, str]]
+    mx_domains: List[MxRecord], ref_mx_domains: List[MxRecord]
 ) -> bool:
     """
     Compare mx_domains with ref_mx_domains to see if they are equivalent.
@@ -18,14 +25,14 @@ def is_mx_equivalent(
     The priority order is taken into account but not the priority number.
     For example, [(1, domain1), (2, domain2)] is equivalent to [(10, domain1), (20, domain2)]
     """
-    mx_domains = sorted(mx_domains, key=lambda x: x[0])
+    mx_domains = sorted(mx_domains, key=lambda x: x.priority)
-    ref_mx_domains = sorted(ref_mx_domains, key=lambda x: x[0])
+    ref_mx_domains = sorted(ref_mx_domains, key=lambda x: x.priority)
 
     if len(mx_domains) < len(ref_mx_domains):
         return False
 
-    for i in range(len(ref_mx_domains)):
+    for actual, expected in zip(mx_domains, ref_mx_domains):
-        if mx_domains[i][1] != ref_mx_domains[i][1]:
+        if actual.domain != expected.domain:
             return False
 
     return True
@@ -37,7 +44,7 @@ class DNSClient(ABC):
         pass
 
     @abstractmethod
-    def get_mx_domains(self, hostname: str) -> List[Tuple[int, str]]:
+    def get_mx_domains(self, hostname: str) -> List[MxRecord]:
         pass
 
     def get_spf_domain(self, hostname: str) -> List[str]:
@@ -81,7 +88,7 @@ class NetworkDNSClient(DNSClient):
         except Exception:
             return None
 
-    def get_mx_domains(self, hostname: str) -> List[Tuple[int, str]]:
+    def get_mx_domains(self, hostname: str) -> List[MxRecord]:
         """
         return list of (priority, domain name) sorted by priority (lowest priority first)
         domain name ends with a "." at the end.
@@ -92,14 +99,14 @@ class NetworkDNSClient(DNSClient):
             for a in answers:
                 record = a.to_text()  # for ex '20 alt2.aspmx.l.google.com.'
                 parts = record.split(" ")
-                ret.append((int(parts[0]), parts[1]))
+                ret.append(MxRecord(priority=int(parts[0]), domain=parts[1]))
-            return sorted(ret, key=lambda x: x[0])
+            return sorted(ret, key=lambda x: x.priority)
         except Exception:
             return []
 
     def get_txt_record(self, hostname: str) -> List[str]:
         try:
-            answers = self._resolver.resolve(hostname, "TXT", search=True)
+            answers = self._resolver.resolve(hostname, "TXT", search=False)
             ret = []
             for a in answers:  # type: dns.rdtypes.ANY.TXT.TXT
                 for record in a.strings:
@@ -112,14 +119,14 @@ class NetworkDNSClient(DNSClient):
 class InMemoryDNSClient(DNSClient):
     def __init__(self):
         self.cname_records: dict[str, Optional[str]] = {}
-        self.mx_records: dict[str, List[Tuple[int, str]]] = {}
+        self.mx_records: dict[str, List[MxRecord]] = {}
         self.spf_records: dict[str, List[str]] = {}
         self.txt_records: dict[str, List[str]] = {}
 
     def set_cname_record(self, hostname: str, cname: str):
         self.cname_records[hostname] = cname
 
-    def set_mx_records(self, hostname: str, mx_list: List[Tuple[int, str]]):
+    def set_mx_records(self, hostname: str, mx_list: List[MxRecord]):
         self.mx_records[hostname] = mx_list
 
     def set_txt_record(self, hostname: str, txt_list: List[str]):
@@ -128,9 +135,9 @@ class InMemoryDNSClient(DNSClient):
     def get_cname_record(self, hostname: str) -> Optional[str]:
         return self.cname_records.get(hostname)
 
-    def get_mx_domains(self, hostname: str) -> List[Tuple[int, str]]:
+    def get_mx_domains(self, hostname: str) -> List[MxRecord]:
         mx_list = self.mx_records.get(hostname, [])
-        return sorted(mx_list, key=lambda x: x[0])
+        return sorted(mx_list, key=lambda x: x.priority)
 
     def get_txt_record(self, hostname: str) -> List[str]:
         return self.txt_records.get(hostname, [])
@@ -140,5 +147,5 @@ def get_network_dns_client() -> NetworkDNSClient:
     return NetworkDNSClient(NAMESERVERS)
 
 
-def get_mx_domains(hostname: str) -> [(int, str)]:
+def get_mx_domains(hostname: str) -> List[MxRecord]:
     return get_network_dns_client().get_mx_domains(hostname)

app/app/email_utils.py

@@ -592,7 +592,7 @@ def email_can_be_used_as_mailbox(email_address: str) -> bool:
 
     from app.models import CustomDomain
 
-    if CustomDomain.get_by(domain=domain, verified=True):
+    if CustomDomain.get_by(domain=domain, is_sl_subdomain=True, verified=True):
         LOG.d("domain %s is a SimpleLogin custom domain", domain)
         return False
 
@@ -657,7 +657,7 @@ def get_mx_domain_list(domain) -> [str]:
     """
     priority_domains = get_mx_domains(domain)
 
-    return [d[:-1] for _, d in priority_domains]
+    return [d.domain[:-1] for d in priority_domains]
 
 
 def personal_email_already_used(email_address: str) -> bool:

app/app/events/event_dispatcher.py

@@ -64,15 +64,9 @@ class EventDispatcher:
             )
             return
 
-        if config.EVENT_WEBHOOK_ENABLED_USER_IDS is not None:
-            if user.id not in config.EVENT_WEBHOOK_ENABLED_USER_IDS:
-                return
-
         partner_user = EventDispatcher.__partner_user(user.id)
         if not partner_user:
-            LOG.i(
+            LOG.i(f"Not sending events because there's no partner user for user {user}")
-                f"Not sending events because there's no partner user for  user {user}"
-            )
             return
 
         event = event_pb2.Event(
@@ -84,7 +78,9 @@ class EventDispatcher:
 
         serialized = event.SerializeToString()
         dispatcher.send(serialized)
-        newrelic.agent.record_custom_metric("Custom/events_stored", 1)
+
+        event_type = content.WhichOneof("content")
+        newrelic.agent.record_custom_event("EventStoredToDb", {"type": event_type})
         LOG.i("Sent event to the dispatcher")
 
     @staticmethod

app/app/mailbox_utils.py

@@ -213,6 +213,9 @@ def generate_activation_code(
 ) -> MailboxActivation:
     clear_activation_codes_for_mailbox(mailbox)
     if use_digit_code:
+        if config.MAILBOX_VERIFICATION_OVERRIDE_CODE:
+            code = config.MAILBOX_VERIFICATION_OVERRIDE_CODE
+        else:
             code = "{:06d}".format(random.randint(1, 999999))
     else:
         code = secrets.token_urlsafe(16)

app/app/models.py

@@ -336,7 +336,7 @@ class Fido(Base, ModelMixin):
 class User(Base, ModelMixin, UserMixin, PasswordOracle):
     __tablename__ = "users"
 
-    FLAG_FREE_DISABLE_CREATE_ALIAS = 1 << 0
+    FLAG_DISABLE_CREATE_CONTACTS = 1 << 0
     FLAG_CREATED_FROM_PARTNER = 1 << 1
     FLAG_FREE_OLD_ALIAS_LIMIT = 1 << 2
     FLAG_CREATED_ALIAS_FROM_PARTNER = 1 << 3
@@ -543,7 +543,7 @@ class User(Base, ModelMixin, UserMixin, PasswordOracle):
     # bitwise flags. Allow for future expansion
     flags = sa.Column(
         sa.BigInteger,
-        default=FLAG_FREE_DISABLE_CREATE_ALIAS,
+        default=FLAG_DISABLE_CREATE_CONTACTS,
         server_default="0",
         nullable=False,
     )
@@ -1168,7 +1168,7 @@ class User(Base, ModelMixin, UserMixin, PasswordOracle):
     def can_create_contacts(self) -> bool:
         if self.is_premium():
             return True
-        if self.flags & User.FLAG_FREE_DISABLE_CREATE_ALIAS == 0:
+        if self.flags & User.FLAG_DISABLE_CREATE_CONTACTS == 0:
             return True
         return not config.DISABLE_CREATE_CONTACTS_FOR_FREE_USERS
 
@@ -1863,6 +1863,8 @@ class Contact(Base, ModelMixin):
 
     MAX_NAME_LENGTH = 512
 
+    FLAG_PARTNER_CREATED = 1 << 0
+
     __tablename__ = "contact"
 
     __table_args__ = (
@@ -1921,6 +1923,9 @@ class Contact(Base, ModelMixin):
     # whether contact is created automatically during the forward phase
     automatic_created = sa.Column(sa.Boolean, nullable=True, default=False)
 
+    # contact flags
+    flags = sa.Column(sa.Integer, nullable=False, default=0, server_default="0")
+
     @property
     def email(self):
         return self.website_email
@@ -2427,6 +2432,10 @@ class CustomDomain(Base, ModelMixin):
         server_default=None,
     )
 
+    pending_deletion = sa.Column(
+        sa.Boolean, nullable=False, default=False, server_default="0"
+    )
+
     __table_args__ = (
         Index(
             "ix_unique_domain",  # Index name
@@ -2434,6 +2443,8 @@ class CustomDomain(Base, ModelMixin):
             unique=True,
             postgresql_where=Column("ownership_verified"),
         ),  # The condition
+        Index("ix_custom_domain_user_id", "user_id"),
+        Index("ix_custom_domain_pending_deletion", "pending_deletion"),
     )
 
     user = orm.relationship(User, foreign_keys=[user_id], backref="custom_domains")
@@ -2755,9 +2766,9 @@ class Mailbox(Base, ModelMixin):
 
         from app.email_utils import get_email_local_part
 
-        mx_domains: [(int, str)] = get_mx_domains(get_email_local_part(self.email))
+        mx_domains = get_mx_domains(get_email_local_part(self.email))
         # Proton is the first domain
-        if mx_domains and mx_domains[0][1] in (
+        if mx_domains and mx_domains[0].domain in (
             "mail.protonmail.ch.",
             "mailsec.protonmail.ch.",
         ):

app/app/sentry_utils.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+from sentry_sdk.types import Event, Hint
+
+_HTTP_CODES_TO_IGNORE = [416]
+
+
+def _should_send(_event: Event, hint: Hint) -> bool:
+    # Check if this is an HTTP Exception event
+    if "exc_info" in hint:
+        exc_type, exc_value, exc_traceback = hint["exc_info"]
+        # Check if it's a Werkzeug HTTPException (raised for HTTP status codes)
+        if hasattr(exc_value, "code") and exc_value.code in _HTTP_CODES_TO_IGNORE:
+            return False
+    return True
+
+
+def sentry_before_send(event: Event, hint: Hint) -> Optional[Event]:
+    if _should_send(event, hint):
+        return event
+    return None

app/cron.py

@@ -14,6 +14,7 @@ from sqlalchemy.sql import Insert, text
 from app import s3, config
 from app.alias_utils import nb_email_log_for_mailbox
 from app.api.views.apple import verify_receipt
+from app.custom_domain_validation import CustomDomainValidation
 from app.db import Session
 from app.dns_utils import get_mx_domains, is_mx_equivalent
 from app.email_utils import (
@@ -905,9 +906,11 @@ def check_custom_domain():
             LOG.i("custom domain has been deleted")
 
 
-def check_single_custom_domain(custom_domain):
+def check_single_custom_domain(custom_domain: CustomDomain):
     mx_domains = get_mx_domains(custom_domain.domain)
-    if not is_mx_equivalent(mx_domains, config.EMAIL_SERVERS_WITH_PRIORITY):
+    validator = CustomDomainValidation(dkim_domain=config.EMAIL_DOMAIN)
+    expected_custom_domains = validator.get_expected_mx_records(custom_domain)
+    if not is_mx_equivalent(mx_domains, expected_custom_domains):
         user = custom_domain.user
         LOG.w(
             "The MX record is not correctly set for %s %s %s",

app/email_handler.py

@@ -52,8 +52,12 @@ from flanker.addresslib import address
 from flanker.addresslib.address import EmailAddress
 from sqlalchemy.exc import IntegrityError
 
-from app import pgp_utils, s3, config
+from app import pgp_utils, s3, config, contact_utils
-from app.alias_utils import try_auto_create, change_alias_status
+from app.alias_utils import (
+    try_auto_create,
+    change_alias_status,
+    get_alias_recipient_name,
+)
 from app.config import (
     EMAIL_DOMAIN,
     URL,
@@ -195,79 +199,16 @@ def get_or_create_contact(from_header: str, mail_from: str, alias: Alias) -> Con
                 mail_from,
             )
             contact_email = mail_from
-
+    contact_result = contact_utils.create_contact(
-    if not is_valid_email(contact_email):
+        email=contact_email,
-        LOG.w(
+        alias=alias,
-            "invalid contact email %s. Parse from %s %s",
-            contact_email,
-            from_header,
-            mail_from,
-        )
-        # either reuse a contact with empty email or create a new contact with empty email
-        contact_email = ""
-
-    contact_email = sanitize_email(contact_email, not_lower=True)
-
-    if contact_name and "\x00" in contact_name:
-        LOG.w("issue with contact name %s", contact_name)
-        contact_name = ""
-
-    contact = Contact.get_by(alias_id=alias.id, website_email=contact_email)
-    if contact:
-        if contact.name != contact_name:
-            LOG.d(
-                "Update contact %s name %s to %s",
-                contact,
-                contact.name,
-                contact_name,
-            )
-            contact.name = contact_name
-            Session.commit()
-
-        # contact created in the past does not have mail_from and from_header field
-        if not contact.mail_from and mail_from:
-            LOG.d(
-                "Set contact mail_from %s: %s to %s",
-                contact,
-                contact.mail_from,
-                mail_from,
-            )
-            contact.mail_from = mail_from
-            Session.commit()
-    else:
-        alias_id = alias.id
-        try:
-            contact_email_for_reply = (
-                contact_email if is_valid_email(contact_email) else ""
-            )
-            contact = Contact.create(
-                user_id=alias.user_id,
-                alias_id=alias_id,
-                website_email=contact_email,
         name=contact_name,
         mail_from=mail_from,
-                reply_email=generate_reply_email(contact_email_for_reply, alias),
+        allow_empty_email=True,
         automatic_created=True,
+        from_partner=False,
     )
-            if not contact_email:
+    return contact_result.contact
-                LOG.d("Create a contact with invalid email for %s", alias)
-                contact.invalid_email = True
-
-            LOG.d(
-                "create contact %s for %s, reverse alias:%s",
-                contact_email,
-                alias,
-                contact.reply_email,
-            )
-
-            Session.commit()
-        except IntegrityError:
-            LOG.info(
-                f"Contact with email {contact_email} for alias_id {alias_id} already existed, fetching from DB"
-            )
-            contact = Contact.get_by(alias_id=alias_id, website_email=contact_email)
-
-    return contact
 
 
 def get_or_create_reply_to_contact(
@@ -292,33 +233,7 @@ def get_or_create_reply_to_contact(
         )
         return None
 
-    contact = Contact.get_by(alias_id=alias.id, website_email=contact_address)
+    return contact_utils.create_contact(contact_address, alias, contact_name).contact
-    if contact:
-        return contact
-    else:
-        LOG.d(
-            "create contact %s for alias %s via reply-to header %s",
-            contact_address,
-            alias,
-            reply_to_header,
-        )
-
-        try:
-            contact = Contact.create(
-                user_id=alias.user_id,
-                alias_id=alias.id,
-                website_email=contact_address,
-                name=contact_name,
-                reply_email=generate_reply_email(contact_address, alias),
-                automatic_created=True,
-            )
-            Session.commit()
-        except IntegrityError:
-            LOG.w("Contact %s %s already exist", alias, contact_address)
-            Session.rollback()
-            contact = Contact.get_by(alias_id=alias.id, website_email=contact_address)
-
-    return contact
 
 
 def replace_header_when_forward(msg: Message, alias: Alias, header: str):
@@ -1250,23 +1165,11 @@ def handle_reply(envelope, msg: Message, rcpt_to: str) -> (bool, str):
 
     Session.commit()
 
-    # make the email comes from alias
+    recipient_name = get_alias_recipient_name(alias)
-    from_header = alias.email
+    if recipient_name.message:
-    # add alias name from alias
+        LOG.d(recipient_name.message)
-    if alias.name:
+    LOG.d("From header is %s", recipient_name.name)
-        LOG.d("Put alias name %s in from header", alias.name)
+    add_or_replace_header(msg, headers.FROM, recipient_name.name)
-        from_header = sl_formataddr((alias.name, alias.email))
-    elif alias.custom_domain:
-        # add alias name from domain
-        if alias.custom_domain.name:
-            LOG.d(
-                "Put domain default alias name %s in from header",
-                alias.custom_domain.name,
-            )
-            from_header = sl_formataddr((alias.custom_domain.name, alias.email))
-
-    LOG.d("From header is %s", from_header)
-    add_or_replace_header(msg, headers.FROM, from_header)
 
     try:
         if str(msg[headers.TO]).lower() == "undisclosed-recipients:;":

app/event_listener.py

@@ -9,7 +9,7 @@ from events.runner import Runner
 from events.event_source import DeadLetterEventSource, PostgresEventSource
 from events.event_sink import ConsoleEventSink, HttpEventSink
 
-_DEFAULT_MAX_RETRIES = 100
+_DEFAULT_MAX_RETRIES = 10
 
 
 class Mode(Enum):

app/events/event_sink.py

@@ -27,7 +27,9 @@ class HttpEventSink(EventSink):
             headers={"Content-Type": "application/x-protobuf"},
             verify=not EVENT_WEBHOOK_SKIP_VERIFY_SSL,
         )
-        newrelic.agent.record_custom_event("event_sent", {"http_code": res.status_code})
+        newrelic.agent.record_custom_event(
+            "EventSentToPartner", {"http_code": res.status_code}
+        )
         if res.status_code != 200:
             LOG.warning(
                 f"Failed to send event to webhook: {res.status_code} {res.text}"

app/events/event_source.py

@@ -72,7 +72,9 @@ class PostgresEventSource(EventSource):
                     Session.close()  # Ensure we get a new connection and we don't leave a dangling tx
 
     def __connect(self):
-        self.__connection = psycopg2.connect(self.__connection_string)
+        self.__connection = psycopg2.connect(
+            self.__connection_string, application_name="sl-event-listen"
+        )
 
         from app.db import Session
 

app/job_runner.py

@@ -3,7 +3,7 @@ Run scheduled jobs.
 Not meant for running job at precise time (+- 1h)
 """
 import time
-from typing import List
+from typing import List, Optional
 
 import arrow
 from sqlalchemy.sql.expression import or_, and_
@@ -240,18 +240,20 @@ def process_job(job: Job):
 
     elif job.name == config.JOB_DELETE_DOMAIN:
         custom_domain_id = job.payload.get("custom_domain_id")
-        custom_domain = CustomDomain.get(custom_domain_id)
+        custom_domain: Optional[CustomDomain] = CustomDomain.get(custom_domain_id)
         if not custom_domain:
             return
 
         domain_name = custom_domain.domain
         user = custom_domain.user
 
+        custom_domain_partner_id = custom_domain.partner_id
         CustomDomain.delete(custom_domain.id)
         Session.commit()
 
         LOG.d("Domain %s deleted", domain_name)
 
+        if custom_domain_partner_id is None:
             send_email(
                 user.email,
                 f"Your domain {domain_name} has been deleted",

app/migrations/versions/2024_091915_88dd7a0abf54_contact_flags_custom_domain_pending_deletion.py

@@ -0,0 +1,31 @@
+"""contact.flags and custom_domain.pending_deletion
+
+Revision ID: 88dd7a0abf54
+Revises: 2441b7ff5da9
+Create Date: 2024-09-19 15:41:20.910374
+
+"""
+import sqlalchemy_utils
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '88dd7a0abf54'
+down_revision = '2441b7ff5da9'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('contact', sa.Column('flags', sa.Integer(), server_default='0', nullable=False))
+    op.add_column('custom_domain', sa.Column('pending_deletion', sa.Boolean(), server_default='0', nullable=False))
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('custom_domain', 'pending_deletion')
+    op.drop_column('contact', 'flags')
+    # ### end Alembic commands ###

app/migrations/versions/2024_093011_62afa3a10010_custom_domain_indices.py

@@ -0,0 +1,27 @@
+"""custom domain indices
+
+Revision ID: 62afa3a10010
+Revises: 88dd7a0abf54
+Create Date: 2024-09-30 11:40:04.127791
+
+"""
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = '62afa3a10010'
+down_revision = '88dd7a0abf54'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.get_context().autocommit_block():
+        op.create_index('ix_custom_domain_pending_deletion', 'custom_domain', ['pending_deletion'], unique=False, postgresql_concurrently=True)
+        op.create_index('ix_custom_domain_user_id', 'custom_domain', ['user_id'], unique=False, postgresql_concurrently=True)
+
+
+def downgrade():
+    with op.get_context().autocommit_block():
+        op.drop_index('ix_custom_domain_user_id', table_name='custom_domain', postgresql_concurrently=True)
+        op.drop_index('ix_custom_domain_pending_deletion', table_name='custom_domain', postgresql_concurrently=True)

app/monitoring.py

@@ -94,6 +94,20 @@ def log_nb_db_connection():
     newrelic.agent.record_custom_metric("Custom/nb_db_connections", nb_connection)
 
 
+@newrelic.agent.background_task()
+def log_nb_db_connection_by_app_name():
+    # get the number of connections to the DB
+    rows = Session.execute(
+        "SELECT application_name, count(datid) FROM pg_stat_activity group by application_name"
+    )
+    for row in rows:
+        if row[0].find("sl-") == 0:
+            LOG.d("number of db connections for app %s = %s", row[0], row[1])
+            newrelic.agent.record_custom_metric(
+                f"Custom/nb_db_app_connection/{row[0]}", row[1]
+            )
+
+
 @newrelic.agent.background_task()
 def log_pending_to_process_events():
     r = Session.execute("select count(*) from sync_event WHERE taken_time IS NULL;")
@@ -125,6 +139,21 @@ def log_events_pending_dead_letter():
     )
 
 
+@newrelic.agent.background_task()
+def log_failed_events():
+    r = Session.execute(
+        """
+        SELECT COUNT(*)
+        FROM sync_event
+        WHERE retry_count >= 10;
+        """,
+    )
+    failed_events = list(r)[0][0]
+
+    LOG.d("number of failed events %s", failed_events)
+    newrelic.agent.record_custom_metric("Custom/sync_events_failed", failed_events)
+
+
 if __name__ == "__main__":
     exporter = MetricExporter(get_newrelic_license())
     while True:
@@ -132,6 +161,8 @@ if __name__ == "__main__":
         log_nb_db_connection()
         log_pending_to_process_events()
         log_events_pending_dead_letter()
+        log_failed_events()
+        log_nb_db_connection_by_app_name()
         Session.close()
 
         exporter.run()

app/poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.7.0 and should not be changed by hand.
 
 [[package]]
 name = "aiohttp"
@@ -360,35 +360,41 @@ files = [
 
 [[package]]
 name = "boto3"
-version = "1.15.9"
+version = "1.35.37"
 description = "The AWS SDK for Python"
 optional = false
-python-versions = "*"
+python-versions = ">=3.8"
 files = [
-    {file = "boto3-1.15.9-py2.py3-none-any.whl", hash = "sha256:e0a1dbc0a0e460dc6de2f4144b5015edad3ab5c17ee83c6194b1a010d815bc60"},
+    {file = "boto3-1.35.37-py3-none-any.whl", hash = "sha256:385ca77bf8ea4ab2d97f6e2435bdb29f77d9301e2f7ac796c2f465753c2adf3c"},
-    {file = "boto3-1.15.9.tar.gz", hash = "sha256:02f5f7a2b1349760b030c34f90a9cb4600bf8fe3cbc76b801d122bc4cecf3a7f"},
+    {file = "boto3-1.35.37.tar.gz", hash = "sha256:470d981583885859fed2fd1c185eeb01cc03e60272d499bafe41b12625b158c8"},
 ]
 
 [package.dependencies]
-botocore = ">=1.18.9,<1.19.0"
+botocore = ">=1.35.37,<1.36.0"
-jmespath = ">=0.7.1,<1.0.0"
+jmespath = ">=0.7.1,<2.0.0"
-s3transfer = ">=0.3.0,<0.4.0"
+s3transfer = ">=0.10.0,<0.11.0"
+
+[package.extras]
+crt = ["botocore[crt] (>=1.21.0,<2.0a0)"]
 
 [[package]]
 name = "botocore"
-version = "1.18.9"
+version = "1.35.37"
 description = "Low-level, data-driven core of boto 3."
 optional = false
-python-versions = "*"
+python-versions = ">=3.8"
 files = [
-    {file = "botocore-1.18.9-py2.py3-none-any.whl", hash = "sha256:dc3244170254cbba7dfde00b0489f830069d93dd6a9e555178d989072d7ee7c2"},
+    {file = "botocore-1.35.37-py3-none-any.whl", hash = "sha256:64f965d4ba7adb8d79ce044c3aef7356e05dd74753cf7e9115b80f477845d920"},
-    {file = "botocore-1.18.9.tar.gz", hash = "sha256:35b06b8801eb2dd7e708de35581f9c0304740645874f3af5b8b0c1648f8d6365"},
+    {file = "botocore-1.35.37.tar.gz", hash = "sha256:b2b4d29bafd95b698344f2f0577bb67064adbf1735d8a0e3c7473daa59c23ba6"},
 ]
 
 [package.dependencies]
-jmespath = ">=0.7.1,<1.0.0"
+jmespath = ">=0.7.1,<2.0.0"
 python-dateutil = ">=2.1,<3.0.0"
-urllib3 = {version = ">=1.20,<1.26", markers = "python_version != \"3.4\""}
+urllib3 = {version = ">=1.25.4,<2.2.0 || >2.2.0,<3", markers = "python_version >= \"3.10\""}
+
+[package.extras]
+crt = ["awscrt (==0.22.0)"]
 
 [[package]]
 name = "cachetools"
@@ -1372,6 +1378,7 @@ files = [
     {file = "greenlet-2.0.2-cp27-cp27m-win32.whl", hash = "sha256:6c3acb79b0bfd4fe733dff8bc62695283b57949ebcca05ae5c129eb606ff2d74"},
     {file = "greenlet-2.0.2-cp27-cp27m-win_amd64.whl", hash = "sha256:283737e0da3f08bd637b5ad058507e578dd462db259f7f6e4c5c365ba4ee9343"},
     {file = "greenlet-2.0.2-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:d27ec7509b9c18b6d73f2f5ede2622441de812e7b1a80bbd446cb0633bd3d5ae"},
+    {file = "greenlet-2.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:d967650d3f56af314b72df7089d96cda1083a7fc2da05b375d2bc48c82ab3f3c"},
     {file = "greenlet-2.0.2-cp310-cp310-macosx_11_0_x86_64.whl", hash = "sha256:30bcf80dda7f15ac77ba5af2b961bdd9dbc77fd4ac6105cee85b0d0a5fcf74df"},
     {file = "greenlet-2.0.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:26fbfce90728d82bc9e6c38ea4d038cba20b7faf8a0ca53a9c07b67318d46088"},
     {file = "greenlet-2.0.2-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9190f09060ea4debddd24665d6804b995a9c122ef5917ab26e1566dcc712ceeb"},
@@ -1380,6 +1387,7 @@ files = [
     {file = "greenlet-2.0.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:76ae285c8104046b3a7f06b42f29c7b73f77683df18c49ab5af7983994c2dd91"},
     {file = "greenlet-2.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:2d4686f195e32d36b4d7cf2d166857dbd0ee9f3d20ae349b6bf8afc8485b3645"},
     {file = "greenlet-2.0.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:c4302695ad8027363e96311df24ee28978162cdcdd2006476c43970b384a244c"},
+    {file = "greenlet-2.0.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d4606a527e30548153be1a9f155f4e283d109ffba663a15856089fb55f933e47"},
     {file = "greenlet-2.0.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c48f54ef8e05f04d6eff74b8233f6063cb1ed960243eacc474ee73a2ea8573ca"},
     {file = "greenlet-2.0.2-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a1846f1b999e78e13837c93c778dcfc3365902cfb8d1bdb7dd73ead37059f0d0"},
     {file = "greenlet-2.0.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3a06ad5312349fec0ab944664b01d26f8d1f05009566339ac6f63f56589bc1a2"},
@@ -1409,6 +1417,7 @@ files = [
     {file = "greenlet-2.0.2-cp37-cp37m-win32.whl", hash = "sha256:3f6ea9bd35eb450837a3d80e77b517ea5bc56b4647f5502cd28de13675ee12f7"},
     {file = "greenlet-2.0.2-cp37-cp37m-win_amd64.whl", hash = "sha256:7492e2b7bd7c9b9916388d9df23fa49d9b88ac0640db0a5b4ecc2b653bf451e3"},
     {file = "greenlet-2.0.2-cp38-cp38-macosx_10_15_x86_64.whl", hash = "sha256:b864ba53912b6c3ab6bcb2beb19f19edd01a6bfcbdfe1f37ddd1778abfe75a30"},
+    {file = "greenlet-2.0.2-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:1087300cf9700bbf455b1b97e24db18f2f77b55302a68272c56209d5587c12d1"},
     {file = "greenlet-2.0.2-cp38-cp38-manylinux2010_x86_64.whl", hash = "sha256:ba2956617f1c42598a308a84c6cf021a90ff3862eddafd20c3333d50f0edb45b"},
     {file = "greenlet-2.0.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fc3a569657468b6f3fb60587e48356fe512c1754ca05a564f11366ac9e306526"},
     {file = "greenlet-2.0.2-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8eab883b3b2a38cc1e050819ef06a7e6344d4a990d24d45bc6f2cf959045a45b"},
@@ -1417,6 +1426,7 @@ files = [
     {file = "greenlet-2.0.2-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:b0ef99cdbe2b682b9ccbb964743a6aca37905fda5e0452e5ee239b1654d37f2a"},
     {file = "greenlet-2.0.2-cp38-cp38-win32.whl", hash = "sha256:b80f600eddddce72320dbbc8e3784d16bd3fb7b517e82476d8da921f27d4b249"},
     {file = "greenlet-2.0.2-cp38-cp38-win_amd64.whl", hash = "sha256:4d2e11331fc0c02b6e84b0d28ece3a36e0548ee1a1ce9ddde03752d9b79bba40"},
+    {file = "greenlet-2.0.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:8512a0c38cfd4e66a858ddd1b17705587900dd760c6003998e9472b77b56d417"},
     {file = "greenlet-2.0.2-cp39-cp39-macosx_11_0_x86_64.whl", hash = "sha256:88d9ab96491d38a5ab7c56dd7a3cc37d83336ecc564e4e8816dbed12e5aaefc8"},
     {file = "greenlet-2.0.2-cp39-cp39-manylinux2010_x86_64.whl", hash = "sha256:561091a7be172ab497a3527602d467e2b3fbe75f9e783d8b8ce403fa414f71a6"},
     {file = "greenlet-2.0.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:971ce5e14dc5e73715755d0ca2975ac88cfdaefcaab078a284fea6cfabf866df"},
@@ -2891,50 +2901,72 @@ files = [
 
 [[package]]
 name = "s3transfer"
-version = "0.3.3"
+version = "0.10.3"
 description = "An Amazon S3 Transfer Manager"
 optional = false
-python-versions = "*"
+python-versions = ">=3.8"
 files = [
-    {file = "s3transfer-0.3.3-py2.py3-none-any.whl", hash = "sha256:2482b4259524933a022d59da830f51bd746db62f047d6eb213f2f8855dcb8a13"},
+    {file = "s3transfer-0.10.3-py3-none-any.whl", hash = "sha256:263ed587a5803c6c708d3ce44dc4dfedaab4c1a32e8329bab818933d79ddcf5d"},
-    {file = "s3transfer-0.3.3.tar.gz", hash = "sha256:921a37e2aefc64145e7b73d50c71bb4f26f46e4c9f414dc648c6245ff92cf7db"},
+    {file = "s3transfer-0.10.3.tar.gz", hash = "sha256:4f50ed74ab84d474ce614475e0b8d5047ff080810aac5d01ea25231cfc944b0c"},
 ]
 
 [package.dependencies]
-botocore = ">=1.12.36,<2.0a.0"
+botocore = ">=1.33.2,<2.0a.0"
+
+[package.extras]
+crt = ["botocore[crt] (>=1.33.2,<2.0a.0)"]
 
 [[package]]
 name = "sentry-sdk"
-version = "1.5.11"
+version = "2.16.0"
 description = "Python client for Sentry (https://sentry.io)"
 optional = false
-python-versions = "*"
+python-versions = ">=3.6"
 files = [
-    {file = "sentry-sdk-1.5.11.tar.gz", hash = "sha256:6c01d9d0b65935fd275adc120194737d1df317dce811e642cbf0394d0d37a007"},
+    {file = "sentry_sdk-2.16.0-py2.py3-none-any.whl", hash = "sha256:49139c31ebcd398f4f6396b18910610a0c1602f6e67083240c33019d1f6aa30c"},
-    {file = "sentry_sdk-1.5.11-py2.py3-none-any.whl", hash = "sha256:c17179183cac614e900cbd048dab03f49a48e2820182ec686c25e7ce46f8548f"},
+    {file = "sentry_sdk-2.16.0.tar.gz", hash = "sha256:90f733b32e15dfc1999e6b7aca67a38688a567329de4d6e184154a73f96c6892"},
 ]
 
 [package.dependencies]
 certifi = "*"
-urllib3 = ">=1.10.0"
+urllib3 = ">=1.26.11"
 
 [package.extras]
 aiohttp = ["aiohttp (>=3.5)"]
+anthropic = ["anthropic (>=0.16)"]
+arq = ["arq (>=0.23)"]
+asyncpg = ["asyncpg (>=0.23)"]
 beam = ["apache-beam (>=2.12)"]
 bottle = ["bottle (>=0.12.13)"]
 celery = ["celery (>=3)"]
+celery-redbeat = ["celery-redbeat (>=2)"]
 chalice = ["chalice (>=1.16.0)"]
+clickhouse-driver = ["clickhouse-driver (>=0.2.0)"]
 django = ["django (>=1.8)"]
 falcon = ["falcon (>=1.4)"]
-flask = ["blinker (>=1.1)", "flask (>=0.11)"]
+fastapi = ["fastapi (>=0.79.0)"]
+flask = ["blinker (>=1.1)", "flask (>=0.11)", "markupsafe"]
+grpcio = ["grpcio (>=1.21.1)", "protobuf (>=3.8.0)"]
+http2 = ["httpcore[http2] (==1.*)"]
 httpx = ["httpx (>=0.16.0)"]
+huey = ["huey (>=2)"]
+huggingface-hub = ["huggingface-hub (>=0.22)"]
+langchain = ["langchain (>=0.0.210)"]
+litestar = ["litestar (>=2.0.0)"]
+loguru = ["loguru (>=0.5)"]
+openai = ["openai (>=1.0.0)", "tiktoken (>=0.3.0)"]
+opentelemetry = ["opentelemetry-distro (>=0.35b0)"]
+opentelemetry-experimental = ["opentelemetry-distro"]
 pure-eval = ["asttokens", "executing", "pure-eval"]
+pymongo = ["pymongo (>=3.1)"]
 pyspark = ["pyspark (>=2.4.4)"]
 quart = ["blinker (>=1.1)", "quart (>=0.16.1)"]
 rq = ["rq (>=0.6)"]
 sanic = ["sanic (>=0.8)"]
 sqlalchemy = ["sqlalchemy (>=1.2)"]
-tornado = ["tornado (>=5)"]
+starlette = ["starlette (>=0.19.1)"]
+starlite = ["starlite (>=1.48)"]
+tornado = ["tornado (>=6)"]
 
 [[package]]
 name = "setuptools"
@@ -3295,18 +3327,18 @@ files = [
 
 [[package]]
 name = "urllib3"
-version = "1.25.10"
+version = "1.26.20"
 description = "HTTP library with thread-safe connection pooling, file post, and more."
 optional = false
-python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, <4"
+python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7"
 files = [
-    {file = "urllib3-1.25.10-py2.py3-none-any.whl", hash = "sha256:e7983572181f5e1522d9c98453462384ee92a0be7fac5f1413a1e35c56cc0461"},
+    {file = "urllib3-1.26.20-py2.py3-none-any.whl", hash = "sha256:0ed14ccfbf1c30a9072c7ca157e4319b70d65f623e91e7b32fadb2853431016e"},
-    {file = "urllib3-1.25.10.tar.gz", hash = "sha256:91056c15fa70756691db97756772bb1eb9678fa585d9184f24534b100dc60f4a"},
+    {file = "urllib3-1.26.20.tar.gz", hash = "sha256:40c2dc0c681e47eb8f90e7e27bf6ff7df2e677421fd46756da1161c39ca70d32"},
 ]
 
 [package.extras]
-brotli = ["brotlipy (>=0.6.0)"]
+brotli = ["brotli (==1.0.9)", "brotli (>=1.0.9)", "brotlicffi (>=0.8.0)", "brotlipy (>=0.6.0)"]
-secure = ["certifi", "cryptography (>=1.3.4)", "idna (>=2.0.0)", "ipaddress", "pyOpenSSL (>=0.14)"]
+secure = ["certifi", "cryptography (>=1.3.4)", "idna (>=2.0.0)", "ipaddress", "pyOpenSSL (>=0.14)", "urllib3-secure-extra"]
 socks = ["PySocks (>=1.5.6,!=1.5.7,<2.0)"]
 
 [[package]]
@@ -3704,4 +3736,4 @@ testing = ["coverage (>=5.0.3)", "zope.event", "zope.testing"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "22b9a61e9999a215aacb889b3790ee1a6840ce249aea2e3d16c6113243d5c126"
+content-hash = "314f199bd50ccbf636ce1c6c753f8c79a1f5a16aa7c1a330a2ec514a13dbad2d"

app/pyproject.toml

@@ -69,7 +69,7 @@ python-dotenv = "^0.14.0"
 ipython = "^7.31.1"
 sqlalchemy_utils = "^0.36.8"
 psycopg2-binary = "^2.9.3"
-sentry_sdk = "^1.5.11"
+sentry_sdk = "^2.16.0"
 blinker = "^1.4"
 arrow = "^0.16.0"
 Flask-WTF = "^0.14.3"

app/server.py

@@ -7,6 +7,7 @@ import arrow
 import click
 import flask_limiter
 import flask_profiler
+import newrelic.agent
 import sentry_sdk
 from coinbase_commerce.error import WebhookInvalidPayload, SignatureVerificationError
 from coinbase_commerce.webhook import Webhook
@@ -116,6 +117,7 @@ from app.oauth.base import oauth_bp
 from app.onboarding.base import onboarding_bp
 from app.phone.base import phone_bp
 from app.redis_services import initialize_redis_services
+from app.sentry_utils import sentry_before_send
 from app.utils import random_string
 
 if SENTRY_DSN:
@@ -127,6 +129,7 @@ if SENTRY_DSN:
             FlaskIntegration(),
             SqlalchemyIntegration(),
         ],
+        before_send=sentry_before_send,
     )
 
 # the app is served behind nginx which uses http and not https
@@ -299,7 +302,9 @@ def set_index_page(app):
                 res.status_code,
                 time.time() - start_time,
             )
-
+            newrelic.agent.record_custom_event(
+                "HttpResponseStatus", {"code": res.status_code}
+            )
         return res
 
 

app/templates/admin/email_search.html

@@ -40,9 +40,9 @@
         </tbody>
     </table>
 {%- endmacro %}
-{% macro list_mailboxes(mbox_count, mboxes) %}
+{% macro list_mailboxes(message, mbox_count, mboxes) %}
     <h4>
-    {{ mbox_count }} Mailboxes found.
+        {{ mbox_count }} {{ message }}.
         {% if mbox_count>10 %}Showing only the last 10.{% endif %}
     </h4>
     <table class="table">
@@ -84,7 +84,7 @@
                 Email
             </th>
             <th>
-          Verified
+                Enabled
             </th>
             <th>
                 Created At
@@ -96,7 +96,7 @@
             <tr>
                 <td>{{ alias.id }}</td>
                 <td><a href="?email={{ alias.email }}">{{ alias.email }}</a></td>
-          <td>{{ "Yes" if alias.verified else "No" }}</td>
+                <td>{{ "Yes" if alias.enabled else "No" }}</td>
                 <td>{{ alias.created_at }}</td>
             </tr>
         {% endfor %}
@@ -126,7 +126,8 @@
 {%- endmacro %}
 {% macro show_domain_deleted_alias(dom_deleted_alias) -%}
     <h4>
-    Domain Deleted Alias {{ dom_deleted_alias.email }} with ID {{ dom_deleted_alias.id }} for domain {{ dom_deleted_alias.domain.domain }}
+        Domain Deleted Alias {{ dom_deleted_alias.email }} with ID {{ dom_deleted_alias.id }} for
+        domain {{ dom_deleted_alias.domain.domain }}
     </h4>
     <table class="table">
         <thead>
@@ -168,25 +169,24 @@
         </form>
     </div>
     {% if data.no_match and email %}
-
         <div class="border border-dark border-2 mt-1 mb-2 p-3 alert alert-warning"
              role="alert">No user, alias or mailbox found for {{ email }}</div>
     {% endif %}
-  {% if data.alias %}
 
+    {% if data.alias %}
         <div class="border border-dark border-2 mt-1 mb-2 p-3">
             <h3 class="mb-3">Found Alias {{ data.alias.email }}</h3>
             {{ list_alias(1,[data.alias]) }}
+            {{ list_mailboxes("Mailboxes for alias", helper.alias_mailbox_count(data.alias), helper.alias_mailboxes(data.alias)) }}
             {{ show_user(data.alias.user) }}
-      {{ list_mailboxes(helper.mailbox_count(data.alias.user) , helper.mailbox_list(data.alias.user) ) }}
         </div>
     {% endif %}
-  {% if data.user %}
 
+    {% if data.user %}
         <div class="border border-dark border-2 mt-1 mb-2 p-3">
             <h3 class="mb-3">Found User {{ data.user.email }}</h3>
             {{ show_user(data.user) }}
-      {{ list_mailboxes(helper.mailbox_count(data.user) , helper.mailbox_list(data.user) ) }}
+            {{ list_mailboxes("Mailboxes for user", helper.mailbox_count(data.user) , helper.mailbox_list(data.user) ) }}
             {{ list_alias(helper.alias_count(data.user) ,helper.alias_list(data.user)) }}
         </div>
     {% endif %}
@@ -199,7 +199,7 @@
 
         <div class="border border-dark mt-1 mb-2 p-3">
             <h3 class="mb-3">Found Mailbox {{ mailbox.email }}</h3>
-      {{ list_mailboxes(1, [mailbox]) }}
+            {{ list_mailboxes("Mailbox found", 1, [mailbox]) }}
             {{ show_user(mailbox.user) }}
         </div>
     {% endfor %}

app/templates/dashboard/domain_detail/dns.html

@@ -91,7 +91,8 @@
           <br />
           Some domain registrars (Namecheap, CloudFlare, etc) might also use <em>@</em> for the root domain.
         </div>
-        {% for priority, email_server in EMAIL_SERVERS_WITH_PRIORITY %}
+
+        {% for record in expected_mx_records %}
 
           <div class="mb-3 p-3 dns-record">
             Record: MX
@@ -99,14 +100,15 @@
             Domain: {{ custom_domain.domain }} or
             <b>@</b>
             <br />
-            Priority: {{ priority }}
+            Priority: {{ record.priority }}
             <br />
             Target: <em data-toggle="tooltip"
      title="Click to copy"
      class="clipboard"
-     data-clipboard-text="{{ email_server }}">{{ email_server }}</em>
+     data-clipboard-text="{{ record.domain }}">{{ record.domain }}</em>
           </div>
         {% endfor %}
+
         <form method="post" action="#mx-form">
           {{ csrf_form.csrf_token }}
           <input type="hidden" name="form-name" value="check-mx">

app/tests/api/test_alias.py

@@ -536,7 +536,7 @@ def test_create_contact_route_free_users(flask_client):
     assert r.status_code == 201
 
     # End trial and disallow for new free users. Config should allow it
-    user.flags = User.FLAG_FREE_DISABLE_CREATE_ALIAS
+    user.flags = User.FLAG_DISABLE_CREATE_CONTACTS
     Session.commit()
     r = flask_client.post(
         url_for("api.create_contact_route", alias_id=alias.id),

app/tests/dashboard/test_alias_contact_manager.py

@@ -4,7 +4,7 @@ from app.models import (
     Alias,
     Contact,
 )
-from tests.utils import login
+from tests.utils import login, random_email
 
 
 def test_add_contact_success(flask_client):
@@ -13,26 +13,28 @@ def test_add_contact_success(flask_client):
 
     assert Contact.filter_by(user_id=user.id).count() == 0
 
+    email = random_email()
     # <<< Create a new contact >>>
     flask_client.post(
         url_for("dashboard.alias_contact_manager", alias_id=alias.id),
         data={
             "form-name": "create",
-            "email": "abcd@gmail.com",
+            "email": email,
         },
         follow_redirects=True,
     )
     # a new contact is added
     assert Contact.filter_by(user_id=user.id).count() == 1
     contact = Contact.filter_by(user_id=user.id).first()
-    assert contact.website_email == "abcd@gmail.com"
+    assert contact.website_email == email
 
     # <<< Create a new contact using a full email format >>>
+    email = random_email()
     flask_client.post(
         url_for("dashboard.alias_contact_manager", alias_id=alias.id),
         data={
             "form-name": "create",
-            "email": "First Last <another@gmail.com>",
+            "email": f"First Last <{email}>",
         },
         follow_redirects=True,
     )
@@ -41,7 +43,7 @@ def test_add_contact_success(flask_client):
     contact = (
         Contact.filter_by(user_id=user.id).filter(Contact.id != contact.id).first()
     )
-    assert contact.website_email == "another@gmail.com"
+    assert contact.website_email == email
     assert contact.name == "First Last"
 
     # <<< Create a new contact with invalid email address >>>

app/tests/dashboard/test_alias_transfer.py

@@ -1,38 +1,72 @@
 import app.alias_utils
+from app import config
 from app.db import Session
+from app.events.event_dispatcher import GlobalDispatcher
 from app.models import (
     Alias,
     Mailbox,
-    User,
     AliasMailbox,
 )
+from tests.events.event_test_utils import (
+    OnMemoryDispatcher,
+    _get_event_from_string,
+    _create_linked_user,
+)
 from tests.utils import login
 
+on_memory_dispatcher = OnMemoryDispatcher()
+
+
+def setup_module():
+    GlobalDispatcher.set_dispatcher(on_memory_dispatcher)
+    config.EVENT_WEBHOOK = "http://test"
+
+
+def teardown_module():
+    GlobalDispatcher.set_dispatcher(None)
+    config.EVENT_WEBHOOK = None
+
 
 def test_alias_transfer(flask_client):
-    user = login(flask_client)
+    (source_user, source_user_pu) = _create_linked_user()
-    mb = Mailbox.create(user_id=user.id, email="mb@gmail.com", commit=True)
+    source_user = login(flask_client, source_user)
+    mb = Mailbox.create(user_id=source_user.id, email="mb@gmail.com", commit=True)
 
-    alias = Alias.create_new_random(user)
+    alias = Alias.create_new_random(source_user)
     Session.commit()
 
     AliasMailbox.create(alias_id=alias.id, mailbox_id=mb.id, commit=True)
 
-    new_user = User.create(
+    (target_user, target_user_pu) = _create_linked_user()
-        email="hey@example.com",
-        password="password",
-        activated=True,
-        commit=True,
-    )
 
     Mailbox.create(
-        user_id=new_user.id, email="hey2@example.com", verified=True, commit=True
+        user_id=target_user.id, email="hey2@example.com", verified=True, commit=True
     )
 
-    app.alias_utils.transfer_alias(alias, new_user, new_user.mailboxes())
+    on_memory_dispatcher.clear()
+    app.alias_utils.transfer_alias(alias, target_user, target_user.mailboxes())
 
     # refresh from db
     alias = Alias.get(alias.id)
-    assert alias.user == new_user
+    assert alias.user == target_user
-    assert set(alias.mailboxes) == set(new_user.mailboxes())
+    assert set(alias.mailboxes) == set(target_user.mailboxes())
     assert len(alias.mailboxes) == 2
+
+    # Check events
+    assert len(on_memory_dispatcher.memory) == 2
+    # 1st delete event
+    event_data = on_memory_dispatcher.memory[0]
+    event_content = _get_event_from_string(event_data, source_user, source_user_pu)
+    assert event_content.alias_deleted is not None
+    alias_deleted = event_content.alias_deleted
+    assert alias_deleted.id == alias.id
+    assert alias_deleted.email == alias.email
+    # 2nd create event
+    event_data = on_memory_dispatcher.memory[1]
+    event_content = _get_event_from_string(event_data, target_user, target_user_pu)
+    assert event_content.alias_created is not None
+    alias_created = event_content.alias_created
+    assert alias.id == alias_created.id
+    assert alias.email == alias_created.email
+    assert alias.note or "" == alias_created.note
+    assert alias.enabled == alias_created.enabled

app/tests/events/event_test_utils.py

@@ -1,4 +1,5 @@
 from app.events.event_dispatcher import Dispatcher
+from app.events.generated import event_pb2
 from app.models import PartnerUser, User
 from app.proton.utils import get_proton_partner
 from tests.utils import create_new_user, random_token
@@ -30,3 +31,14 @@ def _create_linked_user() -> Tuple[User, PartnerUser]:
     )
 
     return user, partner_user
+
+
+def _get_event_from_string(
+    data: str, user: User, pu: PartnerUser
+) -> event_pb2.EventContent:
+    event = event_pb2.Event()
+    event.ParseFromString(data)
+    assert user.id == event.user_id
+    assert pu.external_user_id == event.external_user_id
+    assert pu.partner_id == event.partner_id
+    return event.content

app/tests/events/test_sent_events.py

@@ -1,12 +1,12 @@
 from app import config, alias_utils
 from app.db import Session
 from app.events.event_dispatcher import GlobalDispatcher
-from app.events.generated import event_pb2
+from app.models import Alias
-from app.models import Alias, User, PartnerUser
 from tests.utils import random_token
 from .event_test_utils import (
     OnMemoryDispatcher,
     _create_linked_user,
+    _get_event_from_string,
 )
 
 on_memory_dispatcher = OnMemoryDispatcher()
@@ -26,17 +26,6 @@ def setup_function(func):
     on_memory_dispatcher.clear()
 
 
-def _get_event_from_string(
-    data: str, user: User, pu: PartnerUser
-) -> event_pb2.EventContent:
-    event = event_pb2.Event()
-    event.ParseFromString(data)
-    assert user.id == event.user_id
-    assert pu.external_user_id == event.external_user_id
-    assert pu.partner_id == event.partner_id
-    return event.content
-
-
 def test_fire_event_on_alias_creation():
     (user, pu) = _create_linked_user()
     alias = Alias.create_new_random(user)

app/tests/test_alias_utils.py

@@ -4,6 +4,7 @@ from app.alias_utils import (
     delete_alias,
     check_alias_prefix,
     get_user_if_alias_would_auto_create,
+    get_alias_recipient_name,
     try_auto_create,
 )
 from app.config import ALIAS_DOMAINS
@@ -18,7 +19,8 @@ from app.models import (
     User,
     DomainDeletedAlias,
 )
-from tests.utils import create_new_user, random_domain, random_token
+from app.utils import random_string
+from tests.utils import create_new_user, random_domain, random_token, random_email
 
 
 def test_delete_alias(flask_client):
@@ -131,3 +133,91 @@ def test_auto_create_alias(flask_client):
             assert result, f"Case {test_id} - Failed address {address}"
         else:
             assert result is None, f"Case {test_id} - Failed address {address}"
+
+
+# get_alias_recipient_name
+def test_get_alias_recipient_name_no_overrides():
+    user = create_new_user()
+    alias = Alias.create(
+        user_id=user.id,
+        email=random_email(),
+        mailbox_id=user.default_mailbox_id,
+        commit=True,
+    )
+    res = get_alias_recipient_name(alias)
+    assert res.message is None
+    assert res.name == alias.email
+
+
+def test_get_alias_recipient_name_alias_name():
+    user = create_new_user()
+    alias = Alias.create(
+        user_id=user.id,
+        email=random_email(),
+        mailbox_id=user.default_mailbox_id,
+        name=random_string(),
+        commit=True,
+    )
+    res = get_alias_recipient_name(alias)
+    assert res.message is not None
+    assert res.name == f"{alias.name} <{alias.email}>"
+
+
+def test_get_alias_recipient_alias_with_name_and_custom_domain_name():
+    user = create_new_user()
+    custom_domain = CustomDomain.create(
+        user_id=user.id,
+        domain=random_domain(),
+        name=random_string(),
+        verified=True,
+    )
+    alias = Alias.create(
+        user_id=user.id,
+        email=random_email(),
+        mailbox_id=user.default_mailbox_id,
+        name=random_string(),
+        custom_domain_id=custom_domain.id,
+        commit=True,
+    )
+    res = get_alias_recipient_name(alias)
+    assert res.message is not None
+    assert res.name == f"{alias.name} <{alias.email}>"
+
+
+def test_get_alias_recipient_alias_without_name_and_custom_domain_without_name():
+    user = create_new_user()
+    custom_domain = CustomDomain.create(
+        user_id=user.id,
+        domain=random_domain(),
+        verified=True,
+    )
+    alias = Alias.create(
+        user_id=user.id,
+        email=random_email(),
+        mailbox_id=user.default_mailbox_id,
+        custom_domain_id=custom_domain.id,
+        commit=True,
+    )
+    res = get_alias_recipient_name(alias)
+    assert res.message is None
+    assert res.name == alias.email
+
+
+def test_get_alias_recipient_alias_without_name_and_custom_domain_name():
+    user = create_new_user()
+    custom_domain = CustomDomain.create(
+        user_id=user.id,
+        domain=random_domain(),
+        name=random_string(),
+        verified=True,
+    )
+    alias = Alias.create(
+        user_id=user.id,
+        email=random_email(),
+        mailbox_id=user.default_mailbox_id,
+        custom_domain_id=custom_domain.id,
+        commit=True,
+    )
+    res = get_alias_recipient_name(alias)
+    assert res.message is not None
+    assert res.name == f"{custom_domain.name} <{alias.email}>"

app/tests/test_contact_utils.py

@@ -0,0 +1,196 @@
+from typing import Optional
+
+import pytest
+
+from app import config
+from app.contact_utils import create_contact, ContactCreateError
+from app.db import Session
+from app.models import (
+    Alias,
+    Contact,
+    User,
+)
+from tests.utils import create_new_user, random_email, random_token
+
+
+def setup_module(module):
+    config.DISABLE_CREATE_CONTACTS_FOR_FREE_USERS = True
+
+
+def teardown_module(module):
+    config.DISABLE_CREATE_CONTACTS_FOR_FREE_USERS = False
+
+
+def create_provider():
+    # name auto_created from_partner
+    yield ["name", "a@b.c", True, True]
+    yield [None, None, True, True]
+    yield [None, None, False, True]
+    yield [None, None, True, False]
+    yield [None, None, False, False]
+
+
+@pytest.mark.parametrize(
+    "name, mail_from, automatic_created, from_partner", create_provider()
+)
+def test_create_contact(
+    name: Optional[str],
+    mail_from: Optional[str],
+    automatic_created: bool,
+    from_partner: bool,
+):
+    user = create_new_user()
+    alias = Alias.create_new_random(user)
+    Session.commit()
+    email = random_email()
+    contact_result = create_contact(
+        email,
+        alias,
+        name=name,
+        mail_from=mail_from,
+        automatic_created=automatic_created,
+        from_partner=from_partner,
+    )
+    assert contact_result.error is None
+    contact = contact_result.contact
+    assert contact.user_id == user.id
+    assert contact.alias_id == alias.id
+    assert contact.website_email == email
+    assert contact.name == name
+    assert contact.mail_from == mail_from
+    assert contact.automatic_created == automatic_created
+    assert not contact.invalid_email
+    expected_flags = Contact.FLAG_PARTNER_CREATED if from_partner else 0
+    assert contact.flags == expected_flags
+
+
+def test_create_contact_email_email_not_allowed():
+    user = create_new_user()
+    alias = Alias.create_new_random(user)
+    Session.commit()
+    contact_result = create_contact("", alias)
+    assert contact_result.contact is None
+    assert contact_result.error == ContactCreateError.InvalidEmail
+
+
+def test_create_contact_email_email_allowed():
+    user = create_new_user()
+    alias = Alias.create_new_random(user)
+    Session.commit()
+    contact_result = create_contact("", alias, allow_empty_email=True)
+    assert contact_result.error is None
+    assert contact_result.contact is not None
+    assert contact_result.contact.website_email == ""
+    assert contact_result.contact.invalid_email
+
+
+def test_create_contact_name_overrides_email_name():
+    user = create_new_user()
+    alias = Alias.create_new_random(user)
+    Session.commit()
+    email = random_email()
+    name = random_token()
+    contact_result = create_contact(f"superseeded <{email}>", alias, name=name)
+    assert contact_result.error is None
+    assert contact_result.contact is not None
+    assert contact_result.contact.website_email == email
+    assert contact_result.contact.name == name
+
+
+def test_create_contact_name_taken_from_email():
+    user = create_new_user()
+    alias = Alias.create_new_random(user)
+    Session.commit()
+    email = random_email()
+    name = random_token()
+    contact_result = create_contact(f"{name} <{email}>", alias)
+    assert contact_result.error is None
+    assert contact_result.contact is not None
+    assert contact_result.contact.website_email == email
+    assert contact_result.contact.name == name
+
+
+def test_create_contact_empty_name_is_none():
+    user = create_new_user()
+    alias = Alias.create_new_random(user)
+    Session.commit()
+    email = random_email()
+    contact_result = create_contact(email, alias, name="")
+    assert contact_result.error is None
+    assert contact_result.contact is not None
+    assert contact_result.contact.website_email == email
+    assert contact_result.contact.name is None
+
+
+def test_create_contact_free_user():
+    user = create_new_user()
+    user.trial_end = None
+    user.flags = 0
+    alias = Alias.create_new_random(user)
+    Session.flush()
+    # Free users without the FREE_DISABLE_CREATE_CONTACTS
+    result = create_contact(random_email(), alias)
+    assert result.error is None
+    assert result.created
+    assert result.contact is not None
+    assert not result.contact.automatic_created
+    # Free users with the flag should be able to still create automatic emails
+    user.flags = User.FLAG_DISABLE_CREATE_CONTACTS
+    Session.flush()
+    result = create_contact(random_email(), alias, automatic_created=True)
+    assert result.error is None
+    assert result.created
+    assert result.contact is not None
+    assert result.contact.automatic_created
+    # Free users with the flag cannot create non-automatic emails
+    result = create_contact(random_email(), alias)
+    assert result.error == ContactCreateError.NotAllowed
+
+
+def test_do_not_allow_invalid_email():
+    user = create_new_user()
+    alias = Alias.create_new_random(user)
+    Session.commit()
+    contact_result = create_contact("potato", alias)
+    assert contact_result.contact is None
+    assert contact_result.error == ContactCreateError.InvalidEmail
+    contact_result = create_contact("asdf\x00@gmail.com", alias)
+    assert contact_result.contact is None
+    assert contact_result.error == ContactCreateError.InvalidEmail
+
+
+def test_update_name_for_existing():
+    user = create_new_user()
+    alias = Alias.create_new_random(user)
+    Session.commit()
+    email = random_email()
+    contact_result = create_contact(email, alias)
+    assert contact_result.error is None
+    assert contact_result.created
+    assert contact_result.contact is not None
+    assert contact_result.contact.name is None
+    name = random_token()
+    contact_result = create_contact(email, alias, name=name)
+    assert contact_result.error is None
+    assert not contact_result.created
+    assert contact_result.contact is not None
+    assert contact_result.contact.name == name
+
+
+def test_update_mail_from_for_existing():
+    user = create_new_user()
+    alias = Alias.create_new_random(user)
+    Session.commit()
+    email = random_email()
+    contact_result = create_contact(email, alias)
+    assert contact_result.error is None
+    assert contact_result.created
+    assert contact_result.contact is not None
+    assert contact_result.contact is not None
+    assert contact_result.contact.mail_from is None
+    mail_from = random_email()
+    contact_result = create_contact(email, alias, mail_from=mail_from)
+    assert contact_result.error is None
+    assert not contact_result.created
+    assert contact_result.contact is not None
+    assert contact_result.contact.mail_from == mail_from

app/tests/test_custom_domain_utils.py

@@ -7,11 +7,13 @@ from app.custom_domain_utils import (
     create_custom_domain,
     is_valid_domain,
     sanitize_domain,
+    set_custom_domain_mailboxes,
     CannotUseDomainReason,
+    CannotSetCustomDomainMailboxesCause,
 )
 from app.db import Session
-from app.models import User, CustomDomain, Mailbox
+from app.models import User, CustomDomain, Mailbox, DomainMailbox
-from tests.utils import get_proton_partner
+from tests.utils import get_proton_partner, random_email
 from tests.utils import create_new_user, random_string, random_domain
 
 user: Optional[User] = None
@@ -147,3 +149,119 @@ def test_creates_custom_domain_with_partner_id():
     assert res.instance.domain == domain
     assert res.instance.user_id == user.id
     assert res.instance.partner_id == proton_partner.id
+
+
+# set_custom_domain_mailboxes
+def test_set_custom_domain_mailboxes_empty_list():
+    domain = CustomDomain.create(user_id=user.id, domain=random_domain(), commit=True)
+    res = set_custom_domain_mailboxes(user.id, domain, [])
+    assert res.success is False
+    assert res.reason == CannotSetCustomDomainMailboxesCause.NoMailboxes
+
+
+def test_set_custom_domain_mailboxes_mailbox_from_another_user():
+    other_user = create_new_user()
+    other_mailbox = Mailbox.create(
+        user_id=other_user.id, email=random_email(), verified=True
+    )
+    domain = CustomDomain.create(user_id=user.id, domain=random_domain(), commit=True)
+
+    res = set_custom_domain_mailboxes(user.id, domain, [other_mailbox.id])
+    assert res.success is False
+    assert res.reason == CannotSetCustomDomainMailboxesCause.InvalidMailbox
+
+
+def test_set_custom_domain_mailboxes_mailbox_from_current_user_and_another_user():
+    other_user = create_new_user()
+    other_mailbox = Mailbox.create(
+        user_id=other_user.id, email=random_email(), verified=True
+    )
+    domain = CustomDomain.create(user_id=user.id, domain=random_domain(), commit=True)
+
+    res = set_custom_domain_mailboxes(
+        user.id, domain, [user.default_mailbox_id, other_mailbox.id]
+    )
+    assert res.success is False
+    assert res.reason == CannotSetCustomDomainMailboxesCause.InvalidMailbox
+
+
+def test_set_custom_domain_mailboxes_success():
+    other_mailbox = Mailbox.create(user_id=user.id, email=random_email(), verified=True)
+    domain = CustomDomain.create(user_id=user.id, domain=random_domain(), commit=True)
+
+    res = set_custom_domain_mailboxes(
+        user.id, domain, [user.default_mailbox_id, other_mailbox.id]
+    )
+    assert res.success is True
+    assert res.reason is None
+
+    domain_mailboxes = DomainMailbox.filter_by(domain_id=domain.id).all()
+    assert len(domain_mailboxes) == 2
+    assert domain_mailboxes[0].domain_id == domain.id
+    assert domain_mailboxes[0].mailbox_id == user.default_mailbox_id
+    assert domain_mailboxes[1].domain_id == domain.id
+    assert domain_mailboxes[1].mailbox_id == other_mailbox.id
+
+
+def test_set_custom_domain_mailboxes_set_twice():
+    other_mailbox = Mailbox.create(user_id=user.id, email=random_email(), verified=True)
+    domain = CustomDomain.create(user_id=user.id, domain=random_domain(), commit=True)
+
+    res = set_custom_domain_mailboxes(
+        user.id, domain, [user.default_mailbox_id, other_mailbox.id]
+    )
+    assert res.success is True
+    assert res.reason is None
+
+    res = set_custom_domain_mailboxes(
+        user.id, domain, [user.default_mailbox_id, other_mailbox.id]
+    )
+    assert res.success is True
+    assert res.reason is None
+
+    domain_mailboxes = DomainMailbox.filter_by(domain_id=domain.id).all()
+    assert len(domain_mailboxes) == 2
+    assert domain_mailboxes[0].domain_id == domain.id
+    assert domain_mailboxes[0].mailbox_id == user.default_mailbox_id
+    assert domain_mailboxes[1].domain_id == domain.id
+    assert domain_mailboxes[1].mailbox_id == other_mailbox.id
+
+
+def test_set_custom_domain_mailboxes_removes_old_association():
+    domain = CustomDomain.create(user_id=user.id, domain=random_domain(), commit=True)
+
+    res = set_custom_domain_mailboxes(user.id, domain, [user.default_mailbox_id])
+    assert res.success is True
+    assert res.reason is None
+
+    other_mailbox = Mailbox.create(
+        user_id=user.id, email=random_email(), verified=True, commit=True
+    )
+    res = set_custom_domain_mailboxes(user.id, domain, [other_mailbox.id])
+    assert res.success is True
+    assert res.reason is None
+
+    domain_mailboxes = DomainMailbox.filter_by(domain_id=domain.id).all()
+    assert len(domain_mailboxes) == 1
+    assert domain_mailboxes[0].domain_id == domain.id
+    assert domain_mailboxes[0].mailbox_id == other_mailbox.id
+
+
+def test_set_custom_domain_mailboxes_with_unverified_mailbox():
+    domain = CustomDomain.create(user_id=user.id, domain=random_domain())
+    verified_mailbox = Mailbox.create(
+        user_id=user.id,
+        email=random_email(),
+        verified=True,
+    )
+    unverified_mailbox = Mailbox.create(
+        user_id=user.id,
+        email=random_email(),
+        verified=False,
+    )
+
+    res = set_custom_domain_mailboxes(
+        user.id, domain, [verified_mailbox.id, unverified_mailbox.id]
+    )
+    assert res.success is False
+    assert res.reason is CannotSetCustomDomainMailboxesCause.InvalidMailbox

app/tests/test_custom_domain_validation.py

@@ -5,7 +5,7 @@ from app.constants import DMARC_RECORD
 from app.custom_domain_validation import CustomDomainValidation
 from app.db import Session
 from app.models import CustomDomain, User
-from app.dns_utils import InMemoryDNSClient
+from app.dns_utils import InMemoryDNSClient, MxRecord
 from app.proton.utils import get_proton_partner
 from app.utils import random_string
 from tests.utils import create_new_user, random_domain
@@ -58,6 +58,123 @@ def test_custom_domain_validation_get_dkim_records_for_partner():
     assert records["dkim._domainkey"] == f"dkim._domainkey.{dkim_domain}"
 
 
+# get_expected_mx_records
+def test_custom_domain_validation_get_expected_mx_records_regular_domain():
+    domain = random_domain()
+    custom_domain = create_custom_domain(domain)
+
+    partner_id = get_proton_partner().id
+
+    dkim_domain = random_domain()
+    validator = CustomDomainValidation(
+        domain, partner_domains={partner_id: dkim_domain}
+    )
+    records = validator.get_expected_mx_records(custom_domain)
+    # As the domain is not a partner_domain,default records should be used even if
+    # there is a config for the partner
+    assert len(records) == len(config.EMAIL_SERVERS_WITH_PRIORITY)
+    for i in range(len(config.EMAIL_SERVERS_WITH_PRIORITY)):
+        config_record = config.EMAIL_SERVERS_WITH_PRIORITY[i]
+        assert records[i].priority == config_record[0]
+        assert records[i].domain == config_record[1]
+
+
+def test_custom_domain_validation_get_expected_mx_records_domain_from_partner():
+    domain = random_domain()
+    custom_domain = create_custom_domain(domain)
+
+    partner_id = get_proton_partner().id
+    custom_domain.partner_id = partner_id
+    Session.commit()
+
+    dkim_domain = random_domain()
+    validator = CustomDomainValidation(dkim_domain)
+    records = validator.get_expected_mx_records(custom_domain)
+    # As the domain is a partner_domain but there is no custom config for partner, default records
+    # should be used
+    assert len(records) == len(config.EMAIL_SERVERS_WITH_PRIORITY)
+    for i in range(len(config.EMAIL_SERVERS_WITH_PRIORITY)):
+        config_record = config.EMAIL_SERVERS_WITH_PRIORITY[i]
+        assert records[i].priority == config_record[0]
+        assert records[i].domain == config_record[1]
+
+
+def test_custom_domain_validation_get_expected_mx_records_domain_from_partner_with_custom_config():
+    domain = random_domain()
+    custom_domain = create_custom_domain(domain)
+
+    partner_id = get_proton_partner().id
+    custom_domain.partner_id = partner_id
+    Session.commit()
+
+    dkim_domain = random_domain()
+    expected_mx_domain = random_domain()
+    validator = CustomDomainValidation(
+        dkim_domain, partner_domains={partner_id: expected_mx_domain}
+    )
+    records = validator.get_expected_mx_records(custom_domain)
+    # As the domain is a partner_domain and there is a custom config for partner, partner records
+    # should be used
+    assert len(records) == 2
+
+    assert records[0].priority == 10
+    assert records[0].domain == f"mx1.{expected_mx_domain}."
+    assert records[1].priority == 20
+    assert records[1].domain == f"mx2.{expected_mx_domain}."
+
+
+# get_expected_spf_records
+def test_custom_domain_validation_get_expected_spf_record_regular_domain():
+    domain = random_domain()
+    custom_domain = create_custom_domain(domain)
+
+    partner_id = get_proton_partner().id
+
+    dkim_domain = random_domain()
+    validator = CustomDomainValidation(
+        domain, partner_domains={partner_id: dkim_domain}
+    )
+    record = validator.get_expected_spf_record(custom_domain)
+    # As the domain is not a partner_domain, default records should be used even if
+    # there is a config for the partner
+    assert record == f"v=spf1 include:{config.EMAIL_DOMAIN} ~all"
+
+
+def test_custom_domain_validation_get_expected_spf_record_domain_from_partner():
+    domain = random_domain()
+    custom_domain = create_custom_domain(domain)
+
+    partner_id = get_proton_partner().id
+    custom_domain.partner_id = partner_id
+    Session.commit()
+
+    dkim_domain = random_domain()
+    validator = CustomDomainValidation(dkim_domain)
+    record = validator.get_expected_spf_record(custom_domain)
+    # As the domain is a partner_domain but there is no custom config for partner, default records
+    # should be used
+    assert record == f"v=spf1 include:{config.EMAIL_DOMAIN} ~all"
+
+
+def test_custom_domain_validation_get_expected_spf_record_domain_from_partner_with_custom_config():
+    domain = random_domain()
+    custom_domain = create_custom_domain(domain)
+
+    partner_id = get_proton_partner().id
+    custom_domain.partner_id = partner_id
+    Session.commit()
+
+    dkim_domain = random_domain()
+    expected_mx_domain = random_domain()
+    validator = CustomDomainValidation(
+        dkim_domain, partner_domains={partner_id: expected_mx_domain}
+    )
+    record = validator.get_expected_spf_record(custom_domain)
+    # As the domain is a partner_domain and there is a custom config for partner, partner records
+    # should be used
+    assert record == f"v=spf1 include:{expected_mx_domain} ~all"
+
+
 # validate_dkim_records
 def test_custom_domain_validation_validate_dkim_records_empty_records_failure():
     dns_client = InMemoryDNSClient()
@@ -253,7 +370,7 @@ def test_custom_domain_validation_validate_mx_records_wrong_records_failure():
 
     wrong_record_1 = random_string()
     wrong_record_2 = random_string()
-    wrong_records = [(10, wrong_record_1), (20, wrong_record_2)]
+    wrong_records = [MxRecord(10, wrong_record_1), MxRecord(20, wrong_record_2)]
     dns_client.set_mx_records(domain.domain, wrong_records)
     res = validator.validate_mx_records(domain)
 
@@ -270,7 +387,7 @@ def test_custom_domain_validation_validate_mx_records_success():
 
     domain = create_custom_domain(random_domain())
 
-    dns_client.set_mx_records(domain.domain, config.EMAIL_SERVERS_WITH_PRIORITY)
+    dns_client.set_mx_records(domain.domain, validator.get_expected_mx_records(domain))
     res = validator.validate_mx_records(domain)
 
     assert res.success is True
@@ -328,6 +445,58 @@ def test_custom_domain_validation_validate_spf_records_success():
     assert db_domain.spf_verified is True
 
 
+def test_custom_domain_validation_validate_spf_records_partner_domain_success():
+    dns_client = InMemoryDNSClient()
+    proton_partner_id = get_proton_partner().id
+
+    expected_domain = random_domain()
+    validator = CustomDomainValidation(
+        dkim_domain=random_domain(),
+        dns_client=dns_client,
+        partner_domains={proton_partner_id: expected_domain},
+    )
+
+    domain = create_custom_domain(random_domain())
+    domain.partner_id = proton_partner_id
+    Session.commit()
+
+    dns_client.set_txt_record(domain.domain, [f"v=spf1 include:{expected_domain}"])
+    res = validator.validate_spf_records(domain)
+
+    assert res.success is True
+    assert len(res.errors) == 0
+
+    db_domain = CustomDomain.get_by(id=domain.id)
+    assert db_domain.spf_verified is True
+
+
+def test_custom_domain_validation_validate_spf_cleans_verification_record():
+    dns_client = InMemoryDNSClient()
+    proton_partner_id = get_proton_partner().id
+
+    expected_domain = random_domain()
+    validator = CustomDomainValidation(
+        dkim_domain=random_domain(),
+        dns_client=dns_client,
+        partner_domains={proton_partner_id: expected_domain},
+    )
+
+    domain = create_custom_domain(random_domain())
+    domain.partner_id = proton_partner_id
+    Session.commit()
+
+    wrong_record = random_string()
+    dns_client.set_txt_record(
+        hostname=domain.domain,
+        txt_list=[wrong_record, validator.get_ownership_verification_record(domain)],
+    )
+    res = validator.validate_spf_records(domain)
+
+    assert res.success is False
+    assert len(res.errors) == 1
+    assert res.errors[0] == wrong_record
+
+
 # validate_dmarc_records
 def test_custom_domain_validation_validate_dmarc_records_empty_failure():
     dns_client = InMemoryDNSClient()

app/tests/test_dns_utils.py

@@ -3,6 +3,7 @@ from app.dns_utils import (
     get_network_dns_client,
     is_mx_equivalent,
     InMemoryDNSClient,
+    MxRecord,
 )
 
 from tests.utils import random_domain
@@ -17,8 +18,8 @@ def test_get_mx_domains():
     assert len(r) > 0
 
     for x in r:
-        assert x[0] > 0
+        assert x.priority > 0
-        assert x[1]
+        assert x.domain
 
 
 def test_get_spf_domain():
@@ -33,20 +34,32 @@ def test_get_txt_record():
 
 def test_is_mx_equivalent():
     assert is_mx_equivalent([], [])
-    assert is_mx_equivalent([(1, "domain")], [(1, "domain")])
     assert is_mx_equivalent(
-        [(10, "domain1"), (20, "domain2")], [(10, "domain1"), (20, "domain2")]
+        mx_domains=[MxRecord(1, "domain")], ref_mx_domains=[MxRecord(1, "domain")]
     )
     assert is_mx_equivalent(
-        [(5, "domain1"), (10, "domain2")], [(10, "domain1"), (20, "domain2")]
+        mx_domains=[MxRecord(10, "domain1"), MxRecord(20, "domain2")],
+        ref_mx_domains=[MxRecord(10, "domain1"), MxRecord(20, "domain2")],
     )
     assert is_mx_equivalent(
-        [(5, "domain1"), (10, "domain2"), (20, "domain3")],
+        mx_domains=[MxRecord(5, "domain1"), MxRecord(10, "domain2")],
-        [(10, "domain1"), (20, "domain2")],
+        ref_mx_domains=[MxRecord(10, "domain1"), MxRecord(20, "domain2")],
+    )
+    assert is_mx_equivalent(
+        mx_domains=[
+            MxRecord(5, "domain1"),
+            MxRecord(10, "domain2"),
+            MxRecord(20, "domain3"),
+        ],
+        ref_mx_domains=[MxRecord(10, "domain1"), MxRecord(20, "domain2")],
     )
     assert not is_mx_equivalent(
-        [(5, "domain1"), (10, "domain2")],
+        mx_domains=[MxRecord(5, "domain1"), MxRecord(10, "domain2")],
-        [(10, "domain1"), (20, "domain2"), (20, "domain3")],
+        ref_mx_domains=[
+            MxRecord(10, "domain1"),
+            MxRecord(20, "domain2"),
+            MxRecord(20, "domain3"),
+        ],
     )
 
 

app/tests/test_email_utils.py

@@ -90,12 +90,19 @@ def test_can_be_used_as_personal_email(flask_client):
     assert not email_can_be_used_as_mailbox("ab@sl.local")
     assert not email_can_be_used_as_mailbox("hey@d1.test")
 
-    # custom domain
+    # custom domain as SL domain
     domain = random_domain()
     user = create_new_user()
-    CustomDomain.create(user_id=user.id, domain=domain, verified=True, commit=True)
+    domain_obj = CustomDomain.create(
+        user_id=user.id, domain=domain, verified=True, is_sl_subdomain=True, flush=True
+    )
     assert not email_can_be_used_as_mailbox(f"hey@{domain}")
 
+    # custom domain is NOT SL domain
+    domain_obj.is_sl_subdomain = False
+    Session.flush()
+    assert email_can_be_used_as_mailbox(f"hey@{domain}")
+
     # disposable domain
     disposable_domain = random_domain()
     InvalidMailboxDomain.create(domain=disposable_domain, commit=True)