Create a pipeline for package risk

Signed-off-by: ziadhany <ziadhany2016@gmail.com>

Author: ziadhany

vulnerabilities/api.py

@@ -359,6 +359,7 @@ class Meta:
             "latest_non_vulnerable_version",
             "affected_by_vulnerabilities",
             "fixing_vulnerabilities",
+            "risk",
         ]
 
 

vulnerabilities/improvers/__init__.py

@@ -14,29 +14,31 @@
 from vulnerabilities.pipelines import enhance_with_kev
 from vulnerabilities.pipelines import enhance_with_metasploit
 from vulnerabilities.pipelines import flag_ghost_packages
+from vulnerabilities.pipelines import risk_package
 
 IMPROVERS_REGISTRY = [
-    valid_versions.GitHubBasicImprover,
-    valid_versions.GitLabBasicImprover,
-    valid_versions.NginxBasicImprover,
-    valid_versions.ApacheHTTPDImprover,
-    valid_versions.DebianBasicImprover,
-    valid_versions.NpmImprover,
-    valid_versions.ElixirImprover,
-    valid_versions.ApacheTomcatImprover,
-    valid_versions.ApacheKafkaImprover,
-    valid_versions.IstioImprover,
-    valid_versions.DebianOvalImprover,
-    valid_versions.UbuntuOvalImprover,
-    valid_versions.OSSFuzzImprover,
-    valid_versions.RubyImprover,
-    valid_versions.GithubOSVImprover,
-    vulnerability_status.VulnerabilityStatusImprover,
-    valid_versions.CurlImprover,
-    flag_ghost_packages.FlagGhostPackagePipeline,
-    enhance_with_kev.VulnerabilityKevPipeline,
-    enhance_with_metasploit.MetasploitImproverPipeline,
-    enhance_with_exploitdb.ExploitDBImproverPipeline,
+    # valid_versions.GitHubBasicImprover,
+    # valid_versions.GitLabBasicImprover,
+    # valid_versions.NginxBasicImprover,
+    # valid_versions.ApacheHTTPDImprover,
+    # valid_versions.DebianBasicImprover,
+    # valid_versions.NpmImprover,
+    # valid_versions.ElixirImprover,
+    # valid_versions.ApacheTomcatImprover,
+    # valid_versions.ApacheKafkaImprover,
+    # valid_versions.IstioImprover,
+    # valid_versions.DebianOvalImprover,
+    # valid_versions.UbuntuOvalImprover,
+    # valid_versions.OSSFuzzImprover,
+    # valid_versions.RubyImprover,
+    # valid_versions.GithubOSVImprover,
+    # vulnerability_status.VulnerabilityStatusImprover,
+    # valid_versions.CurlImprover,
+    # flag_ghost_packages.FlagGhostPackagePipeline,
+    # enhance_with_kev.VulnerabilityKevPipeline,
+    # enhance_with_metasploit.MetasploitImproverPipeline,
+    # enhance_with_exploitdb.ExploitDBImproverPipeline,
+    risk_package.RiskPackagePipeline,
 ]
 
 IMPROVERS_REGISTRY = {

vulnerabilities/migrations/0074_package_risk.py

@@ -0,0 +1,23 @@
+# Generated by Django 4.2.16 on 2024-10-22 06:49
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("vulnerabilities", "0073_delete_packagerelatedvulnerability"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="package",
+            name="risk",
+            field=models.DecimalField(
+                decimal_places=2,
+                help_text="Enter a risk score between 0.00 and 10.00, where higher values indicate greater vulnerability risk for the package.",
+                max_digits=4,
+                null=True,
+            ),
+        ),
+    ]

vulnerabilities/models.py

@@ -636,6 +636,14 @@ class Package(PackageURLMixin):
         help_text="True if the package does not exist in the upstream package manager or its repository.",
     )
 
+    risk = models.DecimalField(
+        null=True,
+        max_digits=4,
+        decimal_places=2,
+        help_text="Enter a risk score between 0.00 and 10.00, where higher values "
+        "indicate greater vulnerability risk for the package.",
+    )
+
     objects = PackageQuerySet.as_manager()
 
     def save(self, *args, **kwargs):

vulnerabilities/pipelines/risk_package.py

@@ -0,0 +1,30 @@
+from vulnerabilities.models import Package
+from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.risk import calculate_pkg_risk
+
+
+class RiskPackagePipeline(VulnerableCodePipeline):
+    """
+    Risk Assessment Pipeline for Package Vulnerabilities: Iterate through the packages and evaluate their associated risk.
+    """
+
+    pipeline_id = "risk_package"
+    license_expression = None
+
+    @classmethod
+    def steps(cls):
+        return (cls.add_risk_package,)
+
+    def add_risk_package(self):
+        self.log(f"Add risk package pipeline ")
+
+        updatables = []
+        for pkg in Package.objects.filter(affected_by_vulnerabilities__isnull=False):
+            risk = calculate_pkg_risk(pkg)
+            pkg.risk = risk
+            updatables.append(pkg)
+
+        # Bulk update the 'risk' field for all packages
+        Package.objects.bulk_update(objs=updatables, fields=["risk"], batch_size=1000)
+
+        self.log(f"Successfully added risk package pipeline ")

