341 lines
12 KiB
Python
341 lines
12 KiB
Python
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published
|
|
# by the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
import difflib
|
|
from pathlib import Path
|
|
import re
|
|
from typing import List, Optional
|
|
|
|
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,
|
|
Footer,
|
|
Header,
|
|
SelectionList,
|
|
Static,
|
|
Switch,
|
|
TextArea,
|
|
)
|
|
|
|
from models.agent import Agent
|
|
|
|
|
|
class MultiAgentSelector(Widget):
|
|
"""Widget for selecting multiple agents from a list."""
|
|
|
|
class AgentsSelected(Message):
|
|
def __init__(self, selected_agents: List[Agent]):
|
|
super().__init__()
|
|
self.selected_agents = selected_agents
|
|
|
|
def __init__(self, all_agents: Optional[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):
|
|
yield Header(show_clock=True, icon="⚙")
|
|
title_text = Static("🖥️ 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_container"):
|
|
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
|
|
|
|
with Horizontal(id="action_buttons_container"):
|
|
load_file = Button("📂 Load File", id="load_file_button")
|
|
load_file.styles.margin = (1, 1, 0, 1)
|
|
yield load_file
|
|
|
|
search = Button("🔍 Search", id="search_button")
|
|
search.styles.margin = (1, 0, 0, 1)
|
|
yield search
|
|
|
|
with Horizontal() as select_buttons:
|
|
select_buttons.styles.margin = (0, 0, 0, 0)
|
|
|
|
select_none_button = Button("🚫 Select None", id="select_none")
|
|
select_none_button.styles.margin = (1, 1, 0, 1)
|
|
yield select_none_button
|
|
|
|
select_all_button = Button("✅ Select All", id="select_all")
|
|
select_all_button.styles.margin = (1, 0, 0, 1)
|
|
yield select_all_button
|
|
|
|
with Horizontal() as button_row:
|
|
button_row.styles.height = "auto"
|
|
button_row.styles.margin = (1, 0, 0, 0)
|
|
|
|
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")
|
|
yield Footer()
|
|
|
|
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 == "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 == "load_file_button":
|
|
self._load_from_file()
|
|
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
|
|
|
|
def _load_from_file(self):
|
|
"""Safely load device names from a text file."""
|
|
try:
|
|
# Import here to avoid issues if tkinter isn't available
|
|
import tkinter as tk
|
|
from tkinter import filedialog
|
|
|
|
# Create file dialog
|
|
root = tk.Tk()
|
|
root.withdraw()
|
|
|
|
file_path = filedialog.askopenfilename(
|
|
title="Select device list file",
|
|
filetypes=[
|
|
("Text files", "*.txt"),
|
|
("CSV files", "*.csv"),
|
|
("All files", "*.*"),
|
|
],
|
|
)
|
|
|
|
if not file_path:
|
|
# User cancelled
|
|
return
|
|
|
|
# Validate file path
|
|
path_obj = Path(file_path)
|
|
if not path_obj.exists():
|
|
self.app.notify("File does not exist", severity="error", timeout=3)
|
|
return
|
|
|
|
if not path_obj.is_file():
|
|
self.app.notify(
|
|
"Selected path is not a file", severity="error", timeout=3
|
|
)
|
|
return
|
|
|
|
# Check file size (limit to 1 MB for safety)
|
|
file_size = path_obj.stat().st_size
|
|
if file_size > 1_000_000: # 1 MB
|
|
self.app.notify(
|
|
f"File too large ({file_size:,} bytes). Maximum 1 MB.",
|
|
severity="error",
|
|
timeout=5,
|
|
)
|
|
return
|
|
|
|
# Read file with proper encoding to preserve emojis
|
|
try:
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
content = f.read()
|
|
except UnicodeDecodeError:
|
|
# Try with different encoding if UTF-8 fails
|
|
try:
|
|
with open(file_path, "r", encoding="latin-1") as f:
|
|
content = f.read()
|
|
self.app.notify(
|
|
"File loaded with Latin-1 encoding (UTF-8 failed)",
|
|
severity="warning",
|
|
timeout=3,
|
|
)
|
|
except Exception as e:
|
|
self.app.notify(
|
|
f"Error reading file: {str(e)}", severity="error", timeout=5
|
|
)
|
|
return
|
|
|
|
# Validate and sanitize content
|
|
lines = content.split("\n")
|
|
valid_lines = []
|
|
invalid_count = 0
|
|
|
|
# Pattern for valid hostnames/device names
|
|
# Allows: letters, numbers, hyphens, underscores, periods, and Unicode chars
|
|
hostname_pattern = re.compile(r"^[\w\-\.\u0080-\uFFFF]+$", re.UNICODE)
|
|
|
|
for line in lines:
|
|
line = line.strip()
|
|
if not line:
|
|
continue # Skip empty lines
|
|
|
|
# Check if line looks like a valid hostname/device name
|
|
if hostname_pattern.match(line):
|
|
valid_lines.append(line)
|
|
else:
|
|
invalid_count += 1
|
|
# Log but don't add invalid entries
|
|
|
|
if not valid_lines:
|
|
self.app.notify(
|
|
"No valid device names found in file", severity="warning", timeout=3
|
|
)
|
|
return
|
|
|
|
# Update text area with validated content
|
|
text_area = self.query_one("#device_input", TextArea)
|
|
text_area.text = "\n".join(valid_lines)
|
|
|
|
# Show notification
|
|
msg = f"✅ Loaded {len(valid_lines)} devices from file"
|
|
if invalid_count > 0:
|
|
msg += f" ({invalid_count} invalid entries skipped)"
|
|
|
|
self.app.notify(msg, severity="information", timeout=5)
|
|
|
|
except ImportError:
|
|
self.app.notify(
|
|
"tkinter not available - cannot open file dialog",
|
|
severity="error",
|
|
timeout=3,
|
|
)
|
|
except Exception as e:
|
|
self.app.notify(
|
|
f"Error loading file: {str(e)}", severity="error", timeout=5
|
|
)
|