refactor: improve login logic and enhance error handling in login.py

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,101 +27,118 @@
 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):
     @staticmethod
     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}"),