added method for creating graphs, minor fixes

4 files changed, 97 insertions, 98 deletions

diff --git a/fips_certificates.py b/fips_certificates.py
index 1d5e3232..eb147d48 100755
--- a/fips_certificates.py
+++ b/fips_certificates.py
@@ -33,70 +33,6 @@ def initialize_entry(current_items_found):
     pass
 
 
-def get_dot_graph(found_items: Dict, output_file_name: str):
-    """
-    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')
-    single_dot.attr('node', style='filled')
-    dot.attr('graph', label='Dependencies', labelloc='t', fontsize='30')
-    dot.attr('node', style='filled')
-
-    def found_interesting_cert(current_key):
-        if found_items[current_key]['fips_vendor'] == highlighted_vendor:
-            dot.attr('node', color='red')
-            if found_items[current_key]['fips_status'] == 'Revoked':
-                dot.attr('node', color='grey32')
-            if found_items[current_key]['fips_status'] == 'Historical':
-                dot.attr('node', color='gold3')
-        if found_items[current_key]['fips_vendor'] == "SUSE, LLC":
-            dot.attr('node', color='lightblue')
-
-    def color_check(current_key):
-        dot.attr('node', color='lightgreen')
-        if found_items[current_key]['fips_status'] == 'Revoked':
-            dot.attr('node', color='lightgrey')
-        if found_items[current_key]['fips_status'] == 'Historical':
-            dot.attr('node', color='gold')
-        found_interesting_cert(current_key)
-        dot.node(current_key, label=current_key + '\n' + found_items[current_key]['fips_vendor'] +
-                                    ('\n' + found_items[current_key]['fips_module_name']
-                                     if 'fips_module_name' in found_items[current_key] else ''))
-
-    keys = 0
-    edges = 0
-
-    highlighted_vendor = 'Red Hat®, Inc.'
-    for key in found_items:
-        if key != 'Not found' and found_items[key]['file_status']:
-            if found_items[key]['Connections']:
-                color_check(key)
-                keys += 1
-            else:
-                single_dot.attr('node', color='lightblue')
-                found_interesting_cert(key)
-                single_dot.node(key, label=key + '\n' + found_items[key]['fips_vendor'] + ('\n' + found_items[key][
-                    'fips_module_name'] if 'fips_module_name' in found_items[key] else ''))
-
-    for key in found_items:
-        if key != 'Not found' and found_items[key]['file_status']:
-            for conn in found_items[key]['Connections']:
-                color_check(conn)
-                dot.edge(key, conn)
-                edges += 1
-
-    print(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)
-
-
 def remove_algorithms_from_extracted_data(items, html):
     pass
 
@@ -267,7 +203,7 @@ def main(directory, do_download_meta: bool, do_download_certs: bool, threads: in
     with open(results_dir / 'fips_html_all.json', 'w') as f:
         json.dump(html, f, indent=4, sort_keys=True)
     print("PLOTTING GRAPH")
-    get_dot_graph(html, results_dir / 'output')
+    # get_dot_graph(html, results_dir / 'output')
     end = time.time()
     print("TIME:", end - start)
 
diff --git a/fips_oop_demo.py b/fips_oop_demo.py
index 0b85b4d3..61a90028 100644
--- a/fips_oop_demo.py
+++ b/fips_oop_demo.py
@@ -2,6 +2,7 @@ from sec_certs.dataset import FIPSDataset
 from pathlib import Path
 from datetime import datetime
 import logging
+from sec_certs.helpers import download_parallel
 
 
 def main():
@@ -13,36 +14,35 @@ def main():
 
     # Load metadata for certificates from CSV and HTML sources
     dset.get_certs_from_web()
-    algs = dset.present_algorithms()
-    print(algs)
-    a = {
-        'CVL', 'AES', 'DRBG', 'HMAC', 'SHS', 'Component', 'ECDSA', 'RSA', 'DSA', 'C'
 
-    }
-    # 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(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')
     end = datetime.now()
     logging.info(f'The computation took {(end - start)} seconds.')
 
diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index f5fb3f83..744e56c5 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -489,7 +489,7 @@ class CommonCriteriaCert(Certificate):
 
             report_link = CommonCriteriaCert.cc_url + links[1].get('href')
             security_target_link = CommonCriteriaCert.cc_url + \
-                links[2].get('href')
+                                   links[2].get('href')
 
             return report_link, security_target_link
 
@@ -518,10 +518,10 @@ class CommonCriteriaCert(Certificate):
                 for l in links:
                     if l.get('title').startswith('Maintenance Report:'):
                         main_report_link = CommonCriteriaCert.cc_url + \
-                            l.get('href')
+                                           l.get('href')
                     elif l.get('title').startswith('Maintenance ST'):
                         main_st_link = CommonCriteriaCert.cc_url + \
-                            l.get('href')
+                                       l.get('href')
                     else:
                         logging.error('Unknown link in Maintenance part!')
                 maintainance_updates.add(
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 39927412..ca57b146 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -3,7 +3,7 @@ import re
 from datetime import datetime
 import locale
 import logging
-from typing import Dict, List, ClassVar
+from typing import Dict, List, ClassVar, Set
 import json
 from importlib import import_module
 
@@ -12,6 +12,7 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 import shutil
 
+from graphviz import Digraph
 from tabula import read_pdf
 import pandas as pd
 from bs4 import BeautifulSoup
@@ -261,7 +262,7 @@ class CCDataset(Dataset):
             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
+                                 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']:
@@ -562,3 +563,65 @@ class FIPSDataset(Dataset):
                     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
+        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 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')
+        single_dot.attr('node', style='filled')
+        dot.attr('graph', label='Dependencies', labelloc='t', fontsize='30')
+        dot.attr('node', style='filled')
+
+        def found_interesting_cert(current_key):
+            if self.certs[current_key].vendor == highlighted_vendor:
+                dot.attr('node', color='red')
+                if self.certs[current_key].status == 'Revoked':
+                    dot.attr('node', color='grey32')
+                if self.certs[current_key].status == 'Historical':
+                    dot.attr('node', color='gold3')
+            if self.certs[current_key].vendor == "SUSE, LLC":
+                dot.attr('node', color='lightblue')
+
+        def color_check(current_key):
+            dot.attr('node', color='lightgreen')
+            if self.certs[current_key].status == 'Revoked':
+                dot.attr('node', color='lightgrey')
+            if self.certs[current_key].status == 'Historical':
+                dot.attr('node', color='gold')
+            found_interesting_cert(current_key)
+            dot.node(current_key, label=current_key + '\n' + self.certs[current_key].vendor +
+                                        ('\n' + self.certs[current_key].module_name
+                                         if self.certs[current_key].module_name else ''))
+
+        keys = 0
+        edges = 0
+
+        highlighted_vendor = 'Red Hat®, Inc.'
+        for key in self.certs:
+            if key != 'Not found' and self.certs[key].file_status:
+                if self.certs[key].connections:
+                    color_check(key)
+                    keys += 1
+                else:
+                    single_dot.attr('node', color='lightblue')
+                    found_interesting_cert(key)
+                    single_dot.node(key, label=key + '\n' + self.certs[key].vendor + (
+                        '\n' + self.certs[key].module_name if self.certs[key].module_name else ''))
+
+        for key in self.certs:
+            if key != 'Not found' and self.certs[key].file_status:
+                for conn in self.certs[key].connections:
+                    color_check(conn)
+                    dot.edge(key, conn)
+                    edges += 1
+
+        print(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)