# 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 . import datetime import logging import os import re from typing import Dict, List, Optional import pandas as pd from textual.app import ComposeResult from textual.binding import Binding from textual.containers import Horizontal, Vertical from textual.reactive import reactive from textual.screen import Screen from textual.widgets import Button, DataTable, Footer, Header, Input, Static from models.execution import ExecutionHistoryRecord from models.policy import Allowlist, Policy from services.API import AirlockAPIWrapper from TUI.Widgets.policyselector import PolicySelector from utils.configmanager import get_system_list, get_system_value, load_env from utils.utils import formatHTML logger = logging.getLogger(__name__) class PolicyPrepWorkflowScreen(Screen): """ A Textual screen for the Policy Preparation workflow. This screen provides a multi-step workflow: 1. Select source policies to gather execution data from 2. Select destination policy and associated allowlist 3. Fetch and sort execution history 4. Manual review of approved/needs_review files 5. Generate path exclusions and publisher lists 6. Second manual review of paths/publishers 7. Test - preview changes 8. Liftoff - apply changes Attributes: api (AirlockAPIWrapper): API wrapper for Airlock operations policies (List[Policy]): List of all available policies source_policies (List[Policy]): Selected source policies destination_policy (Optional[Policy]): Destination policy destination_allowlist (Optional[Allowlist]): Associated allowlist workflow_stage (str): Current stage of the workflow working_dir (str): Working directory for exports """ DEFAULT_CSS = """ DataTable > .datatable--row.selected { background: $primary 30%; } DataTable:focus > .datatable--cursor { background: $secondary 20%; } #workflow_title { text-style: bold; color: $text; } #workflow_status { color: $accent; } #checklist_area { max-height: 12; margin: 0 1; } #content_area { height: 1fr; overflow-y: auto; scrollbar-gutter: stable; padding: 1 1; } Horizontal { height: auto; min-height: 3; } Button { min-width: 15; } Button.variant-error { background: $error; color: $text; } Button.variant-success { background: $success; color: $text; } """ BINDINGS = [ Binding("escape", "go_back", "Back"), Binding("q", "main_menu", "Main Menu"), Binding("f", "open_folder", "Open Folder"), Binding("d", "delete_rows", "Delete Selected"), Binding("a", "select_all", "Select All"), Binding("n", "select_none", "Select None"), Binding("space", "toggle_selection", "Toggle Selection", show=False), # Note: 'r' key handled in on_key() for range selection mode ] workflow_stage = reactive("select_source") # Tracks current workflow stage def __init__(self, api: AirlockAPIWrapper, policies: List[Policy]): """ Initialize the PolicyPrepWorkflowScreen. 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.source_policies: List[Policy] = [] self.destination_policy: Optional[Policy] = None self.destination_allowlist: Optional[Allowlist] = None self.working_dir = load_env("WORKING_DIR") or os.getcwd() self.history_days: Optional[int] = None self.path_split: str = "\\" # Path separator for Windows (single backslash) # Data storage self.approved_df: Optional[pd.DataFrame] = None self.needs_review_df: Optional[pd.DataFrame] = None self.unapproved_df: Optional[pd.DataFrame] = None self.primary_paths_df: Optional[pd.DataFrame] = None self.secondary_paths_df: Optional[pd.DataFrame] = None self.publishers_df: Optional[pd.DataFrame] = None self.remaining_hashes_df: Optional[pd.DataFrame] = None # Test data for preview self.test_results: Optional[Dict] = None # Multi-select tracking self.last_clicked_row: Optional[str] = None self.last_clicked_table: Optional[str] = None # Range selection mode (activated by 'r' key) self._range_mode = False # Track if we're navigating with keyboard (to prevent selection) self._keyboard_navigation = False # Tab review tracking for Step 5 (First Review) self.approved_tab_reviewed = False self.needs_review_tab_reviewed = False # Tab review tracking for Step 6 (Path Review) self.paths_tab_reviewed = False self.publishers_tab_reviewed = False # Lock to prevent concurrent table creation self._creating_review_table = False self._creating_path_table = False def compose(self) -> ComposeResult: """Build the UI layout for the workflow screen.""" yield Header(show_clock=True, icon="⚙️") # Title area title = Static("Policy Preparation Workflow", id="workflow_title") title.styles.text_align = "center" title.styles.margin = (0, 0, 0, 1) yield title # Status area status = Static("Step 1: Select Source Policies", id="workflow_status") status.styles.margin = (0, 0, 1, 1) yield status # Main content area - dynamically populated based on workflow stage yield Vertical(id="content_area") # Checklist area - always visible yield Vertical(id="checklist_area") yield Footer() def on_mount(self) -> None: """Initialize the screen when mounted.""" self._update_checklist() self._show_introduction() 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() self._update_checklist() 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_source": "Step 1: Select Source Policies", "select_destination": "Step 2: Select Destination Policy", "select_allowlist": "Step 3: Select Destination Allowlist", "fetch_data": "Step 4: Fetch Execution History", "fetching": "Fetching and sorting execution data...", "first_review": "Step 5: First Manual Review", "build_paths": "Step 6: Building Path Exclusions", "second_review": "Step 7: Second Manual Review", "test": "Step 8: Test - Preview Changes", "liftoff": "Step 9: Liftoff - Apply Changes", "complete": "Workflow Complete", } status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage")) def _update_checklist(self) -> None: """Update the preparation checklist display.""" checklist = self.query_one("#checklist_area", Vertical) checklist.remove_children() # Checklist container with border checklist_container = Vertical() checklist_container.styles.border = ("round", "blue") checklist_container.styles.margin = (0, 1) checklist_container.styles.padding = (0, 1) # Mount the container to the checklist area FIRST checklist.mount(checklist_container) # Title checklist_title = Static("Prep Checklist") checklist_title.styles.text_style = "bold" checklist_title.styles.text_align = "center" checklist_container.mount(checklist_title) # Create two-column layout row1 = Horizontal() row1.styles.height = "auto" checklist_container.mount(row1) col1 = Vertical() col1.styles.width = "50%" col2 = Vertical() col2.styles.width = "50%" row1.mount(col1) row1.mount(col2) # Step 1: Source Policies step1_status = "✓" if self.source_policies else "✖" step1_count = f" ({len(self.source_policies)})" if self.source_policies else "" step1 = Static(f"{step1_status} Source{step1_count}") if self.source_policies: step1.styles.color = "green" else: step1.styles.text_style = "dim" col1.mount(step1) # Step 2: Destination Policy step2_status = "✓" if self.destination_policy else "✖" step2 = Static(f"{step2_status} Destination") if self.destination_policy: step2.styles.color = "green" else: step2.styles.text_style = "dim" col1.mount(step2) # Step 3: Allowlist step3_status = "✓" if self.destination_allowlist else "✖" step3 = Static(f"{step3_status} Allowlist") if self.destination_allowlist: step3.styles.color = "green" else: step3.styles.text_style = "dim" col1.mount(step3) # Step 4: Data Fetched data_fetched = self.approved_df is not None or self.needs_review_df is not None step4_status = "✓" if data_fetched else "✖" if data_fetched: total = 0 if self.approved_df is not None: total += len(self.approved_df) if self.needs_review_df is not None: total += len(self.needs_review_df) step4 = Static(f"{step4_status} Data ({total})") else: step4 = Static(f"{step4_status} Data") if data_fetched: step4.styles.color = "green" else: step4.styles.text_style = "dim" col1.mount(step4) # Step 5: First Review first_review_path = os.path.join(self.working_dir, "Approved") first_review_done = False if self.source_policies: approved_file = os.path.join( first_review_path, f"{self.source_policies[0].name}_approved_executions.csv", ) first_review_done = os.path.exists(approved_file) step5_status = "✓" if first_review_done else "✖" step5 = Static(f"{step5_status} Review 1") if first_review_done: step5.styles.color = "green" else: step5.styles.text_style = "dim" col2.mount(step5) # Step 6: Paths Generated paths_generated = self.primary_paths_df is not None step6_status = "✓" if paths_generated else "✖" if paths_generated: step6 = Static(f"{step6_status} Paths ({len(self.primary_paths_df)})") else: step6 = Static(f"{step6_status} Paths") if paths_generated: step6.styles.color = "green" else: step6.styles.text_style = "dim" col2.mount(step6) # Step 7: Second Review step7_status = ( "✓" if self.workflow_stage in ["test", "liftoff", "complete"] else "✖" ) step7 = Static(f"{step7_status} Review 2") if step7_status == "✓": step7.styles.color = "green" else: step7.styles.text_style = "dim" col2.mount(step7) # Step 8: Tested step8_status = "✓" if self.workflow_stage in ["liftoff", "complete"] else "✖" step8 = Static(f"{step8_status} Tested") if step8_status == "✓": step8.styles.color = "green" else: step8.styles.text_style = "dim" col2.mount(step8) def _show_introduction(self) -> None: """Show workflow introduction and overview.""" self.workflow_stage = "introduction" content = self.query_one("#content_area", Vertical) content.remove_children() # Title title = Static("Welcome to Policy Preparation Workflow") title.styles.margin = (1, 1) title.styles.text_style = "bold" title.styles.text_align = "center" content.mount(title) # Description description = Static( "This workflow will help you:\n" " • Fetch execution history from selected policies\n" " • Review and approve safe executions\n" " • Calculate efficient path exclusions\n" " • Generate publisher trust rules\n" " • Apply changes to your destination policy" ) description.styles.margin = (1, 2) content.mount(description) # Process steps steps_title = Static("The Process:") steps_title.styles.margin = (1, 2, 0, 2) steps_title.styles.text_style = "bold" content.mount(steps_title) steps = Static( " 📋 Step 1: Select source policies (data collection)\n" " 🎯 Step 2: Select destination policy (where changes go)\n" " 📝 Step 3: Select destination allowlist\n" " 📊 Step 4: Fetch execution data (may take 1-2 minutes)\n" " ✅ Step 5: Review approved/needs review executions\n" " 📁 Step 6: Review path exclusions and publishers\n" " 🔍 Step 7: Preview changes before applying\n" " 🚀 Step 8: Liftoff - Apply to production" ) steps.styles.margin = (0, 2) content.mount(steps) # Time estimate estimate = Static("⏱️ Estimated Time: 15-30 minutes depending on data size") estimate.styles.margin = (1, 2) estimate.styles.color = "cyan" content.mount(estimate) # Tips tips_title = Static("💡 Tips:") tips_title.styles.margin = (1, 2, 0, 2) tips_title.styles.text_style = "bold" content.mount(tips_title) tips = Static( " • Start with a test policy first\n" " • Review carefully - changes affect all agents\n" " • Use path rules when possible (more efficient)\n" " • Publishers are powerful - use cautiously" ) tips.styles.margin = (0, 2) tips.styles.color = "yellow" content.mount(tips) # Buttons button_container = Horizontal() button_container.styles.margin = (2, 2) button_container.styles.align = ("center", "middle") content.mount(button_container) continue_btn = Button( "Continue to Policy Selection", id="start_workflow", variant="success" ) cancel_btn = Button("Cancel", id="cancel_workflow", variant="default") button_container.mount(continue_btn) button_container.mount(cancel_btn) def _show_source_policy_selection(self) -> None: """Show the source policy selection screen.""" self.workflow_stage = "select_source" content = self.query_one("#content_area", Vertical) content.remove_children() instruction = Static("Select source policies (click rows to toggle selection):") instruction.styles.margin = (0, 1, 0, 1) content.mount(instruction) # Create a DataTable for multi-select table = DataTable(id="source_policy_table") table.styles.height = "30vh" table.styles.overflow_y = "auto" table.cursor_type = "row" table.zebra_stripes = True # Add columns - checkbox first, then data columns table.add_columns("○", "Name", "ID", "Parent") # Sort policies by name for easier selection sorted_policies = sorted(self.policies, key=lambda p: p.name.lower()) # Add rows for policy in sorted_policies: # Skip parent policies if policy.parent == "global-policy-settings": continue checkbox = "○" # All start unchecked table.add_row( checkbox, policy.name, str(policy.groupid), policy.parent or "N/A", key=str(policy.groupid), ) content.mount(table) # Control buttons control_container = Horizontal() control_container.styles.height = "auto" control_container.styles.margin = (0, 1) # Mount the container first content.mount(control_container) # Then add buttons to it select_none_btn = Button("Clear Selection", id="select_none_source") select_none_btn.styles.width = "1fr" select_none_btn.styles.margin = (0, 1, 0, 0) continue_btn = Button( "→ Continue", id="continue_source_selection", variant="primary" ) continue_btn.styles.width = "1fr" continue_btn.styles.margin = (0, 0, 0, 1) control_container.mount(select_none_btn) control_container.mount(continue_btn) # Track selected policies if not hasattr(self, "selected_source_policy_ids"): self.selected_source_policy_ids = set() else: self.selected_source_policy_ids.clear() def _show_destination_policy_selection(self) -> None: """Show the destination policy selection screen.""" self.workflow_stage = "select_destination" content = self.query_one("#content_area", Vertical) content.remove_children() instruction = Static( f"Selected Source: {', '.join([p.name for p in self.source_policies])}\n\n" "Select the destination policy for enforcement:" ) instruction.styles.margin = (0, 1, 1, 1) content.mount(instruction) # Create policy selector with policies sorted alphabetically by name sorted_policies = sorted(self.policies, key=lambda p: p.name.lower()) policy_selector = PolicySelector(sorted_policies) content.mount(policy_selector) def _show_allowlist_selection(self) -> None: """Show the allowlist selection screen.""" self.workflow_stage = "select_allowlist" content = self.query_one("#content_area", Vertical) content.remove_children() instruction = Static( f"Destination Policy: {self.destination_policy.name}\n\n" "Select the allowlist to use:" ) instruction.styles.margin = (0, 1, 1, 1) content.mount(instruction) # Fetch allowlists for the destination policy try: allowlists_df = self.api.policy_list_allowlists( self.destination_policy.groupid ) allowlists = [ Allowlist(**row.to_dict()) for _, row in allowlists_df.iterrows() ] if not allowlists: content.mount( Static("No allowlists found for this policy!", id="no_allowlists") ) return # Sort allowlists alphabetically by name allowlists = sorted(allowlists, key=lambda al: al.name.lower()) # Add instruction instruction = Static("Click a row to select the allowlist for this policy:") instruction.styles.margin = (0, 1, 1, 1) content.mount(instruction) # Create table for allowlist selection table = DataTable(id="allowlist_table") table.styles.height = "auto" table.styles.max_height = "50%" table.cursor_type = "row" table.add_columns("ID", "Name", "Version") for al in allowlists: table.add_row(str(al.applicationid), al.name, str(al.version)) content.mount(table) # Store allowlists for reference self.allowlists = allowlists except Exception as e: logger.error(f"Failed to fetch allowlists: {e}") content.mount(Static(f"Error fetching allowlists: {str(e)}")) def _show_fetch_data(self) -> None: """Show the data fetching options screen.""" self.workflow_stage = "fetch_data" content = self.query_one("#content_area", Vertical) content.remove_children() # Validate we have source policies before showing this screen if not hasattr(self, "source_policies") or not self.source_policies: logger.error("_show_fetch_data called but no source_policies set!") self.app.notify( "Error: No source policies selected. Returning to policy selection.", severity="error", ) self._show_source_policy_selection() return instruction = Static( f"Ready to fetch execution history from: {', '.join([p.name for p in self.source_policies])}\n\n" "Enter the number of days of history to fetch (or press Enter to use default 150):" ) instruction.styles.margin = (0, 1, 1, 1) content.mount(instruction) # Days input days_container = Horizontal() days_container.styles.margin = (1, 1) days_container.styles.height = "auto" content.mount(days_container) days_label = Static("History Days (1-365): ") days_label.styles.width = "auto" days_input = Input( value="150", placeholder="150", id="history_days_input", type="integer" ) days_input.styles.width = 30 days_input.styles.min_width = 20 days_container.mount(days_label) days_container.mount(days_input) # Type selection type_instruction = Static("\nSelect execution types to include:") type_instruction.styles.margin = (1, 1, 0, 1) content.mount(type_instruction) type_info = Static( "Default: Types 1, 2, 6, 7 (Standard executions)\n" "You can customize this if needed." ) type_info.styles.margin = (0, 1, 1, 1) type_info.styles.text_style = "dim" content.mount(type_info) # Fetch button button_container = Horizontal() button_container.styles.margin = (2, 1) content.mount(button_container) fetch_btn = Button("Fetch Data", id="fetch_data_btn", variant="primary") fetch_btn.styles.margin = (0, 1, 0, 0) skip_btn = Button("Skip (Use Existing)", id="skip_fetch_btn") button_container.mount(fetch_btn) button_container.mount(skip_btn) # Set focus to the input field so it's ready for typing def focus_input(): try: days_input.focus() except Exception as e: logger.debug(f"Could not focus input: {e}") self.call_after_refresh(focus_input) def _fetch_execution_data(self, history_days: int) -> None: """Fetch and sort execution data.""" self.workflow_stage = "fetching" # Show notification that fetch is starting self.app.notify( "Starting data fetch - this may take several minutes for large policies", severity="information", timeout=5, ) # 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("") # 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 fetch: {e}") # Delay the fetch start to ensure UI refresh completes first # This prevents Rust output from starting before the screen is cleared self.set_timer(0.5, lambda: self._perform_fetch(history_days)) def _perform_fetch(self, history_days: int) -> None: """Perform the actual data fetching.""" try: # Validate we have source policies logger.info(f"_perform_fetch called with history_days={history_days}") logger.info( f"self.source_policies exists: {hasattr(self, 'source_policies')}" ) if hasattr(self, "source_policies"): logger.info(f"self.source_policies type: {type(self.source_policies)}") logger.info( f"self.source_policies length: {len(self.source_policies) if self.source_policies else 0}" ) if self.source_policies: logger.info( f"First policy: {self.source_policies[0].name if self.source_policies else 'N/A'}" ) if ( not hasattr(self, "source_policies") or not self.source_policies or len(self.source_policies) == 0 ): logger.error( f"No source policies selected. hasattr={hasattr(self, 'source_policies')}, value={getattr(self, 'source_policies', 'ATTR_MISSING')}" ) self.app.notify("No source policies selected!", severity="error") self._show_fetch_data() return logger.info( f"Starting fetch for {len(self.source_policies)} policies, {history_days} days of history" ) # Fetch execution history policy_executions = ExecutionHistoryRecord.from_policies( self.api, self.source_policies, type_=[1, 2, 6, 7], history_days=history_days, ) logger.info(f"Fetched {len(policy_executions)} total execution records") if not policy_executions: logger.warning("No execution records returned from API") self.app.notify( f"No execution history found for the last {history_days} days", severity="warning", ) self._show_fetch_data() return # Enrich with hash data logger.info("Enriching executions with hash data...") enriched_executions = ExecutionHistoryRecord.enrich_with_hashes( self.api, policy_executions ) # Categorize by hash decision logger.info("Categorizing executions by hash decision...") categorized_executions = ( ExecutionHistoryRecord.categorize_executions_by_hash_decision( enriched_executions ) ) # Sort by decision logger.info("Sorting executions by decision...") approved, unapproved, needs_review, unknown = ( ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions) ) logger.info( f"Sorted: {len(approved)} approved, {len(unapproved)} unapproved, " f"{len(needs_review)} needs review, {len(unknown)} unknown" ) # Store the data self.approved_df = ( pd.DataFrame([r.__dict__ for r in approved]) if approved else pd.DataFrame() ) self.unapproved_df = ( pd.DataFrame([r.__dict__ for r in unapproved]) if unapproved else pd.DataFrame() ) self.needs_review_df = ( pd.DataFrame([r.__dict__ for r in needs_review]) if needs_review else pd.DataFrame() ) # Save to files logger.info("Saving fetched data to files...") self._save_fetched_data() # Show results logger.info("Showing results...") self._show_fetch_results() except Exception as e: logger.error(f"Failed to fetch execution data: {e}", exc_info=True) self.app.notify(f"Failed to fetch data: {str(e)}", severity="error") self._show_fetch_data() def _save_fetched_data(self) -> None: """Save fetched data to CSV and HTML files.""" if not self.source_policies: return policy_name = self.source_policies[0].name review_dir = os.path.join(self.working_dir, "Needs_Review", "Review_First") html_dir = os.path.join(self.working_dir, "Needs_Review", "HTML") os.makedirs(review_dir, exist_ok=True) os.makedirs(html_dir, exist_ok=True) # Save each category categories = { "approved": self.approved_df, "needs_review": self.needs_review_df, "unapproved": self.unapproved_df, } for label, df in categories.items(): if df is not None and not df.empty: csv_path = os.path.join( review_dir, f"{policy_name}_{label}_executions.csv" ) html_path = os.path.join(html_dir, f"{policy_name}_{label}.html") df.to_csv(csv_path, index=False) formatHTML(df, html_path) logger.info(f"Saved {label} executions to {csv_path}") def _show_fetch_results(self) -> None: """Show the results of data fetching.""" logger.debug("=== _show_fetch_results called ===") self.workflow_stage = "first_review" content = self.query_one("#content_area", Vertical) content.remove_children() logger.info( f"Approved count: {len(self.approved_df) if self.approved_df is not None else 0}" ) logger.info( f"Needs review count: {len(self.needs_review_df) if self.needs_review_df is not None else 0}" ) logger.info( f"Unapproved count: {len(self.unapproved_df) if self.unapproved_df is not None else 0}" ) # Results summary approved_count = len(self.approved_df) if self.approved_df is not None else 0 review_count = ( len(self.needs_review_df) if self.needs_review_df is not None else 0 ) unapproved_count = ( len(self.unapproved_df) if self.unapproved_df is not None else 0 ) summary = Static( f"Data Fetch Complete!\n\n" f"Approved: {approved_count} executions\n" f"Needs Review: {review_count} executions\n" f"Unapproved: {unapproved_count} executions (automatically excluded)\n" ) summary.styles.margin = (1, 1) content.mount(summary) logger.info("Mounted summary widget") # Tab selection for review tab_container = Horizontal() tab_container.styles.margin = (1, 1) content.mount(tab_container) approved_tab_btn = Button( "Review Approved", id="show_approved_tab", variant="primary" ) approved_tab_btn.styles.margin = (0, 1, 0, 0) review_tab_btn = Button("Review Needs Review", id="show_needs_review_tab") tab_container.mount(approved_tab_btn) tab_container.mount(review_tab_btn) logger.info("Mounted tab buttons") # Show approved table by default logger.debug("About to call _show_review_table('approved')") self._show_review_table("approved") logger.debug("=== _show_fetch_results complete ===") def _show_review_table(self, table_type: str) -> None: """Show an editable DataTable for reviewing executions.""" # Prevent concurrent execution if self._creating_review_table: logger.warning( f"Already creating review table, ignoring duplicate call for {table_type}" ) return self._creating_review_table = True try: self._show_review_table_impl(table_type) finally: self._creating_review_table = False def _show_review_table_impl(self, table_type: str) -> None: """Internal implementation of _show_review_table.""" logger.debug(f"_show_review_table called with type: {table_type}") content = self.query_one("#content_area", Vertical) # Mark tab as reviewed if table_type == "approved": self.approved_tab_reviewed = True logger.debug("Marked approved tab as reviewed") else: self.needs_review_tab_reviewed = True logger.debug("Marked needs_review tab as reviewed") # Determine which dataframe and table ID to show if table_type == "approved": df = self.approved_df table_id = "approved_review_table" title = "Approved Executions - Select rows to REMOVE:" else: df = self.needs_review_df table_id = "needs_review_table" title = "Needs Review Executions - Select rows to REMOVE:" # Sort DataFrame by filename (case-insensitive) and save back if df is not None and not df.empty and "filename" in df.columns: df = df.sort_values(by="filename", key=lambda x: x.str.lower()) # Save sorted DataFrame back if table_type == "approved": self.approved_df = df else: self.needs_review_df = df # Remove the SPECIFIC table we're about to create if it exists try: existing_specific = content.query_one(f"#{table_id}", DataTable) if existing_specific: logger.debug(f"Removing existing table with ID: {table_id}") existing_specific.remove() except Exception as e: logger.debug( f"No existing table with ID {table_id} found (this is normal): {e}" ) # Remove ALL existing DataTables to be safe try: existing_tables = content.query("DataTable") logger.debug(f"Found {len(existing_tables)} existing tables to remove") for table in existing_tables: logger.debug(f"Removing table: {table.id}") table.remove() except Exception as e: logger.debug(f"Error removing existing tables: {e}") # Remove existing instruction and help text (they accumulate without removal) # Remove ALL Static widgets - they're just text that needs to be replaced try: existing_statics = content.query("Static") logger.debug( f"Found {len(existing_statics)} existing Static widgets to remove" ) for static in existing_statics: static.remove() except Exception as e: logger.debug(f"Error removing Static widgets: {e}") # Force a refresh to ensure removals are processed try: content.refresh() except Exception as e: logger.debug(f"Error refreshing content: {e}") # Note: We no longer remove review_controls or review_continue_container # They are reused between tabs to avoid DuplicateIds errors logger.debug( f"DataFrame for {table_type}: {'empty' if df is None or df.empty else f'{len(df)} rows'}" ) if df is None or df.empty: empty_msg = Static(f"No {table_type} executions to review") empty_msg.styles.margin = (2, 1) content.mount(empty_msg) logger.info(f"No data for {table_type}, mounted empty message") return # Instructions (no ID needed - we remove all Statics anyway) instruction = Static(title) instruction.styles.margin = (1, 1) instruction.styles.text_style = "bold" content.mount(instruction) # Help text (no ID needed) help_text = Static( "Click to toggle, 'r' for range select (click start, press 'r', click end)\n" "Space to toggle cursor row, 'd' to delete, 'a' select all, arrows navigate" ) help_text.styles.margin = (0, 1, 1, 1) help_text.styles.text_style = "dim" content.mount(help_text) # CRITICAL: Check if table already exists in content (should not happen after removal above) try: existing_check = content.query_one(f"#{table_id}", DataTable) if existing_check: logger.error( f"Table {table_id} STILL EXISTS after removal! This should not happen." ) # Don't create a new one - just return return except Exception: # Good - table doesn't exist, proceed with creation pass # Create the review table review_table = DataTable(id=table_id) review_table.styles.height = "40vh" # Increased since we removed button rows review_table.cursor_type = "row" review_table.zebra_stripes = True # Specified columns in order important_cols = [ "policyname", "policyver", "hostname", "username", "publisher", "filename", "pprocess", "gprocess", "sha256", "commandline", ] available_cols = [col for col in important_cols if col in df.columns] if available_cols: # Add checkbox column first review_table.add_columns("○", *available_cols) # Add rows with row keys for tracking for idx, row in df.iterrows(): checkbox = "○" # All start unchecked row_data = [str(row.get(col, "")) for col in available_cols] review_table.add_row(checkbox, *row_data, key=str(idx)) logger.debug(f"About to mount {table_id}") # Final safety check - make sure no table with this ID exists before mounting try: final_check = content.query_one(f"#{table_id}", DataTable) if final_check: logger.warning( f"Table {table_id} still exists after removal attempts! Forcing removal..." ) final_check.remove() content.refresh() except Exception: # Good - table doesn't exist pass content.mount(review_table) logger.debug(f"Successfully mounted {table_id} with {len(df)} rows") # Row count display only (removed Select All, Clear, Delete buttons) try: control_container = content.query_one("#review_controls", Horizontal) # Clear existing content control_container.remove_children() except Exception: # Doesn't exist, create it control_container = Horizontal(id="review_controls") control_container.styles.margin = (1, 1) control_container.styles.height = "auto" control_container.styles.min_height = 1 content.mount(control_container) row_count = Static(f"Total rows: {len(df)}") row_count.styles.margin = (0, 1, 0, 1) control_container.mount(row_count) # Continue button (always at bottom) if not content.query("#review_continue_container"): continue_container = Horizontal(id="review_continue_container") continue_container.styles.margin = (2, 1, 1, 1) continue_container.styles.height = "auto" continue_container.styles.min_height = 3 # Removed dock="bottom" - was hiding content above # Mount the container to the content area FIRST content.mount(continue_container) # NOW mount buttons into the container export_btn = Button("Export to CSV", id="export_review") export_btn.styles.margin = (0, 1, 0, 0) continue_btn = Button( "→ Finish Review & Continue", id="continue_from_review", variant="success", ) continue_container.mount(export_btn) continue_container.mount(continue_btn) # Track selected rows if not hasattr(self, "selected_rows"): self.selected_rows = set() else: self.selected_rows.clear() # Store current review type self.current_review_type = table_type def _delete_selected_rows(self) -> None: """Delete selected rows from the current dataframe.""" if not hasattr(self, "selected_rows") or not self.selected_rows: self.app.notify("No rows selected for deletion", severity="warning") return # Determine which dataframe to modify if self.current_review_type == "approved": df = self.approved_df table_id = "approved_review_table" else: df = self.needs_review_df table_id = "needs_review_table" if df is None: return # Get the table try: content = self.query_one("#content_area", Vertical) table = content.query_one(f"#{table_id}", DataTable) except Exception as e: logger.error(f"Could not find table {table_id}: {e}") return # Get indices to delete indices_to_delete = [int(idx) for idx in self.selected_rows] # Remove rows from DataFrame df_filtered = df.drop(index=indices_to_delete, errors="ignore") # Update the dataframe if self.current_review_type == "approved": self.approved_df = df_filtered else: self.needs_review_df = df_filtered # Remove rows from DataTable (don't rebuild entire table) removed_count = 0 failed_keys = [] for idx in self.selected_rows: try: # Try to remove the row using the key table.remove_row(idx) removed_count += 1 except Exception as e: # Log but continue - some keys might not exist after DataFrame operations logger.debug(f"Could not remove row {idx}: {e}") failed_keys.append(idx) # Clear selection self.selected_rows.clear() # Notify user if failed_keys: self.app.notify( f"Deleted {removed_count} rows ({len(failed_keys)} already removed)", severity="information", ) else: self.app.notify(f"Deleted {removed_count} rows", severity="information") def _show_path_building_screen(self) -> None: """Show loading screen before building paths and publishers.""" self.workflow_stage = "building_paths" content = self.query_one("#content_area", Vertical) content.remove_children() # Add spacer to push text to bottom spacer = Static("") spacer.styles.height = "1fr" content.mount(spacer) # Loading message at bottom (above checklist) status = Static( "Building path exclusions and publisher lists...\n" "This may take a moment for large datasets." ) status.styles.margin = (1, 1) status.styles.text_align = "center" status.styles.color = "cyan" content.mount(status) # Force UI refresh to show the loading screen content.refresh() # Schedule the actual build to happen after UI updates # Using set_timer with a small delay ensures the screen renders self.set_timer(0.1, self._perform_path_build) def _build_paths_and_publishers(self) -> None: """Build path exclusions and publisher lists.""" # Note: This is now bypassed - we go straight from _show_path_building_screen to _perform_path_build self.call_later(self._perform_path_build) def _perform_path_build(self) -> None: """Perform the actual path and publisher building.""" try: if not self.source_policies: raise ValueError("No source policies selected") policy_name = self.source_policies[0].name approved_path = os.path.join( self.working_dir, "Approved", f"{policy_name}_approved_executions.csv" ) review_path = os.path.join( self.working_dir, "Approved", f"{policy_name}_needs_review_executions.csv", ) # Load approved files df1 = ( pd.read_csv(approved_path) if os.path.exists(approved_path) else pd.DataFrame() ) df2 = ( pd.read_csv(review_path) if os.path.exists(review_path) else pd.DataFrame() ) if df1.empty and df2.empty: self.app.notify( "No approved files found! Please complete first review.", severity="error", ) self._show_fetch_results() return # Combine dataframes all_approved = pd.concat([df1, df2], ignore_index=True) if "filename" in all_approved.columns: all_approved = all_approved.sort_values(by="filename") # Calculate paths path_exclusion_const = get_system_value( "PATH_EXCLUSION_CONST", cast_type=int ) if path_exclusion_const: # Primary paths self.primary_paths_df = self._calculate_paths( all_approved, path_exclusion_const ) # Secondary paths if ( not self.primary_paths_df.empty and "sha256" in self.primary_paths_df.columns ): remaining = all_approved[ ~all_approved["sha256"].isin(self.primary_paths_df["sha256"]) ] self.secondary_paths_df = self._calculate_paths( remaining, path_exclusion_const - 1 ) else: # If primary paths are empty, all remaining go to secondary logger.warning( "Primary paths DataFrame is empty or missing sha256 column" ) self.secondary_paths_df = pd.DataFrame() remaining = all_approved # Remaining hashes if ( not self.secondary_paths_df.empty and "sha256" in self.secondary_paths_df.columns ): self.remaining_hashes_df = remaining[ ~remaining["sha256"].isin(self.secondary_paths_df["sha256"]) ] else: logger.warning( "Secondary paths DataFrame is empty or missing sha256 column" ) self.remaining_hashes_df = remaining # Extract publishers if not all_approved.empty: publist = all_approved[ all_approved["publisher"] != "Not Signed" ].drop_duplicates(subset=["publisher"]) # Remove bad publishers bad_publishers = get_system_list("BAD_PUBLISHERS") if bad_publishers: pattern = "|".join(bad_publishers) publist = publist[ ~publist["publisher"].str.contains( pattern, case=False, na=False, regex=True ) ] self.publishers_df = publist # Save to Review_Second folder self._save_path_data() # Show results self._show_path_results() except Exception as e: logger.error(f"Failed to build paths: {e}", exc_info=True) self.app.notify(f"Failed to build paths: {str(e)}", severity="error") self._show_fetch_results() def _regulator(self, string_list: List[str], case_insensitive: bool = True) -> str: """ Create a regex pattern from a list of strings. Args: string_list: List of strings to create pattern from case_insensitive: Whether to make pattern case insensitive Returns: Regex pattern string that matches any of the input strings """ if not string_list: return "" # Escape special regex characters in each string escaped = [re.escape(s) for s in string_list] # Join with | (OR operator) pattern = "|".join(escaped) return pattern def _split_filepaths_grouped( self, df: pd.DataFrame, path_exclusion_constant: int, col: str = "filename" ) -> pd.DataFrame: """ Split filepaths, group by common prefix, and extract metadata. This is a port of the splitFilepathsGrouped function from prepPolicy.py. Args: df: DataFrame with filepath column path_exclusion_constant: Depth for path truncation col: Column name containing filepaths Returns: DataFrame with columns: longestcfp, middle, filename_only, file_extension, plus all original columns """ min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int) def clean_split(path): """Split a path into parts, handling various input types.""" if not isinstance(path, (str, bytes, os.PathLike)): return [] parts = str(os.path.normpath(path)).split(os.sep) parts = [p for p in parts if p] # Remove empty strings return parts # Check for non-string entries non_string_entries = df[ ~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike))) ] if not non_string_entries.empty: logger.warning( f"Non-string entries found in column '{col}': {len(non_string_entries)}" ) df = df.copy() split_paths = df[col].apply(clean_split) # Filter by minimum path length if configured if min_files_for_path is not None: df = df[ split_paths.apply(lambda parts: len(parts) >= min_files_for_path) ].copy() split_paths = split_paths[df.index] # Group by path prefix df["group_key"] = split_paths.apply( lambda parts: os.sep.join(parts[:path_exclusion_constant]) ) grouped = df.groupby("group_key") new_rows = [] for _, group_df in grouped: paths = group_df[col].tolist() split_parts = [clean_split(p) for p in paths] def longest_common_prefix(paths): """Find the longest common prefix among a list of path parts.""" if not paths: return [] prefix = paths[0] for path in paths[1:]: prefix = [a for a, b in zip(prefix, path) if a == b] if not prefix: break return prefix common_prefix = longest_common_prefix(split_parts) prefix_str = os.sep.join(common_prefix) # Process each file in the group for i, parts in enumerate(split_parts): filename = parts[-1] middle = ( os.sep.join(parts[len(common_prefix) : -1]) if len(parts) > len(common_prefix) + 1 else "" ) row = group_df.iloc[i].copy() row["longestcfp"] = prefix_str row["middle"] = middle row["filename_only"] = filename row["file_extension"] = os.path.splitext(filename)[1].lower() new_rows.append(row) result = pd.DataFrame(new_rows).drop(columns=["group_key"]) logger.info( f"_split_filepaths_grouped: Processed {len(df)} rows → {len(result)} rows with metadata" ) return result def _calculate_paths( self, df: pd.DataFrame, path_exclusion_constant: int ) -> pd.DataFrame: """ Calculate path exclusions with full metadata including extensions and hash counts. This is a port of the calculatePath function from prepPolicy.py. Args: df: DataFrame with execution data path_exclusion_constant: Depth for path truncation Returns: DataFrame with columns: policyname, longestcfp, middle, filename_only, file_extension, sha256, unique_sha256_count """ logger.info( f"=== _calculate_paths called with path_exclusion_constant={path_exclusion_constant} ===" ) logger.debug(f"Input DataFrame: {len(df)} rows") logger.debug(f"Columns: {list(df.columns) if not df.empty else 'empty'}") if df.empty: logger.warning("Input DataFrame is empty") return pd.DataFrame() # Use 'filename' column (which typically contains full path) if "filename" not in df.columns: logger.error("'filename' column not found in DataFrame") return pd.DataFrame() # Split filepaths and extract metadata haslcp = self._split_filepaths_grouped(df, path_exclusion_constant, "filename") haslcp = haslcp.drop_duplicates() logger.debug(f"After split_filepaths_grouped: {len(haslcp)} rows") # Filter forbidden paths badpathparts = get_system_list("BAD_PATH_PARTS") if badpathparts: forbidden_pattern = self._regulator(badpathparts, True) forbidden_lcfp = haslcp["longestcfp"].str.contains( forbidden_pattern, case=False, na=False, regex=True ) logger.debug( f"Removing forbidden filepaths: {forbidden_lcfp.sum()} paths filtered" ) lcp_not_forbidden = haslcp[~forbidden_lcfp].copy() else: logger.info( "No BAD_PATH_PARTS configured, skipping forbidden path filtering" ) lcp_not_forbidden = haslcp.copy() logger.debug(f"After forbidden filtering: {len(lcp_not_forbidden)} rows") # Select relevant columns if "policyname" in lcp_not_forbidden.columns: columns_to_keep = [ "policyname", "longestcfp", "middle", "filename_only", "file_extension", "sha256", ] else: # If no policyname, skip it columns_to_keep = [ "longestcfp", "middle", "filename_only", "file_extension", "sha256", ] # Only keep columns that exist columns_to_keep = [ col for col in columns_to_keep if col in lcp_not_forbidden.columns ] lcp_not_forbidden_review = lcp_not_forbidden[columns_to_keep] # Count unique SHA256s per path unique_sha_counts = ( lcp_not_forbidden_review.groupby("longestcfp")["sha256"] .nunique() .reset_index() ) unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"] logger.info( f"Calculated unique SHA256 counts for {len(unique_sha_counts)} paths" ) # Merge counts back into main DataFrame lcp_not_forbidden_review = lcp_not_forbidden_review.merge( unique_sha_counts, on="longestcfp", how="left" ) # Filter by minimum files per path min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int) if min_files_for_path is not None: before_filter = len(lcp_not_forbidden_review) lcp_not_forbidden_review = lcp_not_forbidden_review[ lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path ] logger.debug( f"After MIN_FILES_FOR_PATH filter ({min_files_for_path}): {len(lcp_not_forbidden_review)} rows (removed {before_filter - len(lcp_not_forbidden_review)})" ) logger.debug( f"Final result: {len(lcp_not_forbidden_review)} rows with columns: {list(lcp_not_forbidden_review.columns)}" ) return lcp_not_forbidden_review def _save_path_data(self) -> None: """Save path and publisher data to files.""" if not self.source_policies: return policy_name = self.source_policies[0].name review_dir = os.path.join(self.working_dir, "Needs_Review", "Review_Second") html_dir = os.path.join(self.working_dir, "Needs_Review", "HTML") os.makedirs(review_dir, exist_ok=True) os.makedirs(html_dir, exist_ok=True) # Save each dataframe dataframes = { "primary_paths": self.primary_paths_df, "secondary_paths": self.secondary_paths_df, "publishers": self.publishers_df, "remaining_hashes": self.remaining_hashes_df, } for name, df in dataframes.items(): if df is not None and not df.empty: csv_path = os.path.join(review_dir, f"{policy_name}_{name}.csv") html_path = os.path.join(html_dir, f"{policy_name}_{name}.html") df.to_csv(csv_path, index=False) formatHTML(df, html_path) logger.info(f"Saved {name} to {csv_path}") def _show_path_results(self) -> None: """Show the results of path building.""" logger.debug("=== _show_path_results called ===") self.workflow_stage = "second_review" content = self.query_one("#content_area", Vertical) content.remove_children() # Results summary primary_count = ( len(self.primary_paths_df) if self.primary_paths_df is not None else 0 ) secondary_count = ( len(self.secondary_paths_df) if self.secondary_paths_df is not None else 0 ) publishers_count = ( len(self.publishers_df) if self.publishers_df is not None else 0 ) logger.info(f"Primary paths: {primary_count}") logger.info(f"Secondary paths: {secondary_count}") logger.info(f"Publishers: {publishers_count}") summary = Static( f"Path Analysis Complete!\n\n" f"Primary Paths: {primary_count}\n" f"Secondary Paths: {secondary_count}\n" f"Publishers: {publishers_count}" ) summary.styles.margin = (1, 1) content.mount(summary) logger.info("Mounted summary widget") # Tab selection for different review types tab_container = Horizontal() tab_container.styles.margin = (1, 1) content.mount(tab_container) paths_tab_btn = Button("Review Paths", id="show_paths_tab", variant="primary") paths_tab_btn.styles.margin = (0, 1, 0, 0) publishers_tab_btn = Button("Review Publishers", id="show_publishers_tab") publishers_tab_btn.styles.margin = (0, 1, 0, 0) remaining_tab_btn = Button("Remaining Hashes", id="show_remaining_tab") tab_container.mount(paths_tab_btn) tab_container.mount(publishers_tab_btn) tab_container.mount(remaining_tab_btn) logger.info("Mounted tab buttons") # Show paths table by default logger.debug("About to call _show_path_review_table('paths')") self._show_path_review_table("paths") logger.debug("=== _show_path_results complete ===") def _show_path_review_table(self, table_type: str) -> None: """Show an editable DataTable for reviewing paths/publishers.""" # Prevent concurrent execution if self._creating_path_table: logger.warning( f"Already creating path table, ignoring duplicate call for {table_type}" ) return self._creating_path_table = True try: self._show_path_review_table_impl(table_type) finally: self._creating_path_table = False def _show_path_review_table_impl(self, table_type: str) -> None: """Internal implementation of _show_path_review_table.""" logger.debug(f"_show_path_review_table called with type: {table_type}") content = self.query_one("#content_area", Vertical) # Mark tab as reviewed (only paths and publishers, not remaining) if table_type == "paths": self.paths_tab_reviewed = True logger.debug("Marked paths tab as reviewed") elif table_type == "publishers": self.publishers_tab_reviewed = True logger.debug("Marked publishers tab as reviewed") # Determine which dataframe to show if table_type == "paths": # Combine primary and secondary paths for review dfs = [] if self.primary_paths_df is not None and not self.primary_paths_df.empty: df_copy = self.primary_paths_df.copy() df_copy["type"] = "primary" dfs.append(df_copy) if ( self.secondary_paths_df is not None and not self.secondary_paths_df.empty ): df_copy = self.secondary_paths_df.copy() df_copy["type"] = "secondary" dfs.append(df_copy) if dfs: df = pd.concat(dfs, ignore_index=True) # Aggregate by path to show one row per path if not df.empty and "longestcfp" in df.columns: # Group by longestcfp and type, aggregate extensions aggregated_rows = [] for (path, path_type), group in df.groupby(["longestcfp", "type"]): # Get unique extensions and hash count extensions = ( group["file_extension"].unique() if "file_extension" in group.columns else [] ) extensions_str = ", ".join( sorted(set(ext for ext in extensions if ext)) ) # Get hash count (should be same for all rows with same longestcfp) hash_count = ( group["unique_sha256_count"].iloc[0] if "unique_sha256_count" in group.columns else 0 ) aggregated_rows.append( { "longestcfp": path, "file_extension": extensions_str, "unique_sha256_count": hash_count, "type": path_type, } ) df = pd.DataFrame(aggregated_rows) logger.info(f"Aggregated paths: {len(df)} unique paths") else: df = pd.DataFrame() table_id = "paths_review_table" title = "Path Exclusions - Select paths to REMOVE:" # Show: path, extensions, hash count, type (primary/secondary) columns = ( ["longestcfp", "file_extension", "unique_sha256_count", "type"] if not df.empty else [] ) elif table_type == "publishers": df = self.publishers_df table_id = "publishers_review_table" title = "Approved Publishers - Select publishers to REMOVE:" columns = ["publisher"] if df is not None and not df.empty else [] else: # remaining df = self.remaining_hashes_df table_id = "remaining_review_table" title = "Remaining Hashes (not covered by paths) - For reference only:" columns = ( ["filename", "filepath", "sha256"] if df is not None and not df.empty else [] ) # Remove the SPECIFIC table we're about to create if it exists try: existing_specific = content.query_one(f"#{table_id}", DataTable) if existing_specific: logger.debug(f"Removing existing table with ID: {table_id}") existing_specific.remove() except Exception as e: logger.debug( f"No existing table with ID {table_id} found (this is normal): {e}" ) # Remove any other existing tables try: existing_tables = content.query("DataTable") logger.debug(f"Found {len(existing_tables)} existing tables to remove") for table in existing_tables: table.remove() except Exception as e: logger.debug(f"Error removing existing tables: {e}") # Remove existing controls try: existing_controls = content.query("#path_review_controls") logger.debug(f"Found {len(existing_controls)} existing controls to remove") for control in existing_controls: control.remove() except Exception as e: logger.debug(f"Error removing existing controls: {e}") # Remove existing instruction and help text (they accumulate without removal) # Remove ALL Static widgets - they're just text that needs to be replaced try: existing_statics = content.query("Static") logger.debug( f"Found {len(existing_statics)} existing Static widgets to remove" ) for static in existing_statics: static.remove() except Exception as e: logger.debug(f"Error removing Static widgets: {e}") # Force refresh to ensure removals complete try: content.refresh() except Exception as e: logger.debug(f"Error refreshing content: {e}") if df is None or df.empty: empty_msg = Static(f"No {table_type} to review") empty_msg.styles.margin = (2, 1) content.mount(empty_msg) return # Instructions (no ID needed) instruction = Static(title) instruction.styles.margin = (1, 1) instruction.styles.text_style = "bold" content.mount(instruction) # Help text (different for remaining hashes, no ID needed) if table_type != "remaining": help_text = Static( "Click to toggle, 'r' for range select (click start, press 'r', click end)\n" "Space to toggle cursor row, 'd' to delete, 'a' select all, arrows navigate" ) else: help_text = Static( "These hashes cannot be approved via path exclusions.\n" "They will need individual hash approval if required." ) help_text.styles.margin = (0, 1, 1, 1) help_text.styles.text_style = "dim" content.mount(help_text) # Create the review table review_table = DataTable(id=table_id) review_table.styles.height = "35vh" # Increased since we removed button rows review_table.cursor_type = "row" review_table.zebra_stripes = True # Add columns - checkbox first, then data columns if columns: # For DataFrames, also check what columns actually exist available_cols = [col for col in columns if col in df.columns] if available_cols: # Add checkbox column first review_table.add_columns("○", *available_cols) # Add rows with row keys for tracking for idx, row in df.iterrows(): checkbox = "○" # All start unchecked row_data = [] for col in available_cols: value = row.get(col, "") # Format unique_sha256_count with commas if col == "unique_sha256_count" and isinstance( value, (int, float) ): value = f"{int(value):,}" row_data.append(str(value)) review_table.add_row(checkbox, *row_data, key=str(idx)) logger.debug(f"About to mount {table_id}") # Final safety check before mounting table try: final_check = content.query_one(f"#{table_id}", DataTable) if final_check: logger.warning(f"Table {table_id} still exists! Forcing removal...") final_check.remove() content.refresh() except Exception: pass content.mount(review_table) logger.debug(f"Successfully mounted {table_id}") # Row count display only (removed Select All, Clear, Delete buttons) if table_type != "remaining": # Check if controls container already exists, reuse if it does try: control_container = content.query_one( "#path_review_controls", Horizontal ) # Clear existing content control_container.remove_children() logger.debug("Reusing existing path_review_controls container") except Exception: # Doesn't exist, create it control_container = Horizontal(id="path_review_controls") control_container.styles.margin = (1, 1) control_container.styles.height = "auto" control_container.styles.min_height = 1 content.mount(control_container) logger.debug("Created new path_review_controls container") row_count = Static(f"Total items: {len(df)}") row_count.styles.margin = (0, 1, 0, 1) control_container.mount(row_count) # Continue button (always at bottom) try: continue_container = content.query_one( "#path_continue_container", Horizontal ) # Container exists, check if buttons exist try: export_btn = continue_container.query_one("#export_path_review", Button) continue_btn = continue_container.query_one("#build_preflight", Button) logger.debug("Reusing existing path_continue_container with buttons") # Buttons already exist, just reuse them except Exception: # Container exists but buttons don't, clear and create new continue_container.remove_children() logger.debug("Reusing container, creating new buttons") export_btn = Button("Export to CSV", id="export_path_review") export_btn.styles.margin = (0, 1, 0, 0) continue_btn = Button( "Build Preflight", id="build_preflight", variant="success" ) continue_container.mount(export_btn) continue_container.mount(continue_btn) except Exception: # Container doesn't exist, create it with buttons continue_container = Horizontal(id="path_continue_container") continue_container.styles.margin = (2, 1, 1, 1) continue_container.styles.height = "auto" continue_container.styles.min_height = 3 content.mount(continue_container) logger.debug("Created new path_continue_container") # Create and mount buttons export_btn = Button("Export to CSV", id="export_path_review") export_btn.styles.margin = (0, 1, 0, 0) continue_btn = Button( "Build Preflight", id="build_preflight", variant="success" ) continue_container.mount(export_btn) continue_container.mount(continue_btn) # Track selected rows if not hasattr(self, "selected_path_rows"): self.selected_path_rows = set() else: self.selected_path_rows.clear() # Store current review type self.current_path_review_type = table_type def _delete_selected_path_rows(self) -> None: """Delete selected rows from the current path/publisher dataframe.""" if not hasattr(self, "selected_path_rows") or not self.selected_path_rows: self.app.notify("No rows selected for deletion", severity="warning") return indices_to_delete = [int(idx) for idx in self.selected_path_rows] # Get the appropriate table if self.current_path_review_type == "paths": table_id = "paths_review_table" elif self.current_path_review_type == "publishers": table_id = "publishers_review_table" else: table_id = "remaining_review_table" # Get the table try: content = self.query_one("#content_area", Vertical) table = content.query_one(f"#{table_id}", DataTable) except Exception as e: logger.error(f"Could not find table {table_id}: {e}") return # Determine which dataframe to modify if self.current_path_review_type == "paths": # Need to handle primary and secondary paths # For simplicity, rebuild both dataframes # This is a simplified approach - in production you'd track which type each row belongs to if self.primary_paths_df is not None: self.primary_paths_df = self.primary_paths_df.drop( index=[ i for i in indices_to_delete if i < len(self.primary_paths_df) ], errors="ignore", ) if self.secondary_paths_df is not None: offset = ( len(self.primary_paths_df) if self.primary_paths_df is not None else 0 ) self.secondary_paths_df = self.secondary_paths_df.drop( index=[i - offset for i in indices_to_delete if i >= offset], errors="ignore", ) elif self.current_path_review_type == "publishers": if self.publishers_df is not None: self.publishers_df = self.publishers_df.drop( index=indices_to_delete, errors="ignore" ) # Remove rows from DataTable (don't rebuild entire table) for idx in self.selected_path_rows: try: table.remove_row(idx) except Exception as e: logger.debug(f"Could not remove row {idx}: {e}") # Clear selection self.selected_path_rows.clear() self.app.notify( f"Deleted {len(indices_to_delete)} items", severity="information" ) def _build_preflight(self) -> None: """Build preflight files for testing.""" try: # This would contain the logic to build the final preflight files # For now, we'll just show the test screen self._show_test_screen() except Exception as e: logger.error(f"Failed to build preflight: {e}") self.app.notify(f"Failed to build preflight: {str(e)}", severity="error") def _show_test_screen(self) -> None: """Show the test/preview screen with detailed path listings.""" self.workflow_stage = "test" content = self.query_one("#content_area", Vertical) content.remove_children() summary = Static( "Test Mode - Preview Changes\n\n" "Review the paths that will be added to your policy:" ) summary.styles.margin = (1, 1) summary.styles.text_style = "bold" content.mount(summary) # Policy and Allowlist info info_text = "" if self.destination_policy: info_text += f"📋 Policy: {self.destination_policy.name}\n" if self.destination_allowlist: info_text += f"📋 Allowlist: {self.destination_allowlist.name}\n" if info_text: info = Static(info_text) info.styles.margin = (0, 1, 1, 1) content.mount(info) # Show detailed path exclusions if self.primary_paths_df is not None and not self.primary_paths_df.empty: self._show_path_preview( content, "Primary Path Exclusions", self.primary_paths_df ) if self.secondary_paths_df is not None and not self.secondary_paths_df.empty: self._show_path_preview( content, "Secondary Path Exclusions", self.secondary_paths_df ) # Show publishers if self.publishers_df is not None and not self.publishers_df.empty: pub_title = Static( f"\n📝 Trusted Publishers ({len(self.publishers_df)} publishers):" ) pub_title.styles.margin = (1, 1, 0, 1) pub_title.styles.text_style = "bold" content.mount(pub_title) # Create scrollable table for publishers pub_table = DataTable(id="publisher_preview_table") pub_table.styles.height = "15vh" pub_table.styles.margin = (0, 1) pub_table.cursor_type = "row" pub_table.zebra_stripes = True pub_table.add_column("Publisher") for _, row in self.publishers_df.iterrows(): pub_table.add_row(row["publisher"]) content.mount(pub_table) # Show hash count if self.approved_df is not None and not self.approved_df.empty: hash_info = Static( f"\n🔐 Individual Hash Approvals: {len(self.approved_df):,} hashes\n" f" (Files not covered by paths or publishers)" ) hash_info.styles.margin = (1, 1) content.mount(hash_info) # Warning warning = Static( "\n⚠️ WARNING: These changes cannot be easily undone. ⚠️\n" "Please review all paths carefully before proceeding." ) warning.styles.margin = (1, 1) warning.styles.color = "yellow" warning.styles.text_style = "bold" content.mount(warning) # Buttons button_container = Horizontal() button_container.styles.margin = (2, 1) content.mount(button_container) back_btn = Button( "← Back to Review", id="back_to_path_review", variant="default" ) liftoff_btn = Button( "Liftoff - Apply Changes 🚀", id="liftoff", variant="success" ) button_container.mount(back_btn) button_container.mount(liftoff_btn) def _show_path_preview( self, content: Vertical, title: str, paths_df: pd.DataFrame ) -> None: """Show a preview of paths that will be added.""" # Aggregate paths by longestcfp to get unique paths with all extensions path_list = [] for path, group in paths_df.groupby("longestcfp"): # Get all unique extensions for this path extensions = ( group["file_extension"].unique() if "file_extension" in group.columns else [] ) extensions = sorted(set(ext for ext in extensions if ext)) # Create path rules for each extension for ext in extensions: # Format the path as it will appear in Airlock # Example: C:\Program Files\App\**.exe formatted_path = f"{path}\\**{ext}" # Get hash count for this specific path+extension combo hash_count = ( len(group[group["file_extension"] == ext]) if "file_extension" in group.columns else 0 ) path_list.append((formatted_path, hash_count)) # Show title with count path_title = Static(f"\n📁 {title} ({len(path_list)} path rules):") path_title.styles.margin = (1, 1, 0, 1) path_title.styles.text_style = "bold" content.mount(path_title) # Create scrollable table path_table = DataTable(id=f"{title.lower().replace(' ', '_')}_table") path_table.styles.height = "20vh" path_table.styles.margin = (0, 1) path_table.cursor_type = "row" path_table.zebra_stripes = True path_table.add_columns("Path Rule", "Files Covered") # Add rows for path_rule, hash_count in sorted(path_list): path_table.add_row(path_rule, str(hash_count)) content.mount(path_table) def _apply_changes(self) -> None: """Apply the changes to policies and allowlists.""" self.workflow_stage = "liftoff" content = self.query_one("#content_area", Vertical) content.remove_children() status = Static("Applying changes...\nPlease wait...") status.styles.margin = (2, 1) content.mount(status) self.call_later(self._perform_apply) def _perform_apply(self) -> None: """Perform the actual application of changes.""" try: results = [] errors = [] # Apply path exclusions to policy (primary + secondary) if self.destination_policy: path_rules = [] # Process primary paths if ( self.primary_paths_df is not None and not self.primary_paths_df.empty ): logger.info( f"Processing {len(self.primary_paths_df)} primary paths" ) for _, row in self.primary_paths_df.groupby( ["longestcfp", "file_extension"] ): path = row.iloc[0]["longestcfp"] ext = row.iloc[0]["file_extension"] # Format: C:\Path\**.ext path_rule = f"{path}\\**{ext}" path_rules.append(path_rule) # Process secondary paths if ( self.secondary_paths_df is not None and not self.secondary_paths_df.empty ): logger.info( f"Processing {len(self.secondary_paths_df)} secondary paths" ) for _, row in self.secondary_paths_df.groupby( ["longestcfp", "file_extension"] ): path = row.iloc[0]["longestcfp"] ext = row.iloc[0]["file_extension"] path_rule = f"{path}\\**{ext}" path_rules.append(path_rule) # Apply path rules to policy if path_rules: try: logger.info( f"Applying {len(path_rules)} path exclusions to policy {self.destination_policy.name}" ) response = self.api.policy_add_path_exclusions( str(self.destination_policy.groupid), path_rules ) results.append( f"✓ Added {len(path_rules)} path exclusions to policy" ) logger.info(f"Path exclusions applied successfully: {response}") except Exception as e: error_msg = f"✗ Failed to add path exclusions: {str(e)}" errors.append(error_msg) logger.error(error_msg, exc_info=True) # Apply publishers to policy if ( self.destination_policy and self.publishers_df is not None and not self.publishers_df.empty ): try: publishers = self.publishers_df["publisher"].unique().tolist() logger.info( f"Applying {len(publishers)} publishers to policy {self.destination_policy.name}" ) response = self.api.policy_add_publishers( str(self.destination_policy.groupid), publishers ) results.append( f"✓ Added {len(publishers)} trusted publishers to policy" ) logger.info(f"Publishers applied successfully: {response}") except Exception as e: error_msg = f"✗ Failed to add publishers: {str(e)}" errors.append(error_msg) logger.error(error_msg, exc_info=True) # Apply hashes to allowlist if ( self.destination_allowlist and self.approved_df is not None and not self.approved_df.empty ): try: # Get unique hashes hashes = self.approved_df["sha256"].unique().tolist() logger.info( f"Applying {len(hashes)} hashes to allowlist {self.destination_allowlist.name}" ) response = self.api.hash_add_to_allowlist( str(self.destination_allowlist.applicationid), hashes ) results.append( f"✓ Added {len(hashes):,} approved hashes to allowlist" ) logger.info(f"Hashes applied successfully: {response}") except Exception as e: error_msg = f"✗ Failed to add hashes: {str(e)}" errors.append(error_msg) logger.error(error_msg, exc_info=True) # Show completion with both results and errors all_results = results + errors self._show_completion(all_results, has_errors=len(errors) > 0) except Exception as e: logger.error(f"Critical failure in _perform_apply: {e}", exc_info=True) self.app.notify(f"Critical failure: {str(e)}", severity="error") self._show_test_screen() def _show_completion(self, results: List[str], has_errors: bool = False) -> None: """Show completion screen.""" self.workflow_stage = "complete" content = self.query_one("#content_area", Vertical) content.remove_children() # Title depends on whether there were errors if has_errors: title_text = "Policy Preparation Completed with Errors\n\n" "Results:" title_color = "yellow" else: title_text = ( "Policy Preparation Complete!\n\n" "The following changes have been applied:" ) title_color = "green" summary = Static(title_text) summary.styles.margin = (1, 1) summary.styles.text_style = "bold" summary.styles.color = title_color content.mount(summary) for result in results: result_widget = Static(f" {result}") result_widget.styles.margin = (0, 2) # Color based on success/failure if result.startswith("✓"): result_widget.styles.color = "green" elif result.startswith("✗"): result_widget.styles.color = "red" content.mount(result_widget) # Final message if has_errors: final = Static( f"\n⚠️ Policy '{self.destination_policy.name}' was partially updated.\n" "Please review errors above and retry failed operations manually." ) final.styles.margin = (2, 1) final.styles.color = "yellow" else: final = Static( f"\n✅ Policy '{self.destination_policy.name}' is now ready for enforcement!" ) final.styles.margin = (2, 1) final.styles.color = "green" content.mount(final) # Done button done_btn = Button("Done", id="workflow_done") done_btn.styles.margin = (2, 0, 0, 0) done_btn.styles.width = "50%" content.mount(done_btn) # Event handlers def on_policy_selector_policy_selected( self, message: PolicySelector.PolicySelected ) -> None: """Handle policy selection from PolicySelector widget.""" if self.workflow_stage == "select_destination": self.destination_policy = message.policy logger.info(f"Selected destination policy: {self.destination_policy.name}") self._show_allowlist_selection() def _refresh_table_checkboxes(self, table_id: str, selected_keys: set) -> None: """Refresh checkbox column in a table based on selected keys.""" try: table = self.query_one(f"#{table_id}", DataTable) # Update checkboxes in place without rebuilding the table row_index = 0 for row_key in table.rows.keys(): # Get the actual value from the RowKey object row_key_str = ( str(row_key.value) if hasattr(row_key, "value") else str(row_key) ) # Determine if this row should be checked is_selected = row_key_str in selected_keys checkbox = "✓" if is_selected else "○" # Update the checkbox cell (first column, index 0) try: table.update_cell_at((row_index, 0), checkbox) except Exception as e: logger.error(f"Could not update cell at row {row_index}: {e}") row_index += 1 except Exception as e: logger.error(f"Error refreshing table {table_id}: {e}", exc_info=True) def _get_selected_set(self, table_id: str) -> set: """Get the appropriate selection set for a table.""" if table_id == "source_policy_table": return self.selected_source_policy_ids elif table_id in ["approved_review_table", "needs_review_table"]: return self.selected_rows elif table_id in [ "paths_review_table", "publishers_review_table", "remaining_review_table", ]: return self.selected_path_rows return set() def _range_select(self, table: DataTable, start_key: str, end_key: str) -> None: """Toggle all rows between start and end (inclusive).""" # Get all row keys in order all_keys = [ str(k.value if hasattr(k, "value") else k) for k in table.rows.keys() ] try: start_idx = all_keys.index(start_key) end_idx = all_keys.index(end_key) except ValueError: # Key not found, fall back to single toggle logger.warning("Range select failed: keys not found") return # Ensure start < end if start_idx > end_idx: start_idx, end_idx = end_idx, start_idx # Toggle all rows in range selected_set = self._get_selected_set(table.id) range_keys = [all_keys[i] for i in range(start_idx, end_idx + 1)] # Determine if we're selecting or deselecting # If any row in range is unselected, select all; otherwise deselect all any_unselected = any(key not in selected_set for key in range_keys) if any_unselected: # Select all in range for row_key in range_keys: selected_set.add(row_key) action = "selected" else: # Deselect all in range for row_key in range_keys: selected_set.discard(row_key) action = "deselected" # Refresh display self._refresh_table_checkboxes(table.id, selected_set) # Notify user count = end_idx - start_idx + 1 self.app.notify(f"Range {action} ({count} rows)", timeout=2) def _toggle_single(self, table: DataTable, row_key: str) -> None: """Toggle a single row without affecting others (Ctrl+Click).""" selected_set = self._get_selected_set(table.id) # Toggle if row_key in selected_set: selected_set.remove(row_key) else: selected_set.add(row_key) # Refresh self._refresh_table_checkboxes(table.id, selected_set) self.app.notify(f"Selected {len(selected_set)} rows", timeout=1) def on_data_table_row_highlighted(self, event: DataTable.RowHighlighted) -> None: """Handle row highlighting (clicking) in data tables with range selection support.""" table = event.data_table row_key = ( str(event.row_key.value) if hasattr(event.row_key, "value") else str(event.row_key) ) # If this was triggered by keyboard navigation, skip selection and reset flag if self._keyboard_navigation: self._keyboard_navigation = False logger.debug("KEYBOARD NAV: Ignoring row highlight from arrow keys") return logger.info( f"CLICK: table={table.id}, row={row_key}, range_mode={self._range_mode}" ) # Handle source policy selection if table.id == "source_policy_table": if ( self._range_mode and self.last_clicked_row and self.last_clicked_table == table.id ): # Range toggle logger.info(f"RANGE TOGGLE: from {self.last_clicked_row} to {row_key}") self._range_select(table, self.last_clicked_row, row_key) self._range_mode = False # Exit range mode after operation else: # Normal toggle if row_key in self.selected_source_policy_ids: self.selected_source_policy_ids.remove(row_key) else: self.selected_source_policy_ids.add(row_key) self._refresh_table_checkboxes( table.id, self.selected_source_policy_ids ) self.app.notify( f"Selected {len(self.selected_source_policy_ids)} policies", timeout=1, ) # Remember for next range-select self.last_clicked_row = row_key self.last_clicked_table = table.id # Handle path/publisher review selections elif table.id in [ "paths_review_table", "publishers_review_table", "remaining_review_table", ]: if ( self._range_mode and self.last_clicked_row and self.last_clicked_table == table.id ): # Range toggle logger.info(f"RANGE TOGGLE: from {self.last_clicked_row} to {row_key}") self._range_select(table, self.last_clicked_row, row_key) self._range_mode = False # Exit range mode after operation else: # Normal toggle if row_key in self.selected_path_rows: self.selected_path_rows.remove(row_key) else: self.selected_path_rows.add(row_key) self._refresh_table_checkboxes(table.id, self.selected_path_rows) self.app.notify( f"Selected {len(self.selected_path_rows)} items", timeout=1 ) # Remember for next range-select self.last_clicked_row = row_key self.last_clicked_table = table.id # Handle approved/needs review selections elif table.id in ["approved_review_table", "needs_review_table"]: if ( self._range_mode and self.last_clicked_row and self.last_clicked_table == table.id ): # Range toggle logger.info(f"RANGE TOGGLE: from {self.last_clicked_row} to {row_key}") self._range_select(table, self.last_clicked_row, row_key) self._range_mode = False # Exit range mode after operation else: # Normal toggle if row_key in self.selected_rows: self.selected_rows.remove(row_key) else: self.selected_rows.add(row_key) self._refresh_table_checkboxes(table.id, self.selected_rows) self.app.notify(f"Selected {len(self.selected_rows)} rows", timeout=1) # Remember for next range-select self.last_clicked_row = row_key self.last_clicked_table = table.id def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None: """Handle row selection in data tables.""" table = event.data_table logger.debug(f"Row selected in table: {table.id}") # NOTE: Review table selections are handled in on_data_table_row_highlighted (clicks only) # This event (row_selected) is triggered by arrow key navigation, which should NOT select # Only handle special cases like allowlist selection # Ignore all review tables - they use row_highlighted for selection if table.id in [ "source_policy_table", "approved_review_table", "needs_review_table", "paths_review_table", "publishers_review_table", "remaining_review_table", ]: return # Handle allowlist selection (this one uses row selection, not highlighting) if table.id == "allowlist_table": # Get selected allowlist row_index = table.cursor_row if hasattr(self, "allowlists") and row_index < len(self.allowlists): self.destination_allowlist = self.allowlists[row_index] logger.info(f"Selected allowlist: {self.destination_allowlist.name}") self._show_fetch_data() def on_input_submitted(self, event: Input.Submitted) -> None: """Handle input submission (Enter key press).""" if event.input.id == "history_days_input": # Trigger the fetch when user presses Enter in the days input try: history_days = int(event.input.value) if 1 <= history_days <= 365: self.history_days = history_days self._fetch_execution_data(history_days) else: self.app.notify( "Please enter a value between 1 and 365", severity="warning" ) except (ValueError, TypeError): self.app.notify("Please enter a valid number", severity="warning") def on_button_pressed(self, event: Button.Pressed) -> None: """Handle button presses.""" button_id = event.button.id # Introduction screen buttons if button_id == "start_workflow": self._show_source_policy_selection() elif button_id == "cancel_workflow": self.app.pop_screen() # Source policy selection buttons elif button_id == "select_none_source": self.selected_source_policy_ids.clear() # Refresh checkbox display self._refresh_table_checkboxes( "source_policy_table", self.selected_source_policy_ids ) self.app.notify("Cleared selection", timeout=1) elif button_id == "continue_source_selection": if not self.selected_source_policy_ids: self.app.notify( "Please select at least one source policy", severity="warning" ) else: # Get the actual policy objects self.source_policies = [ p for p in self.policies if str(p.groupid) in self.selected_source_policy_ids ] logger.info( f"Selected source policies: {[p.name for p in self.source_policies]}" ) logger.info( f"self.source_policies set to: {len(self.source_policies)} policies" ) if not self.source_policies: logger.error("source_policies list is empty after selection!") self.app.notify( "Error: Could not load selected policies. Please try again.", severity="error", ) else: self._show_destination_policy_selection() # Tab switching buttons elif button_id == "show_approved_tab": self._show_review_table("approved") elif button_id == "show_needs_review_tab": self._show_review_table("needs_review") elif button_id == "show_paths_tab": self._show_path_review_table("paths") elif button_id == "show_publishers_tab": self._show_path_review_table("publishers") elif button_id == "show_remaining_tab": self._show_path_review_table("remaining") # Row selection buttons elif button_id == "select_all_rows": self._select_all_rows() elif button_id == "select_none_rows": self._select_none_rows() elif button_id == "delete_selected_rows": self._delete_selected_rows() elif button_id == "select_all_path_rows": self._select_all_path_rows() elif button_id == "select_none_path_rows": self._select_none_path_rows() elif button_id == "delete_selected_path_rows": self._delete_selected_path_rows() # Export buttons elif button_id == "export_review": self._export_review_data() elif button_id == "export_path_review": self._export_path_review_data() # Original button handlers elif button_id == "fetch_data_btn": # Get history days from input try: days_input = self.query_one("#history_days_input", Input) history_days = int(days_input.value) if 1 <= history_days <= 365: self.history_days = history_days self._fetch_execution_data(history_days) else: self.app.notify( "Please enter a value between 1 and 365", severity="warning" ) except (ValueError, TypeError): self.app.notify("Please enter a valid number", severity="warning") elif button_id == "skip_fetch_btn": # Check if data already exists if self.source_policies: policy_name = self.source_policies[0].name approved_path = os.path.join( self.working_dir, "Needs_Review", "Review_First", f"{policy_name}_approved_executions.csv", ) if os.path.exists(approved_path): # Load existing data self.approved_df = pd.read_csv(approved_path) review_path = approved_path.replace("approved", "needs_review") if os.path.exists(review_path): self.needs_review_df = pd.read_csv(review_path) self._show_fetch_results() else: self.app.notify( "No existing data found. Please fetch new data.", severity="warning", ) elif button_id == "continue_from_review": # Check if both tabs have been reviewed if not self.approved_tab_reviewed or not self.needs_review_tab_reviewed: self.app.notify( "Please review both 'Approved' and 'Needs Review' tabs before continuing.", severity="warning", timeout=5, ) return # Validate that review is complete if (self.approved_df is None or self.approved_df.empty) and ( self.needs_review_df is None or self.needs_review_df.empty ): self.app.notify( "No data to continue with! Please review and keep some executions.", severity="error", ) else: # Save the reviewed data before continuing self._save_reviewed_data() # Show loading screen then build paths self._show_path_building_screen() elif button_id == "build_preflight": # Check if both required tabs have been reviewed if not self.paths_tab_reviewed or not self.publishers_tab_reviewed: self.app.notify( "Please review both 'Paths' and 'Publishers' tabs before continuing.", severity="warning", timeout=5, ) return # Validate that path review is complete if (self.primary_paths_df is None or self.primary_paths_df.empty) and ( self.publishers_df is None or self.publishers_df.empty ): self.app.notify( "No paths or publishers to build preflight with!", severity="error" ) else: self._build_preflight() elif button_id == "back_to_path_review": # Go back to path review screen self._show_path_review_table("paths") elif button_id == "liftoff": # Confirm before applying self.app.notify("Applying changes...", severity="information") self._apply_changes() elif button_id == "workflow_done": self.app.pop_screen() def _select_all_rows(self) -> None: """Select all rows in the current review table.""" table_id = None df = None if self.current_review_type == "approved": table_id = "approved_review_table" df = self.approved_df else: table_id = "needs_review_table" df = self.needs_review_df if table_id and df is not None: # Use actual DataFrame indices, not range(len(df)) self.selected_rows = set(str(i) for i in df.index) # Refresh checkbox display self._refresh_table_checkboxes(table_id, self.selected_rows) self.app.notify(f"Selected all {len(self.selected_rows)} rows", timeout=1) def _select_none_rows(self) -> None: """Clear all row selections in the current review table.""" self.selected_rows.clear() # Refresh checkbox display table_id = ( "approved_review_table" if self.current_review_type == "approved" else "needs_review_table" ) self._refresh_table_checkboxes(table_id, self.selected_rows) self.app.notify("Cleared selection", timeout=1) def _select_all_path_rows(self) -> None: """Select all rows in the current path review table.""" table_id = None if self.current_path_review_type == "paths": table_id = "paths_review_table" # For paths, the combined DataFrame uses ignore_index=True, so indices are 0..n-1 total = 0 if self.primary_paths_df is not None: total += len(self.primary_paths_df) if self.secondary_paths_df is not None: total += len(self.secondary_paths_df) self.selected_path_rows = set(str(i) for i in range(total)) elif ( self.current_path_review_type == "publishers" and self.publishers_df is not None ): table_id = "publishers_review_table" # For publishers, use actual DataFrame indices self.selected_path_rows = set(str(i) for i in self.publishers_df.index) # Refresh checkbox display if table_id: self._refresh_table_checkboxes(table_id, self.selected_path_rows) self.app.notify(f"Selected all {len(self.selected_path_rows)} items", timeout=1) def _select_none_path_rows(self) -> None: """Clear all row selections in the current path review table.""" self.selected_path_rows.clear() # Refresh checkbox display table_id = ( "paths_review_table" if self.current_path_review_type == "paths" else "publishers_review_table" ) self._refresh_table_checkboxes(table_id, self.selected_path_rows) self.app.notify("Cleared selection", timeout=1) def _save_reviewed_data(self) -> None: """Save the reviewed dataframes to the Approved folder.""" if not self.source_policies: return policy_name = self.source_policies[0].name approved_dir = os.path.join(self.working_dir, "Approved") os.makedirs(approved_dir, exist_ok=True) # Save approved executions if self.approved_df is not None and not self.approved_df.empty: filepath = os.path.join( approved_dir, f"{policy_name}_approved_executions.csv" ) self.approved_df.to_csv(filepath, index=False) logger.info(f"Saved approved executions to {filepath}") # Save needs_review as approved (since user reviewed them) if self.needs_review_df is not None and not self.needs_review_df.empty: filepath = os.path.join( approved_dir, f"{policy_name}_needs_review_executions.csv" ) self.needs_review_df.to_csv(filepath, index=False) logger.info(f"Saved reviewed executions to {filepath}") def _export_review_data(self) -> None: """Export current review data to CSV.""" if not self.source_policies: return policy_name = self.source_policies[0].name timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") if self.current_review_type == "approved" and self.approved_df is not None: filepath = os.path.join( self.working_dir, f"{policy_name}_approved_export_{timestamp}.csv" ) self.approved_df.to_csv(filepath, index=False) self.app.notify(f"Exported to: {filepath}", severity="information") elif ( self.current_review_type == "needs_review" and self.needs_review_df is not None ): filepath = os.path.join( self.working_dir, f"{policy_name}_needs_review_export_{timestamp}.csv" ) self.needs_review_df.to_csv(filepath, index=False) self.app.notify(f"Exported to: {filepath}", severity="information") def on_key(self, event) -> None: """Handle keyboard shortcuts including range selection mode.""" key = event.key # Track arrow key navigation to prevent selection if key in ["up", "down", "left", "right", "pageup", "pagedown", "home", "end"]: self._keyboard_navigation = True return # Let the event propagate for navigation # 'r' activates range selection mode if key == "r": if self.last_clicked_row and self.last_clicked_table: self._range_mode = True self.app.notify( "Range mode: Click end row (or press ESC to cancel)", severity="information", timeout=5, ) logger.info( f"RANGE MODE ACTIVATED: starting from row {self.last_clicked_row} in table {self.last_clicked_table}" ) else: self.app.notify( "Click a row first, then press 'r' to start range selection", severity="warning", timeout=3, ) # ESC cancels range mode elif key == "escape": if self._range_mode: self._range_mode = False self.app.notify( "Range mode cancelled", severity="information", timeout=2 ) logger.info("RANGE MODE CANCELLED") def on_key_up(self, event) -> None: """Handle key releases (currently unused but kept for future).""" pass def _export_path_review_data(self) -> None: """Export current path review data to CSV.""" if not self.source_policies: return policy_name = self.source_policies[0].name timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") if self.current_path_review_type == "paths": # Export both primary and secondary paths if self.primary_paths_df is not None: filepath = os.path.join( self.working_dir, f"{policy_name}_primary_paths_export_{timestamp}.csv", ) self.primary_paths_df.to_csv(filepath, index=False) self.app.notify( f"Exported primary paths to: {filepath}", severity="information" ) if self.secondary_paths_df is not None: filepath = os.path.join( self.working_dir, f"{policy_name}_secondary_paths_export_{timestamp}.csv", ) self.secondary_paths_df.to_csv(filepath, index=False) self.app.notify( f"Exported secondary paths to: {filepath}", severity="information" ) elif ( self.current_path_review_type == "publishers" and self.publishers_df is not None ): filepath = os.path.join( self.working_dir, f"{policy_name}_publishers_export_{timestamp}.csv" ) self.publishers_df.to_csv(filepath, index=False) self.app.notify(f"Exported to: {filepath}", severity="information") def _open_folder(self, path: str) -> None: """Open a folder in the system file explorer.""" try: import platform import subprocess os.makedirs(path, exist_ok=True) if platform.system() == "Windows": subprocess.Popen(f'explorer "{path}"') elif platform.system() == "Darwin": # macOS subprocess.Popen(["open", path]) else: # Linux subprocess.Popen(["xdg-open", path]) self.app.notify(f"Opened: {path}", severity="information") except Exception as e: logger.error(f"Failed to open folder: {e}") self.app.notify(f"Failed to open folder: {str(e)}", severity="error") # Action handlers def action_go_back(self) -> None: """Handle back/escape action.""" stage_transitions = { "select_source": lambda: self.app.pop_screen(), "select_destination": self._show_source_policy_selection, "select_allowlist": self._show_destination_policy_selection, "fetch_data": self._show_allowlist_selection, "first_review": self._show_fetch_data, "second_review": self._show_fetch_results, "test": self._show_path_results, "complete": lambda: self.app.pop_screen(), } transition = stage_transitions.get(self.workflow_stage) if transition: transition() 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() def action_open_folder(self) -> None: """Open the working directory.""" self._open_folder(self.working_dir) def action_delete_rows(self) -> None: """Delete selected rows in the current table.""" if self.workflow_stage == "first_review": self._delete_selected_rows() elif self.workflow_stage == "second_review": self._delete_selected_path_rows() def action_select_all(self) -> None: """Select all rows in the current table.""" if self.workflow_stage == "first_review": self._select_all_rows() elif self.workflow_stage == "second_review": self._select_all_path_rows() def action_select_none(self) -> None: """Clear selection in the current table.""" if self.workflow_stage == "first_review": self._select_none_rows() elif self.workflow_stage == "second_review": self._select_none_path_rows() def action_toggle_selection(self) -> None: """Toggle selection on the current row at cursor position.""" # Get the focused widget (should be a DataTable) focused = self.app.focused if not isinstance(focused, DataTable): return table = focused # Get the current cursor row try: cursor_row = table.cursor_row # Get the row key at the cursor position row_keys = list(table.rows.keys()) if cursor_row < len(row_keys): row_key = str( row_keys[cursor_row].value if hasattr(row_keys[cursor_row], "value") else row_keys[cursor_row] ) logger.info(f"SPACE: Toggling row {row_key} in table {table.id}") # Toggle based on table type if table.id == "source_policy_table": if row_key in self.selected_source_policy_ids: self.selected_source_policy_ids.remove(row_key) else: self.selected_source_policy_ids.add(row_key) self._refresh_table_checkboxes( table.id, self.selected_source_policy_ids ) self.app.notify( f"Selected {len(self.selected_source_policy_ids)} policies", timeout=1, ) # Remember for range mode self.last_clicked_row = row_key self.last_clicked_table = table.id elif table.id in ["approved_review_table", "needs_review_table"]: if row_key in self.selected_rows: self.selected_rows.remove(row_key) else: self.selected_rows.add(row_key) self._refresh_table_checkboxes(table.id, self.selected_rows) self.app.notify( f"Selected {len(self.selected_rows)} rows", timeout=1 ) # Remember for range mode self.last_clicked_row = row_key self.last_clicked_table = table.id elif table.id in [ "paths_review_table", "publishers_review_table", "remaining_review_table", ]: if row_key in self.selected_path_rows: self.selected_path_rows.remove(row_key) else: self.selected_path_rows.add(row_key) self._refresh_table_checkboxes(table.id, self.selected_path_rows) self.app.notify( f"Selected {len(self.selected_path_rows)} items", timeout=1 ) # Remember for range mode self.last_clicked_row = row_key self.last_clicked_table = table.id except Exception as e: logger.error(f"Error toggling selection: {e}")