diff --git a/Loxide.py b/Loxide.py index 47c9562..7fa448c 100644 --- a/Loxide.py +++ b/Loxide.py @@ -30,8 +30,8 @@ import urllib3 from services.API import AirlockAPIWrapper from services.security import getAPI +from TUI.TUI import run_Loxide from utils.setup import get_base_directory, setup -from utils.TUI import run_Loxide from utils.utils import irtang urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) diff --git a/Loxide_Icon.ico b/Loxide_Icon.ico new file mode 100644 index 0000000..dce491f Binary files /dev/null and b/Loxide_Icon.ico differ diff --git a/widgets/OTP_generate.py b/TUI/OTP_generate.py similarity index 100% rename from widgets/OTP_generate.py rename to TUI/OTP_generate.py diff --git a/utils/tui.py b/TUI/TUI.py similarity index 92% rename from utils/tui.py rename to TUI/TUI.py index 9f32133..990cde3 100644 --- a/utils/tui.py +++ b/TUI/TUI.py @@ -22,25 +22,25 @@ from textual.widgets import ( from flows.otp import otp_revoke from flows.prepPolicy import menu_policy_enforce -from flows.quietAgent import findQuietAgents from models.agent import Agent from models.policy import Policy -from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen -from screens.otpactivityscreen import OTPActivitiesScreen -from screens.otpworkflowscreen import OTPWorkflowScreen from services.API import AirlockAPIWrapper from services.policyhandler import confirmUpdateAfromE -from themes.amber_terminal_theme import get_amber_terminal_theme -from themes.retro_terminal_theme import get_retro_terminal_theme +from TUI.agentmoveoperations import AgentMoveOperations +from TUI.moveagentworkflowscreen import MoveAgentWorkflowScreen +from TUI.multiagentselector import MultiAgentSelector +from TUI.OTP_generate import OTPGenerator +from TUI.otpactivityscreen import OTPActivitiesScreen +from TUI.otpworkflowscreen import OTPWorkflowScreen +from TUI.policytreewidget import PolicyTreeWidget +from TUI.quietagentworkflowscreen import QuietAgentWorkflowScreen +from TUI.resultsdisplay import ResultsDisplay +from TUI.theme_amber_terminal import get_amber_terminal_theme +from TUI.theme_retro_terminal import get_retro_terminal_theme +from TUI.themeselector import ThemeSelector from utils.configmanager import load_env from utils.setup import get_base_directory, load_user_config from utils.utils import open_directory -from widgets.agentmoveoperations import AgentMoveOperations -from widgets.multiagentselector import MultiAgentSelector -from widgets.OTP_generate import OTPGenerator -from widgets.policytreewidget import PolicyTreeWidget -from widgets.resultsdisplay import ResultsDisplay -from widgets.themeselector import ThemeSelector dotenv.load_dotenv() @@ -114,8 +114,8 @@ class MainMenuScreen(Screen): "šŸ–„ļø - Find, Move, or Generate OTP for Agents", "move_agent_workflow_button", ), - ("šŸ“Š - OTP Activities By Agent", "otp_activities_button"), - ("šŸ”‡ - Find Quiet Hosts", "find_quiet_button"), + ("šŸ“Š - Review and appove OTP Activities", "otp_activities_button"), + ("šŸ”‡ - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"), ], "policy": [ ("šŸ”’ - Prepare Policy For Enforcement", "policy_prep_button"), @@ -242,7 +242,6 @@ class MainMenuScreen(Screen): """Handle OTP generation request from the workflow.""" global _PENDING_JOB - # Log what we received logger.info( "OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d", len(message.devices), @@ -251,7 +250,6 @@ class MainMenuScreen(Screen): message.duration, ) - # Set up the job to run the OTP generation _PENDING_JOB = ( "otp_workflow", message.devices, @@ -321,23 +319,21 @@ class MainMenuScreen(Screen): match button_id: case "move_agent_workflow_button": - # Push Move Agent workflow screen self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices)) event.stop() - return # Don't exit the app case "otp_generate_button": - # NEW: Push OTP workflow screen instead of legacy function self.app.push_screen(OTPWorkflowScreen(self.app.devices)) event.stop() - return # Don't exit the app case "find_quiet_button": - _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {}) + self.app.push_screen( + QuietAgentWorkflowScreen(self.app.api, self.app.policies) + ) + event.stop() + return case "otp_activities_button": - # === FIXED: push the Textual OTPActivitiesScreen and return immediately === - # This must return so we don't fall through to the code that exits the app. self.app.push_screen(OTPActivitiesScreen()) event.stop() return diff --git a/widgets/agentmoveoperations.py b/TUI/agentmoveoperations.py similarity index 96% rename from widgets/agentmoveoperations.py rename to TUI/agentmoveoperations.py index 03e5da4..d3615f0 100644 --- a/widgets/agentmoveoperations.py +++ b/TUI/agentmoveoperations.py @@ -1,23 +1,3 @@ -""" -Agent Move Operations Widget Module - -This module provides a Textual-based UI widget for performing bulk operations on -agent devices in the Airlock system. It allows users to: -- View selected agents and their current policy assignments -- Move agents to local approval mode with OTP enforcement -- Toggle agents between audit and enforcement policy modes -- Select and move agents to alternate policies (future implementation) - -The widget tracks operation state, manages button availability, and displays -results with success/failure summaries that can be copied to clipboard. - -Dependencies: - - textual: TUI framework for building the widget and UI components - - models.agent: Agent model class - - services.agenthandler: Core agent operation functions - - flows.localApproval: Local approval workflow handling -""" - from dataclasses import asdict from datetime import datetime import logging @@ -33,9 +13,9 @@ from textual.widget import Widget from textual.widgets import Button, DataTable, Header, Static, TextArea from models.agent import Agent -from screens.otpworkflowscreen import OTPWorkflowScreen -from screens.policyselectorscreen import PolicySelectorScreen -from widgets.OTP_generate import OTPGenerator +from TUI.OTP_generate import OTPGenerator +from TUI.otpworkflowscreen import OTPWorkflowScreen +from TUI.policyselectorscreen import PolicySelectorScreen logger = logging.getLogger(__name__) diff --git a/TUI/allowlistselectionscreen.py b/TUI/allowlistselectionscreen.py new file mode 100644 index 0000000..fd84f20 --- /dev/null +++ b/TUI/allowlistselectionscreen.py @@ -0,0 +1,616 @@ +from __future__ import annotations + +import logging +from typing import Optional + +import pandas as pd +from textual.app import ComposeResult +from textual.binding import Binding +from textual.containers import Horizontal, Vertical +from textual.screen import Screen +from textual.widgets import ( + Button, + DataTable, + Footer, + Header, + Static, + TextArea, +) + +logger = logging.getLogger(__name__) + + +class AllowlistSelectionWidget(Static): + """ + Widget for selecting an allowlist and adding hashes to it. + Can be reused in different workflows. + """ + + DEFAULT_CSS = """ + AllowlistSelectionWidget { + height: 1fr; + layout: vertical; + } + #allowlist_main { + height: 100%; + width: 100%; + } + #left_panel { + width: 50%; + padding: 1; + border: solid $primary; + } + #right_panel { + width: 50%; + padding: 1; + border: solid $primary; + } + #allowlist_table { + height: 70%; + margin: 1 0; + } + #allowlist_table > .datatable--header { + text-style: bold; + background: $boost; + } + #allowlist_table Row { + height: 1; + } + #preview_area { + height: 60%; + margin: 1 0; + } + #action_buttons { + height: 10%; + padding: 1; + content-align: center middle; + } + .panel-title { + text-style: bold; + margin: 0 0 1 0; + } + .info-text { + margin: 1 0; + } + """ + + def __init__( + self, + selected_data: pd.DataFrame, + api=None, + hostname: Optional[str] = None, + otpid: Optional[str] = None, + hash_column: str = "sha256", # Default hash column name + ): + """ + Initialize the allowlist selection widget. + + Args: + selected_data: DataFrame containing the selected activities + api: API instance for making allowlist calls + hostname: Optional hostname for context + otpid: Optional OTP ID for context + hash_column: Name of the column containing hashes (default: "sha256") + """ + super().__init__() + self.selected_data = selected_data + self.api = api + self.hostname = hostname + self.otpid = otpid + self.hash_column = hash_column + self.allowlists = [] + self.selected_allowlist = None + self.hashes_to_add = [] + + def compose(self) -> ComposeResult: + with Horizontal(id="allowlist_main"): + # Left panel - Allowlist selection + with Vertical(id="left_panel"): + yield Static("Select Allowlist", classes="panel-title") + yield Static( + f"Choose an allowlist to add {len(self.selected_data)} selected items", + classes="info-text", + ) + + # Allowlist table + self.allowlist_table = DataTable(id="allowlist_table") + self.allowlist_table.cursor_type = "row" + yield self.allowlist_table + + # Refresh button + self.refresh_btn = Button( + "šŸ”„ Refresh Allowlists", id="refresh_allowlists_btn" + ) + yield self.refresh_btn + + # Right panel - Preview and actions + with Vertical(id="right_panel"): + yield Static("Preview", classes="panel-title") + + # Context information + context_text = [] + if self.hostname: + context_text.append(f"Host: {self.hostname}") + if self.otpid: + context_text.append(f"OTP: {self.otpid}") + context_text.append(f"Selected Activities: {len(self.selected_data)}") + + yield Static(" | ".join(context_text), classes="info-text") + + # Preview text area + self.preview_area = TextArea( + id="preview_area", read_only=True, language="markdown" + ) + yield self.preview_area + + # Hash statistics + self.stats_label = Static("", id="stats_label", classes="info-text") + yield self.stats_label + + # Action buttons at bottom + with Horizontal(id="action_buttons"): + self.back_btn = Button("← Back", id="back_btn") + self.add_btn = Button("āž• Add to Allowlist", id="add_to_allowlist_btn") + + self.back_btn.styles.width = "50%" + self.add_btn.styles.width = "50%" + self.add_btn.disabled = True # Disabled until allowlist selected + + yield self.back_btn + yield self.add_btn + + async def on_mount(self) -> None: + """Load allowlists when widget mounts.""" + await self.load_allowlists() + await self.extract_and_preview_hashes() + + async def load_allowlists(self) -> None: + """Load available allowlists from API, grouped by policy association.""" + if not self.api: + logger.error("No API available") + self.allowlist_table.add_column("Error") + self.allowlist_table.add_row("No API available") + return + + try: + # First, try to get the host's policy if hostname is provided + host_policy_allowlists = [] + host_policy_ids = set() + policy_name = None + + if self.hostname: + try: + # Get agent info to find its policy + agents_df = self.api.agent_find_by_hostname(self.hostname) + if not agents_df.empty: + # Get the policy group ID for this host + group_id = agents_df.iloc[0].get("groupid") + policy_name = agents_df.iloc[0].get( + "groupname", "Unknown Policy" + ) + + if group_id: + # Get allowlists for this policy + policy_allowlists_df = self.api.policy_list_allowlists( + group_id + ) + if not policy_allowlists_df.empty: + host_policy_allowlists = policy_allowlists_df.to_dict( + orient="records" + ) + host_policy_ids = { + al.get("applicationid") + for al in host_policy_allowlists + } + logger.info( + f"Found {len(host_policy_allowlists)} allowlists for host's policy" + ) + except Exception as e: + logger.warning(f"Could not get host's policy allowlists: {e}") + + # Get all allowlists + all_allowlists_df = self.api.allowlist_find_all() + + if all_allowlists_df.empty: + self.allowlist_table.add_column("No Allowlists") + self.allowlist_table.add_row("No allowlists found") + return + + all_allowlists = all_allowlists_df.to_dict(orient="records") + + # Separate into two groups: policy-associated and others + other_allowlists = [ + al + for al in all_allowlists + if al.get("applicationid") not in host_policy_ids + ] + + # Sort each group alphabetically by name + host_policy_allowlists.sort(key=lambda x: x.get("name", "").lower()) + other_allowlists.sort(key=lambda x: x.get("name", "").lower()) + + # Combine lists with policy-associated first + self.allowlists = host_policy_allowlists + other_allowlists + + # Setup table columns + self.allowlist_table.clear() + self.allowlist_table.add_columns("Name", "Application ID", "Type") + + # Track which rows are headers vs actual allowlists + self._row_to_allowlist_map = {} + current_row = 0 + + # Add policy-associated allowlists if any + if host_policy_allowlists: + # Add section header + header_text = f"━━━ Policy: {policy_name or 'Host Policy'} ━━━" + self.allowlist_table.add_row(header_text, "", "", key="header_policy") + current_row += 1 + + # Add policy allowlists + for idx, allowlist in enumerate(host_policy_allowlists): + name = allowlist.get("name", "Unknown") + app_id = allowlist.get("applicationid", "Unknown") + + self.allowlist_table.add_row( + f" {name}", # Indent to show grouping + app_id, + "Policy", + key=f"policy_{idx}", + ) + self._row_to_allowlist_map[current_row] = idx + current_row += 1 + + # Add other allowlists + if other_allowlists: + # Add section header + if host_policy_allowlists: + # Add spacer if we have policy allowlists above + self.allowlist_table.add_row("", "", "", key="spacer") + current_row += 1 + + self.allowlist_table.add_row( + "━━━ Other Available Allowlists ━━━", "", "", key="header_other" + ) + current_row += 1 + + # Add other allowlists + for idx, allowlist in enumerate(other_allowlists): + name = allowlist.get("name", "Unknown") + app_id = allowlist.get("applicationid", "Unknown") + + self.allowlist_table.add_row( + f" {name}", # Indent to show grouping + app_id, + "General", + key=f"other_{idx}", + ) + # Map to the correct index in the combined list + actual_idx = len(host_policy_allowlists) + idx + self._row_to_allowlist_map[current_row] = actual_idx + current_row += 1 + + # Log summary + logger.info( + f"Loaded {len(self.allowlists)} total allowlists: " + f"{len(host_policy_allowlists)} policy-associated, " + f"{len(other_allowlists)} others" + ) + + # Update stats label if no allowlists in policy + if self.hostname and not host_policy_allowlists: + self.stats_label.update( + f"Note: No allowlists found for {self.hostname}'s policy | " + + self.stats_label.content.plain + ) + + except Exception as exc: + logger.exception(f"Failed to load allowlists: {exc}") + self.allowlist_table.add_column("Error") + self.allowlist_table.add_row(f"Failed to load: {str(exc)}") + + async def extract_and_preview_hashes(self) -> None: + """Extract hashes from selected data and show preview.""" + preview_lines = ["## Hash Extraction Summary\n"] + + # Check for hash column + if self.hash_column not in self.selected_data.columns: + # Try to find a hash column + possible_hash_cols = [ + "sha256", + "SHA256", + "hash", + "Hash", + "sha1", + "SHA1", + "md5", + "MD5", + "filehash", + "file_hash", + ] + found_col = None + for col in possible_hash_cols: + if col in self.selected_data.columns: + found_col = col + break + + if found_col: + self.hash_column = found_col + preview_lines.append(f"āœ“ Found hash column: **{found_col}**\n") + else: + preview_lines.append("āš ļø **No hash column found**\n") + preview_lines.append("Available columns:\n") + for col in self.selected_data.columns: + if col != "_row_id": + preview_lines.append(f" - {col}\n") + + self.preview_area.text = "".join(preview_lines) + self.stats_label.update("No hashes to add") + return + + # Extract unique hashes + hashes = self.selected_data[self.hash_column].dropna().unique() + self.hashes_to_add = [h for h in hashes if h and str(h).strip()] + + # Build preview + preview_lines.append(f"### Found {len(self.hashes_to_add)} unique hashes\n\n") + + # Show sample of hashes (first 10) + preview_lines.append("**Sample hashes to be added:**\n```\n") + for i, hash_val in enumerate(self.hashes_to_add[:10]): + preview_lines.append(f"{i+1}. {hash_val}\n") + if len(self.hashes_to_add) > 10: + preview_lines.append(f"... and {len(self.hashes_to_add) - 10} more\n") + preview_lines.append("```\n\n") + + # Show sample of source data + preview_lines.append("**Sample source activities:**\n") + sample_cols = [ + col + for col in self.selected_data.columns + if col not in ["_row_id"] and col in ["filename", "path", "action", "user"] + ] + if not sample_cols: + sample_cols = [ + col for col in self.selected_data.columns if col != "_row_id" + ][:3] + + if sample_cols: + preview_lines.append("```\n") + for i, row in self.selected_data[sample_cols].head(5).iterrows(): + row_text = " | ".join([f"{col}: {row[col]}" for col in sample_cols]) + preview_lines.append(f"{row_text}\n") + preview_lines.append("```\n") + + self.preview_area.text = "".join(preview_lines) + + # Update statistics + self.stats_label.update( + f"Ready to add {len(self.hashes_to_add)} unique hashes | " + f"From {len(self.selected_data)} selected activities" + ) + + async def on_data_table_row_selected(self, event) -> None: + """Handle allowlist selection.""" + try: + # Extract row index from event - handle different event structures + row_index = None + + # Try to get row index from coordinate + if hasattr(event, "coordinate") and hasattr(event.coordinate, "row"): + row_index = event.coordinate.row + # Try cursor_row as fallback + elif hasattr(event, "cursor_row"): + row_index = event.cursor_row + # Try getting from the table itself + else: + table = self.allowlist_table + if hasattr(table, "cursor_row"): + row_index = table.cursor_row + + # Validate row index + if row_index is not None and isinstance(row_index, int): + # Account for group headers in the row count + actual_allowlist_index = self._get_allowlist_index_from_row(row_index) + + if ( + actual_allowlist_index is not None + and 0 <= actual_allowlist_index < len(self.allowlists) + ): + self.selected_allowlist = self.allowlists[actual_allowlist_index] + self.add_btn.disabled = False + self.add_btn.label = ( + f"āž• Add to '{self.selected_allowlist.get('name', 'Unknown')}'" + ) + + # Update preview with selection + await self._update_preview_with_selection() + + logger.info( + f"Selected allowlist: {self.selected_allowlist.get('name')}" + ) + else: + logger.debug(f"Row {row_index} is a header or invalid") + else: + logger.warning(f"Could not extract valid row index from event: {event}") + + except Exception as exc: + logger.exception(f"Failed to select allowlist: {exc}") + + def _get_allowlist_index_from_row(self, row_index: int) -> Optional[int]: + """Convert table row index to allowlist list index, accounting for group headers.""" + # This will be updated when we have group headers + if hasattr(self, "_row_to_allowlist_map"): + return self._row_to_allowlist_map.get(row_index) + return row_index + + async def _update_preview_with_selection(self) -> None: + """Update preview when an allowlist is selected.""" + if not self.selected_allowlist: + return + + current_text = self.preview_area.text + # Remove any existing selection header + if "### Selected Allowlist:" in current_text: + lines = current_text.split("\n") + # Find and remove the selection lines + new_lines = [] + skip_next = False + for line in lines: + if line.startswith("### Selected Allowlist:"): + skip_next = True + continue + if skip_next and line.startswith("Application ID:"): + skip_next = False + continue + if not skip_next: + new_lines.append(line) + current_text = "\n".join(new_lines) + + # Add new selection at the top + selection_text = ( + f"### Selected Allowlist: **{self.selected_allowlist.get('name')}**\n" + f"Application ID: {self.selected_allowlist.get('applicationid')}\n\n" + ) + self.preview_area.text = selection_text + current_text + + async def on_button_pressed(self, event) -> None: + """Handle button presses.""" + btn = getattr(event, "button", None) or getattr(event, "sender", None) + btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None) + + if btn is self.back_btn or btn_id == "back_btn": + await self.app.pop_screen() + event.stop() + return + + if btn is self.refresh_btn or btn_id == "refresh_allowlists_btn": + await self.load_allowlists() + event.stop() + return + + if btn is self.add_btn or btn_id == "add_to_allowlist_btn": + await self.add_hashes_to_allowlist() + event.stop() + return + + async def add_hashes_to_allowlist(self) -> None: + """Add the extracted hashes to the selected allowlist.""" + if not self.selected_allowlist or not self.hashes_to_add: + self.app.notify( + "No allowlist selected or no hashes to add", severity="warning" + ) + return + + if not self.api: + self.app.notify("API not available", severity="error") + return + + try: + # Disable button during operation + self.add_btn.disabled = True + self.add_btn.label = "ā³ Adding hashes..." + + # Call API to add hashes + app_id = self.selected_allowlist.get("applicationid") + allowlist_name = self.selected_allowlist.get("name", "Unknown") + + logger.info( + f"Adding {len(self.hashes_to_add)} hashes to allowlist {allowlist_name} (ID: {app_id})" + ) + + result = self.api.hash_add_to_allowlist(app_id, self.hashes_to_add) + + # Success notification + self.app.notify( + f"āœ… Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'", + title="Success", + severity="information", + timeout=5, + ) + + # Update preview to show success + self.preview_area.text = ( + f"## āœ… SUCCESS\n\n" + f"Added **{len(self.hashes_to_add)} hashes** to allowlist:\n" + f"**{allowlist_name}** (ID: {app_id})\n\n" + f"### Operation Details:\n" + f"- Source: {self.hostname or 'Multiple hosts'}\n" + f"- OTP ID: {self.otpid or 'N/A'}\n" + f"- Activities processed: {len(self.selected_data)}\n" + f"- Unique hashes added: {len(self.hashes_to_add)}\n" + ) + + # Change button to "Done" + self.add_btn.label = "āœ… Done - Close" + self.add_btn.disabled = False + + # When clicked again, close the screen + self.add_btn_success = True + + except Exception as exc: + logger.exception(f"Failed to add hashes to allowlist: {exc}") + self.app.notify( + f"āŒ Failed to add hashes: {str(exc)}", + title="Error", + severity="error", + timeout=10, + ) + + # Re-enable button + self.add_btn.disabled = False + self.add_btn.label = "āž• Retry Add to Allowlist" + + +class AllowlistSelectionScreen(Screen): + """ + Screen wrapper for the AllowlistSelectionWidget. + """ + + BINDINGS = [ + Binding("b", "back", "Back"), + Binding("r", "refresh", "Refresh Allowlists"), + Binding("enter", "confirm", "Add to Allowlist"), + ] + + def __init__( + self, + selected_data: pd.DataFrame, + api=None, + hostname: Optional[str] = None, + otpid: Optional[str] = None, + hash_column: str = "sha256", + ): + super().__init__() + self.selected_data = selected_data + self.api = api + self.hostname = hostname + self.otpid = otpid + self.hash_column = hash_column + + def compose(self) -> ComposeResult: + yield Header(show_clock=True) + self.widget = AllowlistSelectionWidget( + self.selected_data, + api=self.api, + hostname=self.hostname, + otpid=self.otpid, + hash_column=self.hash_column, + ) + yield self.widget + yield Footer() + + async def action_back(self) -> None: + """Go back to previous screen.""" + await self.app.pop_screen() + + async def action_refresh(self) -> None: + """Refresh the allowlists.""" + if hasattr(self, "widget") and self.widget: + await self.widget.load_allowlists() + + async def action_confirm(self) -> None: + """Confirm and add to allowlist.""" + if hasattr(self, "widget") and self.widget: + if self.widget.selected_allowlist and self.widget.hashes_to_add: + await self.widget.add_hashes_to_allowlist() diff --git a/screens/moveagentworkflowscreen.py b/TUI/moveagentworkflowscreen.py similarity index 92% rename from screens/moveagentworkflowscreen.py rename to TUI/moveagentworkflowscreen.py index da96331..8398ddf 100644 --- a/screens/moveagentworkflowscreen.py +++ b/TUI/moveagentworkflowscreen.py @@ -4,9 +4,9 @@ from textual.app import ComposeResult from textual.screen import Screen from models.agent import Agent -from widgets.agentmoveoperations import AgentMoveOperations -from widgets.multiagentselector import MultiAgentSelector -from widgets.resultsdisplay import ResultsDisplay +from TUI.agentmoveoperations import AgentMoveOperations +from TUI.multiagentselector import MultiAgentSelector +from TUI.resultsdisplay import ResultsDisplay class MoveAgentWorkflowScreen(Screen): diff --git a/widgets/multiagentselector.py b/TUI/multiagentselector.py similarity index 100% rename from widgets/multiagentselector.py rename to TUI/multiagentselector.py diff --git a/screens/otpactivityscreen.py b/TUI/otpactivityscreen.py similarity index 69% rename from screens/otpactivityscreen.py rename to TUI/otpactivityscreen.py index cc98655..84847e0 100644 --- a/screens/otpactivityscreen.py +++ b/TUI/otpactivityscreen.py @@ -11,6 +11,7 @@ from textual.containers import Horizontal, Vertical from textual.screen import Screen from textual.widgets import Button, DataTable, Footer, Header, Static +from TUI.allowlistselectionscreen import AllowlistSelectionScreen from utils.configmanager import load_env logger = logging.getLogger(__name__) @@ -134,14 +135,12 @@ class OTPActivitiesWidget(Static): or getattr(event, "button_id", None) or getattr(event, "id", None) ) - # ---- Back ---- if btn is self.back_btn or btn_id == getattr(self.back_btn, "id", None): while len(self.app.screen_stack) > 2: self.app.pop_screen() event.stop() return - # ---- Continue ---- if btn is self.continue_btn or btn_id == getattr(self.continue_btn, "id", None): if self._activities_df is None or self._activities_df.empty: @@ -172,7 +171,6 @@ class OTPActivitiesWidget(Static): except Exception as exc: logger.exception("Failed to push ActivityDetailScreen: %s", exc) return - # Unknown button on widget logger.debug( "Unhandled OTPActivitiesWidget button pressed (resolved btn=%r, id=%r)", @@ -212,7 +210,6 @@ class OTPActivitiesWidget(Static): row_key = getattr(event, attr, None) if row_key is not None: break - # If coordinate: try to extract .row or tuple[0] if row_key is None: coord = getattr(event, "coordinate", None) or getattr( @@ -223,7 +220,6 @@ class OTPActivitiesWidget(Static): row_key = coord.row elif isinstance(coord, (tuple, list)) and len(coord) >= 1: row_key = coord[0] - # If still nothing, maybe the event provides the row's cell values directly row_values = None for attr in ("values", "cells", "row", "row_values", "selected_row_values"): @@ -232,7 +228,6 @@ class OTPActivitiesWidget(Static): # Prefer actual sequence of cell values row_values = val break - # If we have row_values, try to map them back to the sessions DataFrame if row_values is not None: # Normalize into list of strings for comparison @@ -323,7 +318,6 @@ class OTPActivitiesWidget(Static): # Helpful debug hint for you to paste back if still failing: logger.debug("Event repr for debugging: %r", event) return - # At this point we should have an integer idx try: idx = int(idx) @@ -332,7 +326,6 @@ class OTPActivitiesWidget(Static): "Final normalization of selected row index failed: %r", idx ) return - # Validate sessions df if self._sessions_df is None or self._sessions_df.empty: logger.warning("Sessions DataFrame empty; nothing to select.") @@ -479,6 +472,7 @@ class ActivityDetailWidget(Static): """ Interactive widget for Activity Detail screen. Shows the provided DataFrame in a DataTable and offers Export + Back buttons. + Now includes Select All/None and Add to Allowlist functionality. """ DEFAULT_CSS = """ @@ -487,9 +481,14 @@ class ActivityDetailWidget(Static): layout: vertical; } #detail_table_container { - height: 85%; + height: 75%; padding: 1 1; } + #selection_buttons { + height: 10%; + padding: 1 1; + content-align: center middle; + } #detail_buttons { height: 15%; padding: 1 1; @@ -504,8 +503,16 @@ class ActivityDetailWidget(Static): if isinstance(activities_df, pd.DataFrame) else pd.DataFrame(activities_df) ) + # Add a unique identifier column if not present + if "_row_id" not in self.activities_df.columns: + self.activities_df["_row_id"] = range(len(self.activities_df)) + self.otpid = otpid self.hostname = hostname + self.selected_row_ids = set() # Track selected rows by unique ID + self.row_key_to_id = {} # Map DataTable row keys to unique row IDs + self.table_row_to_id = {} # Map table row indices to unique row IDs + self._last_sort = None # Track last sort column and order def compose(self) -> ComposeResult: yield Static( @@ -516,73 +523,267 @@ class ActivityDetailWidget(Static): with Vertical(id="detail_table_container"): self.detail_table = DataTable(id="detail_table") yield self.detail_table - # Buttons at bottom + + # Original buttons at bottom with Horizontal(id="detail_buttons"): self.detail_back_btn = Button("Back", id="detail_back_btn") - self.detail_export_btn = Button("Export (CSV)", id="detail_export_btn") - # Make them stretch equally - self.detail_back_btn.styles.width = "50%" - self.detail_export_btn.styles.width = "50%" + self.add_allowlist_btn = Button( + "šŸ“‹ Add Selected to Allowlist", id="add_allowlist_btn" + ) + yield self.add_allowlist_btn yield self.detail_back_btn - yield self.detail_export_btn async def on_mount(self) -> None: - # Populate table from activities_df - self.detail_table.clear() - if self.activities_df is None or self.activities_df.empty: - logger.info("ActivityDetailWidget mounted with empty dataframe.") - return - # Add columns - for col in self.activities_df.columns: - self.detail_table.add_column(col) + await self._build_table(rebuild=True) + self._update_button_states() + + def _update_button_states(self) -> None: + """Update button states based on selection.""" + has_selection = len(self.selected_row_ids) > 0 + self.add_allowlist_btn.disabled = not has_selection + + # Update button labels with count + count = len(self.selected_row_ids) + total = len(self.activities_df) + + if has_selection: + self.add_allowlist_btn.label = f"šŸ“‹ Add {count} Selected to Allowlist" + else: + self.add_allowlist_btn.label = "šŸ“‹ Add Selected to Allowlist" + + async def _build_table(self, rebuild: bool = True) -> None: + """Rebuild the DataTable. If rebuild=False, only refresh rows.""" + if rebuild: + # Full rebuild: clear columns and rows + self.detail_table.clear() + self.detail_table.columns.clear() + self.row_key_to_id.clear() + self.table_row_to_id.clear() + + if self.activities_df is None or self.activities_df.empty: + logger.info("ActivityDetailWidget mounted with empty dataframe.") + return + + # Add columns (checkbox + data columns, excluding internal _row_id) + self.detail_table.add_column("Select", key="select") + for col in self.activities_df.columns: + if col != "_row_id": # Don't display the internal ID column + self.detail_table.add_column(col) + else: + # Partial rebuild: clear rows only + self.detail_table.clear() + self.row_key_to_id.clear() + self.table_row_to_id.clear() + # Add rows - for _, row in self.activities_df.iterrows(): - vals = ["" if pd.isna(v) else v for v in row.to_list()] - self.detail_table.add_row(*[str(v) for v in vals]) - # Allow sorting / cursor - self.detail_table.cursor_type = "row" + for table_idx, (df_idx, row) in enumerate(self.activities_df.iterrows()): + # Get the unique row ID + row_id = row["_row_id"] + + # Build values list (excluding _row_id column) + vals = [] + for col in self.activities_df.columns: + if col != "_row_id": + v = row[col] + vals.append("" if pd.isna(v) else str(v)) + + # Check if this row is selected + checkbox = "ā˜‘" if row_id in self.selected_row_ids else "☐" + + # Add row to table + row_key = self.detail_table.add_row(checkbox, *vals) + + # Map the row key and table index to the unique row ID + self.row_key_to_id[row_key] = row_id + self.table_row_to_id[table_idx] = row_id + + async def on_data_table_cell_selected(self, event: DataTable.CellSelected) -> None: + # Toggle selection when the "Select" column is clicked + if event.cell_key.column_key.value == "select": + table_row_index = event.coordinate.row + + # Get the unique row ID for this table row + row_id = self.table_row_to_id.get(table_row_index) + if row_id is not None: + # Get the row key for updating the cell + row_key = event.cell_key.row_key + + if row_id in self.selected_row_ids: + self.selected_row_ids.remove(row_id) + self.detail_table.update_cell(row_key, "select", "☐") # Unchecked + else: + self.selected_row_ids.add(row_id) + self.detail_table.update_cell(row_key, "select", "ā˜‘") # Checked + + self._update_button_states() + + async def on_data_table_header_selected( + self, event: DataTable.HeaderSelected + ) -> None: + column_key = event.column_key.value if event.column_key else None + if not column_key: + col_index = event.column_index + if col_index == 0: # First column is "Select" + return + # Adjust for hidden _row_id column + visible_cols = [ + col for col in self.activities_df.columns if col != "_row_id" + ] + if col_index - 1 < len(visible_cols): + column_key = visible_cols[col_index - 1] + else: + return + if column_key == "select" or column_key == "_row_id": + return + + ascending = True + if self._last_sort == (column_key, True): + ascending = False + self._last_sort = (column_key, ascending) + + try: + self.activities_df.sort_values( + by=column_key, ascending=ascending, inplace=True + ) + except Exception as exc: + logger.exception("Failed to sort by column %s: %s", column_key, exc) + return + + # āœ… Only refresh rows, not columns + await self._build_table(rebuild=False) async def on_button_pressed(self, event) -> None: btn = getattr(event, "button", None) or getattr(event, "sender", None) btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None) - # Back button in ActivityDetailWidget if btn is self.detail_back_btn or btn_id == "detail_back_btn": - # Pop screens until only the main menu remains while len(self.app.screen_stack) > 2: self.app.pop_screen() event.stop() return - # Export button - if btn is self.detail_export_btn or btn_id == "detail_export_btn": - await self._export_detail_activities() + + if btn is self.add_allowlist_btn or btn_id == "add_allowlist_btn": + await self._open_allowlist_screen() return - async def _export_detail_activities(self) -> None: + async def _select_all(self) -> None: + """Select all rows in the table.""" + # Add all row IDs to selected set + self.selected_row_ids = set(self.activities_df["_row_id"].tolist()) + # Update all checkboxes in the table + for row_key, row_id in self.row_key_to_id.items(): + self.detail_table.update_cell(row_key, "select", "ā˜‘") + + self._update_button_states() + logger.info(f"Selected all {len(self.selected_row_ids)} rows") + + async def _select_none(self) -> None: + """Deselect all rows in the table.""" + # Clear selected set + self.selected_row_ids.clear() + + # Update all checkboxes in the table + for row_key, row_id in self.row_key_to_id.items(): + self.detail_table.update_cell(row_key, "select", "☐") + + self._update_button_states() + logger.info("Cleared all selections") + + async def _open_allowlist_screen(self) -> None: + """Open the allowlist selection screen with selected activities.""" + if not self.selected_row_ids: + self.app.notify("No rows selected", severity="warning") + return + + # Get selected data + selected_df = self.get_selected_data() + + # Get API from app + api = getattr(self.app, "api", None) + if api is None: + logger.error("No API available on self.app.api") + self.app.notify("API not available", severity="error") + return + + # Create and push AllowlistSelectionScreen + try: + allowlist_screen = AllowlistSelectionScreen( + selected_df, api=api, hostname=self.hostname, otpid=self.otpid + ) + await self.app.push_screen(allowlist_screen) + logger.info( + f"Opened allowlist screen with {len(selected_df)} selected activities" + ) + except ImportError as e: + logger.error(f"Failed to import AllowlistSelectionScreen: {e}") + self.app.notify("Allowlist screen module not found", severity="error") + except Exception as e: + logger.exception(f"Failed to open allowlist screen: {e}") + self.app.notify( + f"Error opening allowlist screen: {str(e)}", severity="error" + ) + + async def _export_detail_activities(self) -> None: if self.activities_df is None or self.activities_df.empty: logger.info("No activities to export.") - notification = Static("āŒ No activities to export.", classes="notification") - self.mount(notification) + await self.mount( + Static("āŒ No activities to export.", classes="notification") + ) + return + if not self.selected_row_ids: + logger.info("No rows selected for export.") + await self.mount( + Static("āŒ No rows selected for export.", classes="notification") + ) return try: working_dir = load_env("WORKING_DIR") or os.getcwd() timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") filename = f"otp_activities_detail_{timestamp}.csv" file_path = os.path.join(working_dir, filename) - self.activities_df.to_csv(file_path, index=False) - logger.info("Exported detail activities to %s", file_path) - # Show success notification - notification = Static( - f"āœ… Exported activities to: {filename}", classes="notification" + selected_df = self.get_selected_data() + selected_df.to_csv(file_path, index=False) + logger.info("Exported selected activities to %s", file_path) + await self.mount( + Static( + f"āœ… Exported selected activities to: {filename}", + classes="notification", + ) ) - self.mount(notification) except Exception as exc: logger.exception("Failed to export detail activities: %s", exc) - notification = Static( - "āŒ Failed to export activities; check logs.", classes="notification" + await self.mount( + Static( + "āŒ Failed to export activities; check logs.", + classes="notification", + ) ) - self.mount(notification) + + # āœ… Helper methods + def get_selected_data(self) -> pd.DataFrame: + """Return a DataFrame of the selected rows.""" + if not self.selected_row_ids: + return pd.DataFrame() + # Filter by selected row IDs and drop the internal _row_id column + selected_df = self.activities_df[ + self.activities_df["_row_id"].isin(self.selected_row_ids) + ].copy() + if "_row_id" in selected_df.columns: + selected_df = selected_df.drop(columns=["_row_id"]) + return selected_df + + def get_selected_records(self) -> list[dict]: + """Return selected rows as a list of dicts.""" + if not self.selected_row_ids: + return [] + # Filter by selected row IDs and drop the internal _row_id column + selected_df = self.activities_df[ + self.activities_df["_row_id"].isin(self.selected_row_ids) + ].copy() + if "_row_id" in selected_df.columns: + selected_df = selected_df.drop(columns=["_row_id"]) + return selected_df.to_dict(orient="records") class ActivityDetailScreen(Screen): @@ -590,7 +791,12 @@ class ActivityDetailScreen(Screen): Screen that wraps ActivityDetailWidget. Expects a DataFrame passed on init. """ - BINDINGS = [Binding("b", "back", "Back"), Binding("e", "export", "Export")] + BINDINGS = [ + Binding("b", "back", "Back"), + Binding("e", "export", "Export"), + Binding("a", "select_all", "Select All"), + Binding("n", "select_none", "Select None"), + ] def __init__(self, activities_df: pd.DataFrame, otpid=None, hostname=None) -> None: super().__init__() @@ -621,6 +827,16 @@ class ActivityDetailScreen(Screen): if hasattr(self, "widget") and self.widget is not None: await self.widget._export_detail_activities() + async def action_select_all(self) -> None: + """Handle 'a' key for select all.""" + if hasattr(self, "widget") and self.widget is not None: + await self.widget._select_all() + + async def action_select_none(self) -> None: + """Handle 'n' key for select none.""" + if hasattr(self, "widget") and self.widget is not None: + await self.widget._select_none() + class OTPActivitiesScreen(Screen): """ diff --git a/screens/otpworkflowscreen.py b/TUI/otpworkflowscreen.py similarity index 94% rename from screens/otpworkflowscreen.py rename to TUI/otpworkflowscreen.py index 519ff6b..d7d7322 100644 --- a/screens/otpworkflowscreen.py +++ b/TUI/otpworkflowscreen.py @@ -6,7 +6,7 @@ from textual.app import ComposeResult from textual.screen import Screen from models.agent import Agent -from widgets.OTP_generate import OTPGenerator +from TUI.OTP_generate import OTPGenerator class OTPWorkflowScreen(Screen): diff --git a/widgets/policyselector.py b/TUI/policyselector.py similarity index 99% rename from widgets/policyselector.py rename to TUI/policyselector.py index fb34218..d6e4bdf 100644 --- a/widgets/policyselector.py +++ b/TUI/policyselector.py @@ -14,7 +14,7 @@ import pandas as pd from textual.containers import Horizontal, Vertical from textual.message import Message from textual.widget import Widget -from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea +from textual.widgets import Button, DataTable, Static, TextArea from models.policy import Policy @@ -93,7 +93,6 @@ class PolicySelector(Widget): - Policy table displaying available policies - Back buttons for navigation """ - yield Header(show_clock=True, icon="āš™") title_text = Static( "šŸŽÆ Select Target Policy", id="policy_selector_title", @@ -164,8 +163,6 @@ class PolicySelector(Widget): policy_table.styles.margin = (1, 0, 1, 0) yield policy_table - yield Footer() - def on_mount(self) -> None: """ Initialize the policy table when the widget is mounted. diff --git a/screens/policyselectorscreen.py b/TUI/policyselectorscreen.py similarity index 98% rename from screens/policyselectorscreen.py rename to TUI/policyselectorscreen.py index 5937860..07ec2ee 100644 --- a/screens/policyselectorscreen.py +++ b/TUI/policyselectorscreen.py @@ -25,7 +25,7 @@ import logging from textual.app import ComposeResult from textual.screen import Screen -from widgets.policyselector import PolicySelector +from TUI.policyselector import PolicySelector logger = logging.getLogger(__name__) diff --git a/widgets/policytreewidget.py b/TUI/policytreewidget.py similarity index 100% rename from widgets/policytreewidget.py rename to TUI/policytreewidget.py diff --git a/TUI/quietagentworkflowscreen.py b/TUI/quietagentworkflowscreen.py new file mode 100644 index 0000000..552f214 --- /dev/null +++ b/TUI/quietagentworkflowscreen.py @@ -0,0 +1,848 @@ +# 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 . + +""" +Quiet Agent Workflow Screen Module + +Provides a TUI workflow for identifying quiet agents and moving them to target policies. +This screen replaces the legacy quietAgent.py with a comprehensive TUI interface that: +1. Allows selection of an initial policy to analyze +2. Categorizes devices into "Enforce Ready" and "Non-Enforce Ready" based on activity +3. Allows users to select target policies for each category +4. Uses the API to move devices to their target policies +""" + +import datetime +import logging +import os +from typing import List, Optional + +import pandas as pd +from textual.app import ComposeResult +from textual.containers import Horizontal, Vertical +from textual.reactive import reactive +from textual.screen import Screen +from textual.widgets import Button, DataTable, Footer, Header, Static + +from models.policy import Policy +from services.API import AirlockAPIWrapper +from services.policyhandler import getPolicyInfo +from TUI.policyselector import PolicySelector +from utils.configmanager import load_env + +logger = logging.getLogger(__name__) + + +class QuietAgentWorkflowScreen(Screen): + """ + A Textual screen for the Quiet Agent analysis and migration workflow. + + This screen provides a multi-step workflow: + 1. Select initial policy to analyze + 2. View categorized agents (enforce ready vs. non-enforce ready) + 3. Select target policies for each category + 4. Execute agent migrations + + Attributes: + api (AirlockAPIWrapper): API wrapper for Airlock operations + policies (List[Policy]): List of all available policies + selected_policy (Optional[Policy]): The initially selected policy to analyze + history_days (int): Number of days of history to pull (default: 150) + quiet_days (int): Number of days without execution to be considered quiet (default: 45) + agents_df (Optional[pd.DataFrame]): DataFrame of all agents with analysis results + enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents ready for enforcement + non_enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents not ready for enforcement + workflow_stage (str): Current stage of the workflow + """ + + BINDINGS = [ + ("escape", "go_back", "Back"), + ] + + workflow_stage = reactive("select_policy") # Tracks current workflow stage + + def __init__(self, api: AirlockAPIWrapper, policies: List[Policy]): + """ + Initialize the QuietAgentWorkflowScreen. + + Args: + api (AirlockAPIWrapper): API wrapper for Airlock operations + policies (List[Policy]): List of all available policies + """ + super().__init__() + self.api = api + self.policies = policies + self.selected_policy: Optional[Policy] = None + self.history_days = 150 # Fixed as per requirements + self.quiet_days = 45 # Default value + self.agents_df: Optional[pd.DataFrame] = None + self.enforce_ready_df: Optional[pd.DataFrame] = None + self.non_enforce_ready_df: Optional[pd.DataFrame] = None + self.enforce_ready_target_policy: Optional[Policy] = None + self.non_enforce_ready_target_policy: Optional[Policy] = None + + def compose(self) -> ComposeResult: + """Build the UI layout for the workflow screen.""" + # Include Header and Footer like other standalone screens + yield Header(show_clock=True, icon="āš™") + + # Title area + title = Static("šŸ”’ Quiet Agent Workflow", id="workflow_title") + title.styles.margin = (0, 0, 0, 1) + yield title + + # Status area + status = Static("Step 1: Select Policy to Analyze", id="workflow_status") + status.styles.margin = (0, 0, 1, 1) + yield status + + # Content area - dynamically populated based on workflow stage + yield Vertical(id="content_area") + + yield Footer() + + def on_mount(self) -> None: + """Initialize the screen when mounted.""" + # Show initial policy selection + self._show_policy_selection() + + def watch_workflow_stage(self, old_value: str, new_value: str) -> None: + """React to workflow stage changes.""" + logger.debug(f"Workflow stage changed from {old_value} to {new_value}") + self._update_status_message() + + def _update_status_message(self) -> None: + """Update the status message based on current workflow stage.""" + status_widget = self.query_one("#workflow_status", Static) + + stage_messages = { + "select_policy": "Step 1: Select Policy to Analyze", + "select_quiet_days": "Step 2: Select Quiet Time Period", + "analyzing": "šŸ“Š Analyzing agent activity...", + "view_results": "Step 3: Review Categorized Agents", + "select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents", + "select_non_enforce_target": "Step 5: Select Target Policy for Non-Enforce Ready Agents", + "confirm_migration": "Step 6: Confirm and Execute Migration", + "executing": "ā³ Executing agent migrations...", + "complete": "āœ… Migration Complete", + } + + status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage")) + + def _show_policy_selection(self) -> None: + """Show the initial policy selection screen.""" + self.workflow_stage = "select_policy" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # Create policy selector widget + policy_selector = PolicySelector(self.policies) + content.mount(policy_selector) + + def on_policy_selector_policy_selected( + self, message: PolicySelector.PolicySelected + ) -> None: + """Handle policy selection from PolicySelector widget.""" + # Handle based on current workflow stage + if self.workflow_stage == "select_policy": + # Initial policy selection for analysis + self.selected_policy = message.policy + logger.info(f"Selected policy for analysis: {self.selected_policy.name}") + self._show_quiet_days_selection() + elif self.workflow_stage == "select_enforce_target": + # Target policy selection for enforce ready agents + self.enforce_ready_target_policy = message.policy + logger.info( + f"Selected target policy for enforce ready: {message.policy.name}" + ) + self._show_non_enforce_target_selection() + elif self.workflow_stage == "select_non_enforce_target": + # Target policy selection for non-enforce ready agents + self.non_enforce_ready_target_policy = message.policy + logger.info( + f"Selected target policy for non-enforce ready: {message.policy.name}" + ) + self._show_migration_confirmation() + + def _show_quiet_days_selection(self) -> None: + """Show the quiet days selection screen.""" + self.workflow_stage = "select_quiet_days" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # Create info text + info_widget = Static( + f"Policy Selected: {self.selected_policy.name}\n\n" + f"History Period: {self.history_days} days\n\n" + "Select quiet time period (days without untrusted execution):", + id="quiet_days_info", + ) + info_widget.styles.margin = (0, 0, 2, 0) + content.mount(info_widget) + + # Create button container and mount it first + button_container = Vertical(id="quiet_days_buttons") + button_container.styles.height = "auto" + content.mount(button_container) + + # Now add buttons to the mounted container + for days in [15, 30, 45, 60]: + btn = Button( + f"{days} days {'(Default)' if days == 45 else ''}", + id=f"quiet_days_{days}", + classes="quiet_day_btn", + ) + btn.styles.width = "100%" + btn.styles.margin = (0, 0, 1, 0) + button_container.mount(btn) + + back_btn = Button("← Back", id="back_to_policy_selection") + back_btn.styles.width = "100%" + back_btn.styles.margin = (2, 0, 0, 0) + button_container.mount(back_btn) + + def on_button_pressed(self, event: Button.Pressed) -> None: + """Handle button press events.""" + button_id = event.button.id + + # Quiet days selection buttons + if button_id and button_id.startswith("quiet_days_"): + days = int(button_id.split("_")[-1]) + self.quiet_days = days + logger.info(f"Selected quiet days: {days}") + self._start_analysis() + return + + # Navigation buttons + if button_id == "back_to_policy_selection": + self._show_policy_selection() + return + + if button_id == "back_to_results": + self._show_results() + return + + if button_id == "select_enforce_target_btn": + self._show_enforce_target_selection() + return + + if button_id == "select_non_enforce_target_btn": + self._show_non_enforce_target_selection() + return + + if button_id == "skip_enforce_target_btn": + # Skip enforce ready target selection + self.enforce_ready_target_policy = None + self._show_non_enforce_target_selection() + return + + if button_id == "skip_non_enforce_target_btn": + # Skip non-enforce ready target selection + self.non_enforce_ready_target_policy = None + self._show_migration_confirmation() + return + + if button_id == "confirm_migration_btn": + self._execute_migration() + return + + if button_id == "cancel_migration_btn": + self._show_results() + return + + if button_id == "export_results_btn": + self._export_results() + return + + if button_id == "start_over_btn": + self._show_policy_selection() + return + + def _start_analysis(self) -> None: + """Start the agent activity analysis.""" + self.workflow_stage = "analyzing" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # Show analyzing message with detailed steps + analyzing_msg = Static( + f"šŸ“Š Analyzing Agent Activity\n" + f"{'=' * 50}\n\n" + f"Policy: {self.selected_policy.name}\n" + f"History Period: {self.history_days} days\n" + f"Quiet Threshold: {self.quiet_days} days\n\n" + f"Progress:\n" + f"ā³ Step 1/4: Fetching agents from policy...\n" + f"ā±ļø Step 2/4: Pulling execution history (this may take a moment)...\n" + f"ā±ļø Step 3/4: Analyzing activity patterns...\n" + f"ā±ļø Step 4/4: Categorizing agents...\n\n" + f"Please wait - this operation cannot be cancelled.", + id="analyzing_message", + ) + analyzing_msg.styles.margin = (2, 1) + content.mount(analyzing_msg) + + # Show notification + self.app.notify( + "Starting analysis - this may take several minutes for large policies", + severity="information", + timeout=5, + ) + + # Perform the analysis asynchronously + self.call_later(self._perform_analysis) + + def _perform_analysis(self) -> None: + """Perform the actual agent activity analysis.""" + try: + # Update status: Fetching agents + self._update_analysis_status("Step 1/4: Fetching agents from policy...") + + # Get agents in the selected policy + agents = self.api.agents_find_by_group(self.selected_policy.groupid) + + if agents.empty: + self.app.notify( + f"No agents found in policy: {self.selected_policy.name}", + severity="warning", + timeout=5, + ) + self._show_policy_selection() + return + + agent_count = len(agents) + self.app.notify( + f"Found {agent_count} agents - fetching execution history...", + severity="information", + timeout=3, + ) + + # Update status: Pulling execution history + self._update_analysis_status( + f"Step 2/4: Pulling execution history for {agent_count} agents...\n" + f"(This may take several minutes - progress shown in terminal)" + ) + + # Get execution history (this shows progress bars in terminal via airlock_libs) + policy_exec_history = getPolicyInfo( + self.api, self.selected_policy, [1, 2, 6, 7], self.history_days + ) + + # Update status: Analyzing patterns + self._update_analysis_status("Step 3/4: Analyzing activity patterns...") + self.app.notify( + "History retrieved - analyzing patterns...", + severity="information", + timeout=2, + ) + + if policy_exec_history.empty: + logger.info( + "No execution history found for the selected policy and time range." + ) + # All agents are quiet (no executions) + agents["execution_count"] = 0 + agents["days_since"] = None + agents["required_quiet"] = self.quiet_days + agents["enforce_ready"] = True + else: + # Convert datetime column + policy_exec_history["datetime"] = pd.to_datetime( + policy_exec_history["datetime"], + format="%Y-%m-%dT%H:%M:%SZ", + utc=True, + ) + + # Calculate days ago + now = datetime.datetime.now(datetime.timezone.utc) + policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply( + lambda dt: (now - dt).days + ) + + # Count total executions per hostname + hostname_counts = policy_exec_history["hostname"].value_counts() + agents["execution_count"] = ( + agents["hostname"].map(hostname_counts).fillna(0).astype(int) + ) + + # Find most recent execution per hostname + most_recent_exec = policy_exec_history.sort_values( + by="days_ago" + ).drop_duplicates(subset="hostname", keep="first") + + # Map most recent execution age to agents + agents["days_since"] = agents["hostname"].map( + most_recent_exec.set_index("hostname")["days_ago"] + ) + + # Check for enforcement readiness + agents["required_quiet"] = self.quiet_days + agents["enforce_ready"] = agents["days_since"].apply( + lambda x: True if pd.isna(x) or x > self.quiet_days else False + ) + + # Update status: Categorizing + self._update_analysis_status("Step 4/4: Categorizing agents...") + + # Sort agents + agents = agents.sort_values( + by=["execution_count", "hostname"], ascending=[True, True] + ) + + # Store the results + self.agents_df = agents + + # Categorize agents into DataFrames + self.enforce_ready_df = agents[agents["enforce_ready"] == True].copy() + self.non_enforce_ready_df = agents[agents["enforce_ready"] == False].copy() + + logger.info( + f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, " + f"{len(self.non_enforce_ready_df)} non-enforce ready" + ) + + self.app.notify( + f"Analysis complete! Found {len(self.enforce_ready_df)} enforce ready, " + f"{len(self.non_enforce_ready_df)} not ready", + severity="success", + timeout=5, + ) + + # Show results + self._show_results() + + except Exception as e: + logger.error(f"Error during analysis: {e}", exc_info=True) + self.app.notify(f"Analysis failed: {str(e)}", severity="error", timeout=5) + self._show_policy_selection() + + def _update_analysis_status(self, status_text: str) -> None: + """Update the analysis status message.""" + try: + analyzing_msg = self.query_one("#analyzing_message", Static) + + # Build updated message + updated_text = ( + f"šŸ“Š Analyzing Agent Activity\n" + f"{'=' * 50}\n\n" + f"Policy: {self.selected_policy.name}\n" + f"History Period: {self.history_days} days\n" + f"Quiet Threshold: {self.quiet_days} days\n\n" + f"Progress:\n" + f"āœ… {status_text}\n\n" + f"Please wait - this operation cannot be cancelled." + ) + + analyzing_msg.update(updated_text) + except Exception as e: + logger.debug(f"Could not update analysis status: {e}") + + def _show_results(self) -> None: + """Show the categorized results.""" + self.workflow_stage = "view_results" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # Create results display container and mount it first + results_container = Vertical(id="results_container") + results_container.styles.height = "auto" + results_container.styles.margin = (1, 1) + content.mount(results_container) + + # Summary statistics + total_agents = len(self.enforce_ready_df) + len(self.non_enforce_ready_df) + ready_count = len(self.enforce_ready_df) + not_ready_count = len(self.non_enforce_ready_df) + ready_percentage = (ready_count / total_agents * 100) if total_agents > 0 else 0 + + summary = Static( + f"Analysis Results for: {self.selected_policy.name}\n\n" + f"šŸ“Š Total Agents: {total_agents}\n" + f"āœ… Enforce Ready: {ready_count} ({ready_percentage:.1f}%)\n" + f"āŒ Not Ready: {not_ready_count} ({100 - ready_percentage:.1f}%)\n\n" + f"Quiet Threshold: {self.quiet_days} days\n" + f"History Period: {self.history_days} days", + id="results_summary", + ) + summary.styles.margin = (0, 0, 2, 0) + results_container.mount(summary) + + # Action buttons + button_container = Horizontal(id="results_buttons") + button_container.styles.height = "auto" + results_container.mount(button_container) + + if ready_count > 0: + enforce_btn = Button( + f"Select Target for Enforce Ready ({ready_count})", + id="select_enforce_target_btn", + ) + enforce_btn.styles.margin = (0, 1, 1, 0) + button_container.mount(enforce_btn) + + if not_ready_count > 0: + non_enforce_btn = Button( + f"Select Target for Non-Enforce Ready ({not_ready_count})", + id="select_non_enforce_target_btn", + ) + non_enforce_btn.styles.margin = (0, 1, 1, 0) + button_container.mount(non_enforce_btn) + + export_btn = Button("šŸ’¾ Export Results", id="export_results_btn") + export_btn.styles.margin = (0, 1, 1, 0) + button_container.mount(export_btn) + + start_over_btn = Button("šŸ”„ Start Over", id="start_over_btn") + start_over_btn.styles.margin = (0, 0, 1, 0) + button_container.mount(start_over_btn) + + # Tables showing agents + tables_container = Horizontal() + tables_container.styles.height = "1fr" + results_container.mount(tables_container) + + # Enforce Ready table + if ready_count > 0: + enforce_col = Vertical() + enforce_col.styles.width = "1fr" + enforce_col.styles.margin = (1, 1, 0, 0) + tables_container.mount(enforce_col) + + enforce_label = Static("āœ… Enforce Ready Agents") + enforce_label.styles.margin = (0, 0, 1, 0) + enforce_col.mount(enforce_label) + + enforce_table = DataTable(id="enforce_ready_table") + enforce_table.styles.height = "1fr" + enforce_table.add_columns("Hostname", "Last Exec (days)") + + # Display first 50 agents + for idx, row in self.enforce_ready_df.head(50).iterrows(): + days_since = row["days_since"] + days_str = f"{int(days_since)}" if not pd.isna(days_since) else "Never" + enforce_table.add_row(row["hostname"], days_str) + + if len(self.enforce_ready_df) > 50: + enforce_table.add_row( + f"... and {len(self.enforce_ready_df) - 50} more", "" + ) + + enforce_col.mount(enforce_table) + + # Non-Enforce Ready table + if not_ready_count > 0: + non_enforce_col = Vertical() + non_enforce_col.styles.width = "1fr" + non_enforce_col.styles.margin = (1, 0, 0, 1) + tables_container.mount(non_enforce_col) + + non_enforce_label = Static("āŒ Non-Enforce Ready Agents") + non_enforce_label.styles.margin = (0, 0, 1, 0) + non_enforce_col.mount(non_enforce_label) + + non_enforce_table = DataTable(id="non_enforce_ready_table") + non_enforce_table.styles.height = "1fr" + non_enforce_table.add_columns("Hostname", "Last Exec (days)") + + # Display first 50 agents + for idx, row in self.non_enforce_ready_df.head(50).iterrows(): + days_since = row["days_since"] + days_str = f"{int(days_since)}" if not pd.isna(days_since) else "N/A" + non_enforce_table.add_row(row["hostname"], days_str) + + if len(self.non_enforce_ready_df) > 50: + non_enforce_table.add_row( + f"... and {len(self.non_enforce_ready_df) - 50} more", "" + ) + + non_enforce_col.mount(non_enforce_table) + + def _show_enforce_target_selection(self) -> None: + """Show policy selection for enforce ready agents.""" + self.workflow_stage = "select_enforce_target" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # Info message + info = Static( + f"Select target policy for {len(self.enforce_ready_df)} Enforce Ready agents\n" + f"Source Policy: {self.selected_policy.name}", + id="enforce_target_info", + ) + info.styles.margin = (0, 0, 2, 0) + content.mount(info) + + # Policy selector + policy_selector = PolicySelector(self.policies) + content.mount(policy_selector) + + # Skip button + skip_btn = Button("ā­•ļø Skip - No Migration", id="skip_enforce_target_btn") + skip_btn.styles.width = "50%" + skip_btn.styles.margin = (2, 0, 0, 0) + content.mount(skip_btn) + + def _show_non_enforce_target_selection(self) -> None: + """Show policy selection for non-enforce ready agents.""" + self.workflow_stage = "select_non_enforce_target" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # Info message + info = Static( + f"Select target policy for {len(self.non_enforce_ready_df)} Non-Enforce Ready agents\n" + f"Source Policy: {self.selected_policy.name}", + id="non_enforce_target_info", + ) + info.styles.margin = (0, 0, 2, 0) + content.mount(info) + + # Policy selector + policy_selector = PolicySelector(self.policies) + content.mount(policy_selector) + + # Skip button + skip_btn = Button("ā­•ļø Skip - No Migration", id="skip_non_enforce_target_btn") + skip_btn.styles.width = "50%" + skip_btn.styles.margin = (2, 0, 0, 0) + content.mount(skip_btn) + + def _show_migration_confirmation(self) -> None: + """Show migration confirmation screen.""" + self.workflow_stage = "confirm_migration" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # Build confirmation message + confirmation_lines = [ + "šŸ” Migration Summary\n", + f"Source Policy: {self.selected_policy.name}\n", + ] + + if self.enforce_ready_target_policy: + confirmation_lines.append( + f"\nāœ… Enforce Ready Migration:\n" + f" • Agents: {len(self.enforce_ready_df)}\n" + f" • Target: {self.enforce_ready_target_policy.name}\n" + ) + + if self.non_enforce_ready_target_policy: + confirmation_lines.append( + f"\nāŒ Non-Enforce Ready Migration:\n" + f" • Agents: {len(self.non_enforce_ready_df)}\n" + f" • Target: {self.non_enforce_ready_target_policy.name}\n" + ) + + if ( + not self.enforce_ready_target_policy + and not self.non_enforce_ready_target_policy + ): + confirmation_lines.append("\nāš ļø No migrations will be performed.") + + confirmation = Static("".join(confirmation_lines), id="migration_confirmation") + confirmation.styles.margin = (1, 1, 2, 1) + content.mount(confirmation) + + # Action buttons - mount container first, then add buttons + button_container = Horizontal(id="confirmation_buttons") + button_container.styles.height = "auto" + button_container.styles.margin = (1, 1) + content.mount(button_container) + + if self.enforce_ready_target_policy or self.non_enforce_ready_target_policy: + confirm_btn = Button("āœ… Confirm Migration", id="confirm_migration_btn") + confirm_btn.styles.margin = (0, 1, 0, 0) + button_container.mount(confirm_btn) + + cancel_btn = Button("āŒ Cancel", id="cancel_migration_btn") + button_container.mount(cancel_btn) + + def _execute_migration(self) -> None: + """Execute the agent migrations.""" + self.workflow_stage = "executing" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # Show executing message + executing_msg = Static( + "ā³ Executing agent migrations...\nPlease wait...", + id="executing_message", + ) + executing_msg.styles.margin = (2, 1) + content.mount(executing_msg) + + # Perform migrations asynchronously + self.call_later(self._perform_migrations) + + def _perform_migrations(self) -> None: + """Perform the actual agent migrations.""" + successful_migrations = [] + failed_migrations = [] + + try: + # Migrate enforce ready agents + if self.enforce_ready_target_policy: + for idx, row in self.enforce_ready_df.iterrows(): + try: + result = self.api.agent_move( + row["agentid"], self.enforce_ready_target_policy.groupid + ) + successful_migrations.append( + (row["hostname"], self.enforce_ready_target_policy.name) + ) + logger.debug( + f"Moved {row['hostname']} to {self.enforce_ready_target_policy.name}" + ) + except Exception as e: + failed_migrations.append((row["hostname"], str(e))) + logger.error(f"Failed to move {row['hostname']}: {e}") + + # Migrate non-enforce ready agents + if self.non_enforce_ready_target_policy: + for idx, row in self.non_enforce_ready_df.iterrows(): + try: + result = self.api.agent_move( + row["agentid"], self.non_enforce_ready_target_policy.groupid + ) + successful_migrations.append( + (row["hostname"], self.non_enforce_ready_target_policy.name) + ) + logger.debug( + f"Moved {row['hostname']} to {self.non_enforce_ready_target_policy.name}" + ) + except Exception as e: + failed_migrations.append((row["hostname"], str(e))) + logger.error(f"Failed to move {row['hostname']}: {e}") + + # Show completion results + self._show_completion_results(successful_migrations, failed_migrations) + + except Exception as e: + logger.error(f"Error during migration execution: {e}", exc_info=True) + self.app.notify(f"Migration failed: {str(e)}", severity="error", timeout=5) + self._show_results() + + def _show_completion_results( + self, successful: List[tuple], failed: List[tuple] + ) -> None: + """Show migration completion results.""" + self.workflow_stage = "complete" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # Results summary + total_attempted = len(successful) + len(failed) + success_rate = ( + (len(successful) / total_attempted * 100) if total_attempted > 0 else 0 + ) + + results = Static( + f"āœ… Migration Complete\n\n" + f"Total Agents Migrated: {len(successful)}\n" + f"Failed Migrations: {len(failed)}\n" + f"Success Rate: {success_rate:.1f}%", + id="completion_summary", + ) + results.styles.margin = (1, 1, 2, 1) + content.mount(results) + + # Details tables + if successful: + success_container = Vertical() + success_container.styles.margin = (0, 1) + content.mount(success_container) + + success_label = Static("āœ… Successful Migrations") + success_label.styles.margin = (0, 0, 1, 0) + success_container.mount(success_label) + + success_table = DataTable(id="success_table") + success_table.styles.height = "auto" + success_table.add_columns("Hostname", "Target Policy") + + for hostname, target_policy in successful[:25]: # Show first 25 + success_table.add_row(hostname, target_policy) + + if len(successful) > 25: + success_table.add_row(f"... and {len(successful) - 25} more", "") + + success_container.mount(success_table) + + if failed: + failed_container = Vertical() + failed_container.styles.margin = (2, 1, 0, 1) + content.mount(failed_container) + + failed_label = Static("āŒ Failed Migrations") + failed_label.styles.margin = (0, 0, 1, 0) + failed_container.mount(failed_label) + + failed_table = DataTable(id="failed_table") + failed_table.styles.height = "auto" + failed_table.add_columns("Hostname", "Error") + + for hostname, error in failed[:25]: # Show first 25 + failed_table.add_row(hostname, error[:50]) # Truncate error + + if len(failed) > 25: + failed_table.add_row(f"... and {len(failed) - 25} more", "") + + failed_container.mount(failed_table) + + # Action button + done_btn = Button("āœ” Done", id="start_over_btn") + done_btn.styles.width = "50%" + done_btn.styles.margin = (2, 0, 0, 0) + content.mount(done_btn) + + def _export_results(self) -> None: + """Export analysis results to CSV.""" + try: + working_dir = load_env("WORKING_DIR") or os.getcwd() + filename = os.path.join( + working_dir, + f"{self.selected_policy.name}_quiet_analysis_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv", + ) + + self.agents_df.to_csv(filename, index=False) + logger.info(f"Exported results to {filename}") + self.app.notify( + f"Results exported to:\n{filename}", + severity="information", + timeout=5, + ) + + except Exception as e: + logger.error(f"Failed to export results: {e}") + self.app.notify(f"Export failed: {str(e)}", severity="error", timeout=5) + + def action_go_back(self) -> None: + """Handle back/escape action.""" + # Depending on stage, go back to previous stage or exit + if self.workflow_stage in ["select_policy", "view_results", "complete"]: + self.app.pop_screen() + elif self.workflow_stage == "select_quiet_days": + self._show_policy_selection() + elif self.workflow_stage == "select_enforce_target": + self._show_results() + elif self.workflow_stage == "select_non_enforce_target": + if self.enforce_ready_target_policy: + self._show_enforce_target_selection() + else: + self._show_results() + elif self.workflow_stage == "confirm_migration": + self._show_non_enforce_target_selection() + else: + self.app.pop_screen() diff --git a/widgets/resultsdisplay.py b/TUI/resultsdisplay.py similarity index 100% rename from widgets/resultsdisplay.py rename to TUI/resultsdisplay.py diff --git a/themes/amber_terminal_theme.py b/TUI/theme_amber_terminal.py similarity index 93% rename from themes/amber_terminal_theme.py rename to TUI/theme_amber_terminal.py index 169c325..3662892 100644 --- a/themes/amber_terminal_theme.py +++ b/TUI/theme_amber_terminal.py @@ -12,7 +12,7 @@ def get_amber_terminal_theme(): success=Color.parse("#ffb733"), warning=Color.parse("#ffff66"), error=Color.parse("#ff3300"), - surface=Color.parse("#3a1f00"), # brighter brown for blending + surface=Color.parse("#49331a"), # brighter brown for blending ) diff --git a/themes/retro_terminal_theme.py b/TUI/theme_retro_terminal.py similarity index 100% rename from themes/retro_terminal_theme.py rename to TUI/theme_retro_terminal.py diff --git a/widgets/themeselector.py b/TUI/themeselector.py similarity index 100% rename from widgets/themeselector.py rename to TUI/themeselector.py diff --git a/flows/otp.py b/flows/otp.py index 720aab3..ee07c98 100644 --- a/flows/otp.py +++ b/flows/otp.py @@ -14,98 +14,18 @@ # along with this program. If not, see . -from datetime import datetime import logging -import os import pandas as pd from services.agenthandler import selectAgents from services.API import AirlockAPIWrapper -from utils.configmanager import load_env from utils.selector import Selector -from utils.utils import colorText, get_sanitized_input +from utils.utils import get_sanitized_input logger = logging.getLogger(__name__) -def otp_activities_by_agent(api: AirlockAPIWrapper): - activeagents = api.otp_find_active() - awaitingagents = api.otp_find_awaiting() - enforcedagents = api.otp_find_enforced() - revokedagents = api.otp_find_revoked() - - # Add a 'status' column to each DataFrame - activeagents["status"] = "active" - awaitingagents["status"] = "awaiting" - enforcedagents["status"] = "enforced" - revokedagents["status"] = "revoked" - - # Combine all into one DataFrame - combined_agents = pd.concat( - [activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True - ) - combined_agents = combined_agents.sort_values(by="otpid", ascending=False) - - # Optionally, select specific hosts - user_input = ( - get_sanitized_input("\nWould you like to search for a specific device? (y/n): ") - .strip() - .lower() - ) - if user_input == "y": - agentnames = [] - agents = selectAgents(api) - for agent in agents: - agentnames.append(agent.hostname) - - combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)] - - # Present and select rows - selected_rows = Selector.select_dataframe_with_mode( - combined_agents, - columns=["otpid", "hostname", "status", "purpose", "granted"], - header="OTP Sessions", - ) - combined_df = pd.DataFrame() - - for row in selected_rows: - otpid = row["otpid"] - hostname = row["hostname"] - result = api.otp_get_activities(otpid) - result["hostname"] = hostname - if not result.empty: - logger.info(f"Activities for {hostname} (otpid: {otpid}):\n{result}") - combined_df = pd.concat([combined_df, result], ignore_index=True) - else: - logger.info(f"No activities found for {hostname} (otpid: {otpid})") - - user_input = ( - get_sanitized_input( - "\nWould you like to export the results to a CSV file? (y/n): " - ) - .strip() - .lower() - ) - if user_input == "y": - working_dir = load_env("WORKING_DIR") - timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - filename = f"otp_activities_{timestamp}.csv" - file_path = os.path.join(str(working_dir), filename) - - combined_df.to_csv(file_path, index=False) - logging.info(f"Exported Data to {file_path}") - - print( - colorText( - f"\nāœ… OTP Activity exported to: {working_dir}\\{filename}", - "green", - ) - ) - else: - logging.debug("User declined to export the DataFrame.") - - def otp_revoke(api: AirlockAPIWrapper): activeagents = api.otp_find_active() diff --git a/services/API.py b/services/API.py index 34ce788..7fab0ac 100644 --- a/services/API.py +++ b/services/API.py @@ -253,6 +253,19 @@ class AirlockAPIWrapper: } return self._post("/v1/group/settings/script_custom", payload) + def policy_set_upgradetarget( + self, + groupid: str, + windows: str, + macos: str, + ) -> dict: + payload = { + "groupid": groupid, + "windows": windows, + "macos": macos, + } + return self._post("/v1/group/settings/selfupgrade/target", payload) + # Execution History def history_logging( self, type: List[str], checkpoint: str, policy: List[str]