mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Improved device selection logic

Browse Source
This commit is contained in:
Zarithas
2025-10-24 15:52:25 -04:00
parent f8302e15c1
commit 629f6b510d
5 changed files with 240 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
flows/prepPolicy.py
+50 -5
View File
@@ -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
services/agenthandler.py
+68 -28
View File
@@ -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).")
return []
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
for agent in matched_agents:
agent.enrich_with_policies(policies)
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,
utils/menus.py
+8 -52
View File
@@ -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))
utils/selector.py
+103 -34
View File
@@ -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(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]
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("Item already selected.")
else:
logger.warning("Selection out of range. Try again.")
except ValueError:
logger.warning("Invalid input. Enter a number or 'Q' to quit.")
# 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],
utils/utils.py
+4
View File
@@ -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))