algorithms can be serialized (finally)

4 files changed, 125 insertions, 33 deletions

diff --git a/fips_oop_demo.py b/fips_oop_demo.py
index ed471a3c..58305dc6 100644
--- a/fips_oop_demo.py
+++ b/fips_oop_demo.py
@@ -24,6 +24,8 @@ def main():
 
     logging.info("Extracting keywords now.")
 
+    dset.convert_all_pdfs()
+
     dset.extract_keywords()
 
     logging.info(f'Finished extracting certificates for {len(dset.keywords)} items.')
@@ -47,7 +49,7 @@ def main():
     dset.finalize_results()
 
     logging.info('dump again')
-    dset.dump_to_json()
+    dset.to_json(dset.root_dir / 'fips_full_dataset.json')
 
     dset.get_dot_graph('different_new')
     end = datetime.now()
diff --git a/sec_certs/cert_processing.py b/sec_certs/cert_processing.py
index 739bfa9f..2e4fb3ee 100644
--- a/sec_certs/cert_processing.py
+++ b/sec_certs/cert_processing.py
@@ -11,7 +11,7 @@ def process_parallel(func: Callable, items: Iterable, max_workers: int, callback
     else:
         pool = Pool(max_workers)
 
-    results = [pool.apply_async(func, (i, ), callback=callback) for i in items]
+    results = [pool.apply_async(func, (*i, ), callback=callback) for i in items]
 
     if progress_bar is True:
         bar = tqdm(total=len(results))
diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index a4026726..d58f59f4 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -87,7 +87,8 @@ class FIPSCertificate(Certificate, ComplexSerializableType):
                  fw_version: Optional[str],
                  tables: bool,
                  file_status: Optional[bool],
-                 connections: List):
+                 connections: List,
+                 txt_state: bool=True):
         super().__init__()
         self.cert_id = cert_id
 
@@ -120,6 +121,7 @@ class FIPSCertificate(Certificate, ComplexSerializableType):
         self.tables_done = tables
         self.file_status = file_status
         self.connections = connections
+        self.txt_state = txt_state
 
     def __str__(self) -> str:
         return str(self.cert_id)
@@ -336,6 +338,15 @@ class FIPSCertificate(Certificate, ComplexSerializableType):
                                None,
                                [])
 
+    @staticmethod
+    def convert_pdf_file(cert: 'FipsCertificate', ds: 'FIPSDataset') -> 'FIPSCertificate':
+        if not cert.txt_state:
+            exit_code = helpers.convert_pdf_file(ds.policies_dir / f'{cert.cert_id}.pdf', ds.policies_dir / f'{cert.cert_id}.pdf.txt', ['-layout'])
+            if exit_code != constants.RETURNCODE_OK:
+                logger.error(f'Cert dgst: {cert.dgst} failed to convert security target pdf->txt')
+                cert.txt_state = False
+        return cert
+
 
 class CommonCriteriaCert(Certificate, ComplexSerializableType):
     cc_url = 'http://www.commoncriteriaportal.org'
@@ -651,7 +662,7 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
             cert.state.st_convert_ok = False
         return cert
 
-class FIPSAlgorithm(Certificate):
+class FIPSAlgorithm(Certificate, ComplexSerializableType):
     @property
     def dgst(self):
         # certs in dataset are in format { id: [FIPSAlgorithm] }, there is only one type of algorithm
@@ -674,5 +685,5 @@ class FIPSAlgorithm(Certificate):
 
     @classmethod
     def from_dict(cls, dct: dict) -> 'FIPSAlgorithm':
-        return FIPSAlgorithm(dct['cert_id'], dct['vendor'], dct['implementation'], dct['alg_type'],
-                             dct['validation_date'])
+        return FIPSAlgorithm(dct['cert_id'], dct['vendor'], dct['implementation'], dct['type'],
+                             dct['date'])
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 1a31dd4b..3f9a4f88 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -533,34 +533,59 @@ class FIPSDataset(Dataset, ComplexSerializableType):
         with open(self.root_dir / "fips_full_keywords.json", 'w') as f:
             f.write(json.dumps(self.keywords, indent=4, sort_keys=True))
 
