From 0b83758625b262dfa525f78a1a2f6c61e7ce1640 Mon Sep 17 00:00:00 2001 From: Mohamed El-Kalioby Date: Wed, 23 Jun 2021 20:43:31 +0300 Subject: [PATCH] Throttling TOTP --- mfa/models.py | 11 +++++++++++ mfa/totp.py | 9 +++++++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/mfa/models.py b/mfa/models.py index 30ced21..ff60416 100644 --- a/mfa/models.py +++ b/mfa/models.py @@ -30,3 +30,14 @@ class UserKey(models.Model): class Meta: app_label = "mfa" + + +class OTPTracker(models.Model): + username = models.CharField(max_length=50) + value = models.CharField(max_length = 6) + success = models.BooleanField(blank=True) + done_on = models.DateTimeField(auto_now=True) + + class Meta: + app_label = 'mfa' + indexes = [models.Index(fields=['username'])] \ No newline at end of file diff --git a/mfa/totp.py b/mfa/totp.py index efa8c6c..3d2b6db 100644 --- a/mfa/totp.py +++ b/mfa/totp.py @@ -11,7 +11,7 @@ from django.utils import timezone from django.views.decorators.cache import never_cache from .Common import get_redirect_url -from .models import UserKey +from .models import UserKey, OTPTracker from .views import login @@ -19,10 +19,14 @@ def verify_login(request, username, token): for key in UserKey.objects.filter(username=username, key_type="TOTP"): totp = pyotp.TOTP(key.properties["secret_key"]) if totp.verify(token, valid_window=30): + if OTPTracker.objects.filter(username=username, value=token).exists(): + return [False, "Used Before, please generate another token"] + TOTP_Tracker.objects.create(username=username,value=token, success=True) key.last_used = timezone.now() key.save() return [True, key.id] - return [False] + TOTP_Tracker.objects.create(username = username, value = token, success = False) + return [False,"Invalid Token"] def recheck(request): @@ -60,6 +64,7 @@ def auth(request): request.session["mfa"] = mfa return login(request) context["invalid"] = True + context["invalid_msg"] = res[1] return render(request, "TOTP/Auth.html", context)