Refactor the risk score calculation for vulnerabilities and packages.

Update the tests for exploits and the simple_risk_pipeline.

Signed-off-by: ziad hany <[email protected]>

Author: ziadhany

vulnerabilities/pipelines/compute_package_risk.py

@@ -6,7 +6,6 @@
 # See https://github.com/aboutcode-org/vulnerablecode for support or download.
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
-
 from aboutcode.pipeline import LoopProgress
 from django.db.models import Prefetch
 
@@ -35,12 +34,14 @@ def steps(cls):
         )
 
     def compute_and_store_vulnerability_risk_score(self):
-        affected_vulnerabilities = Vulnerability.objects.filter(
-            affectedbypackagerelatedvulnerability__isnull=False
-        ).prefetch_related(
-            "references",
-            "severities",
-            "exploits",
+        affected_vulnerabilities = (
+            Vulnerability.objects.filter(affecting_packages__isnull=False)
+            .prefetch_related(
+                "references",
+                "severities",
+                "exploits",
+            )
+            .distinct()
         )
 
         self.log(
@@ -53,35 +54,43 @@ def compute_and_store_vulnerability_risk_score(self):
         updated_vulnerability_count = 0
         batch_size = 5000
 
-        for vulnerability in progress.iter(affected_vulnerabilities.paginated()):
+        for vulnerability in progress.iter(affected_vulnerabilities.paginated(per_page=batch_size)):
             severities = vulnerability.severities.all()
             references = vulnerability.references.all()
+            exploits = vulnerability.exploits.all()
 
-            (
-                vulnerability.weighted_severity,
-                vulnerability.exploitability,
-            ) = compute_vulnerability_risk_factors(references, severities, vulnerability.exploits)
+            weighted_severity, exploitability = compute_vulnerability_risk_factors(
+                references=references,
+                severities=severities,
+                exploits=exploits,
+            )
+            vulnerability.weighted_severity = weighted_severity
+            vulnerability.exploitability = exploitability
 
             updatables.append(vulnerability)
 
             if len(updatables) >= batch_size:
-                updated_vulnerability_count += bulk_update_vulnerability_risk_score(
-                    vulnerabilities=updatables,
+                updated_vulnerability_count += bulk_update(
+                    model=Vulnerability,
+                    items=updatables,
+                    fields=["weighted_severity", "exploitability"],
                     logger=self.log,
                 )
-        updated_vulnerability_count += bulk_update_vulnerability_risk_score(
-            vulnerabilities=updatables,
+
+        updated_vulnerability_count += bulk_update(
+            model=Vulnerability,
+            items=updatables,
+            fields=["weighted_severity", "exploitability"],
             logger=self.log,
         )
+
         self.log(
             f"Successfully added risk score for {updated_vulnerability_count:,d} vulnerability"
         )
 
     def compute_and_store_package_risk_score(self):
         affected_packages = (
-            Package.objects.filter(affected_by_vulnerabilities__isnull=False)
-            .only("id")
-            .prefetch_related(
+            Package.objects.filter(affected_by_vulnerabilities__isnull=False).prefetch_related(
                 Prefetch(
                     "affectedbypackagerelatedvulnerability_set__vulnerability",
                     queryset=Vulnerability.objects.only("weighted_severity", "exploitability"),
@@ -111,38 +120,28 @@ def compute_and_store_package_risk_score(self):
             updatables.append(package)
 
             if len(updatables) >= batch_size:
-                updated_package_count += bulk_update_package_risk_score(
-                    packages=updatables,
+                updated_package_count += bulk_update(
+                    model=Package,
+                    items=updatables,
+                    fields=["risk_score"],
                     logger=self.log,
                 )
-        updated_package_count += bulk_update_package_risk_score(
-            packages=updatables,
+        updated_package_count += bulk_update(
+            model=Package,
+            items=updatables,
+            fields=["risk_score"],
             logger=self.log,
         )
         self.log(f"Successfully added risk score for {updated_package_count:,d} package")
 
 
-def bulk_update_package_risk_score(packages, logger):
-    package_count = 0
-    if packages:
-        try:
-            Package.objects.bulk_update(objs=packages, fields=["risk_score"])
-            package_count += len(packages)
-        except Exception as e:
-            logger(f"Error updating packages: {e}")
-        packages.clear()
-    return package_count
-
-
-def bulk_update_vulnerability_risk_score(vulnerabilities, logger):
-    vulnerabilities_count = 0
-    if vulnerabilities:
+def bulk_update(model, items, fields, logger):
+    item_count = 0
+    if items:
         try:
-            Vulnerability.objects.bulk_update(
-                objs=vulnerabilities, fields=["weighted_severity", "exploitability"]
-            )
-            vulnerabilities_count += len(vulnerabilities)
+            model.objects.bulk_update(objs=items, fields=fields)
+            item_count += len(items)
         except Exception as e:
-            logger(f"Error updating vulnerability: {e}")
-        vulnerabilities.clear()
-    return vulnerabilities_count
+            logger(f"Error updating {model.__name__}: {e}")
+        items.clear()
+    return item_count

vulnerabilities/risk.py

@@ -104,9 +104,8 @@ def compute_package_risk(package):
     and determining the associated risk.
     """
     result = []
-    vulnerabilities = package.vulnerabilities.all()
-    for vulnerability in vulnerabilities:
-        if risk := vulnerability.risk_score:
+    for vulnerability in package.affectedbypackagerelatedvulnerability_set.all():
+        if risk := vulnerability.vulnerability.risk_score:
             result.append(float(risk))
 
     if not result:

vulnerabilities/tests/pipelines/test_compute_package_risk.py

@@ -31,4 +31,4 @@ def test_simple_risk_pipeline(vulnerability):
     improver.execute()
 
     pkg = Package.objects.get(type="pypi", name="foo", version="2.3.0")
-    assert pkg.risk_score == Decimal("10")
+    assert pkg.risk_score == Decimal("3.1")  # max( 6.9 * 9/10 , 6.5 * 9/10 ) * .5 = 3.105

vulnerabilities/tests/test_risk.py

@@ -145,17 +145,18 @@ def test_get_weighted_severity(vulnerability):
 
 
 @pytest.mark.django_db
-def test_compute_vulnerability_risk_factors(vulnerability):
+def test_compute_vulnerability_risk_factors(vulnerability, exploit):
     severities = vulnerability.severities.all()
     references = vulnerability.references.all()
 
-    assert compute_vulnerability_risk_factors(references, severities, vulnerability.exploits) == (
+    assert compute_vulnerability_risk_factors(references, severities, exploit) == (
         6.2,
         2,
     )
 
     assert compute_vulnerability_risk_factors(references, severities, None) == (6.2, 0.5)
-    assert compute_vulnerability_risk_factors(references, None, vulnerability.exploits) == (0, 2)
+
+    assert compute_vulnerability_risk_factors(references, None, exploit) == (0, 2)
 
     assert compute_vulnerability_risk_factors(None, None, None) == (0, 0.5)