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