Refine the vulnerability affectation

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

Author: tdruez

product_portfolio/tests/test_models.py

@@ -601,21 +601,25 @@ def test_product_model_improve_packages_from_purldb(self, mock_update_from_purld
         self.assertEqual("apache-2.0", pp1.license_expression)
 
     def test_product_model_affected_by_vulnerabilities(self):
-        vulnerability1 = make_vulnerability(self.dataspace, risk_score=10.0)
-        vulnerability2 = make_vulnerability(
-            self.dataspace, affecting=[self.product1], risk_score=1.0
-        )
+        vulnerability1 = make_vulnerability(self.dataspace, risk_score=1.0)
+        vulnerability2 = make_vulnerability(self.dataspace, risk_score=10.0)
+        vulnerability3 = make_vulnerability(self.dataspace, risk_score=5.0)
 
+        vulnerability1.add_affected(self.product1)
         affected_by = self.product1.affected_by_vulnerabilities.all()
-        self.assertQuerySetEqual([vulnerability2], affected_by)
+        self.assertQuerySetEqual([vulnerability1], affected_by)
         self.product1.refresh_from_db()
-        # self.assertEqual(1.0, self.product1.risk_score)
+        self.assertEqual(1.0, self.product1.risk_score)
 
-        vulnerability1.add_affected(self.product1)
+        vulnerability2.add_affected(self.product1)
         affected_by = self.product1.affected_by_vulnerabilities.order_by("id")
         self.assertQuerySetEqual([vulnerability1, vulnerability2], affected_by)
         self.product1.refresh_from_db()
-        # self.assertEqual(10.0, self.product1.risk_score)
+        self.assertEqual(10.0, self.product1.risk_score)
+
+        vulnerability3.add_affected(self.product1)
+        self.product1.refresh_from_db()
+        self.assertEqual(10.0, self.product1.risk_score)
 
     def test_product_model_get_vulnerability_qs(self):
         package1 = make_package(self.dataspace)

vulnerabilities/models.py

@@ -205,6 +205,8 @@ 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()
 
     @classmethod
     def create_from_data(cls, dataspace, data, validate=False, affecting=None):
@@ -437,6 +439,12 @@ def update_risk_score(self):
         self.save(update_fields=["risk_score"])
         return self.risk_score
 
+    def add_affected_by(self, vulnerability):
+        """Add ``vulnerability`` as affecting this instance."""
+        through_defaults = {"dataspace_id": self.dataspace_id}
+        self.affected_by_vulnerabilities.add(vulnerability, through_defaults=through_defaults)
+        self.update_risk_score()
+
     def get_entry_for_package(self, vulnerablecode):
         if not self.package_url:
             return

vulnerabilities/tests/test_models.py

@@ -126,6 +126,30 @@ def test_vulnerability_mixin_update_risk_score(self):
         package2.update_risk_score()
         self.assertIsNone(package2.risk_score)
 
+    def test_vulnerability_mixin_add_affected_by(self):
+        package1 = make_package(self.dataspace)
+
+        vulnerability1 = make_vulnerability(self.dataspace, risk_score=1.0)
+        vulnerability2 = make_vulnerability(self.dataspace, risk_score=10.0)
+        vulnerability3 = make_vulnerability(self.dataspace, risk_score=5.0)
+
+        package1.add_affected_by(vulnerability1)
+        package1.refresh_from_db()
+        self.assertEqual("1.0", str(package1.risk_score))
+
+        package1.add_affected_by(vulnerability2)
+        package1.refresh_from_db()
+        self.assertEqual("10.0", str(package1.risk_score))
+
+        package1.add_affected_by(vulnerability3)
+        package1.refresh_from_db()
+        self.assertEqual("10.0", str(package1.risk_score))
+
+        self.assertEqual(package1, vulnerability1.affected_packages.get())
+        self.assertEqual(package1, vulnerability2.affected_packages.get())
+        self.assertEqual(package1, vulnerability3.affected_packages.get())
+        self.assertEqual(3, package1.affected_by_vulnerabilities.count())
+
     def test_vulnerability_model_affected_packages_m2m(self):
         package1 = make_package(self.dataspace)
         vulnerability1 = make_vulnerability(dataspace=self.dataspace, affecting=package1)