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 (
+from utils.utils import colorText, formatHTML, print_x_wide, regulator
-    colorText,
-    formatHTML,
-    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
+                    

services/agenthandler.py

@@ -130,24 +130,16 @@ def findAgents(api, return_dataframe):
     else:
         logging.debug("User declined to export the DataFrame.")
 
-
+def collect_device_names() -> List[str]:
-def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
     print(colorText("🔍 Device Search", "cyan"))
     print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan"))
     print(colorText("When you're done, press Enter twice (Three times if you have a single device).\n", "cyan"))
     print(colorText("Example:", "cyan"))
-    print(colorText("H00000", "cyan"))
+    print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan"))
-    print(colorText("UTN00000", "cyan"))
-    print(colorText("i-hSuperSecretServer", "cyan"))
-    print(colorText("u-hVenderBroke\n", "cyan"))
-
     print(colorText("Paste or type your device names below:", "white"))
-    policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
 
     device_input_lines = []
     empty_line_count = 0
-
-    # Regex to validate each line
     valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$')
 
     while True:
@@ -158,49 +150,97 @@ def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
             empty_line_count += 1
             if empty_line_count == 2:
                 break
-            continue  # Don't validate empty lines
+            continue
         else:
             empty_line_count = 0
 
-        # Validate only non-empty lines
         if valid_line_pattern.match(stripped_line):
             device_input_lines.append(stripped_line)
         else:
             print(colorText(f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", "yellow"))
 
-    device_names = [name for name in device_input_lines if name]
+    return [name for name in device_input_lines if name]
+
+
+def choose_match_type() -> bool:
+    print(colorText("Use exact match? (Y for exact, N for fuzzy):", "white"))
+    return get_sanitized_input("").strip().lower() in ["y", "yes"]
+
+
+def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool) -> List['Agent']:
+    if use_exact:
+        return [
+            agent for agent in agents
+            if agent.hostname.lower() in [name.lower() for name in device_names]
+        ]
+    else:
+        pattern = "|".join(map(re.escape, device_names))
+        regex = re.compile(pattern, re.IGNORECASE)
+        return [agent for agent in agents if regex.search(agent.hostname)]
+
+
+def show_unmatched(device_names: List[str], matched_agents: List['Agent'], use_exact: bool):
+    if use_exact:
+        unmatched = [name for name in device_names if not any(agent.hostname.lower() == name.lower() for agent in matched_agents)]
+    else:
+        unmatched = [name for name in device_names if not any(re.search(re.escape(name), agent.hostname, re.IGNORECASE) for agent in matched_agents)]
+
+    if unmatched:
+        logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
+        print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
+
+
+def enrich_agents(agents: List['Agent'], policies: List['Policy']):
+    for agent in agents:
+        agent.enrich_with_policies(policies)
+
+
+def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']:
+    device_names = collect_device_names()
     if not device_names:
         logger.debug("No device names entered")
         print(colorText("⚠️ No device names entered.", "red"))
         return []
 
-    # Build regex pattern to match hostnames
+    use_exact = choose_match_type()
-    pattern = "|".join(map(re.escape, device_names))
-    regex = re.compile(pattern, re.IGNORECASE)
 
-    # Fetch agents
+    policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
     agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()]
-    matched_agents = [agent for agent in agents if regex.search(agent.hostname)]
+    matched_agents = match_agents(device_names, agents, use_exact)
     matched_agents.sort(key=lambda agent: agent.hostname.lower())
 
-    # Show unmatched
+    show_unmatched(device_names, matched_agents, use_exact)
-    unmatched = [name for name in device_names if not any(regex.search(agent.hostname) for agent in agents)]
-    if unmatched:
-        logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
-        print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
 
     if not matched_agents:
         logger.debug("❌ No matching devices found.")
         print(colorText("❌ No matching devices found.", "red"))
-    else:
+        return []
-        logger.debug(f"✅ Found {len(matched_agents)} matching device(s).")        
+
     print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green"))
+    logger.info("Matched agent hostnames:")
+    rows = (len(matched_agents) + 2) // 3  # 3 columns
+    for row in range(rows):
+        line = ""
+        for col in range(3):
+            idx = row + col * rows
+            if idx < len(matched_agents):
+                line += f"{matched_agents[idx].hostname:<30}"
+        logger.info(line)
         
-    # Enrich each agent using its class method
+    matched_agents = Selector.select_with_mode(
-    for agent in matched_agents:
+        matched_agents,
-        agent.enrich_with_policies(policies)
+        label_func=lambda agent: agent.hostname,
+        header="Matched Devices:"
+    )
 
+    if not matched_agents:
+        logger.debug("❌ No matching devices remain after refinement.")
+        print(colorText("❌ No matching devices remain after refinement.", "red"))
+        return []
+
+    enrich_agents(matched_agents, policies)
     return matched_agents
+
 def moveAgentToRelatedPolicy(
     api: AirlockAPIWrapper,
     agent: Agent,

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"))
+                processed_paths, processed_hashes, processed_publishers = testChange(selected_policies, destination_policy, destination_allowlist)
-                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
             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:
+                if not processed_publishers or processed_hashes or processed_paths:
-                    logger.error(" - Preflight checks were not completed successfully (`tested` is False).")
+                    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))