-    # TODO figure out whether the name of this method shuold not be "get_certs", because we don't download every time
 
-    def get_certs_from_web(self):
-        def download_html_pages() -> Tuple[int, int]:
-            html_items = [
-                (f"https://csrc.nist.gov/projects/cryptographic-module-validation-program/certificate/{cert_id}",
-                 self.web_dir / f"{cert_id}.html") for cert_id in list(self.certs.keys()) if
-                not (self.web_dir / f'{cert_id}.html').exists()]
-            sp_items = [(
-                f"https://csrc.nist.gov/CSRC/media/projects/cryptographic-module-validation-program/documents/security-policies/140sp{cert_id}.pdf",
-                self.policies_dir / f"{cert_id}.pdf") for cert_id in list(self.certs.keys()) if
-                not (self.policies_dir / f'{cert_id}.pdf').exists()]
+    def download_all_pdfs(self):
+        sp_paths, sp_urls = [], []
+        self.policies_dir.mkdir(exist_ok=True)
 
-            logging.info(f"downloading {len(html_items) + len(sp_items)} module html and pdf files")
-            _, self.new_files = helpers.download_parallel(
-                html_items + sp_items, 8), len(html_items) + len(sp_items)
+        for cert_id in list(self.certs.keys()):
+            if not (self.policies_dir / f'{cert_id}.pdf').exists():
+                sp_urls.append(f"https://csrc.nist.gov/CSRC/media/projects/cryptographic-module-validation-program/documents/security-policies/140sp{cert_id}.pdf")
+                sp_paths.append(self.policies_dir / f"{cert_id}.pdf")
 
-            pages = [
-                (
-                    f'https://csrc.nist.gov/projects/cryptographic-algorithm-validation-program/validation-search?searchMode=validation&page={i}',
-                    self.algs_dir / f'page{i}.html'
-                ) for i in range(1, 502) if not (self.algs_dir / f'page{i}.html').exists()
-            ]
+        logging.info(f"downloading {len(sp_urls)} module pdf files")
+        Dataset._download_parallel(sp_urls, sp_paths)
+        self.new_files += len(sp_urls)
 
-            logging.info(f"downloading {len(pages)} algorithm html files")
-            helpers.download_parallel(pages, 8)
+    def download_all_htmls(self):
+        html_paths, html_urls = [], []
 
-            return len(html_items) + len(sp_items), len(pages)
+        self.web_dir.mkdir(exist_ok=True)
+        for cert_id in list(self.certs.keys()):
+            if not (self.web_dir / f'{cert_id}.html').exists():
+                html_urls.append(f"https://csrc.nist.gov/projects/cryptographic-module-validation-program/certificate/{cert_id}")
+                html_paths.append(self.web_dir / f"{cert_id}.html")
+
+        logging.info(f"downloading {len(html_urls)} module html files")
+        Dataset._download_parallel(html_urls, html_paths)
+        self.new_files += len(html_urls)
+
+    def download_all_algs(self):
+        algs_paths, algs_urls = [], []
+
+        self.algs_dir.mkdir(exist_ok=True)
+        for i in range(1, 502):
+            if not (self.algs_dir / f'page{i}.html').exists():
+                algs_urls.append(f'https://csrc.nist.gov/projects/cryptographic-algorithm-validation-program/validation-search?searchMode=validation&page={i}')
+                algs_paths.append(self.algs_dir / f"page{i}.html")
+
+        logging.info(f"downloading {len(algs_urls)} algs html files")
+        Dataset._download_parallel(algs_urls, algs_paths)
+        self.new_files += len(algs_urls)
+    
+
+    def convert_all_pdfs(self):
+        logger.info('Converting CC certificate reports to .txt')
+        for cert in self.certs.values():
+            FIPSCertificate.convert_pdf_file(cert, self)
+
+
+    # TODO figure out whether the name of this method shuold not be "get_certs", because we don't download every time
+    def get_certs_from_web(self):
+        def download_html_pages() -> Tuple[int, int]:
+            self.download_all_pdfs()
+            self.download_all_htmls()
+            self.download_all_algs()
 
         def get_certificates_from_html(html_file: Path) -> None:
             logger.info(f'Getting certificate ids from {html_file}')
