mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Quiet Agent UI improvements

Browse Source
This commit is contained in:
Zarithas
2025-12-05 15:05:46 -05:00
parent b19eeb6c96
commit 3ab803c12e
1 changed files with 183 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TUI/Screens/quietagentworkflowscreen.py
+183 -110
View File
@@ -34,7 +34,7 @@ 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 textual.widgets import Button, DataTable, Footer, Header, Input, Static
from models.policy import Policy
from services.API import AirlockAPIWrapper
@@ -51,16 +51,17 @@ class QuietAgentWorkflowScreen(Screen):
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
2. Configure analysis parameters (history period and quiet time period)
3. View categorized agents (enforce ready vs. non-enforce ready)
4. Select target policies for each category
5. 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)
history_days (int): Number of days of history to pull (default: 150, range: 1-365)
quiet_days (int): Number of days without execution to be considered quiet (default: 45, range: 1-365)
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
@@ -86,7 +87,7 @@ class QuietAgentWorkflowScreen(Screen):
self.api = api
self.policies = policies
self.selected_policy: Optional[Policy] = None
self.history_days = 150 # Fixed as per requirements
self.history_days = 150 # Default value, user-selectable
self.quiet_days = 45 # Default value
self.agents_df: Optional[pd.DataFrame] = None
self.enforce_ready_df: Optional[pd.DataFrame] = None
@@ -130,7 +131,7 @@ class QuietAgentWorkflowScreen(Screen):
stage_messages = {
"select_policy": "Step 1: Select Policy to Analyze",
"select_quiet_days": "Step 2: Select Quiet Time Period",
"select_history_days": "Step 2: Configure Analysis Parameters",
"analyzing": "Analyzing agent activity...",
"view_results": "Step 3: Review Categorized Agents",
"select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents",
@@ -161,7 +162,7 @@ class QuietAgentWorkflowScreen(Screen):
# 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()
self._show_history_days_selection()
elif self.workflow_stage == "select_enforce_target":
# Target policy selection for enforce ready agents
self.enforce_ready_target_policy = message.policy
@@ -177,48 +178,167 @@ class QuietAgentWorkflowScreen(Screen):
)
self._show_migration_confirmation()
def _show_quiet_days_selection(self) -> None:
"""Show the quiet days selection screen."""
self.workflow_stage = "select_quiet_days"
def _show_history_days_selection(self) -> None:
"""Show the history days and quiet days selection screen."""
self.workflow_stage = "select_history_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",
"Configure Analysis Parameters:",
id="analysis_params_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)
# Create input container
input_container = Vertical(id="analysis_params_input_container")
input_container.styles.height = "auto"
content.mount(input_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",
# History days label
history_label = Static("History Period (days of execution history to pull):")
history_label.styles.margin = (0, 0, 1, 0)
input_container.mount(history_label)
# Add history days input field
history_input = Input(
placeholder="Enter days (1-365, default: 150)",
value="150",
id="history_days_input",
)
btn.styles.width = "100%"
btn.styles.margin = (0, 0, 1, 0)
button_container.mount(btn)
history_input.styles.width = "50"
history_input.styles.margin = (0, 0, 2, 0)
input_container.mount(history_input)
# Quiet days label
quiet_label = Static(
"Quiet Time Period (days without execution to be considered quiet):"
)
quiet_label.styles.margin = (0, 0, 1, 0)
input_container.mount(quiet_label)
# Add quiet days input field
quiet_input = Input(
placeholder="Enter days (1-365, default: 45)",
value="45",
id="quiet_days_input",
)
quiet_input.styles.width = "50"
quiet_input.styles.margin = (0, 0, 2, 0)
input_container.mount(quiet_input)
# Add submit button
submit_btn = Button(
"Continue",
id="analysis_params_submit",
variant="primary",
)
submit_btn.styles.width = "50"
submit_btn.styles.margin = (1, 0, 0, 0)
input_container.mount(submit_btn)
# Focus the first input field
history_input.focus()
def _validate_and_submit_history_days(self) -> None:
"""Validate and submit the history days and quiet days inputs."""
try:
history_input = self.query_one("#history_days_input", Input)
quiet_input = self.query_one("#quiet_days_input", Input)
history_value = history_input.value.strip()
quiet_value = quiet_input.value.strip()
# Validate history days
if not history_value:
self.app.notify(
"Please enter a history period value", severity="error", timeout=3
)
history_input.focus()
return
try:
history_days = int(history_value)
except ValueError:
self.app.notify(
"Please enter a valid number for history period",
severity="error",
timeout=3,
)
history_input.focus()
return
if history_days < 1 or history_days > 365:
self.app.notify(
"History period must be between 1 and 365 days",
severity="error",
timeout=3,
)
history_input.focus()
return
# Validate quiet days
if not quiet_value:
self.app.notify(
"Please enter a quiet time period value",
severity="error",
timeout=3,
)
quiet_input.focus()
return
try:
quiet_days = int(quiet_value)
except ValueError:
self.app.notify(
"Please enter a valid number for quiet time period",
severity="error",
timeout=3,
)
quiet_input.focus()
return
if quiet_days < 1 or quiet_days > 365:
self.app.notify(
"Quiet time period must be between 1 and 365 days",
severity="error",
timeout=3,
)
quiet_input.focus()
return
# Check that quiet days doesn't exceed history days
if quiet_days > history_days:
self.app.notify(
"Quiet time period cannot exceed history period",
severity="error",
timeout=3,
)
quiet_input.focus()
return
# All validation passed
self.history_days = history_days
self.quiet_days = quiet_days
logger.info(
f"Selected history days: {history_days}, quiet days: {quiet_days}"
)
self._start_analysis()
except Exception as e:
logger.error(f"Error validating analysis parameters: {e}")
self.app.notify(f"Error: {str(e)}", severity="error", timeout=3)
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()
# Analysis parameters submit button
if button_id == "analysis_params_submit":
self._validate_and_submit_history_days()
return
# Navigation buttons
@@ -258,46 +378,44 @@ class QuietAgentWorkflowScreen(Screen):
self._show_policy_selection()
return
def on_input_submitted(self, event: Input.Submitted) -> None:
"""Handle input submission (Enter key pressed)."""
if event.input.id in ["history_days_input", "quiet_days_input"]:
self._validate_and_submit_history_days()
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
# Show notification that analysis is starting
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)
# Clear the screen to provide a blank canvas for Rust progress output
# (Rust output displays over the TUI, so we clear everything except header/footer)
try:
# Clear title
title_widget = self.query_one("#workflow_title", Static)
title_widget.update("")
def _perform_analysis(self) -> None:
# Clear status
status_widget = self.query_one("#workflow_status", Static)
status_widget.update("")
# Clear content area
content = self.query_one("#content_area", Vertical)
content.remove_children()
except Exception as e:
logger.debug(f"Could not clear screen for analysis: {e}")
# Delay the analysis start to ensure UI refresh completes first
# This prevents Rust output from starting before the screen is cleared
self.set_timer(0.2, self._perform_analysis_worker)
def _perform_analysis_worker(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)
@@ -310,32 +428,11 @@ class QuietAgentWorkflowScreen(Screen):
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."
@@ -381,9 +478,6 @@ class QuietAgentWorkflowScreen(Screen):
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]
@@ -394,7 +488,7 @@ class QuietAgentWorkflowScreen(Screen):
# Categorize agents into DataFrames
self.enforce_ready_df = agents[agents["enforce_ready"]].copy()
self.non_enforce_ready_df = agents[not agents["enforce_ready"]].copy()
self.non_enforce_ready_df = agents[~agents["enforce_ready"]].copy()
logger.info(
f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, "
@@ -416,27 +510,6 @@ class QuietAgentWorkflowScreen(Screen):
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"
@@ -821,7 +894,7 @@ class QuietAgentWorkflowScreen(Screen):
# 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":
elif self.workflow_stage == "select_history_days":
self._show_policy_selection()
elif self.workflow_stage == "select_enforce_target":
self._show_results()