Merge pull request #8 from petrs/fips

fips basic functionality

3 files changed, 111 insertions, 48 deletions

diff --git a/requirements.txt b/requirements.txt
index 43e7ade0..068df665 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,3 +10,6 @@ PyPDF2==1.26.0
 python-dateutil==2.8.1
 six==1.15.0
 tabulate==0.8.7
+
+setuptools~=50.3.2
+pikepdf~=1.19.3
\ No newline at end of file
diff --git a/sec_certs/download.py b/sec_certs/download.py
index eacee1f4..33733634 100644
--- a/sec_certs/download.py
+++ b/sec_certs/download.py
@@ -8,7 +8,6 @@ import requests
 from .extract_certificates import PDF2TEXT_CONVERT
 from .files import search_files, FILE_ERRORS_STRATEGY
 
-
 CC_WEB_URL = 'https://www.commoncriteriaportal.org'
 
 
@@ -16,7 +15,7 @@ def download_file(url: str, output: Path) -> int:
     r = requests.get(url, allow_redirects=True)
     with output.open("wb") as f:
         f.write(r.content)
-        #for chunk in r.iter_content(chunk_size=1024):
+        # for chunk in r.iter_content(chunk_size=1024):
         #    if chunk:
         #        f.write(chunk)
     return r.status_code
@@ -26,8 +25,8 @@ def generate_download_script(file_name, certs_dir, targets_dir, base_url, downlo
     with open(file_name, "w", errors=FILE_ERRORS_STRATEGY) as write_file:
         # certs files
         if certs_dir != '':
-            write_file.write('mkdir \"{}\"\n'.format(certs_dir))
-            write_file.write('cd \"{}\"\n\n'.format(certs_dir))
+            write_file.write(f'mkdir \"{certs_dir}\"\n')
+            write_file.write(f'cd \"{certs_dir}\"\n\n')
         for cert in download_files_certs:
             # double %% is necessary to prevent replacement of %2 within script (second argument of script)
             file_name_short_web = cert[0].replace(' ', '%%20')
@@ -35,42 +34,45 @@ def generate_download_script(file_name, certs_dir, targets_dir, base_url, downlo
             if file_name_short_web.find(base_url) != -1:
                 # base url already included
                 write_file.write(
-                    'curl \"{}\" -o \"{}\"\n'.format(file_name_short_web, cert[1]))
+                    f'curl \"{file_name_short_web}\" -o \"{cert[1]}\"\n')
             else:
                 # insert base url
                 write_file.write(
-                    'curl \"{}{}\" -o \"{}\"\n'.format(base_url, file_name_short_web, cert[1]))
-            write_file.write('{} \"{}\"\n\n'.format(PDF2TEXT_CONVERT, cert[1]))
+                    f'curl \"{base_url}{file_name_short_web}\" -o \"{cert[1]}\"\n')
+            write_file.write(f'{PDF2TEXT_CONVERT} \"{cert[1]}\"\n\n')
 
         if len(download_files_certs) > 0 and len(cert) > 2:
             # security targets file
             if targets_dir != '':
                 write_file.write('\n\ncd ..\n')
-                write_file.write('mkdir \"{}\"\n'.format(targets_dir))
-                write_file.write('cd \"{}\"\n\n'.format(targets_dir))
+                write_file.write(f'mkdir \"{targets_dir}\"\n')
+                write_file.write(f'cd \"{targets_dir}\"\n\n')
             for cert in download_files_certs:
                 # double %% is necessary to prevent replacement of %2 within script (second argument of script)
                 file_name_short_web = cert[2].replace(' ', '%%20')
                 if file_name_short_web.find(base_url) != -1:
                     # base url already included
                     write_file.write(
-                        'curl \"{}\" -o \"{}\"\n'.format(file_name_short_web, cert[3]))
+                        f'curl \"{file_name_short_web}\" -o \"{cert[3]}\"\n')
                 else:
                     # insert base url
                     write_file.write(
-                        'curl \"{}{}\" -o \"{}\"\n'.format(base_url, file_name_short_web, cert[3]))
-                write_file.write('{} \"{}\"\n\n'.format(
-                    PDF2TEXT_CONVERT, cert[3]))
+                        f'curl \"{base_url}{file_name_short_web}\" -o \"{cert[3]}\"\n')
+                write_file.write(f'{PDF2TEXT_CONVERT} \"{cert[3]}\"\n\n')
 
 
 def download_cc_web(web_dir: Path):
     download_file("https://www.commoncriteriaportal.org/products/", web_dir / "cc_products_active.html")
-    download_file("https://www.commoncriteriaportal.org/products/index.cfm?archived=1", web_dir / "cc_products_archived.html")
+    download_file("https://www.commoncriteriaportal.org/products/index.cfm?archived=1",
+                  web_dir / "cc_products_archived.html")
     download_file("https://www.commoncriteriaportal.org/labs/", web_dir / "cc_labs.html")
-    download_file("https://www.commoncriteriaportal.org/products/certified_products.csv", web_dir / "cc_products_active.csv")
-    download_file("https://www.commoncriteriaportal.org/products/certified_products-archived.csv", web_dir / "cc_products_archived.csv")
+    download_file("https://www.commoncriteriaportal.org/products/certified_products.csv",
+                  web_dir / "cc_products_active.csv")
+    download_file("https://www.commoncriteriaportal.org/products/certified_products-archived.csv",
+                  web_dir / "cc_products_archived.csv")
     download_file("https://www.commoncriteriaportal.org/pps/", web_dir / "cc_pp_active.html")
-    download_file("https://www.commoncriteriaportal.org/pps/collaborativePP.cfm?cpp=1", web_dir / "cc_pp_collaborative.html")
+    download_file("https://www.commoncriteriaportal.org/pps/collaborativePP.cfm?cpp=1",
+                  web_dir / "cc_pp_collaborative.html")
     download_file("https://www.commoncriteriaportal.org/pps/index.cfm?archived=1", web_dir / "cc_pp_archived.html")
     download_file("https://www.commoncriteriaportal.org/pps/pps.csv", web_dir / "cc_pp_active.csv")
     download_file("https://www.commoncriteriaportal.org/pps/pps-archived.csv", web_dir / "cc_pp_archived.csv")
@@ -87,6 +89,7 @@ def download_cc(walk_dir: Path, cert_list, num_threads):
                 download_file(cert[2], walk_dir / "targets" / cert[3])
             else:
                 download_file(CC_WEB_URL + cert[2], walk_dir / "targets" / cert[3])
+
     with tqdm(total=len(cert_list)) as pbar:
         for response in ThreadPool(num_threads).imap(download_one, cert_list):
             pbar.update(1)
@@ -123,14 +126,13 @@ def generate_failed_download_script(base_dir: Path):
                 file_name_short = file_name[file_name.rfind(os.sep) + 1:]
                 # double %% is necessary to prevent replacement of %2 within script (second argument of script)
                 file_name_short_web = file_name_short.replace(' ', '%%20')
-                download_link = '/files/epfiles/{}'.format(file_name_short_web)
+                download_link = f'/files/epfiles/{file_name_short_web}'
                 download_again.append((download_link, file_name))
 
     generate_download_script('download_failed_certs.bat',
                              '', '', CC_WEB_URL, download_again)
-    print('*** Number of files to be re-downloaded again (inside \'{}\'): {}'.format(
-        'download_failed_certs.bat', len(download_again)))
-
+    print(
+        f'*** Number of files to be re-downloaded again (inside \'{"download_failed_certs.bat"}\'): {len(download_again)}')
 
 
 def generate_fips_basic_download_script():
@@ -151,15 +153,14 @@ def generate_fips_download_script(file_name, fips_dir):
 
     with open(file_name, 'w', errors=FILE_ERRORS_STRATEGY) as write_file:
         # make directories for both html and security policies, scraping in one go
-        write_file.write('mkdir {}\n'.format(html_dir))
-        write_file.write('mkdir {}\n\n'.format(sp_dir))
+        write_file.write(f'mkdir {html_dir}\n')
+        write_file.write(f'mkdir {sp_dir}\n\n')
 
+        # upper bound for max certs, in reality there is ~ 3730 certificates
         for cert_id in range(1, 4001):
             write_file.write(
-                'curl "https://csrc.nist.gov/projects/cryptographic-module-validation-program/certificate/{}" -o {}{}.html\n'.format(
-                    cert_id, html_dir, cert_id))
+                f'curl "https://csrc.nist.gov/projects/cryptographic-module-validation-program/certificate/{cert_id}" -o {html_dir}{cert_id}.html\n')
             write_file.write(
-                'curl "https://csrc.nist.gov/CSRC/media/projects/cryptographic-module-validation-program/documents'
-                '/security-policies/140sp{}.pdf" -o {}{}.pdf\n'.format(
-                    cert_id, sp_dir, cert_id))
-            write_file.write("{} {}{}.pdf\n".format(PDF2TEXT_CONVERT, sp_dir, cert_id))
+                f'curl "https://csrc.nist.gov/CSRC/media/projects/cryptographic-module-validation-program/documents'
+                '/security-policies/140sp{cert_id}.pdf" -o {sp_dir}{cert_id}.pdf\n')
+            write_file.write(f"{PDF2TEXT_CONVERT} {sp_dir}{cert_id}.pdf\n")
diff --git a/sec_certs/fips_certificates.py b/sec_certs/fips_certificates.py
index fbaea146..a256336a 100755
--- a/sec_certs/fips_certificates.py
+++ b/sec_certs/fips_certificates.py
@@ -22,15 +22,30 @@ FIPS_MODULE_URL = 'https://csrc.nist.gov/projects/cryptographic-module-validatio
 
 
 def extract_filename(file: str) -> str:
+    """
+    Extracts filename from path
+    @param file: UN*X path
+    :return: filename without last extension
+    """
     return os.path.splitext(os.path.basename(file))[0]
 
 
 def parse_ul(text):
+    """
+    Parses content between <ul> tags in FIPS .html CMVP page
+    :param text: text in <ul> tags
+    :return: all <li> elements
+    """
     p = re.compile(r"<li>(.*?)<\/li>")
     return p.findall(text)
 
 
 def parse_table(text):
+    """
+    Parses content of <table> tags in FIPS .html CMVP page
+    :param text: text in <table> tags
+    :return: list of all found algorithm IDs
+    """
     items_found_all = []
 
     # find <tr>, in that look for "text-nowrap" and look if there is a cert mentioned
@@ -56,18 +71,26 @@ def parse_table(text):
     return items_found_all
 
 
-def parse_algorithms(text, in_table=False):
-    # print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
-    # print()
-    # print(text)
+def parse_algorithms(text, in_pdf=False):
+    """
+    Parses table of FIPS (non) allowed algorithms
+    :param text: Contents of the table
+    :param in_pdf: Specifies whether the table was found in a PDF security policies file
+    :return: list of all found algorithm IDs
+    """
     items_found = []
-    for m in re.finditer(r"(?:#{}\s?|Cert\.?[^. ]*?\s?)(?:[Cc]\s)?(?P<id>\d+)".format('?' if in_table else ''), text):
+    for m in re.finditer(rf"(?:#{'?' if in_pdf else ''}\s?|Cert\.?[^. ]*?\s?)(?:[Cc]\s)?(?P<id>\d+)", text):
         items_found.append({'Certificate': m.group()})
 
     return items_found
 
 
 def parse_caveat(text):
+    """
+    Parses content of "Caveat" of FIPS CMVP .html file
+    :param text: text of "Caveat"
+    :return: list of all found algorithm IDs
+    """
     items_found = []
 
     for m in re.finditer(r"(?:#\s?|Cert\.?(?!.\s)\s?|Certificate\s?)(?P<id>\d+)", text):
@@ -77,6 +100,10 @@ def parse_caveat(text):
 
 
 def initialize_entry(input_dictionary):
+    """
+    Initialize input dictionary with elements that shuold be always processed
+    :param input_dictionary: dictionary used as "all_items"
+    """
     input_dictionary['fips_exceptions'] = []
     input_dictionary['fips_tested_conf'] = []
 
@@ -95,10 +122,10 @@ def fips_search_html(base_dir, output_file, dump_to_file=False):
 
     all_found_items = {}
 
-    for file in sec_certs.files.search_files(base_dir):
+    for file in search_files(base_dir):
         items_found = {}
         initialize_entry(items_found)
-        text = sec_certs.files.load_cert_html_file(file)
+        text = extract_certificates.load_cert_html_file(file)
         filename = os.path.splitext(os.path.basename(file))[0]
         all_found_items[filename] = items_found
         items_found['cert_fips_id'] = filename
@@ -106,7 +133,6 @@ def fips_search_html(base_dir, output_file, dump_to_file=False):
         for rule in RE_FIPS_HTMLS:
             m = re.search(rule, text)
             if m is None:
-                # print("ERROR: For rule {} nothing found in file {}.".format(rule, file))
                 continue
 
             group_dict = m.groupdict()
@@ -154,6 +180,13 @@ def fips_search_html(base_dir, output_file, dump_to_file=False):
 
 
 def get_dot_graph(found_items, output_file_name):
+    """
+    Function that plots .dot graph of dependencies between certificates
+    Certificates with at least one dependency are displayed in "{output_file_name}connections.pdf", remaining
+    certificates are displayed in {output_file_name}single.pdf
+    :param found_items: Dictionary of all found items generated in main()
+    :param output_file_name: prefix to "connections", "connections.pdf", "single" and "single.pdf"
+    """
     dot = Digraph(comment='Certificate ecosystem')
     single_dot = Digraph(comment='Modules with no dependencies')
     single_dot.attr('graph', label='Single nodes', labelloc='t', fontsize='30')
@@ -204,13 +237,18 @@ def get_dot_graph(found_items, output_file_name):
                 dot.edge(key, conn)
                 edges += 1
 
-    print("rendering {} keys and {} edges".format(keys, edges))
+    print(f"rendering {keys} keys and {edges} edges")
 
     dot.render(output_file_name + 'connections', view=True)
     single_dot.render(output_file_name + 'single', view=True)
 
 
 def remove_algorithms_from_extracted_data(items, html):
+    """
+    Function that removes all found certificate IDs that are matching any IDs labeled as algorithm IDs
+    :param items: All keyword items found in pdf files
+    :param html: All items extracted from html files
+    """
     for file_name in items:
         items[file_name]['file_status'] = True
         html[file_name]['file_status'] = True
@@ -235,6 +273,11 @@ def remove_algorithms_from_extracted_data(items, html):
 
 
 def validate_results(items, html):
+    """
+    Function that validates results and finds the final connection output
+    :param items: All keyword items found in pdf files
+    :param html: All items extracted from html files - this is where we store connections
+    """
     broken_files = set()
     for file_name in items:
         for rule in items[file_name]['rules_cert_id']:
@@ -265,9 +308,6 @@ def validate_results(items, html):
                     html[file_name]['Connections'].append(cert_id)
 
 
-count = 0
-
-
 def parse_list_of_tables(txt: str) -> Set[str]:
     """
     Parses list of tables from function find_tables(), finds ones that mention algorithms
@@ -305,13 +345,19 @@ def extract_page_number(txt: str) -> Optional[str]:
 
 
 def find_tables(txt, file_name, num_pages):
-    global count
-
+    """
+    Function that tries to pages in security policy pdf files, where it's possible to find a table containing
+    algorithms
+    :param txt: file in .txt format (output of pdftotext)
+    :param file_name: name of the file
+    :param num_pages: number of pages in pdf
+    :return:    list of pages possibly containing a table
+                None if these cannot be found
+    """
     # Look for "List of Tables", where we can find exactly tables with page num
     tables_regex = re.compile(r"^(?:(?:[Tt]able\s|[Ll]ist\s)(?:[Oo]f\s))[Tt]ables[\s\S]+?\f", re.MULTILINE)
     table = tables_regex.search(txt)
     if table:
-        count += 1
         rb = parse_list_of_tables(table.group())
         if rb:
             return list(rb)
@@ -319,8 +365,9 @@ def find_tables(txt, file_name, num_pages):
 
     # Otherwise look for "Table" in text and \f representing footer, then extract page number from footer
     print("~" * 20, file_name, '~' * 20)
-    footer_regex = re.compile(r"(?:Table[^\f]*)(?P<first>^[\S\t ]*$)\n(?P<second>(\f[ \t\S]+)$)(?P<third>\n^[ \t\S]+?$)?",
-                              re.MULTILINE)
+    footer_regex = re.compile(
+        r"(?:Table[^\f]*)(?P<first>^[\S\t ]*$)\n(?P<second>(\f[ \t\S]+)$)(?P<third>\n^[ \t\S]+?$)?",
+        re.MULTILINE)
 
     # We have 2 groups, one is optional - trying to parse 2 lines (just in case)
     footer1 = [m.group('first') for m in footer_regex.finditer(txt)]
@@ -331,7 +378,8 @@ def find_tables(txt, file_name, num_pages):
     #     footer2 += [''] * (len(footer1) - len(footer2))
 
     # zipping them together
-    footer_complete = [m[0] + m[1] + m[2] for m in zip(footer1, footer2, footer3) if m[0] is not None and m[1] is not None and m[2] is not None]
+    footer_complete = [m[0] + m[1] + m[2] for m in zip(footer1, footer2, footer3) if
+                       m[0] is not None and m[1] is not None and m[2] is not None]
 
     # removing None and duplicates
     footers = [extract_page_number(x) for x in footer_complete]
@@ -342,12 +390,24 @@ def find_tables(txt, file_name, num_pages):
 
 
 def repair_pdf_page_count(file: str) -> int:
+    """
+    Some pdfs can't be opened by PyPDF2 - opening them with pikepdf and then saving them fixes this issue.
+    By opening this file in a pdf reader, we can already extract number of pages
+    :param file: file name
+    :return: number of pages in pdf file
+    """
     pdf = pikepdf.Pdf.open(file, allow_overwriting_input=True)
     pdf.save(file)
     return len(pdf.pages)
 
 
 def extract_certs_from_tables(list_of_files, html_items):
+    """
+    Function that extracts algorithm IDs from tables in security policies files.
+    :param list_of_files: iterable containing all files to parse
+    :param html_items: dictionary created by main() containing data extracted from html pages
+    :return: list of files that couldn't have been decoded
+    """
     not_decoded = []
     for cert_file in list_of_files:
         if '.txt' not in cert_file:
@@ -439,7 +499,6 @@ def main(directory):
     get_dot_graph(html, 'output')
     end = time.time()
     print("TIME:", end - start)
-    print("COUNT:", count)
 
 
 if __name__ == '__main__':