# HG changeset patch
# User Gustavo Andres Morero <gustavo@netlandish.com>
# Date 1597178969 10800
#      Tue Aug 11 17:49:29 2020 -0300
# Node ID f8f3c1bc0e57c80fb29e2941ef39a395bad717e0
# Parent  cf9ade31a69a477bf35e572a909378464c22ec20
adding ability to require 2FA setup.
moving urls and excluded urls list to settings.
single otp check middleware required.
updates on qr view.
clear device session key when disabling 2FA.

diff --git a/nlotp/checks.py b/nlotp/checks.py
--- a/nlotp/checks.py
+++ b/nlotp/checks.py
@@ -26,30 +26,31 @@
 
 REQUIRED_SETTINGS = (
     {
-        "key": "NLOTP_EXCLUDED_URLS",
+        "key": "NLOTP_VERIFY_EXCLUDED_URLS",
         "code": "E004",
         "hint": (
-            "A list of urls excluded from redirect if user "
-            "is not OTP verified."
+            "A list of urls excluded from redirect to verify page "
+            "if user is not OTP verified."
+        ),
+    },
+    {
+        "key": "NLOTP_SETUP_EXCLUDED_URLS",
+        "code": "E005",
+        "hint": (
+            "A list of urls excluded from redirect to setup page "
+            "if user didn't setup 2FA and it's required "
+            "(NLOTP_2FA_SETUP_REQUIRED = True)."
         ),
     },
 )
 
 REQUIRED_MIDDLEWARES = (
     {
-        "key": "django_otp.middleware.OTPMiddleware",
-        "code": "E005",
-        "hint": (
-            "Must be installed after "
-            "`django.contrib.auth.middleware.AuthenticationMiddleware` "
-            "and before `nlotp.middleware.OTPCheckMiddleware`."
-        ),
-    },
-    {
         "key": "nlotp.middleware.OTPCheckMiddleware",
         "code": "E006",
         "hint": (
-            "Must be installed after `django_otp.middleware.OTPMiddleware`."
+            "Must be installed after "
+            "`django.contrib.auth.middleware.AuthenticationMiddleware`."
         ),
     },
 )
diff --git a/nlotp/middleware.py b/nlotp/middleware.py
--- a/nlotp/middleware.py
+++ b/nlotp/middleware.py
@@ -1,23 +1,33 @@
-from django.conf import settings
-from django.http import HttpResponseRedirect
-from django.urls import reverse_lazy
-from django.utils.deprecation import MiddlewareMixin
-from django_otp import user_has_device
+import functools
 
-NLOTP_LOGIN_URL = getattr(
-    settings,
-    "NLOTP_LOGIN_URL",
-    reverse_lazy("nlotp:verify-otp"),
-)
+from django.http import HttpResponseRedirect
+from django.utils.deprecation import MiddlewareMixin
+from django.utils.functional import SimpleLazyObject
+from django_otp import user_has_device
+from django_otp.middleware import OTPMiddleware
+
+from . import settings
 
 
-class OTPCheckMiddleware(MiddlewareMixin):
+class OTPCheckMiddleware(OTPMiddleware, MiddlewareMixin):
+    def __call__(self, request):
+        return MiddlewareMixin.__call__(self, request)
+
     def process_request(self, request):
-        user = request.user
-        if user.is_authenticated:
-            if user.is_verified() or not user_has_device(user):
+        request.user = SimpleLazyObject(
+            functools.partial(self._verify_user, request, request.user)
+        )
+        if request.user.is_authenticated:
+            if request.user.is_verified():
                 return None
-            else:
-                if request.path not in settings.NLOTP_EXCLUDED_URLS:
-                    return HttpResponseRedirect(NLOTP_LOGIN_URL)
+            if not user_has_device(request.user):
+                if (
+                    settings.NLOTP_2FA_SETUP_REQUIRED
+                    and request.path not in settings.NLOTP_SETUP_EXCLUDED_URLS
+                ):
+                    return HttpResponseRedirect(settings.NLOTP_2FA_SETUP_URL)
+                else:
+                    return None
+            if request.path not in settings.NLOTP_VERIFY_EXCLUDED_URLS:
+                return HttpResponseRedirect(settings.NLOTP_VERIFY_URL)
         return None
diff --git a/nlotp/settings.py b/nlotp/settings.py
new file mode 100644
--- /dev/null
+++ b/nlotp/settings.py
@@ -0,0 +1,58 @@
+from django.conf import settings
+from django.urls import reverse_lazy
+
+
+# 2FA verify page url
+NLOTP_VERIFY_URL = getattr(
+    settings,
+    "NLOTP_VERIFY_URL",
+    reverse_lazy("nlotp:verify-otp"),
+)
+
+
+# list of urls excluded from redirect to verify page if user is not verified
+# NOTE: login and logout views should be added to this list
+NLOTP_VERIFY_EXCLUDED_URLS = getattr(
+    settings,
+    "NLOTP_VERIFY_EXCLUDED_URLS",
+    (
+        NLOTP_VERIFY_URL,
+    ),
+)
+
+
+# if 2FA setup is mandatory
+NLOTP_2FA_SETUP_REQUIRED = getattr(
+    settings,
+    "NLOTP_2FA_SETUP_REQUIRED",
+    False,
+)
+
+
+# 2FA setup page url
+NLOTP_2FA_SETUP_URL = getattr(
+    settings,
+    "NLOTP_2FA_SETUP_URL",
+    reverse_lazy("nlotp:two-factor-auth"),
+)
+
+
+# user TOTP device QR code generation url
+NLOTP_QR_CODE_URL = getattr(
+    settings,
+    "NLOTP_QR_CODE_URL",
+    reverse_lazy("nlotp:qr-code"),
+)
+
+
+# list of urls excluded from redirect to 2FA setup page
+# when `NLOTP_2FA_SETUP_REQUIRED` is enabled and user didn't setup 2FA
+# NOTE: login and logout views should be added to this list
+NLOTP_SETUP_EXCLUDED_URLS = getattr(
+    settings,
+    "NLOTP_SETUP_EXCLUDED_URLS",
+    (
+        NLOTP_2FA_SETUP_URL,
+        NLOTP_QR_CODE_URL,
+    ),
+)
diff --git a/nlotp/templates/nlotp/two_factor_auth.html b/nlotp/templates/nlotp/two_factor_auth.html
--- a/nlotp/templates/nlotp/two_factor_auth.html
+++ b/nlotp/templates/nlotp/two_factor_auth.html
@@ -14,7 +14,7 @@
       {% endfor %}
     </ul>
     {% endif %}
-    <form method="POST" action="{% url 'nlotp:two-factor-auth' %}" id="manage-two-factor-auth-form" autocomplete="on">
+    <form method="POST" action="." id="manage-two-factor-auth-form" autocomplete="on">
       {% csrf_token %}
 
       {{ form.non_field_errors }}
@@ -37,6 +37,10 @@
       {% endwith %}
       <hr/>
       {% endif %}
+
+      {% if setup_required %}
+      <p><b>You need to enable 2FA to continue.</b></p>
+      {% endif %}
       
       <b>Status:</b> <span>{{ totp_device.confirmed|yesno:"Enabled,Disabled" }}</span>
 
@@ -52,7 +56,7 @@
       {% if not totp_device.confirmed %}
       <p>Please scan the QR code or manually enter the key, then enter an authentication code from your app in order to complete setup.</p>
       <p id="qrcode">
-        <img width="200" height="200" src="{% url 'nlotp:qr-code' totp_device.id %}" />
+        <img width="200" height="200" src="{{ qr_code_url}}" />
       </p>
       <p>
         <label>Key:</label> {{ totp_device.bin_key|b32encode_val }}
diff --git a/nlotp/templates/nlotp/verify_otp.html b/nlotp/templates/nlotp/verify_otp.html
--- a/nlotp/templates/nlotp/verify_otp.html
+++ b/nlotp/templates/nlotp/verify_otp.html
@@ -6,7 +6,7 @@
     <title>Two-Step Verification (2FA) Verification</title>
   </head>
   <body>
-    <form action="{% url 'nlotp:verify-otp' %}?next={{ next }}" method="POST" autocomplete="on">
+    <form action="?next={{ next }}" method="POST" autocomplete="on">
       {% csrf_token %}
       <h2>Two-Step Verification (2FA)</h2>
 
diff --git a/nlotp/urls.py b/nlotp/urls.py
--- a/nlotp/urls.py
+++ b/nlotp/urls.py
@@ -6,7 +6,7 @@
 
 urlpatterns = [
     path(
-        'qr-code/<int:device_id>/',
+        'qr-code/',
         views.QRCodeView.as_view(),
         name='qr-code',
     ),
diff --git a/nlotp/views.py b/nlotp/views.py
--- a/nlotp/views.py
+++ b/nlotp/views.py
@@ -6,19 +6,18 @@
 from django.shortcuts import get_object_or_404
 from django.urls import reverse_lazy
 from django.views.generic.edit import FormView, View
+from django_otp import DEVICE_ID_SESSION_KEY
 from django_otp import login as otp_login
 from django_otp.views import LoginView
 
-from . import forms, utils
+from . import forms, settings, utils
 
 
 class QRCodeView(LoginRequiredMixin, View):
-    """ Generates QR code for the TOTP device """
+    """ Generates QR code for the logged in user TOTP device """
 
-    def get(self, request, device_id, *args, **kwargs):
-        device = get_object_or_404(
-            request.user.totpdevice_set.all(), id=device_id
-        )
+    def get(self, request, *args, **kwargs):
+        device = utils.get_user_totp_device(request.user)
         img = qrcode.make(
             device.config_url, image_factory=qrcode.image.svg.SvgImage
         )
@@ -39,7 +38,7 @@
 
     template_name = "nlotp/two_factor_auth.html"
     form_class = forms.TwoFactorAuthForm
-    success_url = reverse_lazy("nlotp:two-factor-auth")
+    success_url = settings.NLOTP_2FA_SETUP_URL
 
     def setup(self, request, *args, **kwargs):
         super().setup(request, *args, **kwargs)
@@ -55,6 +54,8 @@
                 "totp_device": self.totp_device,
                 "static_device": self.static_device,
                 "show_codes": self.show_codes,
+                "qr_code_url": settings.NLOTP_QR_CODE_URL,
+                "setup_required": settings.NLOTP_2FA_SETUP_REQUIRED,
             }
         )
         return context
@@ -102,9 +103,9 @@
             self.totp_device = utils.get_user_totp_device(self.request.user)
             self.static_device.confirmed = False
             self.static_device.save()
+            del self.request.session[DEVICE_ID_SESSION_KEY]
             messages.success(
-                self.request,
-                "Two-Step Verification has been disabled.",
+                self.request, "Two-Step Verification has been disabled.",
             )
         context = self.get_context_data()
         context["form"] = self.form_class(