Add Expiry for API Tokens

vulnerabilities/migrations/0089_expiringtoken.py

@@ -0,0 +1,33 @@
+# Generated by Django 4.2.17 on 2025-03-09 19:00
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("authtoken", "0004_alter_tokenproxy_options"),
+        ("vulnerabilities", "0088_fix_alpine_purl_type"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="ExpiringToken",
+            fields=[
+                (
+                    "token_ptr",
+                    models.OneToOneField(
+                        auto_created=True,
+                        on_delete=django.db.models.deletion.CASCADE,
+                        parent_link=True,
+                        primary_key=True,
+                        serialize=False,
+                        to="authtoken.token",
+                    ),
+                ),
+                ("expires", models.DateTimeField()),
+            ],
+            bases=("authtoken.token",),
+        ),
+    ]

vulnerabilities/models.py

@@ -13,6 +13,7 @@
 import logging
 import xml.etree.ElementTree as ET
 from contextlib import suppress
+from datetime import timedelta
 from functools import cached_property
 from itertools import groupby
 from operator import attrgetter
@@ -57,6 +58,7 @@
 from vulnerabilities.utils import normalize_purl
 from vulnerabilities.utils import purl_to_dict
 from vulnerablecode import __version__ as VULNERABLECODE_VERSION
+from vulnerablecode import settings
 
 logger = logging.getLogger(__name__)
 
@@ -1384,6 +1386,26 @@ def to_advisory_data(self) -> "AdvisoryData":
 UserModel = get_user_model()
 
 
+class ExpiringToken(Token):
+    """
+    Extends Django Rest Framework's Token model to include an expiration date.
+    """
+
+    expires = models.DateTimeField(null=False)
+
+    def is_expired(self):
+        return timezone.now() > self.expires
+
+    @classmethod
+    def get_or_create(cls, user):
+        token, created = Token._default_manager.get_or_create(user=user)
+
+        if created:
+            token.expires = timezone.now() + timedelta(days=settings.TOKEN_EXPIRY_DAYS)  # 30days
+            token.save()
+        return token, created
+
+
 class ApiUserManager(UserManager):
     def create_api_user(self, username, first_name="", last_name="", **extra_fields):
         """
@@ -1409,7 +1431,7 @@ def create_api_user(self, username, first_name="", last_name="", **extra_fields)
         user.set_unusable_password()
         user.save()
 
-        Token._default_manager.get_or_create(user=user)
+        ExpiringToken.get_or_create(user=user)
 
         return user
 

vulnerabilities/tests/test_token_authentication.py

@@ -0,0 +1,95 @@
+from datetime import timedelta
+
+import pytest
+from django.contrib.auth import get_user_model
+from django.utils import timezone
+from rest_framework import exceptions
+
+from vulnerabilities.models import ExpiringToken
+from vulnerabilities.token_authentication import ExpiringTokenAuthentication
+
+User = get_user_model()
+
+
+@pytest.mark.django_db
+def test_expiring_token_creation():
+    """
+    Test tha ExpiringToken is created with an expiration date
+    """
+    user = User.objects.create_user(username="testuser", email="[email protected]")
+    token, created = ExpiringToken.get_or_create(user=user)
+
+    #token and its expiration date 30days from today
+    #print(f"Token: {token.key}, Expires: {token.expires}")
+
+    assert created is True
+    assert token.expires > timezone.now()
+
+
+@pytest.mark.django_db
+def test_expiring_token_is_expired():
+    """
+    Test that is_expired method
+    """
+    user = User.objects.create_user(username="testuser", email="[email protected]")
+    token = ExpiringToken.objects.create(user=user, expires=timezone.now() - timedelta(days=1))
+    # token and its expiration date,yesterday
+    # print(f"Token: {token.key}, Expires: {token.expires}")
+
+    assert token.is_expired() is True #expired
+
+
+@pytest.mark.django_db
+def test_expiring_token_is_not_expired():
+    """
+    Test the is_expired method for a non-expired token
+    """
+    user = User.objects.create_user(username="testuser", email="[email protected]")
+    token = ExpiringToken.objects.create(user=user, expires=timezone.now() + timedelta(days=1))
+
+    #token and its expiration date,tomorrow
+    # print(f"Token: {token.key}, Expires: {token.expires}")
+
+    assert token.is_expired() is False #not expired
+
+
+@pytest.mark.django_db
+def test_expiring_token_authentication_valid():
+    """
+    Test that a valid token authenticates the user
+    """
+    user = User.objects.create_user(username="testuser", email="[email protected]")
+    token = ExpiringToken.objects.create(user=user, expires=timezone.now() + timedelta(days=1))
+
+    auth = ExpiringTokenAuthentication()
+    authenticated_user, authenticated_token = auth.authenticate_credentials(token.key)
+
+    # print(f'Authenticated User:{authenticated_user} and Authenticated Token:{authenticated_token}')
+
+    assert authenticated_user == user
+    assert authenticated_token == token
+
+
+@pytest.mark.django_db
+def test_expiring_token_authentication_expired():
+    """
+    Test that an expired token raises an AuthenticationFailed error
+    """
+    user = User.objects.create_user(username="testuser", email="[email protected]")
+    token = ExpiringToken.objects.create(user=user, expires=timezone.now() - timedelta(days=1))
+
+    auth = ExpiringTokenAuthentication()
+
+    with pytest.raises(exceptions.AuthenticationFailed):
+        auth.authenticate_credentials(token.key)
+
+
+@pytest.mark.django_db
+def test_expiring_token_authentication_invalid():
+    """
+    Test that an invalid/non-existin token raises an AuthenticationFailed error
+    """
+    auth = ExpiringTokenAuthentication()
+
+    with pytest.raises(exceptions.AuthenticationFailed):
+        auth.authenticate_credentials("invalid-token")

vulnerabilities/token_authentication.py

@@ -0,0 +1,21 @@
+from rest_framework import exceptions
+from rest_framework.authentication import TokenAuthentication
+
+from .models import ExpiringToken
+
+
+class ExpiringTokenAuthentication(TokenAuthentication):
+    model = ExpiringToken
+
+    def authenticate_credentials(self, key):
+        try:
+            #try to fetch the token
+            user, token = super().authenticate_credentials(key)
+        except self.model.DoesNotExist:
+            #if the token does not exist/invalid
+            raise exceptions.AuthenticationFailed("Invalid token")
+
+        if token.is_expired():
+            #if the token has expired
+            raise exceptions.AuthenticationFailed("Token has expired")
+        return user, token

vulnerablecode/settings.py

@@ -213,7 +213,9 @@
 # Django restframework
 
 REST_FRAMEWORK = {
-    "DEFAULT_AUTHENTICATION_CLASSES": ("rest_framework.authentication.TokenAuthentication",),
+    "DEFAULT_AUTHENTICATION_CLASSES": (
+        "vulnerabilities.token_authentication.ExpiringTokenAuthentication",
+    ),
     "DEFAULT_PERMISSION_CLASSES": ("rest_framework.permissions.IsAuthenticated",),
     "DEFAULT_RENDERER_CLASSES": (
         "rest_framework.renderers.JSONRenderer",
@@ -362,3 +364,6 @@
         "handlers": ["console"],
         "level": "ERROR",
     }
+
+# Set the number of days until the API token expires
+TOKEN_EXPIRY_DAYS = 30