Merge branch 'main' into ngi-acknowledgements

Author: pombredanne

vulnerabilities/api.py

@@ -16,6 +16,8 @@
 from rest_framework import viewsets
 from rest_framework.decorators import action
 from rest_framework.response import Response
+from rest_framework.throttling import AnonRateThrottle
+from rest_framework.throttling import UserRateThrottle
 
 from vulnerabilities.models import Alias
 from vulnerabilities.models import Package
@@ -231,11 +233,10 @@ class PackageViewSet(viewsets.ReadOnlyModelViewSet):
     serializer_class = PackageSerializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = PackageFilterSet
-    throttle_classes = [StaffUserRateThrottle]
-    throttle_scope = "packages"
+    throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]
 
     # TODO: Fix the swagger documentation for this endpoint
-    @action(detail=False, methods=["post"], throttle_scope="bulk_search_packages")
+    @action(detail=False, methods=["post"])
     def bulk_search(self, request):
         """
         Lookup for vulnerable packages using many Package URLs at once.
@@ -289,7 +290,7 @@ def bulk_search(self, request):
         vulnerable_purls = [str(package.package_url) for package in vulnerable_purls]
         return Response(data=vulnerable_purls)
 
-    @action(detail=False, methods=["get"], throttle_scope="vulnerable_packages")
+    @action(detail=False, methods=["get"])
     def all(self, request):
         """
         Return the Package URLs of all packages known to be vulnerable.
@@ -341,8 +342,7 @@ def get_queryset(self):
     serializer_class = VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = VulnerabilityFilterSet
-    throttle_classes = [StaffUserRateThrottle]
-    throttle_scope = "vulnerabilities"
+    throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]
 
 
 class CPEFilterSet(filters.FilterSet):
@@ -363,11 +363,10 @@ class CPEViewSet(viewsets.ReadOnlyModelViewSet):
     ).distinct()
     serializer_class = VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
-    throttle_classes = [StaffUserRateThrottle]
+    throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]
     filterset_class = CPEFilterSet
-    throttle_scope = "cpes"
 
-    @action(detail=False, methods=["post"], throttle_scope="bulk_search_cpes")
+    @action(detail=False, methods=["post"])
     def bulk_search(self, request):
         """
         Lookup for vulnerabilities using many CPEs at once.
@@ -409,5 +408,4 @@ class AliasViewSet(viewsets.ReadOnlyModelViewSet):
     serializer_class = VulnerabilitySerializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = AliasFilterSet
-    throttle_classes = [StaffUserRateThrottle]
-    throttle_scope = "aliases"
+    throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]

vulnerabilities/tests/test_auth.py

@@ -29,10 +29,11 @@ def setUp(self):
         self.staff_csrf_client = APIClient(enforce_csrf_checks=True)
         self.staff_csrf_client.credentials(HTTP_AUTHORIZATION=self.staff_auth)
 
-    def test_packages_endpoint_throttling(self):
+        self.csrf_client_anon = APIClient(enforce_csrf_checks=True)
+        self.csrf_client_anon_1 = APIClient(enforce_csrf_checks=True)
 
-        # A basic user can only access /packages endpoint 10 times a day
-        for i in range(0, 10):
+    def test_package_endpoint_throttling(self):
+        for i in range(0, 20):
             response = self.csrf_client.get("/api/packages")
             self.assertEqual(response.status_code, 200)
             response = self.staff_csrf_client.get("/api/packages")
@@ -46,122 +47,36 @@ def test_packages_endpoint_throttling(self):
         # 200 - staff user can access API unlimited times
         self.assertEqual(response.status_code, 200)
 
-    def test_cpes_endpoint_throttling(self):
-
-        # A basic user can only access /cpes endpoint 4 times a day
-        for i in range(0, 4):
-            response = self.csrf_client.get("/api/cpes")
-            self.assertEqual(response.status_code, 200)
-            response = self.staff_csrf_client.get("/api/cpes")
-            self.assertEqual(response.status_code, 200)
-
-        response = self.csrf_client.get("/api/cpes")
-        # 429 - too many requests for basic user
-        self.assertEqual(response.status_code, 429)
-
-        response = self.staff_csrf_client.get("/api/cpes", format="json")
-        # 200 - staff user can access API unlimited times
-        self.assertEqual(response.status_code, 200)
-
-    def test_all_vulnerable_packages_endpoint_throttling(self):
-
-        # A basic user can only access /packages/all 1 time a day
-        for i in range(0, 1):
-            response = self.csrf_client.get("/api/packages/all")
-            self.assertEqual(response.status_code, 200)
-            response = self.staff_csrf_client.get("/api/packages/all")
-            self.assertEqual(response.status_code, 200)
-
-        response = self.csrf_client.get("/api/packages/all")
-        # 429 - too many requests for basic user
-        self.assertEqual(response.status_code, 429)
-
-        response = self.staff_csrf_client.get("/api/packages/all", format="json")
-        # 200 - staff user can access API unlimited times
-        self.assertEqual(response.status_code, 200)
-
-    def test_vulnerabilities_endpoint_throttling(self):
-
-        # A basic user can only access /vulnerabilities 8 times a day
-        for i in range(0, 8):
-            response = self.csrf_client.get("/api/vulnerabilities")
-            self.assertEqual(response.status_code, 200)
-            response = self.staff_csrf_client.get("/api/vulnerabilities")
+        # A anonymous user can only access /packages endpoint 10 times a day
+        for i in range(0, 10):
+            print(i)
+            response = self.csrf_client_anon.get("/api/packages")
             self.assertEqual(response.status_code, 200)
 
-        response = self.csrf_client.get("/api/vulnerabilities")
-        # 429 - too many requests for basic user
+        response = self.csrf_client_anon.get("/api/packages")
+        # 429 - too many requests for anon user
         self.assertEqual(response.status_code, 429)
+        self.assertEqual(
+            response.data.get("message"),
+            "Your request has been throttled. Please contact [email protected]",
+        )
 
-        response = self.staff_csrf_client.get("/api/vulnerabilities", format="json")
-        # 200 - staff user can access API unlimited times
-        self.assertEqual(response.status_code, 200)
-
-    def test_aliases_endpoint_throttling(self):
-
-        # A basic user can only access /alias 2 times a day
-        for i in range(0, 2):
-            response = self.csrf_client.get("/api/aliases")
-            self.assertEqual(response.status_code, 200)
-            response = self.staff_csrf_client.get("/api/aliases")
-            self.assertEqual(response.status_code, 200)
-
-        response = self.csrf_client.get("/api/aliases")
-        # 429 - too many requests for basic user
+        response = self.csrf_client_anon.get("/api/vulnerabilities")
+        # 429 - too many requests for anon user
         self.assertEqual(response.status_code, 429)
+        self.assertEqual(
+            response.data.get("message"),
+            "Your request has been throttled. Please contact [email protected]",
+        )
 
-        response = self.staff_csrf_client.get("/api/aliases", format="json")
-        # 200 - staff user can access API unlimited times
-        self.assertEqual(response.status_code, 200)
-
-    def test_bulk_search_packages_endpoint_throttling(self):
         data = json.dumps({"purls": ["pkg:foo/bar"]})
 
-        # A basic user can only access /packages/bulk_search 6 times a day
-        for i in range(0, 6):
-            response = self.csrf_client.post(
-                "/api/packages/bulk_search", data=data, content_type="application/json"
-            )
-            self.assertEqual(response.status_code, 200)
-            response = self.staff_csrf_client.post(
-                "/api/packages/bulk_search", data=data, content_type="application/json"
-            )
-            self.assertEqual(response.status_code, 200)
-
-        response = self.csrf_client.post(
+        response = self.csrf_client_anon.post(
             "/api/packages/bulk_search", data=data, content_type="application/json"
         )
-        # 429 - too many requests for basic user
+        # 429 - too many requests for anon user
         self.assertEqual(response.status_code, 429)
-
-        response = self.staff_csrf_client.post(
-            "/api/packages/bulk_search", data=data, content_type="application/json"
+        self.assertEqual(
+            response.data.get("message"),
+            "Your request has been throttled. Please contact [email protected]",
         )
-        # 200 - staff user can access API unlimited times
-        self.assertEqual(response.status_code, 200)
-
-    def test_bulk_search_cpes_endpoint_throttling(self):
-        data = json.dumps({"cpes": ["cpe:foo/bar"]})
-
-        # A basic user can only access /cpes/bulk_search 5 times a day
-        for i in range(0, 5):
-            response = self.csrf_client.post(
-                "/api/cpes/bulk_search", data=data, content_type="application/json"
-            )
-            self.assertEqual(response.status_code, 200)
-            response = self.staff_csrf_client.post(
-                "/api/cpes/bulk_search", data=data, content_type="application/json"
-            )
-            self.assertEqual(response.status_code, 200)
-
-        response = self.csrf_client.post(
-            "/api/cpes/bulk_search", data=data, content_type="application/json"
-        )
-        # 429 - too many requests for basic user
-        self.assertEqual(response.status_code, 429)
-
-        response = self.staff_csrf_client.post(
-            "/api/cpes/bulk_search", data=data, content_type="application/json"
-        )
-        # 200 - staff user can access API unlimited times
-        self.assertEqual(response.status_code, 200)

vulnerabilities/tests/test_throttling.py

@@ -6,10 +6,12 @@
 # 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 rest_framework.exceptions import Throttled
+from rest_framework.throttling import UserRateThrottle
+from rest_framework.views import exception_handler
 
 
-class StaffUserRateThrottle(ScopedRateThrottle):
+class StaffUserRateThrottle(UserRateThrottle):
     def allow_request(self, request, view):
         """
         Do not apply throttling for superusers and admins.
@@ -18,3 +20,19 @@ def allow_request(self, request, view):
             return True
 
         return super().allow_request(request, view)
+
+
+def throttled_exception_handler(exception, context):
+    """
+    Return this response whenever a request has been throttled
+    """
+
+    response = exception_handler(exception, context)
+
+    if isinstance(exception, Throttled):
+        response_data = {
+            "message": "Your request has been throttled. Please contact [email protected]"
+        }
+        response.data = response_data
+
+    return response

vulnerabilities/throttling.py

@@ -171,35 +171,11 @@
 LOGIN_REDIRECT_URL = "/"
 LOGOUT_REDIRECT_URL = "/"
 
-REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {
-    "vulnerable_packages": env.str(
-        "VULNERABLECODE_ALL_VULNERABLE_PACKAGES_THROTTLING_RATE", default="1/hour"
-    ),
-    "bulk_search_packages": env.str(
-        "VULNERABLECODE_BULK_SEARCH_PACKAGE_THROTTLING_RATE", default="5/hour"
-    ),
-    "packages": env.str("VULNERABLECODE_PACKAGES_SEARCH_THROTTLING_RATE", default="10/minute"),
-    "vulnerabilities": env.str(
-        "VULNERABLECODE_VULNERABILITIES_SEARCH_THROTTLING_RATE", default="10/minute"
-    ),
-    "aliases": env.str("VULNERABLECODE_ALIASES_SEARCH_THROTTLING_RATE", default="5/minute"),
-    "cpes": env.str("VULNERABLECODE_CPE_SEARCH_THROTTLING_RATE", default="5/minute"),
-    "bulk_search_cpes": env.str(
-        "VULNERABLECODE_BULK_SEARCH_CPE_THROTTLING_RATE", default="5/minute"
-    ),
-}
+REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {"anon": "3600/hour", "user": "10800/hour"}
 
 if IS_TESTS:
-    VULNERABLECODEIO_REQUIRE_AUTHENTICATION = True
-    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",
-    }
+    VULNERABLECODEIO_REQUIRE_AUTHENTICATION = False
+    REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {"anon": "10/day", "user": "20/day"}
 
 
 USE_L10N = True
@@ -237,8 +213,11 @@
     ),
     "DEFAULT_THROTTLE_CLASSES": [
         "vulnerabilities.throttling.StaffUserRateThrottle",
+        "rest_framework.throttling.AnonRateThrottle",
+        "rest_framework.throttling.UserRateThrottle",
     ],
     "DEFAULT_THROTTLE_RATES": REST_FRAMEWORK_DEFAULT_THROTTLE_RATES,
+    "EXCEPTION_HANDLER": "vulnerabilities.throttling.throttled_exception_handler",
     "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,