Add all packages endpoint functionality in V2

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

Author: TG1999

vulnerabilities/api_v2.py

@@ -8,8 +8,14 @@
 #
 
 
+from drf_spectacular.utils import OpenApiParameter
+from drf_spectacular.utils import extend_schema
+from drf_spectacular.utils import extend_schema_view
+from packageurl import PackageURL
 from rest_framework import serializers
+from rest_framework import status
 from rest_framework import viewsets
+from rest_framework.decorators import action
 from rest_framework.response import Response
 from rest_framework.reverse import reverse
 
@@ -18,8 +24,7 @@
 from vulnerabilities.models import VulnerabilityReference
 from vulnerabilities.models import VulnerabilitySeverity
 from vulnerabilities.models import Weakness
-from drf_spectacular.utils import extend_schema_view, extend_schema, OpenApiParameter
-from rest_framework.decorators import action
+
 
 class WeaknessV2Serializer(serializers.ModelSerializer):
     cwe_id = serializers.CharField()
@@ -90,6 +95,7 @@ def get_url(self, obj):
             request=request,
         )
 
+
 @extend_schema_view(
     list=extend_schema(
         parameters=[
@@ -197,6 +203,13 @@ class PackageBulkSearchRequestSerializer(PackageurlListSerializer):
     plain_purl = serializers.BooleanField(required=False, default=False)
 
 
+class LookupRequestSerializer(serializers.Serializer):
+    purl = serializers.CharField(
+        required=True,
+        help_text="PackageURL strings in canonical form.",
+    )
+
+
 class PackageV2ViewSet(viewsets.ReadOnlyModelViewSet):
     queryset = Package.objects.all()
     serializer_class = PackageV2Serializer
@@ -233,7 +246,7 @@ def list(self, request, *args, **kwargs):
         serializer = self.get_serializer(queryset, many=True)
         data = serializer.data
         return Response({"packages": data})
-    
+
     @extend_schema(
         request=PackageurlListSerializer,
         responses={200: PackageV2Serializer(many=True)},
@@ -269,7 +282,6 @@ def bulk_lookup(self, request):
             ).data
         )
 
-
     @extend_schema(
         request=PackageBulkSearchRequestSerializer,
         responses={200: PackageV2Serializer(many=True)},
@@ -333,8 +345,54 @@ def bulk_search(self, request):
         query = Package.objects.filter(package_url__in=purls).distinct().with_is_vulnerable()
 
         if not purl_only:
-            return Response(PackageV2Serializer(query, many=True, context={"request": request}).data)
+            return Response(
+                PackageV2Serializer(query, many=True, context={"request": request}).data
+            )
 
         vulnerable_purls = query.vulnerable().only("package_url")
         vulnerable_purls = [str(package.package_url) for package in vulnerable_purls]
         return Response(data=vulnerable_purls)
+
+    @action(detail=False, methods=["get"])
+    def all(self, request):
+        """
+        Return a list of Package URLs of vulnerable packages.
+        """
+        vulnerable_purls = (
+            Package.objects.vulnerable()
+            .only("package_url")
+            .order_by("package_url")
+            .distinct()
+            .values_list("package_url", flat=True)
+        )
+        return Response(vulnerable_purls)
+
+    @extend_schema(
+        request=LookupRequestSerializer,
+        responses={200: PackageV2Serializer(many=True)},
+    )
+    @action(
+        detail=False,
+        methods=["post"],
+        serializer_class=LookupRequestSerializer,
+        filter_backends=[],
+        pagination_class=None,
+    )
+    def lookup(self, request):
+        """
+        Return the response for exact PackageURL requested for.
+        """
+        serializer = self.serializer_class(data=request.data)
+        if not serializer.is_valid():
+            return Response(
+                status=status.HTTP_400_BAD_REQUEST,
+                data={
+                    "error": serializer.errors,
+                    "message": "A 'purl' is required.",
+                },
+            )
+        validated_data = serializer.validated_data
+        purl = validated_data.get("purl")
+
+        qs = self.get_queryset().for_purls([purl]).with_is_vulnerable()
+        return Response(PackageV2Serializer(qs, many=True, context={"request": request}).data)

vulnerabilities/tests/test_api_v2.py

@@ -305,3 +305,182 @@ def test_get_fixing_vulnerabilities(self):
         serializer = PackageV2Serializer()
         vulnerabilities = serializer.get_fixing_vulnerabilities(package)
         self.assertEqual(vulnerabilities, ["VCID-5678"])
+
+    def test_bulk_lookup_with_valid_purls(self):
+        """
+        Test bulk lookup with valid PURLs.
+        """
+        url = reverse("package-v2-bulk-lookup")
+        data = {"purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"]}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertEqual(len(response.data), 2)
+        # Verify that the returned data matches the packages
+        purls = [package["purl"] for package in response.data]
+        self.assertIn("pkg:pypi/[email protected]", purls)
+        self.assertIn("pkg:npm/[email protected]", purls)
+
+    def test_bulk_lookup_with_invalid_purls(self):
+        """
+        Test bulk lookup with invalid PURLs.
+        """
+        url = reverse("package-v2-bulk-lookup")
+        data = {"purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"]}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        # Since the packages don't exist, the response should be empty
+        self.assertEqual(len(response.data), 0)
+
+    def test_bulk_lookup_with_empty_purls(self):
+        """
+        Test bulk lookup with empty purls list.
+        Should return 400 Bad Request.
+        """
+        url = reverse("package-v2-bulk-lookup")
+        data = {"purls": []}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+        self.assertIn("error", response.data)
+        self.assertIn("message", response.data)
+        self.assertEqual(response.data["message"], "A non-empty 'purls' list of PURLs is required.")
+
+    def test_bulk_search_with_valid_purls(self):
+        """
+        Test bulk search with valid PURLs.
+        """
+        url = reverse("package-v2-bulk-search")
+        data = {"purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"]}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertEqual(len(response.data), 2)
+        purls = [package["purl"] for package in response.data]
+        self.assertIn("pkg:pypi/[email protected]", purls)
+        self.assertIn("pkg:npm/[email protected]", purls)
+
+    def test_bulk_search_with_purl_only_true(self):
+        """
+        Test bulk search with purl_only set to True.
+        Should return only the PURLs of vulnerable packages.
+        """
+        url = reverse("package-v2-bulk-search")
+        data = {"purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"], "purl_only": True}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        # Since purl_only=True, response should be a list of PURLs
+        self.assertIsInstance(response.data, list)
+        # Only vulnerable packages should be included
+        self.assertEqual(len(response.data), 1)
+        self.assertEqual(response.data, ["pkg:pypi/[email protected]"])
+
+    def test_bulk_search_with_plain_purl_true(self):
+        """
+        Test bulk search with plain_purl set to True.
+        """
+        url = reverse("package-v2-bulk-search")
+        data = {"purls": ["pkg:pypi/[email protected]", "pkg:pypi/[email protected]"], "plain_purl": True}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        # Since plain_purl=True, packages with the same name and version are grouped
+        self.assertEqual(len(response.data), 1)
+        purls = [package["purl"] for package in response.data]
+        self.assertIn("pkg:pypi/[email protected]", purls[0] or "pkg:pypi/[email protected]" in purls[0])
+
+    def test_bulk_search_with_purl_only_and_plain_purl_true(self):
+        """
+        Test bulk search with purl_only and plain_purl both set to True.
+        Should return only the plain PURLs of vulnerable packages.
+        """
+        url = reverse("package-v2-bulk-search")
+        data = {
+            "purls": ["pkg:pypi/[email protected]", "pkg:pypi/[email protected]"],
+            "purl_only": True,
+            "plain_purl": True,
+        }
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        # Response should be a list of plain PURLs
+        self.assertIsInstance(response.data, list)
+        # Only one plain PURL should be returned for vulnerable packages
+        self.assertEqual(len(response.data), 1)
+        self.assertEqual(response.data, ["pkg:pypi/[email protected]"])
+
+    def test_bulk_search_with_invalid_purls(self):
+        """
+        Test bulk search with invalid PURLs.
+        """
+        url = reverse("package-v2-bulk-search")
+        data = {"purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"]}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertEqual(len(response.data), 0)
+
+    def test_bulk_search_with_empty_purls(self):
+        """
+        Test bulk search with empty purls list.
+        Should return 400 Bad Request.
+        """
+        url = reverse("package-v2-bulk-search")
+        data = {"purls": []}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+        self.assertIn("error", response.data)
+        self.assertIn("message", response.data)
+        self.assertEqual(response.data["message"], "A non-empty 'purls' list of PURLs is required.")
+
+    def test_all_vulnerable_packages(self):
+        """
+        Test the 'all' endpoint that returns all vulnerable package URLs.
+        """
+        url = reverse("package-v2-all")
+        response = self.client.get(url, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        # Since package1 and package3 are vulnerable, they should be returned
+        expected_purls = ["pkg:pypi/[email protected]"]
+        self.assertEqual(sorted(response.data), sorted(expected_purls))
+
+    def test_lookup_with_valid_purl(self):
+        """
+        Test the 'lookup' endpoint with a valid PURL.
+        """
+        url = reverse("package-v2-lookup")
+        data = {"purl": "pkg:pypi/[email protected]"}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertEqual(len(response.data), 1)
+        self.assertEqual(response.data[0]["purl"], "pkg:pypi/[email protected]")
+        self.assertEqual(response.data[0]["affected_by_vulnerabilities"], ["VCID-1234"])
+
+    def test_lookup_with_invalid_purl(self):
+        """
+        Test the 'lookup' endpoint with a PURL that does not exist.
+        Should return an empty list.
+        """
+        url = reverse("package-v2-lookup")
+        data = {"purl": "pkg:pypi/[email protected]"}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        # No packages should be returned
+        self.assertEqual(len(response.data), 0)
+
+    def test_lookup_with_missing_purl(self):
+        """
+        Test the 'lookup' endpoint without providing a 'purl'.
+        Should return 400 Bad Request.
+        """
+        url = reverse("package-v2-lookup")
+        data = {}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+        self.assertIn("error", response.data)
+        self.assertIn("message", response.data)
+        self.assertEqual(response.data["message"], "A 'purl' is required.")
+
+    def test_lookup_with_invalid_purl_format(self):
+        """
+        Test the 'lookup' endpoint with an invalid PURL format.
+        Should return 400 Bad Request.
+        """
+        url = reverse("package-v2-lookup")
+        data = {"purl": "invalid_purl_format"}
+        response = self.client.post(url, data, format="json")
+        self.assertEqual(response.status_code, status.HTTP_200_OK)