FIPS: Add clean_cert_ids and use them.

Adds two new data fields to the FIPSCertificate subobjects:
 pdf_data.clean_cert_ids and heuristics.clean_cert_ids.

The pdf_data ones are cleaned cert_id rule matches from the
security target. While the heuristics ones were also cleaned
using the algorithm dataset by the FipsDataset class.

2 files changed, 29 insertions, 46 deletions

diff --git a/sec_certs/dataset/fips.py b/sec_certs/dataset/fips.py
index 942dda93..49c6bef8 100644
--- a/sec_certs/dataset/fips.py
+++ b/sec_certs/dataset/fips.py
@@ -297,19 +297,21 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
             certificate.pdf_data.algorithms += algorithms
         return not_decoded
 
-    def _compute_heuristics_keywords(self) -> None:
+    def _compute_heuristics_clean_ids(self) -> None:
         for cert in self.certs.values():
-            cert.compute_heuristics_keywords()
+            self._clean_cert_ids(cert)
 
     def _extract_metadata(self):
         certs_to_process = [x for x in self]
-        cert_processing.process_parallel(
+        res = cert_processing.process_parallel(
             FIPSCertificate.extract_sp_metadata,
             certs_to_process,
             config.n_threads,
             use_threading=False,
             progress_bar_desc="Extracting security policy metadata",
         )
+        for r in res:
+            self.certs[r.dgst] = r
 
     def _unify_algorithms(self) -> None:
         for certificate in self.certs.values():
@@ -353,17 +355,16 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
             or cert_first.year < conn_first.year
         )
 
-    def _remove_false_positives_for_cert(self, current_cert: FIPSCertificate) -> None:
-        if current_cert.heuristics.keywords is None:
-            raise RuntimeError("Dataset was probably not built correctly - this should not be happening.")
-        for rule in current_cert.heuristics.keywords["fips_cert_id"]:
-            matches = current_cert.heuristics.keywords["fips_cert_id"][rule]
-            current_cert.heuristics.keywords["fips_cert_id"][rule] = [
-                cert_id
-                for cert_id in matches
-                if self._validate_id(current_cert, cert_id.replace("Cert.", "").replace("cert.", "").lstrip("#CA0 "))
-                and cert_id != current_cert.cert_id
-            ]
+    def _clean_cert_ids(self, current_cert: FIPSCertificate) -> None:
+        current_cert.clean_cert_ids()
+        if not current_cert.state.txt_state:
+            return
+        current_cert.heuristics.clean_cert_ids = {
+            cert_id: count
+            for cert_id, count in current_cert.pdf_data.clean_cert_ids.items()
+            if self._validate_id(current_cert, cert_id.replace("Cert.", "").replace("cert.", "").lstrip("#CA0 "))
+            and cert_id != current_cert.cert_id
+        }
 
     @staticmethod
     def _match_with_algorithm(processed_cert: FIPSCertificate, cert_candidate_id: str) -> bool:
@@ -412,7 +413,7 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
                     lambda x: x,
                     map(
                         lambda cid: "".join(filter(str.isdigit, cid)),
-                        cert.heuristics.keywords["fips_cert_id"]["Cert"],
+                        cert.heuristics.clean_cert_ids,
                     ),
                 )
             )
@@ -422,13 +423,6 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
                 filter(lambda x: x, map(lambda cid: "".join(filter(str.isdigit, cid)), cert.web_data.mentioned_certs))
             )
 
-        current_cert: FIPSCertificate
-
-        for current_cert in self.certs.values():
-            if not current_cert.state.txt_state:
-                continue
-            self._remove_false_positives_for_cert(current_cert)
-
         finder = DependencyFinder()
         finder.fit(self.certs, lambda cert: cert.cert_id, pdf_lookup)  # type: ignore
 
@@ -452,7 +446,7 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
         logger.info("Entering 'analysis' and building connections between certificates.")
         self._extract_metadata()
         self._unify_algorithms()
-        self._compute_heuristics_keywords()
+        self._compute_heuristics_clean_ids()
         self._compute_dependencies()
         if perform_cpe_heuristics:
             _, _, cve_dset = self.compute_cpe_heuristics()
