From adff5103596f1d70746b386fcf8c6d5b52dd63af Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Thu, 15 Oct 2020 16:45:08 +0200
Subject: [PATCH] use PublicDomain instead if ALIAS_DOMAINS

---
 app/dashboard/views/mailbox.py |  4 +---
 app/email_utils.py             | 10 ++--------
 app/models.py                  |  7 +++++--
 shell.py                       | 15 ++++++---------
 tests/conftest.py              |  4 ++++
 tests/models/__init__.py       |  0
 tests/models/test_user.py      | 13 +++++++++++++
 7 files changed, 31 insertions(+), 22 deletions(-)
 create mode 100644 tests/models/__init__.py
 create mode 100644 tests/models/test_user.py

diff --git a/app/dashboard/views/mailbox.py b/app/dashboard/views/mailbox.py
index 0edd2839..07e6749b 100644
--- a/app/dashboard/views/mailbox.py
+++ b/app/dashboard/views/mailbox.py
@@ -5,7 +5,7 @@ from itsdangerous import Signer
 from wtforms import validators
 from wtforms.fields.html5 import EmailField
 
-from app.config import EMAIL_DOMAIN, ALIAS_DOMAINS, MAILBOX_SECRET, URL
+from app.config import MAILBOX_SECRET, URL
 from app.dashboard.base import dashboard_bp
 from app.email_utils import (
     email_can_be_used_as_mailbox,
@@ -113,8 +113,6 @@ def mailbox_route():
         "dashboard/mailbox.html",
         mailboxes=mailboxes,
         new_mailbox_form=new_mailbox_form,
-        EMAIL_DOMAIN=EMAIL_DOMAIN,
-        ALIAS_DOMAINS=ALIAS_DOMAINS,
     )
 
 
diff --git a/app/email_utils.py b/app/email_utils.py
index adbcf566..5b1b97a5 100644
--- a/app/email_utils.py
+++ b/app/email_utils.py
@@ -383,10 +383,7 @@ def can_create_directory_for_address(address: str) -> bool:
 def is_valid_alias_address_domain(address) -> bool:
     """Return whether an address domain might a domain handled by SimpleLogin"""
     domain = get_email_domain_part(address)
-    if domain in ALIAS_DOMAINS:
-        return True
-
-    if domain in PREMIUM_ALIAS_DOMAINS:
+    if PublicDomain.get_by(domain=domain):
         return True
 
     if CustomDomain.get_by(domain=domain, verified=True):
@@ -407,10 +404,7 @@ def email_can_be_used_as_mailbox(email: str) -> bool:
     if not domain:
         return False
 
-    if domain in ALIAS_DOMAINS:
-        return False
-
-    if domain in PREMIUM_ALIAS_DOMAINS:
+    if PublicDomain.get_by(domain=domain):
         return False
 
     from app.models import CustomDomain
diff --git a/app/models.py b/app/models.py
index ffb6a38f..322d7984 100644
--- a/app/models.py
+++ b/app/models.py
@@ -565,9 +565,12 @@ class User(db.Model, ModelMixin, UserMixin):
         - SimpleLogin public domains, available for all users (ALIAS_DOMAIN)
         - SimpleLogin premium domains, only available for Premium accounts (PREMIUM_ALIAS_DOMAIN)
         """
-        domains = ALIAS_DOMAINS
         if self.is_premium():
-            domains += PREMIUM_ALIAS_DOMAINS
+            query = PublicDomain.query.all()
+        else:
+            query = PublicDomain.filter_by(premium_only=False)
+
+        domains = [public_domain.domain for public_domain in query]
 
         return domains
 
diff --git a/shell.py b/shell.py
index c72044cb..4294a9d8 100644
--- a/shell.py
+++ b/shell.py
@@ -100,18 +100,15 @@ def migrate_domain_trash():
     """Move aliases from global trash to domain trash if applicable"""
     for deleted_alias in DeletedAlias.query.all():
         alias_domain = get_email_domain_part(deleted_alias.email)
-        if (
-            alias_domain not in ALIAS_DOMAINS
-            and alias_domain not in PREMIUM_ALIAS_DOMAINS
-        ):
-            domain = CustomDomain.get_by(domain=alias_domain)
-            if domain:
-                LOG.d("move %s to domain %s trash", deleted_alias, domain)
+        if not PublicDomain.get_by(domain=alias_domain):
+            custom_domain = CustomDomain.get_by(domain=alias_domain)
+            if custom_domain:
+                LOG.d("move %s to domain %s trash", deleted_alias, custom_domain)
                 db.session.add(
                     DomainDeletedAlias(
-                        user_id=domain.user_id,
+                        user_id=custom_domain.user_id,
                         email=deleted_alias.email,
-                        domain_id=domain.id,
+                        domain_id=custom_domain.id,
                         created_at=deleted_alias.created_at,
                     )
                 )
diff --git a/tests/conftest.py b/tests/conftest.py
index 65e078e0..d3cec82e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,5 +1,6 @@
 import os
 
+
 os.environ["CONFIG"] = os.path.abspath(
     os.path.join(os.path.dirname(os.path.dirname(__file__)), "tests/test.env")
 )
@@ -13,6 +14,7 @@ import pytest
 
 from app.extensions import db
 from server import create_app
+from init_app import add_public_domains
 
 
 @pytest.fixture
@@ -27,6 +29,7 @@ def flask_app():
 
     with app.app_context():
         db.create_all()
+        add_public_domains()
 
     yield app
 
@@ -45,4 +48,5 @@ def flask_client():
 
     with app.app_context():
         db.create_all()
+        add_public_domains()
         yield client
diff --git a/tests/models/__init__.py b/tests/models/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/models/test_user.py b/tests/models/test_user.py
new file mode 100644
index 00000000..dc87b519
--- /dev/null
+++ b/tests/models/test_user.py
@@ -0,0 +1,13 @@
+from app.models import User
+
+
+def test_available_sl_domains(flask_client):
+    user = User.create(
+        email="a@b.c",
+        password="password",
+        name="Test User",
+        activated=True,
+        commit=True,
+    )
+
+    assert set(user.available_sl_domains()) == {"d1.test", "d2.test", "sl.local"}