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 django.utils import timezone

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

from dbsettings.functions import getValue

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")

class PWResetView(FormView):
    template_name = f"{settings.EXPEPHALON_BACKEND}/auth/pwreset.html"
    form_class = PWResetForm

    def validate_session(self):
        try:
            token = PWResetToken.objects.get(token=self.kwargs["pk"])
            max_age = int(getValue("core.auth.pwreset.max_age", "86400"))
            assert token.creation > timezone.now() - timezone.timedelta(seconds=max_age)
            return token.user
        except:
            messages.error(self.request, "Incorrect or expired password reset link.")
            raise PermissionDenied()

    def form_valid(self, form):
        user = self.validate_session()
        if not form.cleaned_data["password1"] == form.cleaned_data["password2"]:
            messages.error(self.request, "Entered passwords do not match - please try again.")
            return self.form_invalid(form)
        user.set_password(form.cleaned_data["password1"])
        user.save()
        messages.success(self.request, "Your password has been changed. You can now login with your new password.")
        return redirect("login")