expephalon-smsotp/core/views/auth.py

from django.conf import settings
from django.views.generic import FormView, View
from django.contrib.auth import authenticate, login, logout
from django.shortcuts import redirect
from django.core.exceptions import PermissionDenied
from django.contrib import messages

from core.forms import LoginForm, OTPSelectorForm, OTPVerificationForm
from core.models.otp import LoginSession
from core.helpers.otp import get_user_otps, get_otp_choices, get_otp_by_name

class LoginView(FormView):
    template_name = f"{settings.EXPEPHALON_BACKEND}/auth/login.html"
    form_class = LoginForm

    def get(self, request, *args, **kwargs):
        if request.user.is_authenticated:
            return redirect(request.GET.get("next", "dashboard"))
        return super().get(request, *args, **kwargs)

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        context["title"] = "Login"
        return context

    def form_valid(self, form):
        user = authenticate(username=form.cleaned_data['email'],password=form.cleaned_data['password'])
        if user:
            if not get_user_otps(user):
                login(self.request, user)
                return redirect("dashboard")
            session = LoginSession.objects.create(user=user)
            self.request.session["otpsession"] = str(session.uuid)
            self.request.session["next"] = self.request.GET.get("next", "dashboard")
            return redirect("otpselector")
        return super().form_invalid(form)

class OTPSelectorView(FormView):
    template_name = f"{settings.EXPEPHALON_BACKEND}/auth/otp_selector.html"
    form_class = OTPSelectorForm

    def clean_session(self):
        for key in ("otpsession", "otpprovider", "next"):
            try:
                del self.request.session[key]
            except:
                pass

    def get_form_kwargs(self):
        kwargs = super().get_form_kwargs()
        try:
            assert self.request.session["otpsession"]
        except:
            raise PermissionDenied()
        user = LoginSession.objects.get(uuid=self.request.session["otpsession"]).user
        kwargs["otp_choices"] = get_otp_choices(user)
        return kwargs

    def form_valid(self, form):
        self.request.session["otpprovider"] = form.cleaned_data["provider"]
        return redirect("otpvalidator")
    
    def form_invalid(self, form):
        self.clean_session()
        return redirect("login")

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        context["title"] = "Two-Factor Authentication"
        return context

class OTPValidatorView(FormView):
    template_name = f"{settings.EXPEPHALON_BACKEND}/auth/otp_verifier.html"
    form_class = OTPVerificationForm

    def clean_session(self):
        for key in ("otpsession", "otpprovider", "next"):
            try:
                del self.request.session[key]
            except:
                pass

    def validate_session(self, request):
        try:
            assert request.session["otpsession"]
            assert request.session["otpprovider"]
            user = LoginSession.objects.get(uuid=request.session["otpsession"]).user
            assert request.session["otpprovider"] in get_user_otps(user).keys()
            provider = get_otp_by_name(request.session["otpprovider"])()
            return user, provider
        except:
            self.clean_session()
            raise PermissionDenied()

    def get(self, request, *args, **kwargs):
        user, provider = self.validate_session(request)
        response = provider.start_authentication(user)
        messages.info(request, response)
        return super().get(request, *args, **kwargs)

    def post(self, request, *args, **kwargs):
        self.validate_session(request)
        return super().post(request, *args, **kwargs)

    def form_invalid(self, form):
        self.clean_session()
        return redirect("login")

    def form_valid(self, form):
        user, provider = self.validate_session(self.request)
        if provider.validate_token(user, form.cleaned_data["token"]):
            login(self.request, user)
            ret = redirect(self.request.session.get("next", "dashboard"))
            self.clean_session()
            return ret
        self.clean_session()
        messages.error(self.request, "Incorrect token entered. Please try again. If the issue persists, contact support to regain access to your account.")
        return redirect("login")

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        context["title"] = "Two-Factor Authentication"
        return context

class LogoutView(View):
    def get(self, request, *args, **kwargs):
        logout(request)
        return redirect("login")