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)