增加两步验证基础代码

account/migrations/0016_auto_20151211_2230.py

@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.9 on 2015-12-11 14:30
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('account', '0015_userprofile_student_id'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='user',
+            name='tfa_token',
+            field=models.CharField(blank=True, max_length=10, null=True),
+        ),
+        migrations.AddField(
+            model_name='user',
+            name='two_factor_auth',
+            field=models.BooleanField(default=False),
+        ),
+    ]

account/models.py

@@ -40,6 +40,9 @@ class User(AbstractBaseUser):
     reset_password_token_create_time = models.DateTimeField(blank=True, null=True)
     # 论坛授权token
     auth_token = models.CharField(max_length=40, blank=True, null=True)
+    # 是否开启两步验证
+    two_factor_auth = models.BooleanField(default=False)
+    tfa_token = models.CharField(max_length=10, blank=True, null=True)
 
     USERNAME_FIELD = 'username'
     REQUIRED_FIELDS = []

account/serializers.py

@@ -84,3 +84,7 @@ class UserProfileSerializer(serializers.ModelSerializer):
         model = UserProfile
         fields = ["avatar", "blog", "mood", "hduoj_username", "bestcoder_username", "codeforces_username",
                   "rank", "accepted_number", "submissions_number", "problems_status", "phone_number", "school", "student_id"]
+
+
+class ApplyTwoFactorAuthSerializer(serializers.Serializer):
+    code = serializers.IntegerField()

account/views.py

@@ -1,11 +1,13 @@
 # coding=utf-8
 import codecs
+import qrcode
+import StringIO
 from django import http
 from django.contrib import auth
 from django.shortcuts import render
 from django.db.models import Q
 from django.conf import settings
-from django.http import HttpResponseRedirect
+from django.http import HttpResponse
 from django.core.exceptions import MultipleObjectsReturned
 from django.utils.timezone import now
 
@@ -15,6 +17,7 @@ from utils.shortcuts import (serializer_invalid_response, error_response,
                              success_response, error_page, paginate, rand_str)
 from utils.captcha import Captcha
 from utils.mail import send_email
+from utils.otp_auth import OtpAuth
 
 from .decorators import login_required
 from .models import User, UserProfile
@@ -23,7 +26,8 @@ from .serializers import (UserLoginSerializer, UserRegisterSerializer,
                           UserChangePasswordSerializer,
                           UserSerializer, EditUserSerializer,
                           ApplyResetPasswordSerializer, ResetPasswordSerializer,
-                          SSOSerializer, EditUserProfileSerializer, UserProfileSerializer)
+                          SSOSerializer, EditUserProfileSerializer,
+                          UserProfileSerializer, ApplyTwoFactorAuthSerializer)
 
 from .decorators import super_admin_required
 
@@ -151,9 +155,9 @@ class EmailCheckAPIView(APIView):
         检测邮箱是否存在，用状态码标识结果
         ---
         """
-        #这里是为了适应前端表单验证空间的要求
+        # 这里是为了适应前端表单验证空间的要求
         reset = request.GET.get("reset", None)
-        #如果reset为true说明该请求是重置密码页面发出的，要返回的状态码应正好相反
+        # 如果reset为true说明该请求是重置密码页面发出的，要返回的状态码应正好相反
         if reset:
             existed = 200
             does_not_existed = 400
@@ -375,3 +379,41 @@ def reset_password_page(request, token):
     if (now() - user.reset_password_token_create_time).total_seconds() > 30 * 60:
         return error_page(request, u"链接已过期")
     return render(request, "oj/account/reset_password.html", {"user": user})
+
+
+class TwoFactorAuthAPIView(APIView):
+    @login_required
+    def get(self, request):
+        """
+        获取绑定二维码
+        """
+        user = request.user
+        if user.two_factor_auth:
+            return error_response(u"已经开启两步验证了")
+        token = rand_str()
+        user.tfa_token = token
+        user.save()
+
+        image = qrcode.make(OtpAuth(token).to_uri("totp", "OnlineJudge", "OnlineJudge"))
+        buf = StringIO.StringIO()
+        image.save(buf, 'gif')
+
+        return HttpResponse(buf.getvalue(), 'image/gif')
+
+    @login_required
+    def post(self, request):
+        """
+        开启两步验证
+        """
+        serializer = ApplyTwoFactorAuthSerializer(data=request.data)
+        if serializer.is_valid():
+            code = serializer.data["code"]
+            user = request.user
+            if OtpAuth(user.tfa_token).valid_totp(code):
+                user.two_factor_auth = True
+                user.save()
+                return success_response(u"开启两步验证成功")
+            else:
+                return error_response(u"验证码错误")
+        else:
+            return serializer_invalid_response(serializer)