Use impact table to fetch fixed and affected pkg in API

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

vulnerabilities/api_v2.py

@@ -31,6 +31,7 @@
 from vulnerabilities.models import AdvisoryWeakness
 from vulnerabilities.models import CodeFix
 from vulnerabilities.models import CodeFixV2
+from vulnerabilities.models import ImpactedPackage
 from vulnerabilities.models import Package
 from vulnerabilities.models import PackageV2
 from vulnerabilities.models import PipelineRun
@@ -320,36 +321,28 @@ class Meta:
             "risk_score",
         ]
 
-    def get_affected_by_vulnerabilities(self, obj):
-        """
-        Return a dictionary with vulnerabilities as keys and their details, including fixed_by_packages.
-        """
+    def get_affected_by_vulnerabilities(self, package):
+        """Return a dictionary with advisory as keys and their details, including fixed_by_packages."""
         result = {}
         request = self.context.get("request")
-        for adv in getattr(obj, "prefetched_affected_advisories", []):
-            fixed_by_package = adv.fixed_by_packages.first()
-            purl = None
-            if fixed_by_package:
-                purl = fixed_by_package.package_url
-            # Get code fixed for a vulnerability
-            code_fixes = CodeFixV2.objects.filter(advisory=adv).distinct()
+        for impact in package.affected_in_impacts.all():
+            advisory = impact.advisory
+            fixed_by_packages = [pkg.purl for pkg in impact.fixed_by_packages.all()]
+            code_fixes = CodeFixV2.objects.filter(advisory=advisory).distinct()
             code_fix_urls = [
                 reverse("advisory-codefix-detail", args=[code_fix.id], request=request)
                 for code_fix in code_fixes
             ]
