Updating the files based on the findings on the PR:

moved the function is_domain_blocked into its own utils
refactoring the function can_domain_be_used

Author: tozo

app/blocked_domain_utils.py

@@ -0,0 +1,6 @@
+from app.models import BlockedDomain
+
+
+def is_domain_blocked(user_id: int, domain: str) -> bool:
+    """checks whether the provided domain is blocked for a given user"""
+    return BlockedDomain.filter_by(user_id=user_id, domain=domain).first() is not None

app/custom_domain_utils.py

@@ -1,9 +1,9 @@
-import arrow
 import re
-
 from dataclasses import dataclass
 from enum import Enum
-from typing import List, Optional, Type
+from typing import List, Optional
+
+import arrow
 
 from app.constants import JobType
 from app.db import Session
@@ -17,7 +17,7 @@
     Job,
     DomainMailbox,
     Alias,
-    ModelMixin,
+    BlockedDomain,
 )
 from app.user_audit_log_utils import emit_user_audit_log, UserAuditLogAction
 
@@ -98,15 +98,11 @@ def sanitize_domain(domain: str) -> str:
     return new_domain
 
 
-def can_domain_be_used(
-    user: User, domain: str, model_type: Type[ModelMixin]
-) -> Optional[CannotUseDomainReason]:
+def can_domain_be_used(user: User, domain: str) -> Optional[CannotUseDomainReason]:
     if not is_valid_domain(domain):
         return CannotUseDomainReason.InvalidDomain
     elif SLDomain.get_by(domain=domain):
         return CannotUseDomainReason.BuiltinDomain
