Address review comment

Signed-off-by: Tushar Goel <[email protected]>

Author: TG1999

vulnerabilities/api_v2.py

@@ -0,0 +1,192 @@
+#
+# Copyright (c) nexB Inc. and others. All rights reserved.
+# VulnerableCode is a trademark of nexB Inc.
+# SPDX-License-Identifier: Apache-2.0
+# See http://www.apache.org/licenses/LICENSE-2.0 for the license text.
+# See https://github.com/aboutcode-org/vulnerablecode for support or download.
+# See https://aboutcode.org for more information about nexB OSS projects.
+#
+
+from rest_framework import serializers
+from rest_framework import viewsets
+from rest_framework.response import Response
+from rest_framework.reverse import reverse
+
+from vulnerabilities.api import VulnerabilitySeveritySerializer
+from vulnerabilities.models import Package
+from vulnerabilities.models import Vulnerability
+from vulnerabilities.models import VulnerabilityReference
+from vulnerabilities.models import VulnerabilitySeverity
+from vulnerabilities.models import Weakness
+
+
+class WeaknessV2Serializer(serializers.ModelSerializer):
+    cwe_id = serializers.CharField()
+    name = serializers.CharField()
+    description = serializers.CharField()
+
+    class Meta:
+        model = Weakness
+        fields = ["cwe_id", "name", "description"]
+
+
+class VulnerabilityReferenceV2Serializer(serializers.ModelSerializer):
+    url = serializers.CharField()
+    reference_type = serializers.CharField()
+    reference_id = serializers.CharField()
+
+    class Meta:
+        model = VulnerabilityReference
+        fields = ["url", "reference_type", "reference_id"]
+
+
+class VulnerabilityV2Serializer(serializers.ModelSerializer):
+    aliases = serializers.SerializerMethodField()
+    weaknesses = WeaknessV2Serializer(many=True)
+    references = VulnerabilityReferenceV2Serializer(many=True, source="vulnerabilityreference_set")
+    severities = VulnerabilitySeveritySerializer(many=True)
+
+    class Meta:
+        model = Vulnerability
+        fields = [
+            "vulnerability_id",
+            "aliases",
+            "summary",
+            "severities",
+            "weaknesses",
+            "references",
+        ]
+
+    def get_aliases(self, obj):
+        return [alias.alias for alias in obj.aliases.all()]
+
+    def get_severities(self, obj):
+        return obj.severities
+
+
+class VulnerabilityListSerializer(serializers.ModelSerializer):
+    url = serializers.SerializerMethodField()
+
+    class Meta:
+        model = Vulnerability
+        fields = ["vulnerability_id", "url"]
+
+    def get_url(self, obj):
+        request = self.context.get("request")
+        return reverse(
+            "vulnerability-v2-detail",
+            kwargs={"vulnerability_id": obj.vulnerability_id},
+            request=request,
+        )
+
+
+class VulnerabilityV2ViewSet(viewsets.ReadOnlyModelViewSet):
+    queryset = Vulnerability.objects.all()
+    serializer_class = VulnerabilityV2Serializer
+    lookup_field = "vulnerability_id"
+
+    def get_queryset(self):
+        queryset = super().get_queryset()
+        vulnerability_ids = self.request.query_params.getlist("vulnerability_id")
+        aliases = self.request.query_params.getlist("alias")
+
+        if vulnerability_ids:
+            queryset = queryset.filter(vulnerability_id__in=vulnerability_ids)
+
+        if aliases:
+            queryset = queryset.filter(aliases__alias__in=aliases).distinct()
+
+        return queryset
+
+    def get_serializer_class(self):
+        if self.action == "list":
+            return VulnerabilityListSerializer
+        return super().get_serializer_class()
+
+    def list(self, request, *args, **kwargs):
+        queryset = self.get_queryset()
+        vulnerability_ids = request.query_params.getlist("vulnerability_id")
+
+        # If exactly one vulnerability_id is provided, return the serialized data
+        if len(vulnerability_ids) == 1:
+            try:
+                vulnerability = queryset.get(vulnerability_id=vulnerability_ids[0])
+                serializer = self.get_serializer(vulnerability)
+                return Response(serializer.data)
+            except Vulnerability.DoesNotExist:
+                return Response({"detail": "Not found."}, status=404)
+
+        # Otherwise, return a dictionary of vulnerabilities keyed by vulnerability_id
+        page = self.paginate_queryset(queryset)
+        if page is not None:
+            serializer = self.get_serializer(page, many=True)
+            data = serializer.data
+            vulnerabilities = {item["vulnerability_id"]: item for item in data}
+            return self.get_paginated_response({"vulnerabilities": vulnerabilities})
+
+        serializer = self.get_serializer(queryset, many=True)
+        data = serializer.data
+        vulnerabilities = {item["vulnerability_id"]: item for item in data}
+        return Response({"vulnerabilities": vulnerabilities})
+
+
+class PackageV2Serializer(serializers.ModelSerializer):
+    purl = serializers.CharField(source="package_url")
+    affected_by_vulnerabilities = serializers.SerializerMethodField()
+    fixing_vulnerabilities = serializers.SerializerMethodField()
+    next_non_vulnerable_version = serializers.CharField(read_only=True)
+    latest_non_vulnerable_version = serializers.CharField(read_only=True)
+
+    class Meta:
+        model = Package
+        fields = [
+            "purl",
+            "affected_by_vulnerabilities",
+            "fixing_vulnerabilities",
+            "next_non_vulnerable_version",
+            "latest_non_vulnerable_version",
+        ]
+
+    def get_affected_by_vulnerabilities(self, obj):
+        return [vuln.vulnerability_id for vuln in obj.affected_by_vulnerabilities.all()]
+
+    def get_fixing_vulnerabilities(self, obj):
+        return [vuln.vulnerability_id for vuln in obj.fixing_vulnerabilities.all()]
+
+
+class PackageV2ViewSet(viewsets.ReadOnlyModelViewSet):
+    queryset = Package.objects.all()
+    serializer_class = PackageV2Serializer
+
+    def get_queryset(self):
+        queryset = super().get_queryset()
+        package_purls = self.request.query_params.getlist("purl")
+        affected_by_vulnerability = self.request.query_params.get("affected_by_vulnerability")
+        fixing_vulnerability = self.request.query_params.get("fixing_vulnerability")
+
+        if package_purls:
+            queryset = queryset.filter(package_url__in=package_purls)
+        if affected_by_vulnerability:
+            queryset = queryset.filter(
+                affected_by_vulnerabilities__vulnerability_id=affected_by_vulnerability
+            )
+        if fixing_vulnerability:
+            queryset = queryset.filter(
+                fixing_vulnerabilities__vulnerability_id=fixing_vulnerability
+            )
+        return queryset.with_is_vulnerable()
+
+    def list(self, request, *args, **kwargs):
+        queryset = self.get_queryset()
+        # Apply pagination
+        page = self.paginate_queryset(queryset)
+        if page is not None:
+            serializer = self.get_serializer(page, many=True)
+            data = serializer.data
+            # Use 'self.get_paginated_response' to include pagination data
+            return self.get_paginated_response({"purls": data})
+
+        # If pagination is not applied
+        serializer = self.get_serializer(queryset, many=True)
+        data = serializer.data
+        return Response({"purls": data})

vulnerabilities/tests/test_v2_api.py

@@ -19,8 +19,8 @@
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference
 from vulnerabilities.models import Weakness
-from vulnerabilities.v2_api import PackageV2Serializer
-from vulnerabilities.v2_api import VulnerabilityListSerializer
+from vulnerabilities.api_v2 import PackageV2Serializer
+from vulnerabilities.api_v2 import VulnerabilityListSerializer
 
 
 class VulnerabilityV2ViewSetTest(APITestCase):

vulnerablecode/urls.py

@@ -20,8 +20,8 @@
 from vulnerabilities.api import CPEViewSet
 from vulnerabilities.api import PackageViewSet
 from vulnerabilities.api import VulnerabilityViewSet
-from vulnerabilities.v2_api import PackageV2ViewSet
-from vulnerabilities.v2_api import VulnerabilityV2ViewSet
+from vulnerabilities.api_v2 import PackageV2ViewSet
+from vulnerabilities.api_v2 import VulnerabilityV2ViewSet
 from vulnerabilities.views import ApiUserCreateView
 from vulnerabilities.views import HomePage
 from vulnerabilities.views import PackageDetails