Extraxt vulnerabilities from CPE

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

Author: TG1999

vulnerabilities/api.py

@@ -88,6 +88,20 @@ class Meta:
         fields = "__all__"
 
 
+class VulnerabilitySerializer(serializers.HyperlinkedModelSerializer):
+
+    resolved_packages = MinimalPackageSerializer(many=True, source="resolved_to", read_only=True)
+    unresolved_packages = MinimalPackageSerializer(
+        many=True, source="vulnerable_to", read_only=True
+    )
+
+    references = VulnerabilityReferenceSerializer(many=True, source="vulnerabilityreference_set")
+
+    class Meta:
+        model = Vulnerability
+        fields = "__all__"
+
+
 class PackageSerializer(serializers.HyperlinkedModelSerializer):
 
     unresolved_vulnerabilities = MinimalVulnerabilitySerializer(
@@ -103,6 +117,17 @@ class Meta:
         exclude = ["vulnerabilities"]
 
 
+class CPESerializer(serializers.HyperlinkedModelSerializer):
+
+    vulnerabilities = MinimalVulnerabilitySerializer(
+        many=True, source="get_vulnerabilities_from_cpe", read_only=True
+    )
+
+    class Meta:
+        model = VulnerabilityReference
+        fields = ["vulnerabilities"]
+
+
 class PackageFilterSet(filters.FilterSet):
     purl = filters.CharFilter(method="filter_purl")
 
@@ -177,3 +202,16 @@ class VulnerabilityViewSet(viewsets.ReadOnlyModelViewSet):
     paginate_by = 50
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = VulnerabilityFilterSet
+
+
+class VulnerabilityReferenceFilterSet(filters.FilterSet):
+    class Meta:
+        model = VulnerabilityReference
+        fields = ["reference_id"]
+
+
+class CPEViewSet(viewsets.ReadOnlyModelViewSet):
+    queryset = VulnerabilityReference.objects.all()
+    serializer_class = CPESerializer
+    filter_backends = (filters.DjangoFilterBackend,)
+    filterset_class = VulnerabilityReferenceFilterSet

vulnerabilities/models.py

@@ -89,6 +89,9 @@ def resolved_to(self):
     def __str__(self):
         return self.vulcoid
 
+    def get_vulnerabilities_from_cpe(self, cpe):
+        return Vulnerability.objects.filter(vulnerabilityreference__reference_id__startswith=cpe)
+
     class Meta:
         verbose_name_plural = "Vulnerabilities"
 
@@ -117,6 +120,12 @@ class VulnerabilityReference(models.Model):
     def severities(self):
         return VulnerabilitySeverity.objects.filter(reference=self.id)
 
+    @property
+    def get_vulnerabilities_from_cpe(self):
+        return Vulnerability.objects.filter(
+            vulnerabilityreference__reference_id__startswith=self.reference_id
+        )
+
     class Meta:
         unique_together = (
             "vulnerability",

vulnerablecode/urls.py

@@ -26,6 +26,7 @@
 from django.urls import path
 from rest_framework.routers import DefaultRouter
 
+from vulnerabilities.api import CPEViewSet
 from vulnerabilities.api import PackageViewSet
 from vulnerabilities.api import VulnerabilityViewSet
 from vulnerabilities.views import HomePage
@@ -47,6 +48,7 @@ def __init__(self, *args, **kwargs):
 api_router.register(r"packages", PackageViewSet)
 # `DefaultRouter` requires `basename` when registering viewsets that don't define a queryset.
 api_router.register(r"vulnerabilities", VulnerabilityViewSet, basename="vulnerability")
+api_router.register(r"cpes", CPEViewSet, basename="cpe")
 
 urlpatterns = [
     path("admin/", admin.site.urls),