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

Post Black Linting

Browse Source
This commit is contained in:
Zarithas
2025-11-06 11:04:59 -05:00
parent f33b041ac0
commit b538f12e9a
20 changed files with 1106 additions and 618 deletions
Download Patch File Download Diff File Expand all files Collapse all files
services/agenthandler.py
+93 -28
View File
@@ -47,7 +47,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
print(colorText("No agents selected or invalid history range.", "red"))
return
historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d")
historical_date = (datetime.now() - timedelta(days=history_days)).strftime(
"%Y-%m-%d"
)
today = datetime.now().strftime("%Y-%m-%d")
all_history = []
@@ -56,7 +58,11 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
try:
exechistory = api.history_execution(today, historical_date, agent.hostname)
except Exception as e:
print(colorText(f"❌ Error retrieving history for {agent.hostname}: {e}", "red"))
print(
colorText(
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
)
)
continue
if isinstance(exechistory, list):
@@ -76,7 +82,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
print(colorText(f"{key}: {value}", "green"))
print("\n")
else:
print(colorText(f"No execution history found for {agent.hostname}.", "yellow"))
print(
colorText(f"No execution history found for {agent.hostname}.", "yellow")
)
if outputjson:
print(json.dumps(all_history, indent=2))
@@ -92,6 +100,7 @@ def findAllAgents(api):
return agents
def findAgents(api, return_dataframe):
agents = selectAgents(api)
working_dir = load_env("WORKING_DIR")
@@ -113,8 +122,14 @@ def findAgents(api, return_dataframe):
print(agent_df)
logging.debug("Displayed DataFrame to console.")
user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower()
if user_input == 'y':
user_input = (
get_sanitized_input(
"\nWould you like to export the results to a CSV file? (y/n): "
)
.strip()
.lower()
)
if user_input == "y":
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
filename = f"agentsearch_{timestamp}.csv"
file_path = os.path.join(str(working_dir), filename)
@@ -131,17 +146,27 @@ def findAgents(api, return_dataframe):
else:
logging.debug("User declined to export the DataFrame.")
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(
"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\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan"))
print(colorText("Paste or type your device names below:", "white"))
device_input_lines = []
empty_line_count = 0
valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$')
valid_line_pattern = re.compile(r"^[a-zA-Z0-9_\- ]+$")
while True:
line = get_sanitized_input("")
@@ -158,7 +183,12 @@ def collect_device_names() -> List[str]:
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"))
print(
colorText(
f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
"yellow",
)
)
return [name for name in device_input_lines if name]
@@ -168,10 +198,13 @@ def choose_match_type() -> bool:
return get_sanitized_input("").strip().lower() in ["y", "yes"]
def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool) -> List['Agent']:
def match_agents(
device_names: List[str], agents: List["Agent"], use_exact: bool
) -> List["Agent"]:
if use_exact:
return [
agent for agent in agents
agent
for agent in agents
if agent.hostname.lower() in [name.lower() for name in device_names]
]
else:
@@ -180,23 +213,38 @@ def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool
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):
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)]
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)]
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']):
def enrich_agents(agents: List["Agent"], policies: List["Policy"]):
for agent in agents:
agent.enrich_with_policies(policies)
def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']:
def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
device_names = collect_device_names()
if not device_names:
logger.debug("No device names entered")
@@ -227,11 +275,11 @@ def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']:
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:"
header="Matched Devices:",
)
if not matched_agents:
@@ -263,11 +311,17 @@ def moveAgentToRelatedPolicy(
if agent.groupid in policy_relationship_map:
target_policy = policy_relationship_map[agent.groupid]
elif agent.groupid in policy_relationship_map.values():
logger.debug(f"Agent {agent.hostname} is already in an audit group. No action needed.")
print(f"Agent {agent.hostname} is already in an audit group. No action needed.")
logger.debug(
f"Agent {agent.hostname} is already in an audit group. No action needed."
)
print(
f"Agent {agent.hostname} is already in an audit group. No action needed."
)
return
else:
logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.")
logger.warning(
f"Error: No corresponding audit policy found for groupid: {agent.groupid}."
)
return
elif mode == "enforcement":
@@ -275,10 +329,14 @@ def moveAgentToRelatedPolicy(
if agent.groupid in inverse_map:
target_policy = inverse_map[agent.groupid]
elif agent.groupid in inverse_map.values():
logger.info(f"Agent {agent.hostname} is already in an enforcement group. No action needed.")
logger.info(
f"Agent {agent.hostname} is already in an enforcement group. No action needed."
)
return
else:
logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.")
logger.warning(
f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}."
)
return
else:
@@ -299,25 +357,32 @@ def toggleEnforcement(api: AirlockAPIWrapper):
devices = selectAgents(api)
for device in devices:
print(device.hostname)
confirm = Selector.confirm("Would you like to continue with these devices? Y/N: ")
confirm = Selector.confirm(
"Would you like to continue with these devices? Y/N: "
)
if direction and devices and confirm:
for device in devices:
result = moveAgentToRelatedPolicy(api,device, str(direction).lower())
result = moveAgentToRelatedPolicy(api, device, str(direction).lower())
logger.info(f"{device.hostname}: result: {result}")
get_sanitized_input("Press enter to continue")
def moveAgents(api: AirlockAPIWrapper):
devices = selectAgents(api)
for device in devices:
print(device.hostname)
confirm_devices = Selector.confirm("Would you like to continue with these devices? Y/N: ")
confirm_devices = Selector.confirm(
"Would you like to continue with these devices? Y/N: "
)
if devices and confirm_devices:
policies = selectPolicies(api, False)
confirm_move = Selector.confirm(f"Would you like to move these devices to {policies[0].name}?")
confirm_move = Selector.confirm(
f"Would you like to move these devices to {policies[0].name}?"
)
if confirm_move:
for device in devices:
result = api.agent_move(device.agentid, policies[0].groupid)
logger.info(f"{device.hostname}: result: {result}")
else:
logger.info("Exiting without change")
get_sanitized_input("Press enter to continue")
get_sanitized_input("Press enter to continue")