From 3fc250018d6db529bbb8cd5aa066dea8482642a4 Mon Sep 17 00:00:00 2001
From: devStorm <59678453+developStorm@users.noreply.github.com>
Date: Thu, 13 May 2021 16:13:19 -0700
Subject: [PATCH] basic implementation of random suffix

---
 app/config.py                       |  2 ++
 app/dashboard/views/custom_alias.py | 19 ++++++++++++++++---
 app/models.py                       |  5 +++++
 app/utils.py                        |  5 ++++-
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/app/config.py b/app/config.py
index 4dc11beb..bd501953 100644
--- a/app/config.py
+++ b/app/config.py
@@ -374,3 +374,5 @@ except Exception:
 ALIAS_LIMIT = os.environ.get("ALIAS_LIMIT") or "100/day;50/hour;5/minute"
 
 ENABLE_SPAM_ASSASSIN = "ENABLE_SPAM_ASSASSIN" in os.environ
+
+ALIAS_RAND_SUFFIX_LENGTH = int(os.environ.get("ALIAS_RAND_SUFFIX_LENGTH", 5))
\ No newline at end of file
diff --git a/app/dashboard/views/custom_alias.py b/app/dashboard/views/custom_alias.py
index b17ac447..5a080a19 100644
--- a/app/dashboard/views/custom_alias.py
+++ b/app/dashboard/views/custom_alias.py
@@ -10,6 +10,7 @@ from app.config import (
     DISABLE_ALIAS_SUFFIX,
     CUSTOM_ALIAS_SECRET,
     ALIAS_LIMIT,
+    ALIAS_RAND_SUFFIX_LENGTH,
 )
 from app.dashboard.base import dashboard_bp
 from app.extensions import db, limiter
@@ -23,7 +24,7 @@ from app.models import (
     AliasMailbox,
     DomainDeletedAlias,
 )
-from app.utils import random_word, word_exist
+from app.utils import random_word, word_exist, random_string
 
 signer = TimestampSigner(CUSTOM_ALIAS_SECRET)
 
@@ -54,7 +55,7 @@ def get_available_suffixes(user: User) -> [SuffixInfo]:
     # 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 = "." + get_suffix(user) + "@" + 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)
@@ -77,7 +78,7 @@ def get_available_suffixes(user: User) -> [SuffixInfo]:
     # then SimpleLogin domain
     for sl_domain in user.get_sl_domains():
         suffix = (
-            ("" if DISABLE_ALIAS_SUFFIX else "." + random_word())
+            ("" if DISABLE_ALIAS_SUFFIX else "." + get_suffix(user))
             + "@"
             + sl_domain.domain
         )
@@ -249,6 +250,18 @@ def custom_alias():
         mailboxes=mailboxes,
     )
 
+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:
+        return random_string(ALIAS_RAND_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"""
diff --git a/app/models.py b/app/models.py
index e71ddca5..6ef39f94 100644
--- a/app/models.py
+++ b/app/models.py
@@ -287,6 +287,11 @@ class User(db.Model, ModelMixin, UserMixin):
         db.Boolean, default=False, nullable=False, server_default="0"
     )
 
+    # whether to use random string or random word as suffix
+    random_alias_suffix = db.Column(
+        db.Boolean, default=True, nullable=False, server_default="1"
+    )
+
     @classmethod
     def create(cls, email, name="", password=None, **kwargs):
         user: User = super(User, cls).create(email=email, name=name, **kwargs)
diff --git a/app/utils.py b/app/utils.py
index 6b9f8312..ee91dd88 100644
--- a/app/utils.py
+++ b/app/utils.py
@@ -27,9 +27,12 @@ def random_words():
     return "_".join([random.choice(_words) for i in range(nb_words)])
 
 
-def random_string(length=10):
+def random_string(length=10, include_digits=False):
     """Generate a random string of fixed length """
     letters = string.ascii_lowercase
+    if include_digits:
+        letters += string.digits
+        
     return "".join(random.choice(letters) for _ in range(length))