From c730deb4b3aa2b6af2fc7071b00545bb08c07162 Mon Sep 17 00:00:00 2001
From: Florian Aucomte <33633200+faucomte97@users.noreply.github.com>
Date: Wed, 12 Apr 2023 11:17:35 +0100
Subject: [PATCH] fix: Fix avatar token generator (#1758)

* fix: Decode and truncate avatar auth token

* fix: Factorise code
---
 aimmo/avatar_creator.py                        |  5 ++++-
 aimmo/game_creator.py                          |  8 ++------
 aimmo/migrations/0025_generate_auth_token.py   |  7 +++++--
 .../migrations/0031_alter_avatar_auth_token.py | 18 ++++++++++++++++++
 aimmo/models.py                                | 11 +++++------
 5 files changed, 34 insertions(+), 15 deletions(-)
 create mode 100644 aimmo/migrations/0031_alter_avatar_auth_token.py

diff --git a/aimmo/avatar_creator.py b/aimmo/avatar_creator.py
index 5476fffd5..b1d09fa9e 100644
--- a/aimmo/avatar_creator.py
+++ b/aimmo/avatar_creator.py
@@ -1,5 +1,7 @@
+from aimmo.models import Avatar, Game, generate_auth_token
 
-from aimmo.models import Avatar, Game
+NUM_BYTES_FOR_TOKEN_GENERATOR = 16
+TOKEN_MAX_LENGTH = 24
 
 
 def create_avatar_for_user(user, game_id):
@@ -14,4 +16,5 @@ def create_avatar_for_user(user, game_id):
     game: Game = Game.objects.get(id=game_id)
     initial_code = game.worksheet.starter_code
     avatar = Avatar.objects.create(owner=user, code=initial_code, game_id=game_id)
+    avatar.auth_token = generate_auth_token(NUM_BYTES_FOR_TOKEN_GENERATOR, TOKEN_MAX_LENGTH)
     return avatar
diff --git a/aimmo/game_creator.py b/aimmo/game_creator.py
index 6d144664f..a163844a7 100644
--- a/aimmo/game_creator.py
+++ b/aimmo/game_creator.py
@@ -1,15 +1,11 @@
-import secrets
 from aimmo.avatar_creator import create_avatar_for_user
 from aimmo.game_manager import GameManager
+from aimmo.models import generate_auth_token
 
 NUM_BYTES_FOR_TOKEN_GENERATOR = 32
 TOKEN_MAX_LENGTH = 48
 
 
-def generate_game_token():
-    return secrets.token_urlsafe(nbytes=NUM_BYTES_FOR_TOKEN_GENERATOR)[:TOKEN_MAX_LENGTH]
-
-
 def create_game(main_user, form):
     """
     Creates a Game by:
@@ -24,7 +20,7 @@ def create_game(main_user, form):
     :return: The initialised Game object.
     """
     game = form.save(commit=False)
-    game.auth_token = generate_game_token()
+    game.auth_token = generate_auth_token(NUM_BYTES_FOR_TOKEN_GENERATOR, TOKEN_MAX_LENGTH)
     game.generator = "Main"
     game.owner = game.game_class.teacher.new_user
     game.main_user = game.game_class.teacher.new_user
diff --git a/aimmo/migrations/0025_generate_auth_token.py b/aimmo/migrations/0025_generate_auth_token.py
index 5dc08e221..e58578c57 100644
--- a/aimmo/migrations/0025_generate_auth_token.py
+++ b/aimmo/migrations/0025_generate_auth_token.py
@@ -2,7 +2,7 @@
 
 from django.db import migrations, models
 
-from aimmo.game_creator import generate_game_token
+from aimmo.models import generate_auth_token
 from aimmo.game_manager import GameManager
 
 
@@ -13,9 +13,12 @@ class Migration(migrations.Migration):
     ]
 
     def generate_auth_token_for_games(apps, schema_editor):
+        NUM_BYTES_FOR_TOKEN_GENERATOR = 32
+        TOKEN_MAX_LENGTH = 48
+
         Game = apps.get_model("aimmo", "Game")
         for game in Game.objects.all():
-            game.auth_token = generate_game_token()
+            game.auth_token = generate_auth_token(NUM_BYTES_FOR_TOKEN_GENERATOR, TOKEN_MAX_LENGTH)
             game.save()
             game_manager = GameManager()
             game_manager.create_game_secret(game_id=game.id, token=game.auth_token)
diff --git a/aimmo/migrations/0031_alter_avatar_auth_token.py b/aimmo/migrations/0031_alter_avatar_auth_token.py
new file mode 100644
index 000000000..4680547c1
--- /dev/null
+++ b/aimmo/migrations/0031_alter_avatar_auth_token.py
@@ -0,0 +1,18 @@
+# Generated by Django 3.2.18 on 2023-04-11 14:55
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('aimmo', '0030_update_created_by'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='avatar',
+            name='auth_token',
+            field=models.CharField(blank=True, max_length=24),
+        ),
+    ]
diff --git a/aimmo/models.py b/aimmo/models.py
index 7586d63cb..4dd9a5a14 100755
--- a/aimmo/models.py
+++ b/aimmo/models.py
@@ -1,14 +1,13 @@
-from os import urandom
+import secrets
 
-from base64 import urlsafe_b64encode
 from common.models import Class
+from common.models import Teacher
 from django.contrib.auth.models import User
 from django.db import models
 from django.utils import timezone
 
 from aimmo import app_settings
 from aimmo.worksheets import WORKSHEETS
-from common.models import Teacher
 
 DEFAULT_WORKSHEET_ID = 1
 
@@ -17,8 +16,8 @@
 ]
 
 
-def generate_auth_token():
-    return urlsafe_b64encode(urandom(16))
+def generate_auth_token(nbytes, max_length):
+    return secrets.token_urlsafe(nbytes=nbytes)[:max_length]
 
 
 class Game(models.Model):
@@ -138,7 +137,7 @@ class Avatar(models.Model):
     owner = models.ForeignKey(User, on_delete=models.CASCADE)
     game = models.ForeignKey(Game, on_delete=models.CASCADE)
     code = models.TextField()
-    auth_token = models.CharField(max_length=24, default=generate_auth_token)
+    auth_token = models.CharField(max_length=24, blank=True)
 
     class Meta:
         unique_together = ("owner", "game")