Throttle API users based on user group

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 StaffUserRateThrottle
+from vulnerabilities.throttling import GroupUserRateThrottle
 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 = [StaffUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [AnonRateThrottle, GroupUserRateThrottle]
 
     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 = [StaffUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [AnonRateThrottle, GroupUserRateThrottle]
 
 
 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 StaffUserRateThrottle
+from vulnerabilities.throttling import GroupUserRateThrottle
 
 
 class SerializerExcludeFieldsMixin:
@@ -259,7 +259,7 @@ class V2PackageViewSet(viewsets.ReadOnlyModelViewSet):
     lookup_field = "purl"
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = V2PackageFilterSet
-    throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [GroupUserRateThrottle, 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 = [StaffUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [GroupUserRateThrottle, AnonRateThrottle]
 
     def get_queryset(self):
         """
@@ -381,7 +381,7 @@ class CPEViewSet(viewsets.ReadOnlyModelViewSet):
     ).distinct()
     serializer_class = V2VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
-    throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [GroupUserRateThrottle, 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 = [StaffUserRateThrottle, AnonRateThrottle]
+    throttle_classes = [GroupUserRateThrottle, AnonRateThrottle]

vulnerabilities/models.py

@@ -28,6 +28,7 @@
 from cwe2.mappings import xml_database_path
 from cwe2.weakness import Weakness as DBWeakness
 from django.contrib.auth import get_user_model
+from django.contrib.auth.models import Group
 from django.contrib.auth.models import UserManager
 from django.core import exceptions
 from django.core.exceptions import ValidationError
@@ -1472,6 +1473,10 @@ 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

vulnerabilities/throttling.py

@@ -6,18 +6,31 @@
 # See https://github.com/aboutcode-org/vulnerablecode for support or download.
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
+
 from rest_framework.exceptions import Throttled
 from rest_framework.throttling import UserRateThrottle
 from rest_framework.views import exception_handler
 
 
-class StaffUserRateThrottle(UserRateThrottle):
+class GroupUserRateThrottle(UserRateThrottle):
+    scope = "bronze"
+
     def allow_request(self, request, view):
-        """
-        Do not apply throttling for superusers and admins.
-        """
-        if request.user.is_superuser or request.user.is_staff:
-            return True
+        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]):
+                return True
+
+            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)
 
         return super().allow_request(request, view)
 

vulnerablecode/settings.py

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