diff --git a/account/serializers.py b/account/serializers.py index ebae355..4d1077a 100644 --- a/account/serializers.py +++ b/account/serializers.py @@ -86,5 +86,5 @@ class UserProfileSerializer(serializers.ModelSerializer): "rank", "accepted_number", "submissions_number", "problems_status", "phone_number", "school", "student_id"] -class ApplyTwoFactorAuthSerializer(serializers.Serializer): +class TwoFactorAuthCodeSerializer(serializers.Serializer): code = serializers.IntegerField() diff --git a/account/views.py b/account/views.py index 88919c7..86e28dd 100644 --- a/account/views.py +++ b/account/views.py @@ -27,7 +27,7 @@ from .serializers import (UserLoginSerializer, UserRegisterSerializer, UserSerializer, EditUserSerializer, ApplyResetPasswordSerializer, ResetPasswordSerializer, SSOSerializer, EditUserProfileSerializer, - UserProfileSerializer, ApplyTwoFactorAuthSerializer) + UserProfileSerializer, TwoFactorAuthCodeSerializer) from .decorators import super_admin_required @@ -405,7 +405,7 @@ class TwoFactorAuthAPIView(APIView): """ 开启两步验证 """ - serializer = ApplyTwoFactorAuthSerializer(data=request.data) + serializer = TwoFactorAuthCodeSerializer(data=request.data) if serializer.is_valid(): code = serializer.data["code"] user = request.user @@ -416,4 +416,18 @@ class TwoFactorAuthAPIView(APIView): else: return error_response(u"验证码错误") else: - return serializer_invalid_response(serializer) \ No newline at end of file + return serializer_invalid_response(serializer) + + @login_required + def put(self, request): + serializer = TwoFactorAuthCodeSerializer(data=request.data) + if serializer.is_valid(): + user = request.user + code = serializer.data["code"] + if OtpAuth(user.tfa_token).valid_totp(code): + user.two_factor_auth = False + user.save() + else: + return error_response(u"验证码错误") + else: + return serializer_invalid_response(serializer)