980 lines
35 KiB
Python
980 lines
35 KiB
Python
# 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/>.
|
||
|
||
|
||
# TODO Continue implementing logger
|
||
# TODO Add input sanitation and CSV injection prevention
|
||
# TODO Continue OTP and Local approval rewrites
|
||
# TODO Explore pywin32
|
||
# 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 Horizontal, Vertical
|
||
from textual.message import Message
|
||
from textual.reactive import reactive
|
||
from textual.screen import Screen
|
||
from textual.widgets import (
|
||
Button,
|
||
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
|
||
from services.security import getAPI
|
||
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
|
||
from TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
||
from TUI.Screens.otpactivityscreen import OTPActivitiesScreen
|
||
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.policytreewidget import PolicyTreeWidget
|
||
from TUI.Widgets.resultsdisplay import ResultsDisplay
|
||
from TUI.Widgets.serverlogwidget import ServerLogWidget
|
||
from utils.configmanager import (
|
||
get_system_value,
|
||
get_user_value,
|
||
load_env,
|
||
save_user_config,
|
||
)
|
||
from utils.setup import get_base_directory, setup
|
||
from utils.utils import irtang, open_directory
|
||
|
||
dotenv.load_dotenv()
|
||
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GLOBAL STASH
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_APP_RESTART_REASON = None
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
|
||
# ---------------------------------------------------------------------------
|
||
def _persist_user_theme(theme_name: str) -> None:
|
||
"""
|
||
Store the chosen Textual theme in the user's config using the config manager.
|
||
No need to touch .env - config manager handles everything.
|
||
"""
|
||
base_dir = get_base_directory()
|
||
config_dir = base_dir / "config"
|
||
|
||
try:
|
||
save_user_config(config_dir, {"TEXTUAL_THEME": theme_name})
|
||
logger.debug("Updated user config with TEXTUAL_THEME=%s", theme_name)
|
||
except Exception as exc:
|
||
logger.error("Failed to save TEXTUAL_THEME: %s", exc)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1) SCREEN
|
||
# ---------------------------------------------------------------------------
|
||
class MainMenuScreen(Screen):
|
||
api: AirlockAPIWrapper
|
||
current_tab = reactive("")
|
||
|
||
BUTTON_DEFS = {
|
||
"agent_actions": [
|
||
{
|
||
"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",
|
||
"id": "otp_activities_button",
|
||
},
|
||
{
|
||
"label": "🛑 - Revoke Active OTP Session",
|
||
"id": "otp_revoke_button",
|
||
},
|
||
],
|
||
"policy": [
|
||
{
|
||
"label": "âš–ï¸ - Prepare Policy For Enforcement",
|
||
"id": "policy_prep_button",
|
||
},
|
||
{
|
||
"label": "🔕 - Find and Move Quiet Hosts to Enforcement",
|
||
"id": "find_quiet_button",
|
||
},
|
||
],
|
||
}
|
||
|
||
def __init__(self) -> None:
|
||
super().__init__()
|
||
self.extras = get_user_value("EXTRAS", str, "NOTTODAY")
|
||
wd = load_env("WORKING_DIR") or os.getcwd()
|
||
if not os.path.isdir(wd):
|
||
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 = []
|
||
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)
|
||
|
||
# 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.margin = (0, 0, 1, 0)
|
||
widgets.append(desc_text)
|
||
|
||
return Vertical(*widgets)
|
||
|
||
def compose(self) -> ComposeResult:
|
||
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"),
|
||
]
|
||
|
||
if self.extras == "POLICYPREP":
|
||
tabs.insert(2, Tab("Policy Prep", id="policy"))
|
||
|
||
yield Tabs(*tabs, id="tabs")
|
||
yield Vertical(id="content")
|
||
yield Footer()
|
||
|
||
def on_mount(self) -> None:
|
||
self.switch_tab("agent_actions")
|
||
|
||
def on_key(self, event) -> None:
|
||
"""Handle up/down arrow keys for button navigation."""
|
||
if event.key == "down":
|
||
self._focus_nearby_button(1)
|
||
event.prevent_default()
|
||
event.stop()
|
||
elif event.key == "up":
|
||
self._focus_nearby_button(-1)
|
||
event.prevent_default()
|
||
event.stop()
|
||
# left/right are handled by Textual's default tab navigation
|
||
|
||
# focus helpers
|
||
def _get_content_buttons(self) -> list[Button]:
|
||
content = self.query_one("#content", Vertical)
|
||
return list(content.query(Button))
|
||
|
||
def _focus_first_button(self) -> None:
|
||
buttons = self._get_content_buttons()
|
||
if buttons:
|
||
buttons[0].focus()
|
||
|
||
def _focus_tabs(self) -> None:
|
||
tabs = self.query_one("#tabs", Tabs)
|
||
tabs.focus()
|
||
|
||
def _focus_nearby_button(self, direction: int) -> None:
|
||
buttons = self._get_content_buttons()
|
||
if not buttons:
|
||
return
|
||
|
||
try:
|
||
current = next(i for i, b in enumerate(buttons) if b.has_focus)
|
||
except StopIteration:
|
||
if direction > 0:
|
||
buttons[0].focus()
|
||
else:
|
||
buttons[-1].focus()
|
||
return
|
||
|
||
if direction < 0 and current == 0:
|
||
self._focus_tabs()
|
||
return
|
||
|
||
new_index = current + direction
|
||
if 0 <= new_index < len(buttons):
|
||
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()
|
||
|
||
if tab_id in self.BUTTON_DEFS:
|
||
content.mount(self._make_buttons_for(tab_id))
|
||
elif tab_id == "server_log":
|
||
content.mount(ServerLogWidget(self.app.api))
|
||
elif tab_id == "dir":
|
||
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_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
|
||
logger.info("Selected agents: %s", selected_agents)
|
||
# TODO: Implement actual handling of selected agents
|
||
_APP_RESTART_REASON = ("multi_agent_action", selected_agents)
|
||
self.app.exit()
|
||
|
||
def on_theme_selector_theme_selected(
|
||
self, message: ThemeSelector.ThemeSelected
|
||
) -> None:
|
||
"""Handle theme selection from ThemeSelector."""
|
||
global _APP_RESTART_REASON
|
||
_persist_user_theme(message.theme_name)
|
||
_APP_RESTART_REASON = ("restart",)
|
||
self.app.exit()
|
||
|
||
def on_agent_move_operations_operation_complete(
|
||
self, message: AgentMoveOperations.OperationComplete
|
||
) -> None:
|
||
"""Handle completion of agent move operation - show results."""
|
||
logger.info(
|
||
"Agent move operation completed: %s, %d successful, %d unsuccessful",
|
||
message.operation,
|
||
len(message.successful),
|
||
len(message.unsuccessful),
|
||
)
|
||
|
||
# Format results for display
|
||
successful_text = "\n".join(
|
||
[f"{agent.hostname}" for agent, _ in message.successful]
|
||
)
|
||
unsuccessful_text = "\n".join(
|
||
[f"{agent.hostname}: {error}" for agent, error in message.unsuccessful]
|
||
)
|
||
|
||
# Remove the operations widget
|
||
try:
|
||
ops_widget = self.query_one(AgentMoveOperations)
|
||
ops_widget.remove()
|
||
except Exception:
|
||
pass
|
||
|
||
# Show results
|
||
self.query_one("#content", Vertical).mount(
|
||
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
|
||
)
|
||
|
||
def on_results_display_go_back(self, message: ResultsDisplay.GoBack) -> None:
|
||
"""Handle back button from results display."""
|
||
try:
|
||
results_widget = self.query_one(ResultsDisplay)
|
||
results_widget.remove()
|
||
except Exception:
|
||
pass
|
||
# Return to main menu
|
||
self.app.pop_screen()
|
||
|
||
def on_policy_tree_widget_view_execution_history(
|
||
self, message: PolicyTreeWidget.ViewExecutionHistory
|
||
) -> None:
|
||
"""Handle request to view execution history for a device from tree view."""
|
||
logger.info("Viewing execution history for device: %s", message.device.hostname)
|
||
self.app.push_screen(ExecutionHistoryScreen([message.device]))
|
||
message.stop()
|
||
|
||
def on_policy_tree_widget_generate_otp(
|
||
self, message: PolicyTreeWidget.GenerateOTP
|
||
) -> None:
|
||
"""Handle request to generate OTP for a device from tree view."""
|
||
logger.info("Generating OTP for device: %s", message.device.hostname)
|
||
self.app.push_screen(OTPWorkflowScreen([message.device]))
|
||
message.stop()
|
||
|
||
def on_policy_tree_widget_toggle_enforcement(
|
||
self, message: PolicyTreeWidget.ToggleEnforcement
|
||
) -> None:
|
||
"""Handle request to toggle enforcement for a device from tree view."""
|
||
logger.info("Toggling enforcement for device: %s", message.device.hostname)
|
||
|
||
try:
|
||
from services.agenthandler import moveAgentToRelatedPolicy
|
||
from utils.configmanager import get_system_json
|
||
|
||
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||
|
||
# Determine current mode and toggle
|
||
if message.device.groupid in policy_relationship_map:
|
||
# Currently in enforcement, move to audit
|
||
result = moveAgentToRelatedPolicy(self.app.api, message.device, "audit")
|
||
mode = "audit"
|
||
else:
|
||
# Currently in audit, move to enforcement
|
||
result = moveAgentToRelatedPolicy(
|
||
self.app.api, message.device, "enforcement"
|
||
)
|
||
mode = "enforcement"
|
||
|
||
logger.info(f"Successfully toggled {message.device.hostname} to {mode}")
|
||
|
||
# Refresh data at the app level
|
||
self.app.refresh_data()
|
||
|
||
# Refresh the tree widget with new data
|
||
try:
|
||
tree_widget = self.query_one(PolicyTreeWidget)
|
||
tree_widget.refresh_data(self.app.policies, self.app.devices)
|
||
except:
|
||
pass
|
||
|
||
except Exception as e:
|
||
logger.error(
|
||
f"Failed to toggle enforcement for {message.device.hostname}: {e}"
|
||
)
|
||
self.app.bell()
|
||
|
||
message.stop()
|
||
|
||
def on_directory_tree_file_selected(
|
||
self, event: DirectoryTree.FileSelected
|
||
) -> None:
|
||
path = event.path
|
||
logger.debug("Directory file selected: %s", path)
|
||
try:
|
||
open_directory(str(path))
|
||
except Exception as exc:
|
||
logger.error("Failed to open %s: %s", path, exc)
|
||
self.app.bell()
|
||
|
||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||
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))
|
||
event.stop()
|
||
|
||
case "otp_generate_button":
|
||
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
|
||
event.stop()
|
||
|
||
case "find_quiet_button":
|
||
self.app.push_screen(
|
||
QuietAgentWorkflowScreen(self.app.api, self.app.policies)
|
||
)
|
||
event.stop()
|
||
return
|
||
|
||
case "otp_activities_button":
|
||
self.app.push_screen(OTPActivitiesScreen())
|
||
event.stop()
|
||
return
|
||
|
||
case "otp_revoke_button":
|
||
self.app.push_screen(OTPRevokeScreen())
|
||
event.stop()
|
||
return
|
||
|
||
case "policy_prep_button":
|
||
# Use the new TUI workflow screen instead of legacy
|
||
self.app.push_screen(
|
||
PolicyPrepWorkflowScreen(self.app.api, self.app.policies)
|
||
)
|
||
event.stop()
|
||
return
|
||
|
||
case _:
|
||
self.app.bell()
|
||
logger.warning("Unknown button pressed: %s", button_id)
|
||
return
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2) APP
|
||
# ---------------------------------------------------------------------------
|
||
class Loxide(App[Message]):
|
||
api: AirlockAPIWrapper
|
||
working_dir: str
|
||
policies: Optional[list[Policy]]
|
||
devices: Optional[list[Agent]]
|
||
|
||
CSS = """
|
||
#logo {
|
||
width: 100%;
|
||
content-align: center middle;
|
||
text-align: center;
|
||
}
|
||
"""
|
||
BINDINGS = [
|
||
("q", "quit", "Quit"),
|
||
("f", "open_fe", "Launch Explorer"),
|
||
("r", "refresh", "Refresh"),
|
||
]
|
||
|
||
def __init__(self, api: AirlockAPIWrapper):
|
||
self._textual_theme = get_user_value("TEXTUAL_THEME", str, "textual-dark")
|
||
super().__init__()
|
||
self.api = api
|
||
wd = load_env("WORKING_DIR") or os.getcwd()
|
||
if not os.path.isdir(wd):
|
||
wd = os.getcwd()
|
||
self.working_dir = wd
|
||
# Initial data load
|
||
self.refresh_data()
|
||
|
||
def refresh_data(self) -> None:
|
||
"""Public method to refresh policies and devices from the API."""
|
||
try:
|
||
self.policies = [
|
||
Policy(**row.to_dict())
|
||
for _, row in self.api.policy_find_all().iterrows()
|
||
]
|
||
self.devices = [
|
||
Agent(**row.to_dict())
|
||
for _, row in self.api.agent_find_all().iterrows()
|
||
]
|
||
if self.policies and self.devices:
|
||
for agent in self.devices:
|
||
agent.enrich_with_policies(self.policies)
|
||
logger.debug(
|
||
f"Enriched {len(self.devices)} agents with policy information"
|
||
)
|
||
except Exception as exc:
|
||
logger.error("Failed to load policies/devices: %s", exc)
|
||
self.policies = None
|
||
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())
|
||
|
||
def action_refresh(self) -> None:
|
||
self.refresh_data()
|
||
|
||
def action_quit(self) -> None:
|
||
global _APP_RESTART_REASON
|
||
_APP_RESTART_REASON = None
|
||
self.exit()
|
||
|
||
def action_open_fe(self) -> None:
|
||
"""Open the working directory in the OS file manager (footer binding)."""
|
||
path_to_open = self.working_dir or os.getcwd()
|
||
try:
|
||
open_directory(path_to_open)
|
||
except Exception as exc:
|
||
logger.error("Failed to open directory %s: %s", path_to_open, exc)
|
||
self.bell() # optional feedback
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3) PUBLIC ENTRYPOINT - Updated to accept attach_notification_handler
|
||
# ---------------------------------------------------------------------------
|
||
def run_Loxide(api: AirlockAPIWrapper, attach_notification_handler=None) -> None:
|
||
global _APP_RESTART_REASON
|
||
base_dir = get_base_directory()
|
||
env_path = base_dir / ".env"
|
||
dotenv.load_dotenv(dotenv_path=env_path, override=True)
|
||
|
||
max_attempts = 5
|
||
attempts = 0
|
||
|
||
while attempts < max_attempts:
|
||
attempts += 1
|
||
logger.debug("Starting app loop iteration (attempt %d)", attempts)
|
||
_APP_RESTART_REASON = None
|
||
app = Loxide(api)
|
||
|
||
# Attach the notification handler if provided
|
||
if attach_notification_handler:
|
||
attach_notification_handler(app)
|
||
|
||
try:
|
||
app.run()
|
||
except SystemExit as exc:
|
||
if exc.code != 0:
|
||
logger.debug("Caught SystemExit from Textual: %s", exc)
|
||
raise
|
||
|
||
reason = _APP_RESTART_REASON
|
||
logger.debug("After app.run(), _APP_RESTART_REASON = %r", reason)
|
||
|
||
if not reason:
|
||
logger.debug("No restart reason, exiting loop")
|
||
break
|
||
|
||
if reason[0] == "restart":
|
||
logger.debug("Restarting app loop")
|
||
continue
|
||
|
||
if reason[0] == "multi_agent_action":
|
||
logger.info("Multi-agent action with selected agents: %s", reason[1])
|
||
continue
|
||
|
||
logger.error("Unknown restart reason: %r", reason)
|
||
break
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 4) MAIN FUNCTION - Updated to get and pass attach_notification_handler
|
||
# ---------------------------------------------------------------------------
|
||
def main():
|
||
irtang()
|
||
# Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
|
||
# setup() now returns a function to attach the notification handler
|
||
attach_notification_handler = setup()
|
||
logger = logging.getLogger(__name__)
|
||
|
||
try:
|
||
url = get_system_value("URL")
|
||
username = os.getenv("USERNAME")
|
||
|
||
if not url:
|
||
raise ValueError("Missing URL in environment variables.")
|
||
if not username:
|
||
raise ValueError("Missing USERNAME in environment variables.")
|
||
|
||
logger.debug(f"Retrieved URL: {url}")
|
||
logger.debug(f"Retrieved Username: {username}")
|
||
|
||
except ValueError as e:
|
||
logger.error(f"Configuration error: {e}", exc_info=True)
|
||
raise
|
||
|
||
api_key = getAPI(username, "Loxide")
|
||
if api_key is None:
|
||
raise ValueError("API key for Loxide is missing.")
|
||
|
||
api = AirlockAPIWrapper(
|
||
base_url=str(url),
|
||
api_key=api_key,
|
||
)
|
||
run_Loxide(api, attach_notification_handler)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|