Optimize performance, refactor, and rename the add_vulnerability_risk_score function. Rename the help text for the model.

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

Author: ziadhany

vulnerabilities/migrations/0077_vulnerability_exploitability_and_more.py

@@ -1,4 +1,4 @@
-# Generated by Django 4.2.16 on 2024-11-08 14:07
+# Generated by Django 4.2.16 on 2024-11-12 12:16
 
 from django.db import migrations, models
 
@@ -15,7 +15,9 @@ class Migration(migrations.Migration):
             name="exploitability",
             field=models.DecimalField(
                 decimal_places=2,
-                help_text="Exploitability refers to the potential or probability of a software package vulnerability being \n        exploited by malicious actors to compromise systems, applications, or networks. \n        It is determined automatically by the discovery of exploits.",
+                help_text="""Exploitability indicates the likelihood that a vulnerability in a software package could 
+                 be used by malicious actors to compromise systems, applications, or networks.
+                 This metric is determined automatically based on the discovery of known exploits.""",
                 max_digits=4,
                 null=True,
             ),
@@ -25,7 +27,7 @@ class Migration(migrations.Migration):
             name="weighted_severity",
             field=models.DecimalField(
                 decimal_places=2,
-                help_text="Weighted Severity is the maximum value obtained when each Severity is multiplied by its associated Weight/10.",
+                help_text="Weighted severity is the highest value calculated by multiplying each severity by its corresponding weight, divided by 10.",
                 max_digits=4,
                 null=True,
             ),

vulnerabilities/models.py

@@ -206,16 +206,16 @@ class Vulnerability(models.Model):
         null=True,
         max_digits=4,
         decimal_places=2,
-        help_text="""Exploitability refers to the potential or probability of a software package vulnerability being 
-        exploited by malicious actors to compromise systems, applications, or networks. 
-        It is determined automatically by the discovery of exploits.""",
+        help_text=""""Exploitability indicates the likelihood that a vulnerability in a software package could be used 
+        by malicious actors to compromise systems, applications, or networks. 
+        This metric is determined automatically based on the discovery of known exploits.""",
     )
 
     weighted_severity = models.DecimalField(
         null=True,
         max_digits=4,
         decimal_places=2,
-        help_text="Weighted Severity is the maximum value obtained when each Severity is multiplied by its associated Weight/10.",
+        help_text="Weighted severity is the highest value calculated by multiplying each severity by its corresponding weight, divided by 10.",
     )
 
     @property
@@ -228,8 +228,8 @@ def risk_score(self):
         Risk = min(weighted severity * exploitability, 10)
         """
         if self.exploitability is not None and self.weighted_severity is not None:
-            return f"{min(float(self.exploitability) * float(self.weighted_severity), 10.0):.2f}"
-        return None
+            risk_score = min(float(self.exploitability) * float(self.weighted_severity), 10.0)
+            return f"{risk_score:.2f}".rstrip("0").rstrip(".")
 
     objects = VulnerabilityQuerySet.as_manager()
 

vulnerabilities/pipelines/compute_package_risk.py

@@ -9,12 +9,11 @@
 
 from aboutcode.pipeline import LoopProgress
 
-from vulnerabilities.models import AffectedByPackageRelatedVulnerability
 from vulnerabilities.models import Package
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.pipelines import VulnerableCodePipeline
 from vulnerabilities.risk import compute_package_risk
-from vulnerabilities.risk import compute_vulnerability_risk
+from vulnerabilities.risk import compute_vulnerability_risk_factors
 
 
 class ComputePackageRiskPipeline(VulnerableCodePipeline):
@@ -29,11 +28,16 @@ class ComputePackageRiskPipeline(VulnerableCodePipeline):
 
     @classmethod
     def steps(cls):
-        return (cls.add_vulnerability_risk_score, cls.add_package_risk_score)
+        return (
+            cls.compute_and_store_vulnerability_risk_score,
+            cls.compute_and_store_package_risk_score,
+        )
 
-    def add_vulnerability_risk_score(self):
-        affected_vulnerabilities = Vulnerability.objects.filter(
-            affectedbypackagerelatedvulnerability__isnull=False
+    def compute_and_store_vulnerability_risk_score(self):
+        affected_vulnerabilities = (
+            Vulnerability.objects.filter(affectedbypackagerelatedvulnerability__isnull=False)
+            .prefetch_related("references")
+            .only("references", "exploits")
         )
 
         self.log(
@@ -47,11 +51,13 @@ def add_vulnerability_risk_score(self):
         batch_size = 5000
 
         for vulnerability in progress.iter(affected_vulnerabilities.paginated()):
+            references = vulnerability.references
+            severities = vulnerability.severities.select_related("reference")
 
-            vulnerability = compute_vulnerability_risk(vulnerability)
-
-            if not vulnerability:
-                continue
+            (
+                vulnerability.weighted_severity,
+                vulnerability.exploitability,
+            ) = compute_vulnerability_risk_factors(references, severities, vulnerability.exploits)
 
             updatables.append(vulnerability)
 
@@ -68,7 +74,7 @@ def add_vulnerability_risk_score(self):
             f"Successfully added risk score for {updated_vulnerability_count:,d} vulnerability"
         )
 
-    def add_package_risk_score(self):
+    def compute_and_store_package_risk_score(self):
         affected_packages = Package.objects.filter(
             affected_by_vulnerabilities__isnull=False
         ).distinct()

vulnerabilities/risk.py

@@ -10,6 +10,8 @@
 
 from urllib.parse import urlparse
 
+from django.db.models import Prefetch
+
 from vulnerabilities.models import AffectedByPackageRelatedVulnerability
 from vulnerabilities.models import Exploit
 from vulnerabilities.models import Package
@@ -27,6 +29,8 @@ def get_weighted_severity(severities):
     by its associated Weight/10.
     Example of Weighted Severity: max(7*(10/10), 8*(3/10), 6*(8/10)) = 7
     """
