wip

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/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
@@ -1435,6 +1436,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/tests/test_throttling.py

@@ -9,6 +9,7 @@
 
 import json
 
+from django.contrib.auth.models import Group
 from django.core.cache import cache
 from rest_framework.test import APIClient
 from rest_framework.test import APITestCase
@@ -23,12 +24,20 @@ def setUp(self):
         # See https://www.django-rest-framework.org/api-guide/throttling/#setting-up-the-cache
         cache.clear()
 
-        # create a basic user
+        # create a basic user (silver)
         self.user = ApiUser.objects.create_api_user(username="e@mail.com")
         self.auth = f"Token {self.user.auth_token.key}"
         self.csrf_client = APIClient(enforce_csrf_checks=True)
         self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)
 
+        # create user (gold)
+        self.gold_user = ApiUser.objects.create_api_user(username="g@mail.com")
+        gold, _ = Group.objects.get_or_create(name="gold")
+        self.gold_user.groups.add(gold)
+        self.gold_auth = f"Token {self.gold_user.auth_token.key}"
+        self.gold_csrf_client = APIClient(enforce_csrf_checks=True)
+        self.gold_csrf_client.credentials(HTTP_AUTHORIZATION=self.gold_auth)
+
         # create a staff user
         self.staff_user = ApiUser.objects.create_api_user(username="staff@mail.com", is_staff=True)
         self.staff_auth = f"Token {self.staff_user.auth_token.key}"
@@ -45,6 +54,12 @@ def test_package_endpoint_throttling(self):
             response = self.staff_csrf_client.get("/api/packages")
             self.assertEqual(response.status_code, 200)
 
+        for i in range(0, 25):
+            response = self.gold_csrf_client.get("/api/packages")
+            self.assertEqual(response.status_code, 200)
+            response = self.csrf_client.get("/api/packages")
+            self.assertEqual(response.status_code, 200)
+
         response = self.csrf_client.get("/api/packages")
         # 429 - too many requests for basic user
         self.assertEqual(response.status_code, 429)

vulnerabilities/throttling.py

@@ -6,18 +6,30 @@
 # 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
+
+            if user.groups.filter(name="gold").exists():
+                return True
+
+            if user.groups.filter(name="silver").exists():
+                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",