Fix FIPSalgo serialization.

4 files changed, 39 insertions, 32 deletions

diff --git a/sec_certs/dataset/fips.py b/sec_certs/dataset/fips.py
index 1cb67480..56cacd2c 100644
--- a/sec_certs/dataset/fips.py
+++ b/sec_certs/dataset/fips.py
@@ -356,8 +356,8 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
 
     @staticmethod
     def _match_with_algorithm(processed_cert: FIPSCertificate, cert_candidate_id: str) -> bool:
-        for alg_type, cert_id in processed_cert.heuristics.algorithms:
-            curr_id = "".join(filter(str.isdigit, cert_id))
+        for algo in processed_cert.heuristics.algorithms:
+            curr_id = "".join(filter(str.isdigit, algo.cert_id))
             if curr_id == cert_candidate_id:
                 return False
         return True
@@ -376,13 +376,10 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
         if self.algorithms is None:
             raise RuntimeError("Dataset was probably not built correctly - this should not be happening.")
 
-        if cert_candidate_id not in self.algorithms.certs:
-            return True
-
         if not FIPSDataset._match_with_algorithm(processed_cert, cert_candidate_id):
             return False
 
-        algs = self.algorithms.certs_for_id(int(cert_candidate_id))
+        algs = self.algorithms.certs_for_id(cert_candidate_id)
         for current_alg in algs:
             if current_alg.vendor is None or processed_cert.web_data.vendor is None:
                 raise RuntimeError("Dataset was probably not built correctly - this should not be happening.")
diff --git a/sec_certs/dataset/fips_algorithm.py b/sec_certs/dataset/fips_algorithm.py
index e9e3680c..eb70d112 100644
--- a/sec_certs/dataset/fips_algorithm.py
+++ b/sec_certs/dataset/fips_algorithm.py
@@ -28,7 +28,7 @@ class FIPSAlgorithmDataset(Dataset, ComplexSerializableType):
         description: str = "dataset_description",
     ):
         super().__init__(certs, root_dir, name, description)
-        self._id_map: Dict[int, List[str]] = {}
+        self._id_map: Dict[str, List[str]] = {}
 
     def get_certs_from_web(self):
         self.root_dir.mkdir(exist_ok=True)
@@ -101,7 +101,7 @@ class FIPSAlgorithmDataset(Dataset, ComplexSerializableType):
                 )
 
                 alg_type, alg_id = split_alg(validation)
-                fips_alg = FIPSAlgorithm(alg_id, vendor, product, alg_type, date)
+                fips_alg = FIPSAlgorithm(alg_id, alg_type, vendor, product, date)
                 self.certs[fips_alg.dgst] = fips_alg
         # And now rebuild the id map
         self._build_id_map()
@@ -135,7 +135,7 @@ class FIPSAlgorithmDataset(Dataset, ComplexSerializableType):
     def __setitem__(self, key: str, value: FIPSAlgorithm):
         self.certs.__setitem__(key, value)
 
-    def certs_for_id(self, cert_id: int) -> List[FIPSAlgorithm]:
+    def certs_for_id(self, cert_id: str) -> List[FIPSAlgorithm]:
         if cert_id in self._id_map:
             return [self.certs[x] for x in self._id_map[cert_id]]
         else:
diff --git a/sec_certs/sample/fips.py b/sec_certs/sample/fips.py
index 723bc332..3f660bdb 100644
--- a/sec_certs/sample/fips.py
+++ b/sec_certs/sample/fips.py
@@ -23,6 +23,7 @@ from sec_certs.sample.certificate import Certificate
 from sec_certs.sample.certificate import Heuristics as BaseHeuristics
 from sec_certs.sample.certificate import References, logger
 from sec_certs.sample.cpe import CPE
+from sec_certs.sample.fips_algorithm import FIPSAlgorithm
 from sec_certs.serialization.json import ComplexSerializableType
 from sec_certs.utils.helpers import fips_dgst
 
@@ -156,7 +157,7 @@ class _FIPSHTMLParser:
         return ids_found
 
     @staticmethod
-    def parse_table(element: Union[Tag, NavigableString]) -> Set[Tuple[Optional[str], str]]:
+    def parse_table(element: Union[Tag, NavigableString]) -> Set[FIPSAlgorithm]:
         """
         Parses content of <table> tags in FIPS .html CMVP page
 
@@ -171,13 +172,13 @@ class _FIPSHTMLParser:
             cert_ids = _FIPSHTMLParser.extract_algorithm_certificates(tds[1].text)
             name = tds[0].text
             for cert_id in cert_ids:
-                found_items.add((name, cert_id))
+                found_items.add(FIPSAlgorithm(cert_id, name))
 
         return found_items
 
     @staticmethod
-    def parse_description(current_text: str) -> Set[Tuple[Optional[str], str]]:
-        return set(map(lambda x: (None, x), _FIPSHTMLParser.extract_algorithm_certificates(current_text)))
+    def parse_description(current_text: str) -> Set[FIPSAlgorithm]:
+        return set(map(FIPSAlgorithm, _FIPSHTMLParser.extract_algorithm_certificates(current_text)))
 
     @staticmethod
     def extract_algorithm_certificates(current_text: str) -> Set[str]:
@@ -256,7 +257,7 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         exceptions: Optional[List[str]]
         module_type: Optional[str]
         embodiment: Optional[str]
-        algorithms: Optional[Set[Tuple[Optional[str], str]]]
+        algorithms: Optional[Set[FIPSAlgorithm]]
         tested_conf: Optional[List[str]]
         description: Optional[str]
         mentioned_certs: Optional[Dict[str, Dict[str, int]]]
@@ -294,7 +295,7 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
 
         cert_id: int
         keywords: Dict
-        algorithms: Set[Tuple[Optional[str], str]]
+        algorithms: Set[FIPSAlgorithm]
         clean_cert_ids: Dict[str, int]
         st_metadata: Optional[Dict[str, Any]] = field(default=None)
 
@@ -311,7 +312,7 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         """
 
         keywords: Dict[str, Dict]
-        algorithms: Set[Tuple[Optional[str], str]]
+        algorithms: Set[FIPSAlgorithm]
         unmatched_algs: int
         clean_cert_ids: Dict[str, int]
 
@@ -575,12 +576,14 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         return keywords, cert
 
     @staticmethod
-    def analyze_tables(tup: Tuple[FIPSCertificate, bool]) -> Tuple[bool, FIPSCertificate, Set]:  # noqa: C901
+    def analyze_tables(  # noqa: C901
+        tup: Tuple[FIPSCertificate, bool]
+    ) -> Tuple[bool, FIPSCertificate, Set[FIPSAlgorithm]]:
         """
         Searches for tables in pdf documents of the instance.
 
-        :param Tuple[FIPSCertificate, bool] tup: certificate object, whether to use high precision results or approx. results
-        :return Tuple[bool, FIPSCertificate, Set]: True on success / False otherwise, modified cert object, List of processed tables.
+        :param tup: certificate object, whether to use high precision results or approx. results
+        :return: True on success / False otherwise, modified cert object, List of processed tables.
         """
 
         def extract_algorithm_certificates(current_text):
@@ -588,7 +591,7 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
             reg = r"(?:#?\s?|(?:Cert)\.?[^. ]*?\s?)(?:[CcAa]\s)?(?P<id>[CcAa]? ?\d+)"
             for m in re.finditer(reg, current_text):
                 set_items.add(m.group())
-            return set(map(lambda x: (None, x), set_items))
+            return set(map(FIPSAlgorithm, set_items))
 
         cert, precision = tup
         if (not precision and cert.state.tables_done) or (
@@ -638,9 +641,15 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
             ):
                 to_pop.add(cert)
 
-        for alg_type, cert_no in self.heuristics.algorithms:
-            if int("".join(filter(str.isdigit, cert_no))) == int("".join(filter(str.isdigit, cert))):
-                to_pop.add(cert)
+        this_id = int("".join(filter(str.isdigit, cert)))
+
+        for algo in self.heuristics.algorithms:
+            try:
+                algo_id = int("".join(filter(str.isdigit, algo.cert_id)))
+                if algo_id == this_id:
+                    to_pop.add(cert)
+            except ValueError:
+                continue
 
     def clean_cert_ids(self) -> None:
         """
@@ -656,8 +665,8 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         if self.web_data.algorithms is None:
             raise RuntimeError(f"Algorithms were not found for cert {self.cert_id} - this should not be happening.")
 
-        for alg_type, cert_no in self.web_data.algorithms:
-            alg_set.update(cert_no)
+        for algo in self.web_data.algorithms:
+            alg_set.add(algo.cert_id)
 
         for cert_rule in fips_rules["fips_cert_id"]["Cert"]:
             to_pop = set()
diff --git a/sec_certs/sample/fips_algorithm.py b/sec_certs/sample/fips_algorithm.py
index fcfd42b6..ea74cd33 100644
--- a/sec_certs/sample/fips_algorithm.py
+++ b/sec_certs/sample/fips_algorithm.py
@@ -1,20 +1,21 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Optional
 
 from sec_certs import constants
 from sec_certs.serialization.json import ComplexSerializableType
 
 
-@dataclass(eq=True)
+@dataclass(eq=True, order=True, frozen=True)
 class FIPSAlgorithm(ComplexSerializableType):
     """
     Data structure for algorithm of `FIPSCertificate`
     """
 
-    cert_id: int
-    vendor: str
-    implementation: str
-    algorithm_type: str
-    date: str
+    cert_id: str
+    algorithm_type: Optional[str] = field(default=None)
+    vendor: Optional[str] = field(default=None)
+    implementation: Optional[str] = field(default=None)
+    date: Optional[str] = field(default=None)
 
     @property
     def dgst(self) -> str: