wip2

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

Author: keshav-space

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/throttling.py

@@ -22,10 +22,11 @@ def allow_request(self, request, view):
             if user.is_superuser or user.is_staff:
                 return True
 
-            if user.groups.filter(name="gold").exists():
+            user_groups = user.groups.all()
+            if any([group.name == "gold" for group in user_groups]):
                 return True
 
-            if user.groups.filter(name="silver").exists():
+            if any([group.name == "silver" for group in user_groups]):
                 self.scope = "silver"
 
         self.rate = self.THROTTLE_RATES.get(self.scope)