From 3ab803c12e30d129d3c604779c848572f3950949 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Fri, 5 Dec 2025 15:05:46 -0500 Subject: [PATCH] Quiet Agent UI improvements --- TUI/Screens/quietagentworkflowscreen.py | 293 +++++++++++++++--------- 1 file changed, 183 insertions(+), 110 deletions(-) diff --git a/TUI/Screens/quietagentworkflowscreen.py b/TUI/Screens/quietagentworkflowscreen.py index f0d5bbb..1532b1d 100644 --- a/TUI/Screens/quietagentworkflowscreen.py +++ b/TUI/Screens/quietagentworkflowscreen.py @@ -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", + ) + 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}" ) - btn.styles.width = "100%" - btn.styles.margin = (0, 0, 1, 0) - button_container.mount(btn) + 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()