@@ -595,7 +620,7 @@ class FIPSDataset(Dataset, ComplexSerializableType):
             get_certificates_from_html(self.web_dir / f)
 
         logger.info('Downloading certficate html and security policies')
-        self.new_files, new_algs = download_html_pages()
+        download_html_pages()
 
         logger.info(f"{self.new_files} needed to be downloaded")
 
@@ -605,8 +630,7 @@ class FIPSDataset(Dataset, ComplexSerializableType):
                     self.web_dir / f'{cert}.html')
         else:
             logger.info("Certs loaded from previous scanning")
-            dataset = json.loads(open(self.root_dir / 'fips_full_dataset.json').read(),
-                                 cls=import_module('sec_certs.serialization').CustomJSONDecoder)
+            dataset = self.from_json(self.root_dir / 'fips_full_dataset.json')
             self.certs = dataset.certs
 
     def extract_certs_from_tables(self) -> List[Path]:
@@ -823,7 +847,34 @@ class FIPSDataset(Dataset, ComplexSerializableType):
         single_dot.render(str(output_file_name) + '_single', view=True)
 
 
-class AlgorithmDataset(Dataset):
+    def to_dict(self):
+        return {'timestamp': self.timestamp, 'sha256_digest': self.sha256_digest,
+                'name': self.name, 'description': self.description,
+                'n_certs': len(self), 'certs': self.certs, 'algs': self.algorithms}
+
+    @classmethod
+    def from_dict(cls, dct: Dict):
+        certs = dct['certs']
+        dset = cls(certs, Path('./'), dct['name'], dct['description'])
+        dset.algorithms = dct['algs']
+        if len(dset) != (claimed := dct['n_certs']):
+            logger.error(f'The actual number of certs in dataset ({len(dset)}) does not match the claimed number ({claimed}).')
+        return dset
+
+    def to_json(self, output_path: Union[str, Path]):
+        with Path(output_path).open('w') as handle:
+            json.dump(self, handle, indent=4, cls=CustomJSONEncoder)
+
+    @classmethod
+    def from_json(cls, input_path: Union[str, Path]):
+        input_path = Path(input_path)
+        with input_path.open('r') as handle:
+            dset = json.load(handle, cls=CustomJSONDecoder)
+        dset.root_dir = input_path.parent.absolute()
+        return dset
+
+
+class AlgorithmDataset(Dataset, ComplexSerializableType):
 
     def get_certs_from_web(self):
         pass
@@ -850,3 +901,31 @@ class AlgorithmDataset(Dataset):
                         self.certs[alg_id] = []
                     self.certs[alg_id].append(fips_alg)
 
+    def convert_all_pdfs(self):
+        raise 'Not meant to be implemented'
+
+    def download_all_pdfs(self):
+        raise 'Not meant to be implemented'
+
+
+    def to_dict(self):
+        return {"certs": self.certs}
+
+    @classmethod
+    def from_dict(cls, dct: Dict):
+        certs = dct['certs']
+        dset = cls(certs, Path('./'), 'algorithms', 'algorithms used in dataset')
+        return dset
+
+    def to_json(self, output_path: Union[str, Path]):
+        with Path(output_path).open('w') as handle:
+            json.dump(self, handle, indent=4, cls=CustomJSONEncoder)
+
+    @classmethod
+    def from_json(cls, input_path: Union[str, Path]):
+        input_path = Path(input_path)
+        with input_path.open('r') as handle:
+            dset = json.load(handle, cls=CustomJSONDecoder)
+        dset.root_dir = input_path.parent.absolute()
+        return dset
+