291 lines
11 KiB
Python
291 lines
11 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 datetime
|
|
import logging
|
|
import os
|
|
import re
|
|
import time
|
|
|
|
import dotenv
|
|
import numpy as np
|
|
import pandas as pd
|
|
from utils.setup import get_base_directory
|
|
from models.agent import Agent
|
|
from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
|
|
from services.API import AirlockAPIWrapper
|
|
|
|
from utils.utils import colorText
|
|
from utils.configmanager import get_protected_json, 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 = get_protected_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 = get_protected_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 = get_protected_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
|