[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

@@ -1,23 +1,20 @@
 # -*- 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):
-        if is_block_login(username, ip):
-            msg = _("Log in frequently and try again later")
-            logger.warn(msg + ': ' + username + ':' + ip)
-            raise AuthFailedError(msg)
+    def check_session(self):
+        pass
 
-    def check_user_valid(self):
+    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 = errors.ip_blocked
+            logger.warn(msg + ': ' + username + ':' + ip)
+            raise errors.AuthFailedError(msg, 'blocked')
+
+    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: