Merge pull request #757 from shawnz/shawnz/check-admin-site-otp-required

Fix infinite login redirect on admin site with AdminSiteOTPRequiredMixin

Author: moggers87

tests/test_admin.py

@@ -1,11 +1,15 @@
+from binascii import unhexlify
+
 from django.conf import settings
 from django.shortcuts import resolve_url
 from django.test import TestCase
 from django.test.utils import override_settings
+from django.urls import reverse
+from django_otp.oath import totp
 
 from two_factor.admin import patch_admin, unpatch_admin
 
-from .utils import UserMixin
+from .utils import UserMixin, method_registry
 
 
 @override_settings(ROOT_URLCONF='tests.urls_admin')
@@ -48,15 +52,16 @@ class OTPAdminSiteTest(UserMixin, TestCase):
     def setUp(self):
         super().setUp()
         self.user = self.create_superuser()
-        self.login_user()
 
     def test_otp_admin_without_otp(self):
+        self.login_user()
         response = self.client.get('/otp_admin/', follow=True)
         redirect_to = '%s?next=/otp_admin/' % resolve_url(settings.LOGIN_URL)
         self.assertRedirects(response, redirect_to)
 
     @override_settings(LOGIN_URL='two_factor:login')
     def test_otp_admin_without_otp_named_url(self):
+        self.login_user()
         response = self.client.get('/otp_admin/', follow=True)
         redirect_to = '%s?next=/otp_admin/' % resolve_url(settings.LOGIN_URL)
         self.assertRedirects(response, redirect_to)
@@ -66,3 +71,41 @@ def test_otp_admin_with_otp(self):
         self.login_user()
         response = self.client.get('/otp_admin/')
         self.assertEqual(response.status_code, 200)
+
+    @method_registry(['generator'])
+    def test_otp_admin_login_redirect_without_otp(self):
+        # Open URL that is protected
+        #  assert go to login page
+        #  assert the next param not in the session (but still in url)
+        response = self.client.get('/otp_admin/', follow=True)
+        redirect_to = '%s?next=/otp_admin/' % resolve_url(settings.LOGIN_URL)
+        self.assertRedirects(response, redirect_to)
+        self.assertEqual(self.client.session.get('next'), None)
+
+        # Log in given the last redirect
+        #  assert redirect to setup
+        response = self.client.post(
+            redirect_to,
+            {'auth-username': 'bouke@example.com', 'auth-password': 'secret', 'login_view-current_step': 'auth'},
+        )
+        self.assertRedirects(response, reverse('two_factor:setup'))
+        self.assertEqual(self.client.session.get('next'), '/otp_admin/')
+
+        # Setup the device accordingly
+        #  assert redirect to setup completed
+        #  assert button for redirection to the original page
+        response = self.client.post(reverse('two_factor:setup'), data={'setup_view-current_step': 'welcome'})
+        self.assertContains(response, 'Token:')
+        self.assertContains(response, 'autofocus="autofocus"')
+        self.assertContains(response, 'inputmode="numeric"')
+        self.assertContains(response, 'autocomplete="one-time-code"')
+
+        key = response.context_data['keys'].get('generator')
+        bin_key = unhexlify(key.encode())
+        response = self.client.post(
+            reverse('two_factor:setup'),
+            data={'setup_view-current_step': 'generator', 'generator-token': totp(bin_key)},
+            follow=True,
+        )
+
+        self.assertRedirects(response, '/otp_admin/')

two_factor/views/mixins.py

@@ -4,6 +4,7 @@
 from django.template.response import TemplateResponse
 from django.urls import Resolver404, resolve, reverse
 
+from ..admin import AdminSiteOTPRequiredMixin
 from ..utils import default_device
 
 
@@ -90,4 +91,7 @@ def is_otp_view(cls, view_path):
         return (
             hasattr(next_resolver_match.func, 'view_class') and
             issubclass(next_resolver_match.func.view_class, OTPRequiredMixin)
+        ) or (
+            hasattr(next_resolver_match.func, 'admin_site') and
+            isinstance(next_resolver_match.func.admin_site, AdminSiteOTPRequiredMixin)
         )