Merge pull request #16 from dimagi/pkv/otp-sender-fixes

Add otp generation and sending rate limiting

users/migrations/0009_phonedevice_otp_last_sent.py

@@ -0,0 +1,18 @@
+# Generated by Django 4.1.7 on 2024-06-25 08:00
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("users", "0008_credential_usercredential"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="phonedevice",
+            name="otp_last_sent",
+            field=models.DateTimeField(blank=True, null=True),
+        ),
+    ]

users/models.py

@@ -1,3 +1,4 @@
+from datetime import timedelta
 import base64
 from datetime import timedelta
 import os
@@ -10,6 +11,7 @@
 from django.contrib.sites.models import Site
 from django.db import models
 from django.http import HttpResponse
+from django.utils.timezone import now
 from django.urls import reverse
 from django.utils.timezone import now
 from django_otp.models import SideChannelDevice
@@ -93,13 +95,25 @@ def get_or_create_key_for_user(cls, user):
 class PhoneDevice(SideChannelDevice):
     phone_number = PhoneNumberField()
     user = models.ForeignKey(ConnectUser, on_delete=models.CASCADE)
+    otp_last_sent = models.DateTimeField(null=True, blank=True)
 
     def generate_challenge(self):
-        self.generate_token(valid_secs=600)
+        # generate and send new token if the old token is valid for less than 5 minutes
+        # set he otp_last_sent to None to send the new OTP immediately
+        if self.valid_until - now() <= timedelta(minutes=5):
+            self.otp_last_sent = None
+            self.generate_token(valid_secs=600)
         message = f"Your verification token from commcare connect is {self.token} \n\n {settings.APP_HASH}"
-        if not self.phone_number.raw_input.startswith(TEST_NUMBER_PREFIX):
-            sender = get_sms_sender(self.phone_number.country_code)
-            send_sms(self.phone_number.as_e164, message, sender)
+        # send the OTP if last sent message is not within the last 2 minutes
+        if self.otp_last_sent is None or (
+            self.otp_last_sent and now() - self.otp_last_sent >= timedelta(minutes=2)
+        ):
+            if not self.phone_number.raw_input.startswith(TEST_NUMBER_PREFIX):
+                sender = get_sms_sender(self.phone_number.country_code)
+                send_sms(self.phone_number.as_e164, message, sender)
+            self.otp_last_sent = now()
+            self.save()
+
         return message
 
     @classmethod
@@ -113,16 +127,18 @@ def send_otp_httpresponse(cls, phone_number, user):
 
     class Meta:
         constraints = [
-            models.UniqueConstraint(fields=['phone_number', 'user'], name='phone_number_user')
+            models.UniqueConstraint(
+                fields=["phone_number", "user"], name="phone_number_user"
+            )
         ]
 
 
 class RecoveryStatus(models.Model):
     class RecoverySteps(models.TextChoices):
-        CONFIRM_PRIMARY = 'primary'
-        CONFIRM_SECONDARY = 'secondary'
-        RESET_PASSWORD = 'password'
-        
+        CONFIRM_PRIMARY = "primary"
+        CONFIRM_SECONDARY = "secondary"
+        RESET_PASSWORD = "password"
+
     secret_key = models.TextField()
     user = models.ForeignKey(ConnectUser, on_delete=models.CASCADE, unique=True)
     step = models.TextField(choices=RecoverySteps.choices)

users/tests.py

@@ -1,10 +1,14 @@
+from datetime import timedelta
+from re import A
+from django.utils.timezone import now
 import pytest
 from django.urls import reverse
 from fcm_django.models import FCMDevice
+from unittest import mock
 
 from users.factories import CredentialFactory
 from users.fcm_utils import create_update_device
-from users.models import ConnectUser
+from users.models import ConnectUser, PhoneDevice
 
 
 @pytest.mark.django_db
@@ -77,6 +81,59 @@ def test_create_update_device__update_old_device(user):
     assert response.content == b'{"warning": "Another device is already active"}'
 
 
+def test_otp_generation(user):
+    with mock.patch("users.models.send_sms"):
+        phone_device, _ = PhoneDevice.objects.get_or_create(phone_number=user.phone_number, user=user)
+        phone_device.generate_challenge()
+
+        assert phone_device.token is not None
+        assert phone_device.otp_last_sent is not None
+
+
+def test_otp_generation_after_two_minutes(user):
+    with mock.patch("users.models.send_sms") as send_sms:
+        phone_device, _ = PhoneDevice.objects.get_or_create(phone_number=user.phone_number, user=user)
+        phone_device.generate_challenge()
+        token =  phone_device.token
+        assert token is not None
+        assert phone_device.otp_last_sent is not None
+        assert send_sms.call_count == 1
+
+        phone_device.generate_challenge()
+        assert phone_device.token == token
+        assert send_sms.call_count == 1
+
+        phone_device.otp_last_sent -= timedelta(minutes=2)
+        phone_device.save()
+        phone_device.generate_challenge()
+        assert phone_device.token == token
+        assert send_sms.call_count == 2
+
+
+def test_otp_generation_after_five_minutes(user):
+    with mock.patch("users.models.send_sms") as send_sms:
+        phone_device, _ = PhoneDevice.objects.get_or_create(phone_number=user.phone_number, user=user)
+        phone_device.generate_challenge()
+        token = phone_device.token
+        assert phone_device.token is not None
+        assert phone_device.otp_last_sent is not None
+        assert send_sms.call_count == 1
+
+        phone_device.otp_last_sent -= timedelta(minutes=3)
+        phone_device.save()
+        phone_device.generate_challenge()
+        assert phone_device.token is not None
+        assert phone_device.token == token
+        assert send_sms.call_count == 2
+
+        phone_device.valid_until -= timedelta(minutes=5)
+        phone_device.save()
+        phone_device.generate_challenge()
+        assert phone_device.token is not None
+        assert phone_device.token != token
+        assert send_sms.call_count == 3
+
+
 @pytest.mark.django_db
 class TestFetchCredentials:
 
@@ -99,4 +156,4 @@ def test_fetch_credential_with_org_slug(self, authed_client):
 
     def test_fetch_credential_without_org_slug(self, authed_client):
         response = authed_client.get(self.url)
-        self.assert_statements(response, expected_count=13)
+        self.assert_statements(response, expected_count=13)