Allow vulnerabilities look up from a CPE

Change models and add migrations for same, do some minor nit picks in nvd importer

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

Author: TG1999

vulnerabilities/api.py

@@ -183,26 +183,19 @@ class CPEFilterSet(filters.FilterSet):
     cpe = filters.CharFilter(method="filter_cpe")
 
     class Meta:
-        model = VulnerabilityReference
-        fields = ["reference_id"]
+        model = Vulnerability
+        fields = ["vulnerabilityreference__reference_id"]
 
     def filter_cpe(self, queryset, name, value):
         cpe = unquote(value)
-        return self.queryset.filter(reference_id=cpe)
-
-
-class CPESerializer(serializers.HyperlinkedModelSerializer):
-
-    vulnerability = MinimalVulnerabilitySerializer(read_only=True)
-
-    class Meta:
-        model = VulnerabilityReference
-        fields = ["vulnerability"]
+        return self.queryset.filter(vulnerabilityreference__reference_id__startswith=cpe).distinct()
 
 
 class CPEViewSet(viewsets.ReadOnlyModelViewSet):
-    queryset = VulnerabilityReference.objects.filter(reference_id__startswith="cpe")
-    serializer_class = CPESerializer
+    queryset = Vulnerability.objects.filter(
+        vulnerabilityreference__reference_id__startswith="cpe"
+    ).distinct()
+    serializer_class = VulnerabilitySerializer
     paginate_by = 50
     filter_backends = (filters.DjangoFilterBackend,)
     filterset_class = CPEFilterSet

vulnerabilities/importers/nvd.py

@@ -45,14 +45,12 @@ class NVDImporter(Importer):
     spdx_license_expression = "LicenseRef-scancode-unknown"
 
     def advisory_data(self):
-        advisory_data = []
         current_year = date.today().year
         # NVD json feeds start from 2002.
         for year in range(2002, current_year + 1):
             download_url = f"https://nvd.nist.gov/feeds/json/cve/1.1/nvdcve-1.1-{year}.json.gz"
             data = fetch(download_url)
-            advisory_data.extend(to_advisories(data))
-        return advisory_data
+            yield from to_advisories(data)
 
 
 # Isolating network calls for simplicity of testing

vulnerabilities/improve_runner.py

@@ -63,7 +63,11 @@ def process_inferences(inferences: List[Inference], advisory: Advisory, improver
 
         for ref in inference.references:
             reference, _ = models.VulnerabilityReference.objects.get_or_create(
-                vulnerability=vuln, reference_id=ref.reference_id, url=ref.url
+                reference_id=ref.reference_id, url=ref.url
+            )
+
+            models.VulnerabilityRelatedReference.objects.update_or_create(
+                reference=reference, vulnerability=vuln
             )
 
             for severity in ref.severities:

vulnerabilities/migrations/0010_vulnerabilityrelatedreference_and_more.py

@@ -0,0 +1,52 @@
+# Generated by Django 4.0.3 on 2022-04-22 15:43
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('vulnerabilities', '0009_alter_advisory_summary_and_more'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='VulnerabilityRelatedReference',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+            ],
+        ),
+        migrations.AlterUniqueTogether(
+            name='vulnerabilityreference',
+            unique_together={('url', 'reference_id')},
+        ),
+        migrations.AddField(
+            model_name='vulnerability',
+            name='references',
+            field=models.ManyToManyField(through='vulnerabilities.VulnerabilityRelatedReference', to='vulnerabilities.vulnerabilityreference'),
+        ),
+        migrations.AddField(
+            model_name='vulnerabilityreference',
+            name='vulnerabilities',
+            field=models.ManyToManyField(through='vulnerabilities.VulnerabilityRelatedReference', to='vulnerabilities.vulnerability'),
+        ),
+        migrations.AddField(
+            model_name='vulnerabilityrelatedreference',
+            name='reference',
+            field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='vulnerabilities.vulnerabilityreference'),
+        ),
+        migrations.AddField(
+            model_name='vulnerabilityrelatedreference',
+            name='vulnerability',
+            field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='vulnerabilities.vulnerability'),
+        ),
+        migrations.RemoveField(
+            model_name='vulnerabilityreference',
+            name='vulnerability',
+        ),
+        migrations.AlterUniqueTogether(
+            name='vulnerabilityrelatedreference',
+            unique_together={('vulnerability', 'reference')},
+        ),
+    ]

vulnerabilities/models.py

@@ -60,6 +60,10 @@ class Vulnerability(models.Model):
         blank=True,
     )
 
+    references = models.ManyToManyField(
+        to="VulnerabilityReference", through="VulnerabilityRelatedReference"
+    )
+
     @property
     def vulcoid(self):
         return f"VULCOID-{self.vulnerability_id}"
@@ -92,10 +96,11 @@ class VulnerabilityReference(models.Model):
     package manager.
     """
 
-    vulnerability = models.ForeignKey(
-        Vulnerability,
-        on_delete=models.CASCADE,
+    vulnerabilities = models.ManyToManyField(
+        to="Vulnerability",
+        through="VulnerabilityRelatedReference",
     )
+
     url = models.URLField(
         max_length=1024, help_text="URL to the vulnerability reference", blank=True
     )
@@ -111,7 +116,6 @@ def severities(self):
 
     class Meta:
         unique_together = (
-            "vulnerability",
             "url",
             "reference_id",
         )
@@ -121,6 +125,25 @@ def __str__(self):
         return f"{self.url}{reference_id}"
 
 
+class VulnerabilityRelatedReference(models.Model):
+    """
+    A reference related to a vulnerability.
+    """
+
+    vulnerability = models.ForeignKey(
+        Vulnerability,
+        on_delete=models.CASCADE,
+    )
+
+    reference = models.ForeignKey(
+        VulnerabilityReference,
+        on_delete=models.CASCADE,
+    )
+
+    class Meta:
+        unique_together = ("vulnerability", "reference")
+
+
 class Package(PackageURLMixin):
     """
     A software package with related vulnerabilities.

vulnerabilities/tests/test_fix_api.py

@@ -25,6 +25,8 @@
 
 from vulnerabilities.models import Package
 from vulnerabilities.models import Vulnerability
+from vulnerabilities.models import VulnerabilityReference
+from vulnerabilities.models import VulnerabilityRelatedReference
 
 
 class APITestCaseVulnerability(TestCase):
@@ -85,3 +87,23 @@ def test_api_with_single_vulnerability(self):
         assert response["name"] == "test-vulnDB"
         assert response["version"] == "1.0"
         assert response["type"] == "generic"
+
+
+class CPEApi(TestCase):
+    def setUp(self):
+        self.vulnerability = Vulnerability.objects.create(summary="test")
+        for i in range(0, 10):
+            ref, _ = VulnerabilityReference.objects.get_or_create(
+                reference_id=f"cpe:/a:nginx:{i}",
+            )
+            VulnerabilityRelatedReference.objects.create(
+                reference=ref, vulnerability=self.vulnerability
+            )
+
+    def test_api_status(self):
+        response = self.client.get("/api/cpes/", format="json")
+        self.assertEqual(status.HTTP_200_OK, response.status_code)
+
+    def test_api_response(self):
+        response = self.client.get("/api/cpes/?cpe=cpe:/a:nginx:9", format="json").data
+        self.assertEqual(response["count"], 1)