relaxing version if no good cpe match found

4 files changed, 23 insertions, 13 deletions

diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index 20e73585..5f5526e5 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -1265,7 +1265,12 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
         self.heuristics.cpe_candidate_vendors = cpe_dataset.get_candidate_list_of_vendors(self.manufacturer)
 
     def get_heuristics_cpe_match(self, cpe_dataset: CPEDataset):
-        self.heuristics.cpe_matches = cpe_dataset.get_cpe_matches(self.name, self.heuristics.cpe_candidate_vendors, self.heuristics.extracted_versions)
+        self.get_heuristics_cpe_vendors(cpe_dataset)
+        self.heuristics.cpe_matches = cpe_dataset.get_cpe_matches(self.name,
+                                                                  self.heuristics.cpe_candidate_vendors,
+                                                                  self.heuristics.extracted_versions,
+                                                                  n_max_matches=constants.CPE_MAX_MATCHES,
+                                                                  threshold=constants.CPE_MATCHING_THRESHOLD)
 
     def get_heuristics_related_cves(self, cve_dataset: CVEDataset):
         if self.heuristics.verified_cpe_matches:
diff --git a/sec_certs/constants.py b/sec_certs/constants.py
index 41a4f3b7..610e4cc9 100644
--- a/sec_certs/constants.py
+++ b/sec_certs/constants.py
@@ -6,6 +6,9 @@ RETURNCODE_OK = 'ok'
 RETURNCODE_NOK = 'nok'
 REQUEST_TIMEOUT = 10
 
+CPE_MATCHING_THRESHOLD = 70
+CPE_MAX_MATCHES = 10
+
 MIN_CORRECT_CERT_SIZE = 5000
 
 LOGS_FILENAME = './cert_processing_log.txt'
diff --git a/sec_certs/cpe.py b/sec_certs/cpe.py
index f2b1fd7b..94378fe8 100644
--- a/sec_certs/cpe.py
+++ b/sec_certs/cpe.py
@@ -207,11 +207,11 @@ class CPEDataset:
         candidate_vendor_version_pairs = self.get_candidate_vendor_version_pairs(cert_candidate_cpe_vendors, cert_candidate_versions)
 
         if not candidate_vendor_version_pairs:
-            return None
+            return []
 
         return list(itertools.chain.from_iterable([self.vendor_version_to_cpe[x] for x in candidate_vendor_version_pairs]))
 
-    def get_cpe_matches(self, cert_name: str, cert_candidate_cpe_vendors: List[str], cert_candidate_versions: List, relax_version: bool = False, n_max_matches=5, threshold: int = 60) -> Optional[List[Tuple[float, CPE]]]:
+    def get_cpe_matches(self, cert_name: str, cert_candidate_cpe_vendors: List[str], cert_candidate_versions: List[str], relax_version: bool = False, n_max_matches=10, threshold: int = 60) -> Optional[List[Tuple[float, CPE]]]:
         replace_non_letter_non_numbers_with_space = re.compile(r"(?ui)\W")
 
         def sanitize_matched_string(string: str):
@@ -219,9 +219,6 @@ class CPEDataset:
             return replace_non_letter_non_numbers_with_space.sub(' ', string)
         candidates = self.get_candidate_cpe_items(cert_candidate_cpe_vendors, cert_candidate_versions)
 
-        if not candidates:
-            return None
-
         sanitized_cert_name = sanitize_matched_string(cert_name)
         reasonable_matches = []
         for c in candidates:
@@ -233,5 +230,13 @@ class CPEDataset:
 
         if reasonable_matches:
             reasonable_matches = sorted(reasonable_matches, key=lambda x: x[0], reverse=True)
+
+            # possibly filter short titles to avoid false positives
+            # reasonable_matches = list(filter(lambda x: len(x[1].item_name) > 4, reasonable_matches))
+
             return reasonable_matches[:n_max_matches]
+
+        if not reasonable_matches and not relax_version:
+            return self.get_cpe_matches(cert_name, cert_candidate_cpe_vendors, ['-'], relax_version=True, n_max_matches=n_max_matches, threshold=threshold)
+
         return None
\ No newline at end of file
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 60738b92..07756d81 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -61,7 +61,7 @@ class Dataset(ABC):
     def __iter__(self):
         yield from self.certs.values()
 
-    def __getitem__(self, item: str) -> 'Certificate':
+    def __getitem__(self, item: str):
         return self.certs.__getitem__(item.lower())
 
     def __setitem__(self, key: str, value: 'Certificate'):
@@ -163,6 +163,9 @@ class CCDataset(Dataset, ComplexSerializableType):
             state = self.DatasetInternalState()
         self.state = state
 
+    def __iter__(self) -> CommonCriteriaCert:
+        yield from self.certs.values()
+
     def to_dict(self):
         return {**{'state': self.state}, **super().to_dict()}
 
@@ -710,11 +713,6 @@ class CCDataset(Dataset, ComplexSerializableType):
             for cert in self:
                 cert.get_heuristics_version()
 
-        def compute_candidate_cpe_vendors(cpe_dataset: CPEDataset):
-            logger.info('Computing heuristics: Possible vendors from CPE repository that could match the certificate vendor')
-            for cert in self:
-                cert.get_heuristics_cpe_vendors(cpe_dataset)
-
         def compute_cpe_matches(cpe_dataset: CPEDataset):
             logger.info('Computing heuristics: Finding CPE matches for certificates')
             for cert in self:
@@ -722,7 +720,6 @@ class CCDataset(Dataset, ComplexSerializableType):
 
         compute_candidate_versions()
         cpe_dset = self.prepare_cpe_dataset(download_fresh_cpes)
-        compute_candidate_cpe_vendors(cpe_dset)
         compute_cpe_matches(cpe_dset)
 
         if update_json is True: