chore: Hopefully resolved merge conflict

5 files changed, 71 insertions, 18 deletions

diff --git a/.gitignore b/.gitignore
index 1f8db2ed..f6054f09 100644
--- a/.gitignore
+++ b/.gitignore
@@ -123,3 +123,6 @@ fips_test_dataset/
 # log
 cert_processing_log.txt
 ./cc_processing_log.txt
+
+# Example script may dump classification report into the folder
+examples/classification_report.json
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 e1039099..bbf61ad7 100644
--- a/sec_certs/model/cpe_matching.py
+++ b/sec_certs/model/cpe_matching.py
@@ -75,7 +75,13 @@ class CPEClassifier(BaseEstimator):
         """
         return [self.predict_single_cert(x[0], x[1], x[2]) for x in tqdm.tqdm(X, desc='Predicting')]
 
-    def predict_single_cert(self, vendor: str, product_name: str, versions: Optional[List[str]], relax_version: bool = True) -> Optional[List[str]]:
+    def predict_single_cert(self,
+                            vendor: str,
+                            product_name: str,
+                            versions: Optional[List[str]],
+                            relax_version: bool = False,
+                            relax_title: bool = False
+                            ) -> Optional[List[str]]:
         """
         Predict List of CPE uris for triplet (vendor, product_name, list_of_version). The prediction is made as follows:
         1. Sanitize all strings
@@ -84,10 +90,12 @@ class CPEClassifier(BaseEstimator):
         4. Compute string similarity of the candidate CPE matches and certificate name
         5. Evaluate best string similarity, if above threshold, declare it a match.
         6. If no CPE item is matched, we tried again but relax version and check CPEs that don't have their version specified.
+        Also, we search for 100% CPE matches on item name instead of title.
         @param vendor: manufacturer of the certificate
         @param product_name: name of the certificate
         @param versions: List of versions that appear in the certificate name
         @param relax_version: bool, see step 6 above.
+        @param relax_title: bool
         @return:
         """
         sanitized_vendor = CPEClassifier._discard_trademark_symbols(vendor).lower() if vendor else vendor
@@ -95,16 +103,22 @@ class CPEClassifier(BaseEstimator):
         candidate_vendors = self.get_candidate_list_of_vendors(sanitized_vendor)
 
         candidates = self.get_candidate_cpe_matches(candidate_vendors, versions)
-        ratings = [self.compute_best_match(cpe, sanitized_product_name, candidate_vendors, versions) for cpe in candidates]
+        ratings = [self.compute_best_match(cpe, sanitized_product_name, candidate_vendors, versions, relax_title=relax_title) for cpe in candidates]
         threshold = self.match_threshold if not relax_version else 100
         final_matches = list(filter(lambda x: x[0] >= threshold, zip(ratings, candidates)))
+        final_matches = [x[1].uri for x in final_matches[:self.n_max_matches]]
 
-        if relax_version and not final_matches:
-            return self.predict_single_cert(vendor, product_name, ['-'], relax_version=False)
 
-        return [x[1].uri for x in final_matches[:self.n_max_matches]] if final_matches else None
 
-    def compute_best_match(self, cpe: CPE, product_name: str, candidate_vendors: str, versions: Optional[List[str]]) -> float:
+        if not relax_title and not final_matches:
+            final_matches = self.predict_single_cert(vendor, product_name, versions, relax_version=relax_version, relax_title=True)
+
+        if not relax_version and not final_matches:
+            final_matches = self.predict_single_cert(vendor, product_name, ['-'], relax_version=True, relax_title=relax_title)
+
+        return final_matches if final_matches else None
+
+    def compute_best_match(self, cpe: CPE, product_name: str, candidate_vendors: List[str], versions: Optional[List[str]], relax_title: bool = False) -> float:
         """
         Tries several different settings in which string similarity between CPE and certificate name is tested.
         For definition of string similarity, see rapidfuzz package on GitHub. Both token set ratio and partial ratio are tested,
