Improved device selection logic

utils/menus.py

@@ -15,10 +15,8 @@
 
 import logging
 import os
-import re
 
 import dotenv
-import pandas as pd
 
 import services.policyhandler as policyh
 from flows.otp import generate, otp_activities_by_agent, revoke
@@ -28,6 +26,7 @@ from flows.prepPolicy import (
     selectAllowlists,
     selectPolicies,
     sortHashes,
+    testChange,
 )
 from flows.quietAgent import findQuietAgents
 from services.agenthandler import findAgents, moveAgentToRelatedPolicy, selectAgents
@@ -39,9 +38,9 @@ from utils.utils import (
     colorText,
     displayIntro,
     get_sanitized_input,
+    locked,
     open_directory,
     printEnforceChecklist,
-    locked
 )
 
 logger = logging.getLogger(__name__)
@@ -107,7 +106,6 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7
     processed_paths = []
     processed_hashes = []
     processed_publishers = []
-    tested = False
     working_dir = load_env("WORKING_DIR")
 
     while True:
@@ -154,47 +152,7 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7
                 and destination_policy
                 and destination_allowlist
             ):
-                print(colorText("These path exclusions would be added to:", "yellow"))
-                print(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:
-                    print(path)
-
-                print(colorText("These publishers would added", "yellow"))
-
-                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_three_wide(processed_hashes)
-
-                if processed_paths and processed_hashes:
-                    tested = True
+                processed_paths, processed_hashes, processed_publishers = testChange(selected_policies, destination_policy, destination_allowlist)
             else:
                 # Log which condition(s) failed
                 missing_items = []
@@ -215,7 +173,9 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7
             areYouSure()
             confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
             if (
-                tested
+                processed_paths
+                and processed_hashes
+                and processed_publishers
                 and destination_policy
                 and destination_allowlist
                 and confirmation.strip() == "I AGREE"
@@ -230,8 +190,8 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7
                 
             else:
                 logger.error("Confirmation block failed. Reasons:")
-                if not tested:
-                    logger.error(" - Preflight checks were not completed successfully (`tested` is False).")
+                if not processed_publishers or processed_hashes or processed_paths:
+                    logger.error(" - Test not performed.")
                 if not destination_policy:
                     logger.error(" - `destination_policy` is missing or invalid.")
                 if not destination_allowlist:
@@ -321,7 +281,3 @@ def menu_settings():
         else:
             print("Invalid choice. Please try again.")
 
-def print_three_wide(items):
-    for i in range(0, len(items), 3):
-        row = items[i:i+3]
-        print(" | ".join(row))