Merge pull request #35 from petrs/cc_introduce_dataset_state

Introduce state variable of CCDataset

4 files changed, 83 insertions, 24 deletions

diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index 6f451f01..6f1c35b9 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -1063,7 +1063,7 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
     def download_pdf_report(cert: 'CommonCriteriaCert') -> 'CommonCriteriaCert':
         exit_code = helpers.download_file(cert.report_link, cert.state.report_pdf_path)
         if exit_code != requests.codes.ok:
-            logger.error(f'Failed to download report from {cert.report_link}, code: {exit_code}')
+            logger.error(f'Cert dgst: {cert.dgst} failed to download report from {cert.report_link}, code: {exit_code}')
             cert.state.report_link_ok = False
         return cert
 
@@ -1071,7 +1071,7 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
     def download_pdf_target(cert: 'CommonCriteriaCert') -> 'CommonCriteriaCert':
         exit_code = helpers.download_file(cert.st_link, cert.state.st_pdf_path)
         if exit_code != requests.codes.ok:
-            logger.error(f'Cert dgst: {cert.dgst} failed to download report from {cert.report_link}, code: {exit_code}')
+            logger.error(f'Cert dgst: {cert.dgst} failed to download ST from {cert.report_link}, code: {exit_code}')
             cert.state.st_link_ok = False
         return cert
 
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 351a0f56..56cd3a1b 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -2,8 +2,10 @@ import os
 from datetime import datetime
 import locale
 import logging
-from typing import Dict, List, ClassVar, Collection, Union, Set, Tuple
+from typing import Dict, List, ClassVar, Collection, Union, Set, Tuple, Optional
 from itertools import groupby
+from dataclasses import dataclass
+import copy
 
 import json
 from abc import ABC, abstractmethod
@@ -125,6 +127,35 @@ class Dataset(ABC):
 
 class CCDataset(Dataset, ComplexSerializableType):
     # TODO: Make properties propagate to changing internal state of related certificates
+    @dataclass
+    class DatasetInternalState(ComplexSerializableType):
+        meta_sources_parsed: bool = False
+        pdfs_downloaded: bool = False
+        pdfs_converted: bool = False
+        txt_data_extracted: bool = False
+        certs_analyzed: bool = False
+
+        def to_dict(self):
+            return copy.deepcopy(self.__dict__)
+
+        @classmethod
+        def from_dict(cls, dct: Dict[str, bool]):
+            return cls(*tuple(dct.values()))
+
+    def __init__(self, certs: Dict[str, 'Certificate'], root_dir: Path, name: str = 'dataset name', description: str = 'dataset_description', state: Optional[DatasetInternalState] = None):
+        super().__init__(certs, root_dir, name, description)
+        if state is None:
+            state = self.DatasetInternalState()
+        self.state = state
+
+    def to_dict(self):
+        return {**{'state': self.state}, **super().to_dict()}
+
+    @classmethod
+    def from_dict(cls, dct: Dict):
+        dset = super().from_dict(dct)
+        dset.state = copy.deepcopy(dct['state'])
+        return dset
 
     @Dataset.root_dir.setter
     def root_dir(self, new_dir: Union[str, Path]):
@@ -252,6 +283,7 @@ class CCDataset(Dataset, ComplexSerializableType):
             shutil.rmtree(self.web_dir)
 
         self.set_local_paths()
