|
8 | 8 | # |
9 | 9 |
|
10 | 10 |
|
| 11 | +from django.db.models import Prefetch |
11 | 12 | from django_filters import rest_framework as filters |
12 | 13 | from drf_spectacular.utils import OpenApiParameter |
13 | 14 | from drf_spectacular.utils import extend_schema |
|
20 | 21 | from rest_framework.response import Response |
21 | 22 | from rest_framework.reverse import reverse |
22 | 23 |
|
23 | | -from vulnerabilities.api import PackageFilterSet |
24 | | -from vulnerabilities.api import VulnerabilitySeveritySerializer |
25 | 24 | from vulnerabilities.models import Package |
26 | 25 | from vulnerabilities.models import Vulnerability |
27 | 26 | from vulnerabilities.models import VulnerabilityReference |
@@ -195,7 +194,20 @@ class Meta: |
195 | 194 | ] |
196 | 195 |
|
197 | 196 | def get_affected_by_vulnerabilities(self, obj): |
198 | | - return [vuln.vulnerability_id for vuln in obj.affected_by_vulnerabilities.all()] |
| 197 | + """ |
| 198 | + Return a dictionary with vulnerabilities as keys and their details, including fixed_by_packages. |
| 199 | + """ |
| 200 | + result = {} |
| 201 | + for vuln in getattr(obj, "prefetched_affected_vulnerabilities", []): |
| 202 | + fixed_by_package = vuln.fixed_by_packages.first() |
| 203 | + purl = None |
| 204 | + if fixed_by_package: |
| 205 | + purl = fixed_by_package.package_url |
| 206 | + result[vuln.vulnerability_id] = { |
| 207 | + "vulnerability_id": vuln.vulnerability_id, |
| 208 | + "fixed_by_packages": purl, |
| 209 | + } |
| 210 | + return result |
199 | 211 |
|
200 | 212 | def get_fixing_vulnerabilities(self, obj): |
201 | 213 | # Ghost package should not fix any vulnerability. |
@@ -233,7 +245,13 @@ class PackageV2FilterSet(filters.FilterSet): |
233 | 245 |
|
234 | 246 |
|
235 | 247 | class PackageV2ViewSet(viewsets.ReadOnlyModelViewSet): |
236 | | - queryset = Package.objects.all() |
| 248 | + queryset = Package.objects.all().prefetch_related( |
| 249 | + Prefetch( |
| 250 | + "affected_by_vulnerabilities", |
| 251 | + queryset=Vulnerability.objects.prefetch_related("fixed_by_packages"), |
| 252 | + to_attr="prefetched_affected_vulnerabilities", |
| 253 | + ) |
| 254 | + ) |
237 | 255 | serializer_class = PackageV2Serializer |
238 | 256 | filter_backends = (filters.DjangoFilterBackend,) |
239 | 257 | filterset_class = PackageV2FilterSet |
|
0 commit comments