remove duplicate method in CCDataset

3 files changed, 9 insertions, 27 deletions

diff --git a/sec_certs/dataset/common_criteria.py b/sec_certs/dataset/common_criteria.py
index 460ad787..a826276c 100644
--- a/sec_certs/dataset/common_criteria.py
+++ b/sec_certs/dataset/common_criteria.py
@@ -619,27 +619,6 @@ class CCDataset(Dataset, ComplexSerializableType):
 
         self.state.certs_analyzed = True
 
-    @serialize
-    def compute_related_cves(self, download_fresh_cves: bool = False, download_nist_cpe_matching_dict: bool = True):
-        logger.info('Retrieving related CVEs to verified CPE matches')
-        cve_dset = self._prepare_cve_dataset(download_fresh_cves, download_nist_cpe_matching_dict)
-
-        verified_cpe_rich_certs = [x for x in self if x.heuristics.cpe_matches]
-        if not verified_cpe_rich_certs:
-            logger.error(
-                'No certificates with verified CPE match detected. You must run dset.manually_verify_cpe_matches() first. Returning.')
-            return
-
-        relevant_cpes = set(itertools.chain.from_iterable([x.heuristics.cpe_matches for x in verified_cpe_rich_certs]))
-        cve_dset.filter_related_cpes(relevant_cpes)
-
-        for cert in tqdm(verified_cpe_rich_certs, desc='Computing related CVES'):
-            cert.compute_heuristics_related_cves(cve_dset)
-
-        n_vulnerable = len([x for x in verified_cpe_rich_certs if x.heuristics.related_cves])
-        n_vulnerabilities = sum([len(x.heuristics.related_cves) for x in verified_cpe_rich_certs if x.heuristics.related_cves])
-        logger.info(f'In total, we identified {n_vulnerabilities} vulnerabilities in {n_vulnerable} vulnerable certificates.')
-
     def get_certs_from_name(self, cert_name: str) -> List[CommonCriteriaCert]:
         return [crt for crt in self if crt.name == cert_name]
 
diff --git a/sec_certs/dataset/dataset.py b/sec_certs/dataset/dataset.py
index 02473dea..e2972ecc 100644
--- a/sec_certs/dataset/dataset.py
+++ b/sec_certs/dataset/dataset.py
@@ -261,13 +261,17 @@ class Dataset(ABC, ComplexSerializableType):
                 cert.heuristics.cpe_matches = set(cert.heuristics.cpe_matches).union(set(cert.heuristics.verified_cpe_matches))
 
     @serialize
-    def compute_related_cves(self, download_fresh_cves: bool = False):
+    def compute_related_cves(self, download_fresh_cves: bool = False, download_nist_cpe_matching_dict: bool = True):
         logger.info('Retrieving related CVEs to verified CPE matches')
-        cve_dset = self._prepare_cve_dataset(download_fresh_cves)
+        cve_dset = self._prepare_cve_dataset(download_fresh_cves, download_nist_cpe_matching_dict)
 
         self.enrich_automated_cpes_with_manual_labels()
         cpe_rich_certs = [x for x in self if x.heuristics.cpe_matches]
 
+        if not cpe_rich_certs:
+            logger.error('No certificates with verified CPE match detected. You must run dset.manually_verify_cpe_matches() first. Returning.')
+            return
+
         relevant_cpes = set(itertools.chain.from_iterable([x.heuristics.cpe_matches for x in cpe_rich_certs]))
         cve_dset.filter_related_cpes(relevant_cpes)
 
@@ -277,5 +281,4 @@ class Dataset(ABC, ComplexSerializableType):
         n_vulnerable = len([x for x in cpe_rich_certs if x.heuristics.related_cves])
         n_vulnerabilities = sum(
             [len(x.heuristics.related_cves) for x in cpe_rich_certs if x.heuristics.related_cves])
-        logger.info(
-            f'In total, we identified {n_vulnerabilities} vulnerabilities in {n_vulnerable} vulnerable certificates.')
+        logger.info(f'In total, we identified {n_vulnerabilities} vulnerabilities in {n_vulnerable} vulnerable certificates.')
diff --git a/sec_certs/dataset/fips.py b/sec_certs/dataset/fips.py
index 873da0eb..4f63f7d6 100644
--- a/sec_certs/dataset/fips.py
+++ b/sec_certs/dataset/fips.py
@@ -412,14 +412,14 @@ class FIPSDataset(Dataset, ComplexSerializableType):
             FIPSDataset._find_connections(current_cert)
 
     @serialize
-    def finalize_results(self):
+    def finalize_results(self, download_nist_cpe_matching_dict: bool = True):
         logger.info("Entering 'analysis' and building connections between certificates.")
         self.unify_algorithms()
         self.remove_algorithms_from_extracted_data()
         self.validate_results()
 
         self.compute_cpe_heuristics()
-        self.compute_related_cves()
+        self.compute_related_cves(download_nist_cpe_matching_dict=download_nist_cpe_matching_dict)
 
     def _highlight_vendor_in_dot(self, dot: Digraph, current_key: str, highlighted_vendor: str):
         if self.certs[current_key].web_scan.vendor != highlighted_vendor: