Files
AirlockTools/TUI/Widgets/policytreewidget.py
Zarithas fc17c869fc feat: Multiple UI improvements and new server log functionality
- Add Server Log tab with DataTable display of server activity logs

- Fix keyboard navigation bug in agents tab

- Add execution history viewer for selected agents

- Improve policy tree widget functionality by adding single device operations

- Integrate logging notifications into TUI
  - Add TextualNotificationHandler to setup.py
  - Display ERROR/WARNING/CRITICAL logs as toast notifications
  - Remove terminal output to prevent interference with TUI
  - Logs still written to Loxide.log file

closes #45
2025-12-16 16:16:49 -05:00

389 lines
15 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/>.
from collections import defaultdict
import logging
from rich.text import Text
from textual.containers import Horizontal, Vertical
from textual.message import Message
from textual.widget import Widget
from textual.widgets import Button, Input, OptionList, Static, Switch, Tree
from textual.widgets.option_list import Option
logger = logging.getLogger(__name__)
class PolicyTreeWidget(Widget):
"""Widget for displaying and searching a hierarchical policy tree."""
class ViewExecutionHistory(Message):
"""Message sent when user wants to view execution history for a device."""
def __init__(self, device):
super().__init__()
self.device = device
class GenerateOTP(Message):
"""Message sent when user wants to generate OTP for a device."""
def __init__(self, device):
super().__init__()
self.device = device
class ToggleEnforcement(Message):
"""Message sent when user wants to toggle audit/enforcement for a device."""
def __init__(self, device):
super().__init__()
self.device = device
def __init__(self, policies, devices):
super().__init__()
self.policies = policies
self.devices = devices
self.last_highlighted_node = None
self.leaf_counts = defaultdict(int)
self.match_type = "Count" # Default to sorting by count
self.selected_device = None # Track currently selected device
def compose(self):
# Create the switch and its label
switch = Switch(value=False, id="match_switch")
switch.styles.margin = (0, 0, 0, 0) # top, right, bottom, left
switch.styles.padding = (0, 0, 0, 0)
switch_label = Static("Sort: Count", id="match_switch_label")
switch_label.styles.margin = (1, 0, 0, 0)
switch_label.styles.padding = (0, 0, 0, 0)
# Create the tree
policy_tree = Tree("", id="policy_tree") # Label set in on_mount
policy_tree.styles.width = "2fr"
policy_tree.styles.height = "100%"
# Create the search box and details pane
label = Static("Device Search:")
search_box = Input(
placeholder="Search policies or devices...", id="tree_search"
)
exec_history_button = Button(
"📊 Execution History", id="view_exec_history_button", disabled=True
)
exec_history_button.styles.margin = (0, 1, 0, 0) # Right margin
otp_button = Button("🎫 Generate OTP", id="generate_otp_button", disabled=True)
otp_button.styles.margin = (0, 1, 0, 0) # Right margin
toggle_enforcement_button = Button(
"🔄 Toggle Enforcement/Audit", id="toggle_enforcement_button", disabled=True
)
# No right margin on last button
details_pane = Static("", id="details_pane")
# Layout the UI
with Horizontal():
yield policy_tree
with Vertical() as right_pane:
right_pane.styles.width = "3fr"
# Use a Horizontal container for the switch and label
with Horizontal() as switch_container:
switch_container.styles.height = 3
switch_container.styles.margin = (0, 0, 0, 1)
switch_container.styles.padding = (0, 0, 0, 0)
yield switch
yield switch_label
# Add the search box and details pane
yield label
yield search_box
# Action buttons in a horizontal row
with Horizontal() as button_row:
button_row.styles.height = "auto"
yield exec_history_button
yield otp_button
yield toggle_enforcement_button
yield details_pane
def on_mount(self) -> None:
self._precompute_leaf_counts()
# Update root label with total leaf count
total_leaves = sum(
self.leaf_counts.get(policy.groupid, 0)
for policy in self.policies
if policy.parent == "global-policy-settings"
)
policy_tree = self.query_one("#policy_tree", Tree)
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
self._build_tree()
# Expand the root node
policy_tree.root.expand()
def refresh_data(self, policies, devices):
"""Refresh the widget with new data and rebuild the tree."""
self.policies = policies
self.devices = devices
self.selected_device = None
# Disable all buttons since selection is lost
try:
self.query_one("#view_exec_history_button", Button).disabled = True
self.query_one("#generate_otp_button", Button).disabled = True
self.query_one("#toggle_enforcement_button", Button).disabled = True
except:
pass
# Rebuild tree with new data
self._precompute_leaf_counts()
total_leaves = sum(
self.leaf_counts.get(policy.groupid, 0)
for policy in self.policies
if policy.parent == "global-policy-settings"
)
policy_tree = self.query_one("#policy_tree", Tree)
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
self._build_tree()
policy_tree.root.expand()
def _precompute_leaf_counts(self):
"""Precompute leaf counts for each policy group."""
device_counts = defaultdict(int)
for device in self.devices:
device_counts[device.groupid] += 1
child_map = defaultdict(list)
for policy in self.policies:
child_map[policy.parent].append(policy.groupid)
def count_leaves(groupid):
count = device_counts[groupid]
for child_id in child_map.get(groupid, []):
count += count_leaves(child_id)
self.leaf_counts[groupid] = count
return count
for policy in self.policies:
if policy.parent == "global-policy-settings":
count_leaves(policy.groupid)
def _build_tree(self):
policy_tree = self.query_one("#policy_tree", Tree)
policy_tree.clear() # Clear existing nodes
node_map = {}
# Sort top-level policies
top_policies = [
p for p in self.policies if p.parent == "global-policy-settings"
]
# Sort by count (default) or alphabetically
if getattr(self, "match_type", "Count") == "Count":
top_policies.sort(
key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
)
else: # Alphabetical
top_policies.sort(key=lambda p: p.name.lower())
for policy in top_policies:
label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})"
node = policy_tree.root.add(label=label, data=policy)
node_map[policy.groupid] = node
# Sort and add child policies
children_by_parent = defaultdict(list)
for policy in self.policies:
if policy.parent != "global-policy-settings":
children_by_parent[policy.parent].append(policy)
for parent_id, children in children_by_parent.items():
if getattr(self, "match_type", "Count") == "Count":
children.sort(
key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
)
else: # Alphabetical
children.sort(key=lambda p: p.name.lower())
parent_node = node_map.get(parent_id)
if parent_node:
for policy in children:
label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})"
node = parent_node.add(label=label, data=policy)
node_map[policy.groupid] = node
# Add devices (leaf nodes) - always sort alphabetically
devices_by_group = defaultdict(list)
for device in self.devices:
devices_by_group[device.groupid].append(device)
for group_id, devices in devices_by_group.items():
devices.sort(key=lambda d: d.hostname.lower()) # Always sort alphabetically
parent_node = node_map.get(group_id)
if parent_node:
for device in devices:
parent_node.add(label=device.hostname, data=device)
def _collect_tree_nodes(self, node, all_nodes):
all_nodes.append(node)
for child in node.children:
self._collect_tree_nodes(child, all_nodes)
def _remove_match_selector(self):
try:
existing = self.query("#match_selector")
for widget in existing:
if widget.is_attached:
widget.remove()
except Exception as exc:
logger.debug("Failed to remove match_selector: %s", exc)
def on_tree_node_selected(self, message: Tree.NodeSelected) -> None:
node = message.node
data = node.data
details_pane = self.query_one("#details_pane", Static)
exec_history_button = self.query_one("#view_exec_history_button", Button)
otp_button = self.query_one("#generate_otp_button", Button)
toggle_enforcement_button = self.query_one("#toggle_enforcement_button", Button)
if self.last_highlighted_node is not None:
original_label = str(self.last_highlighted_node.label).strip()
if isinstance(self.last_highlighted_node.label, Text):
original_label = self.last_highlighted_node.label.plain
self.last_highlighted_node.set_label(original_label)
label_text = str(node.label).strip()
if isinstance(node.label, Text):
label_text = node.label.plain
highlighted_label = Text(label_text, style="reverse bold")
node.set_label(highlighted_label)
self.last_highlighted_node = node
# Check if selected node is a device (has Agent data)
from models.agent import Agent
if data and isinstance(data, Agent):
self.selected_device = data
exec_history_button.disabled = False
otp_button.disabled = False
toggle_enforcement_button.disabled = False
else:
self.selected_device = None
exec_history_button.disabled = True
otp_button.disabled = True
toggle_enforcement_button.disabled = True
if data:
details = "\n".join(
f"{key}: {value}" for key, value in data.__dict__.items()
)
else:
details = f"Selected: {node.label}"
details_pane.update(details)
message.stop()
def on_switch_changed(self, event: Switch.Changed):
self.match_type = "Alpha" if event.value else "Count"
self.query_one("#match_switch_label", Static).update(
f"Sort: {self.match_type.capitalize()}"
)
self._build_tree()
def on_input_submitted(self, message: Input.Submitted) -> None:
self._remove_match_selector()
query = message.value.strip().lower()
tree = self.query_one("#policy_tree", Tree)
details_pane = self.query_one("#details_pane", Static)
all_nodes = []
self._collect_tree_nodes(tree.root, all_nodes)
label_to_node = {}
for node in all_nodes:
label_text = str(node.label).lower()
label_to_node[label_text] = node
if node.data:
data_dict = (
node.data.__dict__ if hasattr(node.data, "__dict__") else node.data
)
for key, value in data_dict.items():
if isinstance(value, str):
label_to_node[value.lower()] = node
matches = sorted([label for label in label_to_node if query in label])
if matches:
try:
option_list = self.query_one("#match_selector", OptionList)
option_list.clear_options()
option_list.display = True
except:
option_list = OptionList(id="match_selector")
details_pane.parent.mount(option_list)
for label in matches:
option_list.add_option(Option(label, id=f"match_{label}"))
details_pane.update(f"Found {len(matches)} matches. Select one below.")
else:
self._remove_match_selector()
details_pane.update("No matches found.")
def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None:
selected_id = event.option.id.replace("match_", "")
tree = self.query_one("#policy_tree", Tree)
details_pane = self.query_one("#details_pane", Static)
all_nodes = []
self._collect_tree_nodes(tree.root, all_nodes)
label_to_node = {str(node.label).lower(): node for node in all_nodes}
match_node = label_to_node.get(selected_id.lower())
if match_node:
node = match_node
path = []
while node:
path.insert(0, node)
node = node.parent
for node in path:
node.expand()
tree.select_node(match_node)
tree.scroll_to_node(match_node)
match_node.set_label(Text(str(match_node.label), style="reverse bold"))
details_pane.update(f"Selected: {match_node.label}")
try:
option_list = self.query_one("#match_selector", OptionList)
option_list.remove()
except:
pass
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button presses."""
if event.button.id == "view_exec_history_button":
if self.selected_device:
self.post_message(self.ViewExecutionHistory(self.selected_device))
event.stop()
elif event.button.id == "generate_otp_button":
if self.selected_device:
self.post_message(self.GenerateOTP(self.selected_device))
event.stop()
elif event.button.id == "toggle_enforcement_button":
if self.selected_device:
self.post_message(self.ToggleEnforcement(self.selected_device))
event.stop()