From ecc1c63c2c716fd3fcdea40d56cf7640db99e5fe Mon Sep 17 00:00:00 2001 From: SKairinos Date: Thu, 28 Sep 2023 17:12:07 +0100 Subject: [PATCH] create backup token model --- codeforlife/user/migrations/0001_initial.py | 14 +++- codeforlife/user/models/__init__.py | 1 + codeforlife/user/models/backup_token.py | 75 +++++++++++++++++++ codeforlife/user/models/user.py | 3 +- codeforlife/user/tests/models/__init__.py | 0 .../user/tests/models/test_backup_token.py | 63 ++++++++++++++++ 6 files changed, 154 insertions(+), 2 deletions(-) create mode 100644 codeforlife/user/models/backup_token.py create mode 100644 codeforlife/user/tests/models/__init__.py create mode 100644 codeforlife/user/tests/models/test_backup_token.py diff --git a/codeforlife/user/migrations/0001_initial.py b/codeforlife/user/migrations/0001_initial.py index ea3cab55..faf80104 100644 --- a/codeforlife/user/migrations/0001_initial.py +++ b/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')}, + }, + ), ] diff --git a/codeforlife/user/models/__init__.py b/codeforlife/user/models/__init__.py index 24324834..58d947e7 100644 --- a/codeforlife/user/models/__init__.py +++ b/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 diff --git a/codeforlife/user/models/backup_token.py b/codeforlife/user/models/backup_token.py new file mode 100644 index 00000000..b0d8a367 --- /dev/null +++ b/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 diff --git a/codeforlife/user/models/user.py b/codeforlife/user/models/user.py index 7d67d75c..47cc8e79 100644 --- a/codeforlife/user/models/user.py +++ b/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 diff --git a/codeforlife/user/tests/models/__init__.py b/codeforlife/user/tests/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/codeforlife/user/tests/models/test_backup_token.py b/codeforlife/user/tests/models/test_backup_token.py new file mode 100644 index 00000000..a20643a5 --- /dev/null +++ b/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, + )