增加两步验证基础代码

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)

oj/urls.py

@@ -6,7 +6,8 @@ from django.views.generic import TemplateView
 from account.views import (UserLoginAPIView, UsernameCheckAPIView, UserRegisterAPIView,
                            UserChangePasswordAPIView, EmailCheckAPIView,
                            UserAdminAPIView, UserInfoAPIView, ResetPasswordAPIView,
-                           ApplyResetPasswordAPIView, SSOAPIView, UserProfileAPIView)
+                           ApplyResetPasswordAPIView, SSOAPIView, UserProfileAPIView,
+                           TwoFactorAuthAPIView)
 
 from announcement.views import AnnouncementAdminAPIView
 
@@ -132,7 +133,8 @@ urlpatterns = [
     url(r'^account/sso/$', SSOAPIView.as_view(), name="sso_api"),
     url(r'^api/account/userprofile/$', UserProfileAPIView.as_view(), name="userprofile_api"),
     url(r'^reset_password/$', TemplateView.as_view(template_name="oj/account/apply_reset_password.html"), name="apply_reset_password_page"),
-     url(r'^reset_password/t/(?P<token>\w+)/$', "account.views.reset_password_page", name="reset_password_page")
+    url(r'^reset_password/t/(?P<token>\w+)/$', "account.views.reset_password_page", name="reset_password_page"),
+    url(r'^api/two_factor_auth/$', TwoFactorAuthAPIView.as_view(), name="two_factor_auth_api"),
 ]
 
 

utils/otp_auth.py

@@ -0,0 +1,208 @@
+# -*- coding: utf-8 -*-
+"""
+    otpauth
+    ~~~~~~~
+
+    Implements two-step verification of HOTP/TOTP.
+
+    :copyright: (c) 2013 - 2015 by Hsiaoming Yang.
+    :license: BSD, see LICENSE for more details.
+"""
+import sys
+import time
+import hmac
+import base64
+import struct
+import hashlib
+import warnings
+
+
+if sys.version_info[0] == 3:
+    PY2 = False
+    string_type = str
+else:
+    PY2 = True
+    string_type = unicode
+    range = xrange
+
+
+__author__ = 'Hsiaoming Yang <me@lepture.com>'
+__homepage__ = 'https://github.com/lepture/otpauth'
+__version__ = '1.0.1'
+
+
+__all__ = ['OtpAuth', 'HOTP', 'TOTP', 'generate_hotp', 'generate_totp']
+
+
+HOTP = 'hotp'
+TOTP = 'totp'
+
+
+class OtpAuth(object):
+    """One Time Password Authentication.
+
+    :param secret: A secret token for the authentication.
+    """
+
+    def __init__(self, secret):
+        self.secret = secret
+
+    def hotp(self, counter=4):
+        """Generate a HOTP code.
+
+        :param counter: HOTP is a counter based algorithm.
+        """
+        return generate_hotp(self.secret, counter)
+
+    def totp(self, period=30, timestamp=None):
+        """Generate a TOTP code.
+
+        A TOTP code is an extension of HOTP algorithm.
+
+        :param period: A period that a TOTP code is valid in seconds
+        :param timestamp: Create TOTP at this given timestamp
+        """
+        return generate_totp(self.secret, period, timestamp)
+
+    def valid_hotp(self, code, last=0, trials=100):
+        """Valid a HOTP code.
+
+        :param code: A number that is less than 6 characters.
+        :param last: Guess HOTP code from last + 1 range.
+        :param trials: Guest HOTP code end at last + trials + 1.
+        """
+        if not valid_code(code):
+            return False
+
+        code = bytes(int(code))
+        for i in range(last + 1, last + trials + 1):
+            if compare_digest(bytes(self.hotp(counter=i)), code):
+                return i
+        return False
+
+    def valid_totp(self, code, period=30, timestamp=None):
+        """Valid a TOTP code.
+
+        :param code: A number that is less than 6 characters.
+        :param period: A period that a TOTP code is valid in seconds
+        :param timestamp: Validate TOTP at this given timestamp
+        """
+        if not valid_code(code):
+            return False
+        return compare_digest(
+            bytes(self.totp(period, timestamp)),
+            bytes(int(code))
+        )
+
+    @property
+    def encoded_secret(self):
+        secret = base64.b32encode(to_bytes(self.secret))
+        # bytes to string
+        secret = secret.decode('utf-8')
+        # remove pad string
+        return secret.strip('=')
+
+    def to_uri(self, type, label, issuer, counter=None):
+        """Generate the otpauth protocal string.
+
+        :param type: Algorithm type, hotp or totp.
+        :param label: Label of the identifier.
+        :param issuer: The company, the organization or something else.
+        :param counter: Counter of the HOTP algorithm.
+        """
+        type = type.lower()
+
+        if type not in ('hotp', 'totp'):
+            raise ValueError('type must be hotp or totp')
+
+        if type == 'hotp' and not counter:
+            raise ValueError('HOTP type authentication need counter')
+
+        # https://code.google.com/p/google-authenticator/wiki/KeyUriFormat
+        url = ('otpauth://%(type)s/%(label)s?secret=%(secret)s'
+               '&issuer=%(issuer)s')
+        dct = dict(
+            type=type, label=label, issuer=issuer,
+            secret=self.encoded_secret, counter=counter
+        )
+        ret = url % dct
+        if type == 'hotp':
+            ret = '%s&counter=%s' % (ret, counter)
+        return ret
+
+    def to_google(self, type, label, issuer, counter=None):
+        """Generate the otpauth protocal string for Google Authenticator.
+
+        .. deprecated:: 0.2.0
+           Use :func:`to_uri` instead.
+        """
+        warnings.warn('deprecated, use to_uri instead', DeprecationWarning)
+        return self.to_uri(type, label, issuer, counter)
+
+
+def generate_hotp(secret, counter=4):
+    """Generate a HOTP code.
+
+    :param secret: A secret token for the authentication.
+    :param counter: HOTP is a counter based algorithm.
+    """
+    # https://tools.ietf.org/html/rfc4226
+    msg = struct.pack('>Q', counter)
+    digest = hmac.new(to_bytes(secret), msg, hashlib.sha1).digest()
+
+    ob = digest[19]
+    if PY2:
+        ob = ord(ob)
+
+    pos = ob & 15
+    base = struct.unpack('>I', digest[pos:pos + 4])[0] & 0x7fffffff
+    token = base % 1000000
+    return token
+
+
+def generate_totp(secret, period=30, timestamp=None):
+    """Generate a TOTP code.
+
+    A TOTP code is an extension of HOTP algorithm.
+
+    :param secret: A secret token for the authentication.
+    :param period: A period that a TOTP code is valid in seconds
+    :param timestamp: Current time stamp.
+    """
+    if timestamp is None:
+        timestamp = time.time()
+    counter = int(timestamp) // period
+    return generate_hotp(secret, counter)
+
+
+def to_bytes(text):
+    if isinstance(text, string_type):
+        # Python3 str -> bytes
+        # Python2 unicode -> str
+        text = text.encode('utf-8')
+    return text
+
+
+def valid_code(code):
+    code = string_type(code)
+    return code.isdigit() and len(code) <= 6
+
+
+def compare_digest(a, b):
+    func = getattr(hmac, 'compare_digest', None)
+    if func:
+        return func(a, b)
+
+    # fallback
+    if len(a) != len(b):
+        return False
+
+    rv = 0
+    if PY2:
+        from itertools import izip
+        for x, y in izip(a, b):
+            rv |= ord(x) ^ ord(y)
+    else:
+        for x, y in zip(a, b):
+            rv |= x ^ y
+    return rv == 0