Merge remote-tracking branch 'origin/RustImplementation' into RustImplementation
Build Library / Build Library (push) Successful in 4m11s

This commit is contained in:
brotoskyj
2025-11-07 16:49:30 -05:00
8 changed files with 976 additions and 109 deletions
+3 -3
View File
@@ -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)
+1
View File
@@ -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
+41
View File
@@ -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
+114 -106
View File
@@ -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
+345
View File
@@ -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
+196
View File
@@ -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
+231
View File
@@ -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
+45
View File
@@ -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))