-    elif model_type.get_by(domain=domain):
-        return CannotUseDomainReason.DomainAlreadyUsed
     elif get_email_domain_part(user.email) == domain:
         return CannotUseDomainReason.DomainPartOfUserEmail
     elif Mailbox.filter(
@@ -117,6 +113,34 @@ def can_domain_be_used(
         return None
 
 
+def can_custom_domain_be_used(
+    user: User, domain: str
+) -> Optional[CannotUseDomainReason]:
+    reason = can_domain_be_used(user, domain)
+
+    if reason is not None:
+        return reason
+
+    if CustomDomain.get_by(domain=domain):
+        return CannotUseDomainReason.DomainAlreadyUsed
+    else:
+        return None
+
+
+def can_blocked_domain_be_used(
+    user: User, domain: str
+) -> Optional[CannotUseDomainReason]:
+    reason = can_domain_be_used(user, domain)
+
+    if reason is not None:
+        return reason
+
+    if BlockedDomain.get_by(domain=domain):
+        return CannotUseDomainReason.DomainAlreadyUsed
+    else:
+        return None
+
+
 def create_custom_domain(
     user: User, domain: str, partner_id: Optional[int] = None
 ) -> CreateCustomDomainResult:
@@ -127,7 +151,7 @@ def create_custom_domain(
         )
 
     new_domain = sanitize_domain(domain)
-    domain_forbidden_cause = can_domain_be_used(user, new_domain, CustomDomain)
+    domain_forbidden_cause = can_custom_domain_be_used(user, new_domain)
     if domain_forbidden_cause:
         return CreateCustomDomainResult(
             message=domain_forbidden_cause.message(new_domain), message_category="error"

app/dashboard/views/setting.py

@@ -20,7 +20,7 @@
     ALIAS_RANDOM_SUFFIX_LENGTH,
     CONNECT_WITH_PROTON,
 )
-from app.custom_domain_utils import sanitize_domain, can_domain_be_used
+from app.custom_domain_utils import sanitize_domain, can_blocked_domain_be_used
 from app.dashboard.base import dashboard_bp
 from app.db import Session
 from app.errors import ProtonPartnerNotSetUp
@@ -301,8 +301,8 @@ def setting():
         elif request.form.get("form-name") == "blocked-domains-add":
             domain = request.form.get("domain-name")
             new_domain = sanitize_domain(domain)
-            domain_forbidden_cause = can_domain_be_used(
-                current_user, new_domain, BlockedDomain
+            domain_forbidden_cause = can_blocked_domain_be_used(
+                current_user, new_domain
             )
 
             if domain_forbidden_cause:

app/models.py

@@ -1237,15 +1237,6 @@ def has_used_alias_from_partner(self) -> bool:
             > 0
         )
 
-    def is_domain_blocked(self, domain: str) -> bool:
-        """checks whether the provided domain is blocked for a given user"""
-        domain_names = []
-
-        for blocked_domain in BlockedDomain.filter_by(user_id=self.id):
-            domain_names.append(blocked_domain.domain)
-
-        return any(domain in domain_name for domain_name in domain_names)
-
     def __repr__(self):
         return f"<User {self.id} {self.name} {self.email}>"
 

email_handler.py

@@ -59,6 +59,7 @@
     change_alias_status,
     get_alias_recipient_name,
 )
+from app.blocked_domain_utils import is_domain_blocked
 from app.config import (
     EMAIL_DOMAIN,
     URL,
@@ -584,7 +585,7 @@ def handle_forward(envelope, msg: Message, rcpt_to: str) -> List[Tuple[bool, str
             return [(True, status.E209)]
 
     mail_from_domain = get_email_domain_part(mail_from)
-    if user.is_domain_blocked(mail_from_domain):
+    if is_domain_blocked(user.id, mail_from_domain):
         # by default return 2** instead of 5** to allow user to receive emails again when domain is unblocked
         res_status = status.E200
         if user.block_behaviour == BlockBehaviourEnum.return_5xx:

…ns/versions/2025_020620_83676846e104_.py …ns/versions/2025_040814_c18048c40ed9_.pymigrations/versions/2025_020620_83676846e104_.py renamed to migrations/versions/2025_040814_c18048c40ed9_.py migrations/versions/2025_020620_83676846e104_.py renamed to migrations/versions/2025_040814_c18048c40ed9_.py

@@ -1,8 +1,8 @@
 """empty message
 
-Revision ID: 83676846e104
-Revises: 20e7d3ca289a
-Create Date: 2025-02-06 20:35:26.475120
+Revision ID: c18048c40ed9
+Revises: 07855f9f39b1
+Create Date: 2025-04-08 14:43:10.097244
 
 """
 import sqlalchemy_utils
@@ -11,8 +11,8 @@
 
 
 # revision identifiers, used by Alembic.
-revision = '83676846e104'
-down_revision = '20e7d3ca289a'
+revision = 'c18048c40ed9'
+down_revision = '07855f9f39b1'
 branch_labels = None
 depends_on = None
 

tests/test_blocked_domain_utils.py

@@ -0,0 +1,31 @@
+from app.blocked_domain_utils import is_domain_blocked
+from app.db import Session
+from app.models import BlockedDomain
+from tests.utils import create_new_user
+
+
+def setup():
+    user = create_new_user()
+    BlockedDomain.create(
+        user_id=user.id,
+        domain="example.com",
+        flush=True,
+    )
+    Session.flush()
+
+
+def teardown_module():
+    Session.query(BlockedDomain).delete()
+
+
+def test_domain_blocked_for_user():
+    user = create_new_user()
+    BlockedDomain.create(
+        user_id=user.id,
+        domain="example.com",
+        flush=True,
+    )
+    Session.flush()
+
+    assert is_domain_blocked(user.id, "example.com")
+    assert not is_domain_blocked(user.id, "some-other-example.com")

tests/test_custom_domain_utils.py

@@ -12,6 +12,8 @@
     CannotUseDomainReason,
     CannotSetCustomDomainMailboxesCause,
     count_custom_domain_aliases,
+    can_blocked_domain_be_used,
+    can_custom_domain_be_used,
 )
 from app.db import Session
 from app.models import User, CustomDomain, Mailbox, DomainMailbox, BlockedDomain, Alias
@@ -52,46 +54,110 @@ def test_is_valid_domain():
 # can_domain_be_used
 def test_can_domain_be_used():
     domain = f"{random_string(10)}.com"
-    res = can_domain_be_used(user, domain, CustomDomain)
+    res = can_domain_be_used(user, domain)
     assert res is None
 
 
-def test_can_domain_be_used_existing_domain():
+def test_can_domain_be_used_sl_domain():
+    domain = ALIAS_DOMAINS[0]
+    res = can_domain_be_used(user, domain)
+    assert res is CannotUseDomainReason.BuiltinDomain
+
+
+def test_can_domain_be_used_domain_of_user_email():
+    domain = user.email.split("@")[1]
+    res = can_domain_be_used(user, domain)
+    assert res is CannotUseDomainReason.DomainPartOfUserEmail
+
+
+def test_can_domain_be_used_domain_of_existing_mailbox():
+    domain = random_domain()
+    Mailbox.create(user_id=user.id, email=f"email@{domain}", verified=True, commit=True)
+    res = can_domain_be_used(user, domain)
+    assert res is CannotUseDomainReason.DomainUserInMailbox
+
+
+def test_can_domain_be_used_invalid_domain():
+    domain = f"{random_string(10)}@lol.com"
+    res = can_domain_be_used(user, domain)
+    assert res is CannotUseDomainReason.InvalidDomain
+
+
+# can_custom_domain_be_used
+def test_can_custom_domain_be_used():
+    domain = f"{random_string(10)}.com"
+    res = can_custom_domain_be_used(user, domain)
+    assert res is None
+
+
+def test_can_custom_domain_be_used_sl_domain():
+    domain = ALIAS_DOMAINS[0]
+    res = can_custom_domain_be_used(user, domain)
+    assert res is CannotUseDomainReason.BuiltinDomain
+
+
+def test_can_custom_domain_be_used_domain_of_user_email():
+    domain = user.email.split("@")[1]
+    res = can_custom_domain_be_used(user, domain)
+    assert res is CannotUseDomainReason.DomainPartOfUserEmail
+
+
+def test_can_custom_domain_be_used_domain_of_existing_mailbox():
+    domain = random_domain()
+    Mailbox.create(user_id=user.id, email=f"email@{domain}", verified=True, commit=True)
+    res = can_custom_domain_be_used(user, domain)
+    assert res is CannotUseDomainReason.DomainUserInMailbox
+
+
+def test_can_custom_domain_be_used_invalid_domain():
+    domain = f"{random_string(10)}@lol.com"
+    res = can_custom_domain_be_used(user, domain)
+    assert res is CannotUseDomainReason.InvalidDomain
+
+
+def test_can_custom_domain_be_used_existing_domain():
     domain = random_domain()
     CustomDomain.create(user_id=user.id, domain=domain, commit=True)
-    res = can_domain_be_used(user, domain, CustomDomain)
+    res = can_custom_domain_be_used(user, domain)
     assert res is CannotUseDomainReason.DomainAlreadyUsed
 
 
-def test_can_domain_be_used_sl_domain():
+# can_blocked_domain_be_used
+def test_can_blocked_domain_be_used():
+    domain = f"{random_string(10)}.com"
+    res = can_blocked_domain_be_used(user, domain)
+    assert res is None
+
+
+def test_can_blocked_domain_be_used_sl_domain():
     domain = ALIAS_DOMAINS[0]
-    res = can_domain_be_used(user, domain, CustomDomain)
+    res = can_blocked_domain_be_used(user, domain)
     assert res is CannotUseDomainReason.BuiltinDomain
 
 
-def test_can_domain_be_used_domain_of_user_email():
+def test_can_blocked_domain_be_used_domain_of_user_email():
     domain = user.email.split("@")[1]
-    res = can_domain_be_used(user, domain, CustomDomain)
+    res = can_blocked_domain_be_used(user, domain)
     assert res is CannotUseDomainReason.DomainPartOfUserEmail
 
 
-def test_can_domain_be_used_domain_of_existing_mailbox():
+def test_can_blocked_domain_be_used_domain_of_existing_mailbox():
     domain = random_domain()
     Mailbox.create(user_id=user.id, email=f"email@{domain}", verified=True, commit=True)
-    res = can_domain_be_used(user, domain, CustomDomain)
+    res = can_blocked_domain_be_used(user, domain)
     assert res is CannotUseDomainReason.DomainUserInMailbox
 
 
-def test_can_domain_be_used_invalid_domain():
+def test_can_blocked_domain_be_used_invalid_domain():
     domain = f"{random_string(10)}@lol.com"
-    res = can_domain_be_used(user, domain, CustomDomain)
+    res = can_blocked_domain_be_used(user, domain)
     assert res is CannotUseDomainReason.InvalidDomain
 
 
 def test_can_blocked_domain_be_used_existing_domain():
     domain = random_domain()
     BlockedDomain.create(user_id=user.id, domain=domain, commit=True)
-    res = can_domain_be_used(user, domain, BlockedDomain)
+    res = can_blocked_domain_be_used(user, domain)
     assert res is CannotUseDomainReason.DomainAlreadyUsed
 
 

tests/test_domains.py

@@ -1,13 +1,12 @@
 from app.db import Session
-from app.models import SLDomain, PartnerUser, AliasOptions, BlockedDomain
+from app.models import SLDomain, PartnerUser, AliasOptions
 from app.proton.proton_partner import get_proton_partner
 from init_app import add_sl_domains
 from tests.utils import create_new_user, random_token
 
 
 def setup_module():
     Session.query(SLDomain).delete()
-    Session.query(BlockedDomain).delete()
     SLDomain.create(
         domain="hidden", premium_only=False, flush=True, order=5, hidden=True
     )
@@ -34,7 +33,6 @@ def setup_module():
 
 def teardown_module():
     Session.query(SLDomain).delete()
-    Session.query(BlockedDomain).delete()
     add_sl_domains()
 
 
@@ -229,16 +227,3 @@ def test_get_free_partner_and_premium_partner():
     assert [d.domain for d in domains] == user.available_sl_domains(
         alias_options=options
     )
-
-
-def test_domain_blocked_for_user():
-    user = create_new_user()
-    BlockedDomain.create(
-        user_id=user.id,
-        domain="example.com",
-        flush=True,
-    )
-    Session.flush()
-
-    assert user.is_domain_blocked("example.com")
-    assert not user.is_domain_blocked("some-other-example.com")