Improved device selection logic
This commit is contained in:
+8
-52
@@ -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))
|
||||
+108
-39
@@ -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]
|
||||
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],
|
||||
|
||||
@@ -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))
|
||||
Reference in New Issue
Block a user