Use new severity model for risk calculation

Signed-off-by: Keshav Priyadarshi <[email protected]>

Author: keshav-space

vulnerabilities/migrations/0075_alter_vulnerabilityseverity_options_and_more.py renamed to vulnerabilities/migrations/0077_alter_vulnerabilityseverity_options_and_more.py

@@ -153,7 +153,7 @@ def reverse_populate_vulnerability_model_with_severities(apps, schema_editor):
         pass
 
     dependencies = [
-        ("vulnerabilities", "0074_update_pysec_advisory_created_by"),
+        ("vulnerabilities", "0076_alter_packagechangelog_software_version_and_more"),
     ]
 
     operations = [

vulnerabilities/risk.py

@@ -10,10 +10,6 @@
 
 from urllib.parse import urlparse
 
-from vulnerabilities.models import AffectedByPackageRelatedVulnerability
-from vulnerabilities.models import Exploit
-from vulnerabilities.models import Package
-from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference
 from vulnerabilities.severity_systems import EPSS
 from vulnerabilities.weight_config import WEIGHT_CONFIG
@@ -40,7 +36,7 @@ def get_weighted_severity(severities):
 
     score_list = []
     for severity in severities:
-        parsed_url = urlparse(severity.reference.url)
+        parsed_url = urlparse(severity.url)
         severity_source = parsed_url.netloc.replace("www.", "", 1)
         weight = WEIGHT_CONFIG.get(severity_source, DEFAULT_WEIGHT)
         max_weight = float(weight) / 10
@@ -90,34 +86,32 @@ def get_exploitability_level(exploits, references, severities):
     return exploit_level
 
 
-def compute_vulnerability_risk(vulnerability: Vulnerability):
+def compute_vulnerability_risk(vulnerability):
     """
     Risk may be expressed as a number ranging from 0 to 10.
     Risk is calculated from weighted severity and exploitability values.
     It is the maximum value of (the weighted severity multiplied by its exploitability) or 10
 
     Risk = min(weighted severity * exploitability, 10)
     """
-    references = vulnerability.references
-    severities = vulnerability.severities.select_related("reference")
-    exploits = Exploit.objects.filter(vulnerability=vulnerability)
-    if references.exists() or severities.exists() or exploits.exists():
+    severities = vulnerability.severities.all()
+    exploits = vulnerability.exploits.all()
+    reference = vulnerability.references.all()
+    if reference.exists() or severities.exists() or exploits.exists():
         weighted_severity = get_weighted_severity(severities)
-        exploitability = get_exploitability_level(exploits, references, severities)
+        exploitability = get_exploitability_level(exploits, reference, severities)
         return min(weighted_severity * exploitability, 10)
 
 
-def compute_package_risk(package: Package):
+def compute_package_risk(package):
     """
     Calculate the risk for a package by iterating over all vulnerabilities that affects this package
     and determining the associated risk.
     """
 
     result = []
-    for pkg_related_vul in AffectedByPackageRelatedVulnerability.objects.filter(
-        package=package
-    ).prefetch_related("vulnerability"):
-        if risk := compute_vulnerability_risk(pkg_related_vul.vulnerability):
+    for package_vulnerability in package.affectedbypackagerelatedvulnerability_set.all():
+        if risk := compute_vulnerability_risk(package_vulnerability.vulnerability):
             result.append(risk)
 
     if not result:

vulnerabilities/tests/test_risk.py

@@ -29,25 +29,20 @@ def vulnerability():
     vul = Vulnerability(vulnerability_id="VCID-Existing")
     vul.save()
 
-    reference1 = VulnerabilityReference.objects.create(
-        reference_id="",
+    severity1 = VulnerabilitySeverity.objects.create(
         url="https://nvd.nist.gov/vuln/detail/CVE-xxxx-xxx1",
-    )
-
-    VulnerabilitySeverity.objects.create(
-        reference=reference1,
         scoring_system=CVSSV3.identifier,
         scoring_elements="CVSS:3.0/AV:P/AC:H/PR:H/UI:R/S:C/C:H/I:H/A:N/E:H/RL:O/RC:R/CR:H/MAC:H/MC:L",
         value="6.5",
     )
 
-    VulnerabilitySeverity.objects.create(
-        reference=reference1,
+    severity2 = VulnerabilitySeverity.objects.create(
+        url="https://nvd.nist.gov/vuln/detail/CVE-xxxx-xxx1",
         scoring_system=GENERIC.identifier,
         value="MODERATE",  # 6.9
     )
-
-    VulnerabilityRelatedReference.objects.create(reference=reference1, vulnerability=vul)
+    vul.severities.add(severity1)
+    vul.severities.add(severity2)
 
     weaknesses = Weakness.objects.create(cwe_id=119)
     vul.weaknesses.add(weaknesses)
@@ -84,18 +79,13 @@ def high_epss_score():
     vul = Vulnerability(vulnerability_id="VCID-HIGH-EPSS")
     vul.save()
 
-    reference1 = VulnerabilityReference.objects.create(
-        reference_id="",
+    severity = VulnerabilitySeverity.objects.create(
         url="https://nvd.nist.gov/vuln/detail/CVE-xxxx-xxx3",
-    )
-
-    VulnerabilitySeverity.objects.create(
-        reference=reference1,
         scoring_system=EPSS.identifier,
         value=".9",
     )
+    vul.severities.add(severity)
 
-    VulnerabilityRelatedReference.objects.create(reference=reference1, vulnerability=vul)
     return vul.severities
 
 
@@ -105,28 +95,19 @@ def low_epss_score():
     vul = Vulnerability(vulnerability_id="VCID-LOW-EPSS")
     vul.save()
 
-    reference1 = VulnerabilityReference.objects.create(
-        reference_id="",
+    severity = VulnerabilitySeverity.objects.create(
         url="https://nvd.nist.gov/vuln/detail/CVE-xxxx-xxx4",
-    )
-
-    VulnerabilitySeverity.objects.create(
-        reference=reference1,
         scoring_system=EPSS.identifier,
         value=".3",
     )
+    vul.severities.add(severity)
 
-    VulnerabilityRelatedReference.objects.create(reference=reference1, vulnerability=vul)
     return vul.severities
 
 
 @pytest.mark.django_db
 def test_exploitability_level(
-    exploit,
-    vulnerability_with_exploit_ref,
-    high_epss_score,
-    low_epss_score,
-    vulnerability,
+    exploit, vulnerability_with_exploit_ref, high_epss_score, low_epss_score
 ):
 
     assert get_exploitability_level(exploit, None, None) == 2
@@ -137,9 +118,9 @@ def test_exploitability_level(
 
     assert (
         get_exploitability_level(
-            None,
-            vulnerability_with_exploit_ref.references,
-            vulnerability_with_exploit_ref.severities,
+            exploits=None,
+            references=vulnerability_with_exploit_ref.references.all(),
+            severities=vulnerability_with_exploit_ref.severities.all(),
         )
         == 1
     )
@@ -152,18 +133,13 @@ def test_get_weighted_severity(vulnerability):
     severities = vulnerability.severities.all()
     assert get_weighted_severity(severities) == 6.210000000000001
 
-    reference2 = VulnerabilityReference.objects.create(
-        reference_id="",
+    severity2 = VulnerabilitySeverity.objects.create(
         url="https://security-tracker.debian.org/tracker/CVE-2019-13057",
-    )
-
-    VulnerabilitySeverity.objects.create(
-        reference=reference2,
         scoring_system=GENERIC.identifier,
         value="CRITICAL",
     )
+    vulnerability.severities.add(severity2)
 
-    VulnerabilityRelatedReference.objects.create(reference=reference2, vulnerability=vulnerability)
     new_severities = vulnerability.severities.all()
     assert get_weighted_severity(new_severities) == 7