utils/selector.py

@@ -1,28 +1,16 @@
-#    Copyright (C) 2025 James Brotosky, Brandon Wickline
-#
-#    This program is free software: you can redistribute it and/or modify
-#    it under the terms of the GNU Affero General Public License as published
-#    by the Free Software Foundation, either version 3 of the License, or
-#    (at your option) any later version.
-#
-#    This program is distributed in the hope that it will be useful,
-#    but WITHOUT ANY WARRANTY; without even the implied warranty of
-#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-#    GNU Affero General Public License for more details.
-#
-#    You should have received a copy of the GNU Affero General Public License
-#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-
-
 import logging
 from typing import Any, Callable, List, Optional, Union
 
-from utils.utils import get_sanitized_input
+from utils.utils import colorText, get_sanitized_input
 
 logger = logging.getLogger(__name__)
 
 
 class Selector:
+    @staticmethod
+    def _get_sorted_items(items: List[Any], label_func: Callable[[Any], str]) -> List[Any]:
+        return sorted(items, key=lambda item: label_func(item).lower())
+
     @staticmethod
     def _display_choices(
         items: List[Any],
@@ -30,18 +18,54 @@ class Selector:
         num_columns: int = 4,
         header: str = "Available Choices:"
     ) -> None:
-        sorted_items = sorted(items, key=lambda item: label_func(item).lower())
+        rows = (len(items) + num_columns - 1) // num_columns
-        rows = (len(sorted_items) + num_columns - 1) // num_columns
         print(f"\n{header}")
         for row in range(rows):
             line = ""
             for col in range(num_columns):
                 idx = row + col * rows
-                if idx < len(sorted_items):
+                if idx < len(items):
-                    label = label_func(sorted_items[idx])
+                    label = label_func(items[idx])
                     line += f"{idx + 1}: {label:<30}"
             print(line)
 
+    @staticmethod
+    def _display_selected_items(
+        selected: List[Any],
+        label_func: Callable[[Any], str],
+        num_columns: int = 4
+    ) -> None:
+        print(colorText("\nCurrent selections:", "cyan"))
+        if not selected:
+            print(" (none)")
+            return
+
+        sorted_selected = sorted(selected, key=lambda item: label_func(item).lower())
+        rows = (len(sorted_selected) + num_columns - 1) // num_columns
+        for row in range(rows):
+            line = ""
+            for col in range(num_columns):
+                idx = row + col * rows
+                if idx < len(sorted_selected):
+                    label = label_func(sorted_selected[idx])
+                    line += f"{label:<30}"
+            print(line)
+
+    @staticmethod
+    def _parse_selection_input(input_str: str, max_index: int) -> List[int]:
+        selections = []
+        for part in input_str.split(","):
+            part = part.strip()
+            if "-" in part:
+                try:
+                    start, end = map(int, part.split("-"))
+                    selections.extend(range(start, end + 1))
+                except ValueError:
+                    continue
+            elif part.isdigit():
+                selections.append(int(part))
+        return [i for i in selections if 1 <= i <= max_index]
+
     @staticmethod
     def _select_from_list(
         items: List[Any],
@@ -54,35 +78,48 @@ class Selector:
             logger.warning("No items available for selection.")
             return None
 
-        Selector._display_choices(items, label_func, header=header)
+        full_sorted_items = Selector._get_sorted_items(items, label_func)
-        sorted_items = sorted(items, key=lambda item: label_func(item).lower())
+        remaining_items = full_sorted_items.copy()
         selected = []
 
         if allow_multiple:
             while True:
-                choice = get_sanitized_input("Select an item by number (or Q to finish): ").strip().lower()
+                Selector._display_choices(remaining_items, label_func, header=header)
+                Selector._display_selected_items(selected, label_func)
+
+                choice = get_sanitized_input("Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: ").strip().lower()
+
                 if choice == "q":
                     break
-                try:
+                elif choice == "r":
-                    index = int(choice)
+                    selected.clear()
-                    if 1 <= index <= len(sorted_items):
+                    remaining_items = full_sorted_items.copy()
-                        item = sorted_items[index - 1]
+                    print(colorText("🔄 Selections reset.", "yellow"))
+                    continue
+
+                indices = Selector._parse_selection_input(choice, len(remaining_items))
+                newly_selected = []
+
+                for index in indices:
+                    item = remaining_items[index - 1]
                     if item not in selected:
                         selected.append(item)
+                        newly_selected.append(item)
                         if prompt_each:
                             logger.info(f"Selected: {label_func(item)}")
                     else:
                         logger.warning("Item already selected.")
-                    else:
+
-                        logger.warning("Selection out of range. Try again.")
+                # Remove newly selected items from remaining list
-                except ValueError:
+                remaining_items = [item for item in remaining_items if item not in newly_selected]
-                    logger.warning("Invalid input. Enter a number or 'Q' to quit.")
+
             return selected if selected else None
         else:
+            Selector._display_choices(full_sorted_items, label_func, header=header)
             try:
                 choice = int(get_sanitized_input("Select one item by number: "))
-                if 1 <= choice <= len(sorted_items):
+                if 1 <= choice <= len(full_sorted_items):
-                    selected_item = sorted_items[choice - 1]
+                    selected_item = full_sorted_items[choice - 1]
                     logger.info(f"Selected: {label_func(selected_item)}")
                     return selected_item
                 else:
@@ -91,6 +128,38 @@ class Selector:
                 logger.warning("Invalid input.")
             return None
 
+    @staticmethod
+    def select_with_mode(
+        items: List[Any],
+        label_func: Callable[[Any], str],
+        header: str = "Available Choices:"
+    ) -> List[Any]:
+        print(colorText("Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white"))
+        mode = get_sanitized_input("").strip().lower()
+
+        if mode == "a":
+            return items
+        selected = Selector._select_from_list(
+            items,
+            label_func=label_func,
+            allow_multiple=True,
+            prompt_each=False,
+            header=header
+        )
+
+        if not selected:
+            return items
+
+        if mode == "i":
+            print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
+            return selected
+        elif mode == "e":
+            print(colorText(f"🚫 Excluded {len(selected)} item(s).", "yellow"))
+            return [item for item in items if item not in selected]
+        else:
+            print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
+            return items
+
     @staticmethod
     def select_objects(
         objects: List[Any],

utils/utils.py

@@ -615,3 +615,7 @@ def open_directory(path):
         raise OSError(f"Unsupported operating system: {system}")
 
 
+def print_x_wide(items: list, width: int):
+    for i in range(0, len(items), width):
+        row = items[i:i+width]
+        print(" | ".join(row))