From 26199c72cb72f4492466e761fb73f3cde613e3a9 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Thu, 6 Nov 2025 17:36:45 -0500 Subject: [PATCH] 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))