academon/core/views/auth.py

from django.contrib.auth.views import LoginView as DjangoLoginView
from django.contrib.auth import login
from django.utils.decorators import method_decorator
from django.views.decorators.debug import sensitive_post_parameters
from django.views.decorators.csrf import csrf_protect
from django.views.decorators.cache import never_cache
from django.core.exceptions import PermissionDenied
from django.http.response import HttpResponseRedirect
from django.urls import reverse
from django.contrib import messages

from ..models.auth import OTPSession, Profile
from ..forms.auth import OTPForm


class LoginView(DjangoLoginView):
    template_name = "core/auth/login.html"

    @method_decorator(sensitive_post_parameters())
    @method_decorator(csrf_protect)
    @method_decorator(never_cache)
    def dispatch(self, request, *args, **kwargs):
        if self.request.session.get("OTPSession"):
            return HttpResponseRedirect(self.get_otp_url())
        return super().dispatch(request, *args, **kwargs)

    def get_otp_url(self):
        return reverse("checkotp", kwargs={self.redirect_field_name: self.get_redirect_url()})

    def form_valid(self, form):
        if not (user := form.get_user()).totp:
            login(self.request, user)
            profile, _ = Profile.objects.get_or_create(user=user)
            profile.save()
            return HttpResponseRedirect(self.get_success_url())

        session = OTPSession.objects.create(user=user)
        self.request.session["OTPSession"] = session.uuid

        return HttpResponseRedirect(self.get_otp_url())

    def form_invalid(self, form):
        messages.error(self.request, "The username or password you entered was incorrect. Please try again.")
        return super().form_invalid(form)


class OTPView(DjangoLoginView):
    authentication_form = OTPForm

    @method_decorator(sensitive_post_parameters())
    @method_decorator(csrf_protect)
    @method_decorator(never_cache)
    def dispatch(self, request, *args, **kwargs):
        if not self.get_otpsession().is_alive():
            del request.session["OTPSession"]
            raise PermissionDenied("Your session has expired. Please try again.")
        return super().dispatch(request, *args, **kwargs)

    def get_otpsession(self):
        return OTPSession.objects.get(uuid=self.request.session.get("OTPSession"))

    def form_valid(self, form):
        login(self.request, self.get_otpsession().user)
        del request.session["OTPSession"]
        profile, _ = Profile.objects.get_or_create(user=user)
        profile.save()
        return HttpResponseRedirect(self.get_success_url())

    def form_invalid(self, form):
        messages.error(self.request, "The One-Time Password you entered was incorrect. Please try again.")
        return super().form_invalid(form)