1from django.db.models import Q, QuerySet 2 3from utilities.permissions import permission_is_exempt 4 5 6class RestrictedQuerySet(QuerySet): 7 8 def restrict(self, user, action='view'): 9 """ 10 Filter the QuerySet to return only objects on which the specified user has been granted the specified 11 permission. 12 13 :param user: User instance 14 :param action: The action which must be permitted (e.g. "view" for "dcim.view_site"); default is 'view' 15 """ 16 # Resolve the full name of the required permission 17 app_label = self.model._meta.app_label 18 model_name = self.model._meta.model_name 19 permission_required = f'{app_label}.{action}_{model_name}' 20 21 # Bypass restriction for superusers and exempt views 22 if user.is_superuser or permission_is_exempt(permission_required): 23 qs = self 24 25 # User is anonymous or has not been granted the requisite permission 26 elif not user.is_authenticated or permission_required not in user.get_all_permissions(): 27 qs = self.none() 28 29 # Filter the queryset to include only objects with allowed attributes 30 else: 31 attrs = Q() 32 for perm_attrs in user._object_perm_cache[permission_required]: 33 if type(perm_attrs) is list: 34 for p in perm_attrs: 35 attrs |= Q(**p) 36 elif perm_attrs: 37 attrs |= Q(**perm_attrs) 38 else: 39 # Any permission with null constraints grants access to _all_ instances 40 attrs = Q() 41 break 42 qs = self.filter(attrs) 43 44 return qs 45