From 5b6a6fb63e2d5a67649a9db40450b1e835cda561 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Fri, 29 May 2020 15:09:08 -0400
Subject: [PATCH] Move restrict_queryset() function to RestrictedQuerySet

---
 netbox/utilities/api.py         |  5 ++---
 netbox/utilities/permissions.py | 17 -----------------
 netbox/utilities/querysets.py   | 30 ++++++++++++++++++++++++++++++
 netbox/utilities/views.py       |  4 ++--
 4 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/netbox/utilities/api.py b/netbox/utilities/api.py
index ef2650535..ac21d298c 100644
--- a/netbox/utilities/api.py
+++ b/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
-        self.queryset = restrict_queryset(self.queryset, request.user, permission_required)
+        # Restrict the view's QuerySet to allow only the permitted objects
+        self.queryset = self.queryset.restrict(request.user, permission_required)
 
     def dispatch(self, request, *args, **kwargs):
         logger = logging.getLogger('netbox.api.views.ModelViewSet')
diff --git a/netbox/utilities/permissions.py b/netbox/utilities/permissions.py
index be5c0189e..697e18828 100644
--- a/netbox/utilities/permissions.py
+++ b/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)
diff --git a/netbox/utilities/querysets.py b/netbox/utilities/querysets.py
index 34b7a0cf3..36460310e 100644
--- a/netbox/utilities/querysets.py
+++ b/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)
diff --git a/netbox/utilities/views.py b/netbox/utilities/views.py
index a86b5ccc5..fed774812 100644
--- a/netbox/utilities/views.py
+++ b/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