Checking WIP Agent Movement Workflow
This commit is contained in:
+193
-262
@@ -1,311 +1,242 @@
|
||||
# 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/>.
|
||||
"""
|
||||
This module handles the creation of local approval requests.
|
||||
"""
|
||||
|
||||
|
||||
import datetime
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
|
||||
import dotenv
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from typing import List, Optional
|
||||
|
||||
from models.agent import Agent
|
||||
from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
|
||||
from services.agenthandler import moveAgentToRelatedPolicy, selectAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_protected_json, load_env, load_env_json
|
||||
from utils.setup import get_base_directory
|
||||
from utils.configmanager import get_protected_json
|
||||
from utils.utils import colorText, get_sanitized_input
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
class LocalApprovalRequestor:
|
||||
"""Handles creation of local approval requests in Loxide."""
|
||||
|
||||
def getLocalApprovals(api: AirlockAPIWrapper):
|
||||
base_dir = get_base_directory
|
||||
result = api.otp_find_awaiting()
|
||||
local_approval = pd.DataFrame(result["response"]["otpusage"])
|
||||
if os.path.exists(f"{base_dir}\\cache\\newest_local_approval.parquet"):
|
||||
previous_run = pd.read_parquet(
|
||||
f"{base_dir}\\cache\\newest_local_approval.parquet"
|
||||
)
|
||||
previous_run.to_parquet(
|
||||
f"{base_dir}\\cache\\last_local_approval.parquet", index=False
|
||||
)
|
||||
os.remove(f"{base_dir}\\cache\\newest_local_approval.parquet")
|
||||
def __init__(self, api: AirlockAPIWrapper, username: str = None):
|
||||
"""
|
||||
Initialize the local approval requestor.
|
||||
|
||||
# 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"{base_dir}\\cache\\newest_local_approval.parquet", index=False
|
||||
Args:
|
||||
api: AirlockAPIWrapper instance
|
||||
username: Username creating the approvals (for tracking)
|
||||
"""
|
||||
self.api = api
|
||||
self.policy_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
self.username = (
|
||||
username or os.getenv("USERNAME") or os.getenv("USER") or "unknown"
|
||||
)
|
||||
|
||||
return local_approval
|
||||
def create_local_approval(
|
||||
self, agent_id: str, duration_minutes: int, batch_id: Optional[int] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Create a single local approval request.
|
||||
|
||||
Args:
|
||||
agent_id: Agent ID to create approval for
|
||||
duration_minutes: Duration of approval in minutes
|
||||
batch_id: Optional batch identifier (defaults to timestamp)
|
||||
|
||||
def scheduleAddingLAHashes(api: AirlockAPIWrapper):
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if batch_id is None:
|
||||
batch_id = int(time.time())
|
||||
|
||||
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)
|
||||
purpose = (
|
||||
f"🎫 Local Approval 🎫 - {duration_minutes} mins - "
|
||||
f"batch:{batch_id} Client:{agent_id} User:{self.username}"
|
||||
)
|
||||
|
||||
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)
|
||||
self.api.otp_generate(agent_id, duration_minutes, purpose)
|
||||
logger.info(
|
||||
f"Generated local approval for {agent_id}, batch {batch_id}, by {self.username}"
|
||||
)
|
||||
|
||||
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}"
|
||||
)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to process batch {batchid}: {e}")
|
||||
logger.error(f"Failed to generate local approval for {agent_id}: {e}")
|
||||
return False
|
||||
|
||||
def move_agent_to_audit(self, agent: Agent) -> bool:
|
||||
"""
|
||||
Move an agent to its corresponding audit policy.
|
||||
|
||||
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())
|
||||
Args:
|
||||
agent: Agent object to move
|
||||
|
||||
# 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(get_sanitized_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:
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
addLocalApproval(api, batch, duration_selected, agent.agentid)
|
||||
moveAgentToRelatedPolicy(api, agent, "audit")
|
||||
moveAgentToRelatedPolicy(self.api, agent, "audit")
|
||||
logger.info(f"Moved {agent.hostname} to audit policy")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(colorText(f"❌ Error processing agent {agent.hostname}: {e}", "red"))
|
||||
logger.error(f"Failed to move {agent.hostname} to audit: {e}")
|
||||
return False
|
||||
|
||||
def create_local_approval_batch(
|
||||
self,
|
||||
agents: List[Agent],
|
||||
duration_minutes: int,
|
||||
) -> tuple[int, int, int]:
|
||||
"""
|
||||
Create local approvals for multiple agents and move them to audit.
|
||||
|
||||
def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid):
|
||||
Args:
|
||||
agents: List of Agent objects
|
||||
duration_minutes: Duration of approval in minutes
|
||||
db_path: Optional path to database for history tracking
|
||||
|
||||
purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
|
||||
api.otp_generate(agentid, duration_selected, purpose)
|
||||
Returns:
|
||||
Tuple of (batch_id, success_count, failure_count)
|
||||
"""
|
||||
batch_id = int(time.time())
|
||||
success_count = 0
|
||||
failure_count = 0
|
||||
|
||||
print(colorText(f"\n📦 Processing batch {batch_id}...", "cyan"))
|
||||
print(colorText(f"👤 Requested by: {self.username}", "cyan"))
|
||||
print(
|
||||
colorText(f"📊 Moving {len(agents)} agent(s) to local approval\n", "cyan")
|
||||
)
|
||||
|
||||
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", "{}")
|
||||
for agent in agents:
|
||||
try:
|
||||
# Create local approval
|
||||
approval_success = self.create_local_approval(
|
||||
agent.agentid, duration_minutes, batch_id
|
||||
)
|
||||
|
||||
# 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()
|
||||
)
|
||||
if not approval_success:
|
||||
raise Exception("Failed to create local approval")
|
||||
|
||||
# Index last_agents by hostname for quick lookup
|
||||
last_agent_map = {agent.hostname: agent for agent in last_agents}
|
||||
# Move to audit policy
|
||||
move_success = self.move_agent_to_audit(agent)
|
||||
|
||||
# Result buckets
|
||||
newly_added = []
|
||||
same_policy = []
|
||||
moved_to_audit = []
|
||||
moved_to_enforcement = []
|
||||
unusual_move = []
|
||||
if not move_success:
|
||||
raise Exception("Failed to move to audit policy")
|
||||
|
||||
for current in current_agents:
|
||||
previous = last_agent_map.get(current.hostname)
|
||||
print(colorText(f"✓ {agent.hostname}", "green"))
|
||||
success_count += 1
|
||||
|
||||
if not previous:
|
||||
newly_added.append(current)
|
||||
continue
|
||||
except Exception as e:
|
||||
print(colorText(f"✗ {agent.hostname}: {e}", "red"))
|
||||
logger.error(f"Error processing agent {agent.hostname}: {e}")
|
||||
failure_count += 1
|
||||
|
||||
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 batch_id, success_count, failure_count
|
||||
|
||||
# Return all five DataFrames
|
||||
return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move
|
||||
def interactive_local_approval(self):
|
||||
"""
|
||||
Interactive workflow to create local approvals for selected agents.
|
||||
|
||||
This prompts the user to select a duration and agents, then creates
|
||||
the local approvals and moves agents to audit policies.
|
||||
"""
|
||||
# Duration options in minutes
|
||||
duration_options = [
|
||||
(15, "15 minutes"),
|
||||
(60, "1 hour"),
|
||||
(360, "6 hours"),
|
||||
(1440, "1 day"),
|
||||
(10080, "1 week"),
|
||||
]
|
||||
|
||||
def getNewLocalApprovals(api: AirlockAPIWrapper):
|
||||
# Display duration options
|
||||
print(colorText("\n⏱️ Select Local Approval Duration:", "white"))
|
||||
print(colorText("=" * 50, "white"))
|
||||
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
current_la = getLocalApprovals(api)
|
||||
for i, (minutes, label) in enumerate(duration_options, start=1):
|
||||
print(f" {i}. {label} ({minutes} minutes)")
|
||||
|
||||
# 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)
|
||||
print(colorText("=" * 50, "white"))
|
||||
|
||||
# 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)
|
||||
# Get user selection
|
||||
try:
|
||||
choice = int(get_sanitized_input("\nEnter the number of your choice: "))
|
||||
|
||||
# Find new entries
|
||||
new_entries = current_la[~current_la["key"].isin(old_la["key"])]
|
||||
if 1 <= choice <= len(duration_options):
|
||||
duration_minutes, duration_label = duration_options[choice - 1]
|
||||
print(colorText(f"✓ Selected: {duration_label}", "green"))
|
||||
logger.info(f"User selected duration: {duration_minutes} minutes")
|
||||
else:
|
||||
print(colorText("❌ Invalid choice.", "red"))
|
||||
logger.warning("Invalid duration choice")
|
||||
return
|
||||
|
||||
# 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]
|
||||
except ValueError:
|
||||
print(colorText("❌ Invalid input. Please enter a number.", "red"))
|
||||
logger.warning("Invalid input for duration selection")
|
||||
return
|
||||
|
||||
# Save current approvals for next run
|
||||
current_la.drop(columns=["key"], inplace=True)
|
||||
current_la.to_parquet(old_la_path, index=False)
|
||||
# Select agents
|
||||
print(colorText("\n🎯 Select Agents for Local Approval:", "white"))
|
||||
agents = selectAgents(self.api)
|
||||
|
||||
return recent_entries
|
||||
if not agents:
|
||||
print(colorText("❌ No agents found or error retrieving agents.", "red"))
|
||||
logger.warning("No agents selected or error retrieving agents")
|
||||
return
|
||||
|
||||
# Confirm with user
|
||||
print(colorText("\n📋 Summary:", "cyan"))
|
||||
print(colorText(f" Duration: {duration_label}", "white"))
|
||||
print(colorText(f" Agents: {len(agents)}", "white"))
|
||||
|
||||
confirm = get_sanitized_input("\nProceed? (y/n): ").lower()
|
||||
|
||||
if confirm != "y":
|
||||
print(colorText("❌ Operation cancelled.", "yellow"))
|
||||
return
|
||||
|
||||
# Process the batch
|
||||
batch_id, success_count, failure_count = self.create_local_approval_batch(
|
||||
agents, duration_minutes
|
||||
)
|
||||
|
||||
# Display summary
|
||||
self._display_summary(batch_id, duration_label, success_count, failure_count)
|
||||
|
||||
def _display_summary(
|
||||
self, batch_id: int, duration_label: str, success_count: int, failure_count: int
|
||||
):
|
||||
"""
|
||||
Display operation summary.
|
||||
|
||||
Args:
|
||||
batch_id: Batch identifier
|
||||
duration_label: Human-readable duration
|
||||
success_count: Number of successful operations
|
||||
failure_count: Number of failed operations
|
||||
"""
|
||||
print(colorText(f"\n{'=' * 60}", "white"))
|
||||
print(colorText("📊 Local Approval Summary", "cyan"))
|
||||
print(colorText("=" * 60, "white"))
|
||||
|
||||
print(colorText(f"✓ Successfully processed: {success_count}", "green"))
|
||||
|
||||
if failure_count > 0:
|
||||
print(colorText(f"✗ Failed: {failure_count}", "red"))
|
||||
|
||||
print(colorText(f"\n📦 Batch ID: {batch_id}", "cyan"))
|
||||
print(colorText(f"⏱️ Duration: {duration_label}", "cyan"))
|
||||
|
||||
print(colorText("=" * 60, "white"))
|
||||
print(colorText("\n💡 Next Steps:", "yellow"))
|
||||
print(colorText(" • Agents have been moved to audit policies", "white"))
|
||||
print(colorText(" • Local approvals are active", "white"))
|
||||
print(
|
||||
colorText(
|
||||
f" • Agents will return to enforcement after {duration_label}", "white"
|
||||
)
|
||||
)
|
||||
print(colorText("=" * 60 + "\n", "white"))
|
||||
|
||||
Reference in New Issue
Block a user