Address review comments

Signed-off-by: Tushar Goel <[email protected]>

Author: TG1999

vulnerabilities/api.py

@@ -23,13 +23,7 @@
 from vulnerabilities.models import VulnerabilityReference
 from vulnerabilities.models import VulnerabilitySeverity
 from vulnerabilities.models import get_purl_query_lookups
-from vulnerabilities.throttling import AliasesAPIThrottle
-from vulnerabilities.throttling import BulkSearchCPEAPIThrottle
-from vulnerabilities.throttling import BulkSearchPackagesAPIThrottle
-from vulnerabilities.throttling import CPEAPIThrottle
-from vulnerabilities.throttling import PackagesAPIThrottle
-from vulnerabilities.throttling import VulnerabilitiesAPIThrottle
-from vulnerabilities.throttling import VulnerablePackagesAPIThrottle
+from vulnerabilities.throttling import StaffUserRateThrottle
 
 
 class VulnerabilitySeveritySerializer(serializers.ModelSerializer):
@@ -227,18 +221,11 @@ class PackageViewSet(viewsets.ReadOnlyModelViewSet):
     serializer_class = PackageSerializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = PackageFilterSet
-
-    def get_throttles(self):
-        if self.action == "bulk_search":
-            throttle_classes = [BulkSearchPackagesAPIThrottle]
-        elif self.action == "all":
-            throttle_classes = [VulnerablePackagesAPIThrottle]
-        else:
-            throttle_classes = [PackagesAPIThrottle]
-        return [throttle() for throttle in throttle_classes]
+    throttle_classes = [StaffUserRateThrottle]
+    throttle_scope = "packages"
 
     # TODO: Fix the swagger documentation for this endpoint
-    @action(detail=False, methods=["post"])
+    @action(detail=False, methods=["post"], throttle_scope="bulk_search_packages")
     def bulk_search(self, request):
         """
         See https://github.com/nexB/vulnerablecode/pull/369#issuecomment-796877606 for docs
@@ -270,7 +257,7 @@ def bulk_search(self, request):
 
         return Response(response)
 
-    @action(detail=False, methods=["get"])
+    @action(detail=False, methods=["get"], throttle_scope="vulnerable_packages")
     def all(self, request):
         """
         Return all the vulnerable Package URLs.
@@ -318,7 +305,8 @@ def get_queryset(self):
     serializer_class = VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = VulnerabilityFilterSet
-    throttle_classes = [VulnerabilitiesAPIThrottle]
+    throttle_classes = [StaffUserRateThrottle]
+    throttle_scope = "vulnerabilities"
 
 
 class CPEFilterSet(filters.FilterSet):
@@ -335,16 +323,11 @@ class CPEViewSet(viewsets.ReadOnlyModelViewSet):
     ).distinct()
     serializer_class = VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
+    throttle_classes = [StaffUserRateThrottle]
     filterset_class = CPEFilterSet
+    throttle_scope = "cpes"
 
-    def get_throttles(self):
-        if self.action == "bulk_search":
-            throttle_classes = [BulkSearchCPEAPIThrottle]
-        else:
-            throttle_classes = [CPEAPIThrottle]
-        return [throttle() for throttle in throttle_classes]
-
-    @action(detail=False, methods=["post"])
+    @action(detail=False, methods=["post"], throttle_scope="bulk_search_cpes")
     def bulk_search(self, request):
         """
         This endpoint is used to search for vulnerabilities by more than one CPE.
@@ -381,4 +364,5 @@ class AliasViewSet(viewsets.ReadOnlyModelViewSet):
     serializer_class = VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = AliasFilterSet
-    throttle_classes = [AliasesAPIThrottle]
+    throttle_classes = [StaffUserRateThrottle]
+    throttle_scope = "aliases"

vulnerabilities/throttling.py

@@ -6,14 +6,10 @@
 # See https://github.com/nexB/vulnerablecode for support or download.
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
+from rest_framework.throttling import ScopedRateThrottle
 
-from django.contrib.auth import get_user_model
-from rest_framework.throttling import SimpleRateThrottle
 
-User = get_user_model()
-
-
-class StaffUserRateThrottle(SimpleRateThrottle):
+class StaffUserRateThrottle(ScopedRateThrottle):
     def allow_request(self, request, view):
         """
         Do not apply throttling for superusers and admins.
@@ -22,42 +18,3 @@ def allow_request(self, request, view):
             return True
 
         return super().allow_request(request, view)
