improve matching -> more strict

2 files changed, 47 insertions, 10 deletions

diff --git a/sec_certs/model/cpe_matching.py b/sec_certs/model/cpe_matching.py
index e1039099..fb7c211e 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,27 +90,38 @@ 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:
         """
+        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)
 
         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 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 relax_version and not final_matches:
-            return self.predict_single_cert(vendor, product_name, ['-'], relax_version=False)
+        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 [x[1].uri for x in final_matches[:self.n_max_matches]] if final_matches else None
+        return final_matches if final_matches else None
 
-    def compute_best_match(self, cpe: CPE, product_name: str, candidate_vendors: str, versions: Optional[List[str]]) -> float:
+    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 +132,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:
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