Major step towards unification of the UI Implementation of the Back Feature, splitting of TUI files back into subfolders
This commit is contained in:
@@ -0,0 +1,668 @@
|
||||
# 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/>.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
import pandas as pd
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import (
|
||||
Button,
|
||||
DataTable,
|
||||
Footer,
|
||||
Header,
|
||||
Static,
|
||||
TextArea,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AllowlistSelectionWidget(Static):
|
||||
"""
|
||||
Widget for selecting an allowlist and adding hashes to it.
|
||||
Can be reused in different workflows.
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
AllowlistSelectionWidget {
|
||||
height: 1fr;
|
||||
layout: vertical;
|
||||
}
|
||||
#allowlist_main {
|
||||
height: 1fr;
|
||||
width: 100%;
|
||||
}
|
||||
#left_panel {
|
||||
width: 50%;
|
||||
padding: 1;
|
||||
border: solid $primary;
|
||||
}
|
||||
#right_panel {
|
||||
width: 50%;
|
||||
padding: 1;
|
||||
border: solid $primary;
|
||||
}
|
||||
#allowlist_table {
|
||||
height: 70%;
|
||||
margin: 1 0;
|
||||
}
|
||||
#allowlist_table > .datatable--header {
|
||||
text-style: bold;
|
||||
background: $boost;
|
||||
}
|
||||
#allowlist_table Row {
|
||||
height: 1;
|
||||
}
|
||||
#preview_area {
|
||||
height: 60%;
|
||||
margin: 1 0;
|
||||
}
|
||||
#action_buttons {
|
||||
height: auto;
|
||||
min-height: 3;
|
||||
padding: 1;
|
||||
content-align: center middle;
|
||||
}
|
||||
.panel-title {
|
||||
text-style: bold;
|
||||
margin: 0 0 1 0;
|
||||
}
|
||||
.info-text {
|
||||
margin: 1 0;
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
selected_data: pd.DataFrame,
|
||||
api=None,
|
||||
hostname: Optional[str] = None,
|
||||
otpid: Optional[str] = None,
|
||||
hash_column: str = "sha256", # Default hash column name
|
||||
):
|
||||
"""
|
||||
Initialize the allowlist selection widget.
|
||||
|
||||
Args:
|
||||
selected_data: DataFrame containing the selected activities
|
||||
api: API instance for making allowlist calls
|
||||
hostname: Optional hostname for context
|
||||
otpid: Optional OTP ID for context
|
||||
hash_column: Name of the column containing hashes (default: "sha256")
|
||||
"""
|
||||
super().__init__()
|
||||
self.selected_data = selected_data
|
||||
self.api = api
|
||||
self.hostname = hostname
|
||||
self.otpid = otpid
|
||||
self.hash_column = hash_column
|
||||
self.allowlists = []
|
||||
self.selected_allowlist = None
|
||||
self.hashes_to_add = []
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Horizontal(id="allowlist_main"):
|
||||
# Left panel - Allowlist selection
|
||||
with Vertical(id="left_panel"):
|
||||
yield Static("Select Allowlist", classes="panel-title")
|
||||
yield Static(
|
||||
f"Choose an allowlist to add {len(self.selected_data)} selected items",
|
||||
classes="info-text",
|
||||
)
|
||||
|
||||
# Allowlist table
|
||||
self.allowlist_table = DataTable(id="allowlist_table")
|
||||
self.allowlist_table.cursor_type = "row"
|
||||
yield self.allowlist_table
|
||||
|
||||
# Refresh button
|
||||
self.refresh_btn = Button(
|
||||
"🔄 Refresh Allowlists", id="refresh_allowlists_btn"
|
||||
)
|
||||
yield self.refresh_btn
|
||||
|
||||
# Right panel - Preview and actions
|
||||
with Vertical(id="right_panel"):
|
||||
yield Static("Preview", classes="panel-title")
|
||||
|
||||
# Context information
|
||||
context_text = []
|
||||
if self.hostname:
|
||||
context_text.append(f"Host: {self.hostname}")
|
||||
if self.otpid:
|
||||
context_text.append(f"OTP: {self.otpid}")
|
||||
context_text.append(f"Selected Activities: {len(self.selected_data)}")
|
||||
|
||||
yield Static(" | ".join(context_text), classes="info-text")
|
||||
|
||||
# Preview text area
|
||||
self.preview_area = TextArea(
|
||||
id="preview_area", read_only=True, language="markdown"
|
||||
)
|
||||
yield self.preview_area
|
||||
|
||||
# Hash statistics
|
||||
self.stats_label = Static("", id="stats_label", classes="info-text")
|
||||
yield self.stats_label
|
||||
|
||||
# Action buttons at bottom
|
||||
with Horizontal(id="action_buttons"):
|
||||
self.add_btn = Button("➕ Add to Allowlist", id="add_to_allowlist_btn")
|
||||
|
||||
self.add_btn.styles.width = "100%"
|
||||
self.add_btn.disabled = True # Disabled until allowlist selected
|
||||
|
||||
yield self.add_btn
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
"""Load allowlists when widget mounts."""
|
||||
await self.load_allowlists()
|
||||
await self.extract_and_preview_hashes()
|
||||
|
||||
async def load_allowlists(self) -> None:
|
||||
"""Load available allowlists from API, grouped by policy association."""
|
||||
if not self.api:
|
||||
logger.error("No API available")
|
||||
self.allowlist_table.add_column("Error")
|
||||
self.allowlist_table.add_row("No API available")
|
||||
return
|
||||
|
||||
try:
|
||||
# First, try to get the host's policy if hostname is provided
|
||||
host_policy_allowlists = []
|
||||
host_policy_ids = set()
|
||||
policy_name = "Unknown Policy" # Default value
|
||||
group_id = None
|
||||
|
||||
if self.hostname:
|
||||
try:
|
||||
# Get agent info to find its policy
|
||||
agents_df = self.api.agent_find_by_hostname(self.hostname)
|
||||
if not agents_df.empty:
|
||||
# Get the policy group ID for this host
|
||||
group_id = agents_df.iloc[0].get("groupid")
|
||||
|
||||
# Look up the policy name from app's cached policies
|
||||
if (
|
||||
group_id
|
||||
and hasattr(self.app, "policies")
|
||||
and self.app.policies
|
||||
):
|
||||
for policy in self.app.policies:
|
||||
if policy.groupid == group_id:
|
||||
policy_name = policy.name
|
||||
logger.info(
|
||||
f"Found policy name: '{policy_name}' for group_id: {group_id}"
|
||||
)
|
||||
break
|
||||
|
||||
logger.info(
|
||||
f"Found host '{self.hostname}' in policy '{policy_name}' (group_id: {group_id})"
|
||||
)
|
||||
|
||||
if group_id:
|
||||
# Get allowlists for this policy
|
||||
policy_allowlists_df = self.api.policy_list_allowlists(
|
||||
group_id
|
||||
)
|
||||
if not policy_allowlists_df.empty:
|
||||
host_policy_allowlists = policy_allowlists_df.to_dict(
|
||||
orient="records"
|
||||
)
|
||||
host_policy_ids = {
|
||||
al.get("applicationid")
|
||||
for al in host_policy_allowlists
|
||||
}
|
||||
logger.info(
|
||||
f"Found {len(host_policy_allowlists)} allowlists for host's policy"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not get host's policy allowlists: {e}")
|
||||
|
||||
# If we still don't have a policy name, try to get it from the first allowlist or use a default
|
||||
if not policy_name:
|
||||
# Get all policies and try to find which one has allowlists
|
||||
try:
|
||||
all_policies_df = self.api.policy_find_all()
|
||||
if not all_policies_df.empty:
|
||||
# If we have a group_id from somewhere, use it
|
||||
if group_id:
|
||||
policy_row = all_policies_df[
|
||||
all_policies_df["groupid"] == group_id
|
||||
]
|
||||
if not policy_row.empty:
|
||||
policy_name = policy_row.iloc[0].get(
|
||||
"groupname", "Unknown Policy"
|
||||
)
|
||||
else:
|
||||
# Use the first policy as fallback
|
||||
policy_name = all_policies_df.iloc[0].get(
|
||||
"groupname", "Default Policy"
|
||||
)
|
||||
logger.info(f"Using first available policy: {policy_name}")
|
||||
else:
|
||||
policy_name = "Unknown Policy"
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not fetch policies: {e}")
|
||||
policy_name = "Unknown Policy"
|
||||
|
||||
# Get all allowlists
|
||||
all_allowlists_df = self.api.allowlist_find_all()
|
||||
|
||||
if all_allowlists_df.empty:
|
||||
self.allowlist_table.add_column("No Allowlists")
|
||||
self.allowlist_table.add_row("No allowlists found")
|
||||
return
|
||||
|
||||
all_allowlists = all_allowlists_df.to_dict(orient="records")
|
||||
|
||||
# Separate into two groups: policy-associated and others
|
||||
other_allowlists = [
|
||||
al
|
||||
for al in all_allowlists
|
||||
if al.get("applicationid") not in host_policy_ids
|
||||
]
|
||||
|
||||
# Sort each group alphabetically by name
|
||||
host_policy_allowlists.sort(key=lambda x: x.get("name", "").lower())
|
||||
other_allowlists.sort(key=lambda x: x.get("name", "").lower())
|
||||
|
||||
# Combine lists with policy-associated first
|
||||
self.allowlists = host_policy_allowlists + other_allowlists
|
||||
|
||||
# Setup table columns
|
||||
self.allowlist_table.clear()
|
||||
self.allowlist_table.add_columns("Name", "Application ID", "Type")
|
||||
|
||||
# Track which rows are headers vs actual allowlists
|
||||
self._row_to_allowlist_map = {}
|
||||
current_row = 0
|
||||
|
||||
# Add policy-associated allowlists if any
|
||||
if host_policy_allowlists:
|
||||
# Add section header
|
||||
header_text = f"=== Policy: {policy_name or 'Host Policy'} ==="
|
||||
self.allowlist_table.add_row(header_text, "", "", key="header_policy")
|
||||
current_row += 1
|
||||
|
||||
# Add policy allowlists
|
||||
for idx, allowlist in enumerate(host_policy_allowlists):
|
||||
name = allowlist.get("name", "Unknown")
|
||||
app_id = allowlist.get("applicationid", "Unknown")
|
||||
|
||||
self.allowlist_table.add_row(
|
||||
f" {name}", # Indent to show grouping
|
||||
app_id,
|
||||
"Policy",
|
||||
key=f"policy_{idx}",
|
||||
)
|
||||
self._row_to_allowlist_map[current_row] = idx
|
||||
current_row += 1
|
||||
|
||||
# Add other allowlists
|
||||
if other_allowlists:
|
||||
# Add section header
|
||||
if host_policy_allowlists:
|
||||
# Add spacer if we have policy allowlists above
|
||||
self.allowlist_table.add_row("", "", "", key="spacer")
|
||||
current_row += 1
|
||||
|
||||
self.allowlist_table.add_row(
|
||||
"=== Other Available Allowlists ===", "", "", key="header_other"
|
||||
)
|
||||
current_row += 1
|
||||
|
||||
# Add other allowlists
|
||||
for idx, allowlist in enumerate(other_allowlists):
|
||||
name = allowlist.get("name", "Unknown")
|
||||
app_id = allowlist.get("applicationid", "Unknown")
|
||||
|
||||
self.allowlist_table.add_row(
|
||||
f" {name}", # Indent to show grouping
|
||||
app_id,
|
||||
"General",
|
||||
key=f"other_{idx}",
|
||||
)
|
||||
# Map to the correct index in the combined list
|
||||
actual_idx = len(host_policy_allowlists) + idx
|
||||
self._row_to_allowlist_map[current_row] = actual_idx
|
||||
current_row += 1
|
||||
|
||||
# Log summary
|
||||
logger.info(
|
||||
f"Loaded {len(self.allowlists)} total allowlists: "
|
||||
f"{len(host_policy_allowlists)} policy-associated, "
|
||||
f"{len(other_allowlists)} others"
|
||||
)
|
||||
|
||||
# Update stats label if no allowlists in policy
|
||||
if self.hostname and not host_policy_allowlists:
|
||||
self.stats_label.update(
|
||||
f"Note: No allowlists found for {self.hostname}'s policy | "
|
||||
+ self.stats_label.content.plain
|
||||
)
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception(f"Failed to load allowlists: {exc}")
|
||||
self.allowlist_table.add_column("Error")
|
||||
self.allowlist_table.add_row(f"Failed to load: {str(exc)}")
|
||||
|
||||
async def extract_and_preview_hashes(self) -> None:
|
||||
"""Extract hashes from selected data and show preview."""
|
||||
preview_lines = ["## Hash Extraction Summary\n"]
|
||||
|
||||
# Check for hash column
|
||||
if self.hash_column not in self.selected_data.columns:
|
||||
# Try to find a hash column
|
||||
possible_hash_cols = [
|
||||
"sha256",
|
||||
"SHA256",
|
||||
"hash",
|
||||
"Hash",
|
||||
"sha1",
|
||||
"SHA1",
|
||||
"md5",
|
||||
"MD5",
|
||||
"filehash",
|
||||
"file_hash",
|
||||
]
|
||||
found_col = None
|
||||
for col in possible_hash_cols:
|
||||
if col in self.selected_data.columns:
|
||||
found_col = col
|
||||
break
|
||||
|
||||
if found_col:
|
||||
self.hash_column = found_col
|
||||
preview_lines.append(f"✅ Found hash column: **{found_col}**\n")
|
||||
else:
|
||||
preview_lines.append("❌ **No hash column found**\n")
|
||||
preview_lines.append("Available columns:\n")
|
||||
for col in self.selected_data.columns:
|
||||
if col != "_row_id":
|
||||
preview_lines.append(f" - {col}\n")
|
||||
|
||||
self.preview_area.text = "".join(preview_lines)
|
||||
self.stats_label.update("No hashes to add")
|
||||
return
|
||||
|
||||
# Extract unique hashes
|
||||
hashes = self.selected_data[self.hash_column].dropna().unique()
|
||||
self.hashes_to_add = [h for h in hashes if h and str(h).strip()]
|
||||
|
||||
# Build preview
|
||||
preview_lines.append(f"### Found {len(self.hashes_to_add)} unique hashes\n\n")
|
||||
|
||||
# Show sample of hashes (first 10)
|
||||
preview_lines.append("**Sample hashes to be added:**\n```\n")
|
||||
for i, hash_val in enumerate(self.hashes_to_add[:10]):
|
||||
preview_lines.append(f"{i+1}. {hash_val}\n")
|
||||
if len(self.hashes_to_add) > 10:
|
||||
preview_lines.append(f"... and {len(self.hashes_to_add) - 10} more\n")
|
||||
preview_lines.append("```\n\n")
|
||||
|
||||
# Show sample of source data
|
||||
preview_lines.append("**Sample source activities:**\n")
|
||||
sample_cols = [
|
||||
col
|
||||
for col in self.selected_data.columns
|
||||
if col not in ["_row_id"] and col in ["filename", "path", "action", "user"]
|
||||
]
|
||||
if not sample_cols:
|
||||
sample_cols = [
|
||||
col for col in self.selected_data.columns if col != "_row_id"
|
||||
][:3]
|
||||
|
||||
if sample_cols:
|
||||
preview_lines.append("```\n")
|
||||
for i, row in self.selected_data[sample_cols].head(5).iterrows():
|
||||
row_text = " | ".join([f"{col}: {row[col]}" for col in sample_cols])
|
||||
preview_lines.append(f"{row_text}\n")
|
||||
preview_lines.append("```\n")
|
||||
|
||||
self.preview_area.text = "".join(preview_lines)
|
||||
|
||||
# Update statistics
|
||||
self.stats_label.update(
|
||||
f"Ready to add {len(self.hashes_to_add)} unique hashes | "
|
||||
f"From {len(self.selected_data)} selected activities"
|
||||
)
|
||||
|
||||
async def on_data_table_row_selected(self, event) -> None:
|
||||
"""Handle allowlist selection."""
|
||||
try:
|
||||
# Extract row index from event - handle different event structures
|
||||
row_index = None
|
||||
|
||||
# Try to get row index from coordinate
|
||||
if hasattr(event, "coordinate") and hasattr(event.coordinate, "row"):
|
||||
row_index = event.coordinate.row
|
||||
# Try cursor_row as fallback
|
||||
elif hasattr(event, "cursor_row"):
|
||||
row_index = event.cursor_row
|
||||
# Try getting from the table itself
|
||||
else:
|
||||
table = self.allowlist_table
|
||||
if hasattr(table, "cursor_row"):
|
||||
row_index = table.cursor_row
|
||||
|
||||
# Validate row index
|
||||
if row_index is not None and isinstance(row_index, int):
|
||||
# Account for group headers in the row count
|
||||
actual_allowlist_index = self._get_allowlist_index_from_row(row_index)
|
||||
|
||||
if (
|
||||
actual_allowlist_index is not None
|
||||
and 0 <= actual_allowlist_index < len(self.allowlists)
|
||||
):
|
||||
self.selected_allowlist = self.allowlists[actual_allowlist_index]
|
||||
self.add_btn.disabled = False
|
||||
self.add_btn.label = (
|
||||
f"➕ Add to '{self.selected_allowlist.get('name', 'Unknown')}'"
|
||||
)
|
||||
|
||||
# Update preview with selection
|
||||
await self._update_preview_with_selection()
|
||||
|
||||
logger.info(
|
||||
f"Selected allowlist: {self.selected_allowlist.get('name')}"
|
||||
)
|
||||
else:
|
||||
logger.debug(f"Row {row_index} is a header or invalid")
|
||||
else:
|
||||
logger.warning(f"Could not extract valid row index from event: {event}")
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception(f"Failed to select allowlist: {exc}")
|
||||
|
||||
def _get_allowlist_index_from_row(self, row_index: int) -> Optional[int]:
|
||||
"""Convert table row index to allowlist list index, accounting for group headers."""
|
||||
# This will be updated when we have group headers
|
||||
if hasattr(self, "_row_to_allowlist_map"):
|
||||
return self._row_to_allowlist_map.get(row_index)
|
||||
return row_index
|
||||
|
||||
async def _update_preview_with_selection(self) -> None:
|
||||
"""Update preview when an allowlist is selected."""
|
||||
if not self.selected_allowlist:
|
||||
return
|
||||
|
||||
current_text = self.preview_area.text
|
||||
# Remove any existing selection header
|
||||
if "### Selected Allowlist:" in current_text:
|
||||
lines = current_text.split("\n")
|
||||
# Find and remove the selection lines
|
||||
new_lines = []
|
||||
skip_next = False
|
||||
for line in lines:
|
||||
if line.startswith("### Selected Allowlist:"):
|
||||
skip_next = True
|
||||
continue
|
||||
if skip_next and line.startswith("Application ID:"):
|
||||
skip_next = False
|
||||
continue
|
||||
if not skip_next:
|
||||
new_lines.append(line)
|
||||
current_text = "\n".join(new_lines)
|
||||
|
||||
# Add new selection at the top
|
||||
selection_text = (
|
||||
f"### Selected Allowlist: **{self.selected_allowlist.get('name')}**\n"
|
||||
f"Application ID: {self.selected_allowlist.get('applicationid')}\n\n"
|
||||
)
|
||||
self.preview_area.text = selection_text + current_text
|
||||
|
||||
async def on_button_pressed(self, event) -> None:
|
||||
"""Handle button presses."""
|
||||
btn = getattr(event, "button", None) or getattr(event, "sender", None)
|
||||
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None)
|
||||
|
||||
if btn is self.refresh_btn or btn_id == "refresh_allowlists_btn":
|
||||
await self.load_allowlists()
|
||||
event.stop()
|
||||
return
|
||||
|
||||
if btn is self.add_btn or btn_id == "add_to_allowlist_btn":
|
||||
await self.add_hashes_to_allowlist()
|
||||
event.stop()
|
||||
return
|
||||
|
||||
async def add_hashes_to_allowlist(self) -> None:
|
||||
"""Add the extracted hashes to the selected allowlist."""
|
||||
if not self.selected_allowlist or not self.hashes_to_add:
|
||||
self.app.notify(
|
||||
"No allowlist selected or no hashes to add", severity="warning"
|
||||
)
|
||||
return
|
||||
|
||||
if not self.api:
|
||||
self.app.notify("API not available", severity="error")
|
||||
return
|
||||
|
||||
try:
|
||||
# Disable button during operation
|
||||
self.add_btn.disabled = True
|
||||
self.add_btn.label = "Adding hashes..."
|
||||
|
||||
# Call API to add hashes
|
||||
app_id = self.selected_allowlist.get("applicationid")
|
||||
allowlist_name = self.selected_allowlist.get("name", "Unknown")
|
||||
|
||||
logger.info(
|
||||
f"Adding {len(self.hashes_to_add)} hashes to allowlist {allowlist_name} (ID: {app_id})"
|
||||
)
|
||||
|
||||
result = self.api.hash_add_to_allowlist(app_id, self.hashes_to_add)
|
||||
logger.debug(f"Hash adding api call: {result}")
|
||||
# Success notification
|
||||
self.app.notify(
|
||||
f"Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'",
|
||||
title="Success",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
# Update preview to show success
|
||||
self.preview_area.text = (
|
||||
f"## SUCCESS\n\n"
|
||||
f"Added **{len(self.hashes_to_add)} hashes** to allowlist:\n"
|
||||
f"**{allowlist_name}** (ID: {app_id})\n\n"
|
||||
f"### Operation Details:\n"
|
||||
f"- Source: {self.hostname or 'Multiple hosts'}\n"
|
||||
f"- OTP ID: {self.otpid or 'N/A'}\n"
|
||||
f"- Activities processed: {len(self.selected_data)}\n"
|
||||
f"- Unique hashes added: {len(self.hashes_to_add)}\n"
|
||||
)
|
||||
|
||||
# Change button to "Done"
|
||||
self.add_btn.label = "Done - Press q to return to main menu"
|
||||
self.add_btn.disabled = True
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception(f"Failed to add hashes to allowlist: {exc}")
|
||||
self.app.notify(
|
||||
f"Failed to add hashes: {str(exc)}",
|
||||
title="Error",
|
||||
severity="error",
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
# Re-enable button
|
||||
self.add_btn.disabled = False
|
||||
self.add_btn.label = "Retry Add to Allowlist"
|
||||
|
||||
|
||||
class AllowlistSelectionScreen(Screen):
|
||||
"""
|
||||
Screen wrapper for the AllowlistSelectionWidget.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("r", "refresh", "Refresh Allowlists"),
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
selected_data: pd.DataFrame,
|
||||
api=None,
|
||||
hostname: Optional[str] = None,
|
||||
otpid: Optional[str] = None,
|
||||
hash_column: str = "sha256",
|
||||
):
|
||||
super().__init__()
|
||||
self.selected_data = selected_data
|
||||
self.api = api
|
||||
self.hostname = hostname
|
||||
self.otpid = otpid
|
||||
self.hash_column = hash_column
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header(show_clock=True)
|
||||
self.widget = AllowlistSelectionWidget(
|
||||
self.selected_data,
|
||||
api=self.api,
|
||||
hostname=self.hostname,
|
||||
otpid=self.otpid,
|
||||
hash_column=self.hash_column,
|
||||
)
|
||||
yield self.widget
|
||||
yield Footer()
|
||||
|
||||
async def action_go_back(self) -> None:
|
||||
"""Go back to previous screen."""
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_refresh(self) -> None:
|
||||
"""Refresh the allowlists."""
|
||||
if hasattr(self, "widget") and self.widget:
|
||||
await self.widget.load_allowlists()
|
||||
|
||||
async def action_confirm(self) -> None:
|
||||
"""Confirm and add to allowlist."""
|
||||
if hasattr(self, "widget") and self.widget:
|
||||
if self.widget.selected_allowlist and self.widget.hashes_to_add:
|
||||
await self.widget.add_hashes_to_allowlist()
|
||||
@@ -0,0 +1,114 @@
|
||||
# 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/>.
|
||||
from typing import List, Optional
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.css.query import NoMatches
|
||||
from textual.screen import Screen
|
||||
|
||||
from models.agent import Agent
|
||||
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
|
||||
from TUI.Widgets.multiagentselector import MultiAgentSelector
|
||||
from TUI.Widgets.resultsdisplay import ResultsDisplay
|
||||
|
||||
|
||||
class MoveAgentWorkflowScreen(Screen):
|
||||
"""Screen that handles the agent movement workflow."""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
def __init__(self, all_agents: Optional[List[Agent]]):
|
||||
super().__init__()
|
||||
self.all_agents = all_agents
|
||||
self.selected_agents = None
|
||||
self.workflow_stage = "select_agents" # Track current stage
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Start with the multi-agent selector."""
|
||||
yield MultiAgentSelector(self.all_agents)
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle escape key to go back one step within the workflow."""
|
||||
if self.workflow_stage == "select_agents":
|
||||
# At first stage, go back to main menu
|
||||
self.app.pop_screen()
|
||||
elif self.workflow_stage == "operations":
|
||||
# Go back to agent selection
|
||||
try:
|
||||
ops_widget = self.query_one(AgentMoveOperations)
|
||||
ops_widget.remove()
|
||||
except NoMatches:
|
||||
pass
|
||||
self.mount(MultiAgentSelector(self.all_agents))
|
||||
self.workflow_stage = "select_agents"
|
||||
elif self.workflow_stage == "results":
|
||||
# Go back to operations
|
||||
try:
|
||||
results_widget = self.query_one(ResultsDisplay)
|
||||
results_widget.remove()
|
||||
except NoMatches:
|
||||
pass
|
||||
self.mount(AgentMoveOperations(self.selected_agents))
|
||||
self.workflow_stage = "operations"
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Handle q key to go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_multi_agent_selector_agents_selected(
|
||||
self, message: MultiAgentSelector.AgentsSelected
|
||||
) -> None:
|
||||
"""Handle selected agents - switch to operations screen."""
|
||||
self.selected_agents = message.selected_agents
|
||||
|
||||
# Remove the MultiAgentSelector
|
||||
selector = self.query_one(MultiAgentSelector)
|
||||
selector.remove()
|
||||
|
||||
# Mount the AgentMoveOperations with the selected Agent objects
|
||||
self.mount(AgentMoveOperations(self.selected_agents))
|
||||
self.workflow_stage = "operations"
|
||||
|
||||
def on_agent_move_operations_operation_complete(
|
||||
self, message: AgentMoveOperations.OperationComplete
|
||||
) -> None:
|
||||
"""Handle completion of move operation - transition to results screen."""
|
||||
# Format successful results
|
||||
success_lines = []
|
||||
for agent, result in message.successful:
|
||||
success_lines.append(f"✔ {agent.hostname}")
|
||||
|
||||
# Format unsuccessful results
|
||||
failure_lines = []
|
||||
for agent, error in message.unsuccessful:
|
||||
failure_lines.append(f"❌ — {agent.hostname}: {error}")
|
||||
|
||||
successful_text = "\n".join(success_lines) if success_lines else "(none)"
|
||||
unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)"
|
||||
|
||||
# Remove the operations widget
|
||||
ops_widget = self.query_one(AgentMoveOperations)
|
||||
ops_widget.remove()
|
||||
|
||||
# Mount the results display
|
||||
self.mount(
|
||||
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
|
||||
)
|
||||
self.workflow_stage = "results"
|
||||
@@ -0,0 +1,901 @@
|
||||
# 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/>.
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
import logging
|
||||
import os
|
||||
|
||||
import pandas as pd
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Static
|
||||
|
||||
from TUI.Screens.allowlistselectionscreen import AllowlistSelectionScreen
|
||||
from utils.configmanager import load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _load_working_dir() -> str:
|
||||
"""
|
||||
Load the working directory from environment variables or use the current working directory.
|
||||
"""
|
||||
wd = os.environ.get("WORKING_DIR")
|
||||
if wd:
|
||||
return wd
|
||||
return os.getcwd()
|
||||
|
||||
|
||||
class OTPActivitiesWidget(Static):
|
||||
"""
|
||||
Reusable widget that contains the sessions table (left) and an Activity Preview (right).
|
||||
The right side shows an Activity Preview that takes ~75% vertical space, and a lower area
|
||||
with Continue button.
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
OTPActivitiesWidget {
|
||||
height: 1fr;
|
||||
}
|
||||
#main_row {
|
||||
width: 100%;
|
||||
height: 100%;
|
||||
layout: horizontal;
|
||||
}
|
||||
#left_panel {
|
||||
width: 60%;
|
||||
min-width: 60;
|
||||
border: none;
|
||||
}
|
||||
#right_panel {
|
||||
width: 40%;
|
||||
min-width: 40;
|
||||
border: none;
|
||||
layout: vertical;
|
||||
}
|
||||
#activity_preview_container {
|
||||
height: 1fr;
|
||||
border: none;
|
||||
padding: 1 1;
|
||||
}
|
||||
#activity_buttons {
|
||||
height: auto;
|
||||
min-height: 3;
|
||||
padding: 1 1;
|
||||
content-align: center middle;
|
||||
}
|
||||
"""
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
# Layout: horizontal main row with left & right panels
|
||||
with Horizontal(id="main_row"):
|
||||
# Left: sessions area
|
||||
with Vertical(id="left_panel"):
|
||||
yield Static("OTP Sessions", classes="panel-title")
|
||||
with Vertical(id="sessions_table_container"):
|
||||
self.sessions_table = DataTable(id="sessions_table")
|
||||
self.sessions_table.styles.width = "100%"
|
||||
yield self.sessions_table
|
||||
# Right: Activity Preview (top 3/4) + buttons (bottom 1/4)
|
||||
with Vertical(id="right_panel"):
|
||||
# Activity preview area (takes ~75% of right panel)
|
||||
yield Static("Activity Preview", classes="panel-title")
|
||||
with Vertical(id="activity_preview_container"):
|
||||
self.activities_table = DataTable(id="activity_preview_table")
|
||||
yield self.activities_table
|
||||
# Button area at the bottom (Continue)
|
||||
with Horizontal(id="activity_buttons"):
|
||||
self.continue_btn = Button("Continue", id="activity_continue_btn")
|
||||
self.continue_btn.styles.width = "100%"
|
||||
yield self.continue_btn
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
# Configure sessions table and activities preview
|
||||
self.sessions_table.clear()
|
||||
self.sessions_table.add_columns(
|
||||
"otpid", "hostname", "status", "purpose", "granted"
|
||||
)
|
||||
self.activities_table.clear()
|
||||
# activities_table columns are dynamically added when activities are loaded.
|
||||
# Selection behavior
|
||||
self.sessions_table.cursor_type = "row"
|
||||
try:
|
||||
self.sessions_table.zebra_stripes = True
|
||||
except Exception:
|
||||
pass
|
||||
self.activities_table.cursor_type = "row"
|
||||
try:
|
||||
self.activities_table.zebra_stripes = True
|
||||
except Exception:
|
||||
pass
|
||||
# Store state
|
||||
self._sessions_df: pd.DataFrame | None = None
|
||||
self._activities_df: pd.DataFrame | None = None
|
||||
self._selected_session_otpid: str | int | None = None
|
||||
|
||||
async def on_button_pressed(self, event) -> None: # type: ignore[override]
|
||||
"""
|
||||
Handle Continue button for the Activity Preview area.
|
||||
"""
|
||||
# Try to resolve the button object from the event
|
||||
btn = (
|
||||
getattr(event, "button", None)
|
||||
or getattr(event, "sender", None)
|
||||
or getattr(event, "control", None)
|
||||
or getattr(event, "widget", None)
|
||||
)
|
||||
btn_id = (
|
||||
getattr(btn, "id", None)
|
||||
or getattr(event, "button_id", None)
|
||||
or getattr(event, "id", None)
|
||||
)
|
||||
# ---- Continue ----
|
||||
if btn is self.continue_btn or btn_id == getattr(self.continue_btn, "id", None):
|
||||
if self._activities_df is None or self._activities_df.empty:
|
||||
logger.info("Continue pressed but no activities loaded.")
|
||||
self.app.notify(
|
||||
"No activities loaded to continue with.", severity="warning"
|
||||
)
|
||||
return
|
||||
# Copy activities DataFrame to pass to new screen
|
||||
activities_copy = self._activities_df.copy()
|
||||
otpid = self._selected_session_otpid
|
||||
# Optionally include hostname if available
|
||||
hostname = None
|
||||
try:
|
||||
if self._sessions_df is not None:
|
||||
df = self._sessions_df.reset_index(drop=True)
|
||||
match = df[df["otpid"] == otpid]
|
||||
if not match.empty:
|
||||
hostname = match.iloc[0].get("hostname")
|
||||
except Exception:
|
||||
hostname = None
|
||||
# Create and push ActivityDetailScreen, handing the data
|
||||
try:
|
||||
detail_screen = ActivityDetailScreen(
|
||||
activities_copy, otpid=otpid, hostname=hostname
|
||||
)
|
||||
await self.app.push_screen(detail_screen)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to push ActivityDetailScreen: %s", exc)
|
||||
return
|
||||
# Unknown button on widget
|
||||
logger.debug(
|
||||
"Unhandled OTPActivitiesWidget button pressed (resolved btn=%r, id=%r)",
|
||||
btn,
|
||||
btn_id,
|
||||
)
|
||||
|
||||
async def on_data_table_row_selected(self, event) -> None: # type: ignore[override]
|
||||
"""
|
||||
Robust handler for DataTable row-selection across Textual micro-versions.
|
||||
Tries many attribute names and shapes:
|
||||
- numeric index (row_key, row_index, index)
|
||||
- coordinate object or tuple (coordinate.row or (row, col))
|
||||
- direct row values (row, values, cells) -> we try to map those back to the sessions DF
|
||||
- table.cursor_row fallback
|
||||
"""
|
||||
# 1) Determine the sending table (best-effort)
|
||||
sender = None
|
||||
for attr in ("sender", "table", "data_table", "control"):
|
||||
sender = getattr(event, attr, None)
|
||||
if sender is not None:
|
||||
break
|
||||
if sender is None:
|
||||
sender = self.sessions_table # Assume sessions_table if unknown
|
||||
# Only respond to selections in the sessions table
|
||||
if sender is not self.sessions_table:
|
||||
return
|
||||
|
||||
# Helper to log and return
|
||||
def _bad(msg: str, *args):
|
||||
logger.warning(msg, *args)
|
||||
return None
|
||||
|
||||
# 2) Try to extract a numeric index
|
||||
row_key = None
|
||||
for attr in ("row_key", "row", "row_index", "index"):
|
||||
row_key = getattr(event, attr, None)
|
||||
if row_key is not None:
|
||||
break
|
||||
# If coordinate: try to extract .row or tuple[0]
|
||||
if row_key is None:
|
||||
coord = getattr(event, "coordinate", None) or getattr(
|
||||
event, "cursor_coordinate", None
|
||||
)
|
||||
if coord is not None:
|
||||
if hasattr(coord, "row"):
|
||||
row_key = coord.row
|
||||
elif isinstance(coord, (tuple, list)) and len(coord) >= 1:
|
||||
row_key = coord[0]
|
||||
# If still nothing, maybe the event provides the row's cell values directly
|
||||
row_values = None
|
||||
for attr in ("values", "cells", "row", "row_values", "selected_row_values"):
|
||||
val = getattr(event, attr, None)
|
||||
if val:
|
||||
# Prefer actual sequence of cell values
|
||||
row_values = val
|
||||
break
|
||||
# If we have row_values, try to map them back to the sessions DataFrame
|
||||
if row_values is not None:
|
||||
# Normalize into list of strings for comparison
|
||||
try:
|
||||
vals = [
|
||||
"" if pd.isna(v) else str(v)
|
||||
for v in (
|
||||
list(row_values)
|
||||
if not isinstance(row_values, str)
|
||||
else [row_values]
|
||||
)
|
||||
]
|
||||
except Exception:
|
||||
vals = [str(row_values)]
|
||||
# Try to match against the expected columns order we render
|
||||
if self._sessions_df is None or self._sessions_df.empty:
|
||||
logger.warning(
|
||||
"Sessions DataFrame is empty; cannot map selected row values."
|
||||
)
|
||||
return
|
||||
df_ordered = self._sessions_df.reset_index(drop=True)
|
||||
expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]
|
||||
|
||||
# Build stringified candidates for each row in df using the same columns we show
|
||||
def _row_to_vals(sr):
|
||||
out = []
|
||||
for c in expected_cols:
|
||||
if c in sr:
|
||||
v = sr[c]
|
||||
out.append("" if pd.isna(v) else str(v))
|
||||
else:
|
||||
out.append("")
|
||||
return out
|
||||
|
||||
match_idx = None
|
||||
for i, sr in df_ordered.iterrows():
|
||||
cand = _row_to_vals(sr)
|
||||
# Compare prefix: row values might be a subset (e.g. only first 3 cols), so compare prefix only
|
||||
if len(vals) <= len(cand) and all(
|
||||
vals[j] == cand[j] for j in range(len(vals))
|
||||
):
|
||||
match_idx = i
|
||||
break
|
||||
if match_idx is None:
|
||||
# Try looser match: compare first cell only (otpid)
|
||||
first = vals[0] if vals else None
|
||||
if first is not None:
|
||||
for i, sr in df_ordered.iterrows():
|
||||
cand0 = "" if pd.isna(sr.get("otpid")) else str(sr.get("otpid"))
|
||||
if cand0 == first:
|
||||
match_idx = i
|
||||
break
|
||||
if match_idx is None:
|
||||
logger.warning(
|
||||
"Unable to locate DataFrame row matching selected row values: %r",
|
||||
vals,
|
||||
)
|
||||
return
|
||||
idx = int(match_idx)
|
||||
else:
|
||||
# 3) If we have a row_key, try to normalize to an int index
|
||||
if row_key is not None:
|
||||
try:
|
||||
idx = int(row_key)
|
||||
except Exception:
|
||||
# Try converting via string
|
||||
try:
|
||||
idx = int(str(row_key))
|
||||
except Exception:
|
||||
idx = None
|
||||
if idx is None:
|
||||
# Final numeric fallback: use sessions_table.cursor_row if present
|
||||
try:
|
||||
idx = getattr(self.sessions_table, "cursor_row")
|
||||
except Exception:
|
||||
idx = None
|
||||
if idx is None:
|
||||
_bad("Failed to normalize row/key from event: %r", row_key)
|
||||
return
|
||||
else:
|
||||
# 4) Try table cursor_row as last resort
|
||||
try:
|
||||
idx = getattr(self.sessions_table, "cursor_row")
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Could not determine selected row from event: %r", event
|
||||
)
|
||||
# Helpful debug hint for you to paste back if still failing:
|
||||
logger.debug("Event repr for debugging: %r", event)
|
||||
return
|
||||
# At this point we should have an integer idx
|
||||
try:
|
||||
idx = int(idx)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Final normalization of selected row index failed: %r", idx
|
||||
)
|
||||
return
|
||||
# Validate sessions df
|
||||
if self._sessions_df is None or self._sessions_df.empty:
|
||||
logger.warning("Sessions DataFrame empty; nothing to select.")
|
||||
return
|
||||
df_ordered = self._sessions_df.reset_index(drop=True)
|
||||
if idx < 0 or idx >= len(df_ordered):
|
||||
logger.warning(
|
||||
"Selected row index %s out of range (0..%d)", idx, len(df_ordered) - 1
|
||||
)
|
||||
return
|
||||
row_series = df_ordered.iloc[idx]
|
||||
otpid = row_series.get("otpid")
|
||||
hostname = row_series.get("hostname")
|
||||
# Store selected session and fetch activities
|
||||
self._selected_session_otpid = otpid
|
||||
# Obtain api from app (try multiple places)
|
||||
api = (
|
||||
getattr(self.app, "api", None)
|
||||
or getattr(self, "api", None)
|
||||
or getattr(self.app, "airlock_api", None)
|
||||
)
|
||||
if api is None:
|
||||
logger.error("No API available on self.app.api - cannot fetch activities")
|
||||
return
|
||||
logger.info(
|
||||
"Fetching activities for otpid=%s host=%s (selected row=%s)",
|
||||
otpid,
|
||||
hostname,
|
||||
idx,
|
||||
)
|
||||
await self._fetch_activities_for_otpid(api, otpid, hostname=hostname)
|
||||
|
||||
async def load_sessions_from_api(self, api) -> None:
|
||||
"""
|
||||
Pulls OTP session lists, adds status column, concatenates and populates the sessions table.
|
||||
"""
|
||||
try:
|
||||
active = api.otp_find_active()
|
||||
awaiting = api.otp_find_awaiting()
|
||||
enforced = api.otp_find_enforced()
|
||||
revoked = api.otp_find_revoked()
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to fetch OTP session lists: %s", exc)
|
||||
# Present empty
|
||||
active = awaiting = enforced = revoked = pd.DataFrame()
|
||||
|
||||
# Ensure DataFrame objects
|
||||
def _ensure_df(df):
|
||||
return df if isinstance(df, pd.DataFrame) else pd.DataFrame(df)
|
||||
|
||||
active = _ensure_df(active)
|
||||
awaiting = _ensure_df(awaiting)
|
||||
enforced = _ensure_df(enforced)
|
||||
revoked = _ensure_df(revoked)
|
||||
for df, status in [
|
||||
(active, "active"),
|
||||
(awaiting, "awaiting"),
|
||||
(enforced, "enforced"),
|
||||
(revoked, "revoked"),
|
||||
]:
|
||||
if "status" not in df.columns:
|
||||
df["status"] = status
|
||||
combined = pd.concat([active, awaiting, enforced, revoked], ignore_index=True)
|
||||
if "otpid" in combined.columns:
|
||||
combined = combined.sort_values(by="otpid", ascending=False)
|
||||
self._sessions_df = combined
|
||||
# Populate DataTable
|
||||
self.sessions_table.clear()
|
||||
# Ensure columns exist in DF and when missing add empty column
|
||||
expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]
|
||||
for col in expected_cols:
|
||||
if col not in combined.columns:
|
||||
combined[col] = ""
|
||||
self.sessions_table.add_columns(*expected_cols)
|
||||
# Add rows
|
||||
for _, row in combined[expected_cols].iterrows():
|
||||
# Convert values to str for safe insertion
|
||||
vals = ["" if pd.isna(v) else v for v in row.to_list()]
|
||||
self.sessions_table.add_row(*[str(v) for v in vals])
|
||||
logger.info("Loaded %d OTP sessions.", len(combined))
|
||||
|
||||
async def _fetch_activities_for_otpid(self, api, otpid, hostname=None) -> None:
|
||||
"""
|
||||
Fetch activities DataFrame for a given otpid and populate activities_table.
|
||||
"""
|
||||
try:
|
||||
result = api.otp_get_activities(otpid)
|
||||
result_df = (
|
||||
result if isinstance(result, pd.DataFrame) else pd.DataFrame(result)
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to fetch activities for otpid %s: %s", otpid, exc)
|
||||
result_df = pd.DataFrame()
|
||||
# Attach hostname if provided
|
||||
if hostname is not None:
|
||||
result_df["hostname"] = hostname
|
||||
if result_df.empty:
|
||||
logger.info("No activities found for otpid %s (host: %s)", otpid, hostname)
|
||||
self._activities_df = pd.DataFrame()
|
||||
self.activities_table.clear()
|
||||
return
|
||||
# Store and render
|
||||
self._activities_df = result_df.copy()
|
||||
# Rebuild activities_table columns from result_df
|
||||
self.activities_table.clear()
|
||||
# Ensure stable column order
|
||||
for col in result_df.columns:
|
||||
self.activities_table.add_column(col)
|
||||
# Add rows
|
||||
for _, arow in result_df.iterrows():
|
||||
values = ["" if pd.isna(v) else v for v in arow.to_list()]
|
||||
self.activities_table.add_row(*[str(v) for v in values])
|
||||
logger.info(
|
||||
"Loaded %d activity rows for otpid %s (host: %s)",
|
||||
len(result_df),
|
||||
otpid,
|
||||
hostname,
|
||||
)
|
||||
|
||||
async def export_activities(self) -> None:
|
||||
"""
|
||||
Export currently-loaded activities DataFrame to CSV.
|
||||
Can be called directly (programmatically) or from the button handler.
|
||||
"""
|
||||
if self._activities_df is None or self._activities_df.empty:
|
||||
logger.info("No activities loaded to export.")
|
||||
# On-screen short message
|
||||
await self.post_message(Static("No activities to export."))
|
||||
return
|
||||
working_dir = _load_working_dir()
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
filename = f"otp_activities_{self._selected_session_otpid}_{timestamp}.csv"
|
||||
file_path = os.path.join(working_dir, filename)
|
||||
try:
|
||||
self._activities_df.to_csv(file_path, index=False)
|
||||
logger.info("Exported activities to %s", file_path)
|
||||
await self.post_message(Static(f"Exported activities to: {file_path}"))
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to export activities to %s: %s", file_path, exc)
|
||||
await self.post_message(Static("Failed to export activities; check logs."))
|
||||
|
||||
|
||||
class ActivityDetailWidget(Static):
|
||||
"""
|
||||
Interactive widget for Activity Detail screen.
|
||||
Shows the provided DataFrame in a DataTable and offers Export button.
|
||||
Now includes Select All/None and Add to Allowlist functionality.
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
ActivityDetailWidget {
|
||||
height: 1fr;
|
||||
layout: vertical;
|
||||
}
|
||||
#detail_table_container {
|
||||
height: 1fr;
|
||||
padding: 1 1;
|
||||
}
|
||||
#selection_buttons {
|
||||
height: auto;
|
||||
min-height: 3;
|
||||
padding: 1 1;
|
||||
content-align: center middle;
|
||||
}
|
||||
#detail_buttons {
|
||||
height: auto;
|
||||
min-height: 3;
|
||||
padding: 1 1;
|
||||
content-align: center middle;
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(self, activities_df: pd.DataFrame, otpid=None, hostname=None) -> None:
|
||||
super().__init__()
|
||||
self.activities_df = (
|
||||
activities_df.copy()
|
||||
if isinstance(activities_df, pd.DataFrame)
|
||||
else pd.DataFrame(activities_df)
|
||||
)
|
||||
# Add a unique identifier column if not present
|
||||
if "_row_id" not in self.activities_df.columns:
|
||||
self.activities_df["_row_id"] = range(len(self.activities_df))
|
||||
|
||||
self.otpid = otpid
|
||||
self.hostname = hostname
|
||||
self.selected_row_ids = set() # Track selected rows by unique ID
|
||||
self.row_key_to_id = {} # Map DataTable row keys to unique row IDs
|
||||
self.table_row_to_id = {} # Map table row indices to unique row IDs
|
||||
self._last_sort = None # Track last sort column and order
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Static(
|
||||
f"Activity Detail (otpid={self.otpid} host={self.hostname})",
|
||||
classes="panel-title",
|
||||
)
|
||||
# Table container
|
||||
with Vertical(id="detail_table_container"):
|
||||
self.detail_table = DataTable(id="detail_table")
|
||||
yield self.detail_table
|
||||
|
||||
# Original buttons at bottom
|
||||
with Horizontal(id="detail_buttons"):
|
||||
self.add_allowlist_btn = Button(
|
||||
"Add Selected to Allowlist", id="add_allowlist_btn"
|
||||
)
|
||||
yield self.add_allowlist_btn
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
await self._build_table(rebuild=True)
|
||||
self._update_button_states()
|
||||
|
||||
def _update_button_states(self) -> None:
|
||||
"""Update button states based on selection."""
|
||||
has_selection = len(self.selected_row_ids) > 0
|
||||
self.add_allowlist_btn.disabled = not has_selection
|
||||
|
||||
# Update button labels with count
|
||||
count = len(self.selected_row_ids)
|
||||
len(self.activities_df)
|
||||
|
||||
if has_selection:
|
||||
self.add_allowlist_btn.label = f"Add {count} Selected to Allowlist"
|
||||
else:
|
||||
self.add_allowlist_btn.label = "Add Selected to Allowlist"
|
||||
|
||||
async def _build_table(self, rebuild: bool = True) -> None:
|
||||
"""Rebuild the DataTable. If rebuild=False, only refresh rows."""
|
||||
if rebuild:
|
||||
# Full rebuild: clear columns and rows
|
||||
self.detail_table.clear()
|
||||
self.detail_table.columns.clear()
|
||||
self.row_key_to_id.clear()
|
||||
self.table_row_to_id.clear()
|
||||
|
||||
if self.activities_df is None or self.activities_df.empty:
|
||||
logger.info("ActivityDetailWidget mounted with empty dataframe.")
|
||||
return
|
||||
|
||||
# Add columns (checkbox + data columns, excluding internal _row_id)
|
||||
self.detail_table.add_column("Select", key="select")
|
||||
for col in self.activities_df.columns:
|
||||
if col != "_row_id": # Don't display the internal ID column
|
||||
self.detail_table.add_column(col)
|
||||
else:
|
||||
# Partial rebuild: clear rows only
|
||||
self.detail_table.clear()
|
||||
self.row_key_to_id.clear()
|
||||
self.table_row_to_id.clear()
|
||||
|
||||
# Add rows
|
||||
for table_idx, (df_idx, row) in enumerate(self.activities_df.iterrows()):
|
||||
# Get the unique row ID
|
||||
row_id = row["_row_id"]
|
||||
|
||||
# Build values list (excluding _row_id column)
|
||||
vals = []
|
||||
for col in self.activities_df.columns:
|
||||
if col != "_row_id":
|
||||
v = row[col]
|
||||
vals.append("" if pd.isna(v) else str(v))
|
||||
|
||||
# Check if this row is selected
|
||||
checkbox = "☑️" if row_id in self.selected_row_ids else "☐"
|
||||
|
||||
# Add row to table
|
||||
row_key = self.detail_table.add_row(checkbox, *vals)
|
||||
|
||||
# Map the row key and table index to the unique row ID
|
||||
self.row_key_to_id[row_key] = row_id
|
||||
self.table_row_to_id[table_idx] = row_id
|
||||
|
||||
async def on_data_table_cell_selected(self, event: DataTable.CellSelected) -> None:
|
||||
# Toggle selection when the "Select" column is clicked
|
||||
if event.cell_key.column_key.value == "select":
|
||||
table_row_index = event.coordinate.row
|
||||
|
||||
# Get the unique row ID for this table row
|
||||
row_id = self.table_row_to_id.get(table_row_index)
|
||||
if row_id is not None:
|
||||
# Get the row key for updating the cell
|
||||
row_key = event.cell_key.row_key
|
||||
|
||||
if row_id in self.selected_row_ids:
|
||||
self.selected_row_ids.remove(row_id)
|
||||
self.detail_table.update_cell(row_key, "select", "☐") # Unchecked
|
||||
else:
|
||||
self.selected_row_ids.add(row_id)
|
||||
self.detail_table.update_cell(row_key, "select", "☑️") # Checked
|
||||
|
||||
self._update_button_states()
|
||||
|
||||
async def on_data_table_header_selected(
|
||||
self, event: DataTable.HeaderSelected
|
||||
) -> None:
|
||||
column_key = event.column_key.value if event.column_key else None
|
||||
if not column_key:
|
||||
col_index = event.column_index
|
||||
if col_index == 0: # First column is "Select"
|
||||
return
|
||||
# Adjust for hidden _row_id column
|
||||
visible_cols = [
|
||||
col for col in self.activities_df.columns if col != "_row_id"
|
||||
]
|
||||
if col_index - 1 < len(visible_cols):
|
||||
column_key = visible_cols[col_index - 1]
|
||||
else:
|
||||
return
|
||||
if column_key == "select" or column_key == "_row_id":
|
||||
return
|
||||
|
||||
ascending = True
|
||||
if self._last_sort == (column_key, True):
|
||||
ascending = False
|
||||
self._last_sort = (column_key, ascending)
|
||||
|
||||
try:
|
||||
self.activities_df.sort_values(
|
||||
by=column_key, ascending=ascending, inplace=True
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to sort by column %s: %s", column_key, exc)
|
||||
return
|
||||
|
||||
# Only refresh rows, not columns
|
||||
await self._build_table(rebuild=False)
|
||||
|
||||
async def on_button_pressed(self, event) -> None:
|
||||
btn = getattr(event, "button", None) or getattr(event, "sender", None)
|
||||
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None)
|
||||
|
||||
if btn is self.add_allowlist_btn or btn_id == "add_allowlist_btn":
|
||||
await self._open_allowlist_screen()
|
||||
return
|
||||
|
||||
async def _select_all(self) -> None:
|
||||
"""Select all rows in the table."""
|
||||
# Add all row IDs to selected set
|
||||
self.selected_row_ids = set(self.activities_df["_row_id"].tolist())
|
||||
|
||||
# Update all checkboxes in the table
|
||||
for row_key, row_id in self.row_key_to_id.items():
|
||||
self.detail_table.update_cell(row_key, "select", "☑️")
|
||||
|
||||
self._update_button_states()
|
||||
logger.info(f"Selected all {len(self.selected_row_ids)} rows")
|
||||
|
||||
async def _select_none(self) -> None:
|
||||
"""Deselect all rows in the table."""
|
||||
# Clear selected set
|
||||
self.selected_row_ids.clear()
|
||||
|
||||
# Update all checkboxes in the table
|
||||
for row_key, row_id in self.row_key_to_id.items():
|
||||
self.detail_table.update_cell(row_key, "select", "☑️")
|
||||
|
||||
self._update_button_states()
|
||||
logger.info("Cleared all selections")
|
||||
|
||||
async def _open_allowlist_screen(self) -> None:
|
||||
"""Open the allowlist selection screen with selected activities."""
|
||||
if not self.selected_row_ids:
|
||||
self.app.notify("No rows selected", severity="warning")
|
||||
return
|
||||
|
||||
# Get selected data
|
||||
selected_df = self.get_selected_data()
|
||||
|
||||
# Get API from app
|
||||
api = getattr(self.app, "api", None)
|
||||
if api is None:
|
||||
logger.error("No API available on self.app.api")
|
||||
self.app.notify("API not available", severity="error")
|
||||
return
|
||||
|
||||
# Create and push AllowlistSelectionScreen
|
||||
try:
|
||||
allowlist_screen = AllowlistSelectionScreen(
|
||||
selected_df, api=api, hostname=self.hostname, otpid=self.otpid
|
||||
)
|
||||
await self.app.push_screen(allowlist_screen)
|
||||
logger.info(
|
||||
f"Opened allowlist screen with {len(selected_df)} selected activities"
|
||||
)
|
||||
except ImportError as e:
|
||||
logger.error(f"Failed to import AllowlistSelectionScreen: {e}")
|
||||
self.app.notify("Allowlist screen module not found", severity="error")
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to open allowlist screen: {e}")
|
||||
self.app.notify(
|
||||
f"Error opening allowlist screen: {str(e)}", severity="error"
|
||||
)
|
||||
|
||||
async def _export_detail_activities(self) -> None:
|
||||
if self.activities_df is None or self.activities_df.empty:
|
||||
logger.info("No activities to export.")
|
||||
await self.mount(Static("No activities to export.", classes="notification"))
|
||||
return
|
||||
if not self.selected_row_ids:
|
||||
logger.info("No rows selected for export.")
|
||||
await self.mount(
|
||||
Static("No rows selected for export.", classes="notification")
|
||||
)
|
||||
return
|
||||
try:
|
||||
working_dir = load_env("WORKING_DIR") or os.getcwd()
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
filename = f"otp_activities_detail_{timestamp}.csv"
|
||||
file_path = os.path.join(working_dir, filename)
|
||||
selected_df = self.get_selected_data()
|
||||
selected_df.to_csv(file_path, index=False)
|
||||
logger.info("Exported selected activities to %s", file_path)
|
||||
await self.mount(
|
||||
Static(
|
||||
f"Exported selected activities to: {filename}",
|
||||
classes="notification",
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to export detail activities: %s", exc)
|
||||
await self.mount(
|
||||
Static(
|
||||
"¢ Failed to export activities; check logs.",
|
||||
classes="notification",
|
||||
)
|
||||
)
|
||||
|
||||
# Helper methods
|
||||
def get_selected_data(self) -> pd.DataFrame:
|
||||
"""Return a DataFrame of the selected rows."""
|
||||
if not self.selected_row_ids:
|
||||
return pd.DataFrame()
|
||||
# Filter by selected row IDs and drop the internal _row_id column
|
||||
selected_df = self.activities_df[
|
||||
self.activities_df["_row_id"].isin(self.selected_row_ids)
|
||||
].copy()
|
||||
if "_row_id" in selected_df.columns:
|
||||
selected_df = selected_df.drop(columns=["_row_id"])
|
||||
return selected_df
|
||||
|
||||
def get_selected_records(self) -> list[dict]:
|
||||
"""Return selected rows as a list of dicts."""
|
||||
if not self.selected_row_ids:
|
||||
return []
|
||||
# Filter by selected row IDs and drop the internal _row_id column
|
||||
selected_df = self.activities_df[
|
||||
self.activities_df["_row_id"].isin(self.selected_row_ids)
|
||||
].copy()
|
||||
if "_row_id" in selected_df.columns:
|
||||
selected_df = selected_df.drop(columns=["_row_id"])
|
||||
return selected_df.to_dict(orient="records")
|
||||
|
||||
|
||||
class ActivityDetailScreen(Screen):
|
||||
"""
|
||||
Screen that wraps ActivityDetailWidget. Expects a DataFrame passed on init.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("e", "export", "Export"),
|
||||
Binding("a", "select_all", "Select All"),
|
||||
Binding("n", "select_none", "Select None"),
|
||||
]
|
||||
|
||||
def __init__(self, activities_df: pd.DataFrame, otpid=None, hostname=None) -> None:
|
||||
super().__init__()
|
||||
self._activities_df = (
|
||||
activities_df.copy()
|
||||
if isinstance(activities_df, pd.DataFrame)
|
||||
else pd.DataFrame(activities_df)
|
||||
)
|
||||
self._otpid = otpid
|
||||
self._hostname = hostname
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
self.widget = ActivityDetailWidget(
|
||||
self._activities_df, otpid=self._otpid, hostname=self._hostname
|
||||
)
|
||||
yield Header(show_clock=True)
|
||||
yield self.widget
|
||||
yield Footer()
|
||||
|
||||
async def action_go_back(self) -> None:
|
||||
try:
|
||||
await self.app.pop_screen()
|
||||
except Exception:
|
||||
logger.debug("ActivityDetailScreen.action_go_back pop_screen failed.")
|
||||
|
||||
async def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_export(self) -> None:
|
||||
# Delegate to widget export helper
|
||||
if hasattr(self, "widget") and self.widget is not None:
|
||||
await self.widget._export_detail_activities()
|
||||
|
||||
async def action_select_all(self) -> None:
|
||||
"""Handle 'a' key for select all."""
|
||||
if hasattr(self, "widget") and self.widget is not None:
|
||||
await self.widget._select_all()
|
||||
|
||||
async def action_select_none(self) -> None:
|
||||
"""Handle 'n' key for select none."""
|
||||
if hasattr(self, "widget") and self.widget is not None:
|
||||
await self.widget._select_none()
|
||||
|
||||
|
||||
class OTPActivitiesScreen(Screen):
|
||||
"""
|
||||
A Screen intended to be pushed into an existing Textual App.
|
||||
Usage:
|
||||
app.push_screen(OTPActivitiesScreen())
|
||||
or create this screen and call `await screen.load()` inside your app lifecycle.
|
||||
The screen expects `self.app.api` to exist and be an AirlockAPIWrapper instance.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("r", "refresh_sessions", "Refresh Sessions"),
|
||||
Binding("e", "export_activities", "Export activities"),
|
||||
]
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header()
|
||||
self.widget = OTPActivitiesWidget()
|
||||
yield self.widget
|
||||
yield Footer()
|
||||
|
||||
async def on_show(self) -> None:
|
||||
"""Restore focus to the left sessions table when the screen becomes visible."""
|
||||
if hasattr(self, "widget") and hasattr(self.widget, "sessions_table"):
|
||||
self.widget.sessions_table.focus()
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
# Try to load sessions immediately
|
||||
api = getattr(self.app, "api", None)
|
||||
if api is None:
|
||||
logger.warning("OTPActivitiesScreen mounted but no self.app.api found.")
|
||||
else:
|
||||
await self.widget.load_sessions_from_api(api)
|
||||
|
||||
async def action_go_back(self) -> None:
|
||||
"""Go back one screen."""
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
await self.app.pop_screen()
|
||||
|
||||
# Simple actions bound to keys
|
||||
async def action_refresh_sessions(self) -> None:
|
||||
api = getattr(self.app, "api", None)
|
||||
if api is None:
|
||||
logger.error("No API on app; cannot refresh sessions.")
|
||||
return
|
||||
logger.info("Refreshing OTP sessions via API.")
|
||||
await self.widget.load_sessions_from_api(api)
|
||||
|
||||
# If you want an explicit method to fetch activities for a particular otpid from outside:
|
||||
async def fetch_activities_for_otpid(self, otpid, hostname=None) -> None:
|
||||
api = getattr(self.app, "api", None)
|
||||
if api is None:
|
||||
logger.error("No API on app; cannot fetch activities.")
|
||||
return
|
||||
await self.widget._fetch_activities_for_otpid(api, otpid, hostname=hostname)
|
||||
@@ -0,0 +1,393 @@
|
||||
# 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/>.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import pandas as pd
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.message import Message
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Static
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OTPRevokeWidget(Static):
|
||||
"""
|
||||
Widget for managing OTP session revocation.
|
||||
Displays active OTP sessions and allows selection for revocation.
|
||||
"""
|
||||
|
||||
class SessionsRevoked(Message):
|
||||
"""Message sent when sessions are revoked."""
|
||||
|
||||
def __init__(self, revoked_sessions: List[dict]):
|
||||
super().__init__()
|
||||
self.revoked_sessions = revoked_sessions
|
||||
|
||||
DEFAULT_CSS = """
|
||||
OTPRevokeWidget {
|
||||
height: 1fr;
|
||||
}
|
||||
#main_container {
|
||||
width: 100%;
|
||||
height: 100%;
|
||||
layout: vertical;
|
||||
}
|
||||
#sessions_container {
|
||||
height: 1fr;
|
||||
border: none;
|
||||
padding: 1;
|
||||
}
|
||||
#button_container {
|
||||
height: auto;
|
||||
padding: 1;
|
||||
align: center middle;
|
||||
}
|
||||
#result_container {
|
||||
height: auto;
|
||||
max-height: 10;
|
||||
border: solid #444444;
|
||||
padding: 1;
|
||||
margin: 1;
|
||||
overflow-y: auto;
|
||||
}
|
||||
.panel-title {
|
||||
text-style: bold;
|
||||
margin: 0 0 1 0;
|
||||
}
|
||||
"""
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Vertical(id="main_container"):
|
||||
# Sessions table
|
||||
yield Static("OTP Sessions", classes="panel-title")
|
||||
with Vertical(id="sessions_container"):
|
||||
self.sessions_table = DataTable(id="sessions_table")
|
||||
self.sessions_table.styles.width = "100%"
|
||||
self.sessions_table.styles.height = "1fr"
|
||||
yield self.sessions_table
|
||||
|
||||
# Action buttons
|
||||
with Horizontal(id="button_container"):
|
||||
self.refresh_button = Button("🔄 Refresh", id="refresh_btn")
|
||||
self.refresh_button.styles.width = "15%"
|
||||
self.refresh_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.refresh_button
|
||||
|
||||
self.select_all_button = Button("☑️ Select All", id="select_all_btn")
|
||||
self.select_all_button.styles.width = "15%"
|
||||
self.select_all_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.select_all_button
|
||||
|
||||
self.select_none_button = Button(
|
||||
"❌ Clear Selection", id="select_none_btn"
|
||||
)
|
||||
self.select_none_button.styles.width = "20%"
|
||||
self.select_none_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.select_none_button
|
||||
|
||||
self.revoke_button = Button(
|
||||
"🛑 Revoke Selected", id="revoke_btn", variant="error"
|
||||
)
|
||||
self.revoke_button.styles.width = "20%"
|
||||
self.revoke_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.revoke_button
|
||||
|
||||
# Results display
|
||||
with Vertical(id="result_container"):
|
||||
yield Static("Revocation Results", classes="panel-title")
|
||||
self.results_display = Static("No actions performed yet.")
|
||||
yield self.results_display
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
"""Initialize the widget when mounted."""
|
||||
# Configure sessions table
|
||||
self.sessions_table.clear()
|
||||
self.sessions_table.add_columns(
|
||||
"☐", "OTP ID", "Hostname", "Status", "Purpose", "Granted"
|
||||
)
|
||||
|
||||
# Enable row selection with checkbox column
|
||||
self.sessions_table.cursor_type = "row"
|
||||
try:
|
||||
self.sessions_table.zebra_stripes = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Initialize state
|
||||
self._sessions_df: Optional[pd.DataFrame] = None
|
||||
self._filtered_df: Optional[pd.DataFrame] = None
|
||||
self._selected_otpids: set = set()
|
||||
|
||||
async def load_sessions_from_api(self, api) -> None:
|
||||
"""Load active OTP sessions from the API."""
|
||||
try:
|
||||
# Fetch only active sessions
|
||||
active_df = api.otp_find_active()
|
||||
|
||||
# Ensure we have a DataFrame
|
||||
if not isinstance(active_df, pd.DataFrame):
|
||||
active_df = pd.DataFrame(active_df)
|
||||
|
||||
# Add status column
|
||||
active_df["status"] = "active"
|
||||
|
||||
# Sort by otpid if column exists
|
||||
if "otpid" in active_df.columns and not active_df.empty:
|
||||
active_df = active_df.sort_values(by="otpid", ascending=False)
|
||||
|
||||
# Store the full dataframe
|
||||
self._sessions_df = active_df
|
||||
self._filtered_df = active_df.copy()
|
||||
|
||||
# Display in table
|
||||
await self._refresh_table()
|
||||
|
||||
# Update status
|
||||
active_count = len(active_df)
|
||||
|
||||
status_msg = f"Loaded {active_count} active sessions"
|
||||
logger.info(status_msg)
|
||||
self.results_display.update(status_msg)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to load OTP sessions: {e}")
|
||||
self.results_display.update(f"Error loading sessions: {str(e)}")
|
||||
|
||||
async def _refresh_table(self) -> None:
|
||||
"""Refresh the table display with current filtered data."""
|
||||
if self._filtered_df is None or self._filtered_df.empty:
|
||||
self.sessions_table.clear()
|
||||
return
|
||||
|
||||
# Ensure expected columns exist
|
||||
expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]
|
||||
for col in expected_cols:
|
||||
if col not in self._filtered_df.columns:
|
||||
self._filtered_df[col] = ""
|
||||
|
||||
# Clear and repopulate table
|
||||
self.sessions_table.clear(columns=False)
|
||||
|
||||
for _, row in self._filtered_df.iterrows():
|
||||
otpid = str(row.get("otpid", ""))
|
||||
# Check if this row is selected
|
||||
checkbox = "☑️" if otpid in self._selected_otpids else "☐"
|
||||
|
||||
self.sessions_table.add_row(
|
||||
checkbox,
|
||||
str(otpid),
|
||||
str(row.get("hostname", "")),
|
||||
str(row.get("status", "")),
|
||||
str(row.get("purpose", "")),
|
||||
str(row.get("granted", "")),
|
||||
)
|
||||
|
||||
async def on_button_pressed(self, event) -> None:
|
||||
"""Handle button presses."""
|
||||
btn = event.button
|
||||
|
||||
if btn.id == "refresh_btn":
|
||||
# Refresh sessions
|
||||
api = getattr(self.app, "api", None)
|
||||
if api:
|
||||
await self.load_sessions_from_api(api)
|
||||
|
||||
elif btn.id == "select_all_btn":
|
||||
# Select all visible rows
|
||||
if self._filtered_df is not None:
|
||||
self._selected_otpids = set(str(x) for x in self._filtered_df["otpid"])
|
||||
await self._refresh_table()
|
||||
|
||||
elif btn.id == "select_none_btn":
|
||||
# Clear selection
|
||||
self._selected_otpids.clear()
|
||||
await self._refresh_table()
|
||||
|
||||
elif btn.id == "revoke_btn":
|
||||
# Revoke selected sessions
|
||||
await self._revoke_selected()
|
||||
|
||||
async def on_data_table_row_selected(self, event) -> None:
|
||||
"""Handle row selection in the table."""
|
||||
if event.data_table != self.sessions_table:
|
||||
return
|
||||
|
||||
try:
|
||||
# Get the row index from the cursor row
|
||||
row_index = self.sessions_table.cursor_row
|
||||
|
||||
if self._filtered_df is not None and row_index < len(self._filtered_df):
|
||||
# Get the OTP ID for this row
|
||||
otpid = str(self._filtered_df.iloc[row_index]["otpid"])
|
||||
|
||||
# Toggle selection
|
||||
if otpid in self._selected_otpids:
|
||||
self._selected_otpids.remove(otpid)
|
||||
else:
|
||||
self._selected_otpids.add(otpid)
|
||||
|
||||
# Refresh table to update checkbox
|
||||
await self._refresh_table()
|
||||
|
||||
# Restore cursor position
|
||||
self.sessions_table.move_cursor(row=row_index)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Error handling row selection: {e}")
|
||||
|
||||
async def _revoke_selected(self) -> None:
|
||||
"""Revoke the selected OTP sessions."""
|
||||
if not self._selected_otpids:
|
||||
self.results_display.update("❌ No sessions selected for revocation")
|
||||
return
|
||||
|
||||
api = getattr(self.app, "api", None)
|
||||
if not api:
|
||||
self.results_display.update("❌ API not available")
|
||||
return
|
||||
|
||||
# Collect results
|
||||
results = []
|
||||
success_count = 0
|
||||
failure_count = 0
|
||||
|
||||
for otpid in self._selected_otpids:
|
||||
try:
|
||||
# Get hostname for this session
|
||||
hostname = "Unknown"
|
||||
if self._sessions_df is not None:
|
||||
# Convert otpid to same type as in DataFrame for comparison
|
||||
otpid_compare = otpid
|
||||
if len(self._sessions_df) > 0:
|
||||
first_otpid = self._sessions_df["otpid"].iloc[0]
|
||||
if isinstance(first_otpid, int):
|
||||
otpid_compare = int(otpid)
|
||||
|
||||
match = self._sessions_df[
|
||||
self._sessions_df["otpid"] == otpid_compare
|
||||
]
|
||||
if not match.empty:
|
||||
hostname = match.iloc[0].get("hostname", "Unknown")
|
||||
|
||||
# Revoke the session
|
||||
result = api.otp_revoke(otpid)
|
||||
|
||||
if result and result.get("status") != "error":
|
||||
success_count += 1
|
||||
results.append(f"Revoked OTP {otpid} for {hostname}")
|
||||
logger.info(f"Revoked OTP {otpid} for {hostname}: {result}")
|
||||
else:
|
||||
failure_count += 1
|
||||
error_msg = (
|
||||
result.get("message", "Unknown error")
|
||||
if result
|
||||
else "No response"
|
||||
)
|
||||
results.append(
|
||||
f"❌ Failed to revoke OTP {otpid} for {hostname}: {error_msg}"
|
||||
)
|
||||
logger.error(f"Failed to revoke OTP {otpid}: {error_msg}")
|
||||
|
||||
except Exception as e:
|
||||
failure_count += 1
|
||||
results.append(f"❌ Error revoking OTP {otpid}: {str(e)}")
|
||||
logger.exception(f"Exception revoking OTP {otpid}: {e}")
|
||||
|
||||
# Update results display
|
||||
summary = (
|
||||
f"Revocation complete: {success_count} succeeded, {failure_count} failed\n"
|
||||
)
|
||||
details = "\n".join(results[-5:]) # Show last 5 results
|
||||
if len(results) > 5:
|
||||
details = f"... (showing last 5 of {len(results)} results)\n" + details
|
||||
|
||||
self.results_display.update(summary + details)
|
||||
|
||||
# Clear selection and refresh
|
||||
self._selected_otpids.clear()
|
||||
await self.load_sessions_from_api(api)
|
||||
|
||||
# Post message about revoked sessions
|
||||
if success_count > 0:
|
||||
self.post_message(self.SessionsRevoked(results))
|
||||
|
||||
|
||||
class OTPRevokeScreen(Screen):
|
||||
"""
|
||||
Main screen for OTP session revocation workflow.
|
||||
This replaces the otp_revoke function from otp.py.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("r", "refresh", "Refresh"),
|
||||
Binding("a", "select_all", "Select All"),
|
||||
Binding("n", "select_none", "Clear Selection"),
|
||||
Binding("d", "revoke", "Revoke Selected"),
|
||||
]
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header(show_clock=True)
|
||||
self.widget = OTPRevokeWidget()
|
||||
yield self.widget
|
||||
yield Footer()
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
"""Load sessions when screen mounts."""
|
||||
api = getattr(self.app, "api", None)
|
||||
if api:
|
||||
await self.widget.load_sessions_from_api(api)
|
||||
else:
|
||||
logger.warning("OTPRevokeScreen mounted but no self.app.api found.")
|
||||
|
||||
async def action_refresh(self) -> None:
|
||||
"""Refresh the sessions list."""
|
||||
api = getattr(self.app, "api", None)
|
||||
if api:
|
||||
await self.widget.load_sessions_from_api(api)
|
||||
|
||||
async def action_select_all(self) -> None:
|
||||
"""Select all visible sessions."""
|
||||
if self.widget._filtered_df is not None:
|
||||
self.widget._selected_otpids = set(
|
||||
str(x) for x in self.widget._filtered_df["otpid"]
|
||||
)
|
||||
await self.widget._refresh_table()
|
||||
|
||||
async def action_select_none(self) -> None:
|
||||
"""Clear all selections."""
|
||||
self.widget._selected_otpids.clear()
|
||||
await self.widget._refresh_table()
|
||||
|
||||
async def action_revoke(self) -> None:
|
||||
"""Revoke selected sessions."""
|
||||
await self.widget._revoke_selected()
|
||||
|
||||
async def action_go_back(self) -> None:
|
||||
"""Go back to previous screen."""
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
await self.app.pop_screen()
|
||||
@@ -0,0 +1,52 @@
|
||||
# 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/>.
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.screen import Screen
|
||||
|
||||
from models.agent import Agent
|
||||
from TUI.Widgets.OTP_generate import OTPGenerator
|
||||
|
||||
|
||||
class OTPWorkflowScreen(Screen):
|
||||
"""Screen that handles the OTP generation workflow without agent selection."""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
def __init__(self, selected_agents: Optional[List[Agent]]):
|
||||
super().__init__()
|
||||
self.selected_agents = selected_agents
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Directly show the OTP generator for the selected agents."""
|
||||
yield OTPGenerator(self.selected_agents)
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle escape key to go back one screen."""
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Handle q key to go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
|
||||
"""Handle OTP generation request - pass it up to the app level if needed."""
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,109 @@
|
||||
# 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/>.
|
||||
|
||||
"""
|
||||
Policy Selector Screen Module
|
||||
|
||||
Provides a Textual Screen wrapper for the PolicySelector widget that manages
|
||||
the policy selection workflow.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Footer, Header
|
||||
|
||||
from TUI.Widgets.policyselector import PolicySelector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PolicySelectorScreen(Screen):
|
||||
"""
|
||||
A Textual Screen for policy selection in agent move operations.
|
||||
|
||||
This screen wraps the PolicySelector widget and manages the workflow
|
||||
of selecting a target policy for bulk agent movements.
|
||||
|
||||
Attributes:
|
||||
policies: List of available policies (Policy objects or DataFrame).
|
||||
agent_move_operations: Reference to the parent AgentMoveOperations widget.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
CSS = """
|
||||
Screen {
|
||||
layout: vertical;
|
||||
background: $surface;
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
policies,
|
||||
agent_move_operations=None,
|
||||
):
|
||||
"""
|
||||
Initialize the PolicySelectorScreen.
|
||||
|
||||
Args:
|
||||
policies: List of available policies to display.
|
||||
agent_move_operations: Reference to parent AgentMoveOperations widget.
|
||||
Used to call back when policy selection is confirmed.
|
||||
"""
|
||||
super().__init__()
|
||||
self.policies = policies
|
||||
self.agent_move_operations = agent_move_operations
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Create the PolicySelector widget."""
|
||||
yield Header(show_clock=True)
|
||||
yield PolicySelector(self.policies)
|
||||
yield Footer()
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle escape key to go back one screen."""
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Handle q key to go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_policy_selector_policy_selected(
|
||||
self, message: PolicySelector.PolicySelected
|
||||
) -> None:
|
||||
"""
|
||||
Handle policy selection from the PolicySelector widget.
|
||||
|
||||
When a policy is selected, this handler:
|
||||
1. Closes the selector screen
|
||||
2. Calls the parent AgentMoveOperations to execute the move
|
||||
|
||||
Args:
|
||||
message (PolicySelector.PolicySelected): Contains the selected policy.
|
||||
"""
|
||||
# Pop this screen to return to AgentMoveOperations
|
||||
self.app.pop_screen()
|
||||
|
||||
# Call parent widget's method to execute the move
|
||||
if self.agent_move_operations:
|
||||
self.agent_move_operations._execute_move_to_policy(message.policy)
|
||||
@@ -0,0 +1,841 @@
|
||||
# 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/>.
|
||||
|
||||
"""
|
||||
Quiet Agent Workflow Screen Module
|
||||
|
||||
Provides a TUI workflow for identifying quiet agents and moving them to target policies.
|
||||
This screen replaces the legacy quietAgent.py with a comprehensive TUI interface that:
|
||||
1. Allows selection of an initial policy to analyze
|
||||
2. Categorizes devices into "Enforce Ready" and "Non-Enforce Ready" based on activity
|
||||
3. Allows users to select target policies for each category
|
||||
4. Uses the API to move devices to their target policies
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import logging
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
import pandas as pd
|
||||
from textual.app import ComposeResult
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.reactive import reactive
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Static
|
||||
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.policyhandler import getPolicyInfo
|
||||
from TUI.Widgets.policyselector import PolicySelector
|
||||
from utils.configmanager import load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QuietAgentWorkflowScreen(Screen):
|
||||
"""
|
||||
A Textual screen for the Quiet Agent analysis and migration workflow.
|
||||
|
||||
This screen provides a multi-step workflow:
|
||||
1. Select initial policy to analyze
|
||||
2. View categorized agents (enforce ready vs. non-enforce ready)
|
||||
3. Select target policies for each category
|
||||
4. Execute agent migrations
|
||||
|
||||
Attributes:
|
||||
api (AirlockAPIWrapper): API wrapper for Airlock operations
|
||||
policies (List[Policy]): List of all available policies
|
||||
selected_policy (Optional[Policy]): The initially selected policy to analyze
|
||||
history_days (int): Number of days of history to pull (default: 150)
|
||||
quiet_days (int): Number of days without execution to be considered quiet (default: 45)
|
||||
agents_df (Optional[pd.DataFrame]): DataFrame of all agents with analysis results
|
||||
enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents ready for enforcement
|
||||
non_enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents not ready for enforcement
|
||||
workflow_stage (str): Current stage of the workflow
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
("escape", "go_back", "Back"),
|
||||
("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
workflow_stage = reactive("select_policy") # Tracks current workflow stage
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper, policies: List[Policy]):
|
||||
"""
|
||||
Initialize the QuietAgentWorkflowScreen.
|
||||
|
||||
Args:
|
||||
api (AirlockAPIWrapper): API wrapper for Airlock operations
|
||||
policies (List[Policy]): List of all available policies
|
||||
"""
|
||||
super().__init__()
|
||||
self.api = api
|
||||
self.policies = policies
|
||||
self.selected_policy: Optional[Policy] = None
|
||||
self.history_days = 150 # Fixed as per requirements
|
||||
self.quiet_days = 45 # Default value
|
||||
self.agents_df: Optional[pd.DataFrame] = None
|
||||
self.enforce_ready_df: Optional[pd.DataFrame] = None
|
||||
self.non_enforce_ready_df: Optional[pd.DataFrame] = None
|
||||
self.enforce_ready_target_policy: Optional[Policy] = None
|
||||
self.non_enforce_ready_target_policy: Optional[Policy] = None
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Build the UI layout for the workflow screen."""
|
||||
# Include Header and Footer like other standalone screens
|
||||
yield Header(show_clock=True, icon="⚙️")
|
||||
|
||||
# Title area
|
||||
title = Static("Quiet Agent Workflow", id="workflow_title")
|
||||
title.styles.margin = (0, 0, 0, 1)
|
||||
yield title
|
||||
|
||||
# Status area
|
||||
status = Static("Step 1: Select Policy to Analyze", id="workflow_status")
|
||||
status.styles.margin = (0, 0, 1, 1)
|
||||
yield status
|
||||
|
||||
# Content area - dynamically populated based on workflow stage
|
||||
yield Vertical(id="content_area")
|
||||
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Initialize the screen when mounted."""
|
||||
# Show initial policy selection
|
||||
self._show_policy_selection()
|
||||
|
||||
def watch_workflow_stage(self, old_value: str, new_value: str) -> None:
|
||||
"""React to workflow stage changes."""
|
||||
logger.debug(f"Workflow stage changed from {old_value} to {new_value}")
|
||||
self._update_status_message()
|
||||
|
||||
def _update_status_message(self) -> None:
|
||||
"""Update the status message based on current workflow stage."""
|
||||
status_widget = self.query_one("#workflow_status", Static)
|
||||
|
||||
stage_messages = {
|
||||
"select_policy": "Step 1: Select Policy to Analyze",
|
||||
"select_quiet_days": "Step 2: Select Quiet Time Period",
|
||||
"analyzing": "Analyzing agent activity...",
|
||||
"view_results": "Step 3: Review Categorized Agents",
|
||||
"select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents",
|
||||
"select_non_enforce_target": "Step 5: Select Target Policy for Non-Enforce Ready Agents",
|
||||
"confirm_migration": "Step 6: Confirm and Execute Migration",
|
||||
"executing": "Executing agent migrations...",
|
||||
"complete": "Migration Complete",
|
||||
}
|
||||
|
||||
status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage"))
|
||||
|
||||
def _show_policy_selection(self) -> None:
|
||||
"""Show the initial policy selection screen."""
|
||||
self.workflow_stage = "select_policy"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Create policy selector widget
|
||||
policy_selector = PolicySelector(self.policies)
|
||||
content.mount(policy_selector)
|
||||
|
||||
def on_policy_selector_policy_selected(
|
||||
self, message: PolicySelector.PolicySelected
|
||||
) -> None:
|
||||
"""Handle policy selection from PolicySelector widget."""
|
||||
# Handle based on current workflow stage
|
||||
if self.workflow_stage == "select_policy":
|
||||
# Initial policy selection for analysis
|
||||
self.selected_policy = message.policy
|
||||
logger.info(f"Selected policy for analysis: {self.selected_policy.name}")
|
||||
self._show_quiet_days_selection()
|
||||
elif self.workflow_stage == "select_enforce_target":
|
||||
# Target policy selection for enforce ready agents
|
||||
self.enforce_ready_target_policy = message.policy
|
||||
logger.info(
|
||||
f"Selected target policy for enforce ready: {message.policy.name}"
|
||||
)
|
||||
self._show_non_enforce_target_selection()
|
||||
elif self.workflow_stage == "select_non_enforce_target":
|
||||
# Target policy selection for non-enforce ready agents
|
||||
self.non_enforce_ready_target_policy = message.policy
|
||||
logger.info(
|
||||
f"Selected target policy for non-enforce ready: {message.policy.name}"
|
||||
)
|
||||
self._show_migration_confirmation()
|
||||
|
||||
def _show_quiet_days_selection(self) -> None:
|
||||
"""Show the quiet days selection screen."""
|
||||
self.workflow_stage = "select_quiet_days"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Create info text
|
||||
info_widget = Static(
|
||||
f"Policy Selected: {self.selected_policy.name}\n\n"
|
||||
f"History Period: {self.history_days} days\n\n"
|
||||
"Select quiet time period (days without untrusted execution):",
|
||||
id="quiet_days_info",
|
||||
)
|
||||
info_widget.styles.margin = (0, 0, 2, 0)
|
||||
content.mount(info_widget)
|
||||
|
||||
# Create button container and mount it first
|
||||
button_container = Vertical(id="quiet_days_buttons")
|
||||
button_container.styles.height = "auto"
|
||||
content.mount(button_container)
|
||||
|
||||
# Now add buttons to the mounted container
|
||||
for days in [15, 30, 45, 60]:
|
||||
btn = Button(
|
||||
f"{days} days {'(Default)' if days == 45 else ''}",
|
||||
id=f"quiet_days_{days}",
|
||||
classes="quiet_day_btn",
|
||||
)
|
||||
btn.styles.width = "100%"
|
||||
btn.styles.margin = (0, 0, 1, 0)
|
||||
button_container.mount(btn)
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button press events."""
|
||||
button_id = event.button.id
|
||||
|
||||
# Quiet days selection buttons
|
||||
if button_id and button_id.startswith("quiet_days_"):
|
||||
days = int(button_id.split("_")[-1])
|
||||
self.quiet_days = days
|
||||
logger.info(f"Selected quiet days: {days}")
|
||||
self._start_analysis()
|
||||
return
|
||||
|
||||
# Navigation buttons
|
||||
if button_id == "select_enforce_target_btn":
|
||||
self._show_enforce_target_selection()
|
||||
return
|
||||
|
||||
if button_id == "select_non_enforce_target_btn":
|
||||
self._show_non_enforce_target_selection()
|
||||
return
|
||||
|
||||
if button_id == "skip_enforce_target_btn":
|
||||
# Skip enforce ready target selection
|
||||
self.enforce_ready_target_policy = None
|
||||
self._show_non_enforce_target_selection()
|
||||
return
|
||||
|
||||
if button_id == "skip_non_enforce_target_btn":
|
||||
# Skip non-enforce ready target selection
|
||||
self.non_enforce_ready_target_policy = None
|
||||
self._show_migration_confirmation()
|
||||
return
|
||||
|
||||
if button_id == "confirm_migration_btn":
|
||||
self._execute_migration()
|
||||
return
|
||||
|
||||
if button_id == "cancel_migration_btn":
|
||||
self._show_results()
|
||||
return
|
||||
|
||||
if button_id == "export_results_btn":
|
||||
self._export_results()
|
||||
return
|
||||
|
||||
if button_id == "start_over_btn":
|
||||
self._show_policy_selection()
|
||||
return
|
||||
|
||||
def _start_analysis(self) -> None:
|
||||
"""Start the agent activity analysis."""
|
||||
self.workflow_stage = "analyzing"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Show analyzing message with detailed steps
|
||||
analyzing_msg = Static(
|
||||
f"Analyzing Agent Activity\n"
|
||||
f"{'=' * 50}\n\n"
|
||||
f"Policy: {self.selected_policy.name}\n"
|
||||
f"History Period: {self.history_days} days\n"
|
||||
f"Quiet Threshold: {self.quiet_days} days\n\n"
|
||||
f"Progress:\n"
|
||||
f"Step 1/4: Fetching agents from policy...\n"
|
||||
f"Step 2/4: Pulling execution history (this may take a moment)...\n"
|
||||
f"Step 3/4: Analyzing activity patterns...\n"
|
||||
f"Step 4/4: Categorizing agents...\n\n"
|
||||
f"Please wait - this operation cannot be cancelled.",
|
||||
id="analyzing_message",
|
||||
)
|
||||
analyzing_msg.styles.margin = (2, 1)
|
||||
content.mount(analyzing_msg)
|
||||
|
||||
# Show notification
|
||||
self.app.notify(
|
||||
"Starting analysis - this may take several minutes for large policies",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
# Perform the analysis asynchronously
|
||||
self.call_later(self._perform_analysis)
|
||||
|
||||
def _perform_analysis(self) -> None:
|
||||
"""Perform the actual agent activity analysis."""
|
||||
try:
|
||||
# Update status: Fetching agents
|
||||
self._update_analysis_status("Step 1/4: Fetching agents from policy...")
|
||||
|
||||
# Get agents in the selected policy
|
||||
agents = self.api.agents_find_by_group(self.selected_policy.groupid)
|
||||
|
||||
if agents.empty:
|
||||
self.app.notify(
|
||||
f"No agents found in policy: {self.selected_policy.name}",
|
||||
severity="warning",
|
||||
timeout=5,
|
||||
)
|
||||
self._show_policy_selection()
|
||||
return
|
||||
|
||||
agent_count = len(agents)
|
||||
self.app.notify(
|
||||
f"Found {agent_count} agents - fetching execution history...",
|
||||
severity="information",
|
||||
timeout=3,
|
||||
)
|
||||
|
||||
# Update status: Pulling execution history
|
||||
self._update_analysis_status(
|
||||
f"Step 2/4: Pulling execution history for {agent_count} agents...\n"
|
||||
f"(This may take several minutes - progress shown in terminal)"
|
||||
)
|
||||
|
||||
# Get execution history (this shows progress bars in terminal via airlock_libs)
|
||||
policy_exec_history = getPolicyInfo(
|
||||
self.api, self.selected_policy, [1, 2, 6, 7], self.history_days
|
||||
)
|
||||
|
||||
# Update status: Analyzing patterns
|
||||
self._update_analysis_status("Step 3/4: Analyzing activity patterns...")
|
||||
self.app.notify(
|
||||
"History retrieved - analyzing patterns...",
|
||||
severity="information",
|
||||
timeout=2,
|
||||
)
|
||||
|
||||
if policy_exec_history.empty:
|
||||
logger.info(
|
||||
"No execution history found for the selected policy and time range."
|
||||
)
|
||||
# All agents are quiet (no executions)
|
||||
agents["execution_count"] = 0
|
||||
agents["days_since"] = None
|
||||
agents["required_quiet"] = self.quiet_days
|
||||
agents["enforce_ready"] = True
|
||||
else:
|
||||
# Convert datetime column
|
||||
policy_exec_history["datetime"] = pd.to_datetime(
|
||||
policy_exec_history["datetime"],
|
||||
format="%Y-%m-%dT%H:%M:%SZ",
|
||||
utc=True,
|
||||
)
|
||||
|
||||
# Calculate days ago
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
|
||||
lambda dt: (now - dt).days
|
||||
)
|
||||
|
||||
# Count total executions per hostname
|
||||
hostname_counts = policy_exec_history["hostname"].value_counts()
|
||||
agents["execution_count"] = (
|
||||
agents["hostname"].map(hostname_counts).fillna(0).astype(int)
|
||||
)
|
||||
|
||||
# Find most recent execution per hostname
|
||||
most_recent_exec = policy_exec_history.sort_values(
|
||||
by="days_ago"
|
||||
).drop_duplicates(subset="hostname", keep="first")
|
||||
|
||||
# Map most recent execution age to agents
|
||||
agents["days_since"] = agents["hostname"].map(
|
||||
most_recent_exec.set_index("hostname")["days_ago"]
|
||||
)
|
||||
|
||||
# Check for enforcement readiness
|
||||
agents["required_quiet"] = self.quiet_days
|
||||
agents["enforce_ready"] = agents["days_since"].apply(
|
||||
lambda x: True if pd.isna(x) or x > self.quiet_days else False
|
||||
)
|
||||
|
||||
# Update status: Categorizing
|
||||
self._update_analysis_status("Step 4/4: Categorizing agents...")
|
||||
|
||||
# Sort agents
|
||||
agents = agents.sort_values(
|
||||
by=["execution_count", "hostname"], ascending=[True, True]
|
||||
)
|
||||
|
||||
# Store the results
|
||||
self.agents_df = agents
|
||||
|
||||
# Categorize agents into DataFrames
|
||||
self.enforce_ready_df = agents[agents["enforce_ready"]].copy()
|
||||
self.non_enforce_ready_df = agents[not agents["enforce_ready"]].copy()
|
||||
|
||||
logger.info(
|
||||
f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, "
|
||||
f"{len(self.non_enforce_ready_df)} non-enforce ready"
|
||||
)
|
||||
|
||||
self.app.notify(
|
||||
f"Analysis complete! Found {len(self.enforce_ready_df)} enforce ready, "
|
||||
f"{len(self.non_enforce_ready_df)} not ready",
|
||||
severity="success",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
# Show results
|
||||
self._show_results()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during analysis: {e}", exc_info=True)
|
||||
self.app.notify(f"Analysis failed: {str(e)}", severity="error", timeout=5)
|
||||
self._show_policy_selection()
|
||||
|
||||
def _update_analysis_status(self, status_text: str) -> None:
|
||||
"""Update the analysis status message."""
|
||||
try:
|
||||
analyzing_msg = self.query_one("#analyzing_message", Static)
|
||||
|
||||
# Build updated message
|
||||
updated_text = (
|
||||
f"Analyzing Agent Activity\n"
|
||||
f"{'=' * 50}\n\n"
|
||||
f"Policy: {self.selected_policy.name}\n"
|
||||
f"History Period: {self.history_days} days\n"
|
||||
f"Quiet Threshold: {self.quiet_days} days\n\n"
|
||||
f"Progress:\n"
|
||||
f"{status_text}\n\n"
|
||||
f"Please wait - this operation cannot be cancelled."
|
||||
)
|
||||
|
||||
analyzing_msg.update(updated_text)
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not update analysis status: {e}")
|
||||
|
||||
def _show_results(self) -> None:
|
||||
"""Show the categorized results."""
|
||||
self.workflow_stage = "view_results"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Create results display container and mount it first
|
||||
results_container = Vertical(id="results_container")
|
||||
results_container.styles.height = "auto"
|
||||
results_container.styles.margin = (1, 1)
|
||||
content.mount(results_container)
|
||||
|
||||
# Summary statistics
|
||||
total_agents = len(self.enforce_ready_df) + len(self.non_enforce_ready_df)
|
||||
ready_count = len(self.enforce_ready_df)
|
||||
not_ready_count = len(self.non_enforce_ready_df)
|
||||
ready_percentage = (ready_count / total_agents * 100) if total_agents > 0 else 0
|
||||
|
||||
summary = Static(
|
||||
f"Analysis Results for: {self.selected_policy.name}\n\n"
|
||||
f"Total Agents: {total_agents}\n"
|
||||
f"Enforce Ready: {ready_count} ({ready_percentage:.1f}%)\n"
|
||||
f"Not Ready: {not_ready_count} ({100 - ready_percentage:.1f}%)\n\n"
|
||||
f"Quiet Threshold: {self.quiet_days} days\n"
|
||||
f"History Period: {self.history_days} days",
|
||||
id="results_summary",
|
||||
)
|
||||
summary.styles.margin = (0, 0, 2, 0)
|
||||
results_container.mount(summary)
|
||||
|
||||
# Action buttons
|
||||
button_container = Horizontal(id="results_buttons")
|
||||
button_container.styles.height = "auto"
|
||||
results_container.mount(button_container)
|
||||
|
||||
if ready_count > 0:
|
||||
enforce_btn = Button(
|
||||
f"Select Target for Enforce Ready ({ready_count})",
|
||||
id="select_enforce_target_btn",
|
||||
)
|
||||
enforce_btn.styles.margin = (0, 1, 1, 0)
|
||||
button_container.mount(enforce_btn)
|
||||
|
||||
if not_ready_count > 0:
|
||||
non_enforce_btn = Button(
|
||||
f"Select Target for Non-Enforce Ready ({not_ready_count})",
|
||||
id="select_non_enforce_target_btn",
|
||||
)
|
||||
non_enforce_btn.styles.margin = (0, 1, 1, 0)
|
||||
button_container.mount(non_enforce_btn)
|
||||
|
||||
export_btn = Button("Export Results", id="export_results_btn")
|
||||
export_btn.styles.margin = (0, 1, 1, 0)
|
||||
button_container.mount(export_btn)
|
||||
|
||||
start_over_btn = Button("Start Over", id="start_over_btn")
|
||||
start_over_btn.styles.margin = (0, 0, 1, 0)
|
||||
button_container.mount(start_over_btn)
|
||||
|
||||
# Tables showing agents
|
||||
tables_container = Horizontal()
|
||||
tables_container.styles.height = "1fr"
|
||||
results_container.mount(tables_container)
|
||||
|
||||
# Enforce Ready table
|
||||
if ready_count > 0:
|
||||
enforce_col = Vertical()
|
||||
enforce_col.styles.width = "1fr"
|
||||
enforce_col.styles.margin = (1, 1, 0, 0)
|
||||
tables_container.mount(enforce_col)
|
||||
|
||||
enforce_label = Static("Enforce Ready Agents")
|
||||
enforce_label.styles.margin = (0, 0, 1, 0)
|
||||
enforce_col.mount(enforce_label)
|
||||
|
||||
enforce_table = DataTable(id="enforce_ready_table")
|
||||
enforce_table.styles.height = "1fr"
|
||||
enforce_table.add_columns("Hostname", "Last Exec (days)")
|
||||
|
||||
# Display first 50 agents
|
||||
for idx, row in self.enforce_ready_df.head(50).iterrows():
|
||||
days_since = row["days_since"]
|
||||
days_str = f"{int(days_since)}" if not pd.isna(days_since) else "Never"
|
||||
enforce_table.add_row(row["hostname"], days_str)
|
||||
|
||||
if len(self.enforce_ready_df) > 50:
|
||||
enforce_table.add_row(
|
||||
f"... and {len(self.enforce_ready_df) - 50} more", ""
|
||||
)
|
||||
|
||||
enforce_col.mount(enforce_table)
|
||||
|
||||
# Non-Enforce Ready table
|
||||
if not_ready_count > 0:
|
||||
non_enforce_col = Vertical()
|
||||
non_enforce_col.styles.width = "1fr"
|
||||
non_enforce_col.styles.margin = (1, 0, 0, 1)
|
||||
tables_container.mount(non_enforce_col)
|
||||
|
||||
non_enforce_label = Static("Non-Enforce Ready Agents")
|
||||
non_enforce_label.styles.margin = (0, 0, 1, 0)
|
||||
non_enforce_col.mount(non_enforce_label)
|
||||
|
||||
non_enforce_table = DataTable(id="non_enforce_ready_table")
|
||||
non_enforce_table.styles.height = "1fr"
|
||||
non_enforce_table.add_columns("Hostname", "Last Exec (days)")
|
||||
|
||||
# Display first 50 agents
|
||||
for idx, row in self.non_enforce_ready_df.head(50).iterrows():
|
||||
days_since = row["days_since"]
|
||||
days_str = f"{int(days_since)}" if not pd.isna(days_since) else "N/A"
|
||||
non_enforce_table.add_row(row["hostname"], days_str)
|
||||
|
||||
if len(self.non_enforce_ready_df) > 50:
|
||||
non_enforce_table.add_row(
|
||||
f"... and {len(self.non_enforce_ready_df) - 50} more", ""
|
||||
)
|
||||
|
||||
non_enforce_col.mount(non_enforce_table)
|
||||
|
||||
def _show_enforce_target_selection(self) -> None:
|
||||
"""Show policy selection for enforce ready agents."""
|
||||
self.workflow_stage = "select_enforce_target"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Info message
|
||||
info = Static(
|
||||
f"Select target policy for {len(self.enforce_ready_df)} Enforce Ready agents\n"
|
||||
f"Source Policy: {self.selected_policy.name}",
|
||||
id="enforce_target_info",
|
||||
)
|
||||
info.styles.margin = (0, 0, 2, 0)
|
||||
content.mount(info)
|
||||
|
||||
# Policy selector
|
||||
policy_selector = PolicySelector(self.policies)
|
||||
content.mount(policy_selector)
|
||||
|
||||
# Skip button
|
||||
skip_btn = Button("Skip - No Migration", id="skip_enforce_target_btn")
|
||||
skip_btn.styles.width = "50%"
|
||||
skip_btn.styles.margin = (2, 0, 0, 0)
|
||||
content.mount(skip_btn)
|
||||
|
||||
def _show_non_enforce_target_selection(self) -> None:
|
||||
"""Show policy selection for non-enforce ready agents."""
|
||||
self.workflow_stage = "select_non_enforce_target"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Info message
|
||||
info = Static(
|
||||
f"Select target policy for {len(self.non_enforce_ready_df)} Non-Enforce Ready agents\n"
|
||||
f"Source Policy: {self.selected_policy.name}",
|
||||
id="non_enforce_target_info",
|
||||
)
|
||||
info.styles.margin = (0, 0, 2, 0)
|
||||
content.mount(info)
|
||||
|
||||
# Policy selector
|
||||
policy_selector = PolicySelector(self.policies)
|
||||
content.mount(policy_selector)
|
||||
|
||||
# Skip button
|
||||
skip_btn = Button("Skip - No Migration", id="skip_non_enforce_target_btn")
|
||||
skip_btn.styles.width = "50%"
|
||||
skip_btn.styles.margin = (2, 0, 0, 0)
|
||||
content.mount(skip_btn)
|
||||
|
||||
def _show_migration_confirmation(self) -> None:
|
||||
"""Show migration confirmation screen."""
|
||||
self.workflow_stage = "confirm_migration"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Build confirmation message
|
||||
confirmation_lines = [
|
||||
"Migration Summary\n",
|
||||
f"Source Policy: {self.selected_policy.name}\n",
|
||||
]
|
||||
|
||||
if self.enforce_ready_target_policy:
|
||||
confirmation_lines.append(
|
||||
f"\nEnforce Ready Migration:\n"
|
||||
f"Agents: {len(self.enforce_ready_df)}\n"
|
||||
f"Target: {self.enforce_ready_target_policy.name}\n"
|
||||
)
|
||||
|
||||
if self.non_enforce_ready_target_policy:
|
||||
confirmation_lines.append(
|
||||
f"\nNon-Enforce Ready Migration:\n"
|
||||
f"Agents: {len(self.non_enforce_ready_df)}\n"
|
||||
f"Target: {self.non_enforce_ready_target_policy.name}\n"
|
||||
)
|
||||
|
||||
if (
|
||||
not self.enforce_ready_target_policy
|
||||
and not self.non_enforce_ready_target_policy
|
||||
):
|
||||
confirmation_lines.append("\nNo migrations will be performed.")
|
||||
|
||||
confirmation = Static("".join(confirmation_lines), id="migration_confirmation")
|
||||
confirmation.styles.margin = (1, 1, 2, 1)
|
||||
content.mount(confirmation)
|
||||
|
||||
# Action buttons - mount container first, then add buttons
|
||||
button_container = Horizontal(id="confirmation_buttons")
|
||||
button_container.styles.height = "auto"
|
||||
button_container.styles.margin = (1, 1)
|
||||
content.mount(button_container)
|
||||
|
||||
if self.enforce_ready_target_policy or self.non_enforce_ready_target_policy:
|
||||
confirm_btn = Button("Confirm Migration", id="confirm_migration_btn")
|
||||
confirm_btn.styles.margin = (0, 1, 0, 0)
|
||||
button_container.mount(confirm_btn)
|
||||
|
||||
cancel_btn = Button("Cancel", id="cancel_migration_btn")
|
||||
button_container.mount(cancel_btn)
|
||||
|
||||
def _execute_migration(self) -> None:
|
||||
"""Execute the agent migrations."""
|
||||
self.workflow_stage = "executing"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Show executing message
|
||||
executing_msg = Static(
|
||||
"Executing agent migrations...\nPlease wait...",
|
||||
id="executing_message",
|
||||
)
|
||||
executing_msg.styles.margin = (2, 1)
|
||||
content.mount(executing_msg)
|
||||
|
||||
# Perform migrations asynchronously
|
||||
self.call_later(self._perform_migrations)
|
||||
|
||||
def _perform_migrations(self) -> None:
|
||||
"""Perform the actual agent migrations."""
|
||||
successful_migrations = []
|
||||
failed_migrations = []
|
||||
|
||||
try:
|
||||
# Migrate enforce ready agents
|
||||
if self.enforce_ready_target_policy:
|
||||
for idx, row in self.enforce_ready_df.iterrows():
|
||||
try:
|
||||
self.api.agent_move(
|
||||
row["agentid"], self.enforce_ready_target_policy.groupid
|
||||
)
|
||||
successful_migrations.append(
|
||||
(row["hostname"], self.enforce_ready_target_policy.name)
|
||||
)
|
||||
logger.debug(
|
||||
f"Moved {row['hostname']} to {self.enforce_ready_target_policy.name}"
|
||||
)
|
||||
except Exception as e:
|
||||
failed_migrations.append((row["hostname"], str(e)))
|
||||
logger.error(f"Failed to move {row['hostname']}: {e}")
|
||||
|
||||
# Migrate non-enforce ready agents
|
||||
if self.non_enforce_ready_target_policy:
|
||||
for idx, row in self.non_enforce_ready_df.iterrows():
|
||||
try:
|
||||
self.api.agent_move(
|
||||
row["agentid"], self.non_enforce_ready_target_policy.groupid
|
||||
)
|
||||
successful_migrations.append(
|
||||
(row["hostname"], self.non_enforce_ready_target_policy.name)
|
||||
)
|
||||
logger.debug(
|
||||
f"Moved {row['hostname']} to {self.non_enforce_ready_target_policy.name}"
|
||||
)
|
||||
except Exception as e:
|
||||
failed_migrations.append((row["hostname"], str(e)))
|
||||
logger.error(f"Failed to move {row['hostname']}: {e}")
|
||||
|
||||
# Show completion results
|
||||
self._show_completion_results(successful_migrations, failed_migrations)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during migration execution: {e}", exc_info=True)
|
||||
self.app.notify(f"Migration failed: {str(e)}", severity="error", timeout=5)
|
||||
self._show_results()
|
||||
|
||||
def _show_completion_results(
|
||||
self, successful: List[tuple], failed: List[tuple]
|
||||
) -> None:
|
||||
"""Show migration completion results."""
|
||||
self.workflow_stage = "complete"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Results summary
|
||||
total_attempted = len(successful) + len(failed)
|
||||
success_rate = (
|
||||
(len(successful) / total_attempted * 100) if total_attempted > 0 else 0
|
||||
)
|
||||
|
||||
results = Static(
|
||||
f"Migration Complete\n\n"
|
||||
f"Total Agents Migrated: {len(successful)}\n"
|
||||
f"Failed Migrations: {len(failed)}\n"
|
||||
f"Success Rate: {success_rate:.1f}%",
|
||||
id="completion_summary",
|
||||
)
|
||||
results.styles.margin = (1, 1, 2, 1)
|
||||
content.mount(results)
|
||||
|
||||
# Details tables
|
||||
if successful:
|
||||
success_container = Vertical()
|
||||
success_container.styles.margin = (0, 1)
|
||||
content.mount(success_container)
|
||||
|
||||
success_label = Static("Successful Migrations")
|
||||
success_label.styles.margin = (0, 0, 1, 0)
|
||||
success_container.mount(success_label)
|
||||
|
||||
success_table = DataTable(id="success_table")
|
||||
success_table.styles.height = "auto"
|
||||
success_table.add_columns("Hostname", "Target Policy")
|
||||
|
||||
for hostname, target_policy in successful[:25]: # Show first 25
|
||||
success_table.add_row(hostname, target_policy)
|
||||
|
||||
if len(successful) > 25:
|
||||
success_table.add_row(f"... and {len(successful) - 25} more", "")
|
||||
|
||||
success_container.mount(success_table)
|
||||
|
||||
if failed:
|
||||
failed_container = Vertical()
|
||||
failed_container.styles.margin = (2, 1, 0, 1)
|
||||
content.mount(failed_container)
|
||||
|
||||
failed_label = Static("Failed Migrations")
|
||||
failed_label.styles.margin = (0, 0, 1, 0)
|
||||
failed_container.mount(failed_label)
|
||||
|
||||
failed_table = DataTable(id="failed_table")
|
||||
failed_table.styles.height = "auto"
|
||||
failed_table.add_columns("Hostname", "Error")
|
||||
|
||||
for hostname, error in failed[:25]: # Show first 25
|
||||
failed_table.add_row(hostname, error[:50]) # Truncate error
|
||||
|
||||
if len(failed) > 25:
|
||||
failed_table.add_row(f"... and {len(failed) - 25} more", "")
|
||||
|
||||
failed_container.mount(failed_table)
|
||||
|
||||
# Action button
|
||||
done_btn = Button("Done", id="start_over_btn")
|
||||
done_btn.styles.width = "50%"
|
||||
done_btn.styles.margin = (2, 0, 0, 0)
|
||||
content.mount(done_btn)
|
||||
|
||||
def _export_results(self) -> None:
|
||||
"""Export analysis results to CSV."""
|
||||
try:
|
||||
working_dir = load_env("WORKING_DIR") or os.getcwd()
|
||||
filename = os.path.join(
|
||||
working_dir,
|
||||
f"{self.selected_policy.name}_quiet_analysis_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
|
||||
)
|
||||
|
||||
self.agents_df.to_csv(filename, index=False)
|
||||
logger.info(f"Exported results to {filename}")
|
||||
self.app.notify(
|
||||
f"Results exported to:\n{filename}",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to export results: {e}")
|
||||
self.app.notify(f"Export failed: {str(e)}", severity="error", timeout=5)
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle back/escape action."""
|
||||
# Depending on stage, go back to previous stage or exit
|
||||
if self.workflow_stage in ["select_policy", "view_results", "complete"]:
|
||||
self.app.pop_screen()
|
||||
elif self.workflow_stage == "select_quiet_days":
|
||||
self._show_policy_selection()
|
||||
elif self.workflow_stage == "select_enforce_target":
|
||||
self._show_results()
|
||||
elif self.workflow_stage == "select_non_enforce_target":
|
||||
if self.enforce_ready_target_policy:
|
||||
self._show_enforce_target_selection()
|
||||
else:
|
||||
self._show_results()
|
||||
elif self.workflow_stage == "confirm_migration":
|
||||
self._show_non_enforce_target_selection()
|
||||
else:
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
Reference in New Issue
Block a user