diff --git a/tests/test_admin.py b/tests/test_admin.py index 8acb32d27..e6344f1cd 100644 --- a/tests/test_admin.py +++ b/tests/test_admin.py @@ -4,6 +4,7 @@ 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 @@ -52,26 +53,32 @@ class OTPAdminSiteTest(UserMixin, TestCase): def setUp(self): super().setUp() - self.user = self.create_superuser() - 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 + admins without MFA setup should be redirected to the setup page. """ + self.user = self.create_superuser() + self.login_user() + print("user", self.user.is_active, self.user.is_staff) response = self.client.get('/otp_admin/', follow=True) 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): + self.user = self.create_superuser() + self.login_user() + print("user", self.user.is_active, self.user.is_staff) response = self.client.get('/otp_admin/', follow=True) redirect_to = reverse('two_factor:setup') self.assertRedirects(response, redirect_to) def test_otp_admin_with_otp(self): + self.user = self.create_superuser() self.enable_otp() self.login_user() + print("user", self.user.is_active, self.user.is_staff) response = self.client.get('/otp_admin/') self.assertEqual(response.status_code, 200) diff --git a/two_factor/admin.py b/two_factor/admin.py index 2e3ce4f88..7e479f634 100644 --- a/two_factor/admin.py +++ b/two_factor/admin.py @@ -1,3 +1,4 @@ +from functools import update_wrapper from django.conf import settings from django.contrib.admin import AdminSite @@ -6,8 +7,11 @@ 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 @@ -25,33 +29,92 @@ 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() + print("AdminSiteOTPRequiredMixin.has_permission, self.has_admin_permission(request)", self.has_admin_permission(request), request.user.is_verified()) + return self.has_admin_permission(request) and request.user.is_verified() + + def admin_view(self, view, cacheable=False): + """ + Decorator to create an admin view attached to this ``AdminSite``. This + wraps the view and provides permission checking by calling + ``self.has_permission``. + + You'll want to use this from within ``AdminSite.get_urls()``: + 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): + print("AdminSiteOTPRequiredMixin.admin_view.inner", ) + 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) + + @never_cache 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. """ + print("AdminSiteOTPRequiredMixin.login") + # redirect to admin page after login redirect_to = request.POST.get(REDIRECT_FIELD_NAME, request.GET.get(REDIRECT_FIELD_NAME)) + has_admin_access = AdminSite().has_permission(request) + print("AdminSiteOTPRequiredMixin.login", redirect_to, request.method, has_admin_access) # if user (is_active and is_staff) - if request.method == "GET" and AdminSite().has_permission(request): + if request.method == "GET" and has_admin_access: # if user has 2FA setup, go to admin homepage if request.user.is_verified(): + print("User is verified, going to normal index.") index_path = reverse("admin:index", current_app=self.name) # 2FA not setup. redirect to 2FA setup page else: + print("User is not verified. redirecting to two_factor setup.") index_path = reverse("two_factor:setup", current_app=self.name) return HttpResponseRedirect(index_path)