vulnerabilities/risk.py

@@ -1,15 +1,16 @@
 import os
-import re
 
+from vulnerabilities.models import AffectedByPackageRelatedVulnerability
 from vulnerabilities.models import Exploit
 from vulnerabilities.models import Package
-from vulnerabilities.models import PackageRelatedVulnerability
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference
 from vulnerabilities.severity_systems import EPSS
 from vulnerabilities.utils import load_json
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+WEIGHT_CONFIG_PATH = os.path.join(BASE_DIR, "../weight_config.json")
+DEFAULT_WEIGHT = 1
 
 
 def get_weighted_severity(severities):
@@ -18,8 +19,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
     """
-    weight_config_path = os.path.join(BASE_DIR, "..", "weight_config.json")
-    weight_config = load_json(weight_config_path)
+
+    weight_config = load_json(WEIGHT_CONFIG_PATH)
 
     score_map = {
         "low": 3,
@@ -33,11 +34,12 @@ def get_weighted_severity(severities):
 
     score_list = []
     for severity in severities:
-        weights = [
-            value
-            for regex_key, value in weight_config.items()
-            if re.match(regex_key, severity.reference.url)
-        ]
+        weights = []
+        for key, value in weight_config.items():
+            if severity.reference.url.startswith(key):
+                weights.append(value)
+                continue
+            weights.append(DEFAULT_WEIGHT)
 
         if not weights:
             return 0
@@ -113,14 +115,13 @@ def calculate_pkg_risk(package: Package):
     """
 
     result = []
-    for pkg_related_vul in PackageRelatedVulnerability.objects.filter(
-        package=package, fix=False
+    for pkg_related_vul in AffectedByPackageRelatedVulnerability.objects.filter(
+        package=package
     ).prefetch_related("vulnerability"):
-        if pkg_related_vul:
-            risk = calculate_vulnerability_risk(pkg_related_vul.vulnerability)
-            if not risk:
-                continue
-            result.append(risk)
+        risk = calculate_vulnerability_risk(pkg_related_vul.vulnerability)
+        if not risk:
+            continue
+        result.append(risk)
 
     if not result:
         return

vulnerabilities/templates/package_details.html

@@ -117,8 +117,8 @@
                                         Risk
                                     </td>
                                     <td class="two-col-right">
-                                        {% if risk %}
-                                        <a target="_self">{{ risk  }}</a>
+                                        {% if package.risk %}
+                                        <a target="_self">{{ package.risk }}</a>
                                         {% endif %}
                                     </td>
                                 </tr>

vulnerabilities/tests/pipelines/test_risk_pipeline.py

@@ -0,0 +1,24 @@
+import pytest
+
+from vulnerabilities.models import AffectedByPackageRelatedVulnerability
+from vulnerabilities.models import Package
+from vulnerabilities.pipelines.risk_package import RiskPackagePipeline
+from vulnerabilities.tests.test_risk import vulnerability
+
+
+@pytest.mark.django_db
+def test_simple_risk_pipeline(vulnerability):
+    pkg = Package.objects.create(type="pypi", name="foo", version="2.3.0")
+    assert Package.objects.count() == 1
+
+    improver = RiskPackagePipeline()
+    improver.execute()
+
+    assert pkg.risk is None
+
+    AffectedByPackageRelatedVulnerability.objects.create(package=pkg, vulnerability=vulnerability)
+    improver = RiskPackagePipeline()
+    improver.execute()
+
+    pkg = Package.objects.get(type="pypi", name="foo", version="2.3.0")
+    assert str(pkg.risk) == str(3.11)

vulnerabilities/tests/test_api.py

@@ -625,6 +625,7 @@ def test_api_with_lesser_and_greater_fixed_by_packages(self):
                 }
             ],
             "resource_url": "http://testserver/packages/pkg:maven/com.fasterxml.jackson.core/jackson-databind@2.13.1",
+            "risk": None,
         }
 
         assert response == expected

vulnerabilities/views.py

@@ -123,7 +123,6 @@ def get_context_data(self, **kwargs):
         context["fixing_vulnerabilities"] = package.fixing.order_by("vulnerability_id")
         context["package_search_form"] = PackageSearchForm(self.request.GET)
         context["fixed_package_details"] = package.fixed_package_details
-        context["risk"] = calculate_pkg_risk(package)
         context["history"] = list(package.history)
         return context