Merge pull request #154 from kannkyo/feature-add-ge-and-le

add ge and le function

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.
 

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-')