Add CodeFix in API

TG1999 · TG1999 · commit 7bb44be8bd2b · 2025-01-08T20:21:54.000+05:30
Signed-off-by: Tushar Goel &lt;tushar.goel.dav@gmail.com&gt;
diff --git a/vulnerabilities/api_v2.py b/vulnerabilities/api_v2.py
@@ -21,6 +21,7 @@
 from rest_framework.response import Response
 from rest_framework.reverse import reverse
 
+from vulnerabilities.models import CodeFix
 from vulnerabilities.models import Package
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference
@@ -198,14 +199,25 @@ def get_affected_by_vulnerabilities(self, obj):
         Return a dictionary with vulnerabilities as keys and their details, including fixed_by_packages.
         """
         result = {}
+        request = self.context.get("request")
         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
+            # Get code fixed for a vulnerability
+            code_fixes = CodeFix.objects.filter(
+                affected_package_vulnerability__vulnerability=vuln
+            ).distinct()
+            code_fix_urls = [
+                reverse("codefix-detail", args=[code_fix.id], request=request)
+                for code_fix in code_fixes
+            ]
+
             result[vuln.vulnerability_id] = {
                 "vulnerability_id": vuln.vulnerability_id,
                 "fixed_by_packages": purl,
+                "code_fixes": code_fix_urls,
             }
         return result
 
@@ -521,3 +533,81 @@ def lookup(self, request):
 
         qs = self.get_queryset().for_purls([purl]).with_is_vulnerable()
         return Response(PackageV2Serializer(qs, many=True, context={"request": request}).data)
+
+
+from rest_framework import serializers
+
+from vulnerabilities.models import CodeFix
+
+
+class CodeFixSerializer(serializers.ModelSerializer):
+    """
+    Serializer for the CodeFix model.
+    Provides detailed information about a code fix.
+    """
+
+    affected_vulnerability_id = serializers.CharField(
+        source="affected_package_vulnerability.vulnerability.vulnerability_id",
+        read_only=True,
+        help_text="ID of the affected vulnerability.",
+    )
+    affected_package_purl = serializers.CharField(
+        source="affected_package_vulnerability.package.package_url",
+        read_only=True,
+        help_text="PURL of the affected package.",
+    )
+    fixed_package_purl = serializers.CharField(
+        source="fixed_package_vulnerability.package.package_url",
+        read_only=True,
+        help_text="PURL of the fixing package (if available).",
+    )
+    created_at = serializers.DateTimeField(
+        format="%Y-%m-%dT%H:%M:%SZ",
+        read_only=True,
+        help_text="Timestamp when the code fix was created.",
+    )
+    updated_at = serializers.DateTimeField(
+        format="%Y-%m-%dT%H:%M:%SZ",
+        read_only=True,
+        help_text="Timestamp when the code fix was last updated.",
+    )
+
+    class Meta:
+        model = CodeFix
+        fields = [
+            "id",
+            "commits",
+            "pulls",
+            "downloads",
+            "patch",
+            "affected_vulnerability_id",
+            "affected_package_purl",
+            "fixed_package_purl",
+            "notes",
+            "references",
+            "is_reviewed",
+            "created_at",
+            "updated_at",
+        ]
+        read_only_fields = ["created_at", "updated_at"]
+
+
+class CodeFixViewSet(viewsets.ReadOnlyModelViewSet):
+    """
+    API endpoint that allows viewing CodeFix entries.
+    """
+
+    queryset = CodeFix.objects.all()
+    serializer_class = CodeFixSerializer
+
+    def get_queryset(self):
+        """
+        Optionally filter by vulnerability ID.
+        """
+        queryset = super().get_queryset()
+        vulnerability_id = self.request.query_params.get("vulnerability_id")
+        if vulnerability_id:
+            queryset = queryset.filter(
+                affected_package_vulnerability__vulnerability__vulnerability_id=vulnerability_id
+            )
+        return queryset
diff --git a/vulnerabilities/improvers/__init__.py b/vulnerabilities/improvers/__init__.py
@@ -10,6 +10,7 @@
 from vulnerabilities.improvers import valid_versions
 from vulnerabilities.improvers import vulnerability_status
 from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.pipelines import collect_commits
 from vulnerabilities.pipelines import compute_package_risk
 from vulnerabilities.pipelines import compute_package_version_rank
 from vulnerabilities.pipelines import enhance_with_exploitdb
@@ -41,6 +42,7 @@
     enhance_with_exploitdb.ExploitDBImproverPipeline,
     compute_package_risk.ComputePackageRiskPipeline,
     compute_package_version_rank.ComputeVersionRankPipeline,
+    collect_commits.CollectFixCommitsPipeline,
 ]
 
 IMPROVERS_REGISTRY = {
diff --git a/vulnerabilities/models.py b/vulnerabilities/models.py
@@ -1101,6 +1101,8 @@ class AffectedByPackageRelatedVulnerability(PackageRelatedVulnerabilityBase):
         related_name="affected_package_vulnerability_relations",
     )
 
+    objects = BaseQuerySet.as_manager()
+
     class Meta(PackageRelatedVulnerabilityBase.Meta):
         verbose_name_plural = "Affected By Package Related Vulnerabilities"
 
diff --git a/vulnerabilities/pipelines/collect_commits.py b/vulnerabilities/pipelines/collect_commits.py
@@ -20,7 +20,8 @@ def is_vcs_url_already_processed(commit_id):
     """
     Check if a VCS URL exists in a CodeFix entry.
     """
-    return CodeFix.objects.filter(commits__contains=[commit_id]).exists()
+    if "commit" in commit_id:
+        return CodeFix.objects.filter(commits__contains=[commit_id]).exists()
 
 
 class CollectFixCommitsPipeline(VulnerableCodePipeline):
@@ -70,17 +71,19 @@ def collect_and_store_fix_commits(self):
                         f"Skipping already processed reference: {reference.url} with VCS URL {vcs_url}"
                     )
                     continue
-                code_fix, created = CodeFix.objects.get_or_create(
-                    commits=[vcs_url],
-                    affected_package_vulnerability=apv,
-                )
-
-                if created:
-                    created_fix_count += 1
-                    self.log(
-                        f"Created CodeFix entry for reference: {reference.url} with VCS URL {vcs_url}"
+                # check if vcs_url has commit
+                if "/commit/" in vcs_url:
+                    code_fix, created = CodeFix.objects.get_or_create(
+                        commits=[vcs_url],
+                        affected_package_vulnerability=apv,
                     )
 
+                    if created:
+                        created_fix_count += 1
+                        self.log(
+                            f"Created CodeFix entry for reference: {reference.url} with VCS URL {vcs_url}"
+                        )
+
         self.log(f"Successfully created {created_fix_count:,d} CodeFix entries.")
 
 
diff --git a/vulnerabilities/tests/test_api_v2.py b/vulnerabilities/tests/test_api_v2.py
@@ -216,7 +216,7 @@ def test_list_packages(self):
         Should return a list of packages with their details and associated vulnerabilities.
         """
         url = reverse("package-v2-list")
