From 26199c72cb72f4492466e761fb73f3cde613e3a9 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Thu, 6 Nov 2025 17:36:45 -0500 Subject: [PATCH 1/4] Device search is working. Working on The multidevice seach, WIP --- models/agent.py | 6 +- utils/tui.py | 145 ++++++++----------------- widgets/multiagentselector.py | 133 +++++++++++++++++++++++ widgets/policytreewidget.py | 199 ++++++++++++++++++++++++++++++++++ widgets/themeselector.py | 45 ++++++++ 5 files changed, 423 insertions(+), 105 deletions(-) create mode 100644 widgets/multiagentselector.py create mode 100644 widgets/policytreewidget.py create mode 100644 widgets/themeselector.py diff --git a/models/agent.py b/models/agent.py index a41a284..ee6678c 100644 --- a/models/agent.py +++ b/models/agent.py @@ -21,18 +21,18 @@ from models.policy import Policy @dataclass class Agent: + hostname: str agentid: str clientversion: str domain: str freespace: int - groupid: str # Changed to str to match UUID-style IDs - hostname: str + groupid: str ip: str localip: str lastcheckin: str os: str policyversion: str - status: int # raw status code + status: int username: str groupname: Optional[str] = field(default=None) status_text: Optional[str] = field(default=None) diff --git a/utils/tui.py b/utils/tui.py index 5d026b8..a9e0bc1 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -5,7 +5,7 @@ import sys import dotenv from dotenv import set_key from textual.app import App, ComposeResult -from textual.containers import Horizontal, Vertical +from textual.containers import Vertical from textual.reactive import reactive from textual.screen import Screen from textual.widgets import ( @@ -16,7 +16,6 @@ from textual.widgets import ( Static, Tab, Tabs, - Tree, ) from flows.otp import otp_activities_by_agent, otp_generate, otp_revoke @@ -28,6 +27,9 @@ from services.policyhandler import confirmUpdateAfromE from utils.configmanager import load_env from utils.setup import get_base_directory, load_user_config from utils.utils import open_directory +from widgets.multiagentselector import MultiAgentSelector +from widgets.policytreewidget import PolicyTreeWidget +from widgets.themeselector import ThemeSelector dotenv.load_dotenv() @@ -115,21 +117,6 @@ class MainMenuScreen(Screen): ], } - # textual themes to expose - THEME_BUTTONS = [ - ("textual-dark", "textual-dark"), - ("textual-light", "textual-light"), - ("nord", "nord"), - ("gruvbox", "gruvbox"), - ("catppuccin-mocha", "catppuccin-mocha"), - ("dracula", "dracula"), - ("tokyo-night", "tokyo-night"), - ("monokai", "monokai"), - ("flexoki", "flexoki"), - ("catppuccin-latte", "catppuccin-latte"), - ("solarized-light", "solarized-light"), - ] - def __init__(self) -> None: super().__init__() self.extras = load_env("EXTRAS") @@ -143,7 +130,7 @@ class MainMenuScreen(Screen): buttons = [] for label, btn_id in defs: btn = Button(label, id=btn_id) - btn.styles.width = "100%" # Make button span full width of parent + btn.styles.width = "100%" buttons.append(btn) return Vertical(*buttons) @@ -157,6 +144,7 @@ class MainMenuScreen(Screen): Tab("OTP", id="otp"), Tab("Directory", id="dir"), Tab("Settings", id="settings"), + Tab("Multi Select", id="multi_select"), ] if self.extras == "POLICYPREP": @@ -216,88 +204,36 @@ class MainMenuScreen(Screen): elif tab_id == "dir": content.mount(DirectoryTree(self.working_dir, id="dir_tree")) elif tab_id == "p_tree": - layout = Horizontal() - content.mount(layout) - - # Left: Policy Tree - policy_tree = Tree("Policies", id="policy_tree") - policy_tree.styles.width = "2fr" - layout.mount(policy_tree) - - # Right: Details pane - details_pane = Static( - "Select a policy or device to view details", id="details-pane" - ) - details_pane.styles.width = "3fr" - layout.mount(details_pane) - - # Build the tree - node_map = {} - - # Top-level policies - for _, policy in self.app.policies.iterrows(): - if policy["parent"] == "global-policy-settings": - node = policy_tree.root.add( - label=policy["name"], data=policy.to_dict() - ) - node_map[policy["groupid"]] = node - - # Child policies - for _, policy in self.app.policies.iterrows(): - 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.to_dict()) - node_map[policy["groupid"]] = node - - # Devices under policies - for _, device in self.app.devices.iterrows(): - group_id = device["groupid"] - if group_id in node_map: - parent_node = node_map[group_id] - label = device["hostname"] # Keep tree clean - parent_node.add(label=label, data=device.to_dict()) - + content.mount(PolicyTreeWidget(self.app.policies, self.app.devices)) elif tab_id == "settings": - # Create and mount the horizontal container - horizontal_container = Horizontal(id="settings_grid") - horizontal_container.styles.layout = "horizontal" - horizontal_container.styles.height = "auto" - content.mount(Static("Theme Options")) - content.mount(horizontal_container) # Mount the horizontal container first - - # Create 3 columns - for i in range(1): - column = Vertical() - column.styles.width = "1fr" - column.styles.height = "auto" - horizontal_container.mount(column) # Mount each column - - for j in range(i, len(self.THEME_BUTTONS), 1): - if j < len(self.THEME_BUTTONS): - label, btn_id = self.THEME_BUTTONS[j] - button = Button(label, id=f"set_theme_{btn_id}", compact=True) - # button.styles.width = "100%" - column.mount(button) # Mount each button - + content.mount(ThemeSelector()) + elif tab_id == "multi_select": + content.mount(MultiAgentSelector(self.app.devices.to_dict("records"))) else: content.mount(Static(f"Unknown tab: {tab_id}")) def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None: self.switch_tab(event.tab.id) - def on_tree_node_selected(self, message: Tree.NodeSelected) -> None: - node = message.node - data = node.data + def on_multi_agent_selector_agents_selected( + self, message: MultiAgentSelector.AgentsSelected + ) -> None: + """Handle selected agents from MultiAgentSelector.""" + global _PENDING_JOB + selected_agents = message.selected_agents + logger.info("Selected agents: %s", selected_agents) + # TODO: Implement actual handling of selected agents + _PENDING_JOB = ("multi_agent_action", selected_agents) + self.app.exit() - details_pane = self.query_one("#details-pane", Static) - - if data: - details = "\n".join(f"{key}: {value}" for key, value in data.items()) - else: - details = f"Selected: {node.label}" - - details_pane.update(details) + def on_theme_selector_theme_selected( + self, message: ThemeSelector.ThemeSelected + ) -> None: + """Handle theme selection from ThemeSelector.""" + global _PENDING_JOB + _persist_user_theme(message.theme_name) + _PENDING_JOB = ("restart",) + self.app.exit() def on_directory_tree_file_selected( self, event: DirectoryTree.FileSelected @@ -315,14 +251,6 @@ class MainMenuScreen(Screen): button_id = event.button.id logger.debug("Button pressed: %s", button_id) - # theme selection β†’ user config - if button_id.startswith("set_theme_"): - theme_name = button_id.replace("set_theme_", "") - _persist_user_theme(theme_name) - _PENDING_JOB = ("restart",) - self.app.exit() - return - match button_id: case "find_device_button": _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {}) @@ -383,8 +311,15 @@ class Loxide(App): if not os.path.isdir(wd): wd = os.getcwd() self.working_dir = wd - self.policies = api.policy_find_all() - self.devices = api.agent_find_all() + + # Add error handling for API calls + try: + self.policies = api.policy_find_all() + self.devices = api.agent_find_all() + except Exception as exc: + logger.error("Failed to load policies/devices: %s", exc) + self.policies = None + self.devices = None def on_mount(self) -> None: self.theme = self._textual_theme @@ -473,6 +408,12 @@ def run_Loxide(api: AirlockAPIWrapper) -> None: # just loop again; fresh .env was already loaded at the top continue + if job[0] == "multi_agent_action": + # Handle multi-agent selection + # TODO: Implement actual multi-agent action handling + logger.info("Multi-agent action with selected agents: %s", job[1]) + continue + break diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py new file mode 100644 index 0000000..b6521d1 --- /dev/null +++ b/widgets/multiagentselector.py @@ -0,0 +1,133 @@ +import difflib + +from textual.containers import Horizontal +from textual.css.query import NoMatches +from textual.message import Message +from textual.widget import Widget +from textual.widgets import Button, SelectionList, Static, Switch, TextArea + + +class MultiAgentSelector(Widget): + class AgentsSelected(Message): + def __init__(self, selected_agents): + super().__init__() + self.selected_agents = selected_agents + + def __init__(self, all_agents: list[dict]): + super().__init__() + self.all_agents = all_agents + self._match_type = "fuzzy" + + @property + def match_type(self): + return self._match_type + + @match_type.setter + def match_type(self, value): + self._match_type = value + + def compose(self): + yield Static("πŸ” Multi-Agent Selector", id="selector_title") + text_area = TextArea( + id="device_input", + placeholder="Paste device names here (one per line)", + ) + text_area.styles.height = 10 + text_area.styles.overflow_y = "auto" + yield text_area + + with Horizontal() as switch_container: + switch_container.styles.height = "auto" + switch_container.styles.align = ("left", "middle") + switch_label = Static("Match Type: Fuzzy", id="match_switch_label") + switch_label.styles.width = "auto" + switch_label.styles.padding = (0, 1) + yield switch_label + + switch = Switch(value=False, id="match_switch") + switch.styles.width = "auto" + yield switch + + yield Button("Search", id="search_button") + + yield SelectionList(id="match_results") + yield Static(id="unmatched_label") + + yield Horizontal( + Button("Select All", id="select_all"), + Button("Select None", id="select_none"), + ) + + yield Button("Continue with Selected", id="submit_selection", variant="primary") + + def on_switch_changed(self, event: Switch.Changed): + self.match_type = "exact" if event.value else "fuzzy" + self.query_one("#match_switch_label", Static).update( + f"Match Type: {self.match_type.capitalize()}" + ) + + def on_button_pressed(self, event: Button.Pressed): + btn_id = event.button.id # can be None for internal buttons + + # Only query when needed + try: + match_list = self.query_one("#match_results", SelectionList) + except NoMatches: + # UI not mounted yet or id changedβ€”just ignore gracefully + return + + if btn_id == "select_all": + match_list.select_all() + event.stop() + elif btn_id == "select_none": + match_list.deselect_all() + event.stop() + elif btn_id == "submit_selection": + selected = list(match_list.selected) + self.post_message(self.AgentsSelected(selected)) + event.stop() + elif btn_id == "search_button": + self.update_matches() + event.stop() + + def update_matches(self): + raw_input = self.query_one("#device_input", TextArea).text.strip() + device_names = [line.strip() for line in raw_input.split("\n") if line.strip()] + matched, unmatched = self.match_devices(device_names) + match_list = self.query_one("#match_results", SelectionList) + match_list.clear_options() + for name in matched: + match_list.add_option((name, name)) + unmatched_label = self.query_one("#unmatched_label", Static) + if unmatched: + unmatched_label.update(f"⚠️ No matches for: {', '.join(unmatched)}") + else: + unmatched_label.update("") + + def match_devices(self, device_names: list[str]) -> tuple[list[str], list[str]]: + if not self.all_agents or not device_names: + return [], device_names + agent_names = [agent["hostname"] for agent in self.all_agents] + matched = set() + unmatched = [] + for name in device_names: + if self.match_type == "exact": + # Case-insensitive exact match + name_lower = name.lower() + exact_match = None + for agent_name in agent_names: + if agent_name.lower() == name_lower: + exact_match = agent_name + break + if exact_match: + matched.add(exact_match) + else: + unmatched.append(name) + else: + # Fuzzy match + matches = difflib.get_close_matches(name, agent_names, n=5, cutoff=0.5) + if matches: + matched.update(matches) + else: + unmatched.append(name) + return sorted(matched), unmatched diff --git a/widgets/policytreewidget.py b/widgets/policytreewidget.py new file mode 100644 index 0000000..88406e9 --- /dev/null +++ b/widgets/policytreewidget.py @@ -0,0 +1,199 @@ +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, 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 + + def compose(self): + # Left: Policy Tree + policy_tree = Tree("Policies", id="policy_tree") + 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" + ) + details_pane = Static("", id="details_pane") + + with Horizontal(): + yield policy_tree + with Vertical() as right_pane: + right_pane.styles.width = "3fr" + yield label + yield search_box + yield details_pane + + def on_mount(self) -> None: + """Build the tree after mounting.""" + self._build_tree() + + def _build_tree(self) -> None: + """Build the policy tree structure.""" + policy_tree = self.query_one("#policy_tree", Tree) + node_map = {} + + # Top-level policies + for _, policy in self.policies.iterrows(): + if policy["parent"] == "global-policy-settings": + node = policy_tree.root.add(label=policy["name"], data=policy.to_dict()) + node_map[policy["groupid"]] = node + + # Child policies + for _, policy in self.policies.iterrows(): + 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.to_dict()) + node_map[policy["groupid"]] = node + + # Devices under policies + for _, device in self.devices.iterrows(): + 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.to_dict()) + + 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: + 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: + """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 + highlighted_label = Text(label_text, style="reverse bold") + node.set_label(highlighted_label) + self.last_highlighted_node = node + + # Update details pane + if data: + details = "\n".join(f"{key}: {value}" for key, value in data.items()) + else: + 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() + 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: + for key, value in node.data.items(): + 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 + 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: + option_list.add_option(Option(label, id=f"match_{label}")) + + 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) + + 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: + # Expand path (original working logic) + 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}") + + # Remove the match_selector after selection + try: + option_list = self.query_one("#match_selector", OptionList) + option_list.remove() + except: + pass diff --git a/widgets/themeselector.py b/widgets/themeselector.py new file mode 100644 index 0000000..fda7cd1 --- /dev/null +++ b/widgets/themeselector.py @@ -0,0 +1,45 @@ +from textual.containers import Vertical +from textual.message import Message +from textual.widget import Widget +from textual.widgets import Button, Static + + +class ThemeSelector(Widget): + """Widget for selecting and applying Textual themes.""" + + class ThemeSelected(Message): + """Message posted when a theme is selected.""" + + def __init__(self, theme_name: str): + super().__init__() + self.theme_name = theme_name + + AVAILABLE_THEMES = [ + ("textual-dark", "textual-dark"), + ("textual-light", "textual-light"), + ("nord", "nord"), + ("gruvbox", "gruvbox"), + ("catppuccin-mocha", "catppuccin-mocha"), + ("dracula", "dracula"), + ("tokyo-night", "tokyo-night"), + ("monokai", "monokai"), + ("flexoki", "flexoki"), + ("catppuccin-latte", "catppuccin-latte"), + ("solarized-light", "solarized-light"), + ] + + def compose(self): + yield Static("Theme Options", id="theme_title") + + with Vertical() as column: + column.styles.width = "1fr" + column.styles.height = "auto" + + for label, btn_id in self.AVAILABLE_THEMES: + yield Button(label, id=f"set_theme_{btn_id}", compact=True) + + def on_button_pressed(self, event: Button.Pressed) -> None: + button_id = event.button.id + if button_id and button_id.startswith("set_theme_"): + theme_name = button_id.replace("set_theme_", "") + self.post_message(self.ThemeSelected(theme_name)) From 9dfdf8b9eef76dc994ea8c1ea07b4fe82562269e Mon Sep 17 00:00:00 2001 From: Zarithas Date: Thu, 6 Nov 2025 20:26:06 -0500 Subject: [PATCH 2/4] Multiselect appears to be working mostly correctly and is styled. It is only temp a tab, as this is going to be an intermediate/ephem widget --- widgets/multiagentselector.py | 113 ++++++++++++++++++++++++---------- 1 file changed, 82 insertions(+), 31 deletions(-) diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py index b6521d1..a06742d 100644 --- a/widgets/multiagentselector.py +++ b/widgets/multiagentselector.py @@ -1,6 +1,6 @@ import difflib -from textual.containers import Horizontal +from textual.containers import Horizontal, Vertical from textual.css.query import NoMatches from textual.message import Message from textual.widget import Widget @@ -16,7 +16,7 @@ class MultiAgentSelector(Widget): def __init__(self, all_agents: list[dict]): super().__init__() self.all_agents = all_agents - self._match_type = "fuzzy" + self._match_type = "exact" @property def match_type(self): @@ -27,43 +27,73 @@ class MultiAgentSelector(Widget): self._match_type = value def compose(self): - yield Static("πŸ” Multi-Agent Selector", id="selector_title") - text_area = TextArea( - id="device_input", - placeholder="Paste device names here (one per line)", - ) - text_area.styles.height = 10 - text_area.styles.overflow_y = "auto" - yield text_area + title_text = Static("πŸ–§ Multi-Agent Selector", id="selector_title") + title_text.styles.margin = (0, 0, 0, 1) + yield title_text - with Horizontal() as switch_container: - switch_container.styles.height = "auto" - switch_container.styles.align = ("left", "middle") - switch_label = Static("Match Type: Fuzzy", id="match_switch_label") - switch_label.styles.width = "auto" - switch_label.styles.padding = (0, 1) - yield switch_label + with Horizontal() as main_layout: + main_layout.styles.height = "auto" - switch = Switch(value=False, id="match_switch") - switch.styles.width = "auto" - yield switch + # Left side - Input and controls + with Vertical() as left_pane: + left_pane.styles.width = "1fr" + left_pane.styles.height = "auto" - yield Button("Search", id="search_button") + text_area = TextArea( + id="device_input", + placeholder="Paste device names here (one per line). Supports wildcards: * and ?", + ) + text_area.styles.height = 10 + text_area.styles.overflow_y = "auto" + yield text_area - yield SelectionList(id="match_results") - yield Static(id="unmatched_label") + with Horizontal(id="switch_search_container") as switch_search: - yield Horizontal( - Button("Select All", id="select_all"), - Button("Select None", id="select_none"), - ) + switch = Switch(value=False, id="match_switch") + switch.styles.width = "auto" + switch.styles.margin = (1, 0, 0, 0) + switch.styles.padding = (0, 0, 0, 0) + yield switch - yield Button("Continue with Selected", id="submit_selection", variant="primary") + switch_label = Static("Match: Exact", id="match_switch_label") + switch_label.styles.width = "auto" + switch_label.styles.margin = (2, 1, 0, 0) + yield switch_label + + search = Button("πŸ” Search", id="search_button") + search.styles.margin = (1, 0, 0, 0) + + yield search + + with Horizontal() as select_buttons: + select_buttons.styles.margin = (0, 0, 0, 0) + + select_all_button = Button("βœ… Select All", id="select_all") + select_all_button.styles.margin = (1, 1, 0, 1) + yield select_all_button + + select_none_button = Button("🚫 Select None", id="select_none") + select_none_button.styles.margin = (1, 0, 0, 1) + yield select_none_button + + submit_button = Button( + "β–Ί Continue with Selected", id="submit_selection", variant="primary" + ) + submit_button.styles.margin = (0, 5, 2, 1) + submit_button.styles.padding = (0, 6, 0, 0) + yield submit_button + + # Right side - Results + with Vertical() as right_pane: + right_pane.styles.width = "2fr" + + yield SelectionList(id="match_results") + yield Static(id="unmatched_label") def on_switch_changed(self, event: Switch.Changed): - self.match_type = "exact" if event.value else "fuzzy" + self.match_type = "fuzzy" if event.value else "exact" self.query_one("#match_switch_label", Static).update( - f"Match Type: {self.match_type.capitalize()}" + f"Match: {self.match_type.capitalize()}" ) def on_button_pressed(self, event: Button.Pressed): @@ -111,7 +141,28 @@ class MultiAgentSelector(Widget): matched = set() unmatched = [] for name in device_names: - if self.match_type == "exact": + # Check if the name contains wildcards + has_wildcards = "*" in name or "?" in name + + if has_wildcards: + # Use regex for wildcard matching + import re + + # Escape special regex characters except * and ? + pattern = re.escape(name) + # Convert wildcards to regex + pattern = pattern.replace(r"\*", ".*").replace(r"\?", ".") + # Make it case-insensitive and match full string + regex = re.compile(f"^{pattern}$", re.IGNORECASE) + + wildcard_matches = [ + agent_name for agent_name in agent_names if regex.match(agent_name) + ] + if wildcard_matches: + matched.update(wildcard_matches) + else: + unmatched.append(name) + elif self.match_type == "exact": # Case-insensitive exact match name_lower = name.lower() exact_match = None From 8a9b04cb7c320d3b140769e3053148f87e8e2536 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Fri, 7 Nov 2025 09:56:45 -0500 Subject: [PATCH 3/4] Updated OTP Generate with new workflow --- requirements.txt | 1 + screens/otpworkflowscreen.py | 41 ++++ utils/tui.py | 91 +++++++-- widgets/OTP_generate.py | 346 ++++++++++++++++++++++++++++++++++ widgets/multiagentselector.py | 60 +++--- widgets/policytreewidget.py | 35 ++-- 6 files changed, 524 insertions(+), 50 deletions(-) create mode 100644 screens/otpworkflowscreen.py create mode 100644 widgets/OTP_generate.py diff --git a/requirements.txt b/requirements.txt index 1141830..65ee4c7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,7 @@ Requests==2.32.5 textual==6.5.0 tqdm==4.67.1 urllib3==2.5.0 +pyperclip==1.11.0 --extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/ airlock_libs==2.0.0 \ No newline at end of file diff --git a/screens/otpworkflowscreen.py b/screens/otpworkflowscreen.py new file mode 100644 index 0000000..fa5afc1 --- /dev/null +++ b/screens/otpworkflowscreen.py @@ -0,0 +1,41 @@ +from typing import List + +from textual.app import ComposeResult +from textual.screen import Screen + +from models.agent import Agent +from widgets.multiagentselector import MultiAgentSelector +from widgets.OTP_generate import OTPGenerator + + +class OTPWorkflowScreen(Screen): + """Screen that handles the OTP generation workflow.""" + + def __init__(self, all_agents: List[Agent]): + super().__init__() + self.all_agents = all_agents + self.selected_devices = None + + def compose(self) -> ComposeResult: + """Start with the multi-agent selector.""" + yield MultiAgentSelector(self.all_agents) + + def on_multi_agent_selector_agents_selected( + self, message: MultiAgentSelector.AgentsSelected + ) -> None: + """Handle selected agents - switch to OTP generator.""" + self.selected_devices = message.selected_agents + + # Remove the MultiAgentSelector + selector = self.query_one(MultiAgentSelector) + selector.remove() + + # Mount the OTPGenerator with the selected Agent objects + # No need to pass API - it will access self.app.api directly + self.mount(OTPGenerator(self.selected_devices)) + + def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None: + """Handle OTP generation request - call the actual OTP generation function.""" + # This will be handled by the main app, but we can also do it here + # For now, just pass it up to the app level + pass diff --git a/utils/tui.py b/utils/tui.py index a9e0bc1..6eb070d 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -18,9 +18,12 @@ from textual.widgets import ( Tabs, ) -from flows.otp import otp_activities_by_agent, otp_generate, otp_revoke +from flows.otp import otp_activities_by_agent, otp_revoke from flows.prepPolicy import menu_policy_enforce from flows.quietAgent import findQuietAgents +from models.agent import Agent +from models.policy import Policy +from screens.otpworkflowscreen import OTPWorkflowScreen from services.agenthandler import findAgents, moveAgents, toggleEnforcement from services.API import AirlockAPIWrapper from services.policyhandler import confirmUpdateAfromE @@ -28,6 +31,7 @@ from utils.configmanager import load_env from utils.setup import get_base_directory, load_user_config from utils.utils import open_directory from widgets.multiagentselector import MultiAgentSelector +from widgets.OTP_generate import OTPGenerator from widgets.policytreewidget import PolicyTreeWidget from widgets.themeselector import ThemeSelector @@ -107,7 +111,7 @@ class MainMenuScreen(Screen): ("πŸ”€ - Move - Other", "move_other_button"), ], "otp": [ - ("πŸ” - Generate OTPs", "otp_generate_button"), + ("🎫 - Generate OTPs", "otp_generate_button"), ("πŸ“Š - OTP Activities By Agent", "otp_activities_button"), ("❌ - Revoke OTPs", "otp_revoke_button"), ], @@ -117,8 +121,9 @@ class MainMenuScreen(Screen): ], } - def __init__(self) -> None: + def __init__(self, api: AirlockAPIWrapper) -> None: super().__init__() + self.api = api self.extras = load_env("EXTRAS") wd = load_env("WORKING_DIR") or os.getcwd() if not os.path.isdir(wd): @@ -144,7 +149,6 @@ class MainMenuScreen(Screen): Tab("OTP", id="otp"), Tab("Directory", id="dir"), Tab("Settings", id="settings"), - Tab("Multi Select", id="multi_select"), ] if self.extras == "POLICYPREP": @@ -207,8 +211,6 @@ class MainMenuScreen(Screen): content.mount(PolicyTreeWidget(self.app.policies, self.app.devices)) elif tab_id == "settings": content.mount(ThemeSelector()) - elif tab_id == "multi_select": - content.mount(MultiAgentSelector(self.app.devices.to_dict("records"))) else: content.mount(Static(f"Unknown tab: {tab_id}")) @@ -235,6 +237,30 @@ class MainMenuScreen(Screen): _PENDING_JOB = ("restart",) self.app.exit() + def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None: + """Handle OTP generation request from the workflow.""" + global _PENDING_JOB + + # Log what we received + logger.info( + "OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d", + len(message.devices), + message.requestor, + message.reasoning, + message.duration, + ) + + # Set up the job to run the OTP generation + _PENDING_JOB = ( + "otp_workflow", + message.devices, + message.requestor, + message.reasoning, + message.duration, + ) + + self.app.exit() + def on_directory_tree_file_selected( self, event: DirectoryTree.FileSelected ) -> None: @@ -268,7 +294,10 @@ class MainMenuScreen(Screen): case "move_other_button": _PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {}) case "otp_generate_button": - _PENDING_JOB = ("legacy", otp_generate, (self.app.api,), {}) + # NEW: Push OTP workflow screen instead of legacy function + self.app.push_screen(OTPWorkflowScreen(self.app.devices)) + event.stop() + return # Don't exit the app case "otp_activities_button": _PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {}) case "otp_revoke_button": @@ -314,16 +343,20 @@ class Loxide(App): # Add error handling for API calls try: - self.policies = api.policy_find_all() - self.devices = api.agent_find_all() + self.policies = [ + Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows() + ] + self.devices = [ + Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows() + ] except Exception as exc: logger.error("Failed to load policies/devices: %s", exc) self.policies = None self.devices = None - def on_mount(self) -> None: + def on_mount(self, api: AirlockAPIWrapper) -> None: self.theme = self._textual_theme - self.push_screen(MainMenuScreen()) + self.push_screen(MainMenuScreen(api)) def action_quit(self) -> None: global _PENDING_JOB @@ -410,10 +443,44 @@ def run_Loxide(api: AirlockAPIWrapper) -> None: if job[0] == "multi_agent_action": # Handle multi-agent selection - # TODO: Implement actual multi-agent action handling logger.info("Multi-agent action with selected agents: %s", job[1]) continue + # NEW: Handle OTP workflow + if job[0] == "otp_workflow": + _, devices, requestor, reasoning, duration = job + + # Call your OTP generation with the parameters + def otp_generate_with_params(): + + print(f"\n{'='*60}") + print("OTP GENERATION") + print(f"{'='*60}") + print(f"Requestor: {requestor}") + print(f"Reasoning: {reasoning}") + print(f"Duration: {duration} minutes") + print(f"\nGenerating OTPs for {len(devices)} devices:") + print(f"{'='*60}\n") + + # Call your actual OTP generation function + # You'll need to adapt otp_generate to accept these parameters + # For now, this is a placeholder showing the structure + for device in devices: + print(f"Device: {device}") + print(f" Requestor: {requestor}") + print(f" Reason: {reasoning}") + print(f" Duration: {duration} minutes") + # TODO: Actually call your API to generate OTP + # result = api.generate_otp(device, requestor, reasoning, duration) + print() + + print(f"{'='*60}") + print("OTP Generation Complete!") + print(f"{'='*60}") + + _run_legacy_job(otp_generate_with_params, (), {}) + continue + break diff --git a/widgets/OTP_generate.py b/widgets/OTP_generate.py new file mode 100644 index 0000000..b35b39d --- /dev/null +++ b/widgets/OTP_generate.py @@ -0,0 +1,346 @@ +import logging +from typing import List + +from textual.containers import Horizontal, Vertical +from textual.css.query import NoMatches +from textual.message import Message +from textual.reactive import reactive +from textual.widget import Widget +from textual.widgets import Button, Input, RadioButton, RadioSet, Static, TextArea + +from models.agent import Agent + +logger = logging.getLogger(__name__) + + +class OTPGenerator(Widget): + # Reactive properties to track form completion + requestor_filled = reactive(False) + reasoning_filled = reactive(False) + duration_selected = reactive(True) # Default is selected + otp_generated = reactive(False) + + class OTPInfo(Message): + def __init__( + self, devices: List[Agent], requestor: str, reasoning: str, duration: int + ): + super().__init__() + self.devices = devices + self.requestor = requestor + self.reasoning = reasoning + self.duration = duration + + # Duration options in minutes + DURATION_OPTIONS = [ + (15, "15 minutes"), + (60, "1 hour"), + (360, "6 hours"), + (1440, "1 day"), + (10080, "7 days"), + ] + + def __init__(self, devices: List[Agent]): + """Initialize with a list of Agent objects.""" + super().__init__() + self.devices = devices + + def watch_requestor_filled(self, old_value: bool, new_value: bool) -> None: + """Update button state when requestor changes.""" + self._update_button_state() + + def watch_reasoning_filled(self, old_value: bool, new_value: bool) -> None: + """Update button state when reasoning changes.""" + self._update_button_state() + + def watch_otp_generated(self, old_value: bool, new_value: bool) -> None: + """Update button state when OTP is generated.""" + self._update_button_state() + + def _update_button_state(self) -> None: + """Enable/disable the generate button based on form state.""" + try: + button = self.query_one("#generate_button", Button) + # Enable only if all fields filled and OTP not yet generated + button.disabled = not ( + self.requestor_filled + and self.reasoning_filled + and not self.otp_generated + ) + except NoMatches: + pass + + def compose(self): + title_text = Static( + f"🎫 Generate One Time Passes for {len(self.devices)} device(s)", + id="otpgen_title", + ) + title_text.styles.margin = (0, 0, 1, 0) + yield title_text + + with Horizontal() as main_layout: + main_layout.styles.height = "auto" + + # Left side - Inputs and controls + with Vertical() as left_side: + left_side.styles.width = "1fr" + left_side.styles.height = "auto" + + # Requestor input + requestor_label = Static("Who is requesting OTP?") + requestor_label.styles.margin = (0, 0, 0, 0) + yield requestor_label + + requestor_box = Input( + placeholder="Enter requestor name", id="requestor_input" + ) + requestor_box.styles.margin = (0, 0, 1, 0) + yield requestor_box + + # Reasoning input + reasoning_label = Static("What work are they doing?") + reasoning_label.styles.margin = (0, 0, 0, 0) + yield reasoning_label + + reasoning_box = Input( + placeholder="Enter reason for OTP", id="reasoning_input" + ) + reasoning_box.styles.margin = (0, 0, 1, 0) + yield reasoning_box + + # Duration selection + duration_label = Static("Duration:") + duration_label.styles.margin = (0, 0, 0, 0) + yield duration_label + + with RadioSet(id="duration_radio") as radio_set: + radio_set.styles.margin = (0, 0, 1, 0) + for minutes, label in self.DURATION_OPTIONS: + radio = RadioButton(label, id=f"duration_{minutes}") + if minutes == 360: # Default to 6 hours + radio.value = True + yield radio + + # Buttons in a horizontal layout + with Horizontal() as button_row: + button_row.styles.height = "auto" + button_row.styles.margin = (1, 0, 0, 0) + + back_button = Button("← Back", id="back_button") + back_button.styles.width = "1fr" + yield back_button + + generate_button = Button( + "Generate OTP", id="generate_button", variant="primary" + ) + generate_button.styles.width = "2fr" + yield generate_button + + # Right side - Show device list initially, then output after generation + with Vertical() as right_side: + right_side.styles.width = "2fr" + right_side.styles.height = "100%" + + output_label = Static( + f"Selected Devices ({len(self.devices)}):", id="output_label" + ) + output_label.styles.margin = (0, 0, 0, 0) + yield output_label + + # Container for either device list or output + with Vertical(id="output_container") as output_container: + output_container.styles.height = "1fr" + output_container.styles.margin = (1, 0, 0, 0) + output_container.styles.overflow_y = "auto" + output_container.styles.border = ("round", "green") + + # Show device list initially + device_list_text = "\n".join( + f"β€’ {device.hostname}" for device in self.devices + ) + device_display = Static(device_list_text, id="device_display") + yield device_display + + # Copy to clipboard button (hidden initially) + copy_button = Button("πŸ“‹ Copy to Clipboard", id="copy_clipboard_button") + copy_button.styles.margin = (1, 0, 0, 0) + copy_button.styles.display = "none" + yield copy_button + + def on_mount(self) -> None: + """Set initial button state.""" + self._update_button_state() + + def on_input_changed(self, event: Input.Changed) -> None: + """Handle input field changes.""" + input_id = event.input.id + + if input_id == "requestor_input": + self.requestor_filled = bool(event.value.strip()) + elif input_id == "reasoning_input": + self.reasoning_filled = bool(event.value.strip()) + + def on_button_pressed(self, event: Button.Pressed): + btn_id = event.button.id + + if btn_id == "back_button": + self.app.pop_screen() + event.stop() + + elif btn_id == "copy_clipboard_button": + try: + output_area = self.query_one("#otp_output", TextArea) + text_to_copy = output_area.text + + import pyperclip + + pyperclip.copy(text_to_copy) + self.app.notify( + "βœ… Copied to clipboard!", severity="information", timeout=2 + ) + except ImportError: + self.app.notify( + "⚠️ pyperclip not installed. Run: pip install pyperclip", + severity="warning", + ) + except Exception as e: + self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error") + event.stop() + + elif btn_id == "generate_button": + try: + requestor = self.query_one("#requestor_input", Input).value.strip() + reasoning = self.query_one("#reasoning_input", Input).value.strip() + + radio_set = self.query_one("#duration_radio", RadioSet) + selected_button_id = ( + radio_set.pressed_button.id if radio_set.pressed_button else None + ) + + if not selected_button_id: + self._show_error("Please select a duration") + return + + duration = int(selected_button_id.replace("duration_", "")) + + if not requestor or not reasoning: + self._show_error("Please fill in all fields") + return + + self.otp_generated = True + + # Access API from the app - this is the key change! + api = self.app.api + + output_lines = [ + "=" * 60, + "OTP GENERATION RESULTS", + "=" * 60, + ] + + otp_dict = {} + for device in self.devices: + try: + otp_code = api.otp_generate(device.agentid, duration, reasoning) + otp_dict[device.hostname] = otp_code + logger.debug(f"Generated OTP for {device.hostname}: {otp_code}") + except Exception as e: + otp_dict[device.hostname] = f"ERROR: {str(e)}" + logger.error( + f"Failed to generate OTP for {device.hostname}: {e}" + ) + + for hostname, otp_code in otp_dict.items(): + output_lines.append(f"{hostname:30} | {otp_code}") + + output_lines.append("=" * 60) + result_text = "\n".join(output_lines) + self._show_result(result_text) + + # Post message with the OTP info + self.post_message( + self.OTPInfo(self.devices, requestor, reasoning, duration) + ) + event.stop() + + except NoMatches: + self._show_error("UI elements not found") + except Exception as e: + self._show_error(f"Error: {str(e)}") + logger.exception("Error generating OTP") + + def _show_error(self, message: str): + """Display error message in output area.""" + try: + container = self.query_one("#output_container", Vertical) + try: + device_display = self.query_one("#device_display", Static) + device_display.remove() + except NoMatches: + pass + + try: + output_area = self.query_one("#otp_output", TextArea) + except NoMatches: + output_area = TextArea(id="otp_output", read_only=True) + container.mount(output_area) + + output_area.text = f"❌ ERROR: {message}" + except Exception as e: + logger.debug(f"Error showing error message: {e}") + + def _show_result(self, message: str): + """Display result message in output area.""" + try: + container = self.query_one("#output_container", Vertical) + try: + device_display = self.query_one("#device_display", Static) + device_display.remove() + except NoMatches: + pass + + try: + output_area = self.query_one("#otp_output", TextArea) + except NoMatches: + output_area = TextArea(id="otp_output", read_only=True) + container.mount(output_area) + + output_area.text = message + output_label = self.query_one("#output_label", Static) + output_label.update("Generated OTP Details:") + copy_button = self.query_one("#copy_clipboard_button", Button) + copy_button.styles.display = "block" + except Exception as e: + logger.debug(f"Error showing result: {e}") + + def display_otp_result(self, result_text: str): + """Display OTP generation result in the output area.""" + try: + container = self.query_one("#output_container", Vertical) + try: + device_display = self.query_one("#device_display", Static) + device_display.remove() + except NoMatches: + pass + + try: + output_area = self.query_one("#otp_output", TextArea) + except NoMatches: + output_area = TextArea(id="otp_output", read_only=True) + container.mount(output_area) + + output_area.text = result_text + except Exception as e: + logger.debug(f"Error displaying OTP result: {e}") + + def clear_form(self): + """Clear all input fields and reset state.""" + try: + self.query_one("#requestor_input", Input).value = "" + self.query_one("#reasoning_input", Input).value = "" + self.query_one("#otp_output", TextArea).text = "" + self.otp_generated = False + self.requestor_filled = False + self.reasoning_filled = False + self._update_button_state() + except NoMatches: + pass diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py index a06742d..e46c636 100644 --- a/widgets/multiagentselector.py +++ b/widgets/multiagentselector.py @@ -1,4 +1,6 @@ import difflib +import re +from typing import List from textual.containers import Horizontal, Vertical from textual.css.query import NoMatches @@ -6,14 +8,16 @@ from textual.message import Message from textual.widget import Widget from textual.widgets import Button, SelectionList, Static, Switch, TextArea +from models.agent import Agent + class MultiAgentSelector(Widget): class AgentsSelected(Message): - def __init__(self, selected_agents): + def __init__(self, selected_agents: List[Agent]): super().__init__() self.selected_agents = selected_agents - def __init__(self, all_agents: list[dict]): + def __init__(self, all_agents: List[Agent]): super().__init__() self.all_agents = all_agents self._match_type = "exact" @@ -48,7 +52,6 @@ class MultiAgentSelector(Widget): yield text_area with Horizontal(id="switch_search_container") as switch_search: - switch = Switch(value=False, id="match_switch") switch.styles.width = "auto" switch.styles.margin = (1, 0, 0, 0) @@ -62,7 +65,6 @@ class MultiAgentSelector(Widget): search = Button("πŸ” Search", id="search_button") search.styles.margin = (1, 0, 0, 0) - yield search with Horizontal() as select_buttons: @@ -76,17 +78,24 @@ class MultiAgentSelector(Widget): select_none_button.styles.margin = (1, 0, 0, 1) yield select_none_button - submit_button = Button( - "β–Ί Continue with Selected", id="submit_selection", variant="primary" - ) - submit_button.styles.margin = (0, 5, 2, 1) - submit_button.styles.padding = (0, 6, 0, 0) - yield submit_button + with Horizontal() as button_row: + button_row.styles.height = "auto" + button_row.styles.margin = (1, 0, 0, 0) + + back_button = Button("← Back", id="back_button") + back_button.styles.width = "1fr" + yield back_button + + submit_button = Button( + "β–Ά Select & Continue", id="submit_selection", variant="primary" + ) + submit_button.styles.margin = (0, 5, 2, 1) + submit_button.styles.padding = (0, 6, 0, 0) + yield submit_button # Right side - Results with Vertical() as right_pane: right_pane.styles.width = "2fr" - yield SelectionList(id="match_results") yield Static(id="unmatched_label") @@ -97,24 +106,31 @@ class MultiAgentSelector(Widget): ) def on_button_pressed(self, event: Button.Pressed): - btn_id = event.button.id # can be None for internal buttons + btn_id = event.button.id - # Only query when needed try: match_list = self.query_one("#match_results", SelectionList) except NoMatches: - # UI not mounted yet or id changedβ€”just ignore gracefully return - - if btn_id == "select_all": + if btn_id == "back_button": + self.app.pop_screen() + event.stop() + elif btn_id == "select_all": match_list.select_all() event.stop() elif btn_id == "select_none": match_list.deselect_all() event.stop() elif btn_id == "submit_selection": - selected = list(match_list.selected) - self.post_message(self.AgentsSelected(selected)) + # Get selected hostnames + selected_hostnames = list(match_list.selected) + # Convert back to Agent objects + selected_agents = [ + agent + for agent in self.all_agents + if agent.hostname in selected_hostnames + ] + self.post_message(self.AgentsSelected(selected_agents)) event.stop() elif btn_id == "search_button": self.update_matches() @@ -137,22 +153,18 @@ class MultiAgentSelector(Widget): def match_devices(self, device_names: list[str]) -> tuple[list[str], list[str]]: if not self.all_agents or not device_names: return [], device_names - agent_names = [agent["hostname"] for agent in self.all_agents] + agent_names = [agent.hostname for agent in self.all_agents] matched = set() unmatched = [] + for name in device_names: # Check if the name contains wildcards has_wildcards = "*" in name or "?" in name if has_wildcards: # Use regex for wildcard matching - import re - - # Escape special regex characters except * and ? pattern = re.escape(name) - # Convert wildcards to regex pattern = pattern.replace(r"\*", ".*").replace(r"\?", ".") - # Make it case-insensitive and match full string regex = re.compile(f"^{pattern}$", re.IGNORECASE) wildcard_matches = [ diff --git a/widgets/policytreewidget.py b/widgets/policytreewidget.py index 88406e9..3a0efff 100644 --- a/widgets/policytreewidget.py +++ b/widgets/policytreewidget.py @@ -49,26 +49,26 @@ class PolicyTreeWidget(Widget): node_map = {} # Top-level policies - for _, policy in self.policies.iterrows(): - if policy["parent"] == "global-policy-settings": - node = policy_tree.root.add(label=policy["name"], data=policy.to_dict()) - node_map[policy["groupid"]] = node + 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 # Child policies - for _, policy in self.policies.iterrows(): - parent_id = policy["parent"] + 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.to_dict()) - node_map[policy["groupid"]] = node + node = parent_node.add(label=policy.name, data=policy) + node_map[policy.groupid] = node # Devices under policies - for _, device in self.devices.iterrows(): - group_id = device["groupid"] + 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.to_dict()) + label = device.hostname + parent_node.add(label=label, data=device) def _collect_tree_nodes(self, node, all_nodes): """Helper to recursively collect all nodes from a tree.""" @@ -110,7 +110,10 @@ class PolicyTreeWidget(Widget): # Update details pane if data: - details = "\n".join(f"{key}: {value}" for key, value in data.items()) + # Work with dataclass objects using __dict__ + details = "\n".join( + f"{key}: {value}" for key, value in data.__dict__.items() + ) else: details = f"Selected: {node.label}" details_pane.update(details) @@ -135,7 +138,11 @@ class PolicyTreeWidget(Widget): label_text = str(node.label).lower() label_to_node[label_text] = node if node.data: - for key, value in node.data.items(): + # Use __dict__ for dataclass objects + 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 From ed07c79b742711b2d40649a1b5365e074e26896f Mon Sep 17 00:00:00 2001 From: Zarithas Date: Fri, 7 Nov 2025 15:44:11 -0500 Subject: [PATCH 4/4] 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()