849 lines
33 KiB
Python
849 lines
33 KiB
Python
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published
|
|
# by the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
"""
|
|
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.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"),
|
|
]
|
|
|
|
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)
|
|
|
|
back_btn = Button("← Back", id="back_to_policy_selection")
|
|
back_btn.styles.width = "100%"
|
|
back_btn.styles.margin = (2, 0, 0, 0)
|
|
button_container.mount(back_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 == "back_to_policy_selection":
|
|
self._show_policy_selection()
|
|
return
|
|
|
|
if button_id == "back_to_results":
|
|
self._show_results()
|
|
return
|
|
|
|
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"] == True].copy()
|
|
self.non_enforce_ready_df = agents[agents["enforce_ready"] == False].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"\n✅ Enforce 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"\n❌ Non-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("\n⚠️ No 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:
|
|
result = 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:
|
|
result = 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()
|