Slightly more strict CPE matching, less FP

2 files changed, 19 insertions, 5 deletions

diff --git a/sec_certs/dataset/dataset.py b/sec_certs/dataset/dataset.py
index ade17bb5..e02f6e8e 100644
--- a/sec_certs/dataset/dataset.py
+++ b/sec_certs/dataset/dataset.py
@@ -13,6 +13,7 @@ import requests
 import sec_certs.helpers as helpers
 import sec_certs.constants as constants
 import sec_certs.parallel_processing as cert_processing
+from sec_certs.sample.cpe import CPE
 
 from sec_certs.sample.certificate import Certificate
 from sec_certs.serialization.json import ComplexSerializableType
@@ -180,6 +181,16 @@ class Dataset(ABC, ComplexSerializableType):
             cert.compute_heuristics_version()
 
     def _compute_cpe_matches(self, download_fresh_cpes: bool = False) -> Tuple[CPEClassifier, CPEDataset]:
+        def filter_condition(cpe: CPE) -> bool:
+            """
+            Filters out very weak CPE matches that don't improve our database.
+            """
+            if cpe.title and (cpe.version == '-' or cpe.version == '*') and not any(char.isdigit() for char in cpe.title):
+                return False
+            elif not cpe.title and (cpe.version == '-' or cpe.version == '*') and not any(char.isdigit() for char in cpe.item_name):
+                return False
+            return True
+
         logger.info('Computing heuristics: Finding CPE matches for certificates')
         cpe_dset = self._prepare_cpe_dataset(download_fresh_cpes)
         if not cpe_dset.was_enhanced_with_vuln_cpes:
@@ -187,7 +198,7 @@ class Dataset(ABC, ComplexSerializableType):
             cpe_dset.enhance_with_cpes_from_cve_dataset(cve_dset)
 
         clf = CPEClassifier(config.cpe_matching_threshold, config.cpe_n_max_matches)
-        clf.fit([x for x in cpe_dset])
+        clf.fit([x for x in cpe_dset if filter_condition(x)])
 
         for cert in tqdm.tqdm(self, desc='Predicting CPE matches with the classifier'):
             cert.compute_heuristics_cpe_match(clf)
diff --git a/sec_certs/model/cpe_matching.py b/sec_certs/model/cpe_matching.py
index fb7c211e..bbf61ad7 100644
--- a/sec_certs/model/cpe_matching.py
+++ b/sec_certs/model/cpe_matching.py
@@ -98,9 +98,6 @@ class CPEClassifier(BaseEstimator):
         @param relax_title: bool
         @return:
         """
-        if product_name == 'VMware® ESXi 4.0 Update 1 and vCenter Server 4.0 Update 1':
-            print('geee')
-
         sanitized_vendor = CPEClassifier._discard_trademark_symbols(vendor).lower() if vendor else vendor
         sanitized_product_name = CPEClassifier._fully_sanitize_string(product_name) if product_name else product_name
         candidate_vendors = self.get_candidate_list_of_vendors(sanitized_vendor)
@@ -223,11 +220,17 @@ class CPEClassifier(BaseEstimator):
         """
 
         def is_cpe_version_among_cert_versions(cpe_version: Optional[str], cert_versions: List[str]) -> bool:
+            def simple_startswith(seeked_version: str, checked_string: str) -> bool:
+                if seeked_version == checked_string:
+                    return True
+                else:
+                    return checked_string.startswith(seeked_version) and not checked_string[len(seeked_version)].isdigit()
+
             if not cpe_version:
                 return False
             just_numbers = r'(\d{1,5})(\.\d{1,5})' # TODO: The use of this should be double-checked
             for v in cert_versions:
-                if (v.startswith(cpe_version) and re.search(just_numbers, cpe_version)) or cpe_version.startswith(v):
+                if (simple_startswith(v, cpe_version) and re.search(just_numbers, cpe_version)) or simple_startswith(cpe_version, v):
                     return True
             return False