From 629f6b510dd624c03be8b7dcbe5ba5a214db1c5a Mon Sep 17 00:00:00 2001 From: Zarithas Date: Fri, 24 Oct 2025 15:52:25 -0400 Subject: [PATCH] Improved device selection logic --- flows/prepPolicy.py | 55 +++++++++++++-- services/agenthandler.py | 100 ++++++++++++++++++-------- utils/menus.py | 60 +++------------- utils/selector.py | 147 ++++++++++++++++++++++++++++----------- utils/utils.py | 4 ++ 5 files changed, 240 insertions(+), 126 deletions(-) diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py index 53c1f89..6f58b14 100644 --- a/flows/prepPolicy.py +++ b/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 + diff --git a/services/agenthandler.py b/services/agenthandler.py index a04fa3b..a2c0959 100644 --- a/services/agenthandler.py +++ b/services/agenthandler.py @@ -130,24 +130,16 @@ def findAgents(api, return_dataframe): else: logging.debug("User declined to export the DataFrame.") - -def selectAgents(api: AirlockAPIWrapper) -> List[Agent]: +def collect_device_names() -> List[str]: 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("UTN00000", "cyan")) - print(colorText("i-hSuperSecretServer", "cyan")) - print(colorText("u-hVenderBroke\n", "cyan")) - + print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-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 - pattern = "|".join(map(re.escape, device_names)) - regex = re.compile(pattern, re.IGNORECASE) + use_exact = choose_match_type() - # 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 - 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")) + show_unmatched(device_names, matched_agents, use_exact) if not matched_agents: logger.debug("❌ No matching devices found.") print(colorText("❌ No matching devices found.", "red")) - else: - logger.debug(f"✅ Found {len(matched_agents)} matching device(s).") - print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green")) + return [] - # Enrich each agent using its class method - for agent in matched_agents: - agent.enrich_with_policies(policies) + 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) + + matched_agents = Selector.select_with_mode( + matched_agents, + 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, diff --git a/utils/menus.py b/utils/menus.py index 0794d8e..6a2daa7 100644 --- a/utils/menus.py +++ b/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)) \ No newline at end of file diff --git a/utils/selector.py b/utils/selector.py index 87f1bfb..eb7fcaa 100644 --- a/utils/selector.py +++ b/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 . - - 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(sorted_items) + num_columns - 1) // num_columns + rows = (len(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): - label = label_func(sorted_items[idx]) + if idx < len(items): + 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) - sorted_items = sorted(items, key=lambda item: label_func(item).lower()) + full_sorted_items = Selector._get_sorted_items(items, label_func) + 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: - index = int(choice) - if 1 <= index <= len(sorted_items): - item = sorted_items[index - 1] - if item not in selected: - selected.append(item) - if prompt_each: - logger.info(f"Selected: {label_func(item)}") - else: - logger.warning("Item already selected.") + elif choice == "r": + selected.clear() + remaining_items = full_sorted_items.copy() + 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("Selection out of range. Try again.") - except ValueError: - logger.warning("Invalid input. Enter a number or 'Q' to quit.") + logger.warning("Item already selected.") + + # Remove newly selected items from remaining list + remaining_items = [item for item in remaining_items if item not in newly_selected] + 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): - selected_item = sorted_items[choice - 1] + if 1 <= choice <= len(full_sorted_items): + 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], diff --git a/utils/utils.py b/utils/utils.py index 2f48ff3..61a6541 100644 --- a/utils/utils.py +++ b/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)) \ No newline at end of file