feat: Redirect admin users to setup TOTP

When TOTP is required on an admin view and a user does not have a
TOTP device configured, redirect them to the TOTP setup view.

Author: aseem-hegshetye

.gitignore

@@ -10,3 +10,5 @@ example/settings_private.py
 .eggs/
 
 .idea/
+
+venv/

example/urls.py

@@ -5,6 +5,7 @@
 
 from two_factor.gateways.twilio.urls import urlpatterns as tf_twilio_urls
 from two_factor.urls import urlpatterns as tf_urls
+from two_factor.admin import AdminSiteOTPRequired
 
 from .views import (
     ExampleSecretView, HomeView, RegistrationCompleteView, RegistrationView,
@@ -40,6 +41,7 @@
     path('', include(tf_twilio_urls)),
     path('', include('user_sessions.urls', 'user_sessions')),
     path('admin/', admin.site.urls),
+    path('otp_admin/', AdminSiteOTPRequired().urls),
 ]
 
 if settings.DEBUG:

tests/test_admin.py

@@ -1,9 +1,10 @@
 from django.conf import settings
-from django.shortcuts import resolve_url
+from django.shortcuts import resolve_url, reverse
 from django.test import TestCase
 from django.test.utils import override_settings
 
 from two_factor.admin import patch_admin, unpatch_admin
+from two_factor.utils import default_device
 
 from .utils import UserMixin
 
@@ -44,21 +45,29 @@ def test_default_admin(self):
 
 @override_settings(ROOT_URLCONF='tests.urls_otp_admin')
 class OTPAdminSiteTest(UserMixin, TestCase):
+    """
+    otp_admin is admin console that needs OTP for access.
+    Only admin users (is_staff and is_active)
+        with OTP can access it.
+    """
 
     def setUp(self):
         super().setUp()
         self.user = self.create_superuser()
         self.login_user()
-
+       
     def test_otp_admin_without_otp(self):
+        """
+        admins without MFA setup should be redirected to the setup page.
+        """
         response = self.client.get('/otp_admin/', follow=True)
-        redirect_to = '%s?next=/otp_admin/' % resolve_url(settings.LOGIN_URL)
+        redirect_to = reverse('two_factor:setup')
         self.assertRedirects(response, redirect_to)
 
     @override_settings(LOGIN_URL='two_factor:login')
     def test_otp_admin_without_otp_named_url(self):
         response = self.client.get('/otp_admin/', follow=True)
-        redirect_to = '%s?next=/otp_admin/' % resolve_url(settings.LOGIN_URL)
+        redirect_to = reverse('two_factor:setup')
         self.assertRedirects(response, redirect_to)
 
     def test_otp_admin_with_otp(self):

two_factor/admin.py

@@ -1,10 +1,17 @@
+from functools import update_wrapper
+
 from django.conf import settings
 from django.contrib.admin import AdminSite
 from django.contrib.auth import REDIRECT_FIELD_NAME
 from django.contrib.auth.views import redirect_to_login
+from django.http import HttpResponseRedirect
 from django.shortcuts import resolve_url
+from django.urls import reverse
+from django.views.decorators.cache import never_cache
+from django.views.decorators.csrf import csrf_protect
+
 
-from .utils import monkeypatch_method
+from .utils import default_device, monkeypatch_method
 
 try:
     from django.utils.http import url_has_allowed_host_and_scheme
@@ -22,25 +29,64 @@ class AdminSiteOTPRequiredMixin:
     use :meth:`has_permission` in order to secure those views.
     """
 
+    def has_admin_permission(self, request):
+        return super().has_permission(request)
+
     def has_permission(self, request):
         """
         Returns True if the given HttpRequest has permission to view
         *at least one* page in the admin site.
         """
-        if not super().has_permission(request):
-            return False
-        return request.user.is_verified()
+        return self.has_admin_permission(request) and request.user.is_verified()
 
-    def login(self, request, extra_context=None):
+    def admin_view(self, view, cacheable=False):
         """
-        Redirects to the site login page for the given HttpRequest.
-        """
-        redirect_to = request.POST.get(REDIRECT_FIELD_NAME, request.GET.get(REDIRECT_FIELD_NAME))
+        Decorator to create an admin view attached to this ``AdminSite``. This
+        wraps the view and provides permission checking by calling
+        ``self.has_permission``.
 
-        if not redirect_to or not url_has_allowed_host_and_scheme(url=redirect_to, allowed_hosts=[request.get_host()]):
-            redirect_to = resolve_url(settings.LOGIN_REDIRECT_URL)
+        You'll want to use this from within ``AdminSite.get_urls()``:
 
-        return redirect_to_login(redirect_to)
+            class MyAdminSite(AdminSite):
+
+                def get_urls(self):
+                    from django.urls import path
+
+                    urls = super().get_urls()
+                    urls += [
+                        path('my_view/', self.admin_view(some_view))
+                    ]
+                    return urls
+
+        By default, admin_views are marked non-cacheable using the
+        ``never_cache`` decorator. If the view can be safely cached, set
+        cacheable=True.
+        """
+        def inner(request, *args, **kwargs):
+            if not self.has_permission(request):
+                if request.path == reverse('admin:logout', current_app=self.name):
+                    index_path = reverse('admin:index', current_app=self.name)
+                    return HttpResponseRedirect(index_path)
+
+                if (self.has_admin_permission(request) and not default_device(request.user)):
+                    index_path = reverse("two_factor:setup", current_app=self.name)
+                    return HttpResponseRedirect(index_path)
+
+                # Inner import to prevent django.contrib.admin (app) from
+                # importing django.contrib.auth.models.User (unrelated model).
+                from django.contrib.auth.views import redirect_to_login
+                return redirect_to_login(
+                    request.get_full_path(),
+                    reverse('admin:login', current_app=self.name)
+                )
+            return view(request, *args, **kwargs)
+        if not cacheable:
+            inner = never_cache(inner)
+        # We add csrf_protect here so this function can be used as a utility
+        # function for any view, without having to repeat 'csrf_protect'.
+        if not getattr(view, 'csrf_exempt', False):
+            inner = csrf_protect(inner)
+        return update_wrapper(inner, view)
 
 
 class AdminSiteOTPRequired(AdminSiteOTPRequiredMixin, AdminSite):