WIP

Author: dopry

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)

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)