From ed07c79b742711b2d40649a1b5365e074e26896f Mon Sep 17 00:00:00 2001 From: Zarithas Date: Fri, 7 Nov 2025 15:44:11 -0500 Subject: [PATCH] Synching changes to Policytree --- widgets/OTP_generate.py | 9 ++- widgets/policytreewidget.py | 113 ++++++++++++++++++++++-------------- 2 files changed, 73 insertions(+), 49 deletions(-) diff --git a/widgets/OTP_generate.py b/widgets/OTP_generate.py index b35b39d..6aff859 100644 --- a/widgets/OTP_generate.py +++ b/widgets/OTP_generate.py @@ -232,9 +232,8 @@ class OTPGenerator(Widget): api = self.app.api output_lines = [ - "=" * 60, - "OTP GENERATION RESULTS", - "=" * 60, + "Requested OTP Codes:", + "=" * 25, ] otp_dict = {} @@ -250,9 +249,9 @@ class OTPGenerator(Widget): ) for hostname, otp_code in otp_dict.items(): - output_lines.append(f"{hostname:30} | {otp_code}") + output_lines.append(f"{hostname} | {otp_code}") - output_lines.append("=" * 60) + output_lines.append("=" * 25) result_text = "\n".join(output_lines) self._show_result(result_text) diff --git a/widgets/policytreewidget.py b/widgets/policytreewidget.py index 3a0efff..f611d5d 100644 --- a/widgets/policytreewidget.py +++ b/widgets/policytreewidget.py @@ -1,3 +1,4 @@ +from collections import defaultdict import logging from rich.text import Text @@ -17,14 +18,13 @@ class PolicyTreeWidget(Widget): self.policies = policies self.devices = devices self.last_highlighted_node = None + self.leaf_counts = defaultdict(int) def compose(self): - # Left: Policy Tree - policy_tree = Tree("Policies", id="policy_tree") + policy_tree = Tree("", id="policy_tree") # Label set in on_mount policy_tree.styles.width = "2fr" policy_tree.styles.height = "100%" - # Right: Search + Details label = Static("Device Search:") search_box = Input( placeholder="Search policies or devices...", id="tree_search" @@ -40,44 +40,87 @@ class PolicyTreeWidget(Widget): yield details_pane def on_mount(self) -> None: - """Build the tree after mounting.""" + 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() - def _build_tree(self) -> None: - """Build the policy tree structure.""" + 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) node_map = {} - # Top-level policies - for policy in self.policies: - if policy.parent == "global-policy-settings": - node = policy_tree.root.add(label=policy.name, data=policy) - node_map[policy.groupid] = node + # Sort top-level policies + top_policies = [ + p for p in self.policies if p.parent == "global-policy-settings" + ] + top_policies.sort( + key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True + ) - # Child policies - for policy in self.policies: - parent_id = policy.parent - if parent_id in node_map: - parent_node = node_map[parent_id] - node = parent_node.add(label=policy.name, data=policy) - node_map[policy.groupid] = node + 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 - # Devices under policies + # 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(): + children.sort( + key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True + ) + 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) for device in self.devices: group_id = device.groupid - if group_id in node_map: - parent_node = node_map[group_id] - label = device.hostname - parent_node.add(label=label, data=device) + parent_node = node_map.get(group_id) + if parent_node: + parent_node.add(label=device.hostname, data=device) def _collect_tree_nodes(self, node, all_nodes): - """Helper to recursively collect all nodes from a tree.""" all_nodes.append(node) for child in node.children: self._collect_tree_nodes(child, all_nodes) def _remove_match_selector(self): - """Safely remove match selector widgets.""" try: existing = self.query("#match_selector") for widget in existing: @@ -87,20 +130,16 @@ class PolicyTreeWidget(Widget): logger.debug("Failed to remove match_selector: %s", exc) def on_tree_node_selected(self, message: Tree.NodeSelected) -> None: - """Handle tree node selection.""" node = message.node data = node.data details_pane = self.query_one("#details_pane", Static) - # Reset previous highlight if self.last_highlighted_node is not None: original_label = str(self.last_highlighted_node.label).strip() - # Remove any styling if isinstance(self.last_highlighted_node.label, Text): original_label = self.last_highlighted_node.label.plain self.last_highlighted_node.set_label(original_label) - # Apply highlight to current node label_text = str(node.label).strip() if isinstance(node.label, Text): label_text = node.label.plain @@ -108,9 +147,7 @@ class PolicyTreeWidget(Widget): node.set_label(highlighted_label) self.last_highlighted_node = node - # Update details pane if data: - # Work with dataclass objects using __dict__ details = "\n".join( f"{key}: {value}" for key, value in data.__dict__.items() ) @@ -118,12 +155,9 @@ class PolicyTreeWidget(Widget): details = f"Selected: {node.label}" details_pane.update(details) - # Stop event from bubbling message.stop() def on_input_submitted(self, message: Input.Submitted) -> None: - """Handle search input submission.""" - # Remove existing match selector FIRST self._remove_match_selector() query = message.value.strip().lower() @@ -138,7 +172,6 @@ class PolicyTreeWidget(Widget): label_text = str(node.label).lower() label_to_node[label_text] = node if node.data: - # Use __dict__ for dataclass objects data_dict = ( node.data.__dict__ if hasattr(node.data, "__dict__") else node.data ) @@ -146,18 +179,15 @@ class PolicyTreeWidget(Widget): if isinstance(value, str): label_to_node[value.lower()] = node - # Wildcard-style substring match matches = sorted([label for label in label_to_node if query in label]) if matches: - # Try to reuse existing match_selector or create new one try: option_list = self.query_one("#match_selector", OptionList) option_list.clear_options() - option_list.display = True # Ensure it's visible + option_list.display = True except: option_list = OptionList(id="match_selector") - # Mount to the details pane's parent (the Vertical container) details_pane.parent.mount(option_list) for label in matches: @@ -165,17 +195,14 @@ class PolicyTreeWidget(Widget): details_pane.update(f"Found {len(matches)} matches. Select one below.") else: - # Hide or remove the match_selector when no matches self._remove_match_selector() details_pane.update("No matches found.") def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: - """Handle selection from search results.""" selected_id = event.option.id.replace("match_", "") tree = self.query_one("#policy_tree", Tree) details_pane = self.query_one("#details_pane", Static) - # Find the node all_nodes = [] self._collect_tree_nodes(tree.root, all_nodes) @@ -183,7 +210,6 @@ class PolicyTreeWidget(Widget): match_node = label_to_node.get(selected_id.lower()) if match_node: - # Expand path (original working logic) node = match_node path = [] while node: @@ -198,7 +224,6 @@ class PolicyTreeWidget(Widget): match_node.set_label(Text(str(match_node.label), style="reverse bold")) details_pane.update(f"Selected: {match_node.label}") - # Remove the match_selector after selection try: option_list = self.query_one("#match_selector", OptionList) option_list.remove()