Fix tests

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

Author: TG1999

vulnerabilities/api.py

@@ -619,10 +619,6 @@ def get_queryset(self):
             super()
             .get_queryset()
             .prefetch_related(
-                # Prefetch(
-                #     "packages",
-                #     queryset=self.get_packages_qs(),
-                # ),
                 "weaknesses",
                 Prefetch(
                     "fixed_by_packages",

vulnerabilities/api_extension.py

@@ -238,8 +238,6 @@ class Meta:
             "qualifiers",
             "subpath",
             "purl",
-            # this hurts
-            "packagerelatedvulnerability__fix",
         ]
 
     def filter_purl(self, queryset, name, value):

vulnerabilities/import_runner.py

@@ -21,7 +21,9 @@
 from vulnerabilities.improver import Inference
 from vulnerabilities.improvers.default import DefaultImporter
 from vulnerabilities.models import Advisory
+from vulnerabilities.models import AffectedByPackageRelatedVulnerability
 from vulnerabilities.models import Alias
+from vulnerabilities.models import FixingPackageRelatedVulnerability
 from vulnerabilities.models import Package
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityChangeLog

vulnerabilities/improve_runner.py

@@ -17,9 +17,10 @@
 from vulnerabilities.importers import IMPORTERS_REGISTRY
 from vulnerabilities.improver import Inference
 from vulnerabilities.models import Advisory
+from vulnerabilities.models import AffectedByPackageRelatedVulnerability
 from vulnerabilities.models import Alias
+from vulnerabilities.models import FixingPackageRelatedVulnerability
 from vulnerabilities.models import Package
-from vulnerabilities.models import PackageRelatedVulnerability
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityChangeLog
 from vulnerabilities.models import VulnerabilityReference
@@ -148,12 +149,11 @@ def process_inferences(
             fixed_package, created = Package.objects.get_or_create_from_purl(
                 purl=inference.fixed_purl
             )
-            PackageRelatedVulnerability(
+            FixingPackageRelatedVulnerability(
                 vulnerability=vulnerability,
                 package=fixed_package,
                 created_by=improver_name,
                 confidence=inference.confidence,
-                fix=True,
             ).update_or_create(
                 advisory=advisory,
             )

vulnerabilities/management/commands/export.py

@@ -116,7 +116,7 @@ def export_data(self, base_path: Path):
                     }
                     package_vulnerabilities.append(package_data)
 
-                    for vuln in pkg_version.vulnerabilities.all():
+                    for vuln in pkg_version.vulnerabilities:
                         vcid = vuln.vulnerability_id
                         # do not write twice the same file
                         if vcid in seen_vcid:
@@ -158,10 +158,14 @@ def packages_by_type_ns_name():
     qs = (
         Package.objects.order_by("type", "namespace", "name", "version")
         .prefetch_related(
-            "vulnerabilities",
-            "vulnerabilities__references",
-            "vulnerabilities__weaknesses",
-            "vulnerabilities__references__vulnerabilityseverity_set",
+            "affected_by_vulnerabilities",
+            "affected_by_vulnerabilities__references",
+            "affected_by_vulnerabilities__weaknesses",
+            "affected_by_vulnerabilities__references__vulnerabilityseverity_set",
+            "fixing_vulnerabilities",
+            "fixing_vulnerabilities__references",
+            "fixing_vulnerabilities__weaknesses",
+            "fixing_vulnerabilities__references__vulnerabilityseverity_set",
         )
         .paginated()
     )

vulnerabilities/migrations/0073_delete_packagerelatedvulnerability.py

@@ -0,0 +1,16 @@
+# Generated by Django 4.2.15 on 2024-10-15 10:37
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("vulnerabilities", "0072_remove_package_vulnerabilities_and_more"),
+    ]
+
+    operations = [
+        migrations.DeleteModel(
+            name="PackageRelatedVulnerability",
+        ),
+    ]

vulnerabilities/models.py

@@ -87,7 +87,7 @@ def affecting_vulnerabilities(self):
         """
         Return a queryset of Vulnerability that affect a package.
         """
-        return self.filter(packagerelatedvulnerability__fix=False)
+        return self.filter(affecting_packages__isnull=False)
 
     def with_cpes(self):
         """
@@ -187,12 +187,6 @@ class Vulnerability(models.Model):
         to="VulnerabilityReference", through="VulnerabilityRelatedReference"
     )
 
-    # packages = models.ManyToManyField(
-    #     to="Package",
-    #     through="PackageRelatedVulnerability",
-    #     related_name="all_vulnerabilities",
-    # )
-
     affecting_packages = models.ManyToManyField(
         to="Package",
         through="AffectedByPackageRelatedVulnerability",
@@ -452,7 +446,7 @@ def get_fixed_by_package_versions(self, purl: PackageURL, fix=True):
         }
 
         if fix:
-            filter_dict["packagerelatedvulnerability__fix"] = True
+            filter_dict["fixing_vulnerabilities__isnull"] = False
 
         return Package.objects.filter(**filter_dict).distinct()
 
@@ -474,15 +468,15 @@ def affected(self):
         """
         Return only packages affected by a vulnerability.
         """
-        return self.filter(packagerelatedvulnerability__fix=False)
+        return self.filter(affected_by_vulnerabilities__isnull=False)
 
     vulnerable = affected
 
     def fixing(self):
         """
         Return only packages fixing a vulnerability .
         """
-        return self.filter(packagerelatedvulnerability__fix=True)
+        return self.filter(fixing_vulnerabilities__isnull=False)
 
     def with_vulnerability_counts(self):
         return self.annotate(
@@ -585,6 +579,12 @@ def _vulnerable(self, vulnerable=True):
         """
         return self.with_is_vulnerable().filter(is_vulnerable=vulnerable)
 
+    def vulnerable(self):
+        """
+        Return only packages that are vulnerable.
+        """
+        return self.filter(affected_by_vulnerabilities__isnull=False)
+
 
 def get_purl_query_lookups(purl):
     """
@@ -606,12 +606,6 @@ class Package(PackageURLMixin):
     # https://github.com/package-url/packageurl-python/pull/67
     # gets merged
 
-    # vulnerabilities = models.ManyToManyField(
-    #     to="Vulnerability",
-    #     through="PackageRelatedVulnerability",
-    #     related_name="all_packages",
-    # )
-
     affected_by_vulnerabilities = models.ManyToManyField(
         to="Vulnerability",
         through="AffectedByPackageRelatedVulnerability",
@@ -751,6 +745,10 @@ def next_non_vulnerable_version(self):
         next_non_vulnerable, _ = self.get_non_vulnerable_versions()
         return next_non_vulnerable.version if next_non_vulnerable else None
 
+    @property
+    def vulnerabilities(self):
+        return self.affected_by_vulnerabilities.all() | self.fixing_vulnerabilities.all()
+
     @property
     def latest_non_vulnerable_version(self):
         """
@@ -880,107 +878,6 @@ def affecting_vulns(self):
         )
 
 
