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/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 5d026b8..6eb070d 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,18 +16,24 @@ from textual.widgets import ( Static, Tab, Tabs, - Tree, ) -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 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 dotenv.load_dotenv() @@ -105,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"), ], @@ -115,23 +121,9 @@ 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: + 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): @@ -143,7 +135,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) @@ -216,88 +208,58 @@ 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()) 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) + 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() - if data: - details = "\n".join(f"{key}: {value}" for key, value in data.items()) - else: - details = f"Selected: {node.label}" + def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None: + """Handle OTP generation request from the workflow.""" + global _PENDING_JOB - details_pane.update(details) + # 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 @@ -315,14 +277,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), {}) @@ -340,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": @@ -383,12 +340,23 @@ 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() - def on_mount(self) -> None: + # Add error handling for API calls + try: + 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, 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 @@ -473,6 +441,46 @@ 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 + 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..6aff859 --- /dev/null +++ b/widgets/OTP_generate.py @@ -0,0 +1,345 @@ +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 = [ + "Requested OTP Codes:", + "=" * 25, + ] + + 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} | {otp_code}") + + output_lines.append("=" * 25) + 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 new file mode 100644 index 0000000..e46c636 --- /dev/null +++ b/widgets/multiagentselector.py @@ -0,0 +1,196 @@ +import difflib +import re +from typing import List + +from textual.containers import Horizontal, Vertical +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 + +from models.agent import Agent + + +class MultiAgentSelector(Widget): + class AgentsSelected(Message): + def __init__(self, selected_agents: List[Agent]): + super().__init__() + self.selected_agents = selected_agents + + def __init__(self, all_agents: List[Agent]): + super().__init__() + self.all_agents = all_agents + self._match_type = "exact" + + @property + def match_type(self): + return self._match_type + + @match_type.setter + def match_type(self, value): + self._match_type = value + + def compose(self): + title_text = Static("🖧 Multi-Agent Selector", id="selector_title") + title_text.styles.margin = (0, 0, 0, 1) + yield title_text + + with Horizontal() as main_layout: + main_layout.styles.height = "auto" + + # Left side - Input and controls + with Vertical() as left_pane: + left_pane.styles.width = "1fr" + left_pane.styles.height = "auto" + + 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 + + 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) + switch.styles.padding = (0, 0, 0, 0) + yield switch + + 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 + + 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") + + def on_switch_changed(self, event: Switch.Changed): + self.match_type = "fuzzy" if event.value else "exact" + self.query_one("#match_switch_label", Static).update( + f"Match: {self.match_type.capitalize()}" + ) + + def on_button_pressed(self, event: Button.Pressed): + btn_id = event.button.id + + try: + match_list = self.query_one("#match_results", SelectionList) + except NoMatches: + return + 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": + # 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() + 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: + # Check if the name contains wildcards + has_wildcards = "*" in name or "?" in name + + if has_wildcards: + # Use regex for wildcard matching + pattern = re.escape(name) + pattern = pattern.replace(r"\*", ".*").replace(r"\?", ".") + 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 + 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..f611d5d --- /dev/null +++ b/widgets/policytreewidget.py @@ -0,0 +1,231 @@ +from collections import defaultdict +import logging + +from rich.text import Text +from textual.containers import Horizontal, Vertical +from textual.widget import Widget +from textual.widgets import Input, OptionList, Static, Tree +from textual.widgets.option_list import Option + +logger = logging.getLogger(__name__) + + +class PolicyTreeWidget(Widget): + """Widget for displaying and searching a hierarchical policy tree.""" + + def __init__(self, policies, devices): + super().__init__() + self.policies = policies + self.devices = devices + self.last_highlighted_node = None + self.leaf_counts = defaultdict(int) + + def compose(self): + policy_tree = Tree("", id="policy_tree") # Label set in on_mount + policy_tree.styles.width = "2fr" + policy_tree.styles.height = "100%" + + 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: + 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 _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 = {} + + # 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 + ) + + for policy in top_policies: + label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})" + node = policy_tree.root.add(label=label, data=policy) + node_map[policy.groupid] = node + + # Sort and add child policies + children_by_parent = defaultdict(list) + for policy in self.policies: + if policy.parent != "global-policy-settings": + children_by_parent[policy.parent].append(policy) + + for parent_id, children in children_by_parent.items(): + 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 + 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): + all_nodes.append(node) + for child in node.children: + self._collect_tree_nodes(child, all_nodes) + + def _remove_match_selector(self): + try: + existing = self.query("#match_selector") + for widget in existing: + if widget.is_attached: + widget.remove() + except Exception as exc: + logger.debug("Failed to remove match_selector: %s", exc) + + def on_tree_node_selected(self, message: Tree.NodeSelected) -> None: + node = message.node + data = node.data + details_pane = self.query_one("#details_pane", Static) + + if self.last_highlighted_node is not None: + original_label = str(self.last_highlighted_node.label).strip() + if isinstance(self.last_highlighted_node.label, Text): + original_label = self.last_highlighted_node.label.plain + self.last_highlighted_node.set_label(original_label) + + label_text = str(node.label).strip() + if isinstance(node.label, Text): + label_text = node.label.plain + highlighted_label = Text(label_text, style="reverse bold") + node.set_label(highlighted_label) + self.last_highlighted_node = node + + if data: + details = "\n".join( + f"{key}: {value}" for key, value in data.__dict__.items() + ) + else: + details = f"Selected: {node.label}" + details_pane.update(details) + + message.stop() + + def on_input_submitted(self, message: Input.Submitted) -> None: + self._remove_match_selector() + + query = message.value.strip().lower() + tree = self.query_one("#policy_tree", Tree) + details_pane = self.query_one("#details_pane", Static) + + all_nodes = [] + self._collect_tree_nodes(tree.root, all_nodes) + + label_to_node = {} + for node in all_nodes: + label_text = str(node.label).lower() + label_to_node[label_text] = node + if node.data: + data_dict = ( + node.data.__dict__ if hasattr(node.data, "__dict__") else node.data + ) + for key, value in data_dict.items(): + if isinstance(value, str): + label_to_node[value.lower()] = node + + matches = sorted([label for label in label_to_node if query in label]) + + if matches: + try: + option_list = self.query_one("#match_selector", OptionList) + option_list.clear_options() + option_list.display = True + except: + option_list = OptionList(id="match_selector") + details_pane.parent.mount(option_list) + + for label in matches: + option_list.add_option(Option(label, id=f"match_{label}")) + + details_pane.update(f"Found {len(matches)} matches. Select one below.") + else: + self._remove_match_selector() + details_pane.update("No matches found.") + + def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: + selected_id = event.option.id.replace("match_", "") + tree = self.query_one("#policy_tree", Tree) + details_pane = self.query_one("#details_pane", Static) + + all_nodes = [] + self._collect_tree_nodes(tree.root, all_nodes) + + label_to_node = {str(node.label).lower(): node for node in all_nodes} + match_node = label_to_node.get(selected_id.lower()) + + if match_node: + node = match_node + path = [] + while node: + path.insert(0, node) + node = node.parent + for node in path: + node.expand() + + tree.select_node(match_node) + tree.scroll_to_node(match_node) + + match_node.set_label(Text(str(match_node.label), style="reverse bold")) + details_pane.update(f"Selected: {match_node.label}") + + try: + option_list = self.query_one("#match_selector", OptionList) + option_list.remove() + except: + pass 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))