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

Initial commit for statistics feature

Browse Source
This commit is contained in:
Zarithas
2025-12-17 20:17:32 -05:00
parent a7b659c951
commit fbd5b8b4b8
4 changed files with 446 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Loxide.py
+384 -11
View File
@@ -21,13 +21,17 @@
# TODO Fix Requirements.txt
# TODO Create Generic system_config.json for gitea
from collections import Counter
from datetime import UTC, datetime, timedelta
import json
import logging
import os
from typing import Optional
from bson import ObjectId
import dotenv
from textual.app import App, ComposeResult
from textual.containers import Vertical
from textual.containers import Horizontal, Vertical
from textual.message import Message
from textual.reactive import reactive
from textual.screen import Screen
@@ -36,12 +40,14 @@ from textual.widgets import (
DirectoryTree,
Footer,
Header,
Select,
Static,
Tab,
Tabs,
)
import urllib3
import airlock_libs
from models.agent import Agent
from models.policy import Policy
from services.API import AirlockAPIWrapper
@@ -57,7 +63,6 @@ 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
@@ -110,26 +115,26 @@ class MainMenuScreen(Screen):
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",
},
{
"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",
},
],
@@ -143,6 +148,10 @@ class MainMenuScreen(Screen):
wd = os.getcwd()
self.working_dir = wd
# Statistics cache: {days: {"data": stats_dict, "timestamp": datetime}}
self.statistics_cache = {}
self.current_stats_days = 1 # Default to 1 day
def _make_buttons_for(self, tab_id: str) -> Vertical:
defs = self.BUTTON_DEFS.get(tab_id, [])
widgets = []
@@ -173,12 +182,13 @@ class MainMenuScreen(Screen):
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"),
Tab("Tree View", id="p_tree"),
Tab("Server Log", id="server_log"),
Tab("Statistics", id="statistics"),
Tab("Directory", id="dir"),
Tab("Settings", id="settings"),
]
@@ -254,17 +264,374 @@ class MainMenuScreen(Screen):
content.mount(DirectoryTree(self.working_dir, id="dir_tree"))
elif tab_id == "p_tree":
content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
elif tab_id == "statistics":
content.mount(self._create_statistics_widget())
elif tab_id == "settings":
content.mount(ThemeSelector())
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 = []
# Header
header = Static("📊 System Statistics", classes="stats_header")
header.styles.text_align = "center"
header.styles.text_style = "bold"
header.styles.margin = (1, 0, 1, 0)
widgets_to_mount.append(header)
# Controls row
controls = Horizontal()
controls.styles.height = "auto"
controls.styles.margin = (0, 2, 1, 2)
# Time period selector
time_label = Static("Time Period: ")
time_label.styles.width = "auto"
time_label.styles.margin = (0, 1, 0, 0)
time_select = Select(
options=[
("1 Day", 1),
("3 Days", 3),
("7 Days", 7),
],
value=self.current_stats_days,
id="stats_time_select",
)
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)
# Compose controls
controls.compose_add_child(time_label)
controls.compose_add_child(time_select)
controls.compose_add_child(refresh_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)
widgets_to_mount.append(status_container)
# Data display area - using horizontal layout
data_container = Horizontal(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:
self.call_after_refresh(self._update_statistics_display)
return stats_container
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)
timestamp = int(date_days_ago.timestamp())
hex_timestamp = format(timestamp, "08x")
objectid_hex = hex_timestamp + "0000000000000000"
return ObjectId(objectid_hex)
def _fetch_execution_statistics(self, days: int) -> None:
"""Fetch execution history and calculate statistics using airlock_libs."""
try:
status_area = self.query_one("#stats_status_area", Vertical)
except Exception as e:
logger.warning(f"Could not find status area: {e}")
return
status_area.remove_children()
status = Static("⏳ Fetching execution data...")
status.styles.text_style = "italic"
status_area.mount(status)
try:
# Non-trusted execution types
exec_types = [1, 2, 3, 6, 7, 12, 13, 14, 15, 16]
logger.info(
f"Fetching {days} days of execution history for all policies..."
)
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
execs = airlock_libs.pull_policy_exec_histories(
self.app.api, None, str(exec_types), days # Empty string = all policies
)
logger.debug(f"Received response, length: {len(execs) if execs else 0}")
if not execs:
logger.warning("No execution data returned")
status_area.remove_children()
error = Static("No execution data found for the selected time period.")
error.styles.color = "yellow"
status_area.mount(error)
return
# Parse JSON response
data = json.loads(execs)
logger.debug(f"Parsed JSON, keys: {data.keys()}")
all_executions = data.get("response", {}).get("exechistories", [])
logger.debug(f"Found {len(all_executions)} executions in response")
if not all_executions:
logger.warning("No executions in response")
status_area.remove_children()
error = Static("No execution records found.")
error.styles.color = "yellow"
status_area.mount(error)
return
logger.info(f"Retrieved {len(all_executions)} execution records")
# Process statistics
stats = self._process_execution_data(all_executions)
# Cache the results
self.statistics_cache[days] = {
"data": stats,
"timestamp": datetime.now(UTC),
}
# Update display
self._update_statistics_display()
except Exception as e:
logger.error(f"Failed to fetch execution statistics: {e}", exc_info=True)
status_area.remove_children()
error = Static(f"❌ Error fetching data: {str(e)}")
error.styles.color = "red"
status_area.mount(error)
def _process_execution_data(self, executions: list) -> dict:
"""Process raw execution data into statistics."""
# Execution type names
type_names = {
0: "Trusted Execution",
1: "Blocked Execution",
2: "Untrusted Execution [Audit]",
3: "Untrusted Execution [OTP]",
4: "Trusted Path Execution",
5: "Trusted Publisher Execution",
6: "Blocklist Execution",
7: "Blocklist Execution [Audit]",
8: "Trusted Process Execution",
9: "Constrained Execution",
10: "Trusted Metadata Execution",
11: "Trusted Browser Execution",
12: "Blocked Browser Execution",
13: "Untrusted Browser Execution [Audit]",
14: "Untrusted Browser Execution [OTP]",
15: "Blocklist Browser Execution [Audit]",
16: "Blocklist Browser Execution",
17: "Trusted Installer Execution",
18: "Trusted Browser Metadata Execution",
}
# Count execution types
type_counter = Counter()
policy_counter = Counter()
hostname_counter = 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)
# Format execution type counts
type_counts = []
for exec_type, count in sorted(type_counter.items()):
type_name = type_names.get(exec_type, f"Unknown Type {exec_type}")
type_counts.append((type_name, count))
return {
"total_executions": len(executions),
"type_counts": type_counts,
"top_policies": top_policies,
"top_hostnames": top_hostnames,
}
def _update_statistics_display(self) -> None:
"""Update the statistics display with cached data using horizontal layout."""
try:
data_area = self.query_one("#stats_data_area", Horizontal)
status_area = self.query_one("#stats_status_area", Vertical)
except Exception as e:
logger.warning(f"Could not find statistics display areas: {e}")
return
data_area.remove_children()
status_area.remove_children()
# Get cached data
cache_entry = self.statistics_cache.get(self.current_stats_days)
if not cache_entry:
status = Static("No data available. Click Refresh to fetch.")
status.styles.text_style = "italic"
status_area.mount(status)
return
stats = cache_entry["data"]
timestamp = cache_entry["timestamp"]
# Timestamp display
timestamp_str = timestamp.strftime("%Y-%m-%d %H:%M:%S UTC")
timestamp_label = Static(f"Last updated: {timestamp_str}")
timestamp_label.styles.text_style = "dim"
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)
# 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
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)
)
online_agents = sum(
1 for agent in self.app.devices if getattr(agent, "online", False)
)
offline_agents = total_agents - online_agents
else:
enforced_agents = audit_agents = online_agents = offline_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_data = [
f"Agents: {total_agents} ({online_agents} online)",
f"Policies: {total_policies}",
f"Enforced: {enforced_agents}",
f"Audit: {audit_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)
# Mount all columns
data_area.mount(col1)
data_area.mount(col2)
data_area.mount(col3)
def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
self.switch_tab(event.tab.id)
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
def on_select_changed(self, event: Select.Changed) -> None:
"""Handle time period selection change."""
if event.select.id == "stats_time_select":
new_days = event.value
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:
self._update_statistics_display()
def on_multi_agent_selector_agents_selected() -> None:
"""Handle selected agents from AgentSelector."""
global _APP_RESTART_REASON
selected_agents = message.selected_agents
@@ -398,6 +765,12 @@ class MainMenuScreen(Screen):
button_id = event.button.id
logger.debug("Button pressed: %s", button_id)
# Handle statistics refresh button
if button_id == "stats_refresh_btn":
self._fetch_execution_statistics(self.current_stats_days)
event.stop()
return
match button_id:
case "move_agent_workflow_button":
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))