fix start_version, end_version problems

3 files changed, 38 insertions, 14 deletions

diff --git a/sec_certs/dataset/cpe.py b/sec_certs/dataset/cpe.py
index 5e217159..43ae3365 100644
--- a/sec_certs/dataset/cpe.py
+++ b/sec_certs/dataset/cpe.py
@@ -1,3 +1,4 @@
+import copy
 import itertools
 import logging
 import tempfile
@@ -9,7 +10,7 @@ from typing import Any, ClassVar, Dict, Iterator, List, Set, Tuple, Union, cast
 
 import pandas as pd
 
-import sec_certs.helpers as helpers
+from sec_certs import constants, helpers
 from sec_certs.dataset.cve import CVEDataset
 from sec_certs.sample.cpe import CPE, cached_cpe
 from sec_certs.serialization.json import ComplexSerializableType, serialize
@@ -142,6 +143,24 @@ class CPEDataset(ComplexSerializableType):
 
     @serialize
     def enhance_with_cpes_from_cve_dataset(self, cve_dset: Union[CVEDataset, str, Path]) -> None:
+        def adding_condition(
+            considered_cpe: CPE,
+            vndr_item_lookup: Set[Tuple[str, str]],
+            vndr_item_version_lookup: Set[Tuple[str, str, str]],
+        ) -> bool:
+            if (
+                considered_cpe.version == constants.CPE_VERSION_NA
+                and (considered_cpe.vendor, considered_cpe.item_name) not in vndr_item_lookup
+            ):
+                return True
+            elif (
+                considered_cpe.version != constants.CPE_VERSION_NA
+                and (considered_cpe.vendor, considered_cpe.item_name, considered_cpe.version)
+                not in vndr_item_version_lookup
+            ):
+                return True
+            return False
+
         if isinstance(cve_dset, (str, Path)):
             cve_dset = CVEDataset.from_json(cve_dset)
 
@@ -151,10 +170,15 @@ class CPEDataset(ComplexSerializableType):
 
         old_len = len(self.cpes)
 
+        # We only enrich if tuple (vendor, item_name) is not already in the dataset
+        vendor_item_lookup = {(cpe.vendor, cpe.item_name) for cpe in self}
+        vendor_item_version_lookup = {(cpe.vendor, cpe.item_name, cpe.version) for cpe in self}
         for cpe in helpers.tqdm(all_cpes_in_cve_dset, desc="Enriching CPE dataset with new CPEs"):
-            if cpe not in self:
-                self[cpe.uri] = cpe
-
+            if adding_condition(cpe, vendor_item_lookup, vendor_item_version_lookup):
+                new_cpe = copy.deepcopy(cpe)
+                new_cpe.start_version = None
+                new_cpe.end_version = None
+                self[new_cpe.uri] = new_cpe
         self.build_lookup_dicts()
 
         logger.info(f"Enriched the CPE dataset with {len(self.cpes) - old_len} new CPE records.")
diff --git a/sec_certs/dataset/cve.py b/sec_certs/dataset/cve.py
index aff500d8..a18364d7 100644
--- a/sec_certs/dataset/cve.py
+++ b/sec_certs/dataset/cve.py
@@ -26,7 +26,7 @@ logger = logging.getLogger(__name__)
 @dataclass
 class CVEDataset(ComplexSerializableType):
     cves: Dict[str, CVE]
-    cpe_to_cve_ids_lookup: Dict[str, List[str]] = field(init=False)
+    cpe_to_cve_ids_lookup: Dict[str, Set[str]] = field(init=False)
     cve_url: Final[str] = "https://nvd.nist.gov/feeds/json/cve/1.1/nvdcve-1.1-"
     cpe_match_feed_url: Final[str] = "https://nvd.nist.gov/feeds/json/cpematch/1.0/nvdcpematch-1.0.json.zip"
 
@@ -51,6 +51,7 @@ class CVEDataset(ComplexSerializableType):
 
     def build_lookup_dict(self, use_nist_mapping: bool = True, nist_matching_filepath: Optional[Path] = None):
         """
+        Builds look-up dictionary CPE -> Set[CVE]
         Developer's note: There are 3 CPEs that are present in the cpe matching feed, but are badly processed by CVE
         feed, in which case they won't be found as a key in the dictionary. We intentionally ignore those. Feel free
         to add corner cases and manual fixes. According to our investigation, the suffereing CPEs are:
@@ -77,9 +78,9 @@ class CVEDataset(ComplexSerializableType):
                 vulnerable_configurations = cve.vulnerable_cpes
             for cpe in vulnerable_configurations:
                 if cpe.uri not in self.cpe_to_cve_ids_lookup:
-                    self.cpe_to_cve_ids_lookup[cpe.uri] = [cve.cve_id]
+                    self.cpe_to_cve_ids_lookup[cpe.uri] = {cve.cve_id}
                 else:
-                    self.cpe_to_cve_ids_lookup[cpe.uri].append(cve.cve_id)
+                    self.cpe_to_cve_ids_lookup[cpe.uri].add(cve.cve_id)
 
     @classmethod
     def download_cves(cls, output_path_str: str, start_year: int, end_year: int):
@@ -145,7 +146,7 @@ class CVEDataset(ComplexSerializableType):
             dset = json.load(handle, cls=CustomJSONDecoder)
         return dset
 
-    def get_cve_ids_for_cpe_uri(self, cpe_uri: str) -> Optional[List[str]]:
+    def get_cve_ids_for_cpe_uri(self, cpe_uri: str) -> Optional[Set[str]]:
         return self.cpe_to_cve_ids_lookup.get(cpe_uri, None)
 
     def filter_related_cpes(self, relevant_cpes: Set[CPE]):
@@ -175,6 +176,10 @@ class CVEDataset(ComplexSerializableType):
         return df.set_index("cve_id")
 
     def get_nist_cpe_matching_dict(self, input_filepath: Optional[Path]) -> Dict[CPE, List[CPE]]:
+        """
+        Computes dictionary that maps complex CPEs to list of simple CPEs.
+        """
+
         def parse_key_cpe(field: Dict) -> CPE:
             start_version = None
             if "versionStartIncluding" in field:
diff --git a/sec_certs/sample/cpe.py b/sec_certs/sample/cpe.py
index ae14b4a8..bbbf1c39 100644
--- a/sec_certs/sample/cpe.py
+++ b/sec_certs/sample/cpe.py
@@ -91,12 +91,7 @@ class CPE(PandasSerializableType, ComplexSerializableType):
         return hash((self.uri, self.start_version, self.end_version))
 
     def __eq__(self, other: object) -> bool:
-        return (
-            isinstance(other, self.__class__)
-            and self.uri == other.uri
-            and self.start_version == other.start_version
-            and self.end_version == other.end_version
-        )
+        return isinstance(other, self.__class__) and self.uri == other.uri
 
 
 @lru_cache(maxsize=4096)