216 lines
8.5 KiB
Python
216 lines
8.5 KiB
Python
import asyncio
|
||
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 models.agent import Agent
|
||
from models.policy import Policy
|
||
from services.API import AirlockAPIWrapper
|
||
from services.TaskQueue import AsyncTaskQueue, run_sync_task_in_thread
|
||
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__)
|
||
|
||
async def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||
agents = await selectAgents(api)
|
||
history_days = await 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 = []
|
||
|
||
async def fetch_history(agent):
|
||
try:
|
||
exechistory = await api.history_execution(today, historical_date, agent.hostname)
|
||
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"))
|
||
except Exception as e:
|
||
print(colorText(f"❌ Error retrieving history for {agent.hostname}: {e}", "red"))
|
||
|
||
await asyncio.gather(*(fetch_history(agent) for agent in agents))
|
||
|
||
if outputjson:
|
||
print(json.dumps(all_history, indent=2))
|
||
|
||
async def findAllAgents(api: AirlockAPIWrapper):
|
||
policies_df = await api.policy_find_all()
|
||
agents_df = await api.agent_find_all()
|
||
|
||
policies = [Policy(**row["data"]) for _, row in policies_df.iterrows()]
|
||
agents = [Agent(**row["data"]) for _, row in agents_df.iterrows()]
|
||
|
||
groupid_to_name = {policy.groupid: policy.name for policy in policies}
|
||
|
||
queue = AsyncTaskQueue()
|
||
await queue.start_workers()
|
||
|
||
async def enrich_agent(agent):
|
||
agent.enrich(groupid_to_name)
|
||
|
||
for agent in agents:
|
||
await queue.enqueue(f"enrich_{agent.hostname}", enrich_agent, agent)
|
||
|
||
await asyncio.sleep(1)
|
||
await queue.stop_workers()
|
||
|
||
return agents
|
||
|
||
async def findAgents(api: AirlockAPIWrapper, return_dataframe: bool):
|
||
agents = await selectAgents(api)
|
||
working_dir = await load_env("WORKING_DIR")
|
||
|
||
if not agents:
|
||
logging.warning("No agents or policies found.")
|
||
print("No agents matched the criteria.")
|
||
return
|
||
|
||
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
|
||
|
||
print(agent_df)
|
||
logging.debug("Displayed DataFrame to console.")
|
||
|
||
user_input = await get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ")
|
||
user_input = user_input.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(working_dir, filename)
|
||
|
||
await run_sync_task_in_thread(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.")
|
||
|
||
async def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
|
||
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.\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("Paste or type your device names below:", "white"))
|
||
|
||
policies_df = await api.policy_find_all()
|
||
policies = [Policy(**row.to_dict()) for _, row in policies_df.iterrows()]
|
||
|
||
device_input_lines = []
|
||
empty_line_count = 0
|
||
valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$')
|
||
|
||
while True:
|
||
line = await 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"))
|
||
|
||
device_names = [name for name in device_input_lines if name]
|
||
if not device_names:
|
||
logger.debug("No device names entered")
|
||
print(colorText("⚠️ No device names entered.", "red"))
|
||
return []
|
||
|
||
pattern = "\n".join(map(re.escape, device_names))
|
||
regex = re.compile(pattern, re.IGNORECASE)
|
||
|
||
agents_df = await api.agent_find_all()
|
||
agents = [Agent(**row.to_dict()) for _, row in agents_df.iterrows()]
|
||
matched_agents = [agent for agent in agents if regex.search(agent.hostname)]
|
||
matched_agents.sort(key=lambda agent: agent.hostname.lower())
|
||
|
||
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"))
|
||
|
||
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"))
|
||
|
||
for agent in matched_agents:
|
||
agent.enrich_with_policies(policies)
|
||
|
||
return matched_agents
|
||
|
||
async def moveAgentToRelatedPolicy(api: AirlockAPIWrapper, agent: Agent, mode: str = "audit"):
|
||
policy_relationship_map = await get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
||
#TODO - Have this return the policy name it was moved to instead of the groupid
|
||
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 = await api.agent_move(agent.agentid, target_policy)
|
||
if result == {'error': 'Success'}: logger.info(f"{agent.hostname} has been moved to {target_policy}")
|
||
else: logger.debug(result)
|
||
|