Move restrict_queryset() function to RestrictedQuerySet

netbox/utilities/api.py

@@ -16,7 +16,6 @@ from rest_framework.serializers import Field, ModelSerializer, ValidationError
 from rest_framework.viewsets import ModelViewSet as _ModelViewSet
 
 from netbox.api import TokenPermissions
-from utilities.permissions import restrict_queryset
 from .utils import dict_to_filter_params, dynamic_import
 
 
@@ -339,8 +338,8 @@ class ModelViewSet(_ModelViewSet):
         }
         permission_required = TokenPermissions.perms_map[request.method][0] % kwargs
 
-        # Update the view's QuerySet to filter only the permitted objects
+        # Restrict the view's QuerySet to allow only the permitted objects
-        self.queryset = restrict_queryset(self.queryset, request.user, permission_required)
+        self.queryset = self.queryset.restrict(request.user, permission_required)
 
     def dispatch(self, request, *args, **kwargs):
         logger = logging.getLogger('netbox.api.views.ModelViewSet')

netbox/utilities/permissions.py

@@ -34,20 +34,3 @@ def resolve_permission(name):
         raise ValueError(f"Unknown app/model for {name}")
 
     return content_type, action
-
-
-def restrict_queryset(queryset, user, permission_required):
-    """
-    Filters a QuerySet to return only the objects on which the specified user has been granted the specified
-    permission.
-
-    :param queryset: Base QuerySet to be restricted
-    :param user: User instance
-    :param permission_required: Name of the required permission (e.g. "dcim.view_site")
-    """
-    obj_perm_attrs = user._object_perm_cache[permission_required]
-    attrs = Q()
-    for perm_attrs in obj_perm_attrs:
-        if perm_attrs:
-            attrs |= Q(**perm_attrs)
-    return queryset.filter(attrs)

netbox/utilities/querysets.py

@@ -1,3 +1,6 @@
+from django.db.models import Q, QuerySet
+
+
 class DummyQuerySet:
     """
     A fake QuerySet that can be used to cache relationships to objects that have been deleted.
@@ -7,3 +10,30 @@ class DummyQuerySet:
 
     def all(self):
         return self._cache
+
+
+class RestrictedQuerySet(QuerySet):
+
+    def restrict(self, user, permission_required):
+        """
+        Filter the QuerySet to return only objects on which the specified user has been granted the specified
+        permission.
+
+        :param queryset: Base QuerySet to be restricted
+        :param user: User instance
+        :param permission_required: Name of the required permission (e.g. "dcim.view_site")
+        """
+
+        # Determine what constraints (if any) have been placed on this user for this action and model
+        # TODO: Find a better way to ensure permissions are cached
+        if not hasattr(user, '_object_perm_cache'):
+            user.get_all_permisisons()
+        obj_perm_attrs = user._object_perm_cache[permission_required]
+
+        # Filter the queryset to include only objects with allowed attributes
+        attrs = Q()
+        for perm_attrs in obj_perm_attrs:
+            if perm_attrs:
+                attrs |= Q(**perm_attrs)
+
+        return self.filter(attrs)

netbox/utilities/views.py

@@ -28,7 +28,7 @@ from extras.models import CustomField, CustomFieldValue, ExportTemplate
 from extras.querysets import CustomFieldQueryset
 from utilities.exceptions import AbortTransaction
 from utilities.forms import BootstrapMixin, CSVDataField, TableConfigForm
-from utilities.permissions import get_permission_for_model, restrict_queryset
+from utilities.permissions import get_permission_for_model
 from utilities.utils import csv_format, prepare_cloned_fields
 from .error_handlers import handle_protectederror
 from .forms import ConfirmationForm, ImportForm
@@ -66,7 +66,7 @@ class ObjectPermissionRequiredMixin(AccessMixin):
 
         # Update the view's QuerySet to filter only the permitted objects
         if user.is_authenticated and not user.is_superuser:
-            self.queryset = restrict_queryset(self.queryset, user, permission_required)
+            self.queryset = self.queryset.restrict(user, permission_required)
 
         return True