diff --git a/src/epss_api/epss.py b/src/epss_api/epss.py index 84346f1..c2b75be 100644 --- a/src/epss_api/epss.py +++ b/src/epss_api/epss.py @@ -10,7 +10,7 @@ class Score(object): """ EPSS Score Object""" def __init__(self, cve: str, epss: str, percentile: str): - """Initialize EPSS Score Object + """Initialize EPSS Score Object Args: cve (str): CVE-yyyy-n @@ -136,6 +136,34 @@ def percentile_gt(self, min: float) -> list[Score]: return list(self._sortedScoresByPercentile[i:]) + def epss_ge(self, min: float) -> list[Score]: + """Get CVEs with EPSS score greater than or equal to the parameter + + Args: + min (float): limit of EPSS score + + Returns: + list[Score] | None: EPSS score object list + """ + i = bisect_left(self._sortedScoresByEpss, min, + key=lambda x: x.epss) + + return list(self._sortedScoresByEpss[i:]) + + def percentile_ge(self, min: float) -> list[Score]: + """Get CVEs with percentile greater than or equal to the parameter + + Args: + min (float): limit of percentile + + Returns: + list[Score] | None: EPSS score object list + """ + i = bisect_left(self._sortedScoresByPercentile, min, + key=lambda x: x.percentile) + + return list(self._sortedScoresByPercentile[i:]) + def epss_lt(self, max: float) -> list[Score]: """Get CVEs with EPSS score lower than the parameter @@ -151,7 +179,7 @@ def epss_lt(self, max: float) -> list[Score]: return list(self._sortedScoresByEpss[:i]) def percentile_lt(self, max: float) -> list[Score]: - """Get CVEs with percentile lower than the parameter + """Get CVEs with percentile less than the parameter Args: max (float): limit of percentile @@ -164,6 +192,34 @@ def percentile_lt(self, max: float) -> list[Score]: return list(self._sortedScoresByPercentile[:i]) + def epss_le(self, max: float) -> list[Score]: + """Get CVEs with EPSS score less than or equal to the parameter + + Args: + max (float): limit of EPSS score + + Returns: + list[Score] | []: EPSS score object list + """ + i = bisect_right(self._sortedScoresByEpss, max, + key=lambda x: x.epss) + + return list(self._sortedScoresByEpss[:i]) + + def percentile_le(self, max: float) -> list[Score]: + """Get CVEs with percentile less than or equal to the parameter + + Args: + max (float): limit of percentile + + Returns: + list[Score] | []: EPSS score object list + """ + i = bisect_right(self._sortedScoresByPercentile, max, + key=lambda x: x.percentile) + + return list(self._sortedScoresByPercentile[:i]) + def csv(self) -> list[str]: """Get csv data containing all epss scores. diff --git a/tests/epss_api/test_epss.py b/tests/epss_api/test_epss.py index f15c29d..abfde47 100644 --- a/tests/epss_api/test_epss.py +++ b/tests/epss_api/test_epss.py @@ -42,6 +42,24 @@ def test_percentile_gt(min): assert s.percentile <= min +@pytest.mark.parametrize("min", [-1, 0, 0.5, 1, 2]) +def test_epss_ge(min): + scores = epss.epss_ge(min) + for s in scores: + assert s.epss >= min + for s in list(set(scores) - set(epss.scores())): + assert s.epss < min + + +@pytest.mark.parametrize("min", [-1, 0, 0.5, 1, 2]) +def test_percentile_ge(min): + scores = epss.percentile_ge(min) + for s in scores: + assert s.percentile >= min + for s in list(set(scores) - set(epss.scores())): + assert s.percentile < min + + @pytest.mark.parametrize("max", [-1, 0, 0.5, 1, 2]) def test_epss_lt(max): scores = epss.epss_lt(max) @@ -60,6 +78,24 @@ def test_percentile_lt(max): assert s.percentile >= max +@pytest.mark.parametrize("max", [-1, 0, 0.5, 1, 2]) +def test_epss_le(max): + scores = epss.epss_le(max) + for s in scores: + assert s.epss <= max + for s in list(set(scores) - set(epss.scores())): + assert s.epss > max + + +@pytest.mark.parametrize("max", [-1, 0, 0.5, 1, 2]) +def test_percentile_le(max): + scores = epss.percentile_le(max) + for s in scores: + assert s.percentile <= max + for s in list(set(scores) - set(epss.scores())): + assert s.percentile > max + + def test_score(): score = epss.score(cve_id='CVE-2022-0669') assert score.cve.startswith('CVE-')