Move restrict_queryset() function to RestrictedQuerySet

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)