parsing now depends on algorithms

5 files changed, 122 insertions, 62 deletions

diff --git a/fips_oop_demo.py b/fips_oop_demo.py
index 1778c93e..7db29a9a 100644
--- a/fips_oop_demo.py
+++ b/fips_oop_demo.py
@@ -1,9 +1,10 @@
-from sec_certs.dataset import FIPSDataset
+from sec_certs.dataset import FIPSDataset, AlgorithmDataset
 from pathlib import Path
 from datetime import datetime
 import logging
 from sec_certs.helpers import download_parallel
-
+from sec_certs.serialization import CustomJSONEncoder, CustomJSONDecoder
+import json
 
 def main():
     logging.basicConfig(level=logging.INFO)
@@ -15,35 +16,40 @@ def main():
     # Load metadata for certificates from CSV and HTML sources
     dset.get_certs_from_web()
 
+    logging.info(f'Finished parsing. Have dataset with {len(dset)} certificates.')
+    # Dump dataset into JSON
+
+    dset.dump_to_json()
+    logging.info(f'Dataset saved to {dset.root_dir}/fips_full_dataset.json')
+
+    logging.info("Extracting keywords now.")
+
+    dset.extract_keywords()
+
+    logging.info(f'Finished extracting certificates for {len(dset.keywords)} items.')
+    logging.info(f'Dumping keywords to {dset.root_dir}/fips_full_keywords.json')
+    dset.dump_keywords()
+
+    logging.info("Searching for tables in pdfs")
+
+    not_decoded_files = dset.extract_certs_from_tables()
+
+    logging.info(f"Done. Files not decoded: {not_decoded_files}")
+
+    logging.info("Parsing algorithms")
+    aset = AlgorithmDataset({}, Path('fips_dataset/web/algorithms'), 'algorithms', 'sample algs')
+    aset.parse_html()
+
+    dset.algorithms = aset
+
+    logging.info("finalizing results.")
+
+    dset.finalize_results()
+
+    logging.info('dump again')
+    dset.dump_to_json()
 
-    # logging.info(f'Finished parsing. Have dataset with {len(dset)} certificates.')
-    # # Dump dataset into JSON
-    #
-    # dset.dump_to_json()
-    # logging.info(f'Dataset saved to {dset.root_dir}/fips_full_dataset.json')
-    #
-    # logging.info("Extracting keywords now.")
-    #
-    # dset.extract_keywords()
-    #
-    # logging.info(f'Finished extracting certificates for {len(dset.keywords)} items.')
-    # logging.info(f'Dumping keywords to {dset.root_dir}/fips_full_keywords.json')
-    # dset.dump_keywords()
-    #
-    # logging.info("Searching for tables in pdfs")
-    #
-    # not_decoded_files = dset.extract_certs_from_tables()
-    #
-    # logging.info(f"Done. Files not decoded: {not_decoded_files}")
-    #
-    # logging.info("finalizing results.")
-    #
-    # dset.finalize_results()
-    #
-    # logging.info('dump again')
-    # dset.dump_to_json()
-    #
-    # dset.get_dot_graph('new_oop')
+    dset.get_dot_graph('different_new')
     end = datetime.now()
     logging.info(f'The computation took {(end - start)} seconds.')
 
