some more cve/cpe tests

4 files changed, 33 insertions, 16 deletions

diff --git a/src/sec_certs/dataset/cve.py b/src/sec_certs/dataset/cve.py
index 79016991..7754caf1 100644
--- a/src/sec_certs/dataset/cve.py
+++ b/src/sec_certs/dataset/cve.py
@@ -138,7 +138,8 @@ class CVEDataset(JSONPathDataset, ComplexSerializableType):
                 continue
 
             for cpe_uri in vulnerable_cpe_uris:
-                self._cpe_uri_to_cve_ids_lookup.setdefault(cpe_uri, set()).add(cve.cve_id)
+                if not cpe_uris_of_interest or cpe_uri in cpe_uris_of_interest:
+                    self._cpe_uri_to_cve_ids_lookup.setdefault(cpe_uri, set()).add(cve.cve_id)
 
     def _get_cves_from_exactly_matched_cpes(self, cpe_uris: set[str]) -> set[str]:
         return set(
diff --git a/src/sec_certs/sample/cpe.py b/src/sec_certs/sample/cpe.py
index 2cc267c9..5e9e8b65 100644
--- a/src/sec_certs/sample/cpe.py
+++ b/src/sec_certs/sample/cpe.py
@@ -61,7 +61,7 @@ class CPEMatchCriteriaConfiguration(ComplexSerializableType):
     components: list[list[CPEMatchCriteria]]
     _expanded_components: list[list[str]] = field(default_factory=list)
 
-    def matches(self, cpe_ids: set[str]) -> bool:
+    def matches(self, cpe_uris: set[str]) -> bool:
         """
         Returns if given set of cpe_ids matches this configuration.
         """
@@ -69,7 +69,7 @@ class CPEMatchCriteriaConfiguration(ComplexSerializableType):
             raise ValueError(
                 "Cannot match to CPEMatchConfiguration when attribute _expanded_components was not filled-in. That attribute is prepared by `CVEDataset.build_lookup_dict()`."
             )
-        return all(any(x in component for x in cpe_ids) for component in self._expanded_components)
+        return all(any(x in component for x in cpe_uris) for component in self._expanded_components)
 
     @property
     def serialized_attributes(self) -> list[str]:
diff --git a/tests/test_cpe.py b/tests/test_cpe.py
index 22cf69ea..a5096678 100644
--- a/tests/test_cpe.py
+++ b/tests/test_cpe.py
@@ -73,7 +73,3 @@ def test_serialization_missing_path():
 
 def test_enhance_with_nvd_data():
     pass
-
-
-def test_criteria_configuration_expansion():
-    pass
diff --git a/tests/test_cve.py b/tests/test_cve.py
index 35636b82..52bc1b60 100644
--- a/tests/test_cve.py
+++ b/tests/test_cve.py
@@ -6,7 +6,7 @@ from pathlib import Path
 import pytest
 
 from sec_certs.dataset import CVEDataset
-from sec_certs.sample import CVE
+from sec_certs.sample import CPE, CVE
 from sec_certs.serialization.json import SerializationError
 
 
@@ -68,11 +68,31 @@ def test_serialization_missing_path():
 #     pass
 
 
-# def test_dataset_prunning(cve_dataset_path: Path, cpe_match_feed: dict):
-#     cve_dataset = CVEDataset.from_json(cve_dataset_path)
-#     cpes_to_consider = set()
-#     cve_dataset.build_lookup_dict(
-#         cpe_match_feed,
-#     )
-#     # TODO: Note that also CVEs with configurations are being prunned, their prunning must also be tested.
-#     pass
+def test_dataset_prunning(cve_dataset_path: Path, cpe_match_feed: dict):
+    cve_dataset = CVEDataset.from_json(cve_dataset_path)
+    cpes_to_consider = {
+        CPE("D2D8310E-E0B8-4FF1-86EA-24F463E9F175", "cpe:2.3:o:freebsd:freebsd:4.2:*:*:*:*:*:*:*"),
+        CPE("ACFEB8AA-B8FC-49E1-98BE-3D581655FE2E", "cpe:2.3:a:ibm:websphere_application_server:7.0.0.6:*:*:*:*:*:*:*"),
+        CPE("88CA7428-3329-4E07-84CB-428DB1D2BC8E", "cpe:2.3:o:ibm:zos:6.0.1:*:*:*:*:*:*:*"),
+    }
+    cve_dataset.build_lookup_dict(cpe_match_feed, cpes_to_consider)
+
+    assert cve_dataset._cves_with_vulnerable_configurations == [cve_dataset["CVE-2010-2325"]]
+    assert cve_dataset._cpe_uri_to_cve_ids_lookup == {"cpe:2.3:o:freebsd:freebsd:4.2:*:*:*:*:*:*:*": {"CVE-2003-0001"}}
+
+
+def test_criteria_configuration_expansion(cve_dataset_path: Path, cpe_match_feed: dict):
+    cve_dataset = CVEDataset.from_json(cve_dataset_path)
+    cve_dataset.cves = {"CVE-2003-0070": cve_dataset["CVE-2003-0070"]}
+    cve_dataset.build_lookup_dict(cpe_match_feed)
+    assert len(cve_dataset["CVE-2003-0070"].vulnerable_criteria_configurations) == 1
+    expanded_components = cve_dataset["CVE-2003-0070"].vulnerable_criteria_configurations[0]._expanded_components
+    assert len(expanded_components) == 2
+    first, second = expanded_components[0], expanded_components[1]
+    assert len(first) == 10
+    assert "cpe:2.3:a:nalin_dahyabhai:vte:0.16.14:*:*:*:*:*:*:*" in first
+    assert "cpe:2.3:a:nalin_dahyabhai:vte:0.25.1:*:*:*:*:*:*:*" in first
+    assert second == [
+        "cpe:2.3:a:gnome:gnome-terminal:2.0:*:*:*:*:*:*:*",
+        "cpe:2.3:a:gnome:gnome-terminal:2.2:*:*:*:*:*:*:*",
+    ]