Improved device selection logic

flows/prepPolicy.py

@@ -16,6 +16,7 @@
 import logging
 import os
 import os.path
+import re
 from typing import List
 
 import dotenv
@@ -26,11 +27,7 @@ from models.policy import Allowlist, Policy
 from services.API import AirlockAPIWrapper
 from utils.configmanager import get_protected_value, load_env, load_env_json
 from utils.selector import Selector
-from utils.utils import (
-    colorText,
-    formatHTML,
-    regulator,
-)
+from utils.utils import colorText, formatHTML, print_x_wide, regulator
 
 logger = logging.getLogger(__name__)
 
@@ -105,6 +102,9 @@ def sortHashes(
 
  
     for label, records in categories.items():
+        if not records:
+                continue  # Skip empty or falsy categories
+
         csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
         html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html"
 
@@ -369,3 +369,48 @@ def calculatePath(approved_hashes, path_exclusion_constant, split):
     pathExclusions = pd.concat(processed_dfs, ignore_index=True)
 
     return pathExclusions
+
+def testChange(selected_policies, destination_policy, destination_allowlist):
+                working_dir = load_env("WORKING_DIR")
+
+                logger.info("These path exclusions would be added to:")
+                logger.info(destination_policy)
+
+                pathexclusions = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv")
+                hashes = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv")
+
+                unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates()
+
+                drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
+                processed_paths = [
+                    (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
+                    for path, ext in unique_combinations.itertuples(index=False, name=None)
+                ]
+
+                for path in processed_paths:
+                    logger.info(path)
+
+                print(colorText("These publishers would added", "yellow"))
+                processed_publishers = []
+                if os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"):
+                    publishers = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv")
+                    if publishers.empty:
+                        print(colorText("The publishers list is empty.", "red"))
+                    else:
+                        processed_publishers = (
+                            publishers[publishers["publisher"] != "Not Signed"]
+                            ["publisher"]
+                            .drop_duplicates()
+                            .tolist()
+                        )
+                        for publisher in processed_publishers:
+                            print(publisher)
+
+                print(colorText("These hashes would be added to:", "yellow"))
+                print(destination_allowlist)
+
+                processed_hashes = hashes["sha256"].unique().tolist()
+                print_x_wide(processed_hashes, 3)
+
+                return processed_paths, processed_hashes, processed_publishers
+