diff --git a/sec_certs/cert_rules.py b/sec_certs/cert_rules.py
index 2100a60b..2e0ec46c 100644
--- a/sec_certs/cert_rules.py
+++ b/sec_certs/cert_rules.py
@@ -285,7 +285,7 @@ rules_fips_remove_algorithm_ids = [
     r"PAA[: #]*?\d{2}",
     r"PAA[: #]*?\d{1}",
     r"(?:#|Cert\.?|Certificate)[\s#]*?(\d+)?\s*?(?:AES|SHS|SHA|RSA|HMAC|Diffie-Hellman|DRBG|DES|CVL)",
-    r"PKCS[ ]?#?\d+",
+    r"PKCS[\s]?#?\d+",
     r"Survey #192"  # why would they get an address like this /o\ cert 2079
 ]
 rules_fips_cert = [
diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index dd267c14..5c8c06e6 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -161,8 +161,9 @@ class FIPSCertificate(Certificate):
         :return: list of all found algorithm IDs
         """
         set_items = set()
-        for m in re.finditer(rf"(?:#{'?' if in_pdf else 'C?'}\s?|Cert\.?[^. ]*?\s?)(?:[Cc]\s)?(?P<id>\d+)",
-                             current_text):
+        for m in re.finditer(
+                rf"(?:#{'?' if in_pdf else 'C?'}\s?|(?:Cert{'' if in_pdf else '?'})\.?[^. ]*?\s?)(?:[Cc]\s)?(?P<id>\d+)",
+                current_text):
             set_items.add(m.group())
 
         return list(set_items)
@@ -559,9 +560,21 @@ class FIPSAlgorithm(Certificate):
         # for each id
         return self.type
 
-    def __init__(self, vendor, implementation, alg_type, validation_date):
+    def __init__(self, cert_id, vendor, implementation, alg_type, validation_date):
         super().__init__()
+        self.cert_id = cert_id
         self.vendor = vendor
         self.implementation = implementation
         self.type = alg_type
         self.date = validation_date
+
+    def __repr__(self):
+        return self.type + ' algorithm #' + self.cert_id + ' created by ' + self.vendor
+
+    def __str__(self):
+        return str(self.type + ' algorithm #' + self.cert_id + ' created by ' + self.vendor)
+
+    @classmethod
+    def from_dict(cls, dct: dict) -> 'FIPSAlgorithm':
+        return FIPSAlgorithm(dct['cert_id'], dct['vendor'], dct['implementation'], dct['alg_type'],
+                             dct['validation_date'])
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 6b7f4d3b..e53b230c 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -20,7 +20,7 @@ from bs4 import BeautifulSoup
 from sec_certs.files import search_files
 from sec_certs import helpers as helpers
 from sec_certs.helpers import find_tables, repair_pdf
-from sec_certs.certificate import CommonCriteriaCert, Certificate, FIPSCertificate
+from sec_certs.certificate import CommonCriteriaCert, Certificate, FIPSCertificate, FIPSAlgorithm
 from sec_certs.extract_certificates import extract_certificates_keywords
 from sec_certs.constants import FIPS_NOT_AVAILABLE_CERT_SIZE
 
@@ -326,6 +326,7 @@ class FIPSDataset(Dataset):
                  description: str = 'dataset_description'):
         super().__init__(certs, root_dir, name, description)
         self.keywords = {}
+        self.algorithms = None
         self.new_files = 0
 
     @property
@@ -424,15 +425,15 @@ class FIPSDataset(Dataset):
         # Download files containing all available module certs (always)
         html_files = ['fips_modules_active.html',
                       'fips_modules_historical.html', 'fips_modules_revoked.html']
-        # helpers.download_file(
-        #     "https://csrc.nist.gov/projects/cryptographic-module-validation-program/validated-modules/search?SearchMode=Advanced&CertificateStatus=Active&ValidationYear=0",
-        #     self.web_dir / "fips_modules_active.html")
-        # helpers.download_file(
-        #     "https://csrc.nist.gov/projects/cryptographic-module-validation-program/validated-modules/search?SearchMode=Advanced&CertificateStatus=Historical&ValidationYear=0",
-        #     self.web_dir / "fips_modules_historical.html")
-        # helpers.download_file(
-        #     "https://csrc.nist.gov/projects/cryptographic-module-validation-program/validated-modules/search?SearchMode=Advanced&CertificateStatus=Revoked&ValidationYear=0",
-        #     self.web_dir / "fips_modules_revoked.html")
+        helpers.download_file(
+            "https://csrc.nist.gov/projects/cryptographic-module-validation-program/validated-modules/search?SearchMode=Advanced&CertificateStatus=Active&ValidationYear=0",
+            self.web_dir / "fips_modules_active.html")
+        helpers.download_file(
+            "https://csrc.nist.gov/projects/cryptographic-module-validation-program/validated-modules/search?SearchMode=Advanced&CertificateStatus=Historical&ValidationYear=0",
+            self.web_dir / "fips_modules_historical.html")
+        helpers.download_file(
+            "https://csrc.nist.gov/projects/cryptographic-module-validation-program/validated-modules/search?SearchMode=Advanced&CertificateStatus=Revoked&ValidationYear=0",
+            self.web_dir / "fips_modules_revoked.html")
 
         # Parse those files and get list of currently processable files (always)
         for f in html_files:
@@ -444,7 +445,6 @@ class FIPSDataset(Dataset):
         logging.info(f"{self.new_files} needed to be downloaded")
 
         if self.new_files > 0 or not (self.root_dir / 'fips_full_dataset.json').exists():
-            # if False:
             for cert in self.certs:
                 self.certs[cert] = FIPSCertificate.html_from_file(
                     self.web_dir / f'{cert}.html')
@@ -536,10 +536,40 @@ class FIPSDataset(Dataset):
                 self.keywords[file_name]['rules_cert_id'][rule].pop(
                     self.certs[file_name].cert_id, None)
 
+    def unify_algorithms(self):
+        for certificate in self.certs.values():
+            new_algorithms = []
+            for algorithm in certificate.algorithms:
+                if isinstance(algorithm, dict):
+                    new_algorithms.append(algorithm)
+                else:
+                    new_algorithms.append({'Certificate': algorithm})
+            certificate.algorithms = new_algorithms
+
     def validate_results(self):
         """
         Function that validates results and finds the final connection output
         """
+        def validate_id(processed_cert: FIPSCertificate, cert_candidate: str) -> bool:
+            # TODO: do we do this? #1 is used a lot
+            if cert_candidate == '1':
+                return False
+            if cert_candidate not in self.algorithms.certs:
+                return True
+
+            for cert_alg in processed_cert.algorithms:
+                for certificate in cert_alg['Certificate']:
+                    print(certificate)
+                    curr_id = ''.join(filter(str.isdigit, certificate))
+                    if curr_id == cert_candidate:
+                        return False
+
+            algs = self.algorithms.certs[cert_candidate]
+            for current_alg in algs:
+                if processed_cert.vendor[:3] in current_alg.vendor:
+                    return False
+            return True
+
         broken_files = set()
         for file_name in self.keywords:
             for rule in self.keywords[file_name]['rules_cert_id']:
@@ -547,12 +577,11 @@ class FIPSDataset(Dataset):
                     cert_id = ''.join(filter(str.isdigit, cert))
 
                     if cert_id == '' or cert_id not in self.certs:
-                        # TEST
-                        # if cert_id == '' or int(cert_id) > 3730:
                         broken_files.add(file_name)
                         self.keywords[file_name]['file_status'] = False
                         self.certs[file_name].file_status = False
                         break
+
         if broken_files:
             logging.warning("CERTIFICATE FILES WITH WRONG CERTIFICATES PARSED")
             logging.warning(broken_files)
@@ -568,22 +597,14 @@ class FIPSDataset(Dataset):
             for rule in self.keywords[file_name]['rules_cert_id']:
                 for cert in self.keywords[file_name]['rules_cert_id'][rule]:
                     cert_id = ''.join(filter(str.isdigit, cert))
-                    if cert_id not in self.certs[file_name].connections:
+                    if cert_id not in self.certs[file_name].connections and validate_id(self.certs[file_name], cert_id):
                         self.certs[file_name].connections.append(cert_id)
 
     def finalize_results(self):
+        self.unify_algorithms()
         self.remove_algorithms_from_extracted_data()
         self.validate_results()
 
-    def present_algorithms(self) -> Set[str]:
-        found_algs = set()
-        for cert in self.certs.values():
-            for alg in cert.algorithms:
-                if 'Name' in alg:
-                    print(cert, alg['Name'])
-                    found_algs.add(alg['Name'])
-        return found_algs
-
     def get_dot_graph(self, output_file_name: str):
         """
         Function that plots .dot graph of dependencies between certificates
@@ -641,7 +662,7 @@ class FIPSDataset(Dataset):
                     dot.edge(key, conn)
                     edges += 1
 
-        print(f"rendering {keys} keys and {edges} edges")
+        logging.info(f"rendering {keys} keys and {edges} edges")
 
         dot.render(str(output_file_name) + '_connections', view=True)
         single_dot.render(str(output_file_name) + '_single', view=True)
@@ -653,5 +674,24 @@ class AlgorithmDataset(Dataset):
         pass
 
     def parse_html(self):
+        def split_alg(alg_string):
+            cert_type = alg_string.rstrip('0123456789')
+            cert_id = alg_string[len(cert_type):]
+            return cert_type.strip(), cert_id.strip()
+
         for f in search_files(self.root_dir):
-            html_soup = BeautifulSoup(open(f).read(), 'html.parser')
\ No newline at end of file
+            html_soup = BeautifulSoup(open(f).read(), 'html.parser')
+            table = html_soup.find('table', class_='table table-condensed publications-table table-bordered')
+            spans = table.find_all('span')
+            for span in spans:
+                elements = span.find_all('td')
+                vendor, implementation = elements[0].text, elements[1].text
+                elements_sliced = elements[2:]
+                for i in range(0, len(elements_sliced), 2):
+                    alg_type, alg_id = split_alg(elements_sliced[i].text.strip())
+                    validation_date = elements_sliced[i + 1].text.strip()
+                    fips_alg = FIPSAlgorithm(alg_id, vendor, implementation, alg_type, validation_date)
+                    if alg_id not in self.certs:
+                        self.certs[alg_id] = []
+                    self.certs[alg_id].append(fips_alg)
+
diff --git a/sec_certs/serialization.py b/sec_certs/serialization.py
index 15c38a7b..88430920 100644
--- a/sec_certs/serialization.py
+++ b/sec_certs/serialization.py
@@ -3,11 +3,12 @@ from datetime import date
 from pathlib import Path
 
 from sec_certs.dataset import CCDataset, FIPSDataset
-from sec_certs.certificate import CommonCriteriaCert, FIPSCertificate
+from sec_certs.certificate import CommonCriteriaCert, FIPSCertificate, FIPSAlgorithm
 
 serializable_complex_types = (
-CCDataset, FIPSDataset, CommonCriteriaCert, CommonCriteriaCert.MaintainanceReport, CommonCriteriaCert.ProtectionProfile,
-FIPSCertificate)
+    CCDataset, FIPSDataset, CommonCriteriaCert, CommonCriteriaCert.MaintainanceReport,
+    CommonCriteriaCert.ProtectionProfile,
+    FIPSCertificate, FIPSAlgorithm)
 serializable_complex_types_dict = {x.__name__: x for x in serializable_complex_types}