diff options
| -rw-r--r-- | examples/cc_cpe_labeling.py | 4 | ||||
| -rw-r--r-- | sec_certs/dataset/cpe.py | 8 |
2 files changed, 9 insertions, 3 deletions
diff --git a/examples/cc_cpe_labeling.py b/examples/cc_cpe_labeling.py index 89527bbe..752b8052 100644 --- a/examples/cc_cpe_labeling.py +++ b/examples/cc_cpe_labeling.py @@ -23,7 +23,7 @@ def main(): dset.get_certs_from_web(to_download=True) # Automatically match CPEs and CVEs - dset.compute_cpe_heuristics() + _, cpe_dset, _ = dset.compute_cpe_heuristics() dset.compute_related_cves() # Load dataset of ground truth CPE labels @@ -37,7 +37,7 @@ def main(): # Evaluate CPE matching performance metrics (on validation set) and dump classification report into json y_valid = [(x.heuristics.verified_cpe_matches) for x in validation_certs] - evaluate(validation_certs, y_valid, "./my_debug_dataset/classification_report.json") + evaluate(validation_certs, y_valid, "./my_debug_dataset/classification_report.json", cpe_dset) logger.info(f"{dset.json_path} should now contain fully labeled dataset.") diff --git a/sec_certs/dataset/cpe.py b/sec_certs/dataset/cpe.py index b9c04cad..8e552039 100644 --- a/sec_certs/dataset/cpe.py +++ b/sec_certs/dataset/cpe.py @@ -47,13 +47,19 @@ class CPEDataset(ComplexSerializableType): def __contains__(self, item: CPE) -> bool: if not isinstance(item, CPE): raise ValueError(f"{item} is not of CPE class") - return item.uri in self.cpes.keys() + return item.uri in self.cpes.keys() and self.cpes[item.uri] == item + + def __eq__(self, other: object) -> bool: + return isinstance(other, CPEDataset) and self.cpes == other.cpes @property def serialized_attributes(self) -> List[str]: return ["was_enhanced_with_vuln_cpes", "json_path", "cpes"] def __post_init__(self) -> None: + self.build_lookup_dicts() + + def build_lookup_dicts(self) -> None: """ Will build look-up dictionaries that are used for fast matching """ |