+        self.state.meta_sources_parsed = True
 
     def _get_all_certs_from_csv(self, get_active: bool, get_archived: bool) -> Dict[str, 'CommonCriteriaCert']:
         """
@@ -436,6 +468,10 @@ class CCDataset(Dataset, ComplexSerializableType):
         cert_processing.process_parallel(CommonCriteriaCert.download_pdf_target, certs_to_process, constants.N_THREADS)
 
     def download_all_pdfs(self, fresh: bool = True):
+        if self.state.meta_sources_parsed is False:
+            logger.error('Attempting to download pdfs while not having csv/html meta-sources parsed. Returning.')
+            return
+
         logger.info('Downloading CC certificate reports')
         self._download_reports(fresh)
 
@@ -452,25 +488,31 @@ class CCDataset(Dataset, ComplexSerializableType):
                 logger.info('Attempting to re-download failed security target links.')
                 self._download_targets(False)
 
+        self.state.pdfs_downloaded = True
+
     def _convert_reports_to_txt(self, fresh: bool = True):
         self.reports_txt_dir.mkdir(parents=True, exist_ok=True)
 
         if fresh is True:
-            certs_to_process = self.certs.values()
+            certs_to_process = [x for x in self.certs.values() if x.state.report_link_ok]
         else:
-            certs_to_process = [x for x in self.certs.values() if not x.state.report_convert_ok]
+            certs_to_process = [x for x in self.certs.values() if x.state.report_link_ok and not x.state.report_convert_ok]
         cert_processing.process_parallel(CommonCriteriaCert.convert_report_pdf, certs_to_process, constants.N_THREADS)
 
     def _convert_targets_to_txt(self, fresh: bool = True):
         self.targets_txt_dir.mkdir(parents=True, exist_ok=True)
 
         if fresh is True:
-            certs_to_process = self.certs.values()
+            certs_to_process = [x for x in self.certs.values() if x.state.st_link_ok]
         else:
-            certs_to_process = [x for x in self.certs.values() if not x.state.st_convert_ok]
+            certs_to_process = [x for x in self.certs.values() if x.state.st_link_ok and not x.state.st_convert_ok]
         cert_processing.process_parallel(CommonCriteriaCert.convert_target_pdf, certs_to_process, constants.N_THREADS)
 
     def convert_all_pdfs(self, fresh: bool = True):
+        if self.state.pdfs_downloaded is False:
+            logger.info('Attempting to convert pdf while not having them downloaded. Returning.')
+            return
+
         logger.info('Converting CC certificate reports to .txt')
         self._convert_reports_to_txt(fresh)
 
@@ -478,26 +520,28 @@ class CCDataset(Dataset, ComplexSerializableType):
         self._convert_targets_to_txt(fresh)
 
         if fresh is True:
-            # Attempt to re-convert once if some files failed
-            if any(filter(lambda x: not x.state.report_convert_ok, self.certs.values())):
+            # Attempt to re-convert once if some files failed but downloads are ok
+            if any(filter(lambda x: x.state.report_link_ok and not x.state.report_convert_ok, self.certs.values())):
                 logger.info('Attempting to re-convert failed report pdfs')
                 self._convert_reports_to_txt(False)
-            if any(filter(lambda x: not x.state.st_convert_ok, self.certs.values())):
+            if any(filter(lambda x: x.state.st_link_ok and not x.state.st_convert_ok, self.certs.values())):
                 logger.info('Attempting to re-convert failed target pdfs')
                 self._convert_targets_to_txt(False)
 
+        self.state.pdfs_converted = True
+
     def _extract_report_metadata(self, fresh: bool = True):
         if fresh is True:
-            certs_to_process = self.certs.values()
+            certs_to_process = [x for x in self.certs.values() if x.state.report_convert_ok]
         else:
-            certs_to_process = [x for x in self.certs.values() if not x.state.report_extract_ok]
+            certs_to_process = [x for x in self.certs.values() if x.state.report_convert_ok and not x.state.report_extract_ok]
         cert_processing.process_parallel(CommonCriteriaCert.extract_report_pdf_metadata, certs_to_process, constants.N_THREADS)
 
     def _extract_targets_metadata(self, fresh: bool = True):
         if fresh is True:
-            certs_to_process = self.certs.values()
+            certs_to_process = [x for x in self.certs.values() if x.state.st_convert_ok]
         else:
-            certs_to_process = [x for x in self.certs.values() if not x.state.st_extract_ok]
+            certs_to_process = [x for x in self.certs.values() if x.state.st_convert_ok and not x.state.st_extract_ok]
         cert_processing.process_parallel(CommonCriteriaCert.extract_st_pdf_metadata, certs_to_process, constants.N_THREADS)
 
     def extract_pdf_metadata(self, fresh: bool = True):
@@ -507,16 +551,16 @@ class CCDataset(Dataset, ComplexSerializableType):
 
     def _extract_targets_frontpage(self, fresh: bool = True):
         if fresh is True:
-            certs_to_process = self.certs.values()
+            certs_to_process = [x for x in self.certs.values() if x.state.st_convert_ok]
         else:
-            certs_to_process = [x for x in self.certs.values() if not x.state.st_extract_ok]
+            certs_to_process = [x for x in self.certs.values() if x.state.st_convert_ok and not x.state.st_extract_ok]
         cert_processing.process_parallel(CommonCriteriaCert.extract_st_pdf_frontpage, certs_to_process, constants.N_THREADS)
 
     def _extract_report_frontpage(self, fresh: bool = True):
         if fresh is True:
-            certs_to_process = self.certs.values()
+            certs_to_process = [x for x in self.certs.values() if x.state.report_convert_ok]
         else:
-            certs_to_process = [x for x in self.certs.values() if not x.state.report_extract_ok]
+            certs_to_process = [x for x in self.certs.values() if x.state.report_convert_ok and not x.state.report_extract_ok]
         cert_processing.process_parallel(CommonCriteriaCert.extract_report_pdf_frontpage, certs_to_process, constants.N_THREADS)
 
     def extract_pdf_frontpage(self, fresh: bool = True):
@@ -526,16 +570,16 @@ class CCDataset(Dataset, ComplexSerializableType):
 
     def _extract_report_keywords(self, fresh: bool = True):
         if fresh is True:
-            certs_to_process = self.certs.values()
+            certs_to_process = [x for x in self.certs.values() if x.state.report_convert_ok]
         else:
-            certs_to_process = [x for x in self.certs.values() if not x.state.report_extract_ok]
+            certs_to_process = [x for x in self.certs.values() if x.state.report_convert_ok and not x.state.report_extract_ok]
         cert_processing.process_parallel(CommonCriteriaCert.extract_report_pdf_keywords, certs_to_process, constants.N_THREADS)
 
     def _extract_targets_keywords(self, fresh: bool = True):
         if fresh is True:
-            certs_to_process = self.certs.values()
+            certs_to_process = [x for x in self.certs.values() if x.state.st_convert_ok]
         else:
-            certs_to_process = [x for x in self.certs.values() if not x.state.st_extract_ok]
+            certs_to_process = [x for x in self.certs.values() if x.state.st_convert_ok and not x.state.st_extract_ok]
         cert_processing.process_parallel(CommonCriteriaCert.extract_st_pdf_keywords, certs_to_process, constants.N_THREADS)
 
     def extract_pdf_keywords(self, fresh: bool = True):
@@ -544,23 +588,29 @@ class CCDataset(Dataset, ComplexSerializableType):
         self._extract_targets_keywords(fresh)
 
     def extract_data(self, fresh: bool = True):
+        if self.state.pdfs_converted is False:
+            logger.info('Attempting to extract data from txt while not having the pdf->txt conversion done. Returning.')
+            return
+
         logger.info('Extracting various stuff from converted txt filed from CC dataset.')
         self.extract_pdf_metadata(fresh)
         self.extract_pdf_frontpage(fresh)
         self.extract_pdf_keywords(fresh)
 
         if fresh is True:
-            if any(filter(lambda x: not x.state.report_extract_ok, self.certs.values())):
+            if any(filter(lambda x: x.state.report_convert_ok and not x.state.report_extract_ok, self.certs.values())):
                 logger.info('Attempting to re-extract failed data from report txts')
                 self._extract_report_metadata(False)
                 self._extract_report_frontpage(False)
                 self._extract_report_keywords(False)
-            if any(filter(lambda x: not x.state.st_extract_ok, self.certs.values())):
+            if any(filter(lambda x: x.state.st_convert_ok and not x.state.st_extract_ok, self.certs.values())):
                 logger.info('Attempting to re-extract failed data from ST txts')
                 self._extract_targets_metadata(False)
                 self._extract_targets_frontpage(False)
                 self._extract_targets_keywords(False)
 
+        self.state.txt_data_extracted = True
+
 
 class FIPSDataset(Dataset, ComplexSerializableType):
     FIPS_BASE_URL: ClassVar[str] = 'https://csrc.nist.gov'
diff --git a/test/data/test_cc_oop/toy_dataset.json b/test/data/test_cc_oop/toy_dataset.json
index e27076f3..b59327ce 100644
--- a/test/data/test_cc_oop/toy_dataset.json
+++ b/test/data/test_cc_oop/toy_dataset.json
@@ -1,5 +1,13 @@
 {
     "_type": "CCDataset",
+    "state": {
+        "_type": "DatasetInternalState",
+        "meta_sources_parsed": true,
+        "pdfs_downloaded": false,
+        "pdfs_converted": false,
+        "txt_data_extracted": false,
+        "certs_analyzed": false
+    },
     "timestamp": "2020-11-16 17:04:14.770153",
     "sha256_digest": "not implemented",
     "name": "toy dataset",
diff --git a/test/test_cc_oop.py b/test/test_cc_oop.py
index 7bfac882..332ab3f8 100644
--- a/test/test_cc_oop.py
+++ b/test/test_cc_oop.py
@@ -74,6 +74,7 @@ class TestCommonCriteriaOOP(TestCase):
                                                  None)
         self.template_dataset = CCDataset({self.crt_one.dgst: self.crt_one, self.crt_two.dgst: self.crt_two}, Path('/fictional/path/to/dataset'), 'toy dataset', 'toy dataset description')
         self.template_dataset.timestamp = datetime(2020, 11, 16, hour=17, minute=4, second=14, microsecond=770153)
+        self.template_dataset.state.meta_sources_parsed = True
 
         self.template_report_pdf_hashes = {'7ef8227c1aed06bb': '774c41fbba980191ca40ae610b2f61484c5997417b3325b6fd68b345173bde52',
                                           '561a012d9a30e960': '533a5995ef8b736cc48cfda30e8aafec77d285511471e0e5a9e8007c8750203a'}