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)
 

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()