Merge branch 'master' into fix-serialization

8 files changed, 110 insertions, 74 deletions

diff --git a/sec_certs/fips_certificates.py b/fips_certificates.py
index 22b055a8..1d5e3232 100755
--- a/sec_certs/fips_certificates.py
+++ b/fips_certificates.py
@@ -12,13 +12,14 @@ import click
 import pikepdf
 from tabula import read_pdf
 
-from .download import download_fips_web, download_fips
-from . import extract_certificates
-from .files import load_json_files, FILE_ERRORS_STRATEGY, search_files
+from sec_certs.download import download_fips_web, download_fips
+from sec_certs import extract_certificates
+from sec_certs.files import load_json_files, FILE_ERRORS_STRATEGY, search_files
 
 FIPS_BASE_URL = 'https://csrc.nist.gov'
 FIPS_MODULE_URL = 'https://csrc.nist.gov/projects/cryptographic-module-validation-program/certificate/'
 
+
 def extract_filename(file: str) -> str:
     """
     Extracts filename from path
diff --git a/sec_certs/process_certificates.py b/process_certificates.py
index 9ce784a7..143c6be2 100755
--- a/sec_certs/process_certificates.py
+++ b/process_certificates.py
@@ -1,11 +1,10 @@
 #!/usr/bin/env python3
 
 import click
-
-from .files import load_json_files
-from .extract_certificates import *
-from .analyze_certificates import *
-from .download import download_cc_web, download_cc
+from sec_certs.files import load_json_files
+from sec_certs.extract_certificates import *
+from sec_certs.analyze_certificates import *
+from sec_certs.download import download_cc_web, download_cc
 
 
 @click.command()
diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index 5c33f9e9..1ef3491b 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -6,11 +6,12 @@ from pathlib import Path
 import os
 import copy
 
-from . import helpers, extract_certificates
 from abc import ABC, abstractmethod
 from bs4 import Tag, BeautifulSoup, NavigableString
 from typing import Union, Optional, List, Dict, ClassVar, TypeVar, Type
 
+from sec_certs import helpers, extract_certificates
+
 
 class Certificate(ABC):
     T = TypeVar('T', bound='Certificate')
@@ -177,7 +178,8 @@ class FIPSCertificate(Certificate):
         trs = element.find_all('tr')
         for tr in trs:
             tds = tr.find_all('td')
-            found_items.append({'Name': tds[0].text, 'Certificate': parse_algorithms(tds[1].text)})
+            found_items.append(
+                {'Name': tds[0].text, 'Certificate': parse_algorithms(tds[1].text)})
 
         return found_items
 
@@ -192,7 +194,8 @@ class FIPSCertificate(Certificate):
                 html_items_found[pairs[title]] = content.split(';')
             elif 'caveat' in pairs[title]:
                 html_items_found[pairs[title]] = content
-                html_items_found['fips_mentioned_certs'] += FIPSCertificate.parse_caveat(content)
+                html_items_found['fips_mentioned_certs'] += FIPSCertificate.parse_caveat(
+                    content)
 
             elif 'FIPS Algorithms' in title:
                 html_items_found['fips_algorithms'] += FIPSCertificate.parse_table(
@@ -213,7 +216,8 @@ class FIPSCertificate(Certificate):
         vendor_string = current_div.find('div', 'panel-body').find('a')
 
         if not vendor_string:
-            vendor_string = list(current_div.find('div', 'panel-body').children)[0].strip()
+            vendor_string = list(current_div.find(
+                'div', 'panel-body').children)[0].strip()
             html_items_found['fips_vendor_www'] = ''
         else:
             html_items_found['fips_vendor_www'] = vendor_string.get('href')
@@ -225,9 +229,11 @@ class FIPSCertificate(Certificate):
 
     @staticmethod
     def parse_lab(current_div: Tag, html_items_found: Dict, current_file: Path):
-        html_items_found['fips_lab'] = list(current_div.find('div', 'panel-body').children)[0].strip()
+        html_items_found['fips_lab'] = list(
+            current_div.find('div', 'panel-body').children)[0].strip()
         html_items_found['fips_nvlap_code'] = \
-            list(current_div.find('div', 'panel-body').children)[2].strip().split('\n')[1].strip()
+            list(current_div.find(
+                'div', 'panel-body').children)[2].strip().split('\n')[1].strip()
 
         if html_items_found['fips_lab'] == '':
             logging.warning(f"WARNING: NO LAB FOUND{current_file}")
@@ -238,7 +244,7 @@ class FIPSCertificate(Certificate):
     @staticmethod
     def parse_related_files(current_div: Tag, html_items_found: Dict):
         links = current_div.find_all('a')
-        ## TODO: break out of circular imports hell
+        # TODO: break out of circular imports hell
         html_items_found['fips_security_policy_www'] = __import__(
             'sec_certs').certificate.FIPSCertificate.FIPS_BASE_URL + links[0].get('href')
 
@@ -329,9 +335,12 @@ class CommonCriteriaCert(Certificate):
         maintainance_st_link: str
 
         def __post_init__(self):
-            super().__setattr__('maintainance_report_link', helpers.sanitize_link(self.maintainance_report_link))
-            super().__setattr__('maintainance_st_link', helpers.sanitize_link(self.maintainance_st_link))
-            super().__setattr__('maintainance_title', helpers.sanitize_string(self.maintainance_title))
+            super().__setattr__('maintainance_report_link',
+                                helpers.sanitize_link(self.maintainance_report_link))
+            super().__setattr__('maintainance_st_link',
+                                helpers.sanitize_link(self.maintainance_st_link))
+            super().__setattr__('maintainance_title',
+                                helpers.sanitize_string(self.maintainance_title))
             super().__setattr__('maintainance_date', helpers.sanitize_date(self.maintainance_date))
 
         def to_dict(self):
@@ -467,7 +476,8 @@ class CommonCriteriaCert(Certificate):
 
         def get_date(cell: Tag) -> date:
             text = cell.get_text()
-            extracted_date = datetime.strptime(text, '%Y-%m-%d').date() if text else None
+            extracted_date = datetime.strptime(
+                text, '%Y-%m-%d').date() if text else None
             return extracted_date
 
         def get_report_st_links(cell: Tag) -> (str, str):
@@ -477,7 +487,8 @@ class CommonCriteriaCert(Certificate):
             assert links[2].get('title').startswith('Security Target')
 
             report_link = CommonCriteriaCert.cc_url + links[1].get('href')
-            security_target_link = CommonCriteriaCert.cc_url + links[2].get('href')
+            security_target_link = CommonCriteriaCert.cc_url + \
+                links[2].get('href')
 
             return report_link, security_target_link
 
@@ -497,16 +508,19 @@ class CommonCriteriaCert(Certificate):
             maintainance_updates = set()
             for u in possible_updates:
                 text = list(u.stripped_strings)[0]
-                main_date = datetime.strptime(text.split(' ')[0], '%Y-%m-%d').date() if text else None
+                main_date = datetime.strptime(text.split(
+                    ' ')[0], '%Y-%m-%d').date() if text else None
                 main_title = text.split('– ')[1]
                 main_report_link = None
                 main_st_link = None
                 links = u.find_all('a')
                 for l in links:
                     if l.get('title').startswith('Maintenance Report:'):
-                        main_report_link = CommonCriteriaCert.cc_url + l.get('href')
+                        main_report_link = CommonCriteriaCert.cc_url + \
+                            l.get('href')
                     elif l.get('title').startswith('Maintenance ST'):
-                        main_st_link = CommonCriteriaCert.cc_url + l.get('href')
+                        main_st_link = CommonCriteriaCert.cc_url + \
+                            l.get('href')
                     else:
                         logging.error('Unknown link in Maintenance part!')
                 maintainance_updates.add(
@@ -530,7 +544,8 @@ class CommonCriteriaCert(Certificate):
         cert_link = get_cert_link(cells[2])
 
         maintainance_div = get_maintainance_div(cells[0])
-        maintainances = get_maintainance_updates(maintainance_div) if maintainance_div else set()
+        maintainances = get_maintainance_updates(
+            maintainance_div) if maintainance_div else set()
 
         return cls(category, name, manufacturer, scheme, security_level, not_valid_before, not_valid_after, report_link,
                    st_link, 'html', cert_link, manufacturer_web, protection_profiles, maintainances)
diff --git a/sec_certs/constants.py b/sec_certs/constants.py
index d8a80fa3..3ffb867c 100644
--- a/sec_certs/constants.py
+++ b/sec_certs/constants.py
@@ -1,8 +1,6 @@
 from enum import Enum
 
 
-
-
 class CertFramework(Enum):
     CC = 'Common Criteria'
     FIPS = 'FIPS'
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 73efd66a..5158214c 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -1,27 +1,26 @@
 import os
 import re
-from datetime import datetime, date
-import copy
-
-from tabula import read_pdf
+from datetime import datetime
+import locale
+import logging
+from typing import Dict, List, ClassVar
+import json
+from importlib import import_module
 
-from .certificate import CommonCriteriaCert, Certificate, FIPSCertificate
-from .extract_certificates import extract_certificates_keywords
-from .constants import FIPS_NOT_AVAILABLE_CERT_SIZE
 from abc import ABC, abstractmethod
-from . import helpers as helpers
 from pathlib import Path
 import shutil
+
+from tabula import read_pdf
 import pandas as pd
 from bs4 import BeautifulSoup
-import locale
-import logging
-from typing import Dict, List, Optional, Set, ClassVar
-import json
-from importlib import import_module
 
-from .files import search_files
-from .helpers import find_tables, repair_pdf
+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.extract_certificates import extract_certificates_keywords
+from sec_certs.constants import FIPS_NOT_AVAILABLE_CERT_SIZE
 
 
 class Dataset(ABC):
@@ -61,8 +60,8 @@ class Dataset(ABC):
 
     def to_dict(self):
         return copy.deepcopy({'root_dir': self.root_dir, 'timestamp': self.timestamp, 'sha256_digest': self.sha256_digest,
-                'name': self.name,
-                'description': self.description, 'n_certs': len(self), 'certs': list(self.certs.values())})
+                              'name': self.name,
+                              'description': self.description, 'n_certs': len(self), 'certs': list(self.certs.values())})
 
     @classmethod
     def from_dict(cls, dct: Dict):
@@ -94,7 +93,8 @@ class Dataset(ABC):
                 n_merged += 1
 
         self.certs.update(will_be_added)
-        logging.info(f'Added {len(will_be_added)} new and merged further {n_merged} certificates to the dataset.')
+        logging.info(
+            f'Added {len(will_be_added)} new and merged further {n_merged} certificates to the dataset.')
 
 
 class CCDataset(Dataset):
@@ -127,8 +127,10 @@ class CCDataset(Dataset):
         """
         self.web_dir.mkdir(parents=True, exist_ok=True)
 
-        html_items = [(x, self.web_dir / y) for y, x in self.html_products.items()]
-        csv_items = [(x, self.web_dir / y) for y, x in self.csv_products.items()]
+        html_items = [(x, self.web_dir / y)
+                      for y, x in self.html_products.items()]
+        csv_items = [(x, self.web_dir / y)
+                     for y, x in self.csv_products.items()]
 
         if not get_active:
             html_items = [x for x in html_items if 'active' not in str(x[1])]
@@ -162,13 +164,16 @@ class CCDataset(Dataset):
         Creates dictionary of new certificates from csv sources.
         """
         csv_sources = self.csv_products.keys()
-        csv_sources = [x for x in csv_sources if 'active' not in x or get_active]
-        csv_sources = [x for x in csv_sources if 'archived' not in x or get_archived]
+        csv_sources = [
+            x for x in csv_sources if 'active' not in x or get_active]
+        csv_sources = [
+            x for x in csv_sources if 'archived' not in x or get_archived]
 
         new_certs = {}
         for file in csv_sources:
             partial_certs = self.parse_single_csv(self.web_dir / file)
-            logging.info(f'Parsed {len(partial_certs)} certificates from: {file}')
+            logging.info(
+                f'Parsed {len(partial_certs)} certificates from: {file}')
             new_certs.update(partial_certs)
         return new_certs
 
@@ -187,7 +192,8 @@ class CCDataset(Dataset):
                       'maintainance_title', 'maintainance_report_link', 'maintainance_st_link']
 
         df = pd.read_csv(file, engine='python', encoding='windows-1250')
-        df = df.rename(columns={x: y for (x, y) in zip(list(df.columns), csv_header)})
+        df = df.rename(
+            columns={x: y for (x, y) in zip(list(df.columns), csv_header)})
 
         df['is_maintainance'] = ~df.maintainance_title.isnull()
         df = df.fillna(value='')
@@ -195,14 +201,16 @@ class CCDataset(Dataset):
         df[['not_valid_before', 'not_valid_after', 'maintainance_date']] = df[
             ['not_valid_before', 'not_valid_after', 'maintainance_date']].apply(pd.to_datetime)
 
-        df['dgst'] = df.apply(lambda row: helpers.get_first_16_bytes_sha256(get_primary_key_str(row)), axis=1)
+        df['dgst'] = df.apply(lambda row: helpers.get_first_16_bytes_sha256(
+            get_primary_key_str(row)), axis=1)
         df_base = df.loc[df.is_maintainance == False].copy()
         df_main = df.loc[df.is_maintainance == True].copy()
 
         n_all = len(df_base)
         n_deduplicated = len(df_base.drop_duplicates(subset=['dgst']))
         if (n_dup := n_all - n_deduplicated) > 0:
-            logging.warning(f'The CSV {file} contains {n_dup} duplicates by the primary key.')
+            logging.warning(
+                f'The CSV {file} contains {n_dup} duplicates by the primary key.')
 
         df_base = df_base.drop_duplicates(subset=['dgst'])
         df_main = df_main.drop_duplicates()
@@ -227,13 +235,16 @@ class CCDataset(Dataset):
         Prepares dictionary of certificates from all html files.
         """
         html_sources = self.html_products.keys()
-        html_sources = [x for x in html_sources if 'active' not in x or get_active]
-        html_sources = [x for x in html_sources if 'archived' not in x or get_archived]
+        html_sources = [
+            x for x in html_sources if 'active' not in x or get_active]
+        html_sources = [
+            x for x in html_sources if 'archived' not in x or get_archived]
 
         new_certs = {}
         for file in html_sources:
             partial_certs = self.parse_single_html(self.web_dir / file)
-            logging.info(f'Parsed {len(partial_certs)} certificates from: {file}')
+            logging.info(
+                f'Parsed {len(partial_certs)} certificates from: {file}')
             new_certs.update(partial_certs)
         return new_certs
 
@@ -248,7 +259,8 @@ class CCDataset(Dataset):
             footer_text = list(footer.stripped_strings)[0]
             date_string = footer_text.split(',')[1:3]
             time_string = footer_text.split(',')[3].split(' at ')[1]
-            formatted_datetime = date_string[0] + date_string[1] + ' ' + time_string
+            formatted_datetime = date_string[0] + \
+                date_string[1] + ' ' + time_string
             return datetime.strptime(formatted_datetime, ' %B %d %Y %I:%M %p')
 
         def parse_table(soup: BeautifulSoup, table_id: str, category_string: str) -> Dict[str, 'CommonCriteriaCert']:
@@ -268,7 +280,8 @@ class CCDataset(Dataset):
             # TODO: Do we have use for number of expected certs? We get rid of duplicites, so no use for assert expected == actual
             # caption_str = str(table.findAll('caption'))
             # n_expected_certs = int(caption_str.split(category_string + ' – ')[1].split(' Certified Products')[0])
-            table_certs = {x.dgst: x for x in [CommonCriteriaCert.from_html_row(row, category_string) for row in body]}
+            table_certs = {x.dgst: x for x in [
+                CommonCriteriaCert.from_html_row(row, category_string) for row in body]}
 
             return table_certs
 
@@ -347,11 +360,13 @@ class FIPSDataset(Dataset):
                 self.fragments_dir, 'fips', fips_items=self.certs,
                 should_censure_right_away=True)
         else:
-            self.keywords = json.loads(open(self.root_dir / 'fips_full_keywords.json').read())
+            self.keywords = json.loads(
+                open(self.root_dir / 'fips_full_keywords.json').read())
 
     def dump_to_json(self):
         with open(self.root_dir / 'fips_full_dataset.json', 'w') as handle:
-            json.dump(self, handle, cls=import_module('sec_certs.serialization').CustomJSONEncoder, indent=4)
+            json.dump(self, handle, cls=import_module(
+                'sec_certs.serialization').CustomJSONEncoder, indent=4)
 
     def dump_keywords(self):
         with open(self.root_dir / "fips_full_keywords.json", 'w') as f:
@@ -364,7 +379,8 @@ class FIPSDataset(Dataset):
             logging.info(f'Getting certificate ids from {html_file}')
             html = BeautifulSoup(open(html_file).read(), 'html.parser')
 
-            table = [x for x in html.find(id='searchResultsTable').tbody.contents if x != '\n']
+            table = [x for x in html.find(
+                id='searchResultsTable').tbody.contents if x != '\n']
             for entry in table:
                 self.certs[entry.find('a').text] = {}
 
@@ -374,7 +390,8 @@ class FIPSDataset(Dataset):
         self.policies_dir.mkdir(exist_ok=True)
 
         # Download files containing all available module certs (always)
