From 93afc9c6662d10049476e1187e8166452d0611bf Mon Sep 17 00:00:00 2001
From: Markus Holtermann <info@markusholtermann.eu>
Date: Thu, 28 Jul 2016 01:55:23 +0200
Subject: [PATCH] Reworked admin auto-patching

---
 docs/configuration.rst |  3 +++
 two_factor/admin.py    | 41 +++++++++++++++++++++++------------------
 two_factor/utils.py    |  8 --------
 3 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/docs/configuration.rst b/docs/configuration.rst
index 825ab3c58..45d8c522f 100644
--- a/docs/configuration.rst
+++ b/docs/configuration.rst
@@ -11,6 +11,9 @@ General Settings
      The admin currently does not enforce one-time passwords being set for
      admin users.
 
+``TWO_FACTOR_FORCE_OTP_ADMIN`` (default: ``False``)
+  Whether the Django admin will enforce 2 factor authentication.
+
 ``TWO_FACTOR_CALL_GATEWAY`` (default: ``None``)
   Which gateway to use for making phone calls. Should be set to a module or
   object providing a ``make_call`` method. Currently two gateways are bundled:
diff --git a/two_factor/admin.py b/two_factor/admin.py
index 2b7059b72..2d353f603 100644
--- a/two_factor/admin.py
+++ b/two_factor/admin.py
@@ -4,14 +4,12 @@
 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.core.urlresolvers import reverse
 from django.shortcuts import resolve_url
 from django.utils.http import is_safe_url
 from django.utils.translation import ugettext
 
 from .models import PhoneDevice
-from .utils import monkeypatch_method
 from .views import BackupTokensView, LoginView, ProfileView, SetupView
 
 
@@ -130,6 +128,9 @@ def has_permission(self, request):
             return False
         return request.user.is_verified()
 
+
+class AdminSiteOTPMixin(object):
+
     def get_urls(self):
         from django.conf.urls import include, url
 
@@ -148,7 +149,7 @@ def wrapper(*args, **kwargs):
         urlpatterns = [
             url(r'^two_factor/', include(urlpatterns_2fa, namespace='two_factor'))
         ]
-        urlpatterns += super(AdminSiteOTPRequiredMixin, self).get_urls()
+        urlpatterns += super(AdminSiteOTPMixin, self).get_urls()
         return urlpatterns
 
     def login(self, request, extra_context=None):
@@ -164,32 +165,36 @@ def two_factor_backup_tokens(self, request):
         return admin_backup_tokens_view(request)
 
 
-class AdminSiteOTPRequired(AdminSiteOTPRequiredMixin, AdminSite):
+class AdminSiteOTP(AdminSiteOTPMixin, AdminSite):
     """
-    AdminSite enforcing OTP verified staff users.
+    AdminSite using OTP login.
     """
     pass
 
 
-def patch_admin():
-    @monkeypatch_method(AdminSite)
-    def login(self, request, extra_context=None):
-        """
-        Redirects to the site login page for the given HttpRequest.
-        """
-        redirect_to = request.POST.get(REDIRECT_FIELD_NAME, request.GET.get(REDIRECT_FIELD_NAME))
+class AdminSiteOTPRequired(AdminSiteOTPMixin, AdminSiteOTPRequiredMixin, AdminSite):
+    """
+    AdminSite enforcing OTP verified staff users.
+    """
+    pass
 
-        if not redirect_to or not is_safe_url(url=redirect_to, host=request.get_host()):
-            redirect_to = resolve_url(settings.LOGIN_REDIRECT_URL)
 
-        return redirect_to_login(redirect_to)
+__default_admin_site__ = None
 
 
-def unpatch_admin():
-    setattr(AdminSite, 'login', original_login)
+def patch_admin():
+    global __default_admin_site__
+    __default_admin_site__ = admin.site.__class__
+    if getattr(settings, 'TWO_FACTOR_FORCE_OTP_ADMIN', False):
+        admin.site.__class__ = AdminSiteOTPRequired
+    else:
+        admin.site.__class__ = AdminSiteOTP
 
 
-original_login = AdminSite.login
+def unpatch_admin():
+    global __default_admin_site__
+    admin.site.__class__ = __default_admin_site__
+    __default_admin_site__ = None
 
 
 class PhoneDeviceAdmin(admin.ModelAdmin):
diff --git a/two_factor/utils.py b/two_factor/utils.py
index aad926c01..8d8660408 100644
--- a/two_factor/utils.py
+++ b/two_factor/utils.py
@@ -47,14 +47,6 @@ def get_otpauth_url(accountname, secret, issuer=None, digits=None):
     return 'otpauth://totp/%s?%s' % (label, urlencode(query))
 
 
-# from http://mail.python.org/pipermail/python-dev/2008-January/076194.html
-def monkeypatch_method(cls):
-    def decorator(func):
-        setattr(cls, func.__name__, func)
-        return func
-    return decorator
-
-
 def totp_digits():
     """
     Returns the number of digits (as configured by the TWO_FACTOR_TOTP_DIGITS setting)