Simplify the vulnerability affectation system

Signed-off-by: tdruez <[email protected]>

Author: tdruez

product_portfolio/tests/test_models.py

@@ -155,8 +155,11 @@ def test_product_model_all_packages(self):
     def test_product_model_get_vulnerable_packages(self):
         self.assertEqual(0, self.product1.get_vulnerable_packages().count())
 
-        package1 = make_package(self.dataspace, is_vulnerable=True, risk_score=5.0)
+        package1 = make_package(self.dataspace)
+        vulnerability1 = make_vulnerability(self.dataspace, risk_score=5.0)
+        package1.add_affected_by(vulnerability1)
         make_product_package(self.product1, package1)
+
         self.assertEqual(1, self.product1.get_vulnerable_packages().count())
         self.assertEqual(0, self.product1.get_vulnerable_packages(risk_threshold=6.0).count())
         self.assertEqual(1, self.product1.get_vulnerable_packages(risk_threshold=4.0).count())

vulnerabilities/fetch.py

@@ -135,7 +135,7 @@ def create_or_update_vulnerability(
         if updated_fields:
             results["updated"] += 1
 
-    vulnerability.add_affected_packages(affected_packages)
+    vulnerability.add_affected(affected_packages)
     return vulnerability
 
 

vulnerabilities/models.py

@@ -172,41 +172,12 @@ def cve(self):
                 return alias
 
     def add_affected(self, instances):
-        """
-        Assign the ``instances`` (Package, Component, or Product) as affected by this
-        vulnerability.
-        """
-        from component_catalog.models import Component
-        from component_catalog.models import Package
-        from product_portfolio.models import Product
-
-        if not isinstance(instances, list):
+        """Assign the ``instances`` (Package or Product) as affected by this vulnerability."""
+        if not isinstance(instances, (list, tuple, models.QuerySet)):
             instances = [instances]
 
         for instance in instances:
-            if isinstance(instance, Package):
-                self.add_affected_packages([instance])
-            if isinstance(instance, Component):
-                self.add_affected_components([instance])
-            if isinstance(instance, Product):
-                self.add_affected_products([instance])
-
-    def add_affected_packages(self, packages):
-        """Assign the ``packages`` as affected by this vulnerability."""
-        through_defaults = {"dataspace_id": self.dataspace_id}
-        self.affected_packages.add(*packages, through_defaults=through_defaults)
-
-    def add_affected_components(self, components):
-        """Assign the ``components`` as affected by this vulnerability."""
-        through_defaults = {"dataspace_id": self.dataspace_id}
-        self.affected_components.add(*components, through_defaults=through_defaults)
-
-    def add_affected_products(self, products):
-        """Assign the ``products`` as affected by this vulnerability."""
-        through_defaults = {"dataspace_id": self.dataspace_id}
-        self.affected_products.add(*products, through_defaults=through_defaults)
-        for product in products:
-            product.update_risk_score()
+            instance.add_affected_by(vulnerability=self)
 
     @classmethod
     def create_from_data(cls, dataspace, data, validate=False, affecting=None):