Skip to content

Commit e21e691

Browse files
fix: changed metric ids in cvedb to constants and fixed test (#4473)
1 parent 0405d52 commit e21e691

File tree

3 files changed

+14
-22
lines changed

3 files changed

+14
-22
lines changed

cve_bin_tool/cvedb.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,10 @@
4545
DBNAME = "cve.db"
4646
OLD_CACHE_DIR = Path("~") / ".cache" / "cvedb"
4747

48+
EPSS_METRIC_ID = 1
49+
CVSS_2_METRIC_ID = 2
50+
CVSS_3_METRIC_ID = 3
51+
4852

4953
class CVEDB:
5054
"""
@@ -615,9 +619,9 @@ def populate_metrics(self):
615619
# Insert a row without specifying cve_metrics_id
616620
insert_metrics = self.INSERT_QUERIES["insert_metrics"]
617621
data = [
618-
(1, "EPSS"),
619-
(2, "CVSS-2"),
620-
(3, "CVSS-3"),
622+
(EPSS_METRIC_ID, "EPSS"),
623+
(CVSS_2_METRIC_ID, "CVSS-2"),
624+
(CVSS_3_METRIC_ID, "CVSS-3"),
621625
]
622626
# Execute the insert query for each row
623627
for row in data:

cve_bin_tool/data_sources/epss_source.py

Lines changed: 6 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,9 @@ def __init__(self, error_mode=ErrorMode.TruncTrace):
3636
self.backup_cachedir = self.BACKUPCACHEDIR
3737
self.epss_path = str(Path(self.cachedir) / "epss")
3838
self.file_name = os.path.join(self.epss_path, "epss_scores-current.csv")
39-
self.epss_metric_id = None
4039
self.source_name = self.SOURCE
4140

42-
async def update_epss(self, cursor):
41+
async def update_epss(self):
4342
"""
4443
Updates the EPSS data by downloading and parsing the CSV file.
4544
Returns:
@@ -51,7 +50,6 @@ async def update_epss(self, cursor):
5150
"""
5251
self.LOGGER.debug("Fetching EPSS data...")
5352

54-
self.EPSS_id_finder(cursor)
5553
await self.download_epss_data()
5654
self.epss_data = self.parse_epss_data()
5755
return self.epss_data
@@ -110,15 +108,6 @@ async def download_epss_data(self):
110108
except aiohttp.ClientError as e:
111109
self.LOGGER.error(f"An error occurred during downloading epss {e}")
112110

113-
def EPSS_id_finder(self, cursor):
114-
"""Search for metric id in EPSS table"""
115-
query = """
116-
SELECT metrics_id FROM metrics
117-
WHERE metrics_name = "EPSS"
118-
"""
119-
cursor.execute(query)
120-
self.epss_metric_id = cursor.fetchall()[0][0]
121-
122111
def parse_epss_data(self, file_path=None):
123112
"""Parse epss data from the file path given and return the parse data"""
124113
parsed_data = []
@@ -138,9 +127,11 @@ def parse_epss_data(self, file_path=None):
138127
# Parse the data from the remaining rows
139128
for row in reader:
140129
cve_id, epss_score, epss_percentile = row[:3]
141-
parsed_data.append(
142-
(cve_id, self.epss_metric_id, epss_score, epss_percentile)
143-
)
130+
131+
# prevent circular dependency
132+
from cve_bin_tool.cvedb import EPSS_METRIC_ID
133+
134+
parsed_data.append((cve_id, EPSS_METRIC_ID, epss_score, epss_percentile))
144135
return parsed_data
145136

146137
async def get_cve_data(self):

test/test_source_epss.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,15 +23,12 @@ def setup_class(cls):
2323
]
2424

2525
def test_parse_epss(self):
26-
# EPSS need metrics table to populated in the database. To get the EPSS metric id from table.
26+
# EPSS need metrics table to populated in the database. EPSS metric id is a constant.
2727
cvedb = CVEDB()
2828
# creating table
2929
cvedb.init_database()
3030
# populating metrics
3131
cvedb.populate_metrics()
32-
cursor = cvedb.db_open_and_get_cursor()
33-
# seting EPSS_metric_id
34-
self.epss.EPSS_id_finder(cursor)
3532
# parsing the data
3633
self.epss_data = self.epss.parse_epss_data(self.epss.file_name)
3734
cvedb.db_close()

0 commit comments

Comments
 (0)