@@ -115,15 +129,27 @@ class CPEClassifier(BaseEstimator):
         @param versions: versions that appear in the certificate
         @return: Maximal value of the four string similarities discussed above.
         """
-        sanitized_title = CPEClassifier._fully_sanitize_string(cpe.title) if cpe.title else CPEClassifier._fully_sanitize_string(cpe.vendor + ' ' + cpe.item_name + ' ' + cpe.version)
+        if relax_title:
+            sanitized_title = CPEClassifier._fully_sanitize_string(cpe.title) if cpe.title else CPEClassifier._fully_sanitize_string(cpe.vendor + ' ' + cpe.item_name + ' ' + cpe.version + ' ' + cpe.update + ' ' + cpe.target_hw)
+        else:
+            if cpe.title:
+                sanitized_title = CPEClassifier._fully_sanitize_string(cpe.title)
+            else:
+                return 0
+
         sanitized_item_name = CPEClassifier._fully_sanitize_string(cpe.item_name)
         cert_stripped = CPEClassifier._strip_manufacturer_and_version(product_name, candidate_vendors, versions)
 
         token_set_ratio_on_title = fuzz.token_set_ratio(product_name, sanitized_title)
-        token_set_ratio_on_item_name = fuzz.token_set_ratio(cert_stripped, sanitized_item_name)
         partial_ratio_on_title = fuzz.partial_ratio(product_name, sanitized_title)
-        partial_ratio_on_item_name = fuzz.partial_ratio(cert_stripped, sanitized_item_name)
-        return max([token_set_ratio_on_title, partial_ratio_on_title, token_set_ratio_on_item_name, partial_ratio_on_item_name])
+        ratings = [token_set_ratio_on_title, partial_ratio_on_title]
+
+        if relax_title:
+            token_set_ratio_on_item_name = fuzz.token_set_ratio(cert_stripped, sanitized_item_name)
+            partial_ratio_on_item_name = fuzz.partial_ratio(cert_stripped, sanitized_item_name)
+            ratings += [token_set_ratio_on_item_name, partial_ratio_on_item_name]
+
+        return max(ratings)
 
     @staticmethod
     def _fully_sanitize_string(string: str) -> str:
@@ -194,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
 
diff --git a/sec_certs/sample/cpe.py b/sec_certs/sample/cpe.py
index 61da8450..d70f2783 100644
--- a/sec_certs/sample/cpe.py
+++ b/sec_certs/sample/cpe.py
@@ -38,6 +38,14 @@ class CPE(PandasSerializableType, ComplexSerializableType):
         return ['uri', 'title', 'start_version', 'end_version']
 
     @property
+    def update(self) -> str:
+        return ' '.join(self.uri.split(':')[6].split('_'))
+
+    @property
+    def target_hw(self) -> str:
+        return ' '.join(self.uri.split(':')[11].split('_'))
+
+    @property
     def pandas_tuple(self):
         return self.uri, self.vendor, self.item_name, self.version, self.title
 
diff --git a/tests/test_fips_oop.py b/tests/test_fips_oop.py
index a1c39950..67a93c72 100644
--- a/tests/test_fips_oop.py
+++ b/tests/test_fips_oop.py
@@ -51,12 +51,11 @@ class TestFipsOOP(TestCase):
     def setUpClass(cls) -> None:
         config.load(cls.data_dir.parent / 'settings_test.yaml')
 
-    # FIXME - uncomment later - bug in this test - FileNotFound - bad parsing of "microsoft" string
-    # def test_size(self):
-    #     for certs in self.certs_to_parse:
-    #         with TemporaryDirectory() as tmp_dir:
-    #             dataset = _set_up_dataset(tmp_dir, certs)
-    #             self.assertEqual(len(dataset.certs), len(certs), "Wrong number of parsed certs")
+    def test_size(self):
+        for certs in self.certs_to_parse.values():
+            with TemporaryDirectory() as tmp_dir:
+                dataset = _set_up_dataset(tmp_dir, certs)
+                self.assertEqual(len(dataset.certs), len(certs), "Wrong number of parsed certs")
 
     def test_connections_microsoft(self):
         certs = self.certs_to_parse['microsoft']