Add sso and 2fa api

account/serializers.py

@@ -49,3 +49,18 @@ class EditUserSerializer(serializers.Serializer):
 class ApplyResetPasswordSerializer(serializers.Serializer):
     email = serializers.EmailField()
     captcha = serializers.CharField(max_length=4, min_length=4)
+
+
+class ResetPasswordSerializer(serializers.Serializer):
+    token = serializers.CharField(min_length=1, max_length=40)
+    password = serializers.CharField(min_length=6, max_length=30)
+    captcha = serializers.CharField(max_length=4, min_length=4)
+
+
+class SSOSerializer(serializers.Serializer):
+    appkey = serializers.CharField(max_length=35)
+    token = serializers.CharField(max_length=40)
+
+
+class TwoFactorAuthCodeSerializer(serializers.Serializer):
+    code = serializers.IntegerField()

account/urls/user.py

@@ -3,9 +3,12 @@
 
 from django.conf.urls import url
 
-from ..views.user import UserInfoAPI, UserProfileAPI
+from ..views.user import (UserInfoAPI, UserProfileAPI,
+                          SSOAPI, TwoFactorAuthAPI)
 
 urlpatterns = [
-    url(r"^user", UserInfoAPI.as_view(), name="user_info_api"),
-    url(r"^profile$", UserProfileAPI.as_view(), name="user_profile_api")
+    url(r"^user$", UserInfoAPI.as_view(), name="user_info_api"),
+    url(r"^profile$", UserProfileAPI.as_view(), name="user_profile_api"),
+    url(r"^sso$", SSOAPI.as_view(), name="sso_api"),
+    url(r"^two_factor_auth$", TwoFactorAuthAPI.as_view(), name="two_factor_auth_api")
 ]

account/views/oj.py

@@ -1,16 +1,14 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
-import os
-
 from datetime import timedelta
+from otpauth import OtpAuth
 
 from django.contrib import auth
 from django.conf import settings
 from django.core.exceptions import MultipleObjectsReturned
 from django.utils.translation import ugettext as _
 from django.utils.timezone import now
-from otpauth import OtpAuth
 
 from conf.models import WebsiteConfig
 from utils.api import APIView, validate_serializer
@@ -21,7 +19,7 @@ from ..decorators import login_required
 from ..models import User, UserProfile
 from ..serializers import (UserChangePasswordSerializer, UserLoginSerializer,
                            UserRegisterSerializer,
-                           ApplyResetPasswordSerializer)
+                           ApplyResetPasswordSerializer, ResetPasswordSerializer)
 from ..tasks import _send_email
 
 
@@ -112,6 +110,7 @@ class ApplyResetPasswordAPI(APIView):
     def post(self, request):
         data = request.data
         captcha = Captcha(request)
+        config = WebsiteConfig.objects.first()
         if not captcha.check(data["captcha"]):
             return self.error(_("Invalid captcha"))
         try:
@@ -131,7 +130,6 @@ class ApplyResetPasswordAPI(APIView):
             replace("{{ website_name }}", settings.WEBSITE_INFO["website_name"]). \
             replace("{{ link }}", settings.WEBSITE_INFO["url"] + "/reset_password/t/" +
                     user.reset_password_token)
-        config = WebsiteConfig.objects.first()
         _send_email.delay(config.name,
                           user.email,
                           user.username,
@@ -141,6 +139,7 @@ class ApplyResetPasswordAPI(APIView):
 
 
 class ResetPasswordAPI(APIView):
+    @validate_serializer(ResetPasswordSerializer)
     def post(self, request):
         data = request.data
         captcha = Captcha(request)

account/views/user.py

@@ -1,12 +1,23 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
+import qrcode
+
+from io import StringIO
+from otpauth import OtpAuth
+
+from django.conf import settings
+from django.http import HttpResponse
 from django.utils.translation import ugettext as _
 
+from conf.models import WebsiteConfig
 from utils.api import APIView, validate_serializer
+from utils.shortcuts import rand_str
 
 from ..decorators import login_required
-from ..serializers import EditUserSerializer, UserSerializer
+from ..models import User
+from ..serializers import (EditUserSerializer, UserSerializer,
+                           SSOSerializer, TwoFactorAuthCodeSerializer)
 
 
 class UserInfoAPI(APIView):
@@ -43,3 +54,81 @@ class UserProfileAPI(APIView):
             # Timezone & language 暂时不加
         user_profile.save()
         return self.success(_("Succeeded"))
+
+
+class SSOAPI(APIView):
+    @login_required
+    def get(self, request):
+        callback = request.GET.get("callback", None)
+        if not callback:
+            return self.error(_("Parameter Error"))
+        token = rand_str()
+        request.user.auth_token = token
+        request.user.save()
+        return self.success({"redirect_url": callback + "?token=" + token,
+                             "callback": callback})
+
+    @validate_serializer(SSOSerializer)
+    def post(self, request):
+        data = request.data
+        try:
+            User.objects.get(open_api_appkey=data["appkey"])
+        except User.DoesNotExist:
+            return self.error(_("Invalid appkey"))
+        try:
+            user = User.objects.get(auth_token=data["token"])
+            user.auth_token = None
+            user.save()
+            return self.success({"username": user.username,
+                                 "id": user.id,
+                                 "admin_type": user.admin_type,
+                                 "avatar": user.userprofile.avatar})
+        except User.DoesNotExist:
+            return self.error("User does not exist")
+
+
+class TwoFactorAuthAPI(APIView):
+    @login_required
+    def get(self, request):
+        """
+        Get QR code
+        """
+        user = request.user
+        if user.two_factor_auth:
+            return self.error("Already open 2FA")
+        token = rand_str()
+        user.tfa_token = token
+        user.save()
+
+        config = WebsiteConfig.objects.first()
+        image = qrcode.make(OtpAuth(token).to_uri("totp", config.base_url, config.name))
+        buf = StringIO()
+        image.save(buf, 'gif')
+
+        return HttpResponse(buf.getvalue(), 'image/gif')
+
+    @login_required
+    @validate_serializer(TwoFactorAuthCodeSerializer)
+    def post(self, request):
+        """
+        Open 2FA
+        """
+        code = request.data["code"]
+        user = request.user
+        if OtpAuth(user.tfa_token).valid_totp(code):
+            user.two_factor_auth = True
+            user.save()
+            return self.success(_("Succeeded"))
+        else:
+            return self.error(_("Invalid captcha"))
+
+    @login_required
+    @validate_serializer(TwoFactorAuthCodeSerializer)
+    def put(self, request):
+        code = request.data["code"]
+        user = request.user
+        if OtpAuth(user.tfa_token).valid_totp(code):
+            user.two_factor_auth = False
+            user.save()
+        else:
+            return self.error(_("Invalid captcha"))

deploy/requirements.txt

@@ -8,4 +8,5 @@ pytz
 coverage
 python-dateutil
 celery
-Envelopes
+Envelopes
+qrcode

requirements.txt

@@ -6,4 +6,5 @@ python-dateutil
 celery
 Envelopes
 pytz
-jsonfield
+jsonfield
+qrcode