Move restrict_queryset() function to RestrictedQuerySet

This commit is contained in:
Jeremy Stretch 2020-05-29 15:09:08 -04:00
parent 58989b85c8
commit 5b6a6fb63e
4 changed files with 34 additions and 22 deletions

View File

@ -16,7 +16,6 @@ from rest_framework.serializers import Field, ModelSerializer, ValidationError
from rest_framework.viewsets import ModelViewSet as _ModelViewSet from rest_framework.viewsets import ModelViewSet as _ModelViewSet
from netbox.api import TokenPermissions from netbox.api import TokenPermissions
from utilities.permissions import restrict_queryset
from .utils import dict_to_filter_params, dynamic_import 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 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): def dispatch(self, request, *args, **kwargs):
logger = logging.getLogger('netbox.api.views.ModelViewSet') logger = logging.getLogger('netbox.api.views.ModelViewSet')

View File

@ -34,20 +34,3 @@ def resolve_permission(name):
raise ValueError(f"Unknown app/model for {name}") raise ValueError(f"Unknown app/model for {name}")
return content_type, action 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)

View File

@ -1,3 +1,6 @@
from django.db.models import Q, QuerySet
class DummyQuerySet: class DummyQuerySet:
""" """
A fake QuerySet that can be used to cache relationships to objects that have been deleted. 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): def all(self):
return self._cache 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)

View File

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