expephalon-demomodule/core/views/auth.py

from django.conf import settings
from django.views.generic import FormView, View
from django.contrib.auth import authenticate, login, logout, get_user_model
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, PWRequestForm
from core.models.auth import LoginSession, PWResetToken
from core.helpers.otp import get_user_otps, get_otp_choices, get_otp_by_name
from core.helpers.mail import simple_send_mail
from core.helpers.auth import generate_pwreset_mail

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")
        messages.error(self.request, "The credentials you entered are invalid. Please try again.")
        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"

        user = LoginSession.objects.get(uuid=self.request.session["otpsession"]).user
        context["first_name"] = user.profile.get_internal_name
        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"

        user, provider = self.validate_session(self.request)
        context["first_name"] = user.profile.get_internal_name
        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")

class PWRequestView(FormView):
    template_name = f"{settings.EXPEPHALON_BACKEND}/auth/pwrequest.html"
    form_class = PWRequestForm

    def form_valid(self, form):
        try:
            user = get_user_model().objects.get(username=form.cleaned_data["email"])
            token = PWResetToken.objects.create(user=user)
            mail = generate_pwreset_mail(user, token)
            simple_send_mail("Password Reset", mail, user.email)
        finally:
            messages.success(self.request, "If a matching account was found, you should shortly receive an email containing password reset instructions. If you have not received this message after five minutes, please verify that you have entered the correct email address, or contact support.")
            return redirect("login")