querysets.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. from django.db.models import Q, QuerySet
  2. from utilities.permissions import permission_is_exempt
  3. class DummyQuerySet:
  4. """
  5. A fake QuerySet that can be used to cache relationships to objects that have been deleted.
  6. """
  7. def __init__(self, queryset):
  8. self._cache = [obj for obj in queryset.all()]
  9. def all(self):
  10. return self._cache
  11. class RestrictedQuerySet(QuerySet):
  12. def restrict(self, user, action):
  13. """
  14. Filter the QuerySet to return only objects on which the specified user has been granted the specified
  15. permission.
  16. :param user: User instance
  17. :param action: The action which must be permitted (e.g. "view" for "dcim.view_site")
  18. """
  19. # Resolve the full name of the required permission
  20. app_label = self.model._meta.app_label
  21. model_name = self.model._meta.model_name
  22. permission_required = f'{app_label}.{action}_{model_name}'
  23. # Bypass restriction for superusers and exempt views
  24. if user.is_superuser or permission_is_exempt(permission_required):
  25. return self
  26. # User is anonymous or has not been granted the requisite permission
  27. if not user.is_authenticated or permission_required not in user.get_all_permissions():
  28. return self.none()
  29. # Filter the queryset to include only objects with allowed attributes
  30. attrs = Q()
  31. for perm_attrs in user._object_perm_cache[permission_required]:
  32. if perm_attrs:
  33. attrs |= Q(**perm_attrs)
  34. return self.filter(attrs)