-class PackageRelatedVulnerability(models.Model):
-    """
-    Track the relationship between a Package and Vulnerability.
-    """
-
-    # TODO: Fix related_name
-    package = models.ForeignKey(
-        Package,
-        on_delete=models.CASCADE,
-    )
-
-    vulnerability = models.ForeignKey(
-        Vulnerability,
-        on_delete=models.CASCADE,
-    )
-
-    created_by = models.CharField(
-        max_length=100,
-        blank=True,
-        help_text="Fully qualified name of the improver prefixed with the"
-        "module name responsible for creating this relation. Eg:"
-        "vulnerabilities.importers.nginx.NginxBasicImprover",
-    )
-
-    from vulnerabilities.improver import MAX_CONFIDENCE
-
-    confidence = models.PositiveIntegerField(
-        default=MAX_CONFIDENCE,
-        validators=[MinValueValidator(0), MaxValueValidator(MAX_CONFIDENCE)],
-        help_text="Confidence score for this relation",
-    )
-
-    fix = models.BooleanField(
-        default=False,
-        db_index=True,
-        help_text="Does this relation fix the specified vulnerability ?",
-    )
-
-    class Meta:
-        unique_together = ["package", "vulnerability"]
-        verbose_name_plural = "PackageRelatedVulnerabilities"
-        indexes = [models.Index(fields=["fix"])]
-        ordering = ["package", "vulnerability"]
-
-    def __str__(self):
-        return f"{self.package.package_url} {self.vulnerability.vulnerability_id}"
-
-    def update_or_create(self, advisory):
-        """
-        Update if supplied record has more confidence than existing record
-        Create if doesn't exist
-        """
-        try:
-            existing = PackageRelatedVulnerability.objects.get(
-                vulnerability=self.vulnerability, package=self.package
-            )
-            if self.confidence > existing.confidence:
-                existing.created_by = self.created_by
-                existing.confidence = self.confidence
-                existing.fix = self.fix
-                existing.save()
-                # TODO: later we want these to be part of a log field in the DB
-                logger.info(
-                    f"Confidence improved for {self.package} R {self.vulnerability}, "
-                    f"new confidence: {self.confidence}"
-                )
-            self.add_package_vulnerability_changelog(advisory=advisory)
-
-        except self.DoesNotExist:
-            PackageRelatedVulnerability.objects.create(
-                vulnerability=self.vulnerability,
-                created_by=self.created_by,
-                package=self.package,
-                confidence=self.confidence,
-                fix=self.fix,
-            )
-
-            logger.info(
-                f"New relationship {self.package} R {self.vulnerability}, "
-                f"fix: {self.fix}, confidence: {self.confidence}"
-            )
-
-            self.add_package_vulnerability_changelog(advisory=advisory)
-
-    @transaction.atomic
-    def add_package_vulnerability_changelog(self, advisory):
-        from vulnerabilities.utils import get_importer_name
-
-        importer_name = get_importer_name(advisory)
-        if self.fix:
-            change_logger = PackageChangeLog.log_fixing
-        else:
-            change_logger = PackageChangeLog.log_affected_by
-        change_logger(
-            package=self.package,
-            importer=importer_name,
-            source_url=advisory.url or None,
-            related_vulnerability=str(self.vulnerability),
-        )
-
-
 class PackageRelatedVulnerabilityBase(models.Model):
     """
     Abstract base class for package-vulnerability relations.