feat: adding EPSS data (intel#3130)

Co-authored-by: Terri Oda <[email protected]>

Author: Rexbeast2

cve_bin_tool/cvedb.py

@@ -22,7 +22,13 @@
 from rich.progress import track
 
 from cve_bin_tool.async_utils import run_coroutine
-from cve_bin_tool.data_sources import curl_source, gad_source, nvd_source, osv_source
+from cve_bin_tool.data_sources import (
+    curl_source,
+    epss_source,
+    gad_source,
+    nvd_source,
+    osv_source,
+)
 from cve_bin_tool.error_handler import ERROR_CODES, CVEDBError, ErrorMode, SigningError
 from cve_bin_tool.fetch_json_db import Fetch_JSON_DB
 from cve_bin_tool.log import LOGGER
@@ -88,6 +94,8 @@ def __init__(
         self.cve_count = -1
         self.all_cve_entries: list[dict[str, Any]] | None = None
 
+        self.epss_data = None
+
         self.exploits_list: list[Any] = []
         self.exploit_count = 0
 
@@ -126,6 +134,8 @@ async def refresh(self) -> None:
         if self.version_check:
             check_latest_version()
 
+        epss = epss_source.Epss_Source()
+        self.epss_data = await epss.update_epss()
         await self.get_data()
 
     def refresh_cache_and_update_db(self) -> None:
@@ -342,7 +352,29 @@ def insert_queries(self):
         )
         VALUES (?,?,?)
         """
-        return insert_severity, insert_cve_range, insert_exploit
+        insert_cve_metrics = """
+        INSERT or REPLACE INTO cve_metrics (
+            cve_number,
+            metric_id,
+            metric_score,
+            metric_field
+        )
+        VALUES (?, ?, ?, ?)
+        """
+        insert_metrics = """
+            INSERT or REPLACE INTO metrics (
+                metrics_id,
+                metrics_name
+            )
+            VALUES (?, ?)
+        """
+        return (
+            insert_severity,
+            insert_cve_range,
+            insert_exploit,
+            insert_cve_metrics,
+            insert_metrics,
+        )
 
     def init_database(self) -> None:
         """Initialize db tables used for storing cve/version data"""
@@ -429,6 +461,9 @@ def populate_db(self) -> None:
         we'll need a better parser to match those together.
         """
 
+        self.store_epss_data()
+        self.populate_metrics()
+
         for idx, data in enumerate(self.data):
             _, source_name = data
 
@@ -457,7 +492,7 @@ def populate_db(self) -> None:
             self.db_close()
 
     def populate_severity(self, severity_data, cursor, data_source):
-        (insert_severity, _, _) = self.insert_queries()
+        (insert_severity, _, _, _, _) = self.insert_queries()
         del_cve_range = "DELETE from cve_range where CVE_number=?"
 
         for cve in severity_data:
@@ -500,7 +535,7 @@ def populate_severity(self, severity_data, cursor, data_source):
         cursor.executemany(del_cve_range, [(cve["ID"],) for cve in severity_data])
 
     def populate_affected(self, affected_data, cursor, data_source):
-        (_, insert_cve_range, _) = self.insert_queries()
+        (_, insert_cve_range, _, _, _) = self.insert_queries()
         try:
             cursor.executemany(
                 insert_cve_range,
@@ -522,6 +557,21 @@ def populate_affected(self, affected_data, cursor, data_source):
         except Exception as e:
             LOGGER.info(f"Unable to insert data for {data_source} - {e}")
 
+    def populate_metrics(self):
+        cursor = self.db_open_and_get_cursor()
+        # Insert a row without specifying cve_metrics_id
+        (_, _, _, _, insert_metrics) = self.insert_queries()
+        data = [
+            (1, "EPSS"),
+            (2, "CVSS-2"),
+            (3, "CVSS-3"),
+        ]
+        # Execute the insert query for each row
+        for row in data:
+            cursor.execute(insert_metrics, row)
+        self.connection.commit()
+        self.db_close()
+
     def clear_cached_data(self) -> None:
         self.create_cache_backup()
         if self.cachedir.exists():
@@ -712,12 +762,19 @@ def create_exploit_db(self):
         self.db_close()
 
     def populate_exploit_db(self, exploits):
-        (_, _, insert_exploit) = self.insert_queries()
+        (_, _, insert_exploit, _, _) = self.insert_queries()
         cursor = self.db_open_and_get_cursor()
         cursor.executemany(insert_exploit, exploits)
         self.connection.commit()
         self.db_close()
 
+    def store_epss_data(self):
+        (_, _, _, insert_cve_metrics, _) = self.insert_queries()
+        cursor = self.db_open_and_get_cursor()
+        cursor.executemany(insert_cve_metrics, self.epss_data)
+        self.connection.commit()
+        self.db_close()
+
     def dict_factory(self, cursor, row):
         d = {}
         for idx, col in enumerate(cursor.description):
@@ -866,7 +923,13 @@ def db_to_json(self, path, private_key, passphrase):
             shutil.rmtree(temp_gnupg_home)
 
     def json_to_db(self, cursor, db_column, json_data):
-        (insert_severity, insert_cve_range, insert_exploit) = self.insert_queries()
+        (
+            insert_severity,
+            insert_cve_range,
+            insert_exploit,
+            insert_cve_metrics,
+            insert_metrics,
+        ) = self.insert_queries()
         columns = []
         for data in json_data:
             column = list(data.keys())
@@ -887,6 +950,10 @@ def json_to_db(self, cursor, db_column, json_data):
             cursor.executemany(insert_cve_range, values)
         elif db_column == "cve_severity":
             cursor.executemany(insert_severity, values)
+        elif db_column == "cve_metrics":
+            cursor.executemany(insert_cve_metrics, values)
+        elif db_column == "metrics":
+            cursor.executemany(insert_metrics, values)
 
     def json_to_db_wrapper(self, path, pubkey, ignore_signature, log_signature_error):
         try:

test/test_cvedb.py

@@ -66,3 +66,28 @@ def test_import_export_json(self):
         self.cvedb.db_close()
 
         assert cve_entries_before == cve_entries_after
+
+    def test_new_database_schema(self):
+        # Check if the new schema is created in the database
+        self.cvedb.init_database()
+        cursor = self.cvedb.db_open_and_get_cursor()
+
+        tables_to_check = ["cve_metrics", "metrics"]
+        required_columns = {
+            "cve_metrics": ["cve_number", "metric_id", "metric_score", "metric_field"],
+            "metrics": ["metrics_id", "metrics_name"],
+        }
+
+        for table in tables_to_check:
+            cursor.execute(
+                f"SELECT name FROM sqlite_master WHERE type='table' AND name='{table}'"
+            )
+            result = cursor.fetchone()
+            assert result is not None  # Assert that the table exists
+
+            cursor.execute(f"PRAGMA table_info({table})")
+            columns = cursor.fetchall()
+            column_names = [column[1] for column in columns]
+            assert all(column in column_names for column in required_columns[table])
+
+        self.cvedb.db_close()