Major refactor: security enhancements, modularization, config integration, reduced Parquet reliance
- Migrated codebase to class-based architecture for better modularity and maintainability - Introduced system_config.json for centralized configuration (required for runtime) - Added structured working directories for improved file organization - Significantly reduced reliance on Parquet; replaced with alternative data handling - Implemented security improvements across modules - Several TODOs remain in the main script for future enhancements - Linter formatting affected readability in some files (e.g., utils); cleanup is on the agenda
This commit is contained in:
@@ -0,0 +1,292 @@
|
||||
# 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 datetime
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
|
||||
import dotenv
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from models.agent import Agent
|
||||
from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.scheduler import (
|
||||
register_function,
|
||||
run_once_job,
|
||||
)
|
||||
from utils.utils import colorText, load_env, load_env_json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
def getLocalApprovals(api: AirlockAPIWrapper):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
result = api.otp_find_awaiting()
|
||||
local_approval = pd.DataFrame(result["response"]["otpusage"])
|
||||
if os.path.exists(f"{working_dir}\\Scheduling\\newest_local_approval.parquet"):
|
||||
previous_run = pd.read_parquet(f"{working_dir}\\Scheduling\\newest_local_approval.parquet")
|
||||
previous_run.to_parquet(
|
||||
f"{working_dir}\\Scheduling\\last_local_approval.parquet", index=False
|
||||
)
|
||||
os.remove(f"{working_dir}\\Scheduling\\newest_local_approval.parquet")
|
||||
|
||||
# Only keep rows presumably created by the generate local approval function
|
||||
local_approval = local_approval[
|
||||
local_approval["purpose"].str.startswith("🎫 Local Approval 🎫")
|
||||
]
|
||||
|
||||
local_approval["batchid"] = local_approval["purpose"].apply(
|
||||
lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1)
|
||||
)
|
||||
|
||||
if not local_approval.empty:
|
||||
local_approval.to_parquet(
|
||||
f"{working_dir}\\Scheduling\\newest_local_approval.parquet", index=False
|
||||
)
|
||||
|
||||
return local_approval
|
||||
|
||||
|
||||
def scheduleAddingLAHashes(api: AirlockAPIWrapper):
|
||||
|
||||
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD","{}")
|
||||
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
|
||||
pups = load_env_json("PUPS", "[]")
|
||||
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type = int)
|
||||
|
||||
try:
|
||||
register_function("add_hash", returnFromLocalApproval)
|
||||
register_function("move_device", moveAgentToRelatedPolicy)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to register functions: {e}")
|
||||
return
|
||||
|
||||
try:
|
||||
approvals_df = getNewLocalApprovals(api)
|
||||
if approvals_df.empty:
|
||||
logger.debug("No new local approvals found. Nothing to schedule.")
|
||||
return
|
||||
batches = approvals_df.groupby("batchid")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to retrieve or group local approvals: {e}")
|
||||
return
|
||||
|
||||
for batchid, batch_df in batches:
|
||||
try:
|
||||
duration_minutes = int(batch_df["duration"].iloc[0])
|
||||
start_time = datetime.datetime.now()
|
||||
run_time = start_time + datetime.timedelta(minutes=duration_minutes)
|
||||
early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95))
|
||||
|
||||
early_timestamp = early_time.timestamp()
|
||||
run_timestamp = run_time.timestamp()
|
||||
|
||||
# Schedule add_hash job
|
||||
try:
|
||||
run_once_job(
|
||||
f"add_hash_{batchid}",
|
||||
"add_hash",
|
||||
early_timestamp,
|
||||
[
|
||||
api,
|
||||
batch_df,
|
||||
policy_relationship_map,
|
||||
bad_publisher_list,
|
||||
pups,
|
||||
threat_tolerance_constant,
|
||||
],
|
||||
None,
|
||||
)
|
||||
logger.debug(f"Scheduled add_hash for batch {batchid} at {early_time}")
|
||||
except Exception:
|
||||
logger.debug("Failed to schedule add_hash for batch {batchid}: {e}")
|
||||
|
||||
# Schedule move_device jobs
|
||||
devices = batch_df["agentid"].drop_duplicates().tolist()
|
||||
agents = []
|
||||
|
||||
for device in devices:
|
||||
rows = api.agent_find_by_hostname(device).iterrows()
|
||||
agents += [Agent(**row["data"]) for _, row in rows]
|
||||
|
||||
for agent in agents:
|
||||
try:
|
||||
run_once_job(
|
||||
f"move_device_{agent.hostame}_{batchid}",
|
||||
"move_device",
|
||||
run_timestamp,
|
||||
[api, agent, policy_relationship_map],
|
||||
"enforcement",
|
||||
)
|
||||
|
||||
print(
|
||||
f"Scheduled move_device for device {agent.hostname} in batch {batchid} at {run_time}"
|
||||
)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"Failed to schedule move_device for device {agent.hostname} in batch {batchid}: {e}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to process batch {batchid}: {e}")
|
||||
|
||||
|
||||
def returnFromLocalApproval(api, device_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant
|
||||
):
|
||||
"""
|
||||
# Get unique policy names from device list
|
||||
policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist())
|
||||
|
||||
# Create inverse map to go from Audit to Enforcement
|
||||
inverse_map = {v: k for k, v in policy_relationship_map.items()}
|
||||
|
||||
# Fetch all policies
|
||||
all_policies = [Policy(row['groupid'], row['hidden'], row['name'], row['parent']) for _, row in api.policy_find_all().iterrows()]
|
||||
|
||||
# Define policy types
|
||||
policy_types = [1, 2, 6, 7]
|
||||
|
||||
#TODO finish logic for adding hashes
|
||||
"""
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD","{}")
|
||||
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
|
||||
pups = load_env_json("PUPS", "[]")
|
||||
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE")
|
||||
print(f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}")
|
||||
|
||||
def moveToLocalApproval(api: AirlockAPIWrapper):
|
||||
possible_durations = [15, 60, 360, 1440, 10080]
|
||||
duration_selected = None
|
||||
|
||||
print(colorText("Please select a duration:", "white"))
|
||||
for i, option in enumerate(possible_durations, start=1):
|
||||
print(f"{i}. {option}")
|
||||
|
||||
try:
|
||||
choice = int(input("Enter the number of your choice: "))
|
||||
if 1 <= choice <= len(possible_durations):
|
||||
duration_selected = possible_durations[choice - 1]
|
||||
print(colorText(f"You selected: {duration_selected}", "yellow"))
|
||||
logger.debug(f"You selected: {duration_selected}")
|
||||
else:
|
||||
print(colorText("❌ Invalid choice.", "red"))
|
||||
logger.debug("Invalid Input")
|
||||
return
|
||||
except ValueError:
|
||||
print(colorText("❌ Invalid input. Please enter a number.", "red"))
|
||||
logger.debug("Invalid Input")
|
||||
return
|
||||
|
||||
agents = selectAgents(api)
|
||||
batch = int(time.time())
|
||||
|
||||
if not agents:
|
||||
print(colorText("❌ No agents found or error retrieving agents.", "red"))
|
||||
logger.debug("No agents found or error retrieving agents")
|
||||
return
|
||||
|
||||
for agent in agents:
|
||||
try:
|
||||
addLocalApproval(api, batch, duration_selected, agent.agentid)
|
||||
moveAgentToRelatedPolicy(api, agent, "audit")
|
||||
except Exception as e:
|
||||
print(colorText(f"❌ Error processing agent {agent.hostname}: {e}", "red"))
|
||||
|
||||
|
||||
def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid):
|
||||
|
||||
purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
|
||||
api.otp_generate(agentid, duration_selected, purpose)
|
||||
|
||||
|
||||
|
||||
def monitorAuditStatus(api: AirlockAPIWrapper):
|
||||
current_agents = findAllAgents(api)
|
||||
last_agents = []
|
||||
if not last_agents:
|
||||
last_agents = current_agents
|
||||
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD","{}")
|
||||
|
||||
# Reverse map for audit → enforcement
|
||||
reverse_policy_map = {v: k for k, v in policy_relationship_map.items()}
|
||||
known_transitions = set(policy_relationship_map.items()) | set(reverse_policy_map.items())
|
||||
|
||||
# Index last_agents by hostname for quick lookup
|
||||
last_agent_map = {agent.hostname: agent for agent in last_agents}
|
||||
|
||||
# Result buckets
|
||||
newly_added = []
|
||||
same_policy = []
|
||||
moved_to_audit = []
|
||||
moved_to_enforcement = []
|
||||
unusual_move = []
|
||||
|
||||
for current in current_agents:
|
||||
previous = last_agent_map.get(current.hostname)
|
||||
|
||||
if not previous:
|
||||
newly_added.append(current)
|
||||
continue
|
||||
|
||||
if current.groupid == previous.groupid:
|
||||
same_policy.append(current)
|
||||
elif (previous.groupid, current.groupid) in known_transitions:
|
||||
moved_to_audit.append(current)
|
||||
elif (current.groupid, previous.groupid) in known_transitions:
|
||||
moved_to_enforcement.append(current)
|
||||
else:
|
||||
unusual_move.append(current)
|
||||
|
||||
# Return all five DataFrames
|
||||
return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move
|
||||
|
||||
|
||||
def getNewLocalApprovals(api: AirlockAPIWrapper):
|
||||
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
current_la = getLocalApprovals(api)
|
||||
|
||||
# Load old approval list
|
||||
old_la_path = f"{working_dir}\\Scheduling\\last_local_approval.parquet"
|
||||
if os.path.exists(old_la_path):
|
||||
old_la = pd.read_parquet(old_la_path)
|
||||
else:
|
||||
old_la = pd.DataFrame(columns=current_la.columns)
|
||||
|
||||
# Create composite keys
|
||||
current_la["key"] = current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str)
|
||||
old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str)
|
||||
|
||||
# Find new entries
|
||||
new_entries = current_la[~current_la["key"].isin(old_la["key"])]
|
||||
|
||||
# Convert 'granted' to datetime and filter by last 10 minutes
|
||||
new_entries["granted"] = pd.to_datetime(new_entries["granted"], utc=True, errors="coerce")
|
||||
ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10)
|
||||
recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago]
|
||||
|
||||
# Save current approvals for next run
|
||||
current_la.drop(columns=["key"], inplace=True)
|
||||
current_la.to_parquet(old_la_path, index=False)
|
||||
|
||||
return recent_entries
|
||||
+220
@@ -0,0 +1,220 @@
|
||||
# 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 logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
"""
|
||||
def getActiveOTP(url):
|
||||
|
||||
endpoint = url + f'/v1/otp/usage'
|
||||
payload = {
|
||||
"status" : "1"
|
||||
}
|
||||
|
||||
headers = {"X-APIKey": load_env('APIKEY')}
|
||||
payload = json.dumps(payload)
|
||||
|
||||
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
|
||||
result = json.loads(response.text)
|
||||
otp = pd.DataFrame(result["response"]["otpusage"])
|
||||
if os.path.exists("OTP\\PARQ\\newest_active_OTP.parquet"):
|
||||
previous_run = pd.read_parquet("OTP\\PARQ\\newest_active_OTP.parquet")
|
||||
previous_run.to_parquet("OTP\\PARQ\\old_active_OTP.parquet", index=False)
|
||||
os.remove("OTP\\PARQ\\newest_active_OTP.parquet")
|
||||
otp.to_parquet("OTP\\PARQ\\newest_active_OTP.parquet", index=False)
|
||||
if not otp.empty:
|
||||
formatHTML(otp, f"OTP\\HTML\\newest_active_OTP.html")
|
||||
|
||||
def getOTPActivities(url, otpid):
|
||||
endpoint = url + f'/v1/otp/activities'
|
||||
payload = {"otpid": f"{otpid}"}
|
||||
headers = {"X-APIKey": load_env('APIKEY')}
|
||||
payload = json.dumps(payload)
|
||||
|
||||
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
|
||||
result = json.loads(response.text)
|
||||
new_data = pd.DataFrame(result["response"]["otpactivities"])
|
||||
|
||||
# Define file path
|
||||
parquet_path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet"
|
||||
|
||||
# Check if file exists and read it
|
||||
if os.path.exists(parquet_path):
|
||||
existing_data = pd.read_parquet(parquet_path)
|
||||
combined_data = pd.concat([existing_data, new_data], ignore_index=True)
|
||||
combined_data.drop_duplicates(inplace=True)
|
||||
else:
|
||||
combined_data = new_data
|
||||
|
||||
# Save combined data
|
||||
combined_data.to_parquet(parquet_path, index=False)
|
||||
|
||||
# Optional: generate styled HTML if there's data
|
||||
if not combined_data.empty:
|
||||
formatHTML(combined_data, f"OTP/HTML/OTP_activities_{otpid}.html")
|
||||
|
||||
def monitorOTP(url, pups):
|
||||
|
||||
getActiveOTP(url)
|
||||
|
||||
|
||||
old_otp_path = "OTP\\PARQ\\old_active_OTP.parquet"
|
||||
new_otp_path = "OTP\\PARQ\\newest_active_OTP.parquet"
|
||||
|
||||
if os.path.exists(old_otp_path):
|
||||
old_active_OTP = pd.read_parquet(old_otp_path)
|
||||
else:
|
||||
old_active_OTP = pd.DataFrame(columns=['otpid']) # Ensure expected column exists
|
||||
|
||||
current_active_OTP = pd.read_parquet(new_otp_path)
|
||||
|
||||
|
||||
if 'otpid' not in current_active_OTP.columns: current_active_OTP = pd.DataFrame(columns=['otpid'])
|
||||
if 'otpid' not in old_active_OTP.columns: old_active_OTP = pd.DataFrame(columns=['otpid'])
|
||||
|
||||
newly_added = current_active_OTP[~current_active_OTP['otpid'].isin(old_active_OTP['otpid'])]
|
||||
still_in_OTP = old_active_OTP[old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
|
||||
no_longer_OTP = old_active_OTP[~old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
|
||||
|
||||
|
||||
register_function("addhash", addOTPHashes)
|
||||
|
||||
for _, row in newly_added.iterrows():
|
||||
clientid = row['clientid']
|
||||
duration = (int(row['duration']) * 60)
|
||||
hostname = row['hostname']
|
||||
purpose = row ['purpose']
|
||||
pid = row['otpid']
|
||||
early = math.floor(duration * .95)
|
||||
#If newly added to the list - schedule adding the majority of the executions prior to the expiration of OTP period.
|
||||
run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}", "addhash", time.time() + early, [url, clientid, pid, pups], None)
|
||||
print(f"Processing: {pid} with other data: {row}")
|
||||
|
||||
|
||||
for _, row in still_in_OTP.iterrows():
|
||||
pid = row['otpid']
|
||||
#While still in OTP, continue to update activities list
|
||||
getOTPActivities(url,pid)
|
||||
|
||||
for _, row in no_longer_OTP.iterrows():
|
||||
clientid = row['clientid']
|
||||
hostname = row['hostname']
|
||||
purpose = row ['purpose']
|
||||
pid = row['otpid']
|
||||
allowlist = clientf.getDestAllowlistFromClientID(url,clientid)
|
||||
policy, policyid = clientf.getPolicyFromClientID(url,clientid)
|
||||
|
||||
"""
|
||||
|
||||
# Devices can come out of OTP either by timeout, or by early move out of OTP. If they are manually moved out prior to the job to add hashes can run, we want to accelerate the job.
|
||||
# But first, we want to update the OTP activities one final time for the pid, then move up any jobs if they exist, then add the hashes to the local approval allowlist
|
||||
"""
|
||||
|
||||
getOTPActivities(url,pid)
|
||||
find_and_prioritize_jobs_by_pid(pid, 1)
|
||||
addOTPHashes(url, clientid,pid, pups)
|
||||
|
||||
finalhashesadded = pd.read_parquet(f"OTP\\PARQ\\otp_activities_{pid}.parquet")
|
||||
finalhashesadded['policy'] = policy
|
||||
finalhashesadded['allowlist'] = allowlist
|
||||
finalhashesadded['added_at'] = time.localtime()
|
||||
|
||||
if not os.path.exists(f"OTP\\PARQ\\localapprovalhistory.parquet"):
|
||||
df = pd.DataFrame()
|
||||
df.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
|
||||
|
||||
history = pd.read_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
|
||||
history = pd.concat([history, finalhashesadded], ignore_index=True)
|
||||
formatHTML(history, f"localapproval_history.html")
|
||||
|
||||
history.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
|
||||
|
||||
os.remove(f"OTP\\PARQ\\otp_activities_{pid}.parquet")
|
||||
del finalhashesadded
|
||||
|
||||
def addOTPHashes(url, clientid, otpid, pups):
|
||||
path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet"
|
||||
activities = pd.read_parquet(path)
|
||||
|
||||
pattern = regulator(pups)
|
||||
allowlist = clientf.getDestAllowlistFromClientID(url, clientid)
|
||||
|
||||
# Initialize or preserve 'hash_added' column
|
||||
if "hash_added" not in activities.columns: activities["hash_added"] = None
|
||||
|
||||
# Identify rows that should be added (not matching pattern and not already added)
|
||||
approve_by_hash = activities[
|
||||
~activities["filename"].str.contains(pattern, na=False) & (activities["hash_added"] != "added")
|
||||
]
|
||||
|
||||
hashes_to_add = approve_by_hash["sha256"].tolist()
|
||||
|
||||
# Add hashes to policy
|
||||
if hashes_to_add:
|
||||
#TODO add the api call
|
||||
pass
|
||||
|
||||
# Update 'hash_added' column
|
||||
activities["hash_added"] = activities.apply(
|
||||
lambda row: "do not add" if pd.notna(row["filename"]) and pattern in row["filename"]
|
||||
else ("added" if row["sha256"] in hashes_to_add else row["hash_added"]),
|
||||
axis=1
|
||||
)
|
||||
|
||||
# Save the updated DataFrame
|
||||
activities.to_parquet(path)
|
||||
|
||||
def generateOTP(url, agentid):
|
||||
purpose = input(colorText(" Please enter the purpose for the OTP: ", "white"))
|
||||
|
||||
possible_durations = [15, 60, 360, 1440, 10080]
|
||||
duration_selected = " "
|
||||
|
||||
print(colorText("Please select a duration:", "white"))
|
||||
for i, option in enumerate(possible_durations, start=1):
|
||||
print(f"{i}. {option}")
|
||||
|
||||
try:
|
||||
choice = int(input("Enter the number of your choice: "))
|
||||
if 1 <= choice <= len(possible_durations):
|
||||
duration_selected = possible_durations[choice - 1]
|
||||
print(colorText(f"You selected: {duration_selected}", "yellow"))
|
||||
else:
|
||||
print(colorText("Invalid choice.", "red"))
|
||||
except ValueError:
|
||||
print(colorText("Invalid input. Please enter a number.", "red"))
|
||||
|
||||
endpoint = url + '/v1/otp/retrieve'
|
||||
payload = {
|
||||
"duration" : f"{duration_selected}",
|
||||
"agentid" : f"{agentid}",
|
||||
"purpose" : f"{purpose}"
|
||||
}
|
||||
|
||||
headers = {"X-APIKey": load_env('APIKEY')}
|
||||
payload = json.dumps(payload)
|
||||
|
||||
|
||||
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
|
||||
result = json.loads(response.text)
|
||||
otpcode = result["response"]["otpcode"]
|
||||
print(colorText(f"The OPT code is: {otpcode}", "yellow"))
|
||||
|
||||
"""
|
||||
@@ -0,0 +1,351 @@
|
||||
# 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 logging
|
||||
import os
|
||||
import os.path
|
||||
from typing import List
|
||||
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
from models.execution import ExecutionHistoryRecord, Hash
|
||||
from models.policy import Allowlist, Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.selector import Selector
|
||||
from utils.utils import (
|
||||
colorText,
|
||||
formatHTML,
|
||||
import_to_dataframe,
|
||||
load_env,
|
||||
load_env_json,
|
||||
regulator,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
|
||||
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
|
||||
|
||||
allowlists = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||||
logger.debug("Prompting for Policies")
|
||||
print(colorText("Please select policy/policies", "white"))
|
||||
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
|
||||
|
||||
if selected is None:
|
||||
return []
|
||||
|
||||
# Normalize to always return a list
|
||||
logger.debug("Returning {selected.dict}")
|
||||
return selected if isinstance(selected, list) else [selected]
|
||||
|
||||
|
||||
def selectAllowlists(api: AirlockAPIWrapper, allow_multiple=True) -> List[Allowlist]:
|
||||
|
||||
allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()]
|
||||
logger.debug("Prompting for Allowlist(s)")
|
||||
print(colorText("Please select allowlist(s)", "white"))
|
||||
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
|
||||
|
||||
if selected is None:
|
||||
return []
|
||||
|
||||
# Normalize to always return a list
|
||||
logger.debug(f"Returning {selected}")
|
||||
return selected if isinstance(selected, list) else [selected]
|
||||
|
||||
|
||||
def sortHashes(
|
||||
api: AirlockAPIWrapper,
|
||||
selected_policies: List[Policy],
|
||||
type=[1, 2, 6, 7]
|
||||
):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
|
||||
logger.debug(f"{history_days} day selected for history")
|
||||
|
||||
if history_days is None:
|
||||
logging.warning("No history range selected. Aborting.")
|
||||
return
|
||||
|
||||
executions = []
|
||||
hashes = []
|
||||
|
||||
# Pull execution histories for each policy
|
||||
|
||||
policy_executions = ExecutionHistoryRecord.from_policies(
|
||||
api, selected_policies, type_=type, history_days=history_days
|
||||
)
|
||||
|
||||
logger.debug(f"Policy_executions is {policy_executions}")
|
||||
|
||||
executions.extend(policy_executions)
|
||||
logger.debug(f"Executions contains {executions}")
|
||||
if executions:
|
||||
hashes = [Hash(sha256=row["sha256"], **row["data"]) for _, row in api.hash_query([record.sha256 for record in executions]).iterrows()
|
||||
]
|
||||
|
||||
if hashes:
|
||||
unique_hashes = Hash.deduplicate(hashes)
|
||||
|
||||
needs_review, approved, unapproved = Hash.categorize_hashes(
|
||||
hashes=unique_hashes
|
||||
)
|
||||
|
||||
categories = {
|
||||
"needs_review": needs_review,
|
||||
"approved": approved,
|
||||
"unapproved": unapproved,
|
||||
}
|
||||
|
||||
for label, category in categories.items():
|
||||
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{label}_executions.csv"
|
||||
html_path = f"{working_dir}\\Needs_Review\\HTML\\{label}.html"
|
||||
|
||||
ExecutionHistoryRecord.enrich_with_hashes_and_export(
|
||||
executions, category, f"{working_dir}\\Needs_Review\\Review_First", label=label
|
||||
)
|
||||
df = import_to_dataframe(csv_path)
|
||||
formatHTML(df, html_path)
|
||||
|
||||
def buildPathsandPublishers(split):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
df1 = pd.DataFrame()
|
||||
df2 = pd.DataFrame()
|
||||
all_approved_hashes = pd.DataFrame()
|
||||
path1 = f"{working_dir}\\Approved\\approved_executions.csv"
|
||||
path2 = f"{working_dir}\\Approved\\needs_review_executions.csv"
|
||||
|
||||
if os.path.exists(path1):
|
||||
df1 = pd.read_csv(path1)
|
||||
else:
|
||||
logger.warning(f"File not found: {path1}")
|
||||
|
||||
if os.path.exists(path2):
|
||||
df2 = pd.read_csv(path2)
|
||||
else:
|
||||
logger.warning(f"File not found: {path2}")
|
||||
|
||||
if df1.empty and df2.empty:
|
||||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||
all_approved_hashes = pd.DataFrame()
|
||||
logger.debug(all_approved_hashes.head)
|
||||
else:
|
||||
all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
|
||||
if "filename_exec" in all_approved_hashes.columns:
|
||||
all_approved_hashes = all_approved_hashes.sort_values(by="filename_exec")
|
||||
else:
|
||||
logger.warning("Warning: 'filename_exec' column not found in concatenated DataFrame.")
|
||||
|
||||
if not all_approved_hashes.empty:
|
||||
primary_path_exclusions = calculatePath(
|
||||
all_approved_hashes,
|
||||
split,
|
||||
)
|
||||
remaining_hashes = all_approved_hashes[
|
||||
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
|
||||
]
|
||||
secondary_path_exclusions = calculatePath(
|
||||
remaining_hashes, split
|
||||
)
|
||||
remaining_hashes = remaining_hashes[
|
||||
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
|
||||
]
|
||||
dataframes = {
|
||||
"primary_Paths": primary_path_exclusions,
|
||||
"secondary_Paths": secondary_path_exclusions,
|
||||
"hashes_to_add": remaining_hashes,
|
||||
}
|
||||
|
||||
for name, df in dataframes.items():
|
||||
df.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{name}.csv", index=False)
|
||||
formatHTML(df, f"{working_dir}\\Needs_Review\\HTML\\{name}.html")
|
||||
|
||||
if not all_approved_hashes.empty:
|
||||
# Drop all not signed, only keep unique values
|
||||
publist = all_approved_hashes[
|
||||
all_approved_hashes["publisher_hash"] != "Not Signed"
|
||||
].drop_duplicates(subset=["publisher_hash"])
|
||||
# Remove Bad publisher if somehow they made it this far
|
||||
pattern = regulator(load_env_json("BAD_PUBLISHERS","[]"))
|
||||
publist = publist[~publist["publisher_hash"].str.contains(pattern, na=False)]
|
||||
publist = publist[["publisher_hash"]]
|
||||
publist.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\publishers.csv", index=False)
|
||||
|
||||
def buildPreflights():
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
df1 = pd.DataFrame()
|
||||
df2 = pd.DataFrame()
|
||||
approved_hashes = pd.DataFrame()
|
||||
approved_publishers = pd.DataFrame()
|
||||
|
||||
hash = f"{working_dir}\\Approved\\hashes_to_add.csv"
|
||||
path1 = f"{working_dir}\\Approved\\primary_Paths.csv"
|
||||
path2 = f"{working_dir}\\Approved\\secondary_Paths.csv"
|
||||
publishers = f"{working_dir}\\Approved\\publishers.csv"
|
||||
|
||||
if os.path.exists(hash):
|
||||
approved_hashes = pd.read_csv(hash)
|
||||
|
||||
else:
|
||||
logger.warning(f"File not found: {hash}")
|
||||
|
||||
if os.path.exists(path1):
|
||||
df1 = pd.read_csv(path1)
|
||||
else:
|
||||
logger.warning(f"File not found: {path1}")
|
||||
|
||||
if os.path.exists(path2):
|
||||
df2 = pd.read_csv(path2)
|
||||
else:
|
||||
logger.warning(f"File not found: {path2}")
|
||||
|
||||
if df1.empty and df2.empty:
|
||||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||
approved_paths = pd.DataFrame()
|
||||
else:
|
||||
approved_paths = pd.concat([df1, df2], ignore_index=True)
|
||||
|
||||
if os.path.exists(publishers):
|
||||
approved_publishers = pd.read_csv(publishers)
|
||||
|
||||
else:
|
||||
logger.warning(f"File not found: {publishers}")
|
||||
|
||||
dataframes = {"approved_paths": approved_paths, "approved_hashes": approved_hashes, "approved_publishers": approved_publishers}
|
||||
|
||||
for name, df in dataframes.items():
|
||||
df.to_csv(f"{working_dir}\\Preflight\\{name}.csv", index=False)
|
||||
formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{name}.html")
|
||||
|
||||
def splitFilepathsGrouped(df, col="filename"):
|
||||
path_exclusion_constant = load_env("PATH_EXCLUSION_CONST", cast_type= int)
|
||||
min_files_for_path = load_env("MIN_FILES_FOR_PATH", cast_type= int)
|
||||
|
||||
def clean_split(path):
|
||||
if not isinstance(path, (str, bytes, os.PathLike)):
|
||||
return []
|
||||
parts = os.path.normpath(path).split(os.sep)
|
||||
parts = [p for p in parts if p] # Remove empty strings
|
||||
return parts
|
||||
|
||||
# Diagnostic: log any non-string entries
|
||||
non_string_entries = df[~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))]
|
||||
if not non_string_entries.empty:
|
||||
print(f"[WARNING] Non-string entries found in column '{col}':")
|
||||
print(non_string_entries)
|
||||
|
||||
df = df.copy()
|
||||
split_paths = df[col].apply(clean_split)
|
||||
|
||||
# Filter out paths with fewer than `min_files_for_path` components
|
||||
df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy()
|
||||
split_paths = split_paths[df.index]
|
||||
|
||||
df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:path_exclusion_constant]))
|
||||
grouped = df.groupby("group_key")
|
||||
new_rows = []
|
||||
|
||||
for _, group_df in grouped:
|
||||
paths = group_df[col].tolist()
|
||||
split_parts = [clean_split(p) for p in paths]
|
||||
|
||||
def longest_common_prefix(paths):
|
||||
if not paths:
|
||||
return []
|
||||
prefix = paths[0]
|
||||
for path in paths[1:]:
|
||||
prefix = [a for a, b in zip(prefix, path) if a == b]
|
||||
if not prefix:
|
||||
break
|
||||
return prefix
|
||||
|
||||
common_prefix = longest_common_prefix(split_parts)
|
||||
prefix_str = os.sep.join(common_prefix)
|
||||
|
||||
for i, parts in enumerate(split_parts):
|
||||
filename = parts[-1]
|
||||
middle = (
|
||||
os.sep.join(parts[len(common_prefix):-1])
|
||||
if len(parts) > len(common_prefix) + 1
|
||||
else ""
|
||||
)
|
||||
row = group_df.iloc[i].copy()
|
||||
row["longestcfp"] = prefix_str
|
||||
row["middle"] = middle
|
||||
row["filename_only"] = filename
|
||||
row["file_extension"] = os.path.splitext(filename)[1].lower()
|
||||
new_rows.append(row)
|
||||
|
||||
return pd.DataFrame(new_rows).drop(columns=["group_key"])
|
||||
|
||||
def calculatePath(approved_hashes, split):
|
||||
if split:
|
||||
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
|
||||
else:
|
||||
dfs_by_policy = [approved_hashes]
|
||||
|
||||
badpathparts = load_env_json("BAD_PATH_PARTS", "[]")
|
||||
min_files_for_path = load_env("MIN_FILES_FOR_PATH", cast_type = int)
|
||||
|
||||
processed_dfs = []
|
||||
|
||||
for df in dfs_by_policy:
|
||||
haslcp = splitFilepathsGrouped(df, "filename_exec")
|
||||
haslcp = haslcp.drop_duplicates()
|
||||
|
||||
forbidden = regulator(badpathparts, True)
|
||||
forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
|
||||
|
||||
logger.debug("Removing forbidden filepaths for path exceptions")
|
||||
print(colorText("Removing forbidden filepaths for path exceptions", "green"))
|
||||
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
|
||||
|
||||
lcp_not_forbidden_review = lcp_not_forbidden[
|
||||
[
|
||||
"policyname",
|
||||
"longestcfp",
|
||||
"middle",
|
||||
"filename_only",
|
||||
"file_extension",
|
||||
"sha256",
|
||||
]
|
||||
]
|
||||
|
||||
unique_sha_counts = (
|
||||
lcp_not_forbidden_review.groupby("longestcfp")["sha256"].nunique().reset_index()
|
||||
)
|
||||
unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]
|
||||
|
||||
lcp_not_forbidden_review = lcp_not_forbidden_review.merge(
|
||||
unique_sha_counts, on="longestcfp", how="left"
|
||||
)
|
||||
lcp_not_forbidden_review = lcp_not_forbidden_review[
|
||||
lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path
|
||||
]
|
||||
processed_dfs.append(lcp_not_forbidden_review)
|
||||
|
||||
pathExclusions = pd.concat(processed_dfs, ignore_index=True)
|
||||
|
||||
return pathExclusions
|
||||
@@ -0,0 +1,123 @@
|
||||
# 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 datetime
|
||||
import logging
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from flows.prepPolicy import selectPolicies
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.policyhandler import getPolicyInfo
|
||||
from services.selector import Selector
|
||||
from utils.utils import colorText, load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
import dotenv
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
def findQuietAgents(api: AirlockAPIWrapper):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
# Get policy selection and agent list
|
||||
selected_policy = selectPolicies(api, False)
|
||||
if selected_policy:
|
||||
agents = api.agents_find_by_group(selected_policy[0].groupid)
|
||||
|
||||
# Prompt user for history range
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
required_quiet = Selector.select_value(
|
||||
prompt="Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–365): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 365),
|
||||
)
|
||||
|
||||
# Get execution history as a DataFrame
|
||||
policy_exec_history = getPolicyInfo(
|
||||
api, selected_policy[0], [1, 2, 6, 7], history_days
|
||||
)
|
||||
|
||||
|
||||
if policy_exec_history.empty:
|
||||
logging.info("No execution history found for the selected policy and time range.")
|
||||
return
|
||||
|
||||
|
||||
# Convert 'datetime' column to timezone-aware datetime objects
|
||||
policy_exec_history["datetime"] = pd.to_datetime(
|
||||
policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True
|
||||
)
|
||||
|
||||
# Get current UTC time
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
|
||||
# Calculate days ago
|
||||
policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
|
||||
lambda dt: (now - dt).days
|
||||
)
|
||||
|
||||
# Count total executions per hostname
|
||||
hostname_counts = policy_exec_history["hostname"].value_counts()
|
||||
|
||||
# Map execution counts to agents
|
||||
agents["execution_count"] = agents["hostname"].map(hostname_counts).fillna(0).astype(int)
|
||||
|
||||
# Find most recent execution per hostname
|
||||
most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates(
|
||||
subset="hostname", keep="first"
|
||||
)
|
||||
|
||||
# Map most recent execution age to agents
|
||||
agents["days_since"] = agents["hostname"].map(
|
||||
most_recent_exec.set_index("hostname")["days_ago"]
|
||||
)
|
||||
|
||||
# Check for enforcement readiness
|
||||
agents["required_quiet"] = required_quiet
|
||||
agents["enforce_ready"] = agents["days_since"].apply(
|
||||
lambda x: True if pd.isna(x) or x > required_quiet else False
|
||||
)
|
||||
|
||||
# Sort agents by execution count and hostname
|
||||
agents = agents.sort_values(by=["execution_count", "hostname"], ascending=[True, True])
|
||||
|
||||
# Save to CSV
|
||||
filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv"
|
||||
logging.debug(f"Saving CSV to {filename}")
|
||||
print(colorText(f"Saving CSV to {filename}", "green"))
|
||||
agents.to_csv(filename, index=False)
|
||||
|
||||
# Summary statistics
|
||||
total_agents = len(agents)
|
||||
ready_agents = agents["enforce_ready"].sum()
|
||||
not_ready_agents = total_agents - ready_agents
|
||||
ready_percentage = (ready_agents / total_agents) * 100
|
||||
|
||||
# Print results
|
||||
|
||||
|
||||
message = (
|
||||
f"Total agents: {total_agents}\n"
|
||||
f"Agents marked as 'enforce_ready': {ready_agents}\n"
|
||||
f"Agents not ready: {not_ready_agents}\n"
|
||||
f"Percentage ready for enforcement: {ready_percentage:.2f}%"
|
||||
)
|
||||
logger.debug(message)
|
||||
colorText(message,"green")
|
||||
Reference in New Issue
Block a user