Add fixed by package in V2 API (#1706)

* Add fixed by package in V2 API Signed-off-by: Tushar Goel <[email protected]> * Add tests Signed-off-by: Tushar Goel <[email protected]> * Add tests Signed-off-by: Tushar Goel <[email protected]> --------- Signed-off-by: Tushar Goel <[email protected]>
aboutcode-org · Jan 6, 2025 · cebb5d6 · cebb5d6
1 parent 3cee771
commit cebb5d6
Show file tree

Hide file tree

Showing 2 changed files with 142 additions and 40 deletions.
diff --git a/vulnerabilities/api_v2.py b/vulnerabilities/api_v2.py
@@ -8,6 +8,7 @@
 #
 
 
+from django.db.models import Prefetch
 from django_filters import rest_framework as filters
 from drf_spectacular.utils import OpenApiParameter
 from drf_spectacular.utils import extend_schema
@@ -20,8 +21,6 @@
 from rest_framework.response import Response
 from rest_framework.reverse import reverse
 
-from vulnerabilities.api import PackageFilterSet
-from vulnerabilities.api import VulnerabilitySeveritySerializer
 from vulnerabilities.models import Package
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference
@@ -195,7 +194,20 @@ class Meta:
         ]
 
     def get_affected_by_vulnerabilities(self, obj):
-        return [vuln.vulnerability_id for vuln in obj.affected_by_vulnerabilities.all()]
+        """
+        Return a dictionary with vulnerabilities as keys and their details, including fixed_by_packages.
+        """
+        result = {}
+        for vuln in getattr(obj, "prefetched_affected_vulnerabilities", []):
+            fixed_by_package = vuln.fixed_by_packages.first()
+            purl = None
+            if fixed_by_package:
+                purl = fixed_by_package.package_url
+            result[vuln.vulnerability_id] = {
+                "vulnerability_id": vuln.vulnerability_id,
+                "fixed_by_packages": purl,
+            }
+        return result
 
     def get_fixing_vulnerabilities(self, obj):
         # Ghost package should not fix any vulnerability.
@@ -233,7 +245,13 @@ class PackageV2FilterSet(filters.FilterSet):
 
 
 class PackageV2ViewSet(viewsets.ReadOnlyModelViewSet):
-    queryset = Package.objects.all()
+    queryset = Package.objects.all().prefetch_related(
+        Prefetch(
+            "affected_by_vulnerabilities",
+            queryset=Vulnerability.objects.prefetch_related("fixed_by_packages"),
+            to_attr="prefetched_affected_vulnerabilities",
+        )
+    )
     serializer_class = PackageV2Serializer
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = PackageV2FilterSet