-        with self.assertNumQueries(31):
+        with self.assertNumQueries(32):
             response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertIn("results", response.data)
@@ -238,7 +238,7 @@ def test_filter_packages_by_purl(self):
         Test filtering packages by one or more PURLs.
         """
         url = reverse("package-v2-list")
-        with self.assertNumQueries(19):
+        with self.assertNumQueries(20):
             response = self.client.get(url, {"purl": "pkg:pypi/django@3.2"}, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(len(response.data["results"]["packages"]), 1)
@@ -249,7 +249,7 @@ def test_filter_packages_by_affected_vulnerability(self):
         Test filtering packages by affected_by_vulnerability.
         """
         url = reverse("package-v2-list")
-        with self.assertNumQueries(19):
+        with self.assertNumQueries(20):
             response = self.client.get(
                 url, {"affected_by_vulnerability": "VCID-1234"}, format="json"
             )
@@ -308,7 +308,11 @@ def test_package_serializer_fields(self):
 
         # Verify affected_by_vulnerabilities structure
         expected_affected_by_vulnerabilities = {
-            "VCID-1234": {"vulnerability_id": "VCID-1234", "fixed_by_packages": None}
+            "VCID-1234": {
+                "code_fixes": [],
+                "vulnerability_id": "VCID-1234",
+                "fixed_by_packages": None,
+            }
         }
         self.assertEqual(data["affected_by_vulnerabilities"], expected_affected_by_vulnerabilities)
 
