From f7bef3941a9c21e060aec72c1c70fc03fe2deeae Mon Sep 17 00:00:00 2001
From: devStorm <59678453+developStorm@users.noreply.github.com>
Date: Wed, 26 May 2021 22:29:00 -0700
Subject: [PATCH] replace random_word with get_suffix(user)

---
 app/api/views/alias_options.py      |  8 ++++----
 app/dashboard/views/custom_alias.py | 18 +-----------------
 app/models.py                       |  3 ++-
 app/utils.py                        | 17 ++++++++++++++++-
 4 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/app/api/views/alias_options.py b/app/api/views/alias_options.py
index d8a96599..b5c26f0f 100644
--- a/app/api/views/alias_options.py
+++ b/app/api/views/alias_options.py
@@ -9,7 +9,7 @@ from app.dashboard.views.custom_alias import (
 from app.extensions import db
 from app.log import LOG
 from app.models import AliasUsedOn, Alias, User
-from app.utils import convert_to_id, random_word
+from app.utils import convert_to_id, random_word, get_suffix
 
 
 @api_bp.route("/alias/options")
@@ -76,7 +76,7 @@ def options():
         if DISABLE_ALIAS_SUFFIX:
             ret["custom"]["suffixes"].append(f"@{domain}")
         else:
-            ret["custom"]["suffixes"].append(f".{random_word()}@{domain}")
+            ret["custom"]["suffixes"].append(f".{get_suffix(user)}@{domain}")
 
     for custom_domain in user.verified_custom_domains():
         ret["custom"]["suffixes"].append("@" + custom_domain.domain)
@@ -156,7 +156,7 @@ def options_v2():
         if DISABLE_ALIAS_SUFFIX:
             ret["suffixes"].append(f"@{domain}")
         else:
-            ret["suffixes"].append(f".{random_word()}@{domain}")
+            ret["suffixes"].append(f".{get_suffix(user)}@{domain}")
 
     for custom_domain in user.verified_custom_domains():
         ret["suffixes"].append("@" + custom_domain.domain)
@@ -232,7 +232,7 @@ def options_v3():
         if DISABLE_ALIAS_SUFFIX:
             ret["suffixes"].append(f"@{domain}")
         else:
-            ret["suffixes"].append(f".{random_word()}@{domain}")
+            ret["suffixes"].append(f".{get_suffix(user)}@{domain}")
 
     for custom_domain in user.verified_custom_domains():
         ret["suffixes"].append("@" + custom_domain.domain)
diff --git a/app/dashboard/views/custom_alias.py b/app/dashboard/views/custom_alias.py
index 24deee07..e6da8f4b 100644
--- a/app/dashboard/views/custom_alias.py
+++ b/app/dashboard/views/custom_alias.py
@@ -10,7 +10,6 @@ from app.config import (
     DISABLE_ALIAS_SUFFIX,
     CUSTOM_ALIAS_SECRET,
     ALIAS_LIMIT,
-    ALIAS_RANDOM_SUFFIX_LENGTH,
 )
 from app.dashboard.base import dashboard_bp
 from app.extensions import db, limiter
@@ -23,9 +22,8 @@ from app.models import (
     User,
     AliasMailbox,
     DomainDeletedAlias,
-    AliasSuffixEnum,
 )
-from app.utils import random_word, random_string
+from app.utils import get_suffix
 
 signer = TimestampSigner(CUSTOM_ALIAS_SECRET)
 
@@ -252,20 +250,6 @@ def custom_alias():
     )
 
 
-def get_suffix(user: User) -> str:
-    """Get random suffix for an alias based on user's preference.
-
-    Args:
-        user (User): the user who is trying to create an alias
-
-    Returns:
-        str: the random suffix generated
-    """
-    if user.random_alias_suffix == AliasSuffixEnum.random_string.value:
-        return random_string(ALIAS_RANDOM_SUFFIX_LENGTH, include_digits=True)
-    return random_word()
-
-
 def verify_prefix_suffix(user: User, alias_prefix, alias_suffix) -> bool:
     """verify if user could create an alias with the given prefix and suffix"""
     if not alias_prefix or not alias_suffix:  # should be caught on frontend
diff --git a/app/models.py b/app/models.py
index d5df3cbc..04caf197 100644
--- a/app/models.py
+++ b/app/models.py
@@ -36,6 +36,7 @@ from app.utils import (
     random_words,
     random_word,
     sanitize_email,
+    get_suffix,
 )
 
 
@@ -1124,7 +1125,7 @@ class Alias(db.Model, ModelMixin):
 
         # find the right suffix - avoid infinite loop by running this at max 1000 times
         for i in range(1000):
-            suffix = random_word()
+            suffix = get_suffix(user)
             email = f"{prefix}.{suffix}@{FIRST_ALIAS_DOMAIN}"
 
             if not cls.get_by(email=email) and not DeletedAlias.get_by(email=email):
diff --git a/app/utils.py b/app/utils.py
index d6bcd8d1..429328f3 100644
--- a/app/utils.py
+++ b/app/utils.py
@@ -4,8 +4,9 @@ import urllib.parse
 
 from unidecode import unidecode
 
-from .config import WORDS_FILE_PATH
+from .config import WORDS_FILE_PATH, ALIAS_RANDOM_SUFFIX_LENGTH
 from .log import LOG
+from .models import User, AliasSuffixEnum
 
 with open(WORDS_FILE_PATH) as f:
     LOG.d("load words file: %s", WORDS_FILE_PATH)
@@ -16,6 +17,20 @@ def random_word():
     return random.choice(_words)
 
 
+def get_suffix(user: User) -> str:
+    """Get random suffix for an alias based on user's preference.
+
+    Args:
+        user (User): the user who is trying to create an alias
+
+    Returns:
+        str: the random suffix generated
+    """
+    if user.random_alias_suffix == AliasSuffixEnum.random_string.value:
+        return random_string(ALIAS_RANDOM_SUFFIX_LENGTH, include_digits=True)
+    return random_word()
+
+
 def word_exist(word):
     return word in _words