Skip to content

Commit 8362ef4

Browse files
dogboatMaffooch
authored andcommitted
update and optimize prefetcher
1 parent e780723 commit 8362ef4

5 files changed

Lines changed: 725 additions & 11 deletions

File tree

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
"""
2+
RBAC registry for the ``?prefetch=`` path.
3+
4+
The prefetch mixins resolve a query-string field name through ``getattr`` on a
5+
model instance, find a serializer for the resolved related model, and return
6+
the serialized representation. This module allows us to specify authorization
7+
checks on the related objects when serializing.
8+
9+
``_Prefetcher`` filters every resolved related object through the registered
10+
queryset before serializing it. If no policy is registered for a model, the
11+
field is omitted from the response.
12+
"""
13+
14+
from collections.abc import Callable
15+
16+
from django.db.models import Model, Q, QuerySet
17+
18+
from dojo.authorization.authorization import user_has_configuration_permission
19+
from dojo.models import Engagement, Finding, Notes, Test
20+
21+
_REGISTRY: dict[type[Model], Callable[[object], QuerySet]] = {}
22+
23+
24+
def discard_user(func):
25+
"""
26+
Adapter for auth helpers that don't accept a ``user`` parameter --
27+
wraps them so they can be passed to ``register()`` like any other policy.
28+
"""
29+
30+
def wrapper(*args, user, **kwargs):
31+
return func(*args, **kwargs)
32+
33+
return wrapper
34+
35+
36+
def register(model: type[Model], func: Callable, *args, **kwargs) -> None:
37+
"""Register a policy for ``model``. At lookup, ``func`` is invoked as ``func(*args, user=<requesting user>, **kwargs)``."""
38+
39+
def policy(user):
40+
return func(*args, user=user, **kwargs)
41+
42+
_REGISTRY[model] = policy
43+
44+
45+
def get_authorized_queryset(model: type[Model], user) -> QuerySet | None:
46+
"""
47+
Return the queryset of ``model`` instances visible to ``user``.
48+
49+
Returns ``None`` when no policy has been registered. ``_Prefetcher``
50+
treats ``None`` as "deny" and omits the field from the response.
51+
"""
52+
if policy := _REGISTRY.get(model):
53+
return policy(user)
54+
return None
55+
56+
57+
def superuser_only(model: type[Model], user) -> QuerySet:
58+
"""
59+
Policy for models whose top-level ViewSet enforces ``IsSuperUser``
60+
(strict ``request.user.is_superuser`` check). Only superusers pass.
61+
"""
62+
if user is not None and getattr(user, "is_superuser", False):
63+
return model.objects.all()
64+
return model.objects.none()
65+
66+
67+
def django_view_perm(model: type[Model], user) -> QuerySet:
68+
"""
69+
Policy for models whose top-level ViewSet gates on DRF's ``DjangoModelPermissions``.
70+
71+
Passes all superusers and any user holding ``<app_label>.view_<model_name>``.
72+
"""
73+
if user is None:
74+
return model.objects.none()
75+
perm = f"{model._meta.app_label}.view_{model._meta.model_name}"
76+
if user.has_perm(perm):
77+
return model.objects.all()
78+
return model.objects.none()
79+
80+
81+
def dojo_view_perm(model: type[Model], user) -> QuerySet:
82+
"""
83+
Policy for models whose top-level ViewSet gates on a DefectDojo
84+
``BaseDjangoModelPermission`` subclass that requires GET=view.
85+
86+
Passes all superusers and staff users and any user holding ``<app_label>.view_<model_name>``.
87+
"""
88+
if user is None:
89+
return model.objects.none()
90+
perm = f"{model._meta.app_label}.view_{model._meta.model_name}"
91+
if user_has_configuration_permission(user, perm):
92+
return model.objects.all()
93+
return model.objects.none()
94+
95+
96+
def authenticated_only(model: type[Model], user) -> QuerySet:
97+
"""Policy for models whose top-level ViewSet is reachable by any authenticated user."""
98+
if user is not None and getattr(user, "is_authenticated", False):
99+
return model.objects.all()
100+
return model.objects.none()
101+
102+
103+
def children_via_parent(child_model, parent_model, parent_field, *, user) -> QuerySet:
104+
"""
105+
Authorize ``child_model`` by deferring to the policy registered for
106+
``parent_model`` -- the child is visible iff the parent it points to via
107+
``parent_field`` is visible. Used for models that don't have their own
108+
``get_authorized_*`` helper but logically inherit authorization from a
109+
parent (e.g. ``BurpRawRequestResponse`` -> ``Finding`` via ``finding``).
110+
"""
111+
if (parent_qs := get_authorized_queryset(parent_model, user)) is not None:
112+
return child_model.objects.filter(**{f"{parent_field}__in": parent_qs})
113+
return child_model.objects.none()
114+
115+
116+
def notes_policy(user) -> QuerySet:
117+
"""
118+
Authorization for the ``Notes`` model.
119+
120+
Allows note viewership as follows:
121+
* superuser: every note
122+
* anyone else: a note is visible iff
123+
(its attached Finding / Test / Engagement is visible to ``user``)
124+
AND (the note is non-private OR ``user`` authored it).
125+
"""
126+
if user is None:
127+
return Notes.objects.none()
128+
if getattr(user, "is_superuser", False):
129+
return Notes.objects.all()
130+
131+
# Helper method to avoid unnecessary queryset fetching
132+
def _qs_or_none(model, u):
133+
qs = get_authorized_queryset(model, u)
134+
return model.objects.none() if qs is None else qs
135+
136+
finding_qs = _qs_or_none(Finding, user)
137+
test_qs = _qs_or_none(Test, user)
138+
engagement_qs = _qs_or_none(Engagement, user)
139+
140+
parent_visible = Q(finding__in=finding_qs) | Q(test__in=test_qs) | Q(engagement__in=engagement_qs)
141+
return Notes.objects.filter(
142+
parent_visible,
143+
Q(private=False) | Q(author=user),
144+
).distinct()

dojo/api_v2/prefetch/mixins.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ def list(self, request, *args, **kwargs):
99
prefetch_params = request.GET.get("prefetch", "")
1010
prefetch_params = prefetch_params.split(",") if "," in prefetch_params else request.GET.getlist("prefetch")
1111

12-
prefetcher = _Prefetcher()
12+
prefetcher = _Prefetcher(request=request)
1313

1414
# Apply the same operations as the standard list method defined in the
1515
# django rest framework
@@ -35,7 +35,7 @@ def retrieve(self, request, *args, **kwargs):
3535
prefetch_params = request.GET.get("prefetch", "")
3636
prefetch_params = prefetch_params.split(",") if "," in prefetch_params else request.GET.getlist("prefetch")
3737

38-
prefetcher = _Prefetcher()
38+
prefetcher = _Prefetcher(request=request)
3939

4040
entry = self.get_object()
4141
serializer = self.get_serializer()

dojo/api_v2/prefetch/prefetcher.py

Lines changed: 43 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,11 @@
44
from django.conf import settings
55
from rest_framework.serializers import ModelSerializer
66

7+
from dojo.api_v2.prefetch import (
8+
registrations as _registrations, # noqa: F401 -- side-effect import populates the RBAC registry
9+
)
710
from dojo.api_v2.prefetch import utils
11+
from dojo.api_v2.prefetch.authorized_querysets import get_authorized_queryset
812
from dojo.location.api.serializers import LocationFindingReferenceSerializer, LocationSerializer
913
from dojo.location.models import Location, LocationFindingReference
1014
from dojo.models import FileUpload, Finding
@@ -36,7 +40,8 @@ def _is_model_serializer(obj):
3640
# We process all the serializers found in the module SERIALIZER_DEFS_MODULE. We restrict the scope to avoid
3741
# processing all the classes in the symbol table
3842
available_serializers = inspect.getmembers(
39-
sys.modules[SERIALIZER_DEFS_MODULE], _is_model_serializer,
43+
sys.modules[SERIALIZER_DEFS_MODULE],
44+
_is_model_serializer,
4045
)
4146

4247
for _, serializer in available_serializers:
@@ -51,9 +56,10 @@ def _is_model_serializer(obj):
5156

5257
return serializers
5358

54-
def __init__(self):
59+
def __init__(self, request=None):
5560
self._serializers = _Prefetcher._build_serializers()
5661
self._prefetch_data = {}
62+
self._request = request
5763

5864
def _find_serializer(self, field_type):
5965
"""
@@ -136,6 +142,8 @@ def _prefetch(self, entry, fields_to_fetch):
136142
field_to_fetch (list[string]): fields to prefetch
137143
138144
"""
145+
user = getattr(self._request, "user", None) if self._request else None
146+
139147
for field_to_fetch in fields_to_fetch:
140148
# Get the field from the instance
141149
field_value, many = self.get_field_value(entry, field_to_fetch)
@@ -148,16 +156,42 @@ def _prefetch(self, entry, fields_to_fetch):
148156
if extra_serializer is None:
149157
continue
150158

151-
field_data = extra_serializer(many=many).to_representation(
159+
# Authorization gate: only serialize objects from the requester's
160+
# authorized queryset for this model. Deny-by-default -- models
161+
# with no registered policy are omitted from the prefetch payload.
162+
authorized_qs = get_authorized_queryset(model_type, user)
163+
if authorized_qs is None:
164+
continue
165+
166+
# Match the legacy contract: the field key always appears in the
167+
# prefetch payload (possibly empty) once we have a policy. Tests
168+
# and clients can rely on the key being present whenever the
169+
# primary field is non-null on the entry.
170+
self._prefetch_data.setdefault(field_to_fetch, {})
171+
172+
# Check related object authorizations
173+
if many:
174+
related_qs = field_value.all() if hasattr(field_value, "all") else field_value
175+
related_ids = list(related_qs.values_list("pk", flat=True))
176+
if not related_ids:
177+
continue
178+
# Set `field_value` (what will be serialized later) to the set of objects the user has permission to
179+
field_value = authorized_qs.filter(pk__in=related_ids)
180+
if not field_value.exists():
181+
continue
182+
# Only a single related item; check it's in the set of authorized objects. If not, skip it.
183+
elif not authorized_qs.filter(pk=field_value.pk).exists():
184+
continue
185+
186+
serializer_kwargs = {"many": many}
187+
if self._request is not None:
188+
# Add in the request for the serializers to use if they want
189+
serializer_kwargs["context"] = {"request": self._request}
190+
field_data = extra_serializer(**serializer_kwargs).to_representation(
152191
field_value,
153192
)
154193
# For convenience in processing we store the field data in a list
155-
field_data_list = (
156-
field_data if isinstance(field_data, list) else [field_data]
157-
)
158-
159-
if field_to_fetch not in self._prefetch_data:
160-
self._prefetch_data[field_to_fetch] = {}
194+
field_data_list = field_data if isinstance(field_data, list) else [field_data]
161195

162196
# Should not fail as django always generate an id field
163197
for data in field_data_list:

0 commit comments

Comments
 (0)