FIPS: Extract ST metadata and store it.

3 files changed, 30 insertions, 24 deletions

diff --git a/sec_certs/dataset/fips.py b/sec_certs/dataset/fips.py
index f829da9d..942dda93 100644
--- a/sec_certs/dataset/fips.py
+++ b/sec_certs/dataset/fips.py
@@ -232,6 +232,7 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
     @serialize
     def get_certs_from_web(
         self,
+        # TODO: REMOVE THIS TEST ARGUMENT, OMG!
         test: Optional[Path] = None,
         update: bool = False,
         redo_web_scan=False,
@@ -300,6 +301,16 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
         for cert in self.certs.values():
             cert.compute_heuristics_keywords()
 
+    def _extract_metadata(self):
+        certs_to_process = [x for x in self]
+        cert_processing.process_parallel(
+            FIPSCertificate.extract_sp_metadata,
+            certs_to_process,
+            config.n_threads,
+            use_threading=False,
+            progress_bar_desc="Extracting security policy metadata",
+        )
+
     def _unify_algorithms(self) -> None:
         for certificate in self.certs.values():
             new_algorithms: List[Dict] = []
@@ -394,11 +405,7 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
                 return False
         return True
 
-    def _validate_results(self) -> None:
-        """
-        Function that validates results and finds the final connection output
-        """
-
+    def _compute_dependencies(self) -> None:
         def pdf_lookup(cert):
             return set(
                 filter(
@@ -443,9 +450,10 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
         :param bool perform_cpe_heuristics: If CPE heuristics shall be computed, defaults to True
         """
         logger.info("Entering 'analysis' and building connections between certificates.")
+        self._extract_metadata()
         self._unify_algorithms()
         self._compute_heuristics_keywords()
-        self._validate_results()
+        self._compute_dependencies()
         if perform_cpe_heuristics:
             _, _, cve_dset = self.compute_cpe_heuristics()
             self.compute_related_cves(use_nist_cpe_matching_dict=use_nist_cpe_matching_dict, cve_dset=cve_dset)
@@ -580,9 +588,9 @@ class FIPSDataset(Dataset[FIPSCertificate], ComplexSerializableType):
     def plot_graphs(self, show: bool = False) -> None:
         """
         Plots FIPS graphs.
-        # TODO: Currently broken, see https://github.com/crocs-muni/sec-certs/issues/211
-        :param bool show: If plots should be showed with .show() method, defaults to False
+
+        :param bool show: If plots should be shown with .show() method, defaults to False
         """
         self._create_dot_graph("full_graph", show=show)
         self._create_dot_graph("web_only_graph", "web", show=show)
-        self._create_dot_graph("pdf_only_graph", "pdf", show=show)
+        self._create_dot_graph("st_only_graph", "st", show=show)
diff --git a/sec_certs/sample/fips.py b/sec_certs/sample/fips.py
index 33c25742..3fc230ef 100644
--- a/sec_certs/sample/fips.py
+++ b/sec_certs/sample/fips.py
@@ -106,16 +106,6 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         sw_versions: Optional[str]
         product_url: Optional[str]
 
-        @property
-        def dgst(self) -> str:
-            return helpers.get_first_16_bytes_sha256(
-                self.product_url
-                if self.product_url is not None
-                else "" + self.vendor_www
-                if self.vendor_www is not None
-                else ""
-            )
-
         def __repr__(self) -> str:
             return (
                 self.module_name
@@ -137,12 +127,9 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         cert_id: int
         keywords: Dict
         algorithms: List
+        st_metadata: Optional[Dict[str, Any]] = field(default=None)
         # TODO: Add metadata processing.
 
-        @property
-        def dgst(self) -> str:
-            return helpers.get_first_16_bytes_sha256(str(self.keywords))
-
         def __repr__(self) -> str:
             return str(self.cert_id)
 
@@ -352,11 +339,17 @@ class FIPSCertificate(Certificate["FIPSCertificate", "FIPSCertificate.Heuristics
         """
         Downloads security policy file from web. Staticmethod to allow for parametrization.
         """
-        exit_code = helpers.download_file(*cert, delay=1)
+        exit_code = helpers.download_file(*cert, delay=constants.FIPS_DOWNLOAD_DELAY)
         if exit_code != requests.codes.ok:
             logger.error(f"Failed to download security policy from {cert[0]}, code: {exit_code}")
 
     @staticmethod
+    def extract_sp_metadata(cert: FIPSCertificate) -> FIPSCertificate:
+        """Extract the PDF metadata from the security policy. Staticmethod to allow for parametrization."""
+        response, cert.pdf_data.st_metadata = sec_certs.utils.pdf.extract_pdf_metadata(cert.state.sp_path)
+        return cert
+
+    @staticmethod
     def _initialize_dictionary() -> Dict[str, Any]:
         return {
             "module_name": None,
diff --git a/tests/test_fips_oop.py b/tests/test_fips_oop.py
index 32e34c79..e07eb572 100644
--- a/tests/test_fips_oop.py
+++ b/tests/test_fips_oop.py
@@ -116,6 +116,11 @@ class TestFipsOOP(TestCase):
                 dataset = _set_up_dataset(tmp_dir, certs)
                 self.assertEqual(len(dataset.certs), len(certs), "Wrong number of parsed certs")
 
+    def test_metadata_extraction(self):
+        with TemporaryDirectory() as tmp_dir:
+            dst = _set_up_dataset_for_full(tmp_dir, ["3493"], self.cpe_dset_path, self.cve_dset_path)
+            self.assertIsNotNone(dst.certs[fips_dgst("3493")].pdf_data.st_metadata)
+
     def test_connections_microsoft(self):
         certs = self.certs_to_parse["microsoft"]
         with TemporaryDirectory() as tmp_dir: