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

feat: add version updater and statistics enhancements (fixes #29)

Browse Source 
- Implemented version checking system with update notifications
- Integrated Git for fetching and downloading the latest version
- Added statistics updates
- Removed unused code across the project
- Condensed project structure
- Updated README
- Cleaned up UI
This commit is contained in:
Zarithas
2025-12-19 16:31:36 -05:00
parent 09d2c125cd
commit e1e0cb7ac7
26 changed files with 1906 additions and 3806 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Loxide.py
+465 -151
View File
@@ -30,10 +30,10 @@ from typing import Optional
from bson import ObjectId
import dotenv
import plotext as plt
from textual.app import App, ComposeResult
from textual.containers import Horizontal, Vertical
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.message import Message
from textual.reactive import reactive
from textual.screen import Screen
from textual.widgets import (
Button,
@@ -51,7 +51,6 @@ import airlock_libs
from models.agent import Agent
from models.policy import Policy
from services.API import AirlockAPIWrapper
from services.security import getAPI
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
from TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
from TUI.Screens.otpactivityscreen import OTPActivitiesScreen
@@ -59,19 +58,19 @@ from TUI.Screens.otprevokescreen import OTPRevokeScreen
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
from TUI.Screens.policyprepworkflowscreen import PolicyPrepWorkflowScreen
from TUI.Screens.quietagentworkflowscreen import QuietAgentWorkflowScreen
from TUI.Themes.theme_amber_terminal import get_amber_terminal_theme
from TUI.Themes.theme_retro_terminal import get_retro_terminal_theme
from TUI.Themes.themeselector import ThemeSelector
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
from TUI.Widgets.multiagentselector import MultiAgentSelector
from TUI.Widgets.policytreewidget import PolicyTreeWidget
from TUI.Widgets.resultsdisplay import ResultsDisplay
from TUI.Widgets.serverlogwidget import ServerLogWidget
from TUI.Widgets.settingswidget import SettingsWidget
from utils.configmanager import (
get_system_value,
get_user_value,
load_env,
save_user_config,
)
from utils.security import getAPI
from utils.setup import get_base_directory, setup
from utils.utils import irtang, open_directory
@@ -87,6 +86,81 @@ _APP_RESTART_REASON = None
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Helper for plotext charts
# ---------------------------------------------------------------------------
def create_plotext_chart(
chart_type: str,
labels: list,
values: list,
height: int = 10,
width: int = 60,
title: str = "",
color: str = None,
) -> Static:
"""Create a plotext chart and return it as a Static widget with the rendered output."""
# Completely reset plotext state
plt.clear_data()
plt.clear_figure()
plt.clear_color()
plt.clear_terminal()
plt.clf()
# Set theme and canvas size
plt.theme("clear")
plt.plotsize(width, height)
if chart_type == "simple_bar":
# Simple horizontal bar - clean text-based look
if color:
plt.simple_bar(labels, values, width=width, color=color)
else:
plt.simple_bar(labels, values, width=width)
elif chart_type == "bar_h":
# Calculate integer ticks for bar charts
max_val = max(values) if values else 0
if max_val <= 5:
ticks = list(range(0, max_val + 1))
elif max_val <= 10:
ticks = list(range(0, max_val + 1, 2))
else:
step = max(1, int(max_val / 5))
ticks = list(range(0, max_val + step, step))
plt.bar(labels, values, orientation="h", width=0.5)
plt.xlabel("Count")
plt.xticks(ticks)
plt.grid(False, False)
elif chart_type == "bar_v":
max_val = max(values) if values else 0
if max_val <= 5:
ticks = list(range(0, max_val + 1))
elif max_val <= 10:
ticks = list(range(0, max_val + 1, 2))
else:
step = max(1, int(max_val / 5))
ticks = list(range(0, max_val + step, step))
plt.bar(labels, values, orientation="v", width=0.5)
plt.ylabel("Count")
plt.yticks(ticks)
plt.grid(False, False)
if title:
plt.title(title)
# Use build() to get the chart string
chart_output = plt.build()
# Clear state after building
plt.clf()
plt.clear_data()
# Create Static widget with the chart - disable selection
chart_widget = Static(chart_output, classes="chart_display")
chart_widget.can_focus = False
return chart_widget
# ---------------------------------------------------------------------------
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
# ---------------------------------------------------------------------------
@@ -110,31 +184,30 @@ def _persist_user_theme(theme_name: str) -> None:
# ---------------------------------------------------------------------------
class MainMenuScreen(Screen):
api: AirlockAPIWrapper
current_tab = reactive("")
BUTTON_DEFS = {
"agent_actions": [
{
"label": "🖥️ - Multi-Agent Operations",
"label": "🔧 - Multi-Agent Operations",
"id": "move_agent_workflow_button",
"description": "Select agents to: Move policies, Generate OTPs, Toggle audit/enforcement, View history, Export data",
"description": "Select agents to:\n • Move policies\n • Generate OTPs\n • Toggle audit/enforcement\n • View history\n • Export data",
},
{
"label": "🎫 - Review and approve OTP Activities",
"label": "🎫 - Review and approve OTP Activities",
"id": "otp_activities_button",
},
{
"label": "🛑 - Revoke Active OTP Session",
"label": "🛑 - Revoke Active OTP Session",
"id": "otp_revoke_button",
},
],
"policy": [
{
"label": "⚖️ - Prepare Policy For Enforcement",
"label": "⚙️ - Prepare Policy For Enforcement",
"id": "policy_prep_button",
},
{
"label": "🔕 - Find and Move Quiet Hosts to Enforcement",
"label": "🔍 - Find and Move Quiet Hosts to Enforcement",
"id": "find_quiet_button",
},
],
@@ -154,35 +227,81 @@ class MainMenuScreen(Screen):
def _make_buttons_for(self, tab_id: str) -> Vertical:
defs = self.BUTTON_DEFS.get(tab_id, [])
widgets = []
# Special layout for agent_actions tab
if tab_id == "agent_actions":
left_col = Vertical()
left_col.styles.width = "1fr"
right_col = Vertical()
right_col.styles.width = "1fr"
for item in defs:
# Support both old tuple format and new dict format
if isinstance(item, dict):
label = item["label"]
btn_id = item["id"]
description = item.get("description")
else:
# Old tuple format: (label, id)
label, btn_id = item
description = None
btn = Button(label, id=btn_id)
btn.styles.width = "100%"
widgets.append(btn)
btn.styles.margin = (0, 1, 1, 0)
btn_container = Vertical()
btn_container.styles.height = "auto"
btn_container.compose_add_child(btn)
# Add description text if provided
if description:
desc_text = Static(description, classes="button_description")
desc_text.styles.width = "100%"
desc_text.styles.color = "ansi_bright_black"
desc_text.styles.text_align = "center"
desc_text.styles.text_align = "left"
desc_text.styles.margin = (0, 0, 1, 0)
btn_container.compose_add_child(desc_text)
# Multi-Agent on left, OTP buttons on right
if "otp" in btn_id:
right_col.compose_add_child(btn_container)
else:
left_col.compose_add_child(btn_container)
row = Horizontal()
row.styles.width = "100%"
row.styles.height = "auto"
row.compose_add_child(left_col)
row.compose_add_child(right_col)
return Vertical(row)
# Default single-column layout for other tabs
widgets = []
for item in defs:
if isinstance(item, dict):
label = item["label"]
btn_id = item["id"]
description = item.get("description")
else:
label, btn_id = item
description = None
btn = Button(label, id=btn_id)
btn.styles.width = "100%"
btn.styles.margin = (0, 1, 1, 0)
widgets.append(btn)
if description:
desc_text = Static(description, classes="button_description")
desc_text.styles.width = "100%"
desc_text.styles.color = "ansi_bright_black"
desc_text.styles.text_align = "left"
desc_text.styles.margin = (0, 0, 1, 0)
widgets.append(desc_text)
return Vertical(*widgets)
def compose(self) -> ComposeResult:
yield Header(show_clock=True, icon="âš™")
yield Header(show_clock=True, icon="⚙")
tabs = [
Tab("Agents", id="agent_actions"),
@@ -203,6 +322,50 @@ class MainMenuScreen(Screen):
def on_mount(self) -> None:
self.switch_tab("agent_actions")
# Delay version check to ensure UI is fully ready for notifications
if hasattr(self.app, "_version_checker"):
self.set_timer(1.0, self._do_version_check)
def _do_version_check(self) -> None:
"""Perform version check and show notification if update available."""
logger.info("_do_version_check called")
if not hasattr(self.app, "_version_checker"):
logger.warning(
"No _version_checker on app - was create_update_notifier called?"
)
return
checker = self.app._version_checker
logger.info(f"Checker exists: {checker}")
logger.info("Calling check_now(force=True)...")
try:
result = checker.check_now(force=True)
logger.info(
f"Version check result: update_available={result.update_available}, latest={result.latest_version}, current={result.current_version}, error={result.error}"
)
if result.error:
logger.warning(f"Version check returned error: {result.error}")
return
if result.update_available:
# Always nag on startup - don't check dismissed version here
# User can dismiss from settings but we still want startup reminder
logger.info(f"Showing toast for update {result.latest_version}")
msg = f"🆕 Update available: {result.latest_version}\nGo to Settings to download"
self.app.notify(
msg, title="Loxide Update Available", severity="warning", timeout=15
)
logger.info("Toast notification called")
else:
logger.info(
f"No update available - current {result.current_version} is latest"
)
except Exception as e:
logger.error(f"Version check failed: {e}", exc_info=True)
def on_key(self, event) -> None:
"""Handle up/down arrow keys for button navigation."""
if event.key == "down":
@@ -252,7 +415,6 @@ class MainMenuScreen(Screen):
buttons[new_index].focus()
def switch_tab(self, tab_id: str) -> None:
self.current_tab = tab_id
content = self.query_one("#content", Vertical)
content.remove_children()
@@ -267,14 +429,13 @@ class MainMenuScreen(Screen):
elif tab_id == "statistics":
content.mount(self._create_statistics_widget())
elif tab_id == "settings":
content.mount(ThemeSelector())
content.mount(SettingsWidget())
else:
content.mount(Static(f"Unknown tab: {tab_id}"))
def _create_statistics_widget(self) -> Vertical:
"""Create the enhanced statistics display widget with time period selection."""
stats_container = Vertical(id="statistics_container")
# Build all widgets first, then mount them all at once
widgets_to_mount = []
@@ -306,41 +467,56 @@ class MainMenuScreen(Screen):
)
time_select.styles.width = 20
# Refresh button
refresh_btn = Button("🔄 Refresh", id="stats_refresh_btn", variant="primary")
refresh_btn.styles.margin = (0, 0, 0, 2)
# Fetch/Refresh button - starts as "Fetch"
if self.current_stats_days in self.statistics_cache:
button_label = "🔍„ Refresh"
else:
button_label = "📊 Fetch Statistics"
fetch_btn = Button(button_label, id="stats_refresh_btn", variant="primary")
fetch_btn.styles.margin = (0, 0, 0, 2)
# Compose controls
controls.compose_add_child(time_label)
controls.compose_add_child(time_select)
controls.compose_add_child(refresh_btn)
controls.compose_add_child(fetch_btn)
widgets_to_mount.append(controls)
# Timestamp and status area
status_container = Vertical(id="stats_status_area")
status_container.styles.margin = (0, 2, 1, 2)
status_container.styles.margin = (0, 2, 0, 2)
widgets_to_mount.append(status_container)
# Data display area - using horizontal layout
data_container = Horizontal(id="stats_data_area")
# Data display area - scrollable vertical layout
data_container = VerticalScroll(id="stats_data_area")
data_container.styles.margin = (0, 2, 0, 2)
data_container.styles.height = "auto"
widgets_to_mount.append(data_container)
# Compose all widgets into the container
for widget in widgets_to_mount:
stats_container.compose_add_child(widget)
# Schedule data fetch after the widget is mounted
if self.current_stats_days not in self.statistics_cache:
self.call_after_refresh(
lambda: self._fetch_execution_statistics(self.current_stats_days)
)
else:
# If cached, display it. Otherwise show prompt to fetch
if self.current_stats_days in self.statistics_cache:
self.call_after_refresh(self._update_statistics_display)
else:
self.call_after_refresh(self._show_fetch_prompt)
return stats_container
def _show_fetch_prompt(self) -> None:
"""Show a prompt to click Fetch to load statistics."""
try:
status_area = self.query_one("#stats_status_area", Vertical)
status_area.remove_children()
prompt = Static("Click 'Fetch Statistics' to load execution data.")
prompt.styles.text_style = "italic"
prompt.styles.text_align = "center"
status_area.mount(prompt)
except Exception as e:
logger.warning(f"Could not show fetch prompt: {e}")
def _skipback(self, days: int) -> ObjectId:
"""Generate a MongoDB ObjectId for a given number of days ago."""
date_days_ago = datetime.now(UTC) - timedelta(days=days)
@@ -372,10 +548,9 @@ class MainMenuScreen(Screen):
)
logger.debug(f"Calling with exec_types={exec_types}, days={days}")
# Use pull_policy_exec_histories - must pass all 4 args in order: (api, policy_name, type, days)
# Empty string for policy_name means all policies
# Use pull_policy_exec_histories - pass None for policy_name to get all policies
execs = airlock_libs.pull_policy_exec_histories(
self.app.api, None, str(exec_types), days # Empty string = all policies
self.app.api, None, str(exec_types), days # None = all policies
)
logger.debug(f"Received response, length: {len(execs) if execs else 0}")
@@ -449,24 +624,41 @@ class MainMenuScreen(Screen):
18: "Trusted Browser Metadata Execution",
}
# Categorize execution types
# "Untrusted" types: contain "Untrusted" in name
untrusted_types = {2, 3, 13, 14}
# "Block" types: contain "Block" in name (Blocked, Blocklist)
block_types = {1, 6, 7, 12, 15, 16}
# Count execution types
type_counter = Counter()
policy_counter = Counter()
hostname_counter = Counter()
# Separate counters for untrusted vs block
policy_counter_untrusted = Counter()
policy_counter_block = Counter()
hostname_counter_untrusted = Counter()
hostname_counter_block = Counter()
for exec_record in executions:
exec_type = exec_record.get("type", -1)
type_counter[exec_type] += 1
policy_name = exec_record.get("policyname", "Unknown")
policy_counter[policy_name] += 1
hostname = exec_record.get("hostname", "Unknown")
hostname_counter[hostname] += 1
# Get top 5 policies and hostnames
top_policies = policy_counter.most_common(5)
top_hostnames = hostname_counter.most_common(5)
# Categorize by type
if exec_type in untrusted_types:
policy_counter_untrusted[policy_name] += 1
hostname_counter_untrusted[hostname] += 1
elif exec_type in block_types:
policy_counter_block[policy_name] += 1
hostname_counter_block[hostname] += 1
# Get top 5 for each category
top_policies_untrusted = policy_counter_untrusted.most_common(5)
top_policies_block = policy_counter_block.most_common(5)
top_hostnames_untrusted = hostname_counter_untrusted.most_common(5)
top_hostnames_block = hostname_counter_block.most_common(5)
# Format execution type counts
type_counts = []
@@ -477,22 +669,36 @@ class MainMenuScreen(Screen):
return {
"total_executions": len(executions),
"type_counts": type_counts,
"top_policies": top_policies,
"top_hostnames": top_hostnames,
"top_policies_untrusted": top_policies_untrusted,
"top_policies_block": top_policies_block,
"top_hostnames_untrusted": top_hostnames_untrusted,
"top_hostnames_block": top_hostnames_block,
}
def _update_statistics_display(self) -> None:
"""Update the statistics display with cached data using horizontal layout."""
"""Update the statistics display with cached data using vertical scrollable layout."""
try:
data_area = self.query_one("#stats_data_area", Horizontal)
data_area = self.query_one("#stats_data_area", VerticalScroll)
status_area = self.query_one("#stats_status_area", Vertical)
except Exception as e:
logger.warning(f"Could not find statistics display areas: {e}")
return
# Clear any existing plotext state
plt.clear_data()
plt.clear_figure()
data_area.remove_children()
status_area.remove_children()
# Get theme color for charts
theme_name = (
self.app._textual_theme
if hasattr(self.app, "_textual_theme")
else "textual-dark"
)
chart_color = None # use plotext default
# Get cached data
cache_entry = self.statistics_cache.get(self.current_stats_days)
if not cache_entry:
@@ -511,109 +717,187 @@ class MainMenuScreen(Screen):
timestamp_label.styles.margin = (0, 0, 1, 0)
status_area.mount(timestamp_label)
# Create three columns for horizontal layout
# Column 1: Summary and Execution Types
col1 = Vertical()
col1.styles.width = "1fr"
col1.styles.padding = (0, 1, 0, 0)
# Create single scrollable column for better chart display
main_col = Vertical()
main_col.styles.width = "100%"
main_col.styles.height = "auto"
# Column 2: Top Policies
col2 = Vertical()
col2.styles.width = "1fr"
col2.styles.padding = (0, 1)
# Column 3: Top Machines and System Overview
col3 = Vertical()
col3.styles.width = "1fr"
col3.styles.padding = (0, 0, 0, 1)
# === COLUMN 1: Summary and Execution Types ===
total = Static(f"📊 Total Non-Trusted Executions: {stats['total_executions']}")
total.styles.text_style = "bold"
total.styles.margin = (0, 0, 1, 0)
col1.compose_add_child(total)
# Execution type breakdown
if stats["type_counts"]:
type_header = Static("📋 Execution Types")
type_header.styles.text_style = "bold"
type_header.styles.margin = (1, 0, 0, 0)
col1.compose_add_child(type_header)
for type_name, count in stats["type_counts"]:
# Abbreviate long type names
short_name = type_name.replace("Execution", "Exec").replace(
"[Audit]", "[Aud]"
)
type_line = Static(f" {short_name}: {count}")
type_line.styles.margin = (0, 0, 0, 1)
col1.compose_add_child(type_line)
# === COLUMN 2: Top Policies ===
if stats["top_policies"]:
policy_header = Static("⚖️ Top 5 Policies")
policy_header.styles.text_style = "bold"
policy_header.styles.margin = (0, 0, 0, 0)
col2.compose_add_child(policy_header)
for i, (policy_name, count) in enumerate(stats["top_policies"], 1):
policy_line = Static(f"{i}. {policy_name}: {count}")
policy_line.styles.margin = (0, 0, 0, 1)
col2.compose_add_child(policy_line)
# === COLUMN 3: Top Machines and System Overview ===
if stats["top_hostnames"]:
host_header = Static("🖥️ Top 5 Machines")
host_header.styles.text_style = "bold"
host_header.styles.margin = (0, 0, 0, 0)
col3.compose_add_child(host_header)
for i, (hostname, count) in enumerate(stats["top_hostnames"], 1):
host_line = Static(f"{i}. {hostname}: {count}")
host_line.styles.margin = (0, 0, 0, 1)
col3.compose_add_child(host_line)
# System Overview in column 3
# === SYSTEM OVERVIEW (First) ===
total_agents = len(self.app.devices) if self.app.devices else 0
total_policies = len(self.app.policies) if self.app.policies else 0
if self.app.devices:
enforced_agents = sum(
1 for agent in self.app.devices if getattr(agent, "enforcement", False)
)
audit_agents = sum(
1
for agent in self.app.devices
if not getattr(agent, "enforcement", False)
)
# status: 0=Offline, 1=Online, 2=Hidden, 3=Safemode
online_agents = sum(
1 for agent in self.app.devices if getattr(agent, "online", False)
1 for agent in self.app.devices if getattr(agent, "status", 0) == 1
)
offline_agents = sum(
1 for agent in self.app.devices if getattr(agent, "status", 0) == 0
)
hidden_agents = sum(
1 for agent in self.app.devices if getattr(agent, "status", 0) == 2
)
safemode_agents = sum(
1 for agent in self.app.devices if getattr(agent, "status", 0) == 3
)
offline_agents = total_agents - online_agents
else:
enforced_agents = audit_agents = online_agents = offline_agents = 0
online_agents = offline_agents = hidden_agents = safemode_agents = 0
system_header = Static("ℹ️ System Overview")
system_header.styles.text_style = "bold"
system_header.styles.margin = (2, 0, 0, 0)
col3.compose_add_child(system_header)
system_header.styles.margin = (0, 0, 1, 0)
main_col.compose_add_child(system_header)
system_data = [
f"Agents: {total_agents} ({online_agents} online)",
f"Policies: {total_policies}",
f"Enforced: {enforced_agents}",
f"Audit: {audit_agents}",
# Total counts
totals_line = Static(
f"Total Agents: {total_agents} Total Policies: {total_policies}"
)
totals_line.styles.margin = (0, 0, 1, 1)
main_col.compose_add_child(totals_line)
# Agent status with plotext bar chart and percentages
status_data = [
("Online", online_agents),
("Offline", offline_agents),
("Hidden", hidden_agents),
("Safemode", safemode_agents),
]
for line in system_data:
line_widget = Static(f" {line}")
line_widget.styles.margin = (0, 0, 0, 1)
col3.compose_add_child(line_widget)
labels = []
counts = []
for status_name, count in status_data:
percentage = int((count / total_agents) * 100) if total_agents > 0 else 0
labels.append(f"{status_name} ({percentage}%)")
counts.append(count)
# Mount all columns
data_area.mount(col1)
data_area.mount(col2)
data_area.mount(col3)
chart = create_plotext_chart(
"simple_bar", labels, counts, height=15, width=60, color=chart_color
)
chart.styles.margin = (0, 0, 1, 0)
main_col.compose_add_child(chart)
# === EXECUTION SUMMARY ===
total_execs = stats["total_executions"]
total = Static(f"📊 Total Untrusted Executions: {total_execs}")
total.styles.text_style = "bold"
total.styles.margin = (1, 0, 1, 0)
main_col.compose_add_child(total)
# Execution type breakdown with percentages
if stats["type_counts"]:
type_header = Static("📋 Execution Types")
type_header.styles.text_style = "bold"
type_header.styles.margin = (1, 0, 0, 0)
main_col.compose_add_child(type_header)
labels = []
counts = []
for type_name, count in stats["type_counts"]:
short_name = type_name.replace("Execution", "Exec").replace(
"[Audit]", "[Aud]"
)
if len(short_name) > 18:
short_name = short_name[:15] + "..."
percentage = int((count / total_execs) * 100) if total_execs > 0 else 0
labels.append(f"{short_name} ({percentage}%)")
counts.append(count)
chart = create_plotext_chart(
"simple_bar", labels, counts, height=15, width=60, color=chart_color
)
chart.styles.margin = (0, 0, 1, 0)
main_col.compose_add_child(chart)
# === TOP POLICIES - Untrusted ===
if stats["top_policies_untrusted"]:
policy_header = Static("⚖️ Most Active Policies (Untrusted)")
policy_header.styles.text_style = "bold"
policy_header.styles.margin = (1, 0, 0, 0)
main_col.compose_add_child(policy_header)
policy_total = sum(count for _, count in stats["top_policies_untrusted"])
labels = []
values = []
for policy_name, count in stats["top_policies_untrusted"]:
short_name = policy_name[:18] if len(policy_name) > 18 else policy_name
percentage = (
int((count / policy_total) * 100) if policy_total > 0 else 0
)
labels.append(f"{short_name} ({percentage}%)")
values.append(count)
chart = create_plotext_chart(
"simple_bar", labels, values, height=12, width=60, color=chart_color
)
chart.styles.margin = (0, 0, 1, 0)
main_col.compose_add_child(chart)
# === TOP POLICIES - Blocked ===
if stats["top_policies_block"]:
policy_header = Static("⚖️ Most Active Policies (Blocked)")
policy_header.styles.text_style = "bold"
policy_header.styles.margin = (1, 0, 0, 0)
main_col.compose_add_child(policy_header)
policy_total = sum(count for _, count in stats["top_policies_block"])
labels = []
values = []
for policy_name, count in stats["top_policies_block"]:
short_name = policy_name[:18] if len(policy_name) > 18 else policy_name
percentage = (
int((count / policy_total) * 100) if policy_total > 0 else 0
)
labels.append(f"{short_name} ({percentage}%)")
values.append(count)
chart = create_plotext_chart(
"simple_bar", labels, values, height=12, width=60, color=chart_color
)
chart.styles.margin = (0, 0, 1, 0)
main_col.compose_add_child(chart)
# === TOP MACHINES - Untrusted ===
if stats["top_hostnames_untrusted"]:
host_header = Static("🖥️ Most Active Machines (Untrusted)")
host_header.styles.text_style = "bold"
host_header.styles.margin = (1, 0, 0, 0)
main_col.compose_add_child(host_header)
host_total = sum(count for _, count in stats["top_hostnames_untrusted"])
labels = []
values = []
for hostname, count in stats["top_hostnames_untrusted"]:
short_name = hostname[:18] if len(hostname) > 18 else hostname
percentage = int((count / host_total) * 100) if host_total > 0 else 0
labels.append(f"{short_name} ({percentage}%)")
values.append(count)
chart = create_plotext_chart(
"simple_bar", labels, values, height=12, width=60, color=chart_color
)
chart.styles.margin = (0, 0, 1, 0)
main_col.compose_add_child(chart)
# === TOP MACHINES - Blocked ===
if stats["top_hostnames_block"]:
host_header = Static("🖥️ Most Active Machines (Blocked)")
host_header.styles.text_style = "bold"
host_header.styles.margin = (1, 0, 0, 0)
main_col.compose_add_child(host_header)
host_total = sum(count for _, count in stats["top_hostnames_block"])
labels = []
values = []
for hostname, count in stats["top_hostnames_block"]:
short_name = hostname[:18] if len(hostname) > 18 else hostname
percentage = int((count / host_total) * 100) if host_total > 0 else 0
labels.append(f"{short_name} ({percentage}%)")
values.append(count)
chart = create_plotext_chart(
"simple_bar", labels, values, height=12, width=60, color=chart_color
)
chart.styles.margin = (0, 0, 1, 0)
main_col.compose_add_child(chart)
# Mount the main scrollable column
data_area.mount(main_col)
def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
self.switch_tab(event.tab.id)
@@ -625,13 +909,23 @@ class MainMenuScreen(Screen):
if new_days != self.current_stats_days:
self.current_stats_days = new_days
# Fetch if not cached
if new_days not in self.statistics_cache:
self._fetch_execution_statistics(new_days)
else:
# Update button label based on cache status
try:
button = self.query_one("#stats_refresh_btn", Button)
if new_days in self.statistics_cache:
button.label = "🔍„ Refresh"
# Display cached data
self._update_statistics_display()
else:
button.label = "📊 Fetch Statistics"
# Show fetch prompt
self._show_fetch_prompt()
except Exception as e:
logger.warning(f"Could not update button or display: {e}")
def on_multi_agent_selector_agents_selected() -> None:
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
"""Handle selected agents from AgentSelector."""
global _APP_RESTART_REASON
selected_agents = message.selected_agents
@@ -640,10 +934,10 @@ class MainMenuScreen(Screen):
_APP_RESTART_REASON = ("multi_agent_action", selected_agents)
self.app.exit()
def on_theme_selector_theme_selected(
self, message: ThemeSelector.ThemeSelected
def on_settings_widget_theme_selected(
self, message: SettingsWidget.ThemeSelected
) -> None:
"""Handle theme selection from ThemeSelector."""
"""Handle theme selection from SettingsWidget."""
global _APP_RESTART_REASON
_persist_user_theme(message.theme_name)
_APP_RESTART_REASON = ("restart",)
@@ -713,7 +1007,7 @@ class MainMenuScreen(Screen):
logger.info("Toggling enforcement for device: %s", message.device.hostname)
try:
from services.agenthandler import moveAgentToRelatedPolicy
from TUI.Widgets.agentmoveoperations import moveAgentToRelatedPolicy
from utils.configmanager import get_system_json
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
@@ -768,6 +1062,12 @@ class MainMenuScreen(Screen):
# Handle statistics refresh button
if button_id == "stats_refresh_btn":
self._fetch_execution_statistics(self.current_stats_days)
# Update button label to "Refresh" after first fetch
try:
button = event.button
button.label = "🔍„ Refresh"
except Exception as e:
logger.warning(f"Could not update button label: {e}")
event.stop()
return
@@ -826,6 +1126,22 @@ class Loxide(App[Message]):
content-align: center middle;
text-align: center;
}
.chart_display {
overflow: hidden;
}
.chart_display:hover {
background: transparent;
}
#statistics_container {
height: 1fr;
}
#stats_status_area {
height: auto;
}
#stats_data_area {
height: 1fr;
scrollbar-gutter: stable;
}
"""
BINDINGS = [
("q", "quit", "Quit"),
@@ -867,8 +1183,6 @@ class Loxide(App[Message]):
self.devices = None
def on_mount(self, api: AirlockAPIWrapper) -> None:
self.register_theme(get_retro_terminal_theme())
self.register_theme(get_amber_terminal_theme())
self.theme = self._textual_theme
self.push_screen(MainMenuScreen())
TUI/Screens/quietagentworkflowscreen.py
+2 -2
View File
@@ -38,9 +38,9 @@ from textual.widgets import Button, DataTable, Footer, Header, Input, Static
from models.policy import Policy
from services.API import AirlockAPIWrapper
from services.policyhandler import getPolicyInfo
from TUI.Widgets.policyselector import PolicySelector
from utils.configmanager import load_env
from utils.executionfetcher import getExecutions
logger = logging.getLogger(__name__)
@@ -429,7 +429,7 @@ class QuietAgentWorkflowScreen(Screen):
return
# Get execution history (this shows progress bars in terminal via airlock_libs)
policy_exec_history = getPolicyInfo(
policy_exec_history = getExecutions(
self.api, self.selected_policy, [1, 2, 6, 7], self.history_days
)
TUI/Themes/theme_amber_terminal.py
-50
View File
@@ -1,50 +0,0 @@
# 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/>.
from textual.color import Color
from textual.theme import Theme
def get_amber_terminal_theme():
"""Amber CRT theme with compensated brightness for blending."""
return Theme(
name="amber-terminal",
background=Color.parse("#000000"), # pure black
primary=Color.parse("#ffb733"), # bright amber
secondary=Color.parse("#e69500"), # strong amber
success=Color.parse("#ffb733"),
warning=Color.parse("#ffff66"),
error=Color.parse("#ff3300"),
surface=Color.parse("#49331a"), # brighter brown for blending
)
AMBER_TERMINAL_CSS = """
Screen {
align: center middle;
background: #000000; /* force black */
color: #ffb733; /* force amber text */
}
.widget {
border: tall #ffb733; /* force amber border */
background: #3a1f00; /* compensated surface */
width: 80%;
}
* {
font-family: "Courier New", monospace;
}
"""
TUI/Themes/theme_retro_terminal.py
-53
View File
@@ -1,53 +0,0 @@
# 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/>.
from textual.color import Color
def get_retro_terminal_theme():
from textual.theme import Theme
return Theme(
name="retro-terminal",
background=Color.parse("#000000"),
primary=Color.parse("#00ff00"),
secondary=Color.parse("#00aa00"),
success=Color.parse("#00ff00"),
warning=Color.parse("#ffff00"),
error=Color.parse("#ff0000"),
surface=Color.parse("#071802"),
)
RETRO_TERMINAL_CSS = """
/* Retro terminal CRT effect */
Screen {
align: center middle;
background: $background;
color: $text;
}
/* Blocky, pixelated widgets */
.widget {
border: tall $primary;
background: $surface;
width: 80%;
}
/* Monospaced font */
* {
font-family: "Courier New", monospace;
}
"""
TUI/Widgets/OTP_generate.py
-1
View File
@@ -43,7 +43,6 @@ class OTPGenerator(Widget):
# Reactive properties to track form completion
requestor_filled = reactive(False)
reasoning_filled = reactive(False)
duration_selected = reactive(True) # Default is selected
otp_generated = reactive(False)
class OTPInfo(Message):
TUI/Widgets/agentmoveoperations.py
+77 -25
View File
@@ -29,14 +29,71 @@ from textual.widget import Widget
from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea
from models.agent import Agent
from services.API import AirlockAPIWrapper
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
from TUI.Screens.policyselectorscreen import PolicySelectorScreen
from TUI.Widgets.OTP_generate import OTPGenerator
from utils.configmanager import get_system_json
logger = logging.getLogger(__name__)
def moveAgentToRelatedPolicy(
api: AirlockAPIWrapper,
agent: Agent,
mode: str = "audit",
):
"""
Moves an agent between audit and enforcement policies based on the mode.
Args:
api: AirlockAPIWrapper instance.
agent: Agent object.
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
"""
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
if mode == "audit":
if agent.groupid in policy_relationship_map:
target_policy = policy_relationship_map[agent.groupid]
elif agent.groupid in policy_relationship_map.values():
logger.debug(
f"Agent {agent.hostname} is already in an audit group. No action needed."
)
print(
f"Agent {agent.hostname} is already in an audit group. No action needed."
)
return
else:
logger.warning(
f"Error: No corresponding audit policy found for groupid: {agent.groupid}."
)
return
elif mode == "enforcement":
inverse_map = {v: k for k, v in policy_relationship_map.items()}
if agent.groupid in inverse_map:
target_policy = inverse_map[agent.groupid]
elif agent.groupid in inverse_map.values():
logger.info(
f"Agent {agent.hostname} is already in an enforcement group. No action needed."
)
return
else:
logger.warning(
f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}."
)
return
else:
logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.")
return
result = api.agent_move(agent.agentid, target_policy)
return result
class AgentMoveOperations(Widget):
"""
A Textual widget for managing bulk agent operations and policy migrations.
@@ -216,11 +273,11 @@ class AgentMoveOperations(Widget):
results_lines.append(" (none)")
results_lines.append("")
results_lines.append(f"❌ Failed ({len(unsuccessful)}):")
results_lines.append(f"❌ Failed ({len(unsuccessful)}):")
if unsuccessful:
for agent, error in unsuccessful:
results_lines.append(f" ❌ {agent.hostname}: {error}")
results_lines.append(f" ❌ {agent.hostname}: {error}")
else:
results_lines.append(" (none)")
@@ -255,9 +312,9 @@ class AgentMoveOperations(Widget):
- Operations panel: 1/3 width
- Results area: Initially hidden, shown after operation completion
"""
yield Header(show_clock=True, icon="⚙️")
yield Header(show_clock=True, icon="âš™❗")
title_text = Static(
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
f"🖥❗ Agent Operations - {len(self.agents)} device(s) selected",
id="move_ops_title",
)
title_text.styles.margin = (0, 0, 1, 0)
@@ -292,39 +349,39 @@ class AgentMoveOperations(Widget):
yield operations_label
# Operation buttons
export_csv_btn = Button("📄 Export CSV", id="export_csv_btn")
export_csv_btn = Button("📄 Export CSV", id="export_csv_btn")
export_csv_btn.styles.width = "100%"
export_csv_btn.styles.margin = (0, 0, 1, 0)
yield export_csv_btn
local_approval_btn = Button(
"✔️ Local Approval Mode", id="local_approval_btn"
"✔❗ Local Approval Mode", id="local_approval_btn"
)
local_approval_btn.styles.width = "100%"
local_approval_btn.styles.margin = (0, 0, 1, 0)
yield local_approval_btn
otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn")
otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn")
otp_gen_btn.styles.width = "100%"
otp_gen_btn.styles.margin = (0, 0, 1, 0)
yield otp_gen_btn
toggle_enforcement_btn = Button(
"🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn"
"🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn"
)
toggle_enforcement_btn.styles.width = "100%"
toggle_enforcement_btn.styles.margin = (0, 0, 1, 0)
yield toggle_enforcement_btn
other_policy_btn = Button(
"🔀 Move to Other Policy", id="other_policy_btn"
"🔀 Move to Other Policy", id="other_policy_btn"
)
other_policy_btn.styles.width = "100%"
other_policy_btn.styles.margin = (0, 0, 1, 0)
yield other_policy_btn
exec_history_btn = Button(
"📊 View Execution History", id="exec_history_btn"
"📊 View Execution History", id="exec_history_btn"
)
exec_history_btn.styles.width = "100%"
exec_history_btn.styles.margin = (0, 0, 1, 0)
@@ -391,17 +448,17 @@ class AgentMoveOperations(Widget):
pyperclip.copy(results_text.text)
self.app.notify(
"📋✅ Results copied to clipboard!",
"📋✅ Results copied to clipboard!",
severity="information",
timeout=2,
)
except ImportError:
self.app.notify(
"❌ pyperclip not installed. Run: pip install pyperclip",
"❌ pyperclip not installed. Run: pip install pyperclip",
severity="warning",
)
except Exception as e:
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
event.stop()
elif btn_id == "export_csv_btn":
self._start_export_csv_operation()
@@ -452,7 +509,7 @@ class AgentMoveOperations(Widget):
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("✔️ Moving agents to local approval...")
status_label.update("✔❗ Moving agents to local approval...")
# Get API from app
api = self.app.api
@@ -463,8 +520,6 @@ class AgentMoveOperations(Widget):
try:
import time
from services.agenthandler import moveAgentToRelatedPolicy
# Generate batch ID
batch = int(time.time())
duration = 360 # Default 6 hours, could make this configurable
@@ -487,7 +542,7 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error during local approval operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
return
@@ -537,7 +592,7 @@ class AgentMoveOperations(Widget):
successful.append(file_path)
status_label.update(f"✅ Exported to {file_path}")
except Exception:
status_label.update("❌ Failed")
status_label.update("❌ Failed")
self.operation_in_progress = False
@@ -581,7 +636,7 @@ class AgentMoveOperations(Widget):
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("🔄 Toggling enforcement mode...")
status_label.update("🔄 Toggling enforcement mode...")
# Get API from app
api = self.app.api
@@ -590,9 +645,6 @@ class AgentMoveOperations(Widget):
unsuccessful = []
try:
from services.agenthandler import moveAgentToRelatedPolicy
from utils.configmanager import get_system_json
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
for agent in self.agents:
@@ -617,7 +669,7 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error during toggle enforcement operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
return
@@ -687,7 +739,7 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error loading policies: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
self.selected_operation = ""
self.app.notify(f"Failed to load policies: {str(e)}", severity="error")
@@ -722,7 +774,7 @@ class AgentMoveOperations(Widget):
)
except Exception as e:
logger.error(f"Failed to open execution history viewer: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f"❌ Error: {str(e)}")
self.app.notify(
f"Failed to open execution history: {str(e)}", severity="error"
)
TUI/Widgets/prepPolicy.py
-870
View File
@@ -1,870 +0,0 @@
# 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
import re
from typing import List
import dotenv
import pandas as pd
from models.execution import ExecutionHistoryRecord
from models.policy import Allowlist, Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import get_system_list, get_system_value, load_env
from utils.selector import Selector
from utils.utils import (
areYouSure,
clear_screen,
colorText,
formatHTML,
get_sanitized_input,
locked,
open_directory,
print_x_wide,
regulator,
)
logger = logging.getLogger(__name__)
dotenv.load_dotenv()
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
policies = [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(policies, 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, policy=all, allow_multiple=True
) -> List[Allowlist]:
if policy == "all":
allowlists = [
Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()
]
else:
allowlists = [
Allowlist(**row.to_dict())
for _, row in api.policy_list_allowlists(policy[0].groupid).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-365): ",
value_type=int,
valid_range=(1, 365),
)
logger.debug(f"{history_days} day selected for history")
if history_days is None:
logging.warning("No history range selected. Aborting.")
return
policy_executions = ExecutionHistoryRecord.from_policies(
api, selected_policies, type_=type, history_days=history_days
)
logger.debug(f"Executions contains {policy_executions}")
enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(
api, policy_executions
)
categorized_executions = (
ExecutionHistoryRecord.categorize_executions_by_hash_decision(
enriched_executions
)
)
approved, unapproved, needs_review, unknown = (
ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions)
)
categories = {
"needs_review": needs_review,
"approved": approved,
"unapproved": unapproved,
"leftover": unknown,
}
for label, records in categories.items():
if not records:
continue # Skip empty or falsy categories
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html"
# Convert ExecutionHistoryRecord objects to dictionaries
df = pd.DataFrame([r.__dict__ for r in records])
# Optional: flatten hash_obj if needed
if not df.empty and "hash_obj" in df.columns:
hash_df = df["hash_obj"].apply(lambda h: h.to_dict() if h else {})
df = pd.concat([df.drop(columns=["hash_obj"]), hash_df], axis=1)
# Save to CSV
df.to_csv(csv_path, index=False)
logger.info(f"Saved {label} executions to {csv_path}")
# Generate HTML
formatHTML(df, html_path)
logger.info(f"Generated HTML report at {html_path}")
def buildPathsandPublishers(selected_policies: List[Policy], split):
working_dir = load_env("WORKING_DIR")
df1 = pd.DataFrame()
df2 = pd.DataFrame()
all_approved_hashes = pd.DataFrame()
path1 = (
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
)
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
path_exclusion_constant = get_system_value("PATH_EXCLUSION_CONST", cast_type=int)
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" in all_approved_hashes.columns:
all_approved_hashes = all_approved_hashes.sort_values(by="filename")
else:
logger.warning(
"Warning: 'filename' column not found in concatenated DataFrame."
)
if not all_approved_hashes.empty and path_exclusion_constant:
primary_path_exclusions = calculatePath(
all_approved_hashes,
path_exclusion_constant,
split,
)
remaining_hashes = all_approved_hashes[
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
]
secondary_path_exclusions = calculatePath(
remaining_hashes, (path_exclusion_constant - 1), split
)
remaining_hashes = remaining_hashes[
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
]
dataframes = {
"all_approved_hashes": all_approved_hashes,
"primary_Paths": primary_path_exclusions,
"secondary_Paths": secondary_path_exclusions,
"hashes_not_approvable_by_path": remaining_hashes,
}
logger.debug("Preparing to sort dataframes")
for name, df in dataframes.items():
logger.debug(f" DataFrame headers: {list(df.columns)}")
if "hashes" in name:
df.sort_values(by="filename", inplace=True)
else:
df.sort_values(by="longestcfp", inplace=True)
df.to_csv(
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv",
index=False,
)
formatHTML(
df,
f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html",
)
if not all_approved_hashes.empty:
# Drop all not signed, only keep unique values
publist = all_approved_hashes[
all_approved_hashes["publisher"] != "Not Signed"
].drop_duplicates(subset=["publisher"])
# Remove Bad publisher if somehow they made it this far
pattern = regulator(get_system_list("BAD_PUBLISHERS"))
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
publist = publist[["publisher"]]
publist.sort_values(by="publisher", inplace=True)
publist.to_csv(
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv",
index=False,
)
else:
logger.debug("Approved Hashes list appears empty")
def buildPreflights(selected_policies: List[Policy]):
working_dir = load_env("WORKING_DIR")
df1 = pd.DataFrame()
df2 = pd.DataFrame()
approved_hashes = pd.DataFrame()
approved_publishers = pd.DataFrame()
hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv"
path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv"
publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv"
# Read in and combine the two path generations
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)
approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first")
# We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions.
if os.path.exists(hash):
hashes = pd.read_csv(hash)
approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])]
approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first")
else:
logger.warning(f"File not found: {hash}")
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():
logger.debug(f" DataFrame headers: {list(df.columns)}")
if name == "approved_paths":
df.sort_values(by="longestcfp", inplace=True)
elif name == "approved_hashes":
df.sort_values(by="filename", inplace=True)
elif name == "approved_publishers":
df.sort_values(by="publisher", inplace=True)
df.to_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv",
index=False,
)
formatHTML(
df,
f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html",
)
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
def clean_split(path):
if not isinstance(path, (str, bytes, os.PathLike)):
return []
parts = str(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)
if min_files_for_path is not None:
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, path_exclusion_constant, split):
if split:
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
else:
dfs_by_policy = [approved_hashes]
badpathparts = get_system_list("BAD_PATH_PARTS")
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
processed_dfs = []
for df in dfs_by_policy:
haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename")
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
def testChange(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR")
logger.info("These path exclusions would be added to:")
logger.info(destination_policy)
pathexclusions = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
)
hashes = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
)
unique_combinations = pathexclusions[
["longestcfp", "file_extension"]
].drop_duplicates()
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
processed_paths = [
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
for path, ext in unique_combinations.itertuples(index=False, name=None)
]
for path in processed_paths:
logger.info(path)
print(colorText("These publishers would added", "yellow"))
processed_publishers = []
if os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
):
publishers = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
)
if publishers.empty:
print(colorText("The publishers list is empty.", "red"))
else:
processed_publishers = (
publishers[publishers["publisher"] != "Not Signed"]["publisher"]
.drop_duplicates()
.tolist()
)
for publisher in processed_publishers:
print(publisher)
print(colorText("These hashes would be added to:", "yellow"))
print(destination_allowlist)
processed_hashes = hashes["sha256"].unique().tolist()
print_x_wide(processed_hashes, 3)
return processed_paths, processed_hashes, processed_publishers
def menu_policy_enforce(
api: AirlockAPIWrapper,
): # TODO Need to clean up 6 and 7 into functions
selected_policies = []
destination_policy = []
destination_allowlist = []
processed_paths = []
processed_hashes = []
processed_publishers = []
working_dir = load_env("WORKING_DIR")
while True:
printEnforceChecklist(
selected_policies, destination_policy, destination_allowlist
)
choice = get_sanitized_input("\nEnter your choice: ")
if choice == "1":
clear_screen()
selected_policies = selectPolicies(api, True)
elif choice == "2":
clear_screen()
print(
colorText(
"Please choose destination_name Policy for Path Exclusions", "white"
)
)
destination_policy = selectPolicies(api, False)
print(colorText("Please choose Allowlist for Hashes", "white"))
destination_allowlist = selectAllowlists(api, destination_policy, False)
elif choice == "3":
clear_screen()
sortHashes(
api,
selected_policies,
type=[1, 2, 6, 7],
)
elif choice == "4":
clear_screen()
if os.path.exists(
f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"
):
buildPathsandPublishers(selected_policies, False)
else:
print(
"File not found. Please make sure it's saved correctly and try again."
)
elif choice == "5":
clear_screen()
if os.path.exists(
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
) and os.path.exists(
f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
):
buildPreflights(selected_policies)
else:
print(
"File not found. Please make sure it's saved correctly and try again."
)
elif choice == "6":
clear_screen()
if (
os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
)
and os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
)
and destination_policy
and destination_allowlist
):
processed_paths, processed_hashes, processed_publishers = testChange(
selected_policies, destination_policy, destination_allowlist
)
else:
# Log which condition(s) failed
missing_items = []
if not os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
):
missing_items.append("approved_paths.csv not found")
if not os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
):
missing_items.append("approved_hashes.csv not found")
if not destination_policy:
missing_items.append("destination_policy is empty or None")
if not destination_allowlist:
missing_items.append("destination_allowlist is empty or None")
logger.error("Preflight check failed due to the following:")
for item in missing_items:
logger.error(f" - {item}")
elif choice == "7":
clear_screen()
areYouSure()
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
if (
processed_paths
and processed_hashes
and processed_publishers
and destination_policy
and destination_allowlist
and confirmation.strip() == "I AGREE"
):
print(colorText("Proceeding with the code...", "yellow"))
api.hash_add_to_allowlist(
destination_allowlist[0].applicationid, processed_hashes
)
api.policy_add_path_exclusions(
destination_policy[0].groupid, processed_paths
)
if processed_publishers:
api.policy_add_publishers(
destination_policy[0].groupid, processed_publishers
)
locked()
else:
logger.error("Confirmation block failed. Reasons:")
if not processed_publishers or processed_hashes or processed_paths:
logger.error(" - Test not performed.")
if not destination_policy:
logger.error(" - `destination_policy` is missing or invalid.")
if not destination_allowlist:
logger.error(" - `destination_allowlist` is missing or invalid.")
if confirmation.strip() != "I AGREE":
logger.error(
" - User did not confirm with 'I AGREE'. Received: '%s'",
confirmation.strip(),
)
elif choice.upper() == "F":
open_directory(working_dir)
elif choice.upper() == "B":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
def section_header(title):
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(colorText(f" ------------- {title} -------------", "cyan"))
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR")
section_header("Prepare to Enforce Policy")
print(
colorText(
"\nSequentially follow these steps to prepare a policy for enforcement:",
"white",
)
)
# Step 1: Originating Policies
print(
colorText(
"\n1. Choose which policy or policies to gather execution info from", "cyan"
)
)
if not selected_policies:
print(colorText(" [❌] No policies have been chosen", "red"))
else:
print(colorText("The following policies have been chosen:", "green"))
for policy in selected_policies:
print(colorText(f" [✅] {policy.name}", "green"))
# Step 2: Destination Policy and Allowlist
print(
colorText("2. Choose the destination policy and associated allowlist", "cyan")
)
if destination_policy:
print(
colorText(
f" [✅] {destination_policy[0].name} has been selected as the destination policy",
"green",
)
)
else:
print(colorText(" [❌] No destination policy has been chosen", "red"))
if destination_allowlist:
print(
colorText(
f" [✅] {destination_allowlist[0].name} has been selected as allowlist",
"green",
)
)
else:
print(colorText(" [❌] No allowlist has been chosen", "red"))
# Step 3: Data Preparation
print(
colorText(
f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review",
"cyan",
)
)
if selected_policies:
policy_id = selected_policies[0].name
review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv"
print(
colorText(
(
" [✅] Data has been fetched"
if os.path.exists(review_path)
else " [❌] Data has not been fetched"
),
"green" if os.path.exists(review_path) else "red",
)
)
else:
print(
colorText(
" [❌] No policies selected, cannot check data fetch status", "red"
)
)
# Step 4: Manual Review
print(colorText("4. Manually review the files:", "cyan"))
print(
colorText(
" Remove the rows containing hashes you do not approve of", "cyan"
)
)
print(
colorText(
f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.",
"cyan",
)
)
print(
colorText(
" This will start the process to generate possible filepath approvals",
"cyan",
)
)
if selected_policies:
policy_id = selected_policies[0].name
approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv"
second_review_path = (
f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv"
)
print(
colorText(
(
" [✅] Reviewed hashes have been loaded"
if os.path.exists(approved_path)
else " [❌] Reviewed hashes have not been loaded"
),
"green" if os.path.exists(approved_path) else "red",
)
)
print(
colorText(
(
" [✅] Path review list created"
if os.path.exists(second_review_path)
else " [❌] Path review list has not been created"
),
"green" if os.path.exists(second_review_path) else "red",
)
)
else:
print(
colorText(
" [❌] No policies selected, cannot check reviewed hashes or path list",
"red",
)
)
# Step 5: Path Review
print(
colorText(
f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\",
"cyan",
)
)
print(
colorText(
" Remove the rows containing path exclusions or publishers you do not approve of.",
"cyan",
)
)
print(
colorText(
f" When complete, save the files to {working_dir}\\data\\Approved",
"cyan",
)
)
print(
colorText(" Choose this option when done to build your preflights", "cyan")
)
if selected_policies:
policy_id = selected_policies[0].name
reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv"
preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv"
preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv"
print(
colorText(
(
" [✅] Reviewed path list detected"
if os.path.exists(reviewed_path)
else " [❌] Path review list has not been detected"
),
"green" if os.path.exists(reviewed_path) else "red",
)
)
preflight_ready = os.path.exists(preflight_paths) and os.path.exists(
preflight_hashes
)
print(
colorText(
(
" [✅] Preflight Path Exclusion List has been generated"
if preflight_ready
else " [❌] Preflight Path Exclusion List has not been generated"
),
"green" if preflight_ready else "red",
)
)
else:
print(
colorText(
" [❌] No policies selected, cannot check preflight status", "red"
)
)
# Final Steps
print(
colorText(
"6. Test ------------------------------------------------------", "cyan"
)
)
print(
colorText(
" Prints to console the changes that would be made, must be done to proceed. ",
"cyan",
)
)
print(
colorText(
"7. Liftoff ------------------------------------------------------", "cyan"
)
)
print(
colorText(
" Apply path exclusions and approved publishers to selected policy",
"cyan",
)
)
print(colorText(" Apply approved hashes to allowlist", "cyan"))
# Utility Options
print(colorText("F. Open Working Directory", "cyan"))
print(colorText("B. Back", "cyan"))
TUI/Widgets/settingswidget.py
+567
View File
@@ -0,0 +1,567 @@
# 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/>.
"""
Settings widget combining theme selection and update checking.
"""
import logging
import webbrowser
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.message import Message
from textual.widget import Widget
from textual.widgets import Button, Rule, Static
from utils.versionchecker import (
RELEASES_PAGE_URL,
UpdateCheckResult,
check_for_updates,
get_current_version,
get_version_checker,
)
logger = logging.getLogger(__name__)
class SettingsWidget(Widget):
"""Widget for application settings including themes and updates."""
DEFAULT_CSS = """
SettingsWidget {
height: 1fr;
}
/* Update section buttons - add margin between them */
#update_buttons Button {
margin-right: 1;
}
/* Theme buttons - consistent width within columns, slightly smaller */
.theme_btn {
width: 100%;
margin-bottom: 1;
}
/* Column headers */
.theme_column_header {
text-align: center;
text-style: bold;
margin-bottom: 1;
}
/* Section titles */
.settings_section_title {
text-style: bold;
margin-bottom: 1;
}
/* Theme columns - reduce overall width */
#theme_columns {
width: 80%;
}
/* Theme columns spacing */
#dark_themes_col1, #dark_themes_col2 {
margin-right: 1;
}
#light_themes_col {
margin-left: 1;
}
"""
class ThemeSelected(Message):
"""Message posted when a theme is selected."""
def __init__(self, theme_name: str):
super().__init__()
self.theme_name = theme_name
# Dark themes - Column 1
DARK_THEMES_COL1 = [
("Textual Dark", "textual-dark"),
("Nord", "nord"),
("Gruvbox", "gruvbox"),
("Dracula", "dracula"),
]
# Dark themes - Column 2
DARK_THEMES_COL2 = [
("Catppuccin Mocha", "catppuccin-mocha"),
("Tokyo Night", "tokyo-night"),
("Monokai", "monokai"),
]
# Light themes (third column)
LIGHT_THEMES = [
("Textual Light", "textual-light"),
("Flexoki", "flexoki"),
("Catppuccin Latte", "catppuccin-latte"),
("Solarized Light", "solarized-light"),
]
# Combined for backward compatibility
DARK_THEMES = DARK_THEMES_COL1 + DARK_THEMES_COL2
AVAILABLE_THEMES = DARK_THEMES + LIGHT_THEMES
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._update_result: UpdateCheckResult | None = None
self._checking = False
def compose(self):
# Wrap everything in a scrollable container with auto height children
with VerticalScroll(id="settings_scroll"):
# Version & Updates Section
with Vertical(id="updates_section") as updates:
updates.styles.height = "auto"
yield Static(
"📦 Version & Updates",
id="updates_title",
classes="settings_section_title",
)
version_text = f"Current Version: v{get_current_version()}"
yield Static(version_text, id="current_version")
with Horizontal(id="update_buttons") as btn_row:
btn_row.styles.height = "auto"
yield Button("🔍 Check for Updates", id="check_updates_btn")
yield Button("📥 View Releases", id="view_releases_btn")
yield Static("", id="update_status")
yield Rule()
# Theme Section - Three columns: Dark 1, Dark 2, Light
with Vertical(id="themes_section") as themes:
themes.styles.height = "auto"
yield Static(
"🎨 Theme Options",
id="theme_title",
classes="settings_section_title",
)
with Horizontal(id="theme_columns") as cols:
cols.styles.height = "auto"
# Dark themes section (2 columns under one header)
with Vertical(id="dark_themes_section") as dark_section:
dark_section.styles.width = "2fr"
dark_section.styles.height = "auto"
yield Static(
"🌙 Dark Themes",
classes="theme_column_header",
id="dark_header",
)
with Horizontal(id="dark_columns") as dark_cols:
dark_cols.styles.height = "auto"
# Dark themes column 1
with Vertical(id="dark_themes_col1") as dark_col1:
dark_col1.styles.width = "1fr"
dark_col1.styles.height = "auto"
for label, btn_id in self.DARK_THEMES_COL1:
yield Button(
label,
id=f"set_theme_{btn_id}",
classes="theme_btn",
)
# Dark themes column 2
with Vertical(id="dark_themes_col2") as dark_col2:
dark_col2.styles.width = "1fr"
dark_col2.styles.height = "auto"
for label, btn_id in self.DARK_THEMES_COL2:
yield Button(
label,
id=f"set_theme_{btn_id}",
classes="theme_btn",
)
# Light themes column
with Vertical(id="light_themes_col") as light_col:
light_col.styles.width = "1fr"
light_col.styles.height = "auto"
yield Static("☀️ Light Themes", classes="theme_column_header")
for label, btn_id in self.LIGHT_THEMES:
yield Button(
label, id=f"set_theme_{btn_id}", classes="theme_btn"
)
def on_mount(self) -> None:
"""Check for cached update result on mount."""
checker = get_version_checker()
cached_result = checker.get_last_result()
if cached_result and cached_result.update_available:
self._update_result = cached_result
self._show_update_available(cached_result)
def on_button_pressed(self, event: Button.Pressed) -> None:
button_id = event.button.id
if button_id == "check_updates_btn":
self._check_for_updates()
event.stop()
elif button_id == "view_releases_btn":
self._open_releases_page()
event.stop()
elif button_id == "download_update_btn":
self._download_update()
event.stop()
elif button_id == "dismiss_update_btn":
self._dismiss_update()
event.stop()
elif button_id and button_id.startswith("set_theme_"):
theme_name = button_id.replace("set_theme_", "")
self.post_message(self.ThemeSelected(theme_name))
event.stop()
def _check_for_updates(self) -> None:
"""Check for updates and update UI."""
if self._checking:
return
self._checking = True
status = self.query_one("#update_status", Static)
check_btn = self.query_one("#check_updates_btn", Button)
# Show checking status
check_btn.disabled = True
check_btn.label = "⏳ Checking..."
status.update("🔄 Checking for updates...")
# Run check in worker to avoid blocking UI
self.run_worker(self._do_update_check, exclusive=True)
async def _do_update_check(self) -> None:
"""Worker to perform update check."""
try:
result = check_for_updates()
self._update_result = result
# Since we're in an async worker (not a thread), we can call directly
self._update_check_complete(result)
except Exception as e:
logger.error(f"Update check failed: {e}")
self._update_check_failed(str(e))
finally:
self._checking = False
def _update_check_complete(self, result: UpdateCheckResult) -> None:
"""Handle completed update check."""
check_btn = self.query_one("#check_updates_btn", Button)
check_btn.disabled = False
check_btn.label = "🔍 Check for Updates"
if result.error:
self._update_check_failed(result.error)
return
if result.update_available:
self._show_update_available(result)
self.app.notify(
f"🆕 Update available: {result.latest_version}",
title="Update Available",
severity="information",
timeout=8,
)
else:
status = self.query_one("#update_status", Static)
status.update(f"✅ Loxide is up to date (v{result.current_version})")
self.app.notify(
"✅ Loxide is up to date!",
severity="information",
timeout=5,
)
def _update_check_failed(self, error: str) -> None:
"""Handle failed update check."""
check_btn = self.query_one("#check_updates_btn", Button)
check_btn.disabled = False
check_btn.label = "🔍 Check for Updates"
status = self.query_one("#update_status", Static)
status.update(f"⚠️ Could not check for updates: {error}")
def _show_update_available(self, result: UpdateCheckResult) -> None:
"""Show update available UI with release notes."""
status = self.query_one("#update_status", Static)
msg = f"🆕 New version available: {result.latest_version}\n"
msg += f" Current: v{result.current_version}"
if result.release_info and result.release_info.body:
# Show release notes (truncate if very long)
notes = result.release_info.body.strip()
# Limit to ~500 chars to avoid overwhelming the UI
if len(notes) > 500:
notes = notes[:500] + "\n..."
msg += f"\n\n📋 Release Notes:\n{notes}"
status.update(msg)
# Add download/dismiss buttons if not already there
try:
self.query_one("#download_update_btn")
except Exception:
# Buttons don't exist, add them
button_container = self.query_one("#update_buttons", Horizontal)
download_btn = Button(
"📥 Download Update", id="download_update_btn", variant="success"
)
dismiss_btn = Button(
"✖ Dismiss", id="dismiss_update_btn", variant="default"
)
button_container.mount(download_btn)
button_container.mount(dismiss_btn)
def _open_releases_page(self) -> None:
"""Open the releases page in browser."""
try:
webbrowser.open(RELEASES_PAGE_URL)
self.app.notify("📂 Opened releases page in browser", timeout=3)
except Exception as e:
logger.error(f"Could not open browser: {e}")
self.app.notify(f"⚠️ Could not open browser: {e}", severity="warning")
def _download_update(self) -> None:
"""Download the update exe file."""
import os
from pathlib import Path
if not self._update_result or not self._update_result.release_info:
self.app.notify("⚠️ No update information available", severity="warning")
return
download_url = self._update_result.release_info.download_url
if not download_url:
# Fall back to opening the release page
url = self._update_result.release_info.html_url
try:
webbrowser.open(url)
self.app.notify(
"📥 Opened download page in browser (no direct download available)",
timeout=5,
)
except Exception as e:
logger.error(f"Could not open browser: {e}")
self.app.notify(f"⚠️ Could not open browser: {e}", severity="warning")
return
# Determine destination path
if os.name == "nt": # Windows
downloads_dir = Path.home() / "Downloads"
else:
downloads_dir = Path.home() / "Downloads"
if not downloads_dir.exists():
downloads_dir = Path.home()
# Extract filename from URL
filename = download_url.split("/")[-1]
if not filename.endswith(".exe"):
filename = f"Loxide_{self._update_result.latest_version}.exe"
dest_path = downloads_dir / filename
# Disable the button while downloading
try:
btn = self.query_one("#download_update_btn", Button)
btn.disabled = True
btn.label = "⏳ Downloading..."
except Exception:
pass
self.app.notify(
f"📥 Downloading to:\n{dest_path}", title="Download Starting", timeout=5
)
# Small delay so user sees the "downloading to" toast before download completes
self.set_timer(
0.5,
lambda: self.run_worker(
self._do_download(download_url, dest_path), exclusive=True
),
)
async def _do_download(self, download_url: str, dest_path) -> None:
"""Worker to download the update file."""
try:
# Download the file
import requests
response = requests.get(download_url, stream=True, timeout=60)
response.raise_for_status()
with open(dest_path, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
# Success
self.app.notify(
f"✅ Downloaded to:\n{dest_path}",
title="Download Complete",
severity="information",
timeout=10,
)
logger.info(f"Update downloaded to {dest_path}")
# Re-enable button
try:
btn = self.query_one("#download_update_btn", Button)
btn.disabled = False
btn.label = "📥 Download Again"
except Exception:
pass
except Exception as e:
logger.error(f"Download failed: {e}")
self.app.notify(f"❌ Download failed: {e}", severity="error", timeout=10)
# Re-enable button
try:
btn = self.query_one("#download_update_btn", Button)
btn.disabled = False
btn.label = "📥 Download Update"
except Exception:
pass
def _dismiss_update(self) -> None:
"""Dismiss the current update notification."""
if self._update_result and self._update_result.latest_version:
checker = get_version_checker()
checker.dismiss_update(self._update_result.latest_version)
# Remove the extra buttons
try:
self.query_one("#download_update_btn").remove()
self.query_one("#dismiss_update_btn").remove()
except Exception:
pass
status = self.query_one("#update_status", Static)
status.update(f"✓ Dismissed update {self._update_result.latest_version}")
self._update_result = None
# Keep ThemeSelector as a standalone for backward compatibility
class ThemeSelector(Widget):
"""Widget for selecting and applying Textual themes.
DEPRECATED: Use SettingsWidget instead for combined settings UI.
"""
DEFAULT_CSS = """
ThemeSelector {
height: 1fr;
}
/* Theme buttons - consistent width within columns */
.theme_btn {
width: 100%;
margin-bottom: 1;
}
/* Column headers */
.theme_column_header {
text-align: center;
text-style: bold;
margin-bottom: 1;
}
/* Theme columns - reduce overall width */
#theme_columns {
width: 80%;
}
/* Theme columns spacing */
#dark_themes_col1, #dark_themes_col2 {
margin-right: 1;
}
#light_themes_col {
margin-left: 1;
}
"""
class ThemeSelected(Message):
"""Message posted when a theme is selected."""
def __init__(self, theme_name: str):
super().__init__()
self.theme_name = theme_name
DARK_THEMES_COL1 = SettingsWidget.DARK_THEMES_COL1
DARK_THEMES_COL2 = SettingsWidget.DARK_THEMES_COL2
DARK_THEMES = SettingsWidget.DARK_THEMES
LIGHT_THEMES = SettingsWidget.LIGHT_THEMES
AVAILABLE_THEMES = SettingsWidget.AVAILABLE_THEMES
def compose(self):
with VerticalScroll(id="theme_scroll"):
yield Static("Theme Options", id="theme_title")
with Horizontal(id="theme_columns") as cols:
cols.styles.height = "auto"
# Dark themes section (2 columns under one header)
with Vertical(id="dark_themes_section") as dark_section:
dark_section.styles.width = "2fr"
dark_section.styles.height = "auto"
yield Static(
"🌙 Dark Themes",
classes="theme_column_header",
id="dark_header",
)
with Horizontal(id="dark_columns") as dark_cols:
dark_cols.styles.height = "auto"
# Dark themes column 1
with Vertical(id="dark_themes_col1") as dark_col1:
dark_col1.styles.width = "1fr"
dark_col1.styles.height = "auto"
for label, btn_id in self.DARK_THEMES_COL1:
yield Button(
label, id=f"set_theme_{btn_id}", classes="theme_btn"
)
# Dark themes column 2
with Vertical(id="dark_themes_col2") as dark_col2:
dark_col2.styles.width = "1fr"
dark_col2.styles.height = "auto"
for label, btn_id in self.DARK_THEMES_COL2:
yield Button(
label, id=f"set_theme_{btn_id}", classes="theme_btn"
)
# Light themes column
with Vertical(id="light_themes_col") as light_col:
light_col.styles.width = "1fr"
light_col.styles.height = "auto"
yield Static("☀️ Light Themes", classes="theme_column_header")
for label, btn_id in self.LIGHT_THEMES:
yield Button(
label, id=f"set_theme_{btn_id}", classes="theme_btn"
)
def on_button_pressed(self, event: Button.Pressed) -> None:
button_id = event.button.id
if button_id and button_id.startswith("set_theme_"):
theme_name = button_id.replace("set_theme_", "")
self.post_message(self.ThemeSelected(theme_name))
TUI/Themes/themeselector.py → TUI/Widgets/themeselector.py
View File
airlock_libs-6.2.0-cp313-cp313-win_amd64.whl
BIN
View File
Binary file not shown.
docs/README.md
+130 -16
View File
@@ -1,39 +1,47 @@
# 🛡️ Loxide
Python/Rust/Oxide toolkit for secure, auditable, and automated airlock agent and policy management. Designed for enterprise environments, it supports advanced policy workflows, device tracking, and terminal-based interaction.
Python/Rust/Oxide toolkit for secure, auditable, and automated Airlock agent and policy management. Designed for enterprise environments, it supports advanced policy workflows, device tracking, and terminal-based interaction.
![Python](https://img.shields.io/badge/python-3.10+-blue.svg)
![Rust](https://img.shields.io/badge/rust-backend-orange.svg)
![License](https://img.shields.io/badge/license-AGPL--3.0-green.svg)
![Textual](https://img.shields.io/badge/textual-6.5.0-purple.svg)
---
## 🚀 Features
- 🔍 **Fuzzy Device Search**
Quickly locate devices using partial or approximate matches.
Quickly locate devices using partial or approximate matches with wildcard support (`*`, `?`).
- 📦 **Batch Move Devices**
Move multiple devices between groups or policies easily.
Move multiple devices between groups or policies easily with multi-select and file import.
- 🔄 **Toggle Enforcement/Audit Policies**
Seamlessly switch devices between enforcement and audit modes.
- 🕵️‍♂️ **Device History Search**
Track agent executions.
- 🕵️ **Device History Search**
Track agent executions with configurable date ranges (1-365 days) and CSV export.
- 🧰 **Prepare Policies for Enforcement**
Validate and stage policies before pushing them to enforcement.
Multi-step wizard to validate and stage policies before pushing to enforcement.
- 💤 **Find Quiet Hosts**
Identify devices ready for enforcement.
Identify inactive devices ready for enforcement using Rust-powered analysis.
- 🎛️ **TUI**
Navigate with arrow keys and F-key shortcuts using a custom ANSI-colored terminal UI.
- 🔐 **OTP Management**
Generate, revoke, and monitor one-time passwords for agents.
- 📋 **Server Log Viewer**
Real-time server activity monitoring with DataTable display.
- 🎛️ **Modern TUI**
Full keyboard navigation with Textual framework, toast notifications, and color-coded results.
---
## 🧭 Roadmap
- ⚙️ **Rust-based Async API Calls**
Improve performance and concurrency with a Rust-powered backend.
- ✅ **Carbon Black-style Local Approval**
Enable local user approvals for policy exceptions and enforcement actions.
@@ -43,9 +51,105 @@ Python/Rust/Oxide toolkit for secure, auditable, and automated airlock agent and
---
## 🧑‍💻 Requirements
## 📦 Installation
[airlock_libs](https://git.racooncity.org/brotoskyj/-/packages/pypi/airlock-libs/)
### Prerequisites
- Python 3.10+
- Access to Airlock API
- Network access to private PyPI server
### Install Dependencies
```bash
pip install -r requirements.txt
```
### Requirements
```
# Core TUI
textual==6.5.0
# API and data handling
requests==2.32.5
pandas==2.3.3
numpy==2.3.4
# Database
pymongo==4.15.3
# Security
cryptography==46.0.3
keyring==25.6.0
# Utilities
python-dotenv==1.2.1
tqdm==4.67.1
urllib3==2.5.0
pyperclip==1.11.0
# Private package
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
airlock_libs==5.2.1
```
See also: [airlock_libs](https://git.racooncity.org/brotoskyj/-/packages/pypi/airlock-libs/)
---
## 📖 Usage
```bash
python Loxide.py
```
### Keyboard Shortcuts
| Key | Action |
|-----|--------|
| `Tab` | Navigate between elements |
| `Enter` | Select/Confirm |
| `Escape` | Go back / Cancel |
| `←` `→` | Navigate tabs |
| `r` | Refresh (context-dependent) |
| `e` | Export to CSV |
| `q` | Quit application |
---
## 🏗️ Architecture
```
Loxide/
├── Loxide.py # Main entry point
├── API.py # AirlockAPIWrapper
├── setup.py # Initialization & logging
├── configmanager.py # Config management
│
├── models/ # Data models
├── services/ # Business logic
├── utils/ # Helper functions
│
└── TUI/
├── Screens/ # Workflow screens
│ ├── policyprepworkflowscreen.py
│ ├── quietagentworkflowscreen.py
│ ├── agentmoveoperations.py
│ ├── otprevokescreen.py
│ └── executionhistoryscreen.py
│
└── Widgets/ # Reusable components
├── policyselector.py
├── multiagentselector.py
└── serverlogwidget.py
```
---
## 👥 Authors
- **Brandon Wickline** - *Lead Python Developer*
- **James Brotosky** - *Lead Rust Developer*
---
@@ -55,5 +159,15 @@ Python/Rust/Oxide toolkit for secure, auditable, and automated airlock agent and
You may copy, distribute, and modify the software under the terms of the AGPL-3.0 license.
```
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.
```
See the [LICENSE](LICENSE.md) file for full details, or visit
[https://www.gnu.org/license/agpl-3.0.html](https://www.gnu.org/license/agpl-3.0.html)
[https://www.gnu.org/licenses/agpl-3.0.html](https://www.gnu.org/licenses/agpl-3.0.html)
flows/localApproval.py
-252
View File
@@ -1,252 +0,0 @@
# 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 time
from typing import List, Optional
from models.agent import Agent
from services.agenthandler import moveAgentToRelatedPolicy, selectAgents
from services.API import AirlockAPIWrapper
from utils.configmanager import get_system_json
from utils.utils import colorText, get_sanitized_input
logger = logging.getLogger(__name__)
class LocalApprovalRequestor:
"""Handles creation of local approval requests in Loxide."""
def __init__(self, api: AirlockAPIWrapper, username: str = None):
"""
Initialize the local approval requestor.
Args:
api: AirlockAPIWrapper instance
username: Username creating the approvals (for tracking)
"""
self.api = api
self.policy_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
self.username = (
username or os.getenv("USERNAME") or os.getenv("USER") or "unknown"
)
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)
Returns:
True if successful, False otherwise
"""
if batch_id is None:
batch_id = int(time.time())
purpose = (
f" Local Approval - {duration_minutes} mins - "
f"batch:{batch_id} Client:{agent_id} User:{self.username}"
)
try:
self.api.otp_generate(agent_id, duration_minutes, purpose)
logger.info(
f"Generated local approval for {agent_id}, batch {batch_id}, by {self.username}"
)
return True
except Exception as 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.
Args:
agent: Agent object to move
Returns:
True if successful, False otherwise
"""
try:
moveAgentToRelatedPolicy(self.api, agent, "audit")
logger.info(f"Moved {agent.hostname} to audit policy")
return True
except Exception as e:
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.
Args:
agents: List of Agent objects
duration_minutes: Duration of approval in minutes
db_path: Optional path to database for history tracking
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"))
for agent in agents:
try:
# Create local approval
approval_success = self.create_local_approval(
agent.agentid, duration_minutes, batch_id
)
if not approval_success:
raise Exception("Failed to create local approval")
# Move to audit policy
move_success = self.move_agent_to_audit(agent)
if not move_success:
raise Exception("Failed to move to audit policy")
print(colorText(f" {agent.hostname}", "green"))
success_count += 1
except Exception as e:
print(colorText(f" {agent.hostname}: {e}", "red"))
logger.error(f"Error processing agent {agent.hostname}: {e}")
failure_count += 1
return batch_id, success_count, failure_count
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"),
]
# Display duration options
print(colorText("\n Select Local Approval Duration:", "white"))
print(colorText("=" * 50, "white"))
for i, (minutes, label) in enumerate(duration_options, start=1):
print(f" {i}. {label} ({minutes} minutes)")
print(colorText("=" * 50, "white"))
# Get user selection
try:
choice = int(get_sanitized_input("\nEnter the number of your choice: "))
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
except ValueError:
print(colorText("❌ Invalid input. Please enter a number.", "red"))
logger.warning("Invalid input for duration selection")
return
# Select agents
print(colorText("\nSelect Agents for Local Approval:", "white"))
agents = selectAgents(self.api)
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("\nSummary:", "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"))
flows/otp.py
-69
View File
@@ -1,69 +0,0 @@
# 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 pandas as pd
from services.agenthandler import selectAgents
from services.API import AirlockAPIWrapper
from utils.selector import Selector
from utils.utils import get_sanitized_input
logger = logging.getLogger(__name__)
def otp_revoke(api: AirlockAPIWrapper):
activeagents = api.otp_find_active()
awaitingagents = api.otp_find_awaiting()
activeagents["status"] = "active"
awaitingagents["status"] = "awaiting"
combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True)
combined_agents = combined_agents.sort_values(by="otpid", ascending=False)
# Combine all into one DataFrame
combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True)
combined_agents = combined_agents.sort_values(by="otpid", ascending=False)
# Optionally, select specific hosts
user_input = (
get_sanitized_input("\nWould you like to search for a specific device? (y/n): ")
.strip()
.lower()
)
if user_input == "y":
agentnames = []
agents = selectAgents(api)
for agent in agents:
agentnames.append(agent.hostname)
combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)]
# Present and select rows
selected_rows = Selector.select_dataframe_with_mode(
combined_agents,
columns=["otpid", "hostname", "status", "purpose", "granted"],
header="OTP Sessions",
)
for row in selected_rows:
otpid = row["otpid"]
hostname = row["hostname"]
result = api.otp_revoke(otpid)
logger.info(f"{hostname} (otpid: {otpid}):\n{result}")
flows/prepPolicy.py
-870
View File
@@ -1,870 +0,0 @@
# 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
import re
from typing import List
import dotenv
import pandas as pd
from models.execution import ExecutionHistoryRecord
from models.policy import Allowlist, Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import get_system_list, get_system_value, load_env
from utils.selector import Selector
from utils.utils import (
areYouSure,
clear_screen,
colorText,
formatHTML,
get_sanitized_input,
locked,
open_directory,
print_x_wide,
regulator,
)
logger = logging.getLogger(__name__)
dotenv.load_dotenv()
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
policies = [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(policies, 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, policy=all, allow_multiple=True
) -> List[Allowlist]:
if policy == "all":
allowlists = [
Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()
]
else:
allowlists = [
Allowlist(**row.to_dict())
for _, row in api.policy_list_allowlists(policy[0].groupid).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
policy_executions = ExecutionHistoryRecord.from_policies(
api, selected_policies, type_=type, history_days=history_days
)
logger.debug(f"Executions contains {policy_executions}")
enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(
api, policy_executions
)
categorized_executions = (
ExecutionHistoryRecord.categorize_executions_by_hash_decision(
enriched_executions
)
)
approved, unapproved, needs_review, unknown = (
ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions)
)
categories = {
"needs_review": needs_review,
"approved": approved,
"unapproved": unapproved,
"leftover": unknown,
}
for label, records in categories.items():
if not records:
continue # Skip empty or falsy categories
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html"
# Convert ExecutionHistoryRecord objects to dictionaries
df = pd.DataFrame([r.__dict__ for r in records])
# Optional: flatten hash_obj if needed
if not df.empty and "hash_obj" in df.columns:
hash_df = df["hash_obj"].apply(lambda h: h.to_dict() if h else {})
df = pd.concat([df.drop(columns=["hash_obj"]), hash_df], axis=1)
# Save to CSV
df.to_csv(csv_path, index=False)
logger.info(f"Saved {label} executions to {csv_path}")
# Generate HTML
formatHTML(df, html_path)
logger.info(f"Generated HTML report at {html_path}")
def buildPathsandPublishers(selected_policies: List[Policy], split):
working_dir = load_env("WORKING_DIR")
df1 = pd.DataFrame()
df2 = pd.DataFrame()
all_approved_hashes = pd.DataFrame()
path1 = (
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
)
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
path_exclusion_constant = get_system_value("PATH_EXCLUSION_CONST", cast_type=int)
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" in all_approved_hashes.columns:
all_approved_hashes = all_approved_hashes.sort_values(by="filename")
else:
logger.warning(
"Warning: 'filename' column not found in concatenated DataFrame."
)
if not all_approved_hashes.empty and path_exclusion_constant:
primary_path_exclusions = calculatePath(
all_approved_hashes,
path_exclusion_constant,
split,
)
remaining_hashes = all_approved_hashes[
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
]
secondary_path_exclusions = calculatePath(
remaining_hashes, (path_exclusion_constant - 1), split
)
remaining_hashes = remaining_hashes[
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
]
dataframes = {
"all_approved_hashes": all_approved_hashes,
"primary_Paths": primary_path_exclusions,
"secondary_Paths": secondary_path_exclusions,
"hashes_not_approvable_by_path": remaining_hashes,
}
logger.debug("Preparing to sort dataframes")
for name, df in dataframes.items():
logger.debug(f" DataFrame headers: {list(df.columns)}")
if "hashes" in name:
df.sort_values(by="filename", inplace=True)
else:
df.sort_values(by="longestcfp", inplace=True)
df.to_csv(
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv",
index=False,
)
formatHTML(
df,
f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html",
)
if not all_approved_hashes.empty:
# Drop all not signed, only keep unique values
publist = all_approved_hashes[
all_approved_hashes["publisher"] != "Not Signed"
].drop_duplicates(subset=["publisher"])
# Remove Bad publisher if somehow they made it this far
pattern = regulator(get_system_list("BAD_PUBLISHERS"))
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
publist = publist[["publisher"]]
publist.sort_values(by="publisher", inplace=True)
publist.to_csv(
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv",
index=False,
)
else:
logger.debug("Approved Hashes list appears empty")
def buildPreflights(selected_policies: List[Policy]):
working_dir = load_env("WORKING_DIR")
df1 = pd.DataFrame()
df2 = pd.DataFrame()
approved_hashes = pd.DataFrame()
approved_publishers = pd.DataFrame()
hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv"
path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv"
publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv"
# Read in and combine the two path generations
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)
approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first")
# We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions.
if os.path.exists(hash):
hashes = pd.read_csv(hash)
approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])]
approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first")
else:
logger.warning(f"File not found: {hash}")
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():
logger.debug(f" DataFrame headers: {list(df.columns)}")
if name == "approved_paths":
df.sort_values(by="longestcfp", inplace=True)
elif name == "approved_hashes":
df.sort_values(by="filename", inplace=True)
elif name == "approved_publishers":
df.sort_values(by="publisher", inplace=True)
df.to_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv",
index=False,
)
formatHTML(
df,
f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html",
)
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
def clean_split(path):
if not isinstance(path, (str, bytes, os.PathLike)):
return []
parts = str(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)
if min_files_for_path is not None:
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, path_exclusion_constant, split):
if split:
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
else:
dfs_by_policy = [approved_hashes]
badpathparts = get_system_list("BAD_PATH_PARTS")
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
processed_dfs = []
for df in dfs_by_policy:
haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename")
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
def testChange(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR")
logger.info("These path exclusions would be added to:")
logger.info(destination_policy)
pathexclusions = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
)
hashes = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
)
unique_combinations = pathexclusions[
["longestcfp", "file_extension"]
].drop_duplicates()
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
processed_paths = [
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
for path, ext in unique_combinations.itertuples(index=False, name=None)
]
for path in processed_paths:
logger.info(path)
print(colorText("These publishers would added", "yellow"))
processed_publishers = []
if os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
):
publishers = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
)
if publishers.empty:
print(colorText("The publishers list is empty.", "red"))
else:
processed_publishers = (
publishers[publishers["publisher"] != "Not Signed"]["publisher"]
.drop_duplicates()
.tolist()
)
for publisher in processed_publishers:
print(publisher)
print(colorText("These hashes would be added to:", "yellow"))
print(destination_allowlist)
processed_hashes = hashes["sha256"].unique().tolist()
print_x_wide(processed_hashes, 3)
return processed_paths, processed_hashes, processed_publishers
def menu_policy_enforce(
api: AirlockAPIWrapper,
): # TODO Need to clean up 6 and 7 into functions
selected_policies = []
destination_policy = []
destination_allowlist = []
processed_paths = []
processed_hashes = []
processed_publishers = []
working_dir = load_env("WORKING_DIR")
while True:
printEnforceChecklist(
selected_policies, destination_policy, destination_allowlist
)
choice = get_sanitized_input("\nEnter your choice: ")
if choice == "1":
clear_screen()
selected_policies = selectPolicies(api, True)
elif choice == "2":
clear_screen()
print(
colorText(
"Please choose destination_name Policy for Path Exclusions", "white"
)
)
destination_policy = selectPolicies(api, False)
print(colorText("Please choose Allowlist for Hashes", "white"))
destination_allowlist = selectAllowlists(api, destination_policy, False)
elif choice == "3":
clear_screen()
sortHashes(
api,
selected_policies,
type=[1, 2, 6, 7],
)
elif choice == "4":
clear_screen()
if os.path.exists(
f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"
):
buildPathsandPublishers(selected_policies, False)
else:
print(
"File not found. Please make sure it's saved correctly and try again."
)
elif choice == "5":
clear_screen()
if os.path.exists(
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
) and os.path.exists(
f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
):
buildPreflights(selected_policies)
else:
print(
"File not found. Please make sure it's saved correctly and try again."
)
elif choice == "6":
clear_screen()
if (
os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
)
and os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
)
and destination_policy
and destination_allowlist
):
processed_paths, processed_hashes, processed_publishers = testChange(
selected_policies, destination_policy, destination_allowlist
)
else:
# Log which condition(s) failed
missing_items = []
if not os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
):
missing_items.append("approved_paths.csv not found")
if not os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
):
missing_items.append("approved_hashes.csv not found")
if not destination_policy:
missing_items.append("destination_policy is empty or None")
if not destination_allowlist:
missing_items.append("destination_allowlist is empty or None")
logger.error("Preflight check failed due to the following:")
for item in missing_items:
logger.error(f" - {item}")
elif choice == "7":
clear_screen()
areYouSure()
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
if (
processed_paths
and processed_hashes
and processed_publishers
and destination_policy
and destination_allowlist
and confirmation.strip() == "I AGREE"
):
print(colorText("Proceeding with the code...", "yellow"))
api.hash_add_to_allowlist(
destination_allowlist[0].applicationid, processed_hashes
)
api.policy_add_path_exclusions(
destination_policy[0].groupid, processed_paths
)
if processed_publishers:
api.policy_add_publishers(
destination_policy[0].groupid, processed_publishers
)
locked()
else:
logger.error("Confirmation block failed. Reasons:")
if not processed_publishers or processed_hashes or processed_paths:
logger.error(" - Test not performed.")
if not destination_policy:
logger.error(" - `destination_policy` is missing or invalid.")
if not destination_allowlist:
logger.error(" - `destination_allowlist` is missing or invalid.")
if confirmation.strip() != "I AGREE":
logger.error(
" - User did not confirm with 'I AGREE'. Received: '%s'",
confirmation.strip(),
)
elif choice.upper() == "F":
open_directory(working_dir)
elif choice.upper() == "B":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
def section_header(title):
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(colorText(f" ------------- {title} -------------", "cyan"))
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR")
section_header("Prepare to Enforce Policy")
print(
colorText(
"\nSequentially follow these steps to prepare a policy for enforcement:",
"white",
)
)
# Step 1: Originating Policies
print(
colorText(
"\n1. Choose which policy or policies to gather execution info from", "cyan"
)
)
if not selected_policies:
print(colorText(" [❌] No policies have been chosen", "red"))
else:
print(colorText("The following policies have been chosen:", "green"))
for policy in selected_policies:
print(colorText(f" [✅] {policy.name}", "green"))
# Step 2: Destination Policy and Allowlist
print(
colorText("2. Choose the destination policy and associated allowlist", "cyan")
)
if destination_policy:
print(
colorText(
f" [✅] {destination_policy[0].name} has been selected as the destination policy",
"green",
)
)
else:
print(colorText(" [❌] No destination policy has been chosen", "red"))
if destination_allowlist:
print(
colorText(
f" [✅] {destination_allowlist[0].name} has been selected as allowlist",
"green",
)
)
else:
print(colorText(" [❌] No allowlist has been chosen", "red"))
# Step 3: Data Preparation
print(
colorText(
f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review",
"cyan",
)
)
if selected_policies:
policy_id = selected_policies[0].name
review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv"
print(
colorText(
(
" [✅] Data has been fetched"
if os.path.exists(review_path)
else " [❌] Data has not been fetched"
),
"green" if os.path.exists(review_path) else "red",
)
)
else:
print(
colorText(
" [❌] No policies selected, cannot check data fetch status", "red"
)
)
# Step 4: Manual Review
print(colorText("4. Manually review the files:", "cyan"))
print(
colorText(
" Remove the rows containing hashes you do not approve of", "cyan"
)
)
print(
colorText(
f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.",
"cyan",
)
)
print(
colorText(
" This will start the process to generate possible filepath approvals",
"cyan",
)
)
if selected_policies:
policy_id = selected_policies[0].name
approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv"
second_review_path = (
f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv"
)
print(
colorText(
(
" [✅] Reviewed hashes have been loaded"
if os.path.exists(approved_path)
else " [❌] Reviewed hashes have not been loaded"
),
"green" if os.path.exists(approved_path) else "red",
)
)
print(
colorText(
(
" [✅] Path review list created"
if os.path.exists(second_review_path)
else " [❌] Path review list has not been created"
),
"green" if os.path.exists(second_review_path) else "red",
)
)
else:
print(
colorText(
" [❌] No policies selected, cannot check reviewed hashes or path list",
"red",
)
)
# Step 5: Path Review
print(
colorText(
f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\",
"cyan",
)
)
print(
colorText(
" Remove the rows containing path exclusions or publishers you do not approve of.",
"cyan",
)
)
print(
colorText(
f" When complete, save the files to {working_dir}\\data\\Approved",
"cyan",
)
)
print(
colorText(" Choose this option when done to build your preflights", "cyan")
)
if selected_policies:
policy_id = selected_policies[0].name
reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv"
preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv"
preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv"
print(
colorText(
(
" [✅] Reviewed path list detected"
if os.path.exists(reviewed_path)
else " [❌] Path review list has not been detected"
),
"green" if os.path.exists(reviewed_path) else "red",
)
)
preflight_ready = os.path.exists(preflight_paths) and os.path.exists(
preflight_hashes
)
print(
colorText(
(
" [✅] Preflight Path Exclusion List has been generated"
if preflight_ready
else " [❌] Preflight Path Exclusion List has not been generated"
),
"green" if preflight_ready else "red",
)
)
else:
print(
colorText(
" [❌] No policies selected, cannot check preflight status", "red"
)
)
# Final Steps
print(
colorText(
"6. Test ------------------------------------------------------", "cyan"
)
)
print(
colorText(
" Prints to console the changes that would be made, must be done to proceed. ",
"cyan",
)
)
print(
colorText(
"7. Liftoff ------------------------------------------------------", "cyan"
)
)
print(
colorText(
" Apply path exclusions and approved publishers to selected policy",
"cyan",
)
)
print(colorText(" Apply approved hashes to allowlist", "cyan"))
# Utility Options
print(colorText("F. Open Working Directory", "cyan"))
print(colorText("B. Back", "cyan"))
loxide_tight.ico
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 118 KiB

models/execution.py
+1 -7
View File
@@ -29,7 +29,7 @@ import pandas as pd
import airlock_libs
from services.API import AirlockAPIWrapper
from utils.configmanager import get_system_list, get_system_value
from utils.utils import colorText, regulator
from utils.utils import regulator
logger = logging.getLogger(__name__)
@@ -293,12 +293,6 @@ class ExecutionHistoryRecord:
logger.debug(
f"Staging of Execution history for policy: {policy.name} is complete"
)
print(
colorText(
f"Staging of Execution history for policy: {policy.name} is complete",
"green",
)
)
return executions
requirements.txt
+1 -1
View File
@@ -17,8 +17,8 @@ keyring==25.6.0
python-dotenv==1.2.1
# Utilities
tqdm==4.67.1
urllib3==2.5.0
plotext==5.3.2
pyperclip==1.11.0
# Custom/Private packages
services/agenthandler.py
-388
View File
@@ -1,388 +0,0 @@
# 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/>.
from dataclasses import asdict
from datetime import datetime, timedelta
import json
import logging
import os
import re
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
from utils.configmanager import get_system_json, load_env
from utils.selector import Selector
from utils.utils import colorText, get_sanitized_input
logger = logging.getLogger(__name__)
def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
agents = selectAgents(api)
history_days = Selector.select_value(
prompt="Enter how many days of history to pull (1–365): ",
value_type=int,
valid_range=(1, 365),
)
if not agents or not history_days:
print(colorText("No agents selected or invalid history range.", "red"))
return
historical_date = (datetime.now() - timedelta(days=history_days)).strftime(
"%Y-%m-%d"
)
today = datetime.now().strftime("%Y-%m-%d")
all_history = []
for agent in agents:
try:
exechistory = api.history_execution(today, historical_date, agent.hostname)
except Exception as e:
print(
colorText(
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
)
)
continue
if isinstance(exechistory, list):
for block in exechistory:
record = {
"Command": block.get("commandline", "N/A"),
"Date": block.get("datetime", "N/A"),
"Filename": block.get("filename", "N/A"),
"Policy Name": block.get("policyname", "N/A"),
"Hostname": block.get("hostname", "N/A"),
"Hash": block.get("sha256", "N/A"),
}
all_history.append(record)
if not outputjson:
for key, value in record.items():
print(colorText(f"{key}: {value}", "green"))
print("\n")
else:
print(
colorText(f"No execution history found for {agent.hostname}.", "yellow")
)
if outputjson:
print(json.dumps(all_history, indent=2))
def findAllAgents(api):
# Step 1: Load data from API
policies = [Policy(**row["data"]) for _, row in api.policy_find_all().iterrows()]
agents = [Agent(**row["data"]) for _, row in api.agent_find_all().iterrows()]
for agent in agents:
agent.enrich_with_policies(policies)
return agents
def findAgents(api, return_dataframe):
agents = selectAgents(api)
working_dir = load_env("WORKING_DIR")
if not agents:
logging.warning("No agents or policies found.")
print("No agents matched the criteria.")
return
# Convert enriched agents to DataFrame
agent_dicts = [asdict(agent) for agent in agents]
agent_df = pd.DataFrame(agent_dicts)
if return_dataframe:
logging.debug("Returning DataFrame to caller.")
return agent_df
# Otherwise, print and optionally export
print(agent_df)
logging.debug("Displayed DataFrame to console.")
user_input = (
get_sanitized_input(
"\nWould you like to export the results to a CSV file? (y/n): "
)
.strip()
.lower()
)
if user_input == "y":
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
filename = f"agentsearch_{timestamp}.csv"
file_path = os.path.join(str(working_dir), filename)
agent_df.to_csv(file_path, index=False)
logging.info(f"Exported DataFrame to {file_path}")
print(
colorText(
f"\n✓ Matched devices exported to: {working_dir}\\{filename}",
"green",
)
)
else:
logging.debug("User declined to export the DataFrame.")
def collect_device_names() -> List[str]:
print(colorText("🖥��Â Device Search", "cyan"))
print(
colorText(
"Enter the device hostnames you'd like to search for, one per line.", "cyan"
)
)
print(
colorText(
"When you're done, press Enter twice (Three times if you have a single device).\n",
"cyan",
)
)
print(colorText("Example:", "cyan"))
print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan"))
print(colorText("Paste or type your device names below:", "white"))
device_input_lines = []
empty_line_count = 0
valid_line_pattern = re.compile(r"^[a-zA-Z0-9_\- ]+$")
while True:
line = get_sanitized_input("")
stripped_line = line.strip()
if stripped_line == "":
empty_line_count += 1
if empty_line_count == 2:
break
continue
else:
empty_line_count = 0
if valid_line_pattern.match(stripped_line):
device_input_lines.append(stripped_line)
else:
print(
colorText(
f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
"yellow",
)
)
return [name for name in device_input_lines if name]
def choose_match_type() -> bool:
print(colorText("Use exact match? (Y for exact, N for fuzzy):", "white"))
return get_sanitized_input("").strip().lower() in ["y", "yes"]
def match_agents(
device_names: List[str], agents: List["Agent"], use_exact: bool
) -> List["Agent"]:
if use_exact:
return [
agent
for agent in agents
if agent.hostname.lower() in [name.lower() for name in device_names]
]
else:
pattern = "|".join(map(re.escape, device_names))
regex = re.compile(pattern, re.IGNORECASE)
return [agent for agent in agents if regex.search(agent.hostname)]
def show_unmatched(
device_names: List[str], matched_agents: List["Agent"], use_exact: bool
):
if use_exact:
unmatched = [
name
for name in device_names
if not any(
agent.hostname.lower() == name.lower() for agent in matched_agents
)
]
else:
unmatched = [
name
for name in device_names
if not any(
re.search(re.escape(name), agent.hostname, re.IGNORECASE)
for agent in matched_agents
)
]
if unmatched:
logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
def enrich_agents(agents: List["Agent"], policies: List["Policy"]):
for agent in agents:
agent.enrich_with_policies(policies)
def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
device_names = collect_device_names()
if not device_names:
logger.debug("No device names entered")
print(colorText("⚠️ No device names entered.", "red"))
return []
use_exact = choose_match_type()
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()]
matched_agents = match_agents(device_names, agents, use_exact)
matched_agents.sort(key=lambda agent: agent.hostname.lower())
show_unmatched(device_names, matched_agents, use_exact)
if not matched_agents:
logger.debug("❌ No matching devices found.")
print(colorText("❌ No matching devices found.", "red"))
return []
print(colorText(f"✓ Found {len(matched_agents)} matching device(s).", "green"))
logger.info("Matched agent hostnames:")
rows = (len(matched_agents) + 2) // 3 # 3 columns
for row in range(rows):
line = ""
for col in range(3):
idx = row + col * rows
if idx < len(matched_agents):
line += f"{matched_agents[idx].hostname:<30}"
logger.info(line)
matched_agents = Selector.select_with_mode(
matched_agents,
label_func=lambda agent: agent.hostname,
header="Matched Devices:",
)
if not matched_agents:
logger.debug("❌ No matching devices remain after refinement.")
print(colorText("❌ No matching devices remain after refinement.", "red"))
return []
enrich_agents(matched_agents, policies)
return matched_agents
def moveAgentToRelatedPolicy(
api: AirlockAPIWrapper,
agent: Agent,
mode: str = "audit",
):
"""
Moves an agent between audit and enforcement policies based on the mode.
Args:
api: AirlockAPIWrapper instance.
agent: Agent object.
policy_relationship_map: Dict mapping enforcement â–€ –€™ audit.
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
"""
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
if mode == "audit":
if agent.groupid in policy_relationship_map:
target_policy = policy_relationship_map[agent.groupid]
elif agent.groupid in policy_relationship_map.values():
logger.debug(
f"Agent {agent.hostname} is already in an audit group. No action needed."
)
print(
f"Agent {agent.hostname} is already in an audit group. No action needed."
)
return
else:
logger.warning(
f"Error: No corresponding audit policy found for groupid: {agent.groupid}."
)
return
elif mode == "enforcement":
inverse_map = {v: k for k, v in policy_relationship_map.items()}
if agent.groupid in inverse_map:
target_policy = inverse_map[agent.groupid]
elif agent.groupid in inverse_map.values():
logger.info(
f"Agent {agent.hostname} is already in an enforcement group. No action needed."
)
return
else:
logger.warning(
f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}."
)
return
else:
logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.")
return
result = api.agent_move(agent.agentid, target_policy)
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/policyhandler.py
-234
View File
@@ -1,234 +0,0 @@
# 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 gc
import json
import logging
import os
import sys
from bson import ObjectId
import pandas as pd
import tqdm
from models.policy import Policy
from services.API import AirlockAPIWrapper
from utils.setup import get_base_directory
from utils.utils import colorText
logger = logging.getLogger(__name__)
def pullPolicyExechistories(
api: AirlockAPIWrapper,
policy: Policy,
type: list,
days,
outputjson: bool,
):
file_path = f"{get_base_directory()}\\cache\\chunkinator.json"
# Ensure the file exists
if not os.path.exists(file_path):
with open(file_path, "w") as file:
json.dump({"error": "Success", "response": {"exechistories": []}}, file)
logger.debug(f"File '{file_path}' has been created.")
else:
logger.debug(f"File '{file_path}' already exists.")
checkpoint = str(skipback(days))
json_output = {"error": "Success", "response": {"exechistories": []}}
with tqdm.tqdm(
file=sys.stdout,
leave=True,
total=10000,
desc=f"Checkpoint Progress: {checkpoint}",
colour="blue",
initial=1,
) as filebar:
with tqdm.tqdm(
file=sys.stdout,
leave=True,
total=100,
desc=f"Total of {policy} Complete: ",
) as pbar:
while True:
histories = api.history_logging(
type=type, checkpoint=checkpoint, policy=[policy.name]
)
# Ensure histories is a list of dictionaries
if not isinstance(histories, list) or not all(
isinstance(h, dict) for h in histories
):
logger.error(
"Unexpected response format from API. Expected list of dictionaries."
)
break
filebar.total = len(histories)
if not histories:
break
for index, history_item in enumerate(histories):
if (
"checkpoint" not in history_item
or "datetime" not in history_item
):
continue # Skip malformed entries
# Update checkpoint on last item
if index == len(histories) - 1:
checkpoint = history_item[
"checkpoint"
] # pyright: ignore[reportArgumentType]
filebar.desc = f"Checkpoint Progress: {checkpoint}"
break
try:
history_date = datetime.datetime.strptime(
history_item["datetime"].replace(
" +0000 UTC", ""
), # pyright: ignore[reportArgumentType]
"%Y-%m-%dT%H:%M:%SZ",
).date()
except ValueError:
continue # Skip if date format is invalid
if (
datetime.date.today() - datetime.timedelta(days=days)
) <= history_date:
json_output["response"]["exechistories"].append(history_item)
filebar.update(1)
filebar.refresh()
# Deduplicate entries
seen = {}
if os.path.exists(file_path):
with open(file_path, "r") as file:
existing_data = json.load(file)
combined = (
existing_data["response"]["exechistories"]
+ json_output["response"]["exechistories"]
)
else:
combined = json_output["response"]["exechistories"]
for entry in combined:
key = (
entry.get("sha256"),
entry.get("filename"),
entry.get("hostname"),
)
seen[key] = entry
deduplicated = list(seen.values())
with open(file_path, "w") as file:
json.dump(
{
"error": "Success",
"response": {"exechistories": deduplicated},
},
file,
)
json_output["response"]["exechistories"].clear()
# Update progress bar based on last valid item
try:
last_date = datetime.datetime.strptime(
history_item["datetime"].replace(" +0000 UTC", ""), # type: ignore
"%Y-%m-%dT%H:%M:%SZ",
).date()
date_diff = datetime.date.today() - last_date
percentage_diff = (
((days + 10) - date_diff.days) / (days + 10)
) * 100
pbar.n = round(percentage_diff)
pbar.set_description_str(f"Total of {policy} Complete: ")
pbar.refresh()
except Exception:
pass
filebar.n = 1
# Final output
with open(file_path, "r") as file:
final_output = json.load(file)
os.remove(file_path)
return json.dumps(final_output) if outputjson else None
def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
import airlock_libs
executionhist_policy = pd.DataFrame()
exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days)
if exehist is not None:
data = json.loads(exehist)
executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
if not executionhist_policy.empty:
executionhist_policy = executionhist_policy[
[
"datetime",
"sha256",
"publisher",
"filename",
"hostname",
"username",
"pprocess",
"gprocess",
"commandline",
]
]
executionhist_policy["policy"] = policy # Add policy column here
executionhist_policy = executionhist_policy.drop_duplicates(
subset=["sha256", "filename", "hostname"]
)
executionhist_policy = executionhist_policy.sort_values(
by=["sha256", "filename"]
)
logger.debug(f"Staging of Execution history for policy: {policy} is complete")
print(
colorText(
f"Staging of Execution history for policy: {policy} is complete",
"green",
)
)
del data
del exehist
gc.collect()
return executionhist_policy
def skipback(days):
"""
Generate a MongoDB ObjectId for a given number of days ago from today.
"""
adjusted_days = days
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
days=adjusted_days
)
timestamp = int(date_days_ago.timestamp())
hex_timestamp = format(timestamp, "08x")
objectid_hex = hex_timestamp + "0000000000000000"
return ObjectId(objectid_hex)
utils/configmanager.py
-7
View File
@@ -336,10 +336,3 @@ def load_env_json(key: str, default: str = "[]") -> Any:
except Exception as e:
logger.error(f"Failed to parse {key}: {e}")
return json.loads(default)
# Backwards compatibility aliases (deprecated - use get_system_value instead)
get_protected_value = get_system_value
get_protected_json = get_system_json
load_protected_config = load_system_config
PROTECTED_KEYS = SYSTEM_CONFIG_KEYS # For backwards compatibility
utils/executionfetcher.py
+78
View File
@@ -0,0 +1,78 @@
# 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 gc
import json
import logging
from bson import ObjectId
import pandas as pd
from services.API import AirlockAPIWrapper
logger = logging.getLogger(__name__)
def getExecutions(api: AirlockAPIWrapper, policy, type, days):
import airlock_libs
executionhist_policy = pd.DataFrame()
exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days)
if exehist is not None:
data = json.loads(exehist)
executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
if not executionhist_policy.empty:
executionhist_policy = executionhist_policy[
[
"datetime",
"sha256",
"publisher",
"filename",
"hostname",
"username",
"pprocess",
"gprocess",
"commandline",
]
]
executionhist_policy["policy"] = policy # Add policy column here
executionhist_policy = executionhist_policy.drop_duplicates(
subset=["sha256", "filename", "hostname"]
)
executionhist_policy = executionhist_policy.sort_values(
by=["sha256", "filename"]
)
logger.debug(f"Staging of Execution history for policy: {policy} is complete")
del data
del exehist
gc.collect()
return executionhist_policy
def skipback(days):
"""
Generate a MongoDB ObjectId for a given number of days ago from today.
"""
adjusted_days = days
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
days=adjusted_days
)
timestamp = int(date_days_ago.timestamp())
hex_timestamp = format(timestamp, "08x")
objectid_hex = hex_timestamp + "0000000000000000"
return ObjectId(objectid_hex)
services/security.py → utils/security.py
-14
View File
@@ -169,17 +169,3 @@ def getAPI(USERNAME, SERVICE_NAME):
logging.warning(
"Password does not meet complexity requirements. Try again."
)
class APIKeyManager:
_api_key = None
@classmethod
def load(cls, service: str, username: str, password: str):
cls._api_key = retrieve_api_key(service, username, password)
@classmethod
def get(cls) -> str:
if cls._api_key is None:
raise ValueError("API key not loaded. Call APIKeyManager.load() first.")
return cls._api_key
utils/selector.py
-357
View File
@@ -1,357 +0,0 @@
# 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
from typing import Any, Callable, List, Optional, Union
import pandas as pd
from utils.utils import colorText, get_sanitized_input
logger = logging.getLogger(__name__)
class Selector:
@staticmethod
def _get_sorted_items(
items: List[Any], label_func: Callable[[Any], str]
) -> List[Any]:
return sorted(items, key=lambda item: label_func(item).lower())
@staticmethod
def _display_choices(
items: List[Any],
label_func: Callable[[Any], str],
num_columns: int = 3,
header: str = "Available Choices:",
) -> None:
# Force single column if items are DataFrame rows
if items and isinstance(items[0], (pd.Series, dict)):
num_columns = 1
rows = (len(items) + num_columns - 1) // num_columns
print(f"\n{header}")
for row in range(rows):
line = ""
for col in range(num_columns):
idx = row + col * rows
if idx < len(items):
label = label_func(items[idx])
line += f"{idx + 1}: {label:<30}"
print(line)
@staticmethod
def _display_selected_items(
selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 3
) -> None:
print(colorText("\nCurrent selections:", "cyan"))
if not selected:
print(" (none)")
return
sorted_selected = sorted(selected, key=lambda item: label_func(item).lower())
rows = (len(sorted_selected) + num_columns - 1) // num_columns
for row in range(rows):
line = ""
for col in range(num_columns):
idx = row + col * rows
if idx < len(sorted_selected):
label = label_func(sorted_selected[idx])
line += f"{label:<30}"
print(line)
@staticmethod
def _parse_selection_input(input_str: str, max_index: int) -> List[int]:
selections = []
for part in input_str.split(","):
part = part.strip()
if "-" in part:
try:
start, end = map(int, part.split("-"))
selections.extend(range(start, end + 1))
except ValueError:
continue
elif part.isdigit():
selections.append(int(part))
return [i for i in selections if 1 <= i <= max_index]
@staticmethod
def _select_from_list(
items: List[Any],
label_func: Callable[[Any], str],
allow_multiple: bool = False,
prompt_each: bool = False,
header: str = "Available Choices:",
num_columns: int = 3,
) -> Union[Optional[Any], List[Any]]:
if not items:
logger.warning("No items available for selection.")
return None
full_sorted_items = Selector._get_sorted_items(items, label_func)
remaining_items = full_sorted_items.copy()
selected = []
if allow_multiple:
while True:
Selector._display_choices(
remaining_items, label_func, num_columns=num_columns, header=header
)
Selector._display_selected_items(
selected, label_func, num_columns=num_columns
)
choice = (
get_sanitized_input(
"Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: "
)
.strip()
.lower()
)
if choice == "q":
break
elif choice == "r":
selected.clear()
remaining_items = full_sorted_items.copy()
print(colorText("🔄 Selections reset.", "yellow"))
continue
indices = Selector._parse_selection_input(choice, len(remaining_items))
newly_selected = []
for index in indices:
item = remaining_items[index - 1]
if item not in selected:
selected.append(item)
newly_selected.append(item)
if prompt_each:
logger.info(f"Selected: {label_func(item)}")
else:
logger.warning("Item already selected.")
remaining_items = [
item for item in remaining_items if item not in newly_selected
]
return selected if selected else None
else:
Selector._display_choices(
full_sorted_items, label_func, num_columns=num_columns, header=header
)
try:
choice = int(get_sanitized_input("Select one item by number: "))
if 1 <= choice <= len(full_sorted_items):
selected_item = full_sorted_items[choice - 1]
logger.info(f"Selected: {label_func(selected_item)}")
return selected_item
else:
logger.warning("Selection out of range.")
except ValueError:
logger.warning("Invalid input.")
return None
@staticmethod
def select_with_mode(
items: List[Any],
label_func: Callable[[Any], str],
header: str = "Available Choices:",
) -> List[Any]:
print(
colorText(
"Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):",
"white",
)
)
mode = get_sanitized_input("").strip().lower()
if mode == "a":
return items
selected = Selector._select_from_list(
items,
label_func=label_func,
allow_multiple=True,
prompt_each=False,
header=header,
)
if not selected:
return items
if mode == "i":
print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
return selected
elif mode == "e":
print(colorText(f"👫 Excluded {len(selected)} item(s).", "yellow"))
return [item for item in items if item not in selected]
else:
print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
return items
@staticmethod
def select_objects(
objects: List[Any], allow_multiple: bool = False, prompt_each: bool = False
) -> Union[Optional[Any], List[Any]]:
return Selector._select_from_list(
objects,
label_func=lambda obj: getattr(obj, "name", str(obj)),
allow_multiple=allow_multiple,
prompt_each=prompt_each,
header="Available Objects:",
)
@staticmethod
def select_string(
options: List[str], allow_multiple: bool = False, prompt_each: bool = False
) -> Union[Optional[str], List[str]]:
return Selector._select_from_list(
options,
label_func=str,
allow_multiple=allow_multiple,
prompt_each=prompt_each,
header="Available Options:",
)
@staticmethod
def select_int(
options: List[int], allow_multiple: bool = False, prompt_each: bool = False
) -> Union[Optional[int], List[int]]:
return Selector._select_from_list(
options,
label_func=lambda x: str(x),
allow_multiple=allow_multiple,
prompt_each=prompt_each,
header="Available Integers:",
)
@staticmethod
def select_value(
prompt: str,
value_type: type = int,
valid_range: Optional[tuple] = None,
allow_quit: bool = False,
) -> Optional[Any]:
while True:
user_input = get_sanitized_input(prompt).strip().lower()
if allow_quit and user_input == "q":
logger.info("User opted to quit value selection.")
return None
try:
value = value_type(user_input)
if valid_range:
min_val, max_val = valid_range
if not (min_val <= value <= max_val):
logger.warning(f"Value out of range ({min_val}–{max_val}).")
continue
logger.info(f"User selected value: {value}")
return value
except ValueError:
logger.warning(f"Invalid input. Expected a {value_type.__name__}.")
@staticmethod
def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool:
while True:
response = get_sanitized_input(prompt).strip().lower()
if response in ["y", "yes"]:
logger.info("User confirmed action.")
return True
elif response in ["n", "no"]:
logger.info("User declined action.")
return False
else:
logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")
@staticmethod
def select_dataframe_rows(
df: pd.DataFrame,
columns: Optional[List[str]] = None,
allow_multiple: bool = False,
prompt_each: bool = False,
header: str = "Available Rows:",
) -> List[pd.Series]:
if df.empty:
print("DataFrame is empty.")
return []
if columns:
df = df[columns]
items = [row for _, row in df.iterrows()]
def label_func(row):
return str(row.to_dict())
result = Selector._select_from_list(
items,
label_func=label_func,
allow_multiple=allow_multiple,
prompt_each=prompt_each,
header=header,
)
if isinstance(result, pd.Series):
return [result]
elif isinstance(result, list):
return result
else:
return []
@staticmethod
def select_dataframe_with_mode(
df: pd.DataFrame,
columns: Optional[List[str]] = None,
header: str = "Available Rows:",
) -> List[pd.Series]:
if df.empty:
print("⚠️ DataFrame is empty.")
return []
# Filter columns if specified
if columns:
df = df[columns]
items = df.to_dict("records")
def label_func(row):
return " | ".join(str(row[col]) for col in df.columns)
# Show rows first
print(colorText(header, "cyan"))
for i, row in enumerate(items):
print(f"{i}: {label_func(row)}")
# Prompt for mode once
print(
colorText(
"\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):",
"white",
)
)
mode = get_sanitized_input("").strip().lower()
if mode == "a":
return [pd.Series(row) for row in items]
# Prompt for selection only once
selected = Selector._select_from_list(
items,
label_func=label_func,
allow_multiple=True,
prompt_each=False,
header=header,
)
if not selected:
return [pd.Series(row) for row in items]
if mode == "i":
print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
return [pd.Series(row) for row in selected]
elif mode == "e":
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 no rows.", "yellow"))
return []
utils/setup.py
+26 -3
View File
@@ -83,7 +83,18 @@ def get_base_directory() -> Path:
return home / ".local" / "share" / "Loxide"
def configure_logging(log_dir: Path, log_level: str = "INFO"):
def configure_logging(log_dir: Path, cache_dir: Path, log_level: str = "INFO"):
"""
Configure logging and return function to attach notification handler.
Args:
log_dir: Directory for log files
cache_dir: Directory for cache files (used by version checker)
log_level: Logging level string
Returns:
Function to attach notification handler to Textual app
"""
log_file = log_dir / "Loxide.log"
config = {
@@ -134,7 +145,8 @@ def configure_logging(log_dir: Path, log_level: str = "INFO"):
# Return a function to attach the notification handler once the app is created
def attach_notification_handler(app):
"""Attach the Textual notification handler to the root logger."""
"""Attach the Textual notification handler and version checker to the app."""
# Attach logging handler
handler = TextualNotificationHandler(app)
handler.setLevel(logging.WARNING) # Only WARNING and above
formatter = logging.Formatter("%(name)s: %(message)s")
@@ -142,6 +154,17 @@ def configure_logging(log_dir: Path, log_level: str = "INFO"):
logging.getLogger().addHandler(handler)
logging.getLogger().debug("✅ Textual notification handler attached.")
# Attach version checker (checks in background, notifies if update available)
try:
from utils.versionchecker import create_update_notifier
create_update_notifier(app, cache_dir=cache_dir)
logging.getLogger().debug("✅ Version checker attached.")
except ImportError as e:
logging.getLogger().debug(f"Version checker not available: {e}")
except Exception as e:
logging.getLogger().warning(f"Could not initialize version checker: {e}")
return attach_notification_handler
@@ -173,7 +196,7 @@ def setup():
# Configure logging with system-defined log level
log_level = get_system_value("LOG_LEVEL", str, "INFO")
attach_handler = configure_logging(dirs["logs"], log_level)
attach_handler = configure_logging(dirs["logs"], dirs["cache"], log_level)
# Load user config (mutable)
load_user_config(dirs["config"])
utils/utils.py
-437
View File
@@ -19,81 +19,10 @@ import os
import platform
import re
import subprocess
import tempfile
import tkinter as tk
from tkinter import filedialog
import pandas as pd
logger = logging.getLogger(__name__)
def import_to_dataframe(file_path: str) -> pd.DataFrame:
df = pd.DataFrame()
try:
if not os.path.exists(file_path):
print(colorText(f"Error: File '{file_path}' does not exist.", "red"))
return df
ext = os.path.splitext(file_path)[1].lower()
if ext == ".csv":
df = pd.read_csv(file_path)
elif ext == ".parquet":
df = pd.read_parquet(file_path)
else:
print(colorText(f"Error: Unsupported file extension '{ext}'.", "red"))
return df
if df.empty:
print(colorText("Error: File has headers but no data rows.", "red"))
else:
print(colorText(f"Data loaded successfully from {file_path}", "green"))
return df
except pd.errors.EmptyDataError:
print(
colorText(
"Notice: CSV file is completely empty, falling back to empty frame",
"white",
)
)
return pd.DataFrame()
except Exception as e:
print(colorText(f"Error reading file: {e}", "red"))
return pd.DataFrame()
def choose_directory():
root = tk.Tk()
root.withdraw() # Hide the main window
directory = filedialog.askdirectory(title="Select a Directory")
print("Selected directory:", directory)
return directory
def choose_file(initial_directory=None, required_substring=None):
"""Open a file dialog and ensure the selected file contains a required substring."""
while True:
root = tk.Tk()
root.withdraw() # Hide the main window
file_path = filedialog.askopenfilename(initialdir=initial_directory)
if not file_path:
print("No file selected.")
return None
if required_substring and required_substring not in file_path:
print(
f"The selected file must contain '{required_substring}' in its path or name. Please try again."
)
else:
return file_path
def get_sanitized_input(prompt: str) -> str:
while True:
user_input = input(prompt)
@@ -151,235 +80,6 @@ def irtang():
)
def section_header(title):
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(colorText(f" ------------- {title} -------------", "cyan"))
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
def areYouSure():
print(
colorText(
"🛑****************************************************************************************************************************************🛑",
"red",
)
)
print(
colorText(
"⚠️=========================================================================================================================================⚠️",
"yellow",
)
)
print(
colorText(
"🛑========================================================================================================================================🛑",
"red",
)
)
print(
colorText(
"⚠️-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------⚠️",
"yellow",
)
)
print(
colorText(
"🛑========================================================================================================================================🛑",
"red",
)
)
print(
colorText(
"⚠️=========================================================================================================================================⚠️",
"yellow",
)
)
print(
colorText(
"🛑****************************************************************************************************************************************🛑",
"red",
)
)
def locked():
print(
colorText(
r"""
████████████████████████████████████████████████████████████████
███ ██
██ ██████ ███
██ ████████████ ███
██ ████ ███ ███
██ ███ ███ ███
██ ███ ███ ███
██ ▒████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ███
███ ███
████████████████████████████████████████████████████████████████████
▒██████████████████████████████████████████████████████████████████▒
▒████
▒████
▓██████████████████████████████████████████
█████████████████████████████████████████████░
""",
"yellow",
)
)
def printDeviceEnforceChecklist():
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
"cyan",
)
)
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
"\nSequentually follow these steps to prepare a policy for enforcement:",
"white",
)
)
print(
colorText(
"\n1. Choose which originating policy or policies to move to enforcement",
"cyan",
)
)
print(
colorText(
"2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
"cyan",
)
)
print(colorText("3. Manually review the files:", "cyan"))
print(
colorText(
" 'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'",
"cyan",
)
)
print(
colorText(
" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
"cyan",
)
)
print(
colorText(
" If metarules need to be created, please make note of them, and remove the row from the csv.",
"cyan",
)
)
print(
colorText(
" When complete, save both csv files to the directory 'approved' and choose this option.",
"cyan",
)
)
print(
colorText(
" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
"cyan",
)
)
print(
colorText(
"4. Manually review the file 'needs_approved\\paths_needing_review.csv'",
"cyan",
)
)
print(
colorText(
" Remove the rows containing path exclusions you do not approve of",
"cyan",
)
)
print(
colorText(
" When complete, save the csv file to the directory 'approved'", "cyan"
)
)
print(
colorText(
" Do the same process with the list of publishers forthe same directories",
"cyan",
)
)
print(colorText(" Preflight Lists will be generated", "cyan"))
print(
colorText(
"5. Choose the destination policy and parent and child allow list", "cyan"
)
)
print(
colorText(
"6. Test ------------------------------------------------------", "cyan"
)
)
print(colorText(" Print rather than apply selected data.", "cyan"))
print(
colorText(
"7. Liftoff ------------------------------------------------------", "cyan"
)
)
print(
colorText(
" Apply path exclusions according to allowed and approved paths",
"cyan",
)
)
print(
colorText(" Apply signed or attested hashes to Parent Allow List", "cyan")
)
print(
colorText(
" Apply approved, but unsigned hashes to the Child Allow List", "cyan"
)
)
print(
colorText(
"R. Remove/Reset Generated data - will prompt to allow keeping execution history",
"cyan",
)
)
print(colorText("B. Back", "cyan"))
def colorText(text, color):
colors = {
"red": "\033[91m",
@@ -394,133 +94,6 @@ def colorText(text, color):
return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"
def formatHTML(df, output_html_path=None, overwrite=True):
from datetime import datetime
# Get current date and filename for subtitle
today = datetime.now().strftime("%d %B %Y") # Changed to "Day Month Year"
filename = output_html_path.replace(".html", "") if output_html_path else "Report"
dark_css = """
<style>
body {
background-color: #000000;
margin: 0;
padding: 0;
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
color: #f8f8f2;
}
.header {
text-align: center;
margin: 20px auto;
padding: 10px;
border-bottom: 2px solid #ffd700;
max-width: 95%;
}
.header h1 {
color: #ffd700;
margin: 0;
font-size: 32px;
}
.header p {
color: #00bfff;
margin: 5px 0 0 0;
font-size: 18px;
}
.table-container {
overflow-y: scroll;
margin: 0 auto;
width: 95%;
max-height: calc(80vh - 100px);
display: block;
border: 1px solid #3a3a4d;
margin-bottom: 0;
}
table {
border-collapse: collapse;
font-size: 14px;
background-color: #1e1e2f;
color: #f8f8f2;
width: max-content;
}
th, td {
border: 1px solid #3a3a4d;
text-align: left;
padding: 10px;
max-width: 300px;
word-wrap: break-word;
overflow-wrap: break-word;
}
/* First column: no wrap */
td:nth-child(1), th:nth-child(1) {
white-space: nowrap;
max-width: none !important;
word-wrap: normal !important;
}
th {
background-color: #2e2e40;
color: #ffd700;
position: sticky;
top: 0;
z-index: 10;
}
tr:nth-child(even) {
background-color: #262638;
}
tr:hover {
background-color: #33334d;
color: #00bfff;
}
/* Custom scrollbar styling */
.table-container::-webkit-scrollbar {
width: 12px;
}
.table-container::-webkit-scrollbar-track {
background: #1e1e2f;
}
.table-container::-webkit-scrollbar-thumb {
background-color: #3a3a4d;
border-radius: 6px;
}
</style>
"""
header = f"""
<div class="header">
<h1>Airlock Tools</h1>
<p>{filename} - {today}</p>
</div>
"""
html_table = df.to_html(index=False, escape=False)
styled_html = (
f"<html>\n"
f"<head><title>Airlock Tools Report</title></head>\n"
f"<body>\n"
f"{dark_css}\n"
f"{header}\n"
f"<div class='table-container'>\n"
f" {html_table}\n"
f"</div>\n"
f"</body>\n"
f"</html>"
)
if output_html_path:
with open(output_html_path, "w", encoding="utf-8") as f:
f.write(styled_html)
print(f"✅ Styled table saved to '{output_html_path}'")
elif overwrite:
with tempfile.NamedTemporaryFile(
suffix=".html", delete=False, mode="w", encoding="utf-8"
) as f:
f.write(styled_html)
temp_path = f.name
print(f"✅ Styled table saved to temporary file: {temp_path}")
else:
return styled_html
def open_directory(path):
system = platform.system()
@@ -530,13 +103,3 @@ def open_directory(path):
subprocess.run(["xdg-open", path])
else:
raise OSError(f"Unsupported operating system: {system}")
def print_x_wide(items: list, width: int):
for i in range(0, len(items), width):
row = items[i : i + width]
print(" | ".join(row))
def clear_screen():
os.system("cls" if os.name == "nt" else "clear")
utils/versionchecker.py
+560
View File
@@ -0,0 +1,560 @@
# 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/>.
"""
Version checking and update notification system for Loxide.
Checks against Gitea releases at:
https://git.racooncity.org/brotoskyj/AirlockTools/releases
"""
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
import json
import logging
from pathlib import Path
import re
import threading
from typing import Callable, Optional
import requests
logger = logging.getLogger(__name__)
# Current application version - UPDATE THIS ON EACH RELEASE
__version__ = "0.7.0"
# Gitea release API configuration
GITEA_API_BASE = "https://git.racooncity.org/api/v1"
REPO_OWNER = "brotoskyj"
REPO_NAME = "AirlockTools"
RELEASES_URL = f"{GITEA_API_BASE}/repos/{REPO_OWNER}/{REPO_NAME}/releases"
RELEASES_PAGE_URL = f"https://git.racooncity.org/{REPO_OWNER}/{REPO_NAME}/releases"
# How often to check for updates (in hours)
CHECK_INTERVAL_HOURS = 24
@dataclass
class ReleaseInfo:
"""Information about a release."""
tag_name: str
version: tuple # Parsed semantic version (major, minor, patch)
name: str
body: str # Release notes
published_at: datetime
html_url: str
download_url: Optional[str] = None # URL to download the release asset
is_prerelease: bool = False
@dataclass
class UpdateCheckResult:
"""Result of an update check."""
current_version: str
latest_version: Optional[str]
update_available: bool
release_info: Optional[ReleaseInfo]
error: Optional[str] = None
def parse_version(version_str: str) -> Optional[tuple]:
"""
Parse a version string into a comparable tuple.
Supports formats: v1.2.3, 1.2.3, v1.2, 1.2
Returns (major, minor, patch) tuple or None if parsing fails.
"""
if not version_str:
return None
# Strip 'v' prefix if present
clean = version_str.lstrip("vV").strip()
# Match semantic version pattern
match = re.match(r"^(\d+)(?:\.(\d+))?(?:\.(\d+))?", clean)
if not match:
return None
major = int(match.group(1))
minor = int(match.group(2)) if match.group(2) else 0
patch = int(match.group(3)) if match.group(3) else 0
return (major, minor, patch)
def compare_versions(v1: tuple, v2: tuple) -> int:
"""
Compare two version tuples.
Returns:
-1 if v1 < v2
0 if v1 == v2
1 if v1 > v2
"""
for a, b in zip(v1, v2):
if a < b:
return -1
if a > b:
return 1
return 0
def get_current_version() -> str:
"""Get the current application version."""
return __version__
def _parse_release_response(release_data: dict) -> Optional[ReleaseInfo]:
"""Parse a release from Gitea API response."""
try:
tag_name = release_data.get("tag_name", "")
version = parse_version(tag_name)
if not version:
logger.debug(f"Could not parse version from tag: {tag_name}")
return None
# Parse published date
published_str = release_data.get("published_at", "")
try:
published_at = datetime.fromisoformat(published_str.replace("Z", "+00:00"))
except (ValueError, AttributeError):
published_at = datetime.now(UTC)
# Get download URL from assets if available
download_url = None
assets = release_data.get("assets", [])
for asset in assets:
# Prefer .exe or .zip files
name = asset.get("name", "").lower()
if name.endswith((".exe", ".zip", ".msi")):
download_url = asset.get("browser_download_url")
break
return ReleaseInfo(
tag_name=tag_name,
version=version,
name=release_data.get("name", tag_name),
body=release_data.get("body", ""),
published_at=published_at,
html_url=release_data.get("html_url", RELEASES_PAGE_URL),
download_url=download_url,
is_prerelease=release_data.get("prerelease", False),
)
except Exception as e:
logger.warning(f"Failed to parse release data: {e}")
return None
def fetch_latest_release(
include_prerelease: bool = False, timeout: int = 10
) -> Optional[ReleaseInfo]:
"""
Fetch the latest release from Gitea.
Args:
include_prerelease: Whether to include pre-release versions
timeout: Request timeout in seconds
Returns:
ReleaseInfo for the latest release, or None if fetch fails
"""
try:
response = requests.get(
RELEASES_URL,
params={"limit": 10}, # Get last 10 releases to find latest stable
timeout=timeout,
headers={"Accept": "application/json"},
)
response.raise_for_status()
releases = response.json()
if not releases:
logger.debug("No releases found")
return None
# Find the latest release (first non-prerelease if we're excluding them)
for release_data in releases:
release_info = _parse_release_response(release_data)
if release_info is None:
continue
if include_prerelease or not release_info.is_prerelease:
return release_info
# If all are prereleases and we're excluding them, return the first one anyway
# but log a warning
if releases:
logger.debug("All releases are pre-releases")
return _parse_release_response(releases[0])
return None
except requests.exceptions.Timeout:
logger.warning("Timeout fetching releases from Gitea")
return None
except requests.exceptions.RequestException as e:
logger.warning(f"Failed to fetch releases: {e}")
return None
except (json.JSONDecodeError, KeyError) as e:
logger.warning(f"Failed to parse release response: {e}")
return None
def check_for_updates(include_prerelease: bool = False) -> UpdateCheckResult:
"""
Check if a newer version is available.
Args:
include_prerelease: Whether to consider pre-release versions
Returns:
UpdateCheckResult with the check results
"""
current = get_current_version()
current_parsed = parse_version(current)
if not current_parsed:
return UpdateCheckResult(
current_version=current,
latest_version=None,
update_available=False,
release_info=None,
error="Could not parse current version",
)
release_info = fetch_latest_release(include_prerelease=include_prerelease)
if release_info is None:
return UpdateCheckResult(
current_version=current,
latest_version=None,
update_available=False,
release_info=None,
error="Could not fetch release information",
)
is_newer = compare_versions(release_info.version, current_parsed) > 0
return UpdateCheckResult(
current_version=current,
latest_version=release_info.tag_name,
update_available=is_newer,
release_info=release_info,
)
class VersionChecker:
"""
Background version checker that periodically checks for updates
and can notify the application when updates are available.
"""
def __init__(
self,
cache_dir: Optional[Path] = None,
check_interval_hours: int = CHECK_INTERVAL_HOURS,
on_update_available: Optional[Callable[[UpdateCheckResult], None]] = None,
):
"""
Initialize the version checker.
Args:
cache_dir: Directory to store last check timestamp
check_interval_hours: Hours between automatic checks
on_update_available: Callback when update is available
"""
self.cache_dir = cache_dir
self.check_interval = timedelta(hours=check_interval_hours)
self.on_update_available = on_update_available
self._last_check: Optional[datetime] = None
self._last_result: Optional[UpdateCheckResult] = None
self._check_thread: Optional[threading.Thread] = None
self._dismissed_version: Optional[str] = None
# Load cached state
self._load_cache()
@property
def cache_file(self) -> Optional[Path]:
if self.cache_dir:
return self.cache_dir / "version_check_cache.json"
return None
def _load_cache(self) -> None:
"""Load cached check state."""
if not self.cache_file or not self.cache_file.exists():
return
try:
with open(self.cache_file, "r") as f:
data = json.load(f)
if "last_check" in data:
self._last_check = datetime.fromisoformat(data["last_check"])
# Don't load dismissed_version - dismiss is session-only
except (json.JSONDecodeError, ValueError, OSError) as e:
logger.debug(f"Could not load version check cache: {e}")
def _save_cache(self) -> None:
"""Save check state to cache."""
if not self.cache_file:
return
try:
self.cache_file.parent.mkdir(parents=True, exist_ok=True)
data = {}
if self._last_check:
data["last_check"] = self._last_check.isoformat()
# Don't save dismissed_version - dismiss is session-only
with open(self.cache_file, "w") as f:
json.dump(data, f)
except OSError as e:
logger.debug(f"Could not save version check cache: {e}")
def should_check(self) -> bool:
"""Determine if enough time has passed to check again."""
if self._last_check is None:
return True
elapsed = datetime.now(UTC) - self._last_check
return elapsed >= self.check_interval
def check_now(
self, force: bool = False, include_prerelease: bool = False
) -> UpdateCheckResult:
"""
Check for updates immediately.
Args:
force: Check even if recently checked
include_prerelease: Include pre-release versions
Returns:
UpdateCheckResult
"""
if not force and not self.should_check() and self._last_result:
return self._last_result
result = check_for_updates(include_prerelease=include_prerelease)
self._last_check = datetime.now(UTC)
self._last_result = result
self._save_cache()
# Notify if update available and not dismissed
if (
result.update_available
and self.on_update_available
and result.latest_version != self._dismissed_version
):
self.on_update_available(result)
return result
def check_async(
self, force: bool = False, include_prerelease: bool = False
) -> None:
"""
Check for updates in background thread.
Args:
force: Check even if recently checked
include_prerelease: Include pre-release versions
"""
if self._check_thread and self._check_thread.is_alive():
return # Already checking
if not force and not self.should_check():
return # Too soon to check again
def _check():
try:
self.check_now(force=True, include_prerelease=include_prerelease)
except Exception as e:
logger.debug(f"Background version check failed: {e}")
self._check_thread = threading.Thread(target=_check, daemon=True)
self._check_thread.start()
def dismiss_update(self, version: str) -> None:
"""
Dismiss update notification for a specific version.
Only lasts for the current session - will nag again on next startup.
Args:
version: Version to dismiss (e.g., "v1.2.3")
"""
# Session-only dismiss - don't save to cache
self._dismissed_version = version
def clear_dismissed(self) -> None:
"""Clear the dismissed version so user will be nagged again."""
self._dismissed_version = None
def get_last_result(self) -> Optional[UpdateCheckResult]:
"""Get the result of the last check."""
return self._last_result
# Global instance for easy access
_global_checker: Optional[VersionChecker] = None
def get_version_checker(
cache_dir: Optional[Path] = None,
on_update_available: Optional[Callable[[UpdateCheckResult], None]] = None,
) -> VersionChecker:
"""
Get or create the global version checker instance.
Args:
cache_dir: Directory for caching (only used on first call)
on_update_available: Callback for updates (only used on first call)
Returns:
The global VersionChecker instance
"""
global _global_checker
if _global_checker is None:
_global_checker = VersionChecker(
cache_dir=cache_dir,
on_update_available=on_update_available,
)
return _global_checker
def format_update_message(result: UpdateCheckResult, short: bool = False) -> str:
"""
Format a human-readable update message.
Args:
result: The update check result
short: Whether to use a short format
Returns:
Formatted message string
"""
if not result.update_available:
return f"✅ Loxide is up to date (v{result.current_version})"
if short:
return f"🆕 Update available: {result.latest_version}"
msg = f"🆕 Loxide {result.latest_version} is available! (current: v{result.current_version})"
if result.release_info:
msg += f"\n📥 Download: {result.release_info.html_url}"
# Include release notes preview if available
if result.release_info.body:
notes = result.release_info.body.strip()
# Truncate if too long
if len(notes) > 200:
notes = notes[:200] + "..."
msg += f"\n\n📋 Release Notes:\n{notes}"
return msg
# ---------------------------------------------------------------------------
# Textual TUI Integration
# ---------------------------------------------------------------------------
def create_update_notifier(
app, cache_dir: Optional[Path] = None, nag_on_startup: bool = True
):
"""
Create a version checker that notifies via Textual toast notifications.
This should be called after the Textual app is created.
Args:
app: The Textual App instance
cache_dir: Directory for caching check state
nag_on_startup: Always show notification on startup if update available
Returns:
The VersionChecker instance
"""
def on_update_available(result: UpdateCheckResult):
"""Callback when update is available - show toast notification."""
try:
msg = f"🆕 Update available: {result.latest_version}\nGo to Settings to download"
try:
app.notify(
msg, title="Loxide Update Available", severity="warning", timeout=15
)
except RuntimeError:
app.call_from_thread(
app.notify,
msg,
title="Loxide Update Available",
severity="warning",
timeout=15,
)
except Exception as e:
logger.debug(f"Could not show update notification: {e}")
checker = get_version_checker(
cache_dir=cache_dir,
on_update_available=on_update_available,
)
# Store checker on app so Loxide.on_mount can use it
if nag_on_startup:
app._version_checker = checker
app._version_nag_shown = False
return checker
def check_for_updates_startup(
cache_dir: Optional[Path] = None,
) -> Optional[UpdateCheckResult]:
"""
Check for updates during application startup.
This performs a synchronous check but respects the cache interval,
so it will only actually query the network once per CHECK_INTERVAL_HOURS.
Returns the result if an update is available, None otherwise.
Example usage:
result = check_for_updates_startup(cache_dir)
if result and result.update_available:
print(format_update_message(result))
"""
checker = get_version_checker(cache_dir=cache_dir)
# Only check if enough time has passed (uses cache)
if not checker.should_check():
result = checker.get_last_result()
if result and result.update_available:
return result
return None
result = checker.check_now(force=False)
if result.update_available:
return result
return None