merge fresh main

6 files changed, 149 insertions, 11 deletions

diff --git a/src/sec_certs/utils/extract.py b/src/sec_certs/utils/extract.py
index 933db4b8..e6312b8d 100644
--- a/src/sec_certs/utils/extract.py
+++ b/src/sec_certs/utils/extract.py
@@ -353,7 +353,7 @@ def search_only_headers_bsi(filepath: Path):  # noqa: C901
 
             for m in re.finditer(rule_and_sep, whole_text):
                 # check if previous rules had at least one match
-                if constants.TAG_CERT_ID not in items_found.keys():
+                if constants.TAG_CERT_ID not in items_found:
                     logger.error(f"ERROR: front page not found for file: {filepath}")
 
                 match_groups = m.groups()
diff --git a/src/sec_certs/utils/helpers.py b/src/sec_certs/utils/helpers.py
index 596ecf62..c103a2d3 100644
--- a/src/sec_certs/utils/helpers.py
+++ b/src/sec_certs/utils/helpers.py
@@ -28,7 +28,11 @@ def download_file(
         time.sleep(delay)
         # See https://github.com/psf/requests/issues/3953 for header justification
         r = requests.get(
-            url, allow_redirects=True, timeout=constants.REQUEST_TIMEOUT, stream=True, headers={"Accept-Encoding": None}  # type: ignore
+            url,
+            allow_redirects=True,
+            timeout=constants.REQUEST_TIMEOUT,
+            stream=True,
+            headers={"Accept-Encoding": None},  # type: ignore
         )
         ctx: Any
         if show_progress_bar:
diff --git a/src/sec_certs/utils/pandas.py b/src/sec_certs/utils/pandas.py
index 5b0a668d..661b669c 100644
--- a/src/sec_certs/utils/pandas.py
+++ b/src/sec_certs/utils/pandas.py
@@ -141,7 +141,7 @@ def get_sar_level_from_set(sars: set[SAR], sar_family: str) -> int | None:
     Given a set of SARs and a family name, will return level of the seeked SAR from the set.
     """
     family_sars_dict = {x.family: x for x in sars} if (sars and not pd.isnull(sars)) else {}
-    if sar_family not in family_sars_dict.keys():
+    if sar_family not in family_sars_dict:
         return None
     return family_sars_dict[sar_family].level
 
@@ -211,7 +211,7 @@ def compute_cve_correlations(
     tuples = list(
         zip(n_cves_corrs, n_cves_pvalues, worst_cve_corrs, worst_cve_pvalues, avg_cve_corrs, avg_cve_pvalues, supports)
     )
-    dct = {family: correlations for family, correlations in zip(["eal"] + families, tuples)}
+    dct = dict(zip(["eal"] + families, tuples))
     df_corr = pd.DataFrame.from_dict(
         dct,
         orient="index",
diff --git a/src/sec_certs/utils/pdf.py b/src/sec_certs/utils/pdf.py
index 749a8a5a..f2c2c58e 100644
--- a/src/sec_certs/utils/pdf.py
+++ b/src/sec_certs/utils/pdf.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import glob
 import logging
 import subprocess
 from datetime import datetime, timedelta, timezone
@@ -11,6 +10,8 @@ from typing import Any
 
 import pdftotext
 import pikepdf
+import pytesseract
+from PIL import Image
 
 from sec_certs import constants
 from sec_certs.constants import (
@@ -51,13 +52,16 @@ def ocr_pdf_file(pdf_path: Path) -> str:
         )
         if ppm.returncode != 0:
             raise ValueError(f"pdftoppm failed: {ppm.returncode}")
-        for ppm_path in map(Path, glob.glob(str(tmppath / "image*.ppm"))):
+
+        for ppm_path in tmppath.rglob("image*.ppm"):
             base = ppm_path.with_suffix("")
-            tes = subprocess.run(
-                ["tesseract", "-l", "eng+deu+fra", ppm_path, base], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
-            )
-            if tes.returncode != 0:
-                raise ValueError(f"tesseract failed: {tes.returncode}")
+            content = pytesseract.image_to_string(Image.open(ppm_path), lang="eng+deu+fra")
+
+            if content:
+                with Path(base.with_suffix(".txt")).open("w") as file:
+                    file.write(content)
+            else:
+                raise ValueError(f"OCR failed for document {ppm_path}. Check document manually")
 
         contents = ""
 
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
diff --git a/src/sec_certs/utils/profiling.py b/src/sec_certs/utils/profiling.py
new file mode 100644
index 00000000..7da67070
--- /dev/null
+++ b/src/sec_certs/utils/profiling.py
@@ -0,0 +1,45 @@
+import gc
+from contextlib import contextmanager
+from datetime import datetime
+from functools import wraps
+from logging import Logger
+
+import psutil
+
+
+@contextmanager
+def log_stage(logger: Logger, msg: str, collect_garbage: bool = False):
+    """Contextmanager that logs a message to the logger when it is entered and exited.
+    The message has debug information about memory use. Optionally, it can
+    run garbage collection when exiting.
+    """
+    meminfo = psutil.Process().memory_full_info()
+    logger.info(f">> Starting >> {msg}")
+    logger.debug(str(meminfo))
+    start_time = datetime.now()
+
+    try:
+        yield
+    finally:
+        end_time = datetime.now()
+        duration = end_time - start_time
+        meminfo = psutil.Process().memory_full_info()
+        logger.info(f"<< Finished << {msg} ({duration})")
+        logger.debug(str(meminfo))
+
+        if collect_garbage:
+            gc.collect()
+
+
+def staged(logger: Logger, log_message: str, collect_garbage: bool = False):
+    """Like log_stage but a decorator."""
+
+    def deco(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with log_stage(logger, log_message, collect_garbage):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return deco