[Update] 修改 token api

apps/authentication/api/auth.py

@@ -22,7 +22,7 @@ from users.utils import (
     check_otp_code, increase_login_failed_count,
     is_block_login, clean_failed_count
 )
-from .. import const
+from .. import errors
 from ..utils import check_user_valid
 from ..serializers import OtpVerifySerializer
 from ..signals import post_auth_success, post_auth_failed
@@ -174,7 +174,7 @@ class UserOtpAuthApi(RootOrgViewMixin, APIView):
                 status=401
             )
         if not check_otp_code(user.otp_secret_key, otp_code):
-            self.send_auth_signal(success=False, username=user.username, reason=const.mfa_failed)
+            self.send_auth_signal(success=False, username=user.username, reason=errors.mfa_failed)
             return Response({'msg': _('MFA certification failed')}, status=401)
         self.send_auth_signal(success=True, user=user)
         token, expired_at = user.create_bearer_token(request)

apps/authentication/api/token.py

 # -*- coding: utf-8 -*-
 #
 
-import uuid
-
-from django.core.cache import cache
 from django.utils.translation import ugettext as _
 from rest_framework.permissions import AllowAny
 from rest_framework.response import Response
 from rest_framework.generics import CreateAPIView
-from drf_yasg.utils import swagger_auto_schema
 
-from common.utils import get_request_ip, get_logger
+from common.utils import get_request_ip, get_logger, get_object_or_none
 from users.utils import (
     check_otp_code, increase_login_failed_count,
     is_block_login, clean_failed_count
 )
+from users.models import User
 from ..utils import check_user_valid
 from ..signals import post_auth_success, post_auth_failed
-from .. import serializers
+from .. import serializers, errors
 
 
 logger = get_logger(__name__)
@@ -25,29 +22,41 @@ logger = get_logger(__name__)
 __all__ = ['TokenCreateApi']
 
 
-class AuthFailedError(Exception):
-    def __init__(self, msg, reason=None):
-        self.msg = msg
-        self.reason = reason
-
-
-class MFARequiredError(Exception):
-    pass
-
-
 class TokenCreateApi(CreateAPIView):
     permission_classes = (AllowAny,)
     serializer_class = serializers.BearerTokenSerializer
 
-    @staticmethod
-    def check_is_block(username, ip):
+    def check_session(self):
+        pass
+
+    def get_request_ip(self):
+        ip = self.request.data.get('remote_addr', None)
+        ip = ip or get_request_ip(self.request)
+        return ip
+
+    def check_is_block(self):
+        username = self.request.data.get("username")
+        ip = self.get_request_ip()
         if is_block_login(username, ip):
-            msg = _("Log in frequently and try again later")
+            msg = errors.ip_blocked
             logger.warn(msg + ': ' + username + ':' + ip)
-            raise AuthFailedError(msg)
+            raise errors.AuthFailedError(msg, 'blocked')
 
-    def check_user_valid(self):
+    def get_user_from_session(self):
+        user_id = self.request.session["user_id"]
+        user = get_object_or_none(User, pk=user_id)
+        if not user:
+            error = "Not user in session: {}".format(user_id)
+            raise errors.AuthFailedError(error, 'session_error')
+        return user
+
+    def check_user_auth(self):
         request = self.request
+        if request.session.get("auth_password") and \
+                request.session.get('user_id'):
+            user = self.get_user_from_session()
+            return user
+        self.check_is_block()
         username = request.data.get('username', '')
         password = request.data.get('password', '')
         public_key = request.data.get('public_key', '')
@@ -55,34 +64,76 @@ class TokenCreateApi(CreateAPIView):
             username=username, password=password,
             public_key=public_key
         )
+        ip = self.get_request_ip()
         if not user:
-            raise AuthFailedError(msg)
+            raise errors.AuthFailedError(msg, error='auth_failed', username=username)
+        clean_failed_count(username, ip)
+        request.session['auth_password'] = 1
+        request.session['user_id'] = str(user.id)
         return user
 
+    def check_user_mfa_if_need(self, user):
+        if self.request.session.get('auth_mfa'):
+            return True
+        if not user.otp_enabled or not user.otp_secret_key:
+            return True
+        otp_code = self.request.data.get("otp_code")
+        if not otp_code:
+            raise errors.MFARequiredError()
+        if not check_otp_code(user.otp_secret_key, otp_code):
+            raise errors.AuthFailedError(
+                errors.mfa_failed, error='mfa_failed',
+                username=user.username,
+            )
+        return True
+
+    def check_user_login_confirm_if_need(self, user):
+        from orders.models import LoginConfirmOrder
+        confirm_setting = user.get_login_confirm_setting()
+        if self.request.session.get('auth_confirm') or not confirm_setting:
+            return
+        order = None
+        if self.request.session.get('auth_order_id'):
+            order_id = self.request.session['auth_order_id']
+            order = get_object_or_none(LoginConfirmOrder, pk=order_id)
+        if not order:
+            order = confirm_setting.create_confirm_order(self.request)
+            self.request.session['auth_order_id'] = str(order.id)
+
+        if order.status == "accepted":
+            return
+        elif order.status == "rejected":
+            raise errors.LoginConfirmRejectedError()
+        else:
+            raise errors.LoginConfirmWaitError()
+
     def create(self, request, *args, **kwargs):
-        username = self.request.data.get('username')
-        ip = self.request.data.get('remote_addr', None)
-        ip = ip or get_request_ip(self.request)
-        user = None
+        self.check_session()
+        # 如果认证没有过，检查账号密码
         try:
-            self.check_is_block(username, ip)
-            user = self.check_user_valid()
-            if user.otp_enabled:
-                raise MFARequiredError()
+            user = self.check_user_auth()
+            self.check_user_mfa_if_need(user)
+            self.check_user_login_confirm_if_need(user)
             self.send_auth_signal(success=True, user=user)
-            clean_failed_count(username, ip)
             resp = super().create(request, *args, **kwargs)
             return resp
-        except AuthFailedError as e:
-            increase_login_failed_count(username, ip)
-            self.send_auth_signal(success=False, user=user, username=username, reason=str(e))
-            return Response({'msg': str(e)}, status=401)
-        except MFARequiredError:
+        except errors.AuthFailedError as e:
+            if e.username:
+                increase_login_failed_count(e.username, self.get_request_ip())
+                self.send_auth_signal(
+                    success=False, username=e.username, reason=e.reason
+                )
+            return Response({'msg': e.reason, 'error': e.error}, status=401)
+        except errors.MFARequiredError:
             msg = _("MFA required")
-            seed = uuid.uuid4().hex
-            cache.set(seed, user.username, 300)
-            data = {'msg': msg, "choices": ["otp"], "req": seed}
+            data = {'msg': msg, "choices": ["otp"], "error": 'mfa_required'}
             return Response(data, status=300)
+        except errors.LoginConfirmRejectedError as e:
+            pass
+        except errors.LoginConfirmWaitError as e:
+            pass
+        except errors.LoginConfirmRequiredError as e:
+            pass
 
     def send_auth_signal(self, success=True, user=None, username='', reason=''):
         if success:

apps/authentication/const.py

-# -*- coding: utf-8 -*-
-#
-from django.utils.translation import ugettext_lazy as _
-
-
-password_failed = _('Username/password check failed')
-mfa_failed = _('MFA authentication failed')
-user_not_exist = _("Username does not exist")
-password_expired = _("Password expired")
-user_invalid = _('Disabled or expired')

apps/authentication/errors.py

+# -*- coding: utf-8 -*-
+#
+from django.utils.translation import ugettext_lazy as _
+
+password_failed = _('Username/password check failed')
+mfa_failed = _('MFA authentication failed')
+user_not_exist = _("Username does not exist")
+password_expired = _("Password expired")
+user_invalid = _('Disabled or expired')
+ip_blocked = _("Log in frequently and try again later")
+
+mfa_required = _("MFA required")
+login_confirm_required = _("Login confirm required")
+login_confirm_wait = _("Wait login confirm")
+
+
+class AuthFailedError(Exception):
+    def __init__(self, reason, error=None, username=None):
+        self.reason = reason
+        self.error = error
+        self.username = username
+
+
+class MFARequiredError(Exception):
+    reason = mfa_required
+    error = 'mfa_required'
+
+
+class LoginConfirmRequiredError(Exception):
+    reason = login_confirm_required
+    error = 'login_confirm_required'
+
+
+class LoginConfirmWaitError(Exception):
+    reason = login_confirm_wait
+    error = 'login_confirm_wait'
+
+
+class LoginConfirmRejectedError(Exception):
+    reason = login_confirm_wait
+    error = 'login_confirm_rejected'

apps/authentication/utils.py

@@ -8,7 +8,7 @@ from common.utils import (
     get_ip_city, get_object_or_none, validate_ip, get_request_ip
 )
 from users.models import User
-from . import const
+from . import errors
 
 
 def write_login_log(*args, **kwargs):
@@ -38,11 +38,11 @@ def check_user_valid(**kwargs):
         user = None
 
     if user is None:
-        return None, const.user_not_exist
+        return None, errors.user_not_exist
     elif not user.is_valid:
-        return None, const.user_invalid
+        return None, errors.user_invalid
     elif user.password_has_expired:
-        return None, const.password_expired
+        return None, errors.password_expired
 
     if password and authenticate(username=username, password=password):
         return user, ''
@@ -55,4 +55,4 @@ def check_user_valid(**kwargs):
         elif len(public_key_saved) > 1:
             if public_key == public_key_saved[1]:
                 return user, ''
-    return None, const.password_failed
+    return None, errors.password_failed

apps/authentication/views/login.py

@@ -27,7 +27,7 @@ from users.utils import (
 from ..models import LoginConfirmSetting
 from ..signals import post_auth_success, post_auth_failed
 from .. import forms
-from .. import const
+from .. import errors
 
 
 __all__ = [
@@ -83,7 +83,7 @@ class UserLoginView(FormView):
         user = form.get_user()
         # user password expired
         if user.password_has_expired:
-            reason = const.password_expired
+            reason = errors.password_expired
             self.send_auth_signal(success=False, username=user.username, reason=reason)
             return self.render_to_response(self.get_context_data(password_expired=True))
 
@@ -99,7 +99,7 @@ class UserLoginView(FormView):
         # write login failed log
         username = form.cleaned_data.get('username')
         exist = User.objects.filter(username=username).first()
-        reason = const.password_failed if exist else const.user_not_exist
+        reason = errors.password_failed if exist else errors.user_not_exist
         # limit user login failed count
         ip = get_request_ip(self.request)
         increase_login_failed_count(username, ip)
@@ -150,7 +150,7 @@ class UserLoginOtpView(FormView):
         else:
             self.send_auth_signal(
                 success=False, username=user.username,
-                reason=const.mfa_failed
+                reason=errors.mfa_failed
             )
             form.add_error(
                 'otp_code', _('MFA code invalid, or ntp sync server time')