Add data migration for old pysec advisory

Signed-off-by: Keshav Priyadarshi <[email protected]>

Author: keshav-space

vulnerabilities/migrations/0074_update_pysec_advisory_created_by.py

@@ -0,0 +1,37 @@
+# Generated by Django 4.2.16 on 2024-10-24 13:51
+
+from django.db import migrations
+
+"""
+Update the created_by field on Advisory from the old qualified_name
+to the new pipeline_id.
+"""
+
+
+def update_created_by(apps, schema_editor):
+    from vulnerabilities.pipelines.pysec_importer import PyPIImporterPipeline
+
+    Advisory = apps.get_model("vulnerabilities", "Advisory")
+    Advisory.objects.filter(created_by="vulnerabilities.importers.pysec.PyPIImporter").update(
+        created_by=PyPIImporterPipeline.pipeline_id
+    )
+
+
+def reverse_update_created_by(apps, schema_editor):
+    from vulnerabilities.pipelines.pysec_importer import PyPIImporterPipeline
+
+    Advisory = apps.get_model("vulnerabilities", "Advisory")
+    Advisory.objects.filter(created_by=PyPIImporterPipeline.pipeline_id).update(
+        created_by="vulnerabilities.importers.pysec.PyPIImporter"
+    )
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("vulnerabilities", "0073_delete_packagerelatedvulnerability"),
+    ]
+
+    operations = [
+        migrations.RunPython(update_created_by, reverse_code=reverse_update_created_by),
+    ]

vulnerabilities/pipelines/pysec_importer.py

@@ -15,7 +15,6 @@
 import requests
 
 from vulnerabilities.importer import AdvisoryData
-from vulnerabilities.importers.osv import parse_advisory_data
 from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipeline
 
 
@@ -48,6 +47,8 @@ def advisories_count(self) -> int:
 
     def collect_advisories(self) -> Iterable[AdvisoryData]:
         """Yield AdvisoryData using a zipped data dump of OSV data"""
+        from vulnerabilities.importers.osv import parse_advisory_data
+
         with ZipFile(BytesIO(self.advisory_zip)) as zip_file:
             for file_name in zip_file.namelist():
                 if not file_name.startswith("PYSEC-"):

vulnerabilities/tests/test_data_migrations.py

@@ -672,7 +672,7 @@ def setUpBeforeMigration(self, apps):
             date_collected=timezone.now(),
         )
 
-    def test_removal_of_duped_purls(self):
+    def test_update_npm_pypa_created_by_field(self):
         Advisory = apps.get_model("vulnerabilities", "Advisory")
         adv = Advisory.objects.all()
 
@@ -714,7 +714,7 @@ def setUpBeforeMigration(self, apps):
             date_collected=timezone.now(),
         )
 
-    def test_removal_of_duped_purls(self):
+    def test_update_nginx_created_by_field(self):
         Advisory = apps.get_model("vulnerabilities", "Advisory")
         adv = Advisory.objects.all()
 
@@ -753,7 +753,7 @@ def setUpBeforeMigration(self, apps):
             date_collected=timezone.now(),
         )
 
-    def test_removal_of_duped_purls(self):
+    def test_update_gitlab_created_by_field(self):
         Advisory = apps.get_model("vulnerabilities", "Advisory")
         adv = Advisory.objects.all()
 
@@ -794,7 +794,7 @@ def setUpBeforeMigration(self, apps):
             date_collected=timezone.now(),
         )
 
-    def test_removal_of_duped_purls(self):
+    def test_update_github_created_by_field(self):
         Advisory = apps.get_model("vulnerabilities", "Advisory")
         adv = Advisory.objects.all()
 
@@ -835,9 +835,48 @@ def setUpBeforeMigration(self, apps):
             date_collected=timezone.now(),
         )
 
-    def test_removal_of_duped_purls(self):
+    def test_update_nvd_created_by_field(self):
         Advisory = apps.get_model("vulnerabilities", "Advisory")
         adv = Advisory.objects.all()
 
         assert adv.filter(created_by="vulnerabilities.importers.nvd.NVDImporter").count() == 0
         assert adv.filter(created_by="nvd_importer").count() == 1
+
+
+class TestUpdatePysecAdvisoryCreatedByField(TestMigrations):
+    app_name = "vulnerabilities"
+    migrate_from = "0073_delete_packagerelatedvulnerability"
+    migrate_to = "0074_update_pysec_advisory_created_by"
+
+    advisory_data1 = AdvisoryData(
+        aliases=["CVE-2020-13371337"],
+        summary="vulnerability description here",
+        affected_packages=[
+            AffectedPackage(
+                package=PackageURL(type="pypi", name="foobar"),
+                affected_version_range=VersionRange.from_string("vers:pypi/>=1.0.0|<=2.0.0"),
+            )
+        ],
+        references=[Reference(url="https://example.com/with/more/info/CVE-2020-13371337")],
+        date_published=timezone.now(),
+        url="https://test.com",
+    )
+
+    def setUpBeforeMigration(self, apps):
+        Advisory = apps.get_model("vulnerabilities", "Advisory")
+        adv1 = Advisory.objects.create(
+            aliases=self.advisory_data1.aliases,
+            summary=self.advisory_data1.summary,
+            affected_packages=[pkg.to_dict() for pkg in self.advisory_data1.affected_packages],
+            references=[ref.to_dict() for ref in self.advisory_data1.references],
+            url=self.advisory_data1.url,
+            created_by="vulnerabilities.importers.pysec.PyPIImporter",
+            date_collected=timezone.now(),
+        )
+
+    def test_update_pysec_created_by_field(self):
+        Advisory = apps.get_model("vulnerabilities", "Advisory")
+        adv = Advisory.objects.all()
+
+        assert adv.filter(created_by="vulnerabilities.importers.pysec.PyPIImporter").count() == 0
+        assert adv.filter(created_by="pysec_importer").count() == 1