From 9e8da97ad94b73383da53e9b9dd81bb9be632a0d Mon Sep 17 00:00:00 2001 From: Zarithas Date: Wed, 29 Oct 2025 16:48:19 -0400 Subject: [PATCH] Added Move Agent, UI Tweaks --- Rust/services.rs | 65 ++++++++++++++++++++++ flows/quietAgent.py | 114 ++++++++++++++++++++------------------- services/agenthandler.py | 38 ++++++++++++- services/security.py | 10 ++-- utils/menus.py | 95 ++++++++++++++++++++------------ utils/selector.py | 4 +- utils/utils.py | 1 + 7 files changed, 230 insertions(+), 97 deletions(-) create mode 100644 Rust/services.rs diff --git a/Rust/services.rs b/Rust/services.rs new file mode 100644 index 0000000..9ba17ce --- /dev/null +++ b/Rust/services.rs @@ -0,0 +1,65 @@ +use std::str::FromStr; + +use pyo3::prelude::*; +use reqwest::{Client, header::{HeaderMap, HeaderName, HeaderValue}}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +#[pyfunction] +pub fn pull_policy_exec_histories() { + println!("Hello World from Rust!"); +} + +#[pyfunction] +pub fn api(py: Python<'_>, x: Py) { + let base_url: String = x.getattr(py, "base_url").unwrap().to_string(); + println!("{}", base_url); +} + +#[tokio::main] +#[pyfunction] +pub async fn history_logging(py: Python<'_>, py_self: Py, exec_types: String, checkpoint_number: String, policy_names: String ) -> String { + #[derive(Serialize, Debug)] + struct Payload { + checkpoint: String, + } + let base_url = py_self.getattr(py, "base_url").unwrap().to_string(); + let headers = py_self.getattr(py, "headers").unwrap().to_string(); + let headers_replace = headers.replace('\'', "\""); + let parsed: Value = serde_json::from_str(&headers_replace.as_str()).unwrap(); + let mut header_map = HeaderMap::new(); + if let Some(obj) = parsed.as_object() { + for (key, value) in obj { + if let Some(v) = value.as_str() { + let val = HeaderValue::from_str(v).unwrap(); + header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val); + } + } + } + let exec_types = exec_types.replace(" ", ""); + let payload_dict: Payload = Payload { + checkpoint:exec_types.to_string(), + }; + println!("{:?}", payload_dict); + let client = Client::builder() + .danger_accept_invalid_certs(true) + .default_headers(header_map) + .timeout(std::time::Duration::from_secs(30)) + .build() + .unwrap(); + let res = client + .post(format!("{}/v1/logging/exechistories", base_url)) + .json(&payload_dict) + .send() + .await; + match res { + Ok(res) => { + println!("{:?}", res.text().await.unwrap()); + } + Err(e) => { + eprintln!("{}", e); + } + } + let testingstring: String = "Testing".to_string(); + return testingstring +} \ No newline at end of file diff --git a/flows/quietAgent.py b/flows/quietAgent.py index 6e8a642..fd5e1a3 100644 --- a/flows/quietAgent.py +++ b/flows/quietAgent.py @@ -23,7 +23,7 @@ from flows.prepPolicy import selectPolicies from services.API import AirlockAPIWrapper from services.policyhandler import getPolicyInfo from utils.selector import Selector -from utils.utils import colorText, load_env +from utils.utils import colorText, get_sanitized_input, load_env logger = logging.getLogger(__name__) @@ -50,75 +50,79 @@ def findQuietAgents(api: AirlockAPIWrapper): valid_range=(1, 150), ) + confirm = Selector.confirm(f"Do you wish to proceed to pull history for {selected_policy[0].name}? Y/N : ") # Get execution history as a DataFrame - policy_exec_history = getPolicyInfo( + if confirm: + 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 + if policy_exec_history.empty: + logging.info("No execution history found for the selected policy and time range.") + get_sanitized_input("Press enter to continue") + 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 - ) + # 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) + # 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 - ) + # 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() + # 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) + # 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" - ) + # 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"] - ) + # 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 - ) + # 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]) + # 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) + # 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 + # 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 + # 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") + 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") + get_sanitized_input("Press enter to continue") diff --git a/services/agenthandler.py b/services/agenthandler.py index fcc455e..5a89991 100644 --- a/services/agenthandler.py +++ b/services/agenthandler.py @@ -24,6 +24,7 @@ from typing import List import pandas as pd +from flows.prepPolicy import selectPolicies from models.agent import Agent from models.policy import Policy from services.API import AirlockAPIWrapper @@ -241,6 +242,7 @@ def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']: enrich_agents(matched_agents, policies) return matched_agents + def moveAgentToRelatedPolicy( api: AirlockAPIWrapper, agent: Agent, @@ -284,4 +286,38 @@ def moveAgentToRelatedPolicy( return result = api.agent_move(agent.agentid, target_policy) - return result \ No newline at end of file + return result + + +def toggleEnforcement(api: AirlockAPIWrapper): + choices = ["Audit", "Enforcement", "Exit"] + print(colorText("Move devices to which state?:", "yellow")) + direction = Selector.select_string(choices, False, False) + if direction == "Exit": + pass + else: + devices = selectAgents(api) + for device in devices: + print(device.hostname) + confirm = Selector.confirm("Would you like to continue with these devices? Y/N: ") + if direction and devices and confirm: + for device in devices: + result = moveAgentToRelatedPolicy(api,device, str(direction).lower()) + logger.info(f"{device.hostname}: result: {result}") + get_sanitized_input("Press enter to continue") + +def moveAgents(api: AirlockAPIWrapper): + devices = selectAgents(api) + for device in devices: + print(device.hostname) + confirm_devices = Selector.confirm("Would you like to continue with these devices? Y/N: ") + if devices and confirm_devices: + policies = selectPolicies(api, False) + confirm_move = Selector.confirm(f"Would you like to move these devices to {policies[0].name}?") + if confirm_move: + for device in devices: + result = api.agent_move(device.agentid, policies[0].groupid) + logger.info(f"{device.hostname}: result: {result}") + else: + logger.info("Exiting without change") + get_sanitized_input("Press enter to continue") \ No newline at end of file diff --git a/services/security.py b/services/security.py index bbf2bf4..7574824 100644 --- a/services/security.py +++ b/services/security.py @@ -135,15 +135,15 @@ def getAPI(USERNAME, SERVICE_NAME): raise ValueError("Failed to retrieve API key after 3 incorrect attempts.") else: logging.warning(f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'.") - api_key = getpass(f"🗝️ No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip() + api_key = getpass(f"No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip() print("Please exit and relaunch program after saving your credential to avoid errors") while True: - password = getpass("🔓 Create a password to encrypt your API key: ") - confirm_password = getpass("🔒 Confirm your password: ") + password = getpass("Create a password to encrypt your API key: ") + confirm_password = getpass("Confirm your password: ") if password != confirm_password: - logging.warning("❌ Passwords do not match. Try again.") + logging.warning("Passwords do not match. Try again.") continue if check_password_complexity(password): @@ -155,7 +155,7 @@ def getAPI(USERNAME, SERVICE_NAME): logging.error(f"Failed to store API key: {e}") break else: - logging.warning("❌ Password does not meet complexity requirements. Try again.") + logging.warning("Password does not meet complexity requirements. Try again.") class APIKeyManager: diff --git a/utils/menus.py b/utils/menus.py index ba68a7d..2363a58 100644 --- a/utils/menus.py +++ b/utils/menus.py @@ -29,10 +29,13 @@ from flows.prepPolicy import ( testChange, ) from flows.quietAgent import findQuietAgents -from services.agenthandler import findAgents, moveAgentToRelatedPolicy, selectAgents +from services.agenthandler import ( + findAgents, + moveAgents, + toggleEnforcement, +) from services.API import AirlockAPIWrapper from utils.configmanager import load_env -from utils.selector import Selector from utils.utils import ( areYouSure, clear_screen, @@ -42,6 +45,7 @@ from utils.utils import ( locked, open_directory, printEnforceChecklist, + welcome, ) logger = logging.getLogger(__name__) @@ -54,45 +58,27 @@ def menu_main(api: AirlockAPIWrapper): while True: clear_screen() displayIntro() - # Add Settings, and give option to change working dir - print(colorText("1. ✅ - Move Device(s) to local approval (Placeholder)", "yellow")) - print(colorText("2. 🎫 - OTP", "yellow")) - print(colorText("3. 🔄 - Move to Audit/Enforcement", "yellow")) - print(colorText("4. 🔍 - Device Search", "yellow")) - print(colorText("5. 🔇 - Find Quiet Hosts", "yellow")) - if extras == "POLICYPREP" : print(colorText("6. 🛡️ - Policy Enforcement Tools", "yellow")) + welcome() + print(colorText("1. 🔍 - Device Search", "yellow")) + print(colorText("2. 🔀 - Move Device(s)", "yellow")) + print(colorText("3. 🎫 - One Time Pass (OTP)", "yellow")) + print(colorText("4. 🔇 - Find Quiet Hosts", "yellow")) + if extras == "POLICYPREP" : print(colorText("5. 🛡️ - Policy Enforcement Tools", "yellow")) + print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan")) print(colorText("F. 📂 - Open Working Directory", "yellow")) print(colorText("S. 🛠️ - Settings", "yellow")) print(colorText("Q. 🔚 - Quit", "yellow")) choice = get_sanitized_input("\nEnter Menu Item: ") if choice == "1": - print("This Feature is still in development") - get_sanitized_input("Press enter to continue") + findAgents(api,False) elif choice == "2": - menu_otp(api) + menu_move(api) elif choice == "3": - choices = ["audit", "enforcement", "Exit"] - print(colorText("Move devices to which state?:", "yellow")) - direction = Selector.select_string(choices, False, False) - if direction == "Exit": - pass - else: - devices = selectAgents(api) - print(colorText("Would you like to continue with these devices?","white")) - for device in devices: - print(device.hostname) - confirm = Selector.confirm() - if direction and devices and confirm: - for device in devices: - result = moveAgentToRelatedPolicy(api,device, str(direction)) - logger.info(f"{device.hostname}: result: {result}") - get_sanitized_input("Press enter to continue") + menu_otp(api) elif choice == "4": - findAgents(api,False) + findQuietAgents(api) elif choice == "5": - findQuietAgents(api) - elif choice == "6": if extras == "POLICYPREP": menu_policymanagment(api) elif choice.upper() == "F": open_directory(working_dir) @@ -214,14 +200,47 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 else: print(colorText("Invalid choice. Please try again.", "red")) +def menu_move(api: AirlockAPIWrapper): + working_dir = load_env("WORKING_DIR") + while True: + clear_screen() + displayIntro() + print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= ↔️ Agent Movement ↔️ =-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan")) + print(colorText("1. ✅ - Move to local approval (Placeholder)", "yellow")) + print(colorText("2. 🔄 - Move to Audit/Enforcement", "yellow")) + print(colorText("3. 🔀 - Move - Other", "yellow")) + print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan")) + print(colorText("F. 📂 - Open Working Directory", "yellow")) + print(colorText("S. 🛠️ - Settings", "yellow")) + print(colorText("B. 🔙 - Back", "yellow")) + + choice = get_sanitized_input("Enter your choice: ") + + if choice == "1": + print("This Feature is still in development") + get_sanitized_input("Press enter to continue") + elif choice == "2": + toggleEnforcement(api) + elif choice == "3": + moveAgents(api) + elif choice.upper() == "F": + open_directory(working_dir) + elif choice.upper() == "S": + menu_settings() + elif choice.upper() == "B": + break + + def menu_otp(api: AirlockAPIWrapper): working_dir = load_env("WORKING_DIR") while True: - - print(colorText("\n--- 🎫 OTP Submenu 🎫 ---", "cyan")) + clear_screen() + displayIntro() + print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 🎫 OTP 🎫 =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan")) print(colorText("1. 🔐 -Generate OTPs", "cyan")) print(colorText("2. 📊 -OTP Activities By Agent", "cyan")) print(colorText("3. ❌ -Revoke OTPs", "cyan")) + print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan")) print(colorText("F. 📂 - Open Working Directory", "yellow")) print(colorText("S. 🛠️ - Settings", "yellow")) print(colorText("B. 🔙 - Back", "yellow")) @@ -247,8 +266,12 @@ def menu_otp(api: AirlockAPIWrapper): def menu_policymanagment(api: AirlockAPIWrapper): working_dir = load_env("WORKING_DIR") while True: + clear_screen() + displayIntro() + print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 🛡️ Policy Tools 🛡️ =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan")) print(colorText("1. 🔒 - Prepare Policy For Enforcement", "yellow")) print(colorText("2. 🔄 - Update Audit Policies from Enforcement Policies", "yellow")) + print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan")) print(colorText("F. 📂 - Open Working Directory", "yellow")) print(colorText("S. 🛠️ - Settings", "yellow")) print(colorText("B. 🔙 - Back", "yellow")) @@ -272,8 +295,12 @@ def menu_policymanagment(api: AirlockAPIWrapper): print(colorText("Invalid choice. Please try again.", "red")) def menu_settings(): while True: - print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan")) + clear_screen() + displayIntro() + print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 🛠️ Settings 🛠️ =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan")) print(colorText("This Feature is still in development, if you have ideas for options you would like to see, let us know.", "cyan")) + print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan")) + # print(colorText("2. Sub-option B","cyan")) print(colorText("B. 🔙 - Back", "yellow")) choice = get_sanitized_input("Enter your choice: ") diff --git a/utils/selector.py b/utils/selector.py index 4bc6b76..8659a41 100644 --- a/utils/selector.py +++ b/utils/selector.py @@ -316,5 +316,5 @@ class Selector: print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow")) return [pd.Series(row) for row in items if row not in selected] else: - print(colorText("⚠️ Invalid mode. Returning all rows.", "yellow")) - return [pd.Series(row) for row in items] \ No newline at end of file + print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow")) + return [] \ No newline at end of file diff --git a/utils/utils.py b/utils/utils.py index de52cf1..76df8fe 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -166,6 +166,7 @@ def displayIntro(): "cyan", ) ) +def welcome(): print( colorText( "=================================================================================",