Add sso and 2fa api

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"))