Merge pull request #16 from ninoseki/use-cvss

feat: use cvss lib to calculate rating

pycomponents/cve_search/schemas.py

@@ -11,10 +11,12 @@ def severity(self) -> Optional[str]:
         if self.cvss is None:
             return None
 
+        # since it is impossible to determine whether the score based on CVSS v2 or v3,
+        # use CVSS v2's scale
         if self.cvss >= 7.0:
-            return "high"
+            return "High"
 
         if self.cvss >= 4.0:
-            return "medium"
+            return "Medium"
 
-        return "low"
+        return "Low"

pycomponents/osv/schemas.py

@@ -36,6 +36,11 @@ class AffectedItem(BaseModel):
     database_specific: DatabaseSpecific
 
 
+class Severity(BaseModel):
+    type: str
+    score: str
+
+
 class Vuln(BaseModel):
     id: str
     details: str
@@ -45,6 +50,7 @@ class Vuln(BaseModel):
     references: List[Reference] = Field(default_factory=list)
     affected: List[AffectedItem] = Field(default_factory=list)
     schema_version: str
+    severity: List[Severity] = Field(default_factory=list)
 
     @property
     def cve_id(self) -> Optional[str]:

pycomponents/rating.py

@@ -0,0 +1,47 @@
+from typing import Optional, Union, cast
+
+from cvss import CVSS2, CVSS3
+from cyclonedx.model.vulnerability import VulnerabilityRating
+
+from .cve_search import get_cve_search
+from .exceptions import CVESearchQueryException
+from .osv.schemas import Severity, Vuln
+
+
+def severity_to_rating(severity: Severity) -> Optional[VulnerabilityRating]:
+    severity_type_to_cvss = {
+        "CVSS_V3": CVSS3,
+        "CVSS_V2": CVSS2,
+    }
+
+    klass = severity_type_to_cvss.get(severity.type)
+    if klass is None:
+        return None
+
+    c = cast(Union[CVSS2, CVSS3], klass(severity.score))
+    score, _, _ = c.scores()
+    severity_str, _, _ = c.severities()
+    return VulnerabilityRating(score=score, severity=severity_str)
+
+
+class RatingFactory:
+    @staticmethod
+    def from_osv_vuln(vuln: Vuln) -> Optional[VulnerabilityRating]:
+        rating: Optional[VulnerabilityRating] = None
+
+        if len(vuln.severity) > 0:
+            # TODO: consider a case when there are multi severities
+            severity = vuln.severity[0]
+            rating = severity_to_rating(severity)
+            if rating is not None:
+                return rating
+
+        if vuln.cve_id is not None:
+            try:
+                cve_search = get_cve_search()
+                res = cve_search.query(vuln.cve_id)
+                rating = VulnerabilityRating(score=res.cvss, severity=res.severity)
+            except CVESearchQueryException:
+                pass
+
+        return rating

pycomponents/vulnerability.py

@@ -12,10 +12,10 @@
     VulnerabilitySource,
 )
 
-from .cve_search import get_cve_search
-from .exceptions import CVESearchQueryException, OSVQueryException
+from .exceptions import OSVQueryException
 from .osv import get_osv
 from .osv.schemas import Vuln
+from .rating import RatingFactory
 
 
 def to_datetime(s: str) -> datetime:
@@ -35,14 +35,7 @@ def from_osv_vuln(vuln: Vuln) -> Vulnerability:
             if source is not None:
                 references.append(VulnerabilityReference(source=source))
 
-        rating: Optional[VulnerabilityRating] = None
-        if vuln.cve_id is not None:
-            try:
-                cve_search = get_cve_search()
-                res = cve_search.query(vuln.cve_id)
-                rating = VulnerabilityRating(score=res.cvss, severity=res.severity)
-            except CVESearchQueryException:
-                pass
+        rating = RatingFactory.from_osv_vuln(vuln)
 
         vulnerability = Vulnerability(
             id=vuln.id,

pyproject.toml

@@ -14,6 +14,7 @@ readme = "README.md"
 python = "^3.8"
 arrow = "^1.2.2"
 cachetools = "^5.2.0"
+cvss = "^2.5"
 cyclonedx-python-lib = "^2.7.1"
 environs = "^9.5.0"
 httpx = "^0.23.0"

tests/test_rating.py

@@ -0,0 +1,17 @@
+import pytest
+
+from pycomponents.osv.schemas import Severity
+from pycomponents.rating import severity_to_rating
+
+
+@pytest.fixture
+def severity():
+    return Severity(
+        type="CVSS_V3", score="CVSS:3.0/AV:N/AC:L/PR:N/UI:N/S:C/C:H/I:N/A:N"
+    )
+
+
+def test_severity_to_rating(severity: Severity):
+    rating = severity_to_rating(severity)
+    assert rating.score == 8.6
+    assert rating.severity == "High"