mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Added Move Agent, UI Tweaks

Browse Source
This commit is contained in:
Zarithas
2025-10-29 16:48:19 -04:00
parent 07021fe072
commit 9e8da97ad9
7 changed files with 230 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Rust/services.rs
+65
View File
@@ -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<PyAny>) {
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<PyAny>, 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
}
flows/quietAgent.py
+59 -55
View File
@@ -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")
services/agenthandler.py
+37 -1
View File
@@ -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
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")
services/security.py
+5 -5
View File
@@ -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:
utils/menus.py
+61 -34
View File
@@ -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: ")
utils/selector.py
+2 -2
View File
@@ -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]
print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
return []
utils/utils.py
+1
View File
@@ -166,6 +166,7 @@ def displayIntro():
"cyan",
)
)
def welcome():
print(
colorText(
"=================================================================================",