[Update] 添加OTP认证功能

apps/users/api.py

@@ -2,6 +2,7 @@
 import uuid
 
 from django.core.cache import cache
+from django.urls import reverse
 
 from rest_framework import generics
 from rest_framework.permissions import AllowAny, IsAuthenticated
@@ -139,40 +140,75 @@ class UserProfile(APIView):
         return Response(self.serializer_class(request.user).data)
 
 
-class UserAuthApi(APIView):
+class UserOtpAuthApi(APIView):
     permission_classes = (AllowAny,)
     serializer_class = UserSerializer
 
     def post(self, request):
-        otp_check = request.data.get('otp_check', None)
+        otp_code = request.data.get('otp_code', '')
+        seed = request.data.get('seed', '')
+
+        user = cache.get(seed, None)
+        if not user:
+            return Response({'msg': '请先进行用户名和密码验证'}, status=401)
+
+        if not check_otp_code(user.otp_secret_key, otp_code):
+            return Response({'msg': 'otp认证失败'}, status=401)
+
+        token = generate_token(request, user)
+        self.write_login_log(request, user)
+        return Response(
+            {
+                'token': token,
+                'user': self.serializer_class(user).data
+             }
+        )
 
-        if otp_check:
-            # otp验证
-            return self.check_auth_otp(request)
-        else:
-            # password验证
-            return self.check_auth_password(request)
+    @staticmethod
+    def write_login_log(request, user):
+        login_ip = request.data.get('remote_addr', None)
+        login_type = request.data.get('login_type', '')
+        user_agent = request.data.get('HTTP_USER_AGENT', '')
+
+        if not login_ip:
+            login_ip = get_login_ip(request)
 
-    def check_auth_password(self, request):
+        write_login_log_async.delay(
+            user.username, ip=login_ip,
+            type=login_type, user_agent=user_agent,
+        )
+
+
+class UserAuthApi(APIView):
+    permission_classes = (AllowAny,)
+    serializer_class = UserSerializer
+
+    def post(self, request):
         user, msg = self.check_user_valid(request)
 
-        if user:
-            token = generate_token(request, user)
-            if not user.otp_enabled:
-                self.write_login_log(request, user)
-            return Response({'token': token, 'user': self.serializer_class(user).data})
-        else:
+        if not user:
             return Response({'msg': msg}, status=401)
 
-    def check_auth_otp(self, request):
-        otp_code = request.data.get('otp_code', '')
-        user, msg = self.check_user_valid(request)
-        if user:
+        if not user.otp_enabled:
             token = generate_token(request, user)
-            if check_otp_code(user.otp_secret_key, otp_code):
-                self.write_login_log(request, user)
-                return Response({'token': token, 'user': self.serializer_class(user).data})
-        return Response({'msg': msg}, status=401)
+            self.write_login_log(request, user)
+            return Response(
+                {
+                    'token': token,
+                    'user': self.serializer_class(user).data
+                }
+            )
+
+        seed = uuid.uuid4().hex
+        cache.set(seed, user, 300)
+        return Response(
+            {
+                'code': 101,
+                'msg': '请携带seed值,进行OTP二次认证',
+                'otp_url': reverse('api-users:user-otp-auth'),
+                'seed': seed,
+                'user': self.serializer_class(user).data
+            }, status=300)
 
     @staticmethod
     def check_user_valid(request):

apps/users/models/user.py

@@ -232,7 +232,7 @@ class User(AbstractUser):
 
     def disable_otp(self):
         self.otp_level = 0
-        self.otp_secret_key = ''
+        self.otp_secret_key = None
 
     def to_json(self):
         return OrderedDict({

apps/users/serializers.py

@@ -21,7 +21,7 @@ class UserSerializer(BulkSerializerMixin, serializers.ModelSerializer):
         list_serializer_class = BulkListSerializer
         exclude = [
             'first_name', 'last_name', 'password', '_private_key',
-            '_public_key', 'otp_secret_key', 'user_permissions'
+            '_public_key', '_otp_secret_key', 'user_permissions'
         ]
 
     def get_field_names(self, declared_fields, info):

apps/users/templates/users/_base_otp.html

@@ -25,7 +25,7 @@
 		 	<div>
 		 		<a href="{% url 'index' %}">首页</a>
 		 		<b>丨</b>
-		 		<a href="#">帮助中心</a>
+		 		<a href="http://docs.jumpserver.org/zh/docs/">文档</a>
 		 		<b>丨</b>
 		 		<a href="https://www.github.com/jumpserver/">GitHub</a>
 		 	</div>

apps/users/templates/users/login_otp.html

@@ -66,7 +66,7 @@
                             <button type="submit" class="btn btn-primary block full-width m-b">{% trans 'Next' %}</button>
 
                             <a href="#">
-                                <small>{% trans "Can't provide security? Please contact the administrator" %}</small>
+                                <small>{% trans "Can't provide otp code? Please contact the administrator" %}</small>
                             </a>
 
                         </form>

apps/users/templates/users/user_otp_enable_bind.html

@@ -7,10 +7,8 @@
     <div class="verify">
         <p style="margin:20px auto;"><strong style="color: #000000">使用手机 Google Authenticator 应用扫描以下二维码，获取6位验证码</strong></p>
 
-
         <div id="qr_code"></div>
 
-
         <form class="" role="form" method="post" action="">
             {% csrf_token %}
 
@@ -18,12 +16,13 @@
                 <input type="text" class="" name="otp_code" placeholder="{% trans 'Six figures' %}" required="">
             </div>
 
+            <button type="submit" class="next">{% trans 'Next' %}</button>
+
+
             {% if 'otp_code' in form.errors %}
                 <p style="color: #ed5565">{{ form.otp_code.errors.as_text }}</p>
             {% endif %}
 
-            <button type="submit" class="next">{% trans 'Next' %}</button>
-
         </form>
     </div>
 

apps/users/templates/users/user_password_authentication.html

@@ -5,6 +5,7 @@
 {% block content %}
     <form class="" role="form" method="post" action="">
         {% csrf_token %}
+
         <div class="form-input">
             <input type="text" class="" name="{{ form.username.html_name }}" value="{{ form.username.value }}" readonly="readonly" required="">
         </div>
@@ -13,13 +14,12 @@
             <input type="password" class="" name="{{ form.password.html_name }}" placeholder="{% trans 'Password' %}" required="">
         </div>
 
+        <button type="submit" class="next">{% trans 'Next' %}</button>
+
         {% if 'password' in form.errors %}
             <p class="red-fonts">{{ form.password.errors.as_text }}</p>
         {% endif %}
 
-
-        <button type="submit" class="next">{% trans 'Next' %}</button>
-
     </form>
 {% endblock %}
 

apps/users/templates/users/user_profile.html

@@ -152,8 +152,7 @@
                                                    href="
                                                         {% if request.user.otp_enabled and request.user.otp_secret_key %}
                                                             {% if request.user.otp_force_enabled %}
-                                                                javascript:void(0)
-                                                                "><span style="color:#ed5565">{% trans 'Disable' %}</span>
+                                                                " disabled >{% trans 'Disable' %}
                                                             {% else %}
                                                                 {% url 'users:user-otp-disable-authentication' %}
                                                                 ">{% trans 'Disable' %}

apps/users/urls/api_urls.py

@@ -20,6 +20,7 @@ urlpatterns = [
     url(r'^v1/connection-token/$', api.UserConnectionTokenApi.as_view(), name='connection-token'),
     url(r'^v1/profile/$', api.UserProfile.as_view(), name='user-profile'),
     url(r'^v1/auth/$', api.UserAuthApi.as_view(), name='user-auth'),
+    url(r'^v1/otp/auth/$', api.UserOtpAuthApi.as_view(), name='user-otp-auth'),
     url(r'^v1/users/(?P<pk>[0-9a-zA-Z\-]{36})/password/$',
         api.ChangeUserPasswordApi.as_view(), name='change-user-password'),
     url(r'^v1/users/(?P<pk>[0-9a-zA-Z\-]{36})/password/reset/$',

apps/users/utils.py

@@ -224,16 +224,14 @@ def get_ip_city(ip, timeout=10):
     return city
 
 
-def get_user(request):
-    if is_login(request):
-        user = request.user
-    else:
-        user = cache.get(request.session.session_key)
+def get_tmp_user_from_session(request):
+    user_id = request.session.get('tmp_user_id')
+    user = get_object_or_none(User, pk=user_id)
     return user
 
 
-def is_login(request):
-    return isinstance(request.user, User)
+def set_tmp_user_to_session(request, user):
+    request.session['tmp_user_id'] = str(user.id)
 
 
 def redirect_user_first_login_or_index(request, redirect_field_name):
@@ -244,9 +242,15 @@ def redirect_user_first_login_or_index(request, redirect_field_name):
         request.GET.get(redirect_field_name, reverse('index')))
 
 
-def generate_otp_uri(user, issuer="Jumpserver"):
-    otp_secret_key = base64.b32encode(os.urandom(10)).decode('utf-8')
-    cache.set('otp_secret_key', otp_secret_key, 300)
+def generate_otp_uri(request, issuer="Jumpserver"):
+    if request.user.is_authenticated:
+        user = request.user
+    else:
+        user = get_tmp_user_from_session(request)
+    otp_secret_key = cache.get(request.session.session_key+'otp_key', '')
+    if not otp_secret_key:
+        otp_secret_key = base64.b32encode(os.urandom(10)).decode('utf-8')
+    cache.set(request.session.session_key+'otp_key', otp_secret_key, 600)
     totp = pyotp.TOTP(otp_secret_key)
     return totp.provisioning_uri(name=user.username, issuer_name=issuer)
 

apps/users/views/login.py

@@ -19,12 +19,12 @@ from django.views.generic.base import TemplateView
 from django.views.generic.edit import FormView
 from formtools.wizard.views import SessionWizardView
 from django.conf import settings
-from django.core.cache import cache
 
 from common.utils import get_object_or_none
 from common.mixins import DatetimeSearchMixin, AdminUserRequiredMixin
 from ..models import User, LoginLog
-from ..utils import send_reset_password_mail, check_otp_code , get_login_ip, redirect_user_first_login_or_index
+from ..utils import send_reset_password_mail, check_otp_code, get_login_ip, redirect_user_first_login_or_index, \
+    get_tmp_user_from_session, set_tmp_user_to_session
 from ..tasks import write_login_log_async
 from .. import forms
 
@@ -54,11 +54,12 @@ class UserLoginView(FormView):
     def form_valid(self, form):
         if not self.request.session.test_cookie_worked():
             return HttpResponse(_("Please enable cookies and try again."))
-        cache.set(self.request.session.session_key, form.get_user(), 600)
+
+        set_tmp_user_to_session(self.request, form.get_user())
         return redirect(self.get_success_url())
 
     def get_success_url(self):
-        user = cache.get(self.request.session.session_key)
+        user = get_tmp_user_from_session(self.request)
 
         if user.otp_enabled and user.otp_secret_key:
             # 1,2 & T
@@ -94,7 +95,7 @@ class UserLoginOtpView(FormView):
     redirect_field_name = 'next'
 
     def form_valid(self, form):
-        user = cache.get(self.request.session.session_key)
+        user = get_tmp_user_from_session(self.request)
         otp_code = form.cleaned_data.get('otp_code')
         otp_secret_key = user.otp_secret_key
 

apps/users/views/user.py

@@ -35,7 +35,7 @@ from common.mixins import JSONResponseMixin
 from common.utils import get_logger, get_object_or_none, is_uuid, ssh_key_gen
 from .. import forms
 from ..models import User, UserGroup
-from ..utils import AdminUserRequiredMixin, generate_otp_uri, check_otp_code, get_user, is_login
+from ..utils import AdminUserRequiredMixin, generate_otp_uri, check_otp_code, get_tmp_user_from_session
 from ..signals import post_user_create
 from ..tasks import write_login_log_async
 
@@ -400,20 +400,31 @@ class UserOtpEnableAuthenticationView(FormView):
     form_class = forms.UserCheckPasswordForm
 
     def get_form(self, form_class=None):
+        if self.request.user.is_authenticated:
+            user = self.request.user
+        else:
+            user = get_tmp_user_from_session(self.request)
         form = super().get_form(form_class=form_class)
-        form['username'].initial = get_user(self.request).username
+        form['username'].initial = user.username
         return form
 
     def get_context_data(self, **kwargs):
+        if self.request.user.is_authenticated:
+            user = self.request.user
+        else:
+            user = get_tmp_user_from_session(self.request)
         context = {
-            'user': get_user(self.request)
+            'user': user
         }
         kwargs.update(context)
         return super().get_context_data(**kwargs)
 
     def form_valid(self, form):
+        if self.request.user.is_authenticated:
+            user = self.request.user
+        else:
+            user = get_tmp_user_from_session(self.request)
         password = form.cleaned_data.get('password')
-        user = get_user(self.request)
         user = authenticate(username=user.username, password=password)
         if not user:
             form.add_error("password", _("Password invalid"))
@@ -428,8 +439,12 @@ class UserOtpEnableInstallAppView(TemplateView):
     template_name = 'users/user_otp_enable_install_app.html'
 
     def get_context_data(self, **kwargs):
+        if self.request.user.is_authenticated:
+            user = self.request.user
+        else:
+            user = get_tmp_user_from_session(self.request)
         context = {
-            'user': get_user(self.request)
+            'user': user
         }
         kwargs.update(context)
         return super().get_context_data(**kwargs)
@@ -441,16 +456,20 @@ class UserOtpEnableBindView(TemplateView, FormView):
     success_url = reverse_lazy('users:user-otp-settings-success')
 
     def get_context_data(self, **kwargs):
+        if self.request.user.is_authenticated:
+            user = self.request.user
+        else:
+            user = get_tmp_user_from_session(self.request)
         context = {
-            'otp_uri': generate_otp_uri(user=get_user(self.request)),
-            'user': get_user(self.request)
+            'otp_uri': generate_otp_uri(self.request),
+            'user': user
         }
         kwargs.update(context)
         return super().get_context_data(**kwargs)
 
     def form_valid(self, form):
         otp_code = form.cleaned_data.get('otp_code')
-        otp_secret_key = cache.get('otp_secret_key')
+        otp_secret_key = cache.get(self.request.session.session_key+'otp_key', '')
 
         if check_otp_code(otp_secret_key, otp_code):
             self.save_otp(otp_secret_key)
@@ -461,7 +480,10 @@ class UserOtpEnableBindView(TemplateView, FormView):
             return self.form_invalid(form)
 
     def save_otp(self, otp_secret_key):
-        user = get_user(self.request)
+        if self.request.user.is_authenticated:
+            user = self.request.user
+        else:
+            user = get_tmp_user_from_session(self.request)
         user.enable_otp()
         user.otp_secret_key = otp_secret_key
         user.save()
@@ -489,11 +511,8 @@ class UserOtpDisableAuthenticationView(FormView):
 class UserOtpSettingsSuccessView(TemplateView):
     template_name = 'flash_message_standalone.html'
 
-    def get(self, request, *args, **kwargs):
-        response = super().get(request, *args, **kwargs)
-        if is_login(request):
-            auth_logout(request)
-        return response
+    # def get(self, request, *args, **kwargs):
+    #     return super().get(request, *args, **kwargs)
 
     def get_context_data(self, **kwargs):
         title, describe = self.get_title_describe()
@@ -508,7 +527,11 @@ class UserOtpSettingsSuccessView(TemplateView):
         return super().get_context_data(**kwargs)
 
     def get_title_describe(self):
-        user = get_user(self.request)
+        if self.request.user.is_authenticated:
+            user = self.request.user
+            auth_logout(self.request)
+        else:
+            user = get_tmp_user_from_session(self.request)
         title = _('OTP enable success')
         describe = _('OTP enable success, return login page')
         if not user.otp_enabled:

requirements/requirements.txt

@@ -54,6 +54,7 @@ pyasn1==0.4.2
 pycparser==2.18
 pycrypto==2.6.1
 pyldap==2.4.45
+pyotp==2.2.6
 PyNaCl==1.2.1
 python-dateutil==2.6.1
 python-gssapi==0.6.4