-
-    def get_cache_key(self, request, view):
-        """
-        Return the cache key to use for this request.
-        """
-        if request.user.is_authenticated:
-            ident = request.user.pk
-        else:
-            ident = self.get_ident(request)
-
-        return self.cache_format % {"scope": self.scope, "ident": ident}
-
-
-class VulnerablePackagesAPIThrottle(StaffUserRateThrottle):
-    scope = "vulnerable_packages"
-
-
-class BulkSearchPackagesAPIThrottle(StaffUserRateThrottle):
-    scope = "bulk_search_packages"
-
-
-class PackagesAPIThrottle(StaffUserRateThrottle):
-    scope = "packages"
-
-
-class VulnerabilitiesAPIThrottle(StaffUserRateThrottle):
-    scope = "vulnerabilities"
-
-
-class AliasesAPIThrottle(StaffUserRateThrottle):
-    scope = "aliases"
-
-
-class CPEAPIThrottle(StaffUserRateThrottle):
-    scope = "cpes"
-
-
-class BulkSearchCPEAPIThrottle(StaffUserRateThrottle):
-    scope = "bulk_search_cpes"

vulnerablecode/settings.py

@@ -150,23 +150,28 @@
 
 LOGIN_REDIRECT_URL = "/"
 LOGOUT_REDIRECT_URL = "/"
-TEST_PACKAGE_THROTTLING_RATE = None
-TEST_BULK_SEARCH_PACKAGE_THROTTLING_RATE = None
-TEST_ALL_VULNERABLE_PACKAGE_THROTTLING_RATE = None
-TEST_VULNERABILITIES_THROTTLING_RATE = None
-TEST_CPES_THROTTLING_RATE = None
-TEST_BULK_SEARCH_CPES_THROTTLING_RATE = None
-TEST_ALIASES_THROTTLING_RATE = None
+
+REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {
+    "vulnerable_packages": "1/hour",
+    "bulk_search_packages": "5/hour",
+    "packages": "10/minute",
+    "vulnerabilities": "10/minute",
+    "aliases": "5/minute",
+    "cpes": "5/minute",
+    "bulk_search_cpes": "5/hour",
+}
 
 if IS_TESTS:
     VULNERABLECODEIO_REQUIRE_AUTHENTICATION = True
-    TEST_PACKAGE_THROTTLING_RATE = "10/day"
-    TEST_BULK_SEARCH_PACKAGE_THROTTLING_RATE = "6/day"
-    TEST_ALL_VULNERABLE_PACKAGE_THROTTLING_RATE = "1/day"
-    TEST_VULNERABILITIES_THROTTLING_RATE = "8/day"
-    TEST_CPES_THROTTLING_RATE = "4/day"
-    TEST_BULK_SEARCH_CPES_THROTTLING_RATE = "5/day"
-    TEST_ALIASES_THROTTLING_RATE = "2/day"
+    REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {
+        "vulnerable_packages": "1/day",
+        "bulk_search_packages": "6/day",
+        "packages": "10/day",
+        "vulnerabilities": "8/day",
+        "aliases": "2/day",
+        "cpes": "4/day",
+        "bulk_search_cpes": "5/day",
+    }
 
 
 USE_L10N = True
@@ -202,15 +207,7 @@
     "DEFAULT_THROTTLE_CLASSES": [
         "vulnerabilities.throttling.StaffUserRateThrottle",
     ],
-    "DEFAULT_THROTTLE_RATES": {
-        "vulnerable_packages": TEST_ALL_VULNERABLE_PACKAGE_THROTTLING_RATE or "1/hour",
-        "bulk_search_packages": TEST_BULK_SEARCH_PACKAGE_THROTTLING_RATE or "5/hour",
-        "packages": TEST_PACKAGE_THROTTLING_RATE or "10/minute",
-        "vulnerabilities": TEST_VULNERABILITIES_THROTTLING_RATE or "10/minute",
-        "aliases": TEST_ALIASES_THROTTLING_RATE or "5/minute",
-        "cpes": TEST_CPES_THROTTLING_RATE or "5/minute",
-        "bulk_search_cpes": TEST_BULK_SEARCH_CPES_THROTTLING_RATE or "5/hour",
-    },
+    "DEFAULT_THROTTLE_RATES": REST_FRAMEWORK_DEFAULT_THROTTLE_RATES,
     "DEFAULT_PAGINATION_CLASS": "vulnerabilities.pagination.SmallResultSetPagination",
     # Limit the load on the Database returning a small number of records by default. https://github.com/nexB/vulnerablecode/issues/819
     "PAGE_SIZE": 10,