diff --git a/sec_certs/sample/fips.py b/sec_certs/sample/fips.py
index 3fc230ef..e520624b 100644
--- a/sec_certs/sample/fips.py
+++ b/sec_certs/sample/fips.py
@@ -127,8 +127,8 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         cert_id: int
         keywords: Dict
         algorithms: List
+        clean_cert_ids: Dict[str, int]
         st_metadata: Optional[Dict[str, Any]] = field(default=None)
-        # TODO: Add metadata processing.
 
         def __repr__(self) -> str:
             return str(self.cert_id)
@@ -145,6 +145,7 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         keywords: Dict[str, Dict]
         algorithms: List[Dict[str, Dict]]
         unmatched_algs: int
+        clean_cert_ids: Dict[str, int]
 
         extracted_versions: Optional[Set[str]] = field(default=None)
         cpe_matches: Optional[Set[str]] = field(default=None)
@@ -155,10 +156,6 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         web_references: References = field(default_factory=References)
 
         @property
-        def serialized_attributes(self) -> List[str]:
-            return copy.deepcopy(super().serialized_attributes)
-
-        @property
         def dgst(self) -> str:
             return helpers.get_first_16_bytes_sha256(str(self.keywords))
 
@@ -315,8 +312,9 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
                 items_found["cert_id"],
                 {} if not initialized else initialized.pdf_data.keywords,
                 [] if not initialized else initialized.pdf_data.algorithms,
+                {} if not initialized else initialized.pdf_data.clean_cert_ids,
             ),
-            FIPSCertificate.Heuristics(dict(), [], 0),
+            FIPSCertificate.Heuristics(dict(), [], 0, {}),
             state,
         )
 
@@ -638,7 +636,7 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         return result
 
     def _process_to_pop(self, reg_to_match: Pattern, cert: str, to_pop: Set[str]) -> None:
-        for found in self.heuristics.keywords["fips_certlike"]["Certlike"]:
+        for found in self.pdf_data.keywords["fips_certlike"]["Certlike"]:
             match_in_found = reg_to_match.search(found)
             match_in_cert = reg_to_match.search(cert)
             if (
@@ -653,39 +651,30 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
                 if int("".join(filter(str.isdigit, cert_no))) == int("".join(filter(str.isdigit, cert))):
                     to_pop.add(cert)
 
-    def compute_heuristics_keywords(self) -> None:
+    def clean_cert_ids(self) -> None:
         """
-        Compute the heuristics keywords.
-
-         - Removes algorithm mentions from the cert_id rule matches.
+        Removes algorithm mentions from the cert_id rule matches and stores them into clean_cert_id matches.
         """
         self.state.file_status = True
         if not self.pdf_data.keywords:
             return
 
-        self.heuristics.keywords = copy.deepcopy(self.pdf_data.keywords)
-
-        # XXX: Do not do this, the heuristics keywords need to be from the Security target only,
-        #      because the "st_references" are computed based on these, and the "web_references"
-        #      are computed based on "web_data.mentioned_certs".
-        # Add the mentioned certs
-        # if self.web_data.mentioned_certs:
-        #     for cert_id, value in self.web_data.mentioned_certs.items():
-        #         self.heuristics.keywords["fips_cert_id"]["Cert"]["#" + str(cert_id)] = value["count"]
+        matches = copy.deepcopy(self.pdf_data.keywords["fips_cert_id"]["Cert"])
 
         alg_set = self._create_alg_set()
         for cert_rule in fips_rules["fips_cert_id"]["Cert"]:
             to_pop = set()
-            for cert in self.heuristics.keywords["fips_cert_id"]["Cert"]:
+            for cert in matches:
                 if cert in alg_set:
                     to_pop.add(cert)
                     continue
                 self._process_to_pop(cert_rule, cert, to_pop)
 
             for r in to_pop:
-                self.heuristics.keywords["fips_cert_id"]["Cert"].pop(r, None)
+                matches.pop(r, None)
 
-        self.heuristics.keywords["fips_cert_id"]["Cert"].pop("#" + str(self.cert_id), None)
+        matches.pop("#" + str(self.cert_id), None)
+        self.pdf_data.clean_cert_ids = matches
 
     @staticmethod
     def get_compare(vendor: str) -> str: