jumpserver/apps/users/views/profile/otp.py

# ~*~ coding: utf-8 ~*~

from django.contrib.auth import logout as auth_logout
from django.http.response import HttpResponseRedirect
from django.shortcuts import redirect
from django.urls import reverse
from django.utils.translation import gettext as _
from django.views.generic.base import TemplateView
from django.views.generic.edit import FormView

from authentication.errors import SessionEmptyError
from authentication.mfa import MFAOtp, otp_failed_msg
from authentication.mixins import AuthMixin
from common.permissions import IsValidUser
from common.utils import get_logger, FlashMessageUtil
from common.views.mixins import PermissionsMixin
from ... import forms
from ...utils import (
    generate_otp_uri, check_otp_code,
    get_user_or_pre_auth_user,
)

__all__ = [
    'UserOtpEnableStartView',
    'UserOtpEnableInstallAppView',
    'UserOtpEnableBindView',
    'UserOtpDisableView',
]

logger = get_logger(__name__)


class UserOtpEnableStartView(AuthMixin, TemplateView):
    template_name = 'users/user_otp_check_password.html'

    def get(self, request, *args, **kwargs):
        try:
            self.get_user_from_session()
        except SessionEmptyError:
            url = reverse('authentication:login') + '?_=otp_enable_start'
            return redirect(url)
        return super().get(request, *args, **kwargs)


class UserOtpEnableInstallAppView(TemplateView):
    template_name = 'users/user_otp_enable_install_app.html'

    def get_context_data(self, **kwargs):
        user = get_user_or_pre_auth_user(self.request)
        context = {'user': user}
        kwargs.update(context)
        return super().get_context_data(**kwargs)


class UserOtpEnableBindView(AuthMixin, TemplateView, FormView):
    template_name = 'users/user_otp_enable_bind.html'
    form_class = forms.UserCheckOtpCodeForm

    def get(self, request, *args, **kwargs):
        pre_response = self._pre_check_can_bind()
        if pre_response:
            return pre_response
        return super().get(request, *args, **kwargs)

    def post(self, request, *args, **kwargs):
        pre_response = self._pre_check_can_bind()
        if pre_response:
            return pre_response
        return super().post(request, *args, **kwargs)

    def _pre_check_can_bind(self):
        try:
            user = self.get_user_from_session()
        except Exception as e:
            verify_url = reverse('authentication:user-otp-enable-start') + f'?e={e}'
            return HttpResponseRedirect(verify_url)

        if user.otp_secret_key:
            return self.has_already_bound_message()
        return None

    @staticmethod
    def has_already_bound_message():
        message_data = {
            'title': _('Already bound'),
            'error': _('MFA already bound, disable first, then bound'),
            'interval': 10,
            'redirect_url': reverse('authentication:user-otp-disable'),
        }
        response = FlashMessageUtil.gen_and_redirect_to(message_data)
        return response

    def form_valid(self, form):
        otp_code = form.cleaned_data.get('otp_code')
        otp_secret_key = self.request.session.get('otp_secret_key', '')

        valid = check_otp_code(otp_secret_key, otp_code)
        if not valid:
            form.add_error("otp_code", otp_failed_msg)
            return self.form_invalid(form)

        self.save_otp(otp_secret_key)
        auth_logout(self.request)
        return super().form_valid(form)

    def save_otp(self, otp_secret_key):
        user = get_user_or_pre_auth_user(self.request)
        user.otp_secret_key = otp_secret_key
        user.save(update_fields=['otp_secret_key'])

    def get_success_url(self):
        message_data = {
            'title': _('OTP enable success'),
            'message': _('OTP enable success, return login page'),
            'interval': 5,
            'redirect_url': reverse('authentication:login'),
        }
        url = FlashMessageUtil.gen_message_url(message_data)
        return url

    def get_context_data(self, **kwargs):
        user = get_user_or_pre_auth_user(self.request)
        otp_uri, otp_secret_key = generate_otp_uri(user.username)
        self.request.session['otp_secret_key'] = otp_secret_key
        context = {
            'otp_uri': otp_uri,
            'otp_secret_key': otp_secret_key,
            'user': user
        }
        kwargs.update(context)
        return super().get_context_data(**kwargs)


class UserOtpDisableView(PermissionsMixin, FormView):
    template_name = 'users/user_verify_mfa.html'
    form_class = forms.UserCheckOtpCodeForm
    permission_classes = [IsValidUser]

    def form_valid(self, form):
        user = self.request.user
        otp_code = form.cleaned_data.get('otp_code')
        otp = MFAOtp(user)

        ok, error = otp.check_code(otp_code)
        if not ok:
            form.add_error('otp_code', error)
            return super().form_invalid(form)

        otp.disable()
        auth_logout(self.request)
        return super().form_valid(form)

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        context.update({
            'title': _("Disable OTP")
        })
        return context

    def get_success_url(self):
        message_data = {
            'title': _('OTP disable success'),
            'message': _('OTP disable success, return login page'),
            'interval': 5,
            'redirect_url': reverse('authentication:login'),
        }
        url = FlashMessageUtil.gen_message_url(message_data)
        return url