feat: MFA 登录次数限制

apps/authentication/api/mfa.py

@@ -29,7 +29,7 @@ class MFAChallengeApi(AuthMixin, CreateAPIView):
             if not valid:
                 self.request.session['auth_mfa'] = ''
                 raise errors.MFAFailedError(
-                    username=user.username, request=self.request
+                    username=user.username, request=self.request, ip=self.get_request_ip()
                 )
             else:
                 self.request.session['auth_mfa'] = '1'

apps/authentication/errors.py

@@ -6,9 +6,7 @@ from django.conf import settings
 
 from common.exceptions import JMSException
 from .signals import post_auth_failed
-from users.utils import (
-    increase_login_failed_count, get_login_failed_count
-)
+from users.utils import LoginBlockUtil, MFABlockUtils
 
 reason_password_failed = 'password_failed'
 reason_password_decrypt_failed = 'password_decrypt_failed'
@@ -52,7 +50,15 @@ block_login_msg = _(
     "The account has been locked "
     "(please contact admin to unlock it or try again after {} minutes)"
 )
-mfa_failed_msg = _("MFA code invalid, or ntp sync server time")
+block_mfa_msg = _(
+    "The account has been locked "
+    "(please contact admin to unlock it or try again after {} minutes)"
+)
+mfa_failed_msg = _(
+    "MFA code invalid, or ntp sync server time, "
+    "You can also try {times_try} times "
+    "(The account will be temporarily locked for {block_time} minutes)"
+)
 
 mfa_required_msg = _("MFA required")
 mfa_unset_msg = _("MFA not set, please set it first")
@@ -80,7 +86,7 @@ class AuthFailedNeedBlockMixin:
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        increase_login_failed_count(self.username, self.ip)
+        LoginBlockUtil(self.username, self.ip).incr_failed_count()
 
 
 class AuthFailedError(Exception):
@@ -107,13 +113,12 @@ class AuthFailedError(Exception):
 class CredentialError(AuthFailedNeedLogMixin, AuthFailedNeedBlockMixin, AuthFailedError):
     def __init__(self, error, username, ip, request):
         super().__init__(error=error, username=username, ip=ip, request=request)
-        times_up = settings.SECURITY_LOGIN_LIMIT_COUNT
-        times_failed = get_login_failed_count(username, ip)
-        times_try = int(times_up) - int(times_failed)
+        util = LoginBlockUtil(username, ip)
+        times_remainder = util.get_remainder_times()
         block_time = settings.SECURITY_LOGIN_LIMIT_TIME
 
         default_msg = invalid_login_msg.format(
-            times_try=times_try, block_time=block_time
+            times_try=times_remainder, block_time=block_time
         )
         if error == reason_password_failed:
             self.msg = default_msg
@@ -123,12 +128,32 @@ class CredentialError(AuthFailedNeedLogMixin, AuthFailedNeedBlockMixin, AuthFail
 
 class MFAFailedError(AuthFailedNeedLogMixin, AuthFailedError):
     error = reason_mfa_failed
-    msg = mfa_failed_msg
+    msg: str
 
-    def __init__(self, username, request):
+    def __init__(self, username, request, ip):
+        util = MFABlockUtils(username, ip)
+        util.incr_failed_count()
+
+        times_remainder = util.get_remainder_times()
+        block_time = settings.SECURITY_LOGIN_LIMIT_TIME
+
+        if times_remainder:
+            self.msg = mfa_failed_msg.format(
+                times_try=times_remainder, block_time=block_time
+            )
+        else:
+            self.msg = block_mfa_msg.format(settings.SECURITY_LOGIN_LIMIT_TIME)
         super().__init__(username=username, request=request)
 
 
+class BlockMFAError(AuthFailedNeedLogMixin, AuthFailedError):
+    error = 'block_mfa'
+
+    def __init__(self, username, request, ip):
+        self.msg = block_mfa_msg.format(settings.SECURITY_LOGIN_LIMIT_TIME)
+        super().__init__(username=username, request=request, ip=ip)
+
+
 class MFAUnsetError(AuthFailedNeedLogMixin, AuthFailedError):
     error = reason_mfa_unset
     msg = mfa_unset_msg

apps/authentication/mixins.py

@@ -15,9 +15,7 @@ from django.shortcuts import reverse
 
 from common.utils import get_object_or_none, get_request_ip, get_logger, bulk_get
 from users.models import User
-from users.utils import (
-    is_block_login, clean_failed_count
-)
+from users.utils import LoginBlockUtil, MFABlockUtils
 from . import errors
 from .utils import rsa_decrypt
 from .signals import post_auth_success, post_auth_failed
@@ -117,7 +115,7 @@ class AuthMixin:
         else:
             username = self.request.POST.get("username")
         ip = self.get_request_ip()
-        if is_block_login(username, ip):
+        if LoginBlockUtil(username, ip).is_block():
             logger.warn('Ip was blocked' + ': ' + username + ':' + ip)
             exception = errors.BlockLoginError(username=username, ip=ip)
             if raise_exception:
@@ -197,7 +195,7 @@ class AuthMixin:
         self._check_password_require_reset_or_not(user)
         self._check_passwd_is_too_simple(user, password)
 
-        clean_failed_count(username, ip)
+        LoginBlockUtil(username, ip).clean_failed_count()
         request.session['auth_password'] = 1
         request.session['user_id'] = str(user.id)
         request.session['auto_login'] = auto_login
@@ -253,15 +251,34 @@ class AuthMixin:
             raise errors.MFAUnsetError(user, self.request, url)
         raise errors.MFARequiredError()
 
+    def mark_mfa_ok(self):
+        self.request.session['auth_mfa'] = 1
+        self.request.session['auth_mfa_time'] = time.time()
+        self.request.session['auth_mfa_type'] = 'otp'
+
+    def check_mfa_is_block(self, username, ip, raise_exception=True):
+        if MFABlockUtils(username, ip).is_block():
+            logger.warn('Ip was blocked' + ': ' + username + ':' + ip)
+            exception = errors.BlockMFAError(username=username, request=self.request, ip=ip)
+            if raise_exception:
+                raise exception
+            else:
+                return exception
+
     def check_user_mfa(self, code):
         user = self.get_user_from_session()
+        ip = self.get_request_ip()
+        self.check_mfa_is_block(user.username, ip)
         ok = user.check_mfa(code)
         if ok:
-            self.request.session['auth_mfa'] = 1
-            self.request.session['auth_mfa_time'] = time.time()
-            self.request.session['auth_mfa_type'] = 'otp'
+            self.mark_mfa_ok()
             return
-        raise errors.MFAFailedError(username=user.username, request=self.request)
+
+        raise errors.MFAFailedError(
+            username=user.username,
+            request=self.request,
+            ip=ip
+        )
 
     def get_ticket(self):
         from tickets.models import Ticket

apps/authentication/views/mfa.py

@@ -22,10 +22,12 @@ class UserLoginOtpView(mixins.AuthMixin, FormView):
         try:
             self.check_user_mfa(otp_code)
             return redirect_to_guard_view()
-        except errors.MFAFailedError as e:
+        except (errors.MFAFailedError, errors.BlockMFAError) as e:
             form.add_error('otp_code', e.msg)
             return super().form_invalid(form)
         except Exception as e:
             logger.error(e)
+            import traceback
+            traceback.print_exception()
             return redirect_to_guard_view()