diff --git a/apps/common/utils/random.py b/apps/common/utils/random.py index fae1d80bd..f75e7001b 100644 --- a/apps/common/utils/random.py +++ b/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