fix: 修改 random_string 方法，支持只生成随机数字

apps/common/utils/random.py

@@ -18,25 +18,35 @@ def random_ip():
     return socket.inet_ntoa(struct.pack('>I', random.randint(1, 0xffffffff)))
 
 
-def random_string(length, lower=True, upper=True, digit=True, special_char=False):
-    chars = string.ascii_letters
-    if digit:
-        chars += string.digits
+def random_string(length: int, lower=True, upper=True, digit=True, special_char=False):
+    args_names = ['lower', 'upper', 'digit', 'special_char']
+    args_values = [lower, upper, digit, special_char]
+    args_string = [string.ascii_lowercase, string.ascii_uppercase, string.digits, string_punctuation]
+    args_string_map = dict(zip(args_names, args_string))
+    kwargs = dict(zip(args_names, args_values))
+    kwargs_keys = list(kwargs.keys())
+    kwargs_values = list(kwargs.values())
+    args_true_count = len([i for i in kwargs_values if i])
+    assert any(kwargs_values), f'Parameters {kwargs_keys} must have at least one `True`'
+    assert length >= args_true_count, f'Expected length >= {args_true_count}, bug got {length}'
+
+    can_startswith_special_char = args_true_count == 1 and special_char
+
+    chars = ''.join([args_string_map[k] for k, v in kwargs.items() if v])
 
     while True:
         password = list(random.choice(chars) for i in range(length))
-        if upper and not any(c.upper() for c in password):
-            continue
-        if lower and not any(c.lower() for c in password):
-            continue
-        if digit and not any(c.isdigit() for c in password):
-            continue
-        break
-
-    if special_char:
-        spc = random.choice(string_punctuation)
-        i = random.choice(range(1, len(password)))
-        password[i] = spc
+        for k, v in kwargs.items():
+            if v and not (set(password) & set(args_string_map[k])):
+                # 没有包含指定的字符, retry
+                break
+        else:
+            if not can_startswith_special_char and password[0] in args_string_map['special_char']:
+                # 首位不能为特殊字符, retry
+                continue
+            else:
+                # 满足要求终止 while 循环
+                break
 
     password = ''.join(password)
     return password