create backup token model

codeforlife/user/migrations/0001_initial.py

@@ -1,6 +1,7 @@
-# Generated by Django 3.2.20 on 2023-09-27 15:03
+# Generated by Django 3.2.20 on 2023-09-28 15:43
 
 import django.contrib.auth.models
+import django.core.validators
 from django.db import migrations, models
 import django.db.models.deletion
 
@@ -64,4 +65,15 @@ class Migration(migrations.Migration):
                 'unique_together': {('session', 'auth_factor')},
             },
         ),
+        migrations.CreateModel(
+            name='BackupToken',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('token', models.CharField(max_length=8, validators=[django.core.validators.MinLengthValidator(8)])),
+                ('user', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='backup_tokens', to='user.user')),
+            ],
+            options={
+                'unique_together': {('user', 'token')},
+            },
+        ),
     ]

codeforlife/user/models/__init__.py

@@ -7,6 +7,7 @@
 # from .teacher_invitation import SchoolTeacherInvitation
 # from .teacher import Teacher
 from .auth_factor import AuthFactor
+from .backup_token import BackupToken
 from .session import Session
 from .session_auth_factor import SessionAuthFactor
 from .user import User

codeforlife/user/models/backup_token.py

@@ -0,0 +1,75 @@
+import typing as t
+from itertools import groupby
+
+from django.contrib.auth.hashers import check_password, make_password
+from django.core.exceptions import ValidationError
+from django.core.validators import MinLengthValidator
+from django.db import models
+
+from . import user
+
+
+class BackupToken(models.Model):
+    max_count = 10
+    max_count_validation_error = ValidationError(
+        f"Exceeded max count of {max_count}"
+    )
+
+    class Manager(models.Manager["BackupToken"]):
+        def create(self, token: str, **kwargs):
+            return super().create(token=make_password(token), **kwargs)
+
+        def bulk_create(
+            self,
+            backup_tokens: t.List["BackupToken"],
+            *args,
+            **kwargs,
+        ):
+            def key(backup_token: BackupToken):
+                return backup_token.user.id
+
+            backup_tokens.sort(key=key)
+            for user_id, group in groupby(backup_tokens, key=key):
+                if (
+                    len(list(group))
+                    + BackupToken.objects.filter(user_id=user_id).count()
+                    > BackupToken.max_count
+                ):
+                    raise BackupToken.max_count_validation_error
+
+            for backup_token in backup_tokens:
+                backup_token.token = make_password(backup_token.token)
+
+            return super().bulk_create(backup_tokens, *args, **kwargs)
+
+    objects: Manager = Manager()
+
+    user: "user.User" = models.ForeignKey(
+        "user.User",
+        related_name="backup_tokens",
+        on_delete=models.CASCADE,
+    )
+
+    token = models.CharField(
+        max_length=8,
+        validators=[MinLengthValidator(8)],
+    )
+
+    class Meta:
+        unique_together = ["user", "token"]
+
+    def save(self, *args, **kwargs):
+        if self.id is None:
+            if (
+                BackupToken.objects.filter(user=self.user).count()
+                >= BackupToken.max_count
+            ):
+                raise BackupToken.max_count_validation_error
+
+        return super().save(*args, **kwargs)
+
+    def check_token(self, token: str):
+        if check_password(token, self.token):
+            self.delete()
+            return True
+        return False

codeforlife/user/models/user.py

@@ -42,11 +42,12 @@
 from django.db.models.query import QuerySet
 from django.utils.translation import gettext_lazy as _
 
-from . import auth_factor, session
+from . import auth_factor, backup_token, session
 
 
 class User(_User):
     auth_factors: QuerySet["auth_factor.AuthFactor"]
+    backup_tokens: QuerySet["backup_token.BackupToken"]
     session: "session.Session"
     userprofile: UserProfile
 

codeforlife/user/tests/models/test_backup_token.py

@@ -0,0 +1,63 @@
+from django.contrib.auth.hashers import check_password
+from django.core.exceptions import ValidationError
+from django.test import TestCase
+from django.utils.crypto import get_random_string
+
+from ...models import BackupToken, User
+
+
+class TestBackupToken(TestCase):
+    def setUp(self):
+        self.user = User.objects.get(id=2)
+
+    def test_bulk_create(self):
+        token = get_random_string(8)
+        backup_tokens = BackupToken.objects.bulk_create(
+            [BackupToken(user=self.user, token=token)]
+        )
+
+        assert check_password(token, backup_tokens[0].token)
+        with self.assertRaises(ValidationError):
+            BackupToken.objects.bulk_create(
+                [
+                    BackupToken(
+                        user=self.user,
+                        token=get_random_string(8),
+                    )
+                    for _ in range(BackupToken.max_count)
+                ]
+            )
+
+    def test_create(self):
+        token = get_random_string(8)
+        backup_token = BackupToken.objects.create(user=self.user, token=token)
+
+        assert check_password(token, backup_token.token)
+
+        BackupToken.objects.bulk_create(
+            [
+                BackupToken(
+                    user=self.user,
+                    token=get_random_string(8),
+                )
+                for _ in range(BackupToken.max_count - 1)
+            ]
+        )
+
+        with self.assertRaises(ValidationError):
+            BackupToken.objects.create(
+                user=self.user,
+                token=get_random_string(8),
+            )
+
+    def test_check_token(self):
+        token = get_random_string(8)
+        backup_token = BackupToken.objects.create(user=self.user, token=token)
+
+        assert backup_token.check_token(token)
+        assert backup_token.id is None
+        with self.assertRaises(BackupToken.DoesNotExist):
+            BackupToken.objects.get(
+                user=backup_token.user,
+                token=backup_token.token,
+            )