diff --git a/tests/test_admin.py b/tests/test_admin.py index 57d11225f..4d44d71c3 100644 --- a/tests/test_admin.py +++ b/tests/test_admin.py @@ -1,10 +1,9 @@ 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 .utils import UserMixin @@ -44,6 +43,11 @@ 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() @@ -51,14 +55,18 @@ def setUp(self): self.login_user() def test_otp_admin_without_otp(self): + """ + if user has admin permissions (is_staff and is_active) + but doesnt have OTP setup, redirect the user to OTP 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): diff --git a/two_factor/admin.py b/two_factor/admin.py index 9deb880c4..201229095 100644 --- a/two_factor/admin.py +++ b/two_factor/admin.py @@ -1,9 +1,12 @@ + from django.conf import settings from django.contrib import admin 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.utils.http import is_safe_url from .models import PhoneDevice @@ -30,9 +33,25 @@ def has_permission(self, request): def login(self, request, extra_context=None): """ Redirects to the site login page for the given HttpRequest. + If user has admin permissions but 2FA not setup, then redirect to + 2FA setup page. """ + # redirect to admin page after login redirect_to = request.POST.get(REDIRECT_FIELD_NAME, request.GET.get(REDIRECT_FIELD_NAME)) + # if user (is_active and is_staff) + if request.method == "GET" and AdminSite().has_permission(request): + + # if user has 2FA setup, go to admin homepage + if request.user.is_verified(): + index_path = reverse("admin:index", current_app=self.name) + + # 2FA not setup. redirect to 2FA setup page + else: + index_path = reverse("two_factor:setup", current_app=self.name) + + return HttpResponseRedirect(index_path) + if not redirect_to or not is_safe_url(url=redirect_to, allowed_hosts=[request.get_host()]): redirect_to = resolve_url(settings.LOGIN_REDIRECT_URL)