247 lines
9.0 KiB
Python
247 lines
9.0 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 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 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 (Three times if you have a single device).\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 = [Policy(**row.to_dict()) for _, row in api.policy_find_all().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("")
|
||
stripped_line = line.strip()
|
||
|
||
if stripped_line == "":
|
||
empty_line_count += 1
|
||
if empty_line_count == 2:
|
||
break
|
||
continue # Don't validate empty lines
|
||
else:
|
||
empty_line_count = 0
|
||
|
||
# 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))
|
||
regex = re.compile(pattern, re.IGNORECASE)
|
||
|
||
# Fetch agents
|
||
agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().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)}")
|
||
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"))
|
||
|
||
# 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", "{}")
|
||
|
||
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
|
||
|
||
api.agent_move(agent.agentid, target_policy)
|