+    if not severities:
+        return 0
 
     score_map = {
         "low": 3,
@@ -90,33 +94,21 @@ def get_exploitability_level(exploits, references, severities):
     return exploit_level
 
 
-def compute_vulnerability_risk(vulnerability: Vulnerability):
+def compute_vulnerability_risk_factors(references, severities, exploits):
     """
-    Computes the risk score for a given vulnerability.
-
-    Risk is expressed as a number ranging from 0 to 10 and is calculated based on:
-    - Weighted severity: a value derived from the associated severities of the vulnerability.
-    - Exploitability: a measure of how easily the vulnerability can be exploited.
-
-    The risk score is computed as:
-        Risk = min(weighted_severity * exploitability, 10)
+    Compute weighted severity and exploitability for a vulnerability.
 
     Args:
-        vulnerability (Vulnerability): The vulnerability object to compute the risk for.
+        references (list): References linked to the vulnerability.
+        severities (list): Severity levels of the vulnerability.
+        exploits (list): Exploit details for the vulnerability.
 
     Returns:
-        Vulnerability: The updated vulnerability object with computed risk-related attributes.
-
-    Notes:
-        - If there are no associated references, severities, or exploits, the computation is skipped.
+        tuple: (weighted_severity, exploitability).
     """
-    references = vulnerability.references
-    severities = vulnerability.severities.select_related("reference")
-    exploits = Exploit.objects.filter(vulnerability=vulnerability)
-    if references.exists() or severities.exists() or exploits.exists():
-        vulnerability.weighted_severity = get_weighted_severity(severities)
-        vulnerability.exploitability = get_exploitability_level(exploits, references, severities)
-        return vulnerability
+    weighted_severity = get_weighted_severity(severities)
+    exploitability = get_exploitability_level(exploits, references, severities)
+    return weighted_severity, exploitability
 
 
 def compute_package_risk(package: Package):
@@ -126,9 +118,15 @@ def compute_package_risk(package: Package):
     """
 
     result = []
-    for pkg_related_vul in AffectedByPackageRelatedVulnerability.objects.filter(
+    affected_pkg_related_vul = AffectedByPackageRelatedVulnerability.objects.filter(
         package=package
-    ).prefetch_related("vulnerability"):
+    ).prefetch_related(
+        Prefetch(
+            "vulnerability",
+            queryset=Vulnerability.objects.only("weighted_severity", "exploitability"),
+        )
+    )
+    for pkg_related_vul in affected_pkg_related_vul:
         if risk := pkg_related_vul.vulnerability.risk_score:
             result.append(float(risk))
 

vulnerabilities/templates/vulnerability_details.html

@@ -124,9 +124,9 @@
 
                             <tr>
                                 <td class="two-col-left"
-                                    data-tooltip="Exploitability refers to the potential or probability of a software package vulnerability being 
-                                    exploited by malicious actors to compromise systems, applications, or networks. 
-                                    It is determined automatically by the discovery of exploits.">
+                                    data-tooltip="Exploitability indicates the likelihood that a vulnerability in a software package could be used 
+                                    by malicious actors to compromise systems, applications, or networks. 
+                                    This metric is determined automatically based on the discovery of known exploits.">
                                     Exploitability</td>
                                 <td class="two-col-right wrap-strings">
                                     {{ vulnerability.exploitability }}
@@ -135,7 +135,7 @@
 
                             <tr>
                                 <td class="two-col-left"
-                                    data-tooltip="Weighted Severity is the maximum value obtained when each Severity is multiplied by its associated Weight/10."
+                                    data-tooltip="Weighted severity is the highest value calculated by multiplying each severity by its corresponding weight, divided by 10."
                                 >Weighted Severity</td>
                                 <td class="two-col-right wrap-strings">
                                     {{ vulnerability.weighted_severity }}

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 f"{pkg.risk_score:.2f}" == "3.10"
+    assert pkg.risk_score == Decimal("10")

vulnerabilities/tests/test_risk.py

@@ -15,7 +15,7 @@
 from vulnerabilities.models import VulnerabilityRelatedReference
 from vulnerabilities.models import VulnerabilitySeverity
 from vulnerabilities.models import Weakness
-from vulnerabilities.risk import compute_vulnerability_risk
+from vulnerabilities.risk import compute_vulnerability_risk_factors
 from vulnerabilities.risk import get_exploitability_level
 from vulnerabilities.risk import get_weighted_severity
 from vulnerabilities.severity_systems import CVSSV3
@@ -169,6 +169,44 @@ def test_get_weighted_severity(vulnerability):
 
 
 @pytest.mark.django_db
-def test_compute_vulnerability_risk(vulnerability):
-    vulnerability = compute_vulnerability_risk(vulnerability)
-    assert vulnerability.risk_score == str(3.11)
+def test_compute_vulnerability_risk_factors(vulnerability):
+    assert compute_vulnerability_risk_factors(
+        vulnerability.references, vulnerability.severities, vulnerability.exploits
+    ) == (6.210000000000001, 2)
+    assert compute_vulnerability_risk_factors(
+        vulnerability.references, vulnerability.severities, None
+    ) == (
+        6.210000000000001,
+        0.5,
+    )
+    assert compute_vulnerability_risk_factors(
+        vulnerability.references, None, vulnerability.exploits
+    ) == (
+        0,
+        2,
+    )
+    assert compute_vulnerability_risk_factors(None, None, None) == (0, 0.5)
+
+
+@pytest.mark.django_db
+def test_get_vulnerability_risk_score(vulnerability):
+    vulnerability.weighted_severity = 6.0
+    vulnerability.exploitability = 2
+
+    assert vulnerability.risk_score == "10"  # max risk_score can be reached
+
+    vulnerability.weighted_severity = 6
+    vulnerability.exploitability = 0.5
+    assert vulnerability.risk_score == "3"
+
+    vulnerability.weighted_severity = 5.6
+    vulnerability.exploitability = 0.5
+    assert vulnerability.risk_score == "2.8"
+
+    vulnerability.weighted_severity = None
+    vulnerability.exploitability = 0.5
+    assert vulnerability.risk_score is None
+
+    vulnerability.weighted_severity = None
+    vulnerability.exploitability = None
+    assert vulnerability.risk_score is None