From 6a86ec27c13ff5e413c5a5f96d9b7671646f9388 Mon Sep 17 00:00:00 2001
From: Paul Horton <phorton@sonatype.com>
Date: Wed, 15 Sep 2021 15:40:28 +0100
Subject: [PATCH] feat: helper methods for deriving Severity and SourceType

Signed-off-by: Paul Horton <phorton@sonatype.com>
---
 cyclonedx/model/vulnerability.py  | 49 +++++++++++++++++++++++-
 tests/test_model_vulnerability.py | 62 ++++++++++++++++++++++++++++++-
 2 files changed, 109 insertions(+), 2 deletions(-)

diff --git a/cyclonedx/model/vulnerability.py b/cyclonedx/model/vulnerability.py
index 831935f4..80d8cc2e 100644
--- a/cyclonedx/model/vulnerability.py
+++ b/cyclonedx/model/vulnerability.py
@@ -39,6 +39,29 @@ class VulnerabilitySourceType(Enum):
     OPEN_FAIR = 'Open FAIR'
     OTHER = 'Other'
 
+    @staticmethod
+    def get_from_vector(vector: str):
+        if vector.startswith('CVSS:3.'):
+            return VulnerabilitySourceType.CVSS_V3
+        elif vector.startswith('CVSS:2.'):
+            return VulnerabilitySourceType.CVSS_V2
+        elif vector.startswith('OWASP'):
+            return VulnerabilitySourceType.OWASP
+        else:
+            return VulnerabilitySourceType.OTHER
+
+    def get_localised_vector(self, vector: str) -> str:
+        """
+        This method will remove any Source Scheme type from the supplied vector.
+
+        For example if VulnerabilitySourceType.OWASP
+
+        :param vector:
+        :return:
+        """
+        if self == VulnerabilitySourceType.CVSS_V3:
+            return
+
 
 class VulnerabilitySeverity(Enum):
     """
@@ -51,6 +74,27 @@ class VulnerabilitySeverity(Enum):
     CRITICAL = 'Critical'
     UNKNOWN = 'Unknown'
 
+    @staticmethod
+    def get_from_cvss_scores(scores: tuple[float] = None):
+        if type(scores) is float:
+            scores = (scores, )
+
+        if scores is None:
+            return VulnerabilitySeverity.UNKNOWN
+
+        max_cvss_score = max(scores)
+
+        if max_cvss_score >= 9.0:
+            return VulnerabilitySeverity.CRITICAL
+        elif max_cvss_score >= 7.0:
+            return VulnerabilitySeverity.HIGH
+        elif max_cvss_score >= 4.0:
+            return VulnerabilitySeverity.MEDIUM
+        elif max_cvss_score > 0.0:
+            return VulnerabilitySeverity.LOW
+        else:
+            return VulnerabilitySeverity.NONE
+
 
 class VulnerabilityRating:
     """
@@ -71,7 +115,10 @@ def __init__(self, score_base: float = None, score_impact: float = None, score_e
         self._score_exploitability = score_exploitability
         self._severity = severity
         self._method = method
-        self._vector = vector
+        if self._method:
+            self._vector = self._method.get_localised_vector(vector=vector)
+        else:
+            self._vector = vector
 
     def get_base_score(self) -> float:
         return self._score_base
diff --git a/tests/test_model_vulnerability.py b/tests/test_model_vulnerability.py
index cf538ce7..cea32b2d 100644
--- a/tests/test_model_vulnerability.py
+++ b/tests/test_model_vulnerability.py
@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from cyclonedx.model.vulnerability import VulnerabilityRating
+from cyclonedx.model.vulnerability import VulnerabilityRating, VulnerabilitySeverity, VulnerabilitySourceType
 
 
 class TestModelVulnerability(TestCase):
@@ -16,3 +16,63 @@ def test_v_rating_scores_base_only(self):
     def test_v_rating_scores_all(self):
         vr = VulnerabilityRating(score_base=1.0, score_impact=3.5, score_exploitability=5.6)
         self.assertTrue(vr.has_score())
+
+    def test_v_severity_from_cvss_scores_single_critical(self):
+        self.assertEqual(
+            VulnerabilitySeverity.get_from_cvss_scores(9.1),
+            VulnerabilitySeverity.CRITICAL
+        )
+
+    def test_v_severity_from_cvss_scores_multiple_critical(self):
+        self.assertEqual(
+            VulnerabilitySeverity.get_from_cvss_scores((9.1, 9.5)),
+            VulnerabilitySeverity.CRITICAL
+        )
+
+    def test_v_severity_from_cvss_scores_single_high(self):
+        self.assertEqual(
+            VulnerabilitySeverity.get_from_cvss_scores(8.9),
+            VulnerabilitySeverity.HIGH
+        )
+
+    def test_v_severity_from_cvss_scores_single_medium(self):
+        self.assertEqual(
+            VulnerabilitySeverity.get_from_cvss_scores(4.2),
+            VulnerabilitySeverity.MEDIUM
+        )
+
+    def test_v_severity_from_cvss_scores_single_low(self):
+        self.assertEqual(
+            VulnerabilitySeverity.get_from_cvss_scores(1.1),
+            VulnerabilitySeverity.LOW
+        )
+
+    def test_v_severity_from_cvss_scores_single_none(self):
+        self.assertEqual(
+            VulnerabilitySeverity.get_from_cvss_scores(0.0),
+            VulnerabilitySeverity.NONE
+        )
+
+    def test_v_severity_from_cvss_scores_multiple_high(self):
+        self.assertEqual(
+            VulnerabilitySeverity.get_from_cvss_scores((1.2, 8.9, 2.2, 5.6)),
+            VulnerabilitySeverity.HIGH
+        )
+
+    def test_v_source_parse_cvss3_1(self):
+        self.assertEqual(
+            VulnerabilitySourceType.get_from_vector('CVSS:3.0/AV:L/AC:L/PR:N/UI:R/S:C/C:L/I:N/A:N'),
+            VulnerabilitySourceType.CVSS_V3
+        )
+
+    def test_v_source_parse_cvss2_1(self):
+        self.assertEqual(
+            VulnerabilitySourceType.get_from_vector('CVSS:2.0/AV:N/AC:L/Au:N/C:N/I:N/A:C'),
+            VulnerabilitySourceType.CVSS_V2
+        )
+
+    def test_v_source_parse_owasp_1(self):
+        self.assertEqual(
+            VulnerabilitySourceType.get_from_vector('OWASP/K9:M1:O0:Z2/D1:X1:W1:L3/C2:I1:A1:T1/F1:R1:S2:P3/50'),
+            VulnerabilitySourceType.OWASP
+        )