Turn prefix and hashed_key into proper APIKey fields

Author: florimondmanca

rest_framework_api_key/admin.py

@@ -7,8 +7,8 @@
 
 class APIKeyModelAdmin(admin.ModelAdmin):
     list_display = (
-        "name",
         "prefix",
+        "name",
         "created",
         "expiry_date",
         "_has_expired",

rest_framework_api_key/crypto.py

@@ -7,6 +7,15 @@
     from .models import APIKey
 
 
+def concatenate(left: str, right: str) -> str:
+    return "{}.{}".format(left, right)
+
+
+def split(concatenated: str) -> Tuple[str, str]:
+    left, _, right = concatenated.partition(".")
+    return left, right
+
+
 class KeyGenerator:
     def __init__(self, prefix_length: int = 8, secret_key_length: int = 32):
         self.prefix_length = prefix_length
@@ -18,16 +27,15 @@ def get_prefix(self) -> str:
     def get_secret_key(self) -> str:
         return get_random_string(self.secret_key_length)
 
-    @staticmethod
-    def concatenate(left: str, right: str) -> str:
-        return "{}.{}".format(left, right)
+    def hash(self, value: str) -> str:
+        return make_password(value)
 
-    def generate(self) -> Tuple[str, str]:
+    def generate(self) -> Tuple[str, str, str]:
         prefix = self.get_prefix()
         secret_key = self.get_secret_key()
-        key = self.concatenate(prefix, secret_key)
-        hashed_key = self.concatenate(prefix, make_password(key))
-        return key, hashed_key
+        key = concatenate(prefix, secret_key)
+        hashed_key = self.hash(key)
+        return key, prefix, hashed_key
 
     def verify(self, key: str, hashed_key: str) -> bool:
         return check_password(key, hashed_key)

rest_framework_api_key/migrations/0004_prefix_hashed_key.py

@@ -0,0 +1,48 @@
+# Generated by Django 2.2.2 on 2019-06-29 10:38
+
+from django.db import migrations, models
+
+APP_NAME = "rest_framework_api_key"
+MODEL_NAME = "apikey"
+DEPENDENCIES = [(APP_NAME, "0003_auto_20190623_1952")]
+
+
+def populate_prefix_hashed_key(apps, schema_editor):
+    model = apps.get_model(APP_NAME, MODEL_NAME)
+
+    for api_key in model.objects.all():
+        prefix, _, hashed_key = api_key.id.partition(".")
+        api_key.prefix = prefix
+        api_key.hashed_key = hashed_key
+        api_key.save()
+
+
+class Migration(migrations.Migration):
+
+    dependencies = DEPENDENCIES
+
+    operations = [
+        migrations.AddField(
+            model_name=MODEL_NAME,
+            name="hashed_key",
+            field=models.CharField(max_length=100, null=True),
+        ),
+        migrations.AddField(
+            model_name=MODEL_NAME,
+            name="prefix",
+            field=models.CharField(max_length=8, unique=True, null=True),
+        ),
+        migrations.RunPython(
+            populate_prefix_hashed_key, migrations.RunPython.noop
+        ),
+        migrations.AlterField(
+            model_name=MODEL_NAME,
+            name="hashed_key",
+            field=models.CharField(max_length=100, editable=False),
+        ),
+        migrations.AlterField(
+            model_name=MODEL_NAME,
+            name="prefix",
+            field=models.CharField(max_length=8, unique=True, editable=False),
+        ),
+    ]

rest_framework_api_key/models.py

@@ -4,15 +4,26 @@
 from django.db import models
 from django.utils import timezone
 
-from .crypto import KeyGenerator
+from .crypto import concatenate, KeyGenerator, split
 
 
 class BaseAPIKeyManager(models.Manager):
     key_generator = KeyGenerator()
 
     def assign_key(self, obj: "AbstractAPIKey") -> str:
-        key, hashed_key = self.key_generator.generate()
-        obj.id = hashed_key
+        try:
+            key, prefix, hashed_key = self.key_generator.generate()
+        except TypeError:  # Compatibility with < 1.4
+            key, hashed_key = self.key_generator.generate()
+            pk = hashed_key
+            prefix, hashed_key = split(hashed_key)
+        else:
+            pk = concatenate(prefix, hashed_key)
+
+        obj.id = pk
+        obj.prefix = prefix
+        obj.hashed_key = hashed_key
+
         return key
 
     def create_key(self, **kwargs) -> Tuple["AbstractAPIKey", str]:
@@ -32,10 +43,8 @@ def is_valid(self, key: str) -> bool:
         queryset = self.get_usable_keys()
 
         try:
-            api_key = queryset.get(
-                id__startswith=prefix
-            )  # type: AbstractAPIKey
-        except (self.model.DoesNotExist, self.model.MultipleObjectsReturned):
+            api_key = queryset.get(prefix=prefix)  # type: AbstractAPIKey
+        except self.model.DoesNotExist:
             return False
 
         if not api_key.is_valid(key):
@@ -57,6 +66,8 @@ class AbstractAPIKey(models.Model):
     id = models.CharField(
         max_length=100, unique=True, primary_key=True, editable=False
     )
+    prefix = models.CharField(max_length=8, unique=True, editable=False)
+    hashed_key = models.CharField(max_length=100, editable=False)
     created = models.DateTimeField(auto_now_add=True, db_index=True)
     name = models.CharField(
         max_length=50,
@@ -94,12 +105,6 @@ def __init__(self, *args, **kwargs):
         # Store the initial value of `revoked` to detect changes.
         self._initial_revoked = self.revoked
 
-    def _prefix(self) -> str:
-        return self.pk.partition(".")[0]
-
-    _prefix.short_description = "Prefix"
-    prefix = property(_prefix)
-
     def _has_expired(self) -> bool:
         if self.expiry_date is None:
             return False
@@ -110,8 +115,7 @@ def _has_expired(self) -> bool:
     has_expired = property(_has_expired)
 
     def is_valid(self, key: str) -> bool:
-        _, _, hashed_key = self.pk.partition(".")
-        return type(self).objects.key_generator.verify(key, hashed_key)
+        return type(self).objects.key_generator.verify(key, self.hashed_key)
 
     def clean(self):
         self._validate_revoked()

tests/legacy/__init__.py

@@ -0,0 +1,13 @@
+from typing import Tuple
+from rest_framework_api_key.crypto import concatenate, KeyGenerator
+
+
+class LegacyKeyGenerator(KeyGenerator):
+    """Pre-1.4 key generator."""
+
+    def generate(self) -> Tuple[str, str, str]:
+        prefix = self.get_prefix()
+        secret_key = self.get_secret_key()
+        key = concatenate(prefix, secret_key)
+        hashed_key = concatenate(prefix, self.hash(key))
+        return key, hashed_key

tests/legacy/crypto.py

@@ -0,0 +1,6 @@
+from rest_framework_api_key.permissions import HasAPIKey
+from .crypto import LegacyKeyGenerator
+
+
+class HasAPIKeyWithLegacyKeyGenerator(HasAPIKey):
+    key_generator = LegacyKeyGenerator()

tests/legacy/permissions.py

@@ -0,0 +1,48 @@
+# Generated by Django 2.2.2 on 2019-06-29 10:38
+
+from django.db import migrations, models
+
+APP_NAME = "heroes"
+MODEL_NAME = "heroapikey"
+DEPENDENCIES = [(APP_NAME, "0001_initial")]
+
+
+def populate_prefix_hashed_key(apps, schema_editor):
+    model = apps.get_model(APP_NAME, MODEL_NAME)
+
+    for api_key in model.objects.all():
+        prefix, _, hashed_key = api_key.id.partition(".")
+        api_key.prefix = prefix
+        api_key.hashed_key = hashed_key
+        api_key.save()
+
+
+class Migration(migrations.Migration):
+
+    dependencies = DEPENDENCIES
+
+    operations = [
+        migrations.AddField(
+            model_name=MODEL_NAME,
+            name="hashed_key",
+            field=models.CharField(max_length=100, null=True),
+        ),
+        migrations.AddField(
+            model_name=MODEL_NAME,
+            name="prefix",
+            field=models.CharField(max_length=8, unique=True, null=True),
+        ),
+        migrations.RunPython(
+            populate_prefix_hashed_key, migrations.RunPython.noop
+        ),
+        migrations.AlterField(
+            model_name=MODEL_NAME,
+            name="hashed_key",
+            field=models.CharField(max_length=100, editable=False),
+        ),
+        migrations.AlterField(
+            model_name=MODEL_NAME,
+            name="prefix",
+            field=models.CharField(max_length=8, unique=True, editable=False),
+        ),
+    ]

tests/project/heroes/migrations/0002_prefix_hashed_key.py

@@ -13,7 +13,8 @@
 
 def test_key_generation():
     api_key, generated_key = APIKey.objects.create_key(name="test")
-    prefix, _, hashed_key = api_key.id.partition(".")
+    prefix = api_key.prefix
+    hashed_key = api_key.hashed_key
 
     assert prefix and hashed_key
 

tests/test_model.py

@@ -0,0 +1,16 @@
+import pytest
+
+from .legacy.permissions import HasAPIKeyWithLegacyKeyGenerator
+
+pytestmark = pytest.mark.django_db
+
+
+@pytest.fixture(name="view")
+def fixture_view(view_with_permissions):
+    return view_with_permissions(HasAPIKeyWithLegacyKeyGenerator)
+
+
+def test_api_key_granted(create_request, view):
+    request = create_request()
+    response = view(request)
+    assert response.status_code == 200