First Round Async. Much work left to do, dont trust results of hash categorization presently.
This commit is contained in:
+70
-104
@@ -1,19 +1,4 @@
|
||||
# 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 asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -27,16 +12,16 @@ 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.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(
|
||||
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),
|
||||
@@ -48,63 +33,68 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||||
|
||||
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:
|
||||
async def fetch_history(agent):
|
||||
try:
|
||||
exechistory = api.history_execution(today, historical_date, agent.hostname)
|
||||
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"))
|
||||
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"))
|
||||
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()
|
||||
|
||||
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()]
|
||||
policies = [Policy(**row["data"]) for _, row in policies_df.iterrows()]
|
||||
agents = [Agent(**row["data"]) for _, row in agents_df.iterrows()]
|
||||
|
||||
# Step 2: Create groupid → groupname map
|
||||
groupid_to_name = {policy.groupid: policy.name for policy in policies}
|
||||
|
||||
# Step 3: Enrich agents
|
||||
for agent in agents:
|
||||
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
|
||||
|
||||
def findAgents(api, return_dataframe):
|
||||
agents = selectAgents(api)
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
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
|
||||
|
||||
# Convert enriched agents to DataFrame
|
||||
agent_dicts = [asdict(agent) for agent in agents]
|
||||
agent_df = pd.DataFrame(agent_dicts)
|
||||
|
||||
@@ -112,30 +102,24 @@ def findAgents(api, 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()
|
||||
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)
|
||||
|
||||
agent_df.to_csv(file_path, index=False)
|
||||
logging.info(f"Exported DataFrame to {file_path}")
|
||||
await run_sync_task_in_thread(agent_df.to_csv, file_path, index=False)
|
||||
|
||||
print(
|
||||
colorText(
|
||||
f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
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 selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
|
||||
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"))
|
||||
@@ -144,50 +128,44 @@ def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
|
||||
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 = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||||
|
||||
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
|
||||
|
||||
# Regex to validate each line
|
||||
valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$')
|
||||
|
||||
while True:
|
||||
line = get_sanitized_input("")
|
||||
line = await get_sanitized_input("")
|
||||
stripped_line = line.strip()
|
||||
|
||||
if stripped_line == "":
|
||||
empty_line_count += 1
|
||||
if empty_line_count == 2:
|
||||
break
|
||||
continue # Don't validate empty lines
|
||||
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"))
|
||||
|
||||
# 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]
|
||||
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))
|
||||
pattern = "\n".join(map(re.escape, device_names))
|
||||
regex = re.compile(pattern, re.IGNORECASE)
|
||||
|
||||
# Fetch agents
|
||||
agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()]
|
||||
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())
|
||||
|
||||
# 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)}")
|
||||
@@ -197,30 +175,17 @@ def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
|
||||
logger.debug("❌ No matching devices found.")
|
||||
print(colorText("❌ No matching devices found.", "red"))
|
||||
else:
|
||||
logger.debug(f"✅ Found {len(matched_agents)} matching device(s).")
|
||||
logger.debug(f"✅ Found {len(matched_agents)} matching device(s).")
|
||||
print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green"))
|
||||
|
||||
# Enrich each agent using its class method
|
||||
for agent in matched_agents:
|
||||
agent.enrich_with_policies(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", "{}")
|
||||
|
||||
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]
|
||||
@@ -231,7 +196,6 @@ def moveAgentToRelatedPolicy(
|
||||
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:
|
||||
@@ -242,9 +206,11 @@ def moveAgentToRelatedPolicy(
|
||||
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
|
||||
|
||||
api.agent_move(agent.agentid, target_policy)
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user