Merge pull request #351 from crocs-muni/feat-fips-references

Analysis of FIPS references

3 files changed, 95 insertions, 6 deletions

diff --git a/src/sec_certs/dataset/dataset.py b/src/sec_certs/dataset/dataset.py
index 6466630f..9e49a347 100644
--- a/src/sec_certs/dataset/dataset.py
+++ b/src/sec_certs/dataset/dataset.py
@@ -474,6 +474,10 @@ class Dataset(Generic[CertSubType, AuxiliaryDatasetsSubType], ComplexSerializabl
             self.auxiliary_datasets.cpe_dset = self._prepare_cpe_dataset()
 
         clf = CPEClassifier(config.cpe_matching_threshold, config.cpe_n_max_matches)
+
+        if self.auxiliary_datasets.cpe_dset is None:
+            raise ValueError("CPE dataset cannot be None")
+
         clf.fit([x for x in self.auxiliary_datasets.cpe_dset if filter_condition(x)])
 
         cert: CertSubType
@@ -590,6 +594,9 @@ class Dataset(Generic[CertSubType, AuxiliaryDatasetsSubType], ComplexSerializabl
         if not self.auxiliary_datasets.cve_dset:
             self.auxiliary_datasets.cve_dset = self._prepare_cve_dataset()
 
+        if self.auxiliary_datasets.cve_dset is None:
+            raise ValueError("CVE dataset cannot be None")
+
         if not self.auxiliary_datasets.cve_dset.look_up_dicts_built:
             cpe_match_dict = self._prepare_cpe_match_dict()
             all_cpes = self._get_all_cpes_in_dataset()
diff --git a/src/sec_certs/sample/fips.py b/src/sec_certs/sample/fips.py
index d351c3ce..26510191 100644
--- a/src/sec_certs/sample/fips.py
+++ b/src/sec_certs/sample/fips.py
@@ -83,15 +83,13 @@ class FIPSHTMLParser:
     def _build_vendor_dict(vendor_div: Tag) -> dict[str, Any]:
         if not (link := vendor_div.find("a")):
             return {"vendor_url": None, "vendor": list(vendor_div.find("div", "panel-body").children)[0].strip()}
-        else:
-            return {"vendor_url": link.get("href"), "vendor": link.text.strip()}
+        return {"vendor_url": link.get("href"), "vendor": link.text.strip()}
 
     @staticmethod
     def _build_related_files_dict(related_files_div: Tag) -> dict[str, Any]:
         if cert_link := [x for x in related_files_div.find_all("a") if "Certificate" in x.text]:
             return {"certificate_pdf_url": constants.FIPS_BASE_URL + cert_link[0].get("href")}
-        else:
-            return {"certificate_pdf_url": None}
+        return {"certificate_pdf_url": None}
 
     @staticmethod
     def _build_validation_history_dict(validation_history_div: Tag) -> dict[str, Any]:
@@ -433,8 +431,7 @@ class FIPSCertificate(
                 fips_certlike = self.keywords["fips_certlike"].get("Certlike", {})
                 matches = {re.search(r"#\s{0,1}\d{1,4}", x) for x in fips_certlike}
                 return {"".join([x for x in match.group() if x.isdigit()]) for match in matches if match}
-            else:
-                return set()
+            return set()
 
     @dataclass(eq=True)
     class Heuristics(BaseHeuristics, ComplexSerializableType):
diff --git a/src/sec_certs/utils/plot_utils.py b/src/sec_certs/utils/plot_utils.py
new file mode 100644
index 00000000..b9e7ac7c
--- /dev/null
+++ b/src/sec_certs/utils/plot_utils.py
@@ -0,0 +1,85 @@
+from typing import Dict, List, Tuple
+
+from networkx import DiGraph
+from pandas import DataFrame
+
+
+def get_cert_property(df: DataFrame, cert_id: int, column: str) -> str:
+    if column not in df.columns:
+        raise ValueError(f"Dataset does not have column '{column}'")
+
+    sub_df = df[df["cert_id"] == int(cert_id)]
+
+    if not sub_df.shape[0]:  # Certificate is not in the dataset
+        raise ValueError(f"Cert ID: {cert_id} not in dataset")
+
+    if sub_df.shape[0] > 1:  # There are more than one occurence with same ID
+        raise ValueError(f"Error Cert ID: {cert_id} has {sub_df.shape[0]} occurrences.")
+
+    return sub_df.iloc[0][column]
+
+
+def get_fips_cert_references_graph(
+    df: DataFrame, cert_id: int, colour_mapper: Dict[str, str]
+) -> Tuple[DiGraph, List[str]]:
+    if cert_id not in df["cert_id"].unique():
+        raise ValueError(f"Cert ID: {cert_id} is not in the dataset")
+
+    cert_id_series = df[df["cert_id"] == cert_id].iloc[0]
+    colour_map = [colour_mapper["chosen_cert_colour"]]
+    graph = DiGraph()
+    graph.add_node(cert_id)
+
+    # Display which certificates are directly referenced by the chosen certificate
+    for referenced_cert_id in cert_id_series["module_directly_referencing"]:
+        graph.add_node(referenced_cert_id)
+        graph.add_edge(cert_id, referenced_cert_id)
+        colour_map.append(colour_mapper["referencing_colour"])
+
+    # Display which certificates are directly referencing the chosen certificate
+    for referencing_cert_id in cert_id_series["module_directly_referenced_by"]:
+        graph.add_node(referencing_cert_id)
+        graph.add_edge(referencing_cert_id, cert_id)
+        colour_map.append(colour_mapper["referenced_colour"])
+
+    return graph, colour_map
+
+
+def get_most_referenced_cert_graph(df: DataFrame, status_colour_mapper: Dict[str, str]) -> Tuple[DiGraph, List[str]]:
+    graph = DiGraph()
+    colour_map = []
+    max_referenced_by_num = df["incoming_direct_references_count"].max()
+    most_referenced_certificate = df[df["incoming_direct_references_count"] == max_referenced_by_num].iloc[0]
+
+    origin_cert_id: int = most_referenced_certificate["cert_id"]
+    origin_cert_status: str = most_referenced_certificate["status"]
+    graph.add_node(origin_cert_id)
+    colour_map.append(status_colour_mapper[origin_cert_status])
+
+    for cert_id_str in most_referenced_certificate["module_directly_referenced_by"]:
+        cert_id_int = int(cert_id_str)
+        graph.add_node(cert_id_int)
+        graph.add_edge(cert_id_int, origin_cert_id)
+        cert_status: str = get_cert_property(df, cert_id_int, "status")
+        colour_map.append(status_colour_mapper[cert_status])
+
+    return graph, colour_map
+
+
+def get_most_referencing_cert_graph(df: DataFrame, status_colour_mapper: Dict[str, str]) -> Tuple[DiGraph, List[str]]:
+    graph = DiGraph()
+    colour_map = []
+    max_referencing_num = df["outgoing_direct_references_count"].max()
+    most_referencing_cert = df[df["outgoing_direct_references_count"] == max_referencing_num].iloc[0]
+    origin_cert_id = most_referencing_cert["cert_id"]
+    origin_cert_status = most_referencing_cert["status"]
+    colour_map.append(status_colour_mapper[origin_cert_status])
+
+    for cert_id_str in most_referencing_cert["module_directly_referencing"]:
+        cert_id_int = int(cert_id_str)
+        graph.add_node(cert_id_int)
+        graph.add_edge(origin_cert_id, cert_id_int)
+        cert_status: str = get_cert_property(df, cert_id_int, "status")
+        colour_map.append(status_colour_mapper[cert_status])
+
+    return graph, colour_map