From 6806708820bb1019d2adced7e7eb5b5e27dd974d Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Mon, 9 Dec 2024 15:18:14 +0800
Subject: [PATCH] perf: Random secret string

---
 apps/common/utils/random.py | 52 +++++++++++++++++++++++++------------
 1 file changed, 36 insertions(+), 16 deletions(-)

diff --git a/apps/common/utils/random.py b/apps/common/utils/random.py
index 7978267b9..2fcebdf94 100644
--- a/apps/common/utils/random.py
+++ b/apps/common/utils/random.py
@@ -18,9 +18,8 @@ def random_ip():
     return socket.inet_ntoa(struct.pack('>I', random.randint(1, 0xffffffff)))
 
 
-def random_replace_char(s, chars, length):
+def random_replace_char(seq, chars, length):
     using_index = set()
-    seq = list(s)
 
     while length > 0:
         index = secrets.randbelow(len(seq) - 1)
@@ -29,7 +28,7 @@ def random_replace_char(s, chars, length):
         seq[index] = secrets.choice(chars)
         using_index.add(index)
         length -= 1
-    return ''.join(seq)
+    return seq
 
 
 def remove_exclude_char(s, exclude_chars):
@@ -47,19 +46,40 @@ def random_string(
     if length < 4:
         raise ValueError('The length of the string must be greater than 3')
 
-    chars_map = (
-        (lower, string.ascii_lowercase),
-        (upper, string.ascii_uppercase),
-        (digit, string.digits),
-    )
-    chars = ''.join([i[1] for i in chars_map if i[0]])
-    chars = remove_exclude_char(chars, exclude_chars)
-    texts = list(secrets.choice(chars) for __ in range(length))
-    texts = ''.join(texts)
+    char_list = []
+    if lower:
+
+        lower_chars = remove_exclude_char(string.ascii_lowercase, exclude_chars)
+        if not lower_chars:
+            raise ValueError('After excluding characters, no lowercase letters are available.')
+        char_list.append(lower_chars)
+
+    if upper:
+        upper_chars = remove_exclude_char(string.ascii_uppercase, exclude_chars)
+        if not upper_chars:
+            raise ValueError('After excluding characters, no uppercase letters are available.')
+        char_list.append(upper_chars)
+
+    if digit:
+        digit_chars = remove_exclude_char(string.digits, exclude_chars)
+        if not digit_chars:
+            raise ValueError('After excluding characters, no digits are available.')
+        char_list.append(digit_chars)
+
+    secret_chars = [secrets.choice(chars) for chars in char_list]
+
+    all_chars = ''.join(char_list)
+
+    remaining_length = length - len(secret_chars)
+    seq = [secrets.choice(all_chars) for _ in range(remaining_length)]
 
-    # 控制一下特殊字符的数量, 别随机出来太多
     if special_char:
-        symbols = remove_exclude_char(symbols, exclude_chars)
+        special_chars = remove_exclude_char(symbols, exclude_chars)
+        if not special_chars:
+            raise ValueError('After excluding characters, no special characters are available.')
         symbol_num = length // 16 + 1
-        texts = random_replace_char(texts, symbols, symbol_num)
-    return texts
+        seq = random_replace_char(seq, symbols, symbol_num)
+    secret_chars += seq
+
+    secrets.SystemRandom().shuffle(secret_chars)
+    return ''.join(secret_chars)