diff --git a/.gitignore b/.gitignore index 7cd853fbb..80de2fd45 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,5 @@ example/settings_private.py .eggs/ .idea/ + +venv/ diff --git a/example/urls.py b/example/urls.py index 8e7531f09..c457eb1ee 100644 --- a/example/urls.py +++ b/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: diff --git a/tests/test_admin.py b/tests/test_admin.py index 57d11225f..c34d97636 100644 --- a/tests/test_admin.py +++ b/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): diff --git a/two_factor/admin.py b/two_factor/admin.py index 7d9b29513..08351e814 100644 --- a/two_factor/admin.py +++ b/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):