Skip to content

Commit cebb5d6

Browse files
authored
Add fixed by package in V2 API (#1706)
* Add fixed by package in V2 API Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com> * Add tests Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com> * Add tests Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com> --------- Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>
1 parent 3cee771 commit cebb5d6

File tree

2 files changed

+142
-40
lines changed

2 files changed

+142
-40
lines changed

vulnerabilities/api_v2.py

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#
99

1010

11+
from django.db.models import Prefetch
1112
from django_filters import rest_framework as filters
1213
from drf_spectacular.utils import OpenApiParameter
1314
from drf_spectacular.utils import extend_schema
@@ -20,8 +21,6 @@
2021
from rest_framework.response import Response
2122
from rest_framework.reverse import reverse
2223

23-
from vulnerabilities.api import PackageFilterSet
24-
from vulnerabilities.api import VulnerabilitySeveritySerializer
2524
from vulnerabilities.models import Package
2625
from vulnerabilities.models import Vulnerability
2726
from vulnerabilities.models import VulnerabilityReference
@@ -195,7 +194,20 @@ class Meta:
195194
]
196195

197196
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
199211

200212
def get_fixing_vulnerabilities(self, obj):
201213
# Ghost package should not fix any vulnerability.
@@ -233,7 +245,13 @@ class PackageV2FilterSet(filters.FilterSet):
233245

234246

235247
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+
)
237255
serializer_class = PackageV2Serializer
238256
filter_backends = (filters.DjangoFilterBackend,)
239257
filterset_class = PackageV2FilterSet

0 commit comments

Comments
 (0)