@@ -387,7 +391,13 @@ def test_get_affected_by_vulnerabilities(self):
         vulnerabilities = serializer.get_affected_by_vulnerabilities(package)
         self.assertEqual(
             vulnerabilities,
-            {"VCID-1234": {"vulnerability_id": "VCID-1234", "fixed_by_packages": None}},
+            {
+                "VCID-1234": {
+                    "code_fixes": [],
+                    "vulnerability_id": "VCID-1234",
+                    "fixed_by_packages": None,
+                }
+            },
         )
 
     def test_get_fixing_vulnerabilities(self):
@@ -591,7 +601,7 @@ def test_lookup_with_valid_purl(self):
         """
         url = reverse("package-v2-lookup")
         data = {"purl": "pkg:pypi/django@3.2"}
-        with self.assertNumQueries(12):
+        with self.assertNumQueries(13):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(1, len(response.data))
@@ -603,7 +613,13 @@ def test_lookup_with_valid_purl(self):
         self.assertEqual(response.data[0]["purl"], "pkg:pypi/django@3.2")
         self.assertEqual(
             response.data[0]["affected_by_vulnerabilities"],
-            {"VCID-1234": {"vulnerability_id": "VCID-1234", "fixed_by_packages": None}},
+            {
+                "VCID-1234": {
+                    "code_fixes": [],
+                    "vulnerability_id": "VCID-1234",
+                    "fixed_by_packages": None,
+                }
+            },
         )
         self.assertEqual(response.data[0]["fixing_vulnerabilities"], [])
 
diff --git a/vulnerablecode/urls.py b/vulnerablecode/urls.py
@@ -20,6 +20,7 @@
 from vulnerabilities.api import CPEViewSet
 from vulnerabilities.api import PackageViewSet
 from vulnerabilities.api import VulnerabilityViewSet
+from vulnerabilities.api_v2 import CodeFixViewSet
 from vulnerabilities.api_v2 import PackageV2ViewSet
 from vulnerabilities.api_v2 import VulnerabilityV2ViewSet
 from vulnerabilities.views import ApiUserCreateView
@@ -48,6 +49,8 @@ def __init__(self, *args, **kwargs):
 api_v2_router = OptionalSlashRouter()
 api_v2_router.register("packages", PackageV2ViewSet, basename="package-v2")
 api_v2_router.register("vulnerabilities", VulnerabilityV2ViewSet, basename="vulnerability-v2")
+api_v2_router.register("codefixes", CodeFixViewSet, basename="codefix")
+
 
 urlpatterns = [
     path("api/v2/", include(api_v2_router.urls)),

Original file line number	Diff line number	Diff line change
`@@ -1101,6 +1101,8 @@ class AffectedByPackageRelatedVulnerability(PackageRelatedVulnerabilityBase):`
`1101`	`1101`	`related_name="affected_package_vulnerability_relations",`
`1102`	`1102`	`)`
`1103`	`1103`
	`1104`	`+ objects = BaseQuerySet.as_manager()`
	`1105`	`+`
`1104`	`1106`	`class Meta(PackageRelatedVulnerabilityBase.Meta):`
`1105`	`1107`	`verbose_name_plural = "Affected By Package Related Vulnerabilities"`
`1106`	`1108`