Finish with example

Author: RadoRado

styleguide_example/blog_examples/admin_2fa/models.py

@@ -1,3 +1,4 @@
+import uuid
 from typing import Optional
 
 import pyotp
@@ -11,6 +12,7 @@ class UserTwoFactorAuthData(models.Model):
     user = models.OneToOneField(settings.AUTH_USER_MODEL, related_name="two_factor_auth_data", on_delete=models.CASCADE)
 
     otp_secret = models.CharField(max_length=255)
+    session_identifier = models.UUIDField(blank=True, null=True)
 
     def generate_qr_code(self, name: Optional[str] = None) -> str:
         totp = pyotp.TOTP(self.otp_secret)
@@ -21,3 +23,13 @@ def generate_qr_code(self, name: Optional[str] = None) -> str:
 
         # The result is going to be an HTML <svg> tag
         return qr_code_image.to_string().decode("utf_8")
+
+    def validate_otp(self, otp: str) -> bool:
+        totp = pyotp.TOTP(self.otp_secret)
+
+        return totp.verify(otp)
+
+    def rotate_session_identifier(self):
+        self.session_identifier = uuid.uuid4()
+
+        self.save(update_fields=["session_identifier"])

styleguide_example/blog_examples/admin_2fa/views.py

@@ -1,6 +1,9 @@
+from django import forms
 from django.core.exceptions import ValidationError
-from django.views.generic import TemplateView
+from django.urls import reverse_lazy
+from django.views.generic import FormView, TemplateView
 
+from .models import UserTwoFactorAuthData
 from .services import user_two_factor_auth_data_create
 
 
@@ -21,3 +24,41 @@ def post(self, request):
             context["form_errors"] = exc.messages
 
         return self.render_to_response(context)
+
+
+class AdminConfirmTwoFactorAuthView(FormView):
+    template_name = "admin_2fa/confirm_2fa.html"
+    success_url = reverse_lazy("admin:index")
+
+    class Form(forms.Form):
+        otp = forms.CharField(required=True)
+
+        def clean_otp(self):
+            self.two_factor_auth_data = UserTwoFactorAuthData.objects.filter(user=self.user).first()
+
+            if self.two_factor_auth_data is None:
+                raise ValidationError("2FA not set up.")
+
+            otp = self.cleaned_data.get("otp")
+
+            if not self.two_factor_auth_data.validate_otp(otp):
+                raise ValidationError("Invalid 2FA code.")
+
+            return otp
+
+    def get_form_class(self):
+        return self.Form
+
+    def get_form(self, *args, **kwargs):
+        form = super().get_form(*args, **kwargs)
+
+        form.user = self.request.user
+
+        return form
+
+    def form_valid(self, form):
+        form.two_factor_auth_data.rotate_session_identifier()
+
+        self.request.session["2fa_token"] = str(form.two_factor_auth_data.session_identifier)
+
+        return super().form_valid(form)

styleguide_example/blog_examples/migrations/0003_usertwofactorauthdata.py

@@ -1,4 +1,4 @@
-# Generated by Django 4.1.9 on 2023-07-04 12:17
+# Generated by Django 4.1.9 on 2023-07-05 08:49
 
 from django.conf import settings
 from django.db import migrations, models
@@ -18,6 +18,7 @@ class Migration(migrations.Migration):
             fields=[
                 ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
                 ('otp_secret', models.CharField(max_length=255)),
+                ('session_identifier', models.UUIDField(blank=True, null=True)),
                 ('user', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, related_name='two_factor_auth_data', to=settings.AUTH_USER_MODEL)),
             ],
         ),

styleguide_example/blog_examples/templates/admin_2fa/confirm_2fa.html

@@ -0,0 +1,24 @@
+{% extends "admin/login.html" %}
+
+{% block content %}
+  {% if form.non_field_errors %}
+    {% for error in form.non_field_errors %}
+      <p class="errornote">
+	{{ error }}
+      </p>
+    {% endfor %}
+  {% endif %}
+
+  <form action="" method="post">
+    {% csrf_token %}
+
+    <div class="form-row">
+      {{ form.otp.errors }}
+      {{ form.otp.label_tag }} {{ form.otp }}
+    </div>
+
+    <div class="submit-row">
+      <input type="submit" value="Submit">
+    </div>
+  </form>
+{% endblock %}

styleguide_example/custom_admin/sites.py

@@ -1,15 +1,59 @@
 from django.contrib import admin
-from django.urls import path
+from django.contrib.auth import REDIRECT_FIELD_NAME
+from django.urls import path, reverse
 
-from styleguide_example.blog_examples.admin_2fa.views import AdminSetupTwoFactorAuthView
+from styleguide_example.blog_examples.admin_2fa.models import UserTwoFactorAuthData
+from styleguide_example.blog_examples.admin_2fa.views import (
+    AdminConfirmTwoFactorAuthView,
+    AdminSetupTwoFactorAuthView,
+)
 
 
 class AdminSite(admin.AdminSite):
     def get_urls(self):
         base_urlpatterns = super().get_urls()
 
         extra_urlpatterns = [
-            path("setup-2fa/", self.admin_view(AdminSetupTwoFactorAuthView.as_view()), name="setup-2fa")
+            path("setup-2fa/", self.admin_view(AdminSetupTwoFactorAuthView.as_view()), name="setup-2fa"),
+            path("confirm-2fa/", self.admin_view(AdminConfirmTwoFactorAuthView.as_view()), name="confirm-2fa"),
         ]
 
         return extra_urlpatterns + base_urlpatterns
+
+    def login(self, request, *args, **kwargs):
+        if request.method != "POST":
+            return super().login(request, *args, **kwargs)
+
+        username = request.POST.get("username")
+
+        two_factor_auth_data = UserTwoFactorAuthData.objects.filter(user__email=username).first()
+
+        request.POST._mutable = True
+        request.POST[REDIRECT_FIELD_NAME] = reverse("admin:confirm-2fa")
+
+        if two_factor_auth_data is None:
+            request.POST[REDIRECT_FIELD_NAME] = reverse("admin:setup-2fa")
+
+        request.POST._mutable = False
+
+        return super().login(request, *args, **kwargs)
+
+    def has_permission(self, request):
+        has_perm = super().has_permission(request)
+
+        if not has_perm:
+            return has_perm
+
+        two_factor_auth_data = UserTwoFactorAuthData.objects.filter(user=request.user).first()
+
+        allowed_paths = [reverse("admin:confirm-2fa"), reverse("admin:setup-2fa")]
+
+        if request.path in allowed_paths:
+            return True
+
+        if two_factor_auth_data is not None:
+            two_factor_auth_token = request.session.get("2fa_token")
+
+            return str(two_factor_auth_data.session_identifier) == two_factor_auth_token
+
+        return False