-        html_files = ['fips_modules_active.html', 'fips_modules_historical.html', 'fips_modules_revoked.html']
+        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")
@@ -399,14 +416,16 @@ class FIPSDataset(Dataset):
             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()]
 
-        _, self.new_files = helpers.download_parallel(html_items + sp_items, 8), len(html_items) + len(sp_items)
+        _, self.new_files = helpers.download_parallel(
+            html_items + sp_items, 8), len(html_items) + len(sp_items)
 
         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')
+                self.certs[cert] = FIPSCertificate.html_from_file(
+                    self.web_dir / f'{cert}.html')
         else:
             logging.info("Certs loaded from previous scanning")
             dataset = json.loads(open(self.root_dir / 'fips_full_dataset.json').read(),
@@ -439,11 +458,13 @@ class FIPSDataset(Dataset):
             if tables:
                 lst = []
                 try:
-                    data = read_pdf(cert_file.with_suffix(''), pages=tables, silent=True)
+                    data = read_pdf(cert_file.with_suffix(''),
+                                    pages=tables, silent=True)
                 except Exception:
                     try:
                         repair_pdf(cert_file.with_suffix(''))
-                        data = read_pdf(cert_file.with_suffix(''), pages=tables, silent=True)
+                        data = read_pdf(cert_file.with_suffix(
+                            ''), pages=tables, silent=True)
 
                     except Exception:
                         not_decoded.append(cert_file)
@@ -453,10 +474,12 @@ class FIPSDataset(Dataset):
                 for df in data:
                     for col in range(len(df.columns)):
                         if 'cert' in df.columns[col].lower() or 'algo' in df.columns[col].lower():
-                            lst += FIPSCertificate.parse_algorithms(df.iloc[:, col].to_string(index=False), True)
+                            lst += FIPSCertificate.parse_algorithms(
+                                df.iloc[:, col].to_string(index=False), True)
 
                     # Parse again if someone picks not so descriptive column names
-                    lst += FIPSCertificate.parse_algorithms(df.to_string(index=False))
+                    lst += FIPSCertificate.parse_algorithms(
+                        df.to_string(index=False))
 
                 if lst:
                     self.certs[stem_name].algorithms += lst
@@ -485,7 +508,8 @@ class FIPSDataset(Dataset):
                                     cert).group('id'):
                                 to_pop.add(cert)
                 for r in to_pop:
-                    self.keywords[file_name]['rules_cert_id'][rule].pop(r, None)
+                    self.keywords[file_name]['rules_cert_id'][rule].pop(
+                        r, None)
 
                 self.keywords[file_name]['rules_cert_id'][rule].pop(
                     self.certs[file_name].cert_id, None)
diff --git a/sec_certs/helpers.py b/sec_certs/helpers.py
index 50518911..dd95b0c8 100644
--- a/sec_certs/helpers.py
+++ b/sec_certs/helpers.py
@@ -1,6 +1,5 @@
 import re
 from typing import Sequence, Tuple, Optional, Set, List, Dict
-
 import logging
 import pikepdf
 import requests
@@ -16,7 +15,6 @@ import pandas as pd
 from bs4 import Tag, NavigableString
 
 
-
 def download_file(url: str, output: Path) -> int:
     r = requests.get(url, allow_redirects=True)
     with output.open('wb') as f:
diff --git a/sec_certs/serialization.py b/sec_certs/serialization.py
index 8283ac35..15c38a7b 100644
--- a/sec_certs/serialization.py
+++ b/sec_certs/serialization.py
@@ -1,8 +1,9 @@
 import json
 from datetime import date
 from pathlib import Path
-from .dataset import CCDataset, FIPSDataset
-from .certificate import CommonCriteriaCert, FIPSCertificate
+
+from sec_certs.dataset import CCDataset, FIPSDataset
+from sec_certs.certificate import CommonCriteriaCert, FIPSCertificate
 
 serializable_complex_types = (
 CCDataset, FIPSDataset, CommonCriteriaCert, CommonCriteriaCert.MaintainanceReport, CommonCriteriaCert.ProtectionProfile,
diff --git a/setup.py b/setup.py
index ad7e43ed..f50e45b3 100644
--- a/setup.py
+++ b/setup.py
@@ -41,7 +41,7 @@ setup(
     },
     entry_points="""
         [console_scripts]
-        process-certs=sec_certs.process_certificates:main
-        fips-certs=sec_certs.fips_certificates:main
+        process-certs=process_certificates:main
+        fips-certs=fips_certificates:main
     """
 )