Major step towards unification of the UI Implementation of the Back Feature, splitting of TUI files back into subfolders
This commit is contained in:
@@ -0,0 +1,289 @@
|
||||
# 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.widget import Widget
|
||||
from textual.widgets import 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."""
|
||||
|
||||
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
|
||||
|
||||
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"
|
||||
)
|
||||
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
|
||||
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 _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)
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user