Merge pull request #500 from MWeesenaar/enforce-setup-device-on-otp-required-page

Enforcing a redirect to setup of otp device when none available for user

Author: dopry

CHANGELOG.md

@@ -1,3 +1,8 @@
+## Unreleased
+
+### Added
+- Enforcing a redirect to setup of otp device when none available for user [#550](https://github.com/jazzband/django-two-factor-auth/pull/500)
+
 ## 1.14.0
 
 ### Added

tests/test_views_login.py

@@ -38,7 +38,7 @@ def test_valid_login(self, mock_signal):
         response = self._post({'auth-username': '[email protected]',
                                'auth-password': 'secret',
                                'login_view-current_step': 'auth'})
-        self.assertRedirects(response, resolve_url(settings.LOGIN_REDIRECT_URL))
+        self.assertRedirects(response, reverse('two_factor:setup'))
 
         # No signal should be fired for non-verified user logins.
         self.assertFalse(mock_signal.called)
@@ -80,7 +80,8 @@ def test_valid_login_with_allowed_external_redirect(self):
             {'auth-username': '[email protected]',
              'auth-password': 'secret',
              'login_view-current_step': 'auth'})
-        self.assertRedirects(response, redirect_url, fetch_redirect_response=False)
+        self.assertEqual(self.client.session.get('next'), redirect_url)
+        self.assertRedirects(response, reverse('two_factor:setup'), fetch_redirect_response=False)
 
     def test_valid_login_with_disallowed_external_redirect(self):
         redirect_url = 'https://test.disallowed-success-url.com'
@@ -90,7 +91,7 @@ def test_valid_login_with_disallowed_external_redirect(self):
             {'auth-username': '[email protected]',
              'auth-password': 'secret',
              'login_view-current_step': 'auth'})
-        self.assertRedirects(response, reverse('two_factor:profile'), fetch_redirect_response=False)
+        self.assertRedirects(response, reverse('two_factor:setup'), fetch_redirect_response=False)
 
     @mock.patch('two_factor.views.core.time')
     def test_valid_login_primary_key_stored(self, mock_time):
@@ -395,12 +396,12 @@ def test_login_different_user_on_existing_session(self, mock_logger):
         response = self._post({'auth-username': '[email protected]',
                                'auth-password': 'secret',
                                'login_view-current_step': 'auth'})
-        self.assertRedirects(response, resolve_url(settings.LOGIN_REDIRECT_URL))
+        self.assertRedirects(response, reverse('two_factor:setup'))
 
         response = self._post({'auth-username': '[email protected]',
                                'auth-password': 'secret',
                                'login_view-current_step': 'auth'})
-        self.assertRedirects(response, resolve_url(settings.LOGIN_REDIRECT_URL))
+        self.assertRedirects(response, reverse('two_factor:setup'))
 
     def test_missing_management_data(self):
         # missing management data
@@ -431,8 +432,7 @@ def test_login_different_user_with_otp_on_existing_session(self):
         response = self._post({'auth-username': '[email protected]',
                                'auth-password': 'secret',
                                'login_view-current_step': 'auth'})
-        self.assertRedirects(response,
-                             resolve_url(settings.LOGIN_REDIRECT_URL))
+        self.assertRedirects(response, reverse('two_factor:setup'))
 
         response = self._post({'auth-username': '[email protected]',
                                'auth-password': 'secret',

tests/test_views_mixins.py

@@ -1,8 +1,12 @@
+from binascii import unhexlify
+
 from django.conf import settings
 from django.shortcuts import resolve_url
 from django.test import TestCase
+from django.urls import reverse
+from django_otp.oath import totp
 
-from .utils import UserMixin
+from .utils import UserMixin, method_registry
 
 
 class OTPRequiredMixinTest(UserMixin, TestCase):
@@ -53,3 +57,48 @@ def test_verified(self):
         self.login_user()
         response = self.client.get('/secure/')
         self.assertEqual(response.status_code, 200)
+
+    @method_registry(['generator'])
+    def test_valid_login_with_redirect_field_name_without_device(self):
+        self.create_user()
+        protected_url = '/secure/'
+
+        # 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(protected_url)
+        redirect_url = "%s?%s%s" % (reverse('two_factor:login'), 'next=', protected_url)
+        self.assertRedirects(response, redirect_url)
+        self.assertEqual(self.client.session.get('next'), None)
+
+        # Log in given the last redirect
+        #  assert redirect to setup
+        response = self.client.post(
+            redirect_url,
+            {'auth-username': '[email protected]',
+             'auth-password': 'secret',
+             'login_view-current_step': 'auth'})
+        self.assertRedirects(response, reverse('two_factor:setup'))
+        self.assertEqual(self.client.session.get('next'), protected_url)
+
+        # 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, protected_url)

two_factor/views/core.py

@@ -182,7 +182,13 @@ def done(self, form_list, **kwargs):
                                     samesite=getattr(settings, 'TWO_FACTOR_REMEMBER_COOKIE_SAMESITE', 'Lax'),
                                     )
 
-        return response
+            return response
+
+        # If the user does not have a device.
+        else:
+            if self.request.GET.get('next'):
+                self.request.session['next'] = self.get_success_url()
+            return redirect('two_factor:setup')
 
     # Copied from django.conrib.auth.views.LoginView (Branch: stable/1.11.x)
     # https://github.com/django/django/blob/58df8aa40fe88f753ba79e091a52f236246260b3/django/contrib/auth/views.py#L63
@@ -394,7 +400,7 @@ def dispatch(self, request, *args, **kwargs):
 
 @class_view_decorator(never_cache)
 @class_view_decorator(login_required)
-class SetupView(IdempotentSessionWizardView):
+class SetupView(RedirectURLMixin, IdempotentSessionWizardView):
     """
     View for handling OTP setup using a wizard.
 
@@ -417,6 +423,27 @@ class SetupView(IdempotentSessionWizardView):
         # Other forms are dynamically added in get_form_list()
     )
 
+    # Copied from django.conrib.auth.views.LoginView (Branch: stable/1.11.x)
+    # https://github.com/django/django/blob/58df8aa40fe88f753ba79e091a52f236246260b3/django/contrib/auth/views.py#L63
+    def get_success_url(self):
+        url = self.get_redirect_url()
+        return url or reverse('two_factor:setup_complete')
+
+    # Copied from django.conrib.auth.views.LoginView (Branch: stable/1.11.x)
+    # https://github.com/django/django/blob/58df8aa40fe88f753ba79e091a52f236246260b3/django/contrib/auth/views.py#L67
+    def get_redirect_url(self):
+        """Return the user-originating redirect URL if it's safe."""
+        redirect_to = self.request.POST.get(
+            REDIRECT_FIELD_NAME,
+            self.request.GET.get(REDIRECT_FIELD_NAME, '')
+        )
+        url_is_safe = url_has_allowed_host_and_scheme(
+            url=redirect_to,
+            allowed_hosts=self.get_success_url_allowed_hosts(),
+            require_https=self.request.is_secure(),
+        )
+        return redirect_to if url_is_safe else ''
+
     def get_method(self):
         method_data = self.storage.validated_step_data.get('method', {})
         method_key = method_data.get('method', None)
@@ -427,7 +454,7 @@ def get(self, request, *args, **kwargs):
         Start the setup wizard. Redirect if already enabled.
         """
         if default_device(self.request.user):
-            return redirect(self.success_url)
+            return redirect(self.get_success_url())
         return super().get(request, *args, **kwargs)
 
     def get_form(self, step=None, **kwargs):
@@ -495,7 +522,7 @@ def done(self, form_list, **kwargs):
             raise NotImplementedError("Unknown method '%s'" % method.code)
 
         django_otp.login(self.request, device)
-        return redirect(self.success_url)
+        return redirect(self.get_success_url())
 
     def get_form_kwargs(self, step=None):
         kwargs = {}
@@ -607,6 +634,11 @@ class SetupCompleteView(TemplateView):
     """
     template_name = 'two_factor/core/setup_complete.html'
 
+    def get(self, request, *args, **kwargs):
+        if request.session.get('next'):
+            return redirect(request.session.get('next'))
+        return super().get(request, *args, **kwargs)
+
     def get_context_data(self):
         return {
             'phone_methods': get_available_phone_methods(),