add deterministic cpe-cve matching

3 files changed, 35 insertions, 6 deletions

diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index f4f184ef..79804ef7 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -23,6 +23,7 @@ from sec_certs.extract_certificates import load_cert_file, normalize_match_strin
 from sec_certs.cert_rules import fips_rules, fips_common_rules
 from sec_certs.configuration import config
 from sec_certs.cpe import CPE, CPEDataset
+from sec_certs.cve import CVE, CVEDataset
 
 logger = logging.getLogger(__name__)
 
@@ -873,16 +874,21 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
     class Heuristics(ComplexSerializableType):
         extracted_versions: List[str]
         cpe_candidate_vendors: List[str]
-        cpe_matches: List[str]
+        cpe_matches: Optional[List[Tuple[float, CPE]]]
+        verified_cpe_match: Optional[CPE]
+        related_cves: Optional[List[CVE]]
 
-        def __init__(self, extracted_versions: Optional[List[str]] = None, cpe_matches: Optional[List[str]] = None, cpe_candidate_vendors: Optional[List[str]] = None):
+        def __init__(self, extracted_versions: Optional[List[str]] = None, cpe_matches: Optional[List[str]] = None, cpe_candidate_vendors: Optional[List[str]] = None,
+                     verified_cpe_match: Optional[str] = None,
+                     related_cves: Optional[List[CVE]] = None):
             self.extracted_versions = extracted_versions
             self.cpe_matches = cpe_matches
             self.cpe_candidate_vendors = cpe_candidate_vendors
-
+            self.verified_cpe_match = verified_cpe_match
+            self.related_cves = related_cves
 
         def to_dict(self):
-            return {'extracted_versions': self.extracted_versions, 'cpe_matches': self.cpe_matches}
+            return {'extracted_versions': self.extracted_versions, 'cpe_matches': self.cpe_matches, 'verified_cpe_match': self.verified_cpe_match, 'related_cves': self.related_cves}
 
         @classmethod
         def from_dict(cls, dct: Dict[str, str]):
@@ -1251,3 +1257,13 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
 
     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)
+        if self.heuristics.cpe_matches:
+            self.heuristics.verified_cpe_match = self.heuristics.cpe_matches[0][1]
+        else:
+            self.heuristics.verified_cpe_match = None
+
+    def get_heuristics_related_cves(self, cve_dataset: CVEDataset):
+        if self.heuristics.verified_cpe_match:
+            self.heuristics.related_cves = cve_dataset.get_cves_for_cpe(self.heuristics.verified_cpe_match.uri)
+        else:
+            self.heuristics.related_cves = None
\ No newline at end of file
diff --git a/sec_certs/cve.py b/sec_certs/cve.py
index a9a5891f..53a4a997 100644
--- a/sec_certs/cve.py
+++ b/sec_certs/cve.py
@@ -183,4 +183,4 @@ class CVEDataset(ComplexSerializableType):
     def get_cves_for_cpe(self, cpe_uri: str) -> Optional[List[str]]:
         if not isinstance(cpe_uri, str):
             return None
-        return self.cpes_to_cve_lookup[cpe_uri] if self.cpes_to_cve_lookup[cpe_uri] else None
+        return self.cpes_to_cve_lookup.get(cpe_uri, None)
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 858ac891..e35e8cd9 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -28,6 +28,7 @@ from sec_certs.certificate import CommonCriteriaCert, Certificate, FIPSCertifica
 from sec_certs.serialization import ComplexSerializableType, CustomJSONDecoder, CustomJSONEncoder
 from sec_certs.configuration import config
 from sec_certs.cpe import CPEDataset
+from sec_certs.cve import CVEDataset
 
 logger = logging.getLogger(__name__)
 
@@ -654,7 +655,7 @@ class CCDataset(Dataset, ComplexSerializableType):
         if update_json is True:
             self.to_json(self.json_path)
 
-    def compute_heuristics(self, cpe_xml_path: Optional[str] = None, update_json=True):
+    def compute_heuristics(self, cpe_xml_path: Optional[str] = None, cve_dataset_path: Optional[str] = None, update_json=True):
         def compute_candidate_versions():
             logger.info('Computing heuristics: possible product versions in certificate name')
             for cert in self:
@@ -670,6 +671,11 @@ class CCDataset(Dataset, ComplexSerializableType):
             for cert in self:
                 cert.get_heuristics_cpe_match(cpe_dataset)
 
+        def compute_related_cves(cve_dataset: CVEDataset):
+            logger.info('Retreiving CVEs for verified CPE match.')
+            for cert in self:
+                cert.get_heuristics_related_cves(cve_dataset)
+
         compute_candidate_versions()
 
         if not cpe_xml_path:
@@ -679,6 +685,13 @@ class CCDataset(Dataset, ComplexSerializableType):
         compute_candidate_cpe_vendors(cpe_dset)
         compute_cpe_matches(cpe_dset)
 
+        if not cve_dataset_path:
+            cve_dataset = CVEDataset.from_web()
+        else:
+            cve_dataset = CVEDataset.from_json(cve_dataset_path)
+
+        compute_related_cves(cve_dataset)
+
         if update_json is True:
             self.to_json(self.json_path)