323 lines
12 KiB
Python
323 lines
12 KiB
Python
# 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 json
|
||
import logging
|
||
import os
|
||
import re
|
||
from dataclasses import asdict
|
||
from datetime import datetime, timedelta
|
||
from typing import List
|
||
|
||
import pandas as pd
|
||
|
||
from flows.prepPolicy import selectPolicies
|
||
from models.agent import Agent
|
||
from models.policy import Policy
|
||
from services.API import AirlockAPIWrapper
|
||
from utils.configmanager import get_protected_json, load_env
|
||
from utils.selector import Selector
|
||
from utils.utils import colorText, get_sanitized_input
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||
agents = selectAgents(api)
|
||
history_days = Selector.select_value(
|
||
prompt="Enter how many days of history to pull (1–150): ",
|
||
value_type=int,
|
||
valid_range=(1, 150),
|
||
)
|
||
|
||
if not agents or not history_days:
|
||
print(colorText("No agents selected or invalid history range.", "red"))
|
||
return
|
||
|
||
historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d")
|
||
today = datetime.now().strftime("%Y-%m-%d")
|
||
|
||
all_history = []
|
||
|
||
for agent in agents:
|
||
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"))
|
||
continue
|
||
|
||
if isinstance(exechistory, list):
|
||
for block in exechistory:
|
||
record = {
|
||
"Command": block.get("commandline", "N/A"),
|
||
"Date": block.get("datetime", "N/A"),
|
||
"Filename": block.get("filename", "N/A"),
|
||
"Policy Name": block.get("policyname", "N/A"),
|
||
"Hostname": block.get("hostname", "N/A"),
|
||
"Hash": block.get("sha256", "N/A"),
|
||
}
|
||
all_history.append(record)
|
||
|
||
if not outputjson:
|
||
for key, value in record.items():
|
||
print(colorText(f"{key}: {value}", "green"))
|
||
print("\n")
|
||
else:
|
||
print(colorText(f"No execution history found for {agent.hostname}.", "yellow"))
|
||
|
||
if outputjson:
|
||
print(json.dumps(all_history, indent=2))
|
||
|
||
|
||
def findAllAgents(api):
|
||
# Step 1: Load data from API
|
||
policies = [Policy(**row["data"]) for _, row in api.policy_find_all().iterrows()]
|
||
agents = [Agent(**row["data"]) for _, row in api.agent_find_all().iterrows()]
|
||
|
||
for agent in agents:
|
||
agent.enrich_with_policies(policies)
|
||
|
||
return agents
|
||
|
||
def findAgents(api, return_dataframe):
|
||
agents = selectAgents(api)
|
||
working_dir = load_env("WORKING_DIR")
|
||
|
||
if not agents:
|
||
logging.warning("No agents or policies found.")
|
||
print("No agents matched the criteria.")
|
||
return
|
||
|
||
# Convert enriched agents to DataFrame
|
||
agent_dicts = [asdict(agent) for agent in agents]
|
||
agent_df = pd.DataFrame(agent_dicts)
|
||
|
||
if return_dataframe:
|
||
logging.debug("Returning DataFrame to caller.")
|
||
return agent_df
|
||
|
||
# Otherwise, print and optionally export
|
||
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':
|
||
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||
filename = f"agentsearch_{timestamp}.csv"
|
||
file_path = os.path.join(str(working_dir), filename)
|
||
|
||
agent_df.to_csv(file_path, index=False)
|
||
logging.info(f"Exported DataFrame to {file_path}")
|
||
|
||
print(
|
||
colorText(
|
||
f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
|
||
"green",
|
||
)
|
||
)
|
||
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("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_\- ]+$')
|
||
|
||
while True:
|
||
line = get_sanitized_input("")
|
||
stripped_line = line.strip()
|
||
|
||
if stripped_line == "":
|
||
empty_line_count += 1
|
||
if empty_line_count == 2:
|
||
break
|
||
continue
|
||
else:
|
||
empty_line_count = 0
|
||
|
||
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"))
|
||
|
||
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 []
|
||
|
||
use_exact = choose_match_type()
|
||
|
||
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 = match_agents(device_names, agents, use_exact)
|
||
matched_agents.sort(key=lambda agent: agent.hostname.lower())
|
||
|
||
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"))
|
||
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)
|
||
|
||
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,
|
||
mode: str = "audit",
|
||
):
|
||
"""
|
||
Moves an agent between audit and enforcement policies based on the mode.
|
||
|
||
Args:
|
||
api: AirlockAPIWrapper instance.
|
||
agent: Agent object.
|
||
policy_relationship_map: Dict mapping enforcement → audit.
|
||
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
|
||
"""
|
||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
||
|
||
if mode == "audit":
|
||
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.")
|
||
return
|
||
else:
|
||
logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.")
|
||
return
|
||
|
||
elif mode == "enforcement":
|
||
inverse_map = {v: k for k, v in policy_relationship_map.items()}
|
||
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.")
|
||
return
|
||
else:
|
||
logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.")
|
||
return
|
||
|
||
else:
|
||
logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.")
|
||
return
|
||
|
||
result = api.agent_move(agent.agentid, target_policy)
|
||
return result
|
||
|
||
|
||
def toggleEnforcement(api: AirlockAPIWrapper):
|
||
choices = ["Audit", "Enforcement", "Exit"]
|
||
print(colorText("Move devices to which state?:", "yellow"))
|
||
direction = Selector.select_string(choices, False, False)
|
||
if direction == "Exit":
|
||
pass
|
||
else:
|
||
devices = selectAgents(api)
|
||
for device in devices:
|
||
print(device.hostname)
|
||
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())
|
||
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: ")
|
||
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}?")
|
||
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") |