-
-            result[adv.avid] = {
-                "advisory_id": adv.avid,
-                "fixed_by_packages": purl,
+            result[advisory.avid] = {
+                "advisory_id": advisory.avid,
+                "fixed_by_packages": fixed_by_packages,
                 "code_fixes": code_fix_urls,
             }
+
         return result
 
-    def get_fixing_vulnerabilities(self, obj):
-        # Ghost package should not fix any vulnerability.
-        if obj.is_ghost:
-            return []
-        return [adv.avid for adv in obj.fixing_advisories.all()]
+    def get_fixing_vulnerabilities(self, package):
+        return [impact.advisory.avid for impact in package.fixed_in_impacts.all()]
 
 
 class PackageurlListSerializer(serializers.Serializer):
@@ -968,12 +961,21 @@ def get_view_name(self):
 
 
 class AdvisoriesPackageV2ViewSet(viewsets.ReadOnlyModelViewSet):
-    queryset = PackageV2.objects.all().prefetch_related(
-        Prefetch(
-            "affected_by_advisories",
-            queryset=AdvisoryV2.objects.prefetch_related("fixed_by_packages"),
-            to_attr="prefetched_affected_advisories",
+    queryset = (
+        PackageV2.objects.all()
+        .prefetch_related(
+            Prefetch(
+                "affected_in_impacts",
+                queryset=ImpactedPackage.objects.select_related("advisory").prefetch_related(
+                    "fixed_by_packages",
+                ),
+            ),
+            Prefetch(
+                "fixed_in_impacts",
+                queryset=ImpactedPackage.objects.select_related("advisory"),
+            ),
         )
+        .with_is_vulnerable()
     )
     serializer_class = AdvisoryPackageV2Serializer
     filter_backends = (filters.DjangoFilterBackend,)
@@ -994,30 +996,28 @@ def get_queryset(self):
 
     def list(self, request, *args, **kwargs):
         queryset = self.get_queryset()
-        # Apply pagination
         page = self.paginate_queryset(queryset)
+
+        advisories = set()
         if page is not None:
-            # Collect only vulnerabilities for packages in the current page
-            advisories = set()
             for package in page:
-                advisories.update(package.affected_by_advisories.all())
-                advisories.update(package.fixing_advisories.all())
+                advisories.update({impact.advisory for impact in package.affected_in_impacts.all()})
+                advisories.update({impact.advisory for impact in package.fixed_in_impacts.all()})
 
             # Serialize the vulnerabilities with advisory_id and advisory label as keys
             advisory_data = {f"{adv.avid}": AdvisoryV2Serializer(adv).data for adv in advisories}
 
             # Serialize the current page of packages
             serializer = self.get_serializer(page, many=True)
             data = serializer.data
-            print(data)
+
             # Use 'self.get_paginated_response' to include pagination data
             return self.get_paginated_response({"advisories": advisory_data, "packages": data})
 
         # If pagination is not applied, collect vulnerabilities for all packages
-        advisories = set()
         for package in queryset:
-            advisories.update(package.affected_by_vulnerabilities.all())
-            advisories.update(package.fixing_vulnerabilities.all())
+            advisories.update({impact.advisory for impact in package.affected_in_impacts.all()})
+            advisories.update({impact.advisory for impact in package.fixed_in_impacts.all()})
 
         advisory_data = {f"{adv.avid}": AdvisoryV2Serializer(adv).data for adv in advisories}
 
@@ -1053,13 +1053,28 @@ def bulk_lookup(self, request):
         purls = validated_data.get("purls")
 
         # Fetch packages matching the provided purls
-        packages = PackageV2.objects.for_purls(purls).with_is_vulnerable()
+        packages = (
+            PackageV2.objects.for_purls(purls)
+            .prefetch_related(
+                Prefetch(
+                    "affected_in_impacts",
+                    queryset=ImpactedPackage.objects.select_related("advisory").prefetch_related(
+                        "fixed_by_packages",
+                    ),
+                ),
+                Prefetch(
+                    "fixed_in_impacts",
+                    queryset=ImpactedPackage.objects.select_related("advisory"),
+                ),
+            )
+            .with_is_vulnerable()
+        )
 
         # Collect vulnerabilities associated with these packages
         advisories = set()
         for package in packages:
-            advisories.update(package.affected_by_advisories.all())
-            advisories.update(package.fixing_advisories.all())
+            advisories.update({impact.advisory for impact in package.affected_in_impacts.all()})
+            advisories.update({impact.advisory for impact in package.fixed_in_impacts.all()})
 
         # Serialize vulnerabilities with vulnerability_id as keys
         advisory_data = {adv.avid: AdvisoryV2Serializer(adv).data for adv in advisories}
@@ -1124,6 +1139,20 @@ def bulk_search(self, request):
                 PackageV2.objects.filter(plain_package_url__in=plain_purls)
                 .order_by("plain_package_url")
                 .distinct("plain_package_url")
+                .prefetch_related(
+                    Prefetch(
+                        "affected_in_impacts",
+                        queryset=ImpactedPackage.objects.select_related(
+                            "advisory"
+                        ).prefetch_related(
+                            "fixed_by_packages",
+                        ),
+                    ),
+                    Prefetch(
+                        "fixed_in_impacts",
+                        queryset=ImpactedPackage.objects.select_related("advisory"),
+                    ),
+                )
                 .with_is_vulnerable()
             )
 
@@ -1132,14 +1161,16 @@ def bulk_search(self, request):
             # Collect vulnerabilities associated with these packages
             advisories = set()
             for package in packages:
-                advisories.update(package.affected_by_advisories.all())
-                advisories.update(package.fixing_advisories.all())
+                advisories.update({impact.advisory for impact in package.affected_in_impacts.all()})
+                advisories.update({impact.advisory for impact in package.fixed_in_impacts.all()})
 
             advisory_data = {adv.avid: AdvisoryV2Serializer(adv).data for adv in advisories}
 
             if not purl_only:
                 package_data = AdvisoryPackageV2Serializer(
-                    packages, many=True, context={"request": request}
+                    packages,
+                    many=True,
+                    context={"request": request},
                 ).data
                 return Response(
                     {
@@ -1154,20 +1185,38 @@ def bulk_search(self, request):
             vulnerable_purls = [str(package.plain_package_url) for package in vulnerable_purls]
             return Response(data=vulnerable_purls)
 
-        query = PackageV2.objects.filter(package_url__in=purls).distinct().with_is_vulnerable()
+        query = (
+            PackageV2.objects.filter(package_url__in=purls)
+            .distinct()
+            .prefetch_related(
+                Prefetch(
+                    "affected_in_impacts",
+                    queryset=ImpactedPackage.objects.select_related("advisory").prefetch_related(
+                        "fixed_by_packages",
+                    ),
+                ),
+                Prefetch(
+                    "fixed_in_impacts",
+                    queryset=ImpactedPackage.objects.select_related("advisory"),
+                ),
+            )
+            .with_is_vulnerable()
+        )
         packages = query
 
         # Collect vulnerabilities associated with these packages
         advisories = set()
         for package in packages:
-            advisories.update(package.affected_by_advisories.all())
-            advisories.update(package.fixing_advisories.all())
+            advisories.update({impact.advisory for impact in package.affected_in_impacts.all()})
+            advisories.update({impact.advisory for impact in package.fixed_in_impacts.all()})
 
         advisory_data = {adv.advisory_id: AdvisoryV2Serializer(adv).data for adv in advisories}
 
         if not purl_only:
             package_data = AdvisoryPackageV2Serializer(
-                packages, many=True, context={"request": request}
+                packages,
+                many=True,
+                context={"request": request},
             ).data
             return Response(
                 {