diff --git a/app/dashboard/views/custom_alias.py b/app/dashboard/views/custom_alias.py index 4523f2fb..b7fda5da 100644 --- a/app/dashboard/views/custom_alias.py +++ b/app/dashboard/views/custom_alias.py @@ -11,12 +11,32 @@ from app.dashboard.base import dashboard_bp from app.email_utils import email_belongs_to_alias_domains, get_email_domain_part from app.extensions import db from app.log import LOG -from app.models import Alias, CustomDomain, DeletedAlias, Mailbox +from app.models import Alias, CustomDomain, DeletedAlias, Mailbox, User from app.utils import convert_to_id, random_word, word_exist signer = TimestampSigner(CUSTOM_ALIAS_SECRET) +def available_suffixes(user: User) -> [bool, str, str]: + """Return (is_custom_domain, alias-suffix, time-signed alias-suffix)""" + user_custom_domains = [cd.domain for cd in user.verified_custom_domains()] + + # List of (is_custom_domain, alias-suffix, time-signed alias-suffix) + suffixes = [] + + # put custom domain first + for alias_domain in user_custom_domains: + suffix = "@" + alias_domain + suffixes.append((True, suffix, signer.sign(suffix).decode())) + + # then default domain + for domain in ALIAS_DOMAINS: + suffix = ("" if DISABLE_ALIAS_SUFFIX else "." + random_word()) + "@" + domain + suffixes.append((False, suffix, signer.sign(suffix).decode())) + + return suffixes + + @dashboard_bp.route("/custom_alias", methods=["GET", "POST"]) @login_required def custom_alias(): @@ -32,17 +52,7 @@ def custom_alias(): user_custom_domains = [cd.domain for cd in current_user.verified_custom_domains()] # List of (is_custom_domain, alias-suffix, time-signed alias-suffix) - suffixes = [] - - # put custom domain first - for alias_domain in user_custom_domains: - suffix = "@" + alias_domain - suffixes.append((True, suffix, signer.sign(suffix).decode())) - - # then default domain - for domain in ALIAS_DOMAINS: - suffix = ("" if DISABLE_ALIAS_SUFFIX else "." + random_word()) + "@" + domain - suffixes.append((False, suffix, signer.sign(suffix).decode())) + suffixes = available_suffixes(current_user) mailboxes = [mb.email for mb in current_user.mailboxes()] @@ -105,7 +115,12 @@ def custom_alias(): else: flash("something went wrong", "warning") - return render_template("dashboard/custom_alias.html", **locals()) + return render_template( + "dashboard/custom_alias.html", + user_custom_domains=user_custom_domains, + suffixes=suffixes, + mailboxes=mailboxes, + ) def verify_prefix_suffix(user, alias_prefix, alias_suffix) -> bool: diff --git a/tests/dashboard/test_custom_alias.py b/tests/dashboard/test_custom_alias.py index 7245d385..3e189c3b 100644 --- a/tests/dashboard/test_custom_alias.py +++ b/tests/dashboard/test_custom_alias.py @@ -1,7 +1,11 @@ from flask import url_for from app.config import EMAIL_DOMAIN -from app.dashboard.views.custom_alias import signer, verify_prefix_suffix +from app.dashboard.views.custom_alias import ( + signer, + verify_prefix_suffix, + available_suffixes, +) from app.extensions import db from app.models import Mailbox, CustomDomain from app.utils import random_word @@ -53,3 +57,18 @@ def test_verify_prefix_suffix(flask_client): word = random_word() suffix = f".{word}@{EMAIL_DOMAIN}" assert verify_prefix_suffix(user, "prefix", suffix) + + +def test_available_suffixes(flask_client): + user = login(flask_client) + db.session.commit() + + CustomDomain.create(user_id=user.id, domain="test.com", verified=True) + + assert len(available_suffixes(user)) > 0 + + # first suffix is custom domain + first_suffix = available_suffixes(user)[0] + assert first_suffix[0] + assert first_suffix[1] == "@test.com" + assert first_suffix[2].startswith("@test.com")