diff --git a/app/dashboard/views/custom_alias.py b/app/dashboard/views/custom_alias.py index 0ee94649..a0ef9d6f 100644 --- a/app/dashboard/views/custom_alias.py +++ b/app/dashboard/views/custom_alias.py @@ -52,20 +52,27 @@ def get_available_suffixes(user: User) -> [SuffixInfo]: # put custom domain first # for each user domain, generate both the domain and a random suffix version for custom_domain in user_custom_domains: + if custom_domain.random_prefix_generation: + suffix = "." + random_word() + "@" + custom_domain.domain + suffix_info = SuffixInfo(True, suffix, signer.sign(suffix).decode(), False) + if user.default_alias_custom_domain_id == custom_domain.id: + suffixes.insert(0, suffix_info) + else: + suffixes.append(suffix_info) + suffix = "@" + custom_domain.domain suffix_info = SuffixInfo(True, suffix, signer.sign(suffix).decode(), False) # put the default domain to top - if user.default_alias_custom_domain_id == custom_domain.id: + # only if random_prefix_generation isn't enabled + if ( + user.default_alias_custom_domain_id == custom_domain.id + and not custom_domain.random_prefix_generation + ): suffixes.insert(0, suffix_info) else: suffixes.append(suffix_info) - if custom_domain.random_prefix_generation: - suffix = "." + random_word() + "@" + custom_domain.domain - suffix_info = SuffixInfo(True, suffix, signer.sign(suffix).decode(), False) - suffixes.append(suffix_info) - # then SimpleLogin domain for sl_domain in user.get_sl_domains(): suffix = ( diff --git a/tests/dashboard/test_custom_alias.py b/tests/dashboard/test_custom_alias.py index 062c892e..001c6e91 100644 --- a/tests/dashboard/test_custom_alias.py +++ b/tests/dashboard/test_custom_alias.py @@ -135,6 +135,27 @@ def test_available_suffixes_default_domain(flask_client): assert first_suffix.suffix == "@test.com" +def test_available_suffixes_random_prefix_generation(flask_client): + user = login(flask_client) + + CustomDomain.create(user_id=user.id, domain="test.com", verified=True, commit=True) + cd2 = CustomDomain.create( + user_id=user.id, domain="test2.com", verified=True, commit=True + ) + + user.default_alias_custom_domain_id = cd2.id + + # first suffix is test2.com + first_suffix = get_available_suffixes(user)[0] + assert first_suffix.suffix == "@test2.com" + + cd2.random_prefix_generation = True + # e.g. .meo@test2.com + first_suffix = get_available_suffixes(user)[0] + assert first_suffix.suffix.endswith("@test2.com") + assert first_suffix.suffix.startswith(".") + + def test_add_already_existed_alias(flask_client): user = login(flask_client) db.session.commit()