Throttle API requests based on user permissions

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

vulnerabilities/api.py

@@ -34,7 +34,7 @@
 from vulnerabilities.models import get_purl_query_lookups
 from vulnerabilities.severity_systems import EPSS
 from vulnerabilities.severity_systems import SCORING_SYSTEMS
-from vulnerabilities.throttling import GroupUserRateThrottle
+from vulnerabilities.throttling import PermissionBasedUserRateThrottle
 from vulnerabilities.utils import get_severity_range
 
 
@@ -471,7 +471,7 @@ class PackageViewSet(viewsets.ReadOnlyModelViewSet):
     serializer_class = PackageSerializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = PackageFilterSet
-    throttle_classes = [AnonRateThrottle, GroupUserRateThrottle]
+    throttle_classes = [AnonRateThrottle, PermissionBasedUserRateThrottle]
 
     def get_queryset(self):
         return super().get_queryset().with_is_vulnerable()
@@ -688,7 +688,7 @@ def get_queryset(self):
     serializer_class = VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = VulnerabilityFilterSet
-    throttle_classes = [AnonRateThrottle, GroupUserRateThrottle]
+    throttle_classes = [AnonRateThrottle, PermissionBasedUserRateThrottle]
 
 
 class CPEFilterSet(filters.FilterSet):

vulnerabilities/api_extension.py

@@ -33,7 +33,7 @@
 from vulnerabilities.models import VulnerabilitySeverity
 from vulnerabilities.models import Weakness
 from vulnerabilities.models import get_purl_query_lookups
-from vulnerabilities.throttling import GroupUserRateThrottle
+from vulnerabilities.throttling import PermissionBasedUserRateThrottle
 
 
 class SerializerExcludeFieldsMixin:
@@ -259,7 +259,7 @@ class V2PackageViewSet(viewsets.ReadOnlyModelViewSet):
     lookup_field = "purl"
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = V2PackageFilterSet
-    throttle_classes = [GroupUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [PermissionBasedUserRateThrottle, AnonRateThrottle]
 
     def get_queryset(self):
         return super().get_queryset().with_is_vulnerable().prefetch_related("vulnerabilities")
@@ -345,7 +345,7 @@ class VulnerabilityViewSet(viewsets.ReadOnlyModelViewSet):
     lookup_field = "vulnerability_id"
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = V2VulnerabilityFilterSet
-    throttle_classes = [GroupUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [PermissionBasedUserRateThrottle, AnonRateThrottle]
 
     def get_queryset(self):
         """
@@ -381,7 +381,7 @@ class CPEViewSet(viewsets.ReadOnlyModelViewSet):
     ).distinct()
     serializer_class = V2VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
-    throttle_classes = [GroupUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [PermissionBasedUserRateThrottle, AnonRateThrottle]
     filterset_class = CPEFilterSet
 
     @action(detail=False, methods=["post"])
@@ -420,4 +420,4 @@ class AliasViewSet(viewsets.ReadOnlyModelViewSet):
     serializer_class = V2VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = AliasFilterSet
-    throttle_classes = [GroupUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [PermissionBasedUserRateThrottle, AnonRateThrottle]

vulnerabilities/migrations/0093_alter_apiuser_options.py

@@ -0,0 +1,23 @@
+# Generated by Django 4.2.22 on 2025-06-13 08:07
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("vulnerabilities", "0092_pipelineschedule_pipelinerun"),
+    ]
+
+    operations = [
+        migrations.AlterModelOptions(
+            name="apiuser",
+            options={
+                "permissions": [
+                    ("throttle_unrestricted", "Exempt from API throttling limits"),
+                    ("throttle_18000_hour", "Can make 18000 API requests per hour"),
+                    ("throttle_14400_hour", "Can make 14400 API requests per hour"),
+                ]
+            },
+        ),
+    ]

vulnerabilities/models.py

@@ -1473,10 +1473,6 @@ def create_api_user(self, username, first_name="", last_name="", **extra_fields)
         user.set_unusable_password()
         user.save()
 
-        # Assign the default basic group
-        default_group, _ = Group.objects.get_or_create(name="silver")
-        user.groups.add(default_group)
-
         Token._default_manager.get_or_create(user=user)
 
         return user
@@ -1494,14 +1490,17 @@ def _validate_username(self, email):
 
 
 class ApiUser(UserModel):
-    """
-    A User proxy model to facilitate simplified admin API user creation.
-    """
+    """A User proxy model to facilitate simplified admin API user creation."""
 
     objects = ApiUserManager()
 
     class Meta:
         proxy = True
+        permissions = [
+            ("throttle_unrestricted", "Exempt from API throttling limits"),
+            ("throttle_18000_hour", "Can make 18000 API requests per hour"),
+            ("throttle_14400_hour", "Can make 14400 API requests per hour"),
+        ]
 
 
 class ChangeLog(models.Model):

vulnerabilities/throttling.py

@@ -12,25 +12,19 @@
 from rest_framework.views import exception_handler
 
 
-class GroupUserRateThrottle(UserRateThrottle):
-    scope = "bronze"
-
+class PermissionBasedUserRateThrottle(UserRateThrottle):
     def allow_request(self, request, view):
         user = request.user
 
         if user and user.is_authenticated:
-            if user.is_superuser or user.is_staff:
-                return True
-
-            user_groups = user.groups.all()
-            if any([group.name == "gold" for group in user_groups]):
+            if user.has_perm("vulnerabilities.throttle_unrestricted"):
                 return True
+            elif user.has_perm("vulnerabilities.throttle_18000_hour"):
+                self.rate = "18000/hour"
+            elif user.has_perm("vulnerabilities.throttle_14400_hour"):
+                self.rate = "14400/hour"
 
-            if any([group.name == "silver" for group in user_groups]):
-                self.scope = "silver"
-
-        self.rate = self.THROTTLE_RATES.get(self.scope)
-        self.num_requests, self.duration = self.parse_rate(self.rate)
+            self.num_requests, self.duration = self.parse_rate(self.rate)
 
         return super().allow_request(request, view)
 

vulnerablecode/settings.py

@@ -190,20 +190,11 @@
 LOGIN_REDIRECT_URL = "/"
 LOGOUT_REDIRECT_URL = "/"
 
-REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {
-    # No throttling for users in gold group.
-    "silver": "10800/hour",
-    "bronze": "7200/hour",
-    "anon": "3600/hour",
-}
+REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {"anon": "3600/hour", "user": "10800/hour"}
+
 
 if IS_TESTS:
     VULNERABLECODEIO_REQUIRE_AUTHENTICATION = False
-    REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {
-        "silver": "20/day",
-        "bronze": "15/day",
-        "anon": "10/day",
-    }
 
 USE_L10N = True
 
@@ -243,7 +234,7 @@
         "rest_framework.filters.SearchFilter",
     ),
     "DEFAULT_THROTTLE_CLASSES": [
-        "vulnerabilities.throttling.GroupUserRateThrottle",
+        "vulnerabilities.throttling.PermissionBasedUserRateThrottle",
         "rest_framework.throttling.AnonRateThrottle",
     ],
     "DEFAULT_THROTTLE_RATES": REST_FRAMEWORK_DEFAULT_THROTTLE_RATES,