From bcf3ecc15c2d015c7a97f51b3cf195773efebc32 Mon Sep 17 00:00:00 2001
From: Spitap <dev@asdrip.fr>
Date: Mon, 22 Aug 2022 12:15:08 +0200
Subject: [PATCH] Fixed generation issue, warning when user uses its last
 backup code

---
 mfa/recovery.py                 | 45 ++++++++++++++++++++++++++++-----
 mfa/templates/TOTP/recheck.html | 24 ++++++++++++++++--
 mfa/totp.py                     | 11 --------
 mfa/urls.py                     |  1 +
 mfa/views.py                    |  1 +
 5 files changed, 63 insertions(+), 19 deletions(-)

diff --git a/mfa/recovery.py b/mfa/recovery.py
index 15ab30e..d4b4351 100644
--- a/mfa/recovery.py
+++ b/mfa/recovery.py
@@ -1,19 +1,22 @@
 from django.shortcuts import render
 from django.views.decorators.cache import never_cache
+from django.template.context_processors import csrf
 from django.http import HttpResponse
 from .Common import get_redirect_url
 from .models import *
 import simplejson
 import random
 import string
-
+import datetime
 
 #TODO : 
 # - Show authtificator panel on login everytime if RECOVERY is not deactivated
 # - Generation abuse checks
 
-def token_left(request):
-    uk = User_Keys.objects.filter(username=request.user.username, key_type="RECOVERY", enabled=True)
+def token_left(request, username=None):
+    if not username and request:
+        username = request.user.username
+    uk = User_Keys.objects.filter(username=username, key_type = "RECOVERY")
     keyLeft=0
     for key in uk:
         keyEnabled = key.properties["enabled"]
@@ -58,9 +61,10 @@ def verify_login(request, username,token):
             if token == secret_keys[i] and key.properties["enabled"][i]:
                 key.properties["enabled"][i] = False
                 key.save()
-                if token_left(request) == 0:
-                    newTokens(username)
-                return [True, key.id]
+                lastToken = False
+                if token_left(None, username) == 0:
+                    lastToken = True
+                return [True, key.id, lastToken]
     return [False]
 
 def getTokens(request):
@@ -73,6 +77,35 @@ def getTokens(request):
             enable.append(key.properties["enabled"][i])
     return HttpResponse(simplejson.dumps({"keys":tokens, "enable":enable}))
 
+@never_cache
+def auth(request):
+    from .views import login
+    context=csrf(request)
+    if request.method=="POST":
+        tokenLength = len(request.POST["otp"])
+        if tokenLength == 10 and "RECOVERY" not in settings.MFA_UNALLOWED_METHODS:
+            #Backup code check
+            resBackup=verify_login(request, request.session["base_username"], token=request.POST["otp"])
+            if resBackup[0]:
+                mfa = {"verified": True, "method": "RECOVERY","id":resBackup[1], "lastBackup":resBackup[2]}
+                if getattr(settings, "MFA_RECHECK", False):
+                    mfa["next_check"] = datetime.datetime.timestamp((datetime.datetime.now()
+                                            + datetime.timedelta(
+                                seconds=random.randint(settings.MFA_RECHECK_MIN, settings.MFA_RECHECK_MAX))))
+                request.session["mfa"] = mfa
+                if resBackup[2]:
+                    #If the last bakup code has just been used, we return a response insead of redirecting to login
+                    context["lastBackup"] = True
+                    return render(request,"TOTP/Auth.html", context)                
+                return login(request)
+    elif request.method=="GET":
+        mfa = request.session["mfa"]
+        if mfa and mfa["verified"] and mfa["method"] == "RECOVERY" and "lastBackup":
+            return login(request)
+
+    context["invalid"]=True
+    return render(request,"TOTP/Auth.html", context)
+
 @never_cache
 def start(request):
     """Start Managing recovery tokens"""
diff --git a/mfa/templates/TOTP/recheck.html b/mfa/templates/TOTP/recheck.html
index cd2bb40..e28ed1e 100644
--- a/mfa/templates/TOTP/recheck.html
+++ b/mfa/templates/TOTP/recheck.html
@@ -1,4 +1,14 @@
 <script type="application/javascript">
+    $(document).ready(function showWarningLastBackup() {
+        {% if lastBackup %}
+        $("#modal-title").html("Last backup code used !")
+        $("#modal-body").html("Don't forget to regenerate new backup code after login !")
+        $('#modal-footer').html(`<FORM METHOD="GET" ACTION="{% url 'recovery_auth' %}" Id="confirmLogin" onSubmit="" name="recoveryLastBackupConfirm">
+        <input type='submit'class='btn btn-lg btn-success btn-block' value='Continue'>`)
+        $("#popUpModal").modal('show')
+        {% endif %}
+        return
+    });
     function send_totp() {
         $.ajax({"url":"{% url 'totp_recheck' %}", method:"POST",dataType:"JSON",
             data:{"csrfmiddlewaretoken":"{{ csrf_token }}","otp":$("#otp").val()},
@@ -10,7 +20,16 @@
                             }
                         }
         })
-
+    }
+    function tryToAuth() {
+        if ($("#otp").val().length == 6) {
+            document.getElementById("formLogin").submit();     
+        }
+        else if ($("#otp").val().length == 10) {
+            const form = document.getElementById("formLogin");
+            form.setAttribute("ACTION", "{% url 'recovery_auth' %}")
+            form.submit();
+        }
     }
 </script>
 <div class='container'>
@@ -58,7 +77,7 @@
 
           <div class="form-group d-grid gap-2">
 
-	      <input type="{% if mode == "auth" %}submit{% elif mode == 'recheck'  %}button{% endif %}" {% if mode == "recheck" %}onclick="send_totp()" {% endif %} class="btn btn-lg btn-success btn-block" value="Sign in">
+	      <input type="button" onclick="{% if mode == "recheck" %} send_totp() {% else %} tryToAuth() {% endif %}" class="btn btn-lg btn-success btn-block" value="Sign in">
 	      </div>
       </div>
       </fieldset>
@@ -76,3 +95,4 @@
 	      </div>
 	      </div>
 </div>
+{% include "modal.html" %}
\ No newline at end of file
diff --git a/mfa/totp.py b/mfa/totp.py
index 42a5de6..7781cce 100644
--- a/mfa/totp.py
+++ b/mfa/totp.py
@@ -53,17 +53,6 @@ def auth(request):
                                 seconds=random.randint(settings.MFA_RECHECK_MIN, settings.MFA_RECHECK_MAX))))
                 request.session["mfa"] = mfa
                 return login(request)
-        elif tokenLength == 10 and "RECOVERY" not in settings.MFA_UNALLOWED_METHODS:
-            #Backup code check
-            resBackup=recovery.verify_login(request, request.session["base_username"], token=request.POST["otp"])
-            if resBackup[0]:
-                mfa = {"verified": True, "method": "RECOVERY","id":resBackup[1]}
-                if getattr(settings, "MFA_RECHECK", False):
-                    mfa["next_check"] = datetime.datetime.timestamp((datetime.datetime.now()
-                                            + datetime.timedelta(
-                                seconds=random.randint(settings.MFA_RECHECK_MIN, settings.MFA_RECHECK_MAX))))
-                request.session["mfa"] = mfa
-                return login(request)
         context["invalid"]=True
     return render(request,"TOTP/Auth.html", context)
 
diff --git a/mfa/urls.py b/mfa/urls.py
index e894f77..a85ef5d 100644
--- a/mfa/urls.py
+++ b/mfa/urls.py
@@ -15,6 +15,7 @@ urlpatterns = [
     url(r'recovery/start', recovery.start, name="manage_recovery_codes"),
     url(r'recovery/getTokens', recovery.getTokens, name="get_recovery_tokens"),
     url(r'recovery/genTokens', recovery.genTokens, name="regen_recovery_tokens"),
+    url(r'recovery/auth', recovery.auth, name="recovery_auth"),
 
     url(r'email/start/', Email.start , name="start_email"),
     url(r'email/auth/', Email.auth , name="email_auth"),
diff --git a/mfa/views.py b/mfa/views.py
index 640df25..23c56e7 100644
--- a/mfa/views.py
+++ b/mfa/views.py
@@ -55,6 +55,7 @@ def reset_cookie(request):
     response=HttpResponseRedirect(settings.LOGIN_URL)
     response.delete_cookie("base_username")
     return response
+
 def login(request):
     from django.contrib import auth
     from django.conf import settings