From ed10ffee54085eb31f250ec55af083a666d70676 Mon Sep 17 00:00:00 2001 From: wxg0103 <727495428@qq.com> Date: Tue, 11 Aug 2026 16:17:44 +0800 Subject: [PATCH] refactor: improve login logic and enhance error handling in login.py --- apps/users/serializers/login.py | 183 +++++++++++++++----------------- 1 file changed, 87 insertions(+), 96 deletions(-) diff --git a/apps/users/serializers/login.py b/apps/users/serializers/login.py index b4a0620d049..cc1ac3959da 100644 --- a/apps/users/serializers/login.py +++ b/apps/users/serializers/login.py @@ -11,8 +11,8 @@ import json from application.models import ApplicationAccessToken -from common.auth.common import FileToken from captcha.image import ImageCaptcha +from common.auth.common import FileToken from common.constants.authentication_type import AuthenticationType from common.constants.cache_version import Cache_Version from common.database_model_manage.database_model_manage import DatabaseModelManage @@ -27,56 +27,44 @@ from rest_framework import serializers from users.models import User +system_version, system_get_key = Cache_Version.SYSTEM.value + class LoginRequest(serializers.Serializer): - username = serializers.CharField(required=True, max_length=64, help_text=_("Username"), label=_("Username")) + username = serializers.CharField(required=True, max_length=64, label=_("Username")) password = serializers.CharField(required=True, max_length=128, label=_("Password")) - captcha = serializers.CharField( - required=False, max_length=64, label=_("captcha"), allow_null=True, allow_blank=True - ) - encryptedData = serializers.CharField(required=False, label=_("encryptedData"), allow_null=True, allow_blank=True) - - -system_version, system_get_key = Cache_Version.SYSTEM.value + captcha = serializers.CharField(required=False, max_length=64, allow_null=True, allow_blank=True) + encryptedData = serializers.CharField(required=False, allow_null=True, allow_blank=True) class LoginResponse(serializers.Serializer): - """ - 登录响应对象 - """ - token = serializers.CharField(required=True, label=_("token")) -def record_login_fail(username: str, expire: int = 600): +def _incr_fail_count(cache_key: str, expire: int) -> int: + """原子递增失败计数，key 不存在时初始化并返回当前值""" + try: + return cache.incr(cache_key, 1, version=system_version) + except ValueError: + cache.set(cache_key, 1, timeout=expire, version=system_version) + return 1 + + +def record_login_fail(username: str, expire: int = 600) -> int: """记录登录失败次数（原子）返回当前失败计数""" if not username: return 0 - fail_key = system_get_key(f"system_{username}") - try: - fail_count = cache.incr(fail_key, 1, version=system_version) - except ValueError: - # key 不存在，初始化并设置过期 - cache.set(fail_key, 1, timeout=expire, version=system_version) - fail_count = 1 - return fail_count + return _incr_fail_count(system_get_key(f"system_{username}"), expire) -def record_login_fail_lock(username: str, expire: int = 10): +def record_login_fail_lock(username: str, expire: int = 10) -> int: """ 使用 cache.incr 保证原子递增，并在不存在时初始化计数器并返回当前值。 这里的计数器用于判断是否应当进入"锁定"分支，避免依赖非原子 get -> set 的组合。 """ if not username: return 0 - fail_key = system_get_key(f"system_{username}_lock_count") - try: - fail_count = cache.incr(fail_key, 1, version=system_version) - except ValueError: - # key 不存在，初始化并设置过期（分钟转秒） - cache.set(fail_key, 1, timeout=expire * 60, version=system_version) - fail_count = 1 - return fail_count + return _incr_fail_count(system_get_key(f"system_{username}_lock_count"), expire * 60) class LoginSerializer(serializers.Serializer): @@ -84,44 +72,73 @@ class LoginSerializer(serializers.Serializer): def get_auth_setting(): """获取认证设置""" auth_setting_model = DatabaseModelManage.get_model("auth_setting") - auth_setting = {} - if auth_setting_model: - setting_obj = auth_setting_model.objects.filter(param_key="auth_setting").first() - if setting_obj: - try: - auth_setting = json.loads(setting_obj.param_value) or {} - except Exception: - auth_setting = {} - return auth_setting + if not auth_setting_model: + return {} + setting_obj = auth_setting_model.objects.filter(param_key="auth_setting").first() + if not setting_obj: + return {} + try: + return json.loads(setting_obj.param_value) or {} + except Exception: + return {} @staticmethod - def login(instance): - # 解密数据 + def _decrypt_request_data(instance: dict) -> dict: + """解密并合并 encryptedData，返回更新后的请求数据""" username = instance.get("username", "") encrypted_data = instance.get("encryptedData", "") - - if encrypted_data: - try: - decrypted_raw = decrypt(encrypted_data) - # decrypt 可能返回非 JSON 字符串，防护解析异常 - decrypted_data = json.loads(decrypted_raw) if decrypted_raw else {} - if isinstance(decrypted_data, dict): - instance.update(decrypted_data) - except Exception as e: - maxkb_logger.exception("Failed to decrypt/parse encryptedData for user %s: %s", username, e) - raise AppApiException(500, _("Invalid encrypted data")) + if not encrypted_data: + return instance try: - request_serializer = LoginRequest(data=instance) - request_serializer.is_valid(raise_exception=True) - except serializers.ValidationError: - raise + decrypted_raw = decrypt(encrypted_data) + # decrypt 可能返回非 JSON 字符串，防护解析异常 + decrypted_data = json.loads(decrypted_raw) if decrypted_raw else {} + if isinstance(decrypted_data, dict): + instance.update(decrypted_data) except Exception as e: - raise AppApiException(500, str(e)) + maxkb_logger.exception("Failed to decrypt/parse encryptedData for user %s: %s", username, e) + raise AppApiException(500, _("Invalid encrypted data")) + return instance + + @staticmethod + def _authenticate(username: str, password: str) -> User | None: + """校验用户名密码，失败记录计数并抛异常""" + user = User.objects.filter(username=username).first() + if not user or not password_verify(password, user.password): + return None + # Transparently upgrade legacy MD5 hash to PBKDF2 + if needs_password_upgrade(user.password): + user.password = password_encrypt(password) + user.save(update_fields=["password"]) + return user + + @staticmethod + def _issue_token(user: User) -> str: + """签发登录 token 并写入缓存""" + token = signing.dumps( + { + "username": user.username, + "id": str(user.id), + "email": user.email, + "type": AuthenticationType.SYSTEM_USER.value, + } + ) + version, get_key = Cache_Version.TOKEN.value + cache.set(get_key(token), user, timeout=CONFIG.get_session_timeout(), version=version) + return token + + @staticmethod + def login(instance): + # 解密数据 + instance = LoginSerializer._decrypt_request_data(instance) + + request_serializer = LoginRequest(data=instance) + request_serializer.is_valid(raise_exception=True) validated_data = request_serializer.validated_data - username = validated_data.get("username", "") - password = validated_data.get("password", "") + username = validated_data["username"] + password = validated_data["password"] captcha = validated_data.get("captcha", "") # 获取认证配置 @@ -134,47 +151,28 @@ def login(instance): license_validator = DatabaseModelManage.get_model("license_is_valid") is_license_valid = bool(license_validator()) if license_validator else False - if is_license_valid: + if is_license_valid and LoginSerializer._is_account_locked(username, failed_attempts): # 检查账户是否被锁定 - if LoginSerializer._is_account_locked(username, failed_attempts): - raise AppApiException( - 1005, _("This account has been locked for %s minutes, please try again later") % lock_time - ) - - # 验证验证码 + raise AppApiException( + 1005, _("This account has been locked for %s minutes, please try again later") % lock_time + ) if LoginSerializer._need_captcha(username, max_attempts): + # 验证验证码 LoginSerializer._validate_captcha(username, captcha) # 验证用户凭据：先按用户名查找，再用 password_verify 验证密码 - user = User.objects.filter(username=username).first() - - if not user or not password_verify(password, user.password): + user = LoginSerializer._authenticate(username, password) + if user is None: LoginSerializer._handle_failed_login(username, is_license_valid, failed_attempts, lock_time) raise AppApiException(500, _("The username or password is incorrect")) - # Transparently upgrade legacy MD5 hash to PBKDF2 - if needs_password_upgrade(user.password): - user.password = password_encrypt(password) - user.save(update_fields=["password"]) - if not user.is_active: raise AppApiException(1005, _("The user has been disabled, please contact the administrator!")) # 清除失败计数并生成令牌 cache.delete(system_get_key(f"system_{username}"), version=system_version) cache.delete(system_get_key(f"system_{username}_lock"), version=system_version) - token = signing.dumps( - { - "username": user.username, - "id": str(user.id), - "email": user.email, - "type": AuthenticationType.SYSTEM_USER.value, - } - ) - - version, get_key = Cache_Version.TOKEN.value - timeout = CONFIG.get_session_timeout() - cache.set(get_key(token), user, timeout=timeout, version=version) + token = LoginSerializer._issue_token(user) return {"token": token}, FileToken(str(user.id), AuthenticationType.SYSTEM_USER.value).to_token() @@ -268,10 +266,6 @@ def _handle_failed_login(username: str, is_license_valid: bool, failed_attempts: class CaptchaResponse(serializers.Serializer): - """ - 登录响应对象 - """ - captcha = serializers.CharField(required=True, label=_("captcha")) @@ -299,13 +293,10 @@ def chat_generate(username: str, type: str = "chat", access_token: str = ""): @staticmethod def _generate_captcha_if_needed(username: str, type: str, need_captcha: bool): - """ - 提取的公共验证码生成方法 - """ + """提取的公共验证码生成方法""" if need_captcha: chars = get_random_chars() - image = ImageCaptcha() - data = image.generate(chars) + data = ImageCaptcha().generate(chars) captcha = base64.b64encode(data.getbuffer()) cache.set( Cache_Version.CAPTCHA.get_key(captcha=f"{type}_{username}"),