From 1d9caadaf3af788580e1d7c388479e8ca1cc26b4 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Tue, 2 Dec 2025 16:22:43 -0500 Subject: [PATCH] Major step towards unification of the UI Implementation of the Back Feature, splitting of TUI files back into subfolders --- Loxide.py | 3 +- TUI/{ => Screens}/allowlistselectionscreen.py | 55 +- TUI/Screens/moveagentworkflowscreen.py | 114 ++ TUI/{ => Screens}/otpactivityscreen.py | 110 +- TUI/Screens/otprevokescreen.py | 393 ++++ TUI/Screens/otpworkflowscreen.py | 52 + TUI/Screens/policyprepworkflowscreen.py | 1822 +++++++++++++++++ TUI/{ => Screens}/policyselectorscreen.py | 20 +- TUI/{ => Screens}/quietagentworkflowscreen.py | 101 +- TUI/TUI.py | 193 +- TUI/{ => Themes}/theme_amber_terminal.py | 15 + TUI/{ => Themes}/theme_retro_terminal.py | 15 + TUI/{ => Themes}/themeselector.py | 15 + TUI/{ => Widgets}/OTP_generate.py | 36 +- TUI/{ => Widgets}/agentmoveoperations.py | 87 +- TUI/{ => Widgets}/multiagentselector.py | 31 +- TUI/{ => Widgets}/policyselector.py | 56 +- TUI/{ => Widgets}/policytreewidget.py | 15 + TUI/Widgets/prepPolicy.py | 870 ++++++++ TUI/{ => Widgets}/resultsdisplay.py | 52 +- TUI/moveagentworkflowscreen.py | 61 - TUI/otpworkflowscreen.py | 24 - flows/localApproval.py | 65 +- flows/prepPolicy.py | 46 +- models/execution.py | 8 +- services/agenthandler.py | 28 +- services/policyhandler.py | 17 +- utils/selector.py | 12 +- utils/setup.py | 6 +- 29 files changed, 3761 insertions(+), 561 deletions(-) rename TUI/{ => Screens}/allowlistselectionscreen.py (93%) create mode 100644 TUI/Screens/moveagentworkflowscreen.py rename TUI/{ => Screens}/otpactivityscreen.py (92%) create mode 100644 TUI/Screens/otprevokescreen.py create mode 100644 TUI/Screens/otpworkflowscreen.py create mode 100644 TUI/Screens/policyprepworkflowscreen.py rename TUI/{ => Screens}/policyselectorscreen.py (82%) rename TUI/{ => Screens}/quietagentworkflowscreen.py (91%) rename TUI/{ => Themes}/theme_amber_terminal.py (55%) rename TUI/{ => Themes}/theme_retro_terminal.py (50%) rename TUI/{ => Themes}/themeselector.py (68%) rename TUI/{ => Widgets}/OTP_generate.py (92%) rename TUI/{ => Widgets}/agentmoveoperations.py (90%) rename TUI/{ => Widgets}/multiagentselector.py (88%) rename TUI/{ => Widgets}/policyselector.py (92%) rename TUI/{ => Widgets}/policytreewidget.py (93%) create mode 100644 TUI/Widgets/prepPolicy.py rename TUI/{ => Widgets}/resultsdisplay.py (74%) delete mode 100644 TUI/moveagentworkflowscreen.py delete mode 100644 TUI/otpworkflowscreen.py diff --git a/Loxide.py b/Loxide.py index fae9cc8..7c3e974 100644 --- a/Loxide.py +++ b/Loxide.py @@ -30,7 +30,7 @@ from services.API import AirlockAPIWrapper from services.security import getAPI from TUI.TUI import run_Loxide from utils.configmanager import get_system_value -from utils.setup import get_base_directory, setup +from utils.setup import setup from utils.utils import irtang urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) @@ -40,7 +40,6 @@ def main(): irtang() # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored setup() - base_dir = get_base_directory() logger = logging.getLogger(__name__) try: diff --git a/TUI/allowlistselectionscreen.py b/TUI/Screens/allowlistselectionscreen.py similarity index 93% rename from TUI/allowlistselectionscreen.py rename to TUI/Screens/allowlistselectionscreen.py index 72426f2..cb04f2d 100644 --- a/TUI/allowlistselectionscreen.py +++ b/TUI/Screens/allowlistselectionscreen.py @@ -1,3 +1,17 @@ +# 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 . from __future__ import annotations import logging @@ -150,14 +164,11 @@ class AllowlistSelectionWidget(Static): # 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.styles.width = "100%" self.add_btn.disabled = True # Disabled until allowlist selected - yield self.back_btn yield self.add_btn async def on_mount(self) -> None: @@ -380,9 +391,9 @@ class AllowlistSelectionWidget(Static): if found_col: self.hash_column = found_col - preview_lines.append(f"✓ Found hash column: **{found_col}**\n") + preview_lines.append(f"✅ Found hash column: **{found_col}**\n") else: - preview_lines.append("⚠️ **No hash column found**\n") + 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": @@ -464,7 +475,7 @@ class AllowlistSelectionWidget(Static): 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')}'" + f"➕ Add to '{self.selected_allowlist.get('name', 'Unknown')}'" ) # Update preview with selection @@ -523,11 +534,6 @@ class AllowlistSelectionWidget(Static): 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() @@ -553,7 +559,7 @@ class AllowlistSelectionWidget(Static): try: # Disable button during operation self.add_btn.disabled = True - self.add_btn.label = "⏳ Adding hashes..." + self.add_btn.label = "Adding hashes..." # Call API to add hashes app_id = self.selected_allowlist.get("applicationid") @@ -564,10 +570,10 @@ class AllowlistSelectionWidget(Static): ) result = self.api.hash_add_to_allowlist(app_id, self.hashes_to_add) - + logger.debug(f"Hash adding api call: {result}") # Success notification self.app.notify( - f"✅ Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'", + f"Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'", title="Success", severity="information", timeout=5, @@ -575,7 +581,7 @@ class AllowlistSelectionWidget(Static): # Update preview to show success self.preview_area.text = ( - f"## ✅ SUCCESS\n\n" + 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" @@ -586,13 +592,13 @@ class AllowlistSelectionWidget(Static): ) # Change button to "Done" - self.add_btn.label = "✅ Done" + self.add_btn.label = "Done - Press q to return to main menu" self.add_btn.disabled = 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)}", + f"Failed to add hashes: {str(exc)}", title="Error", severity="error", timeout=10, @@ -600,7 +606,7 @@ class AllowlistSelectionWidget(Static): # Re-enable button self.add_btn.disabled = False - self.add_btn.label = "⟳ Retry Add to Allowlist" + self.add_btn.label = "Retry Add to Allowlist" class AllowlistSelectionScreen(Screen): @@ -609,9 +615,9 @@ class AllowlistSelectionScreen(Screen): """ BINDINGS = [ - Binding("b", "back", "Back"), + Binding("escape", "go_back", "Back"), + Binding("q", "main_menu", "Main Menu"), Binding("r", "refresh", "Refresh Allowlists"), - Binding("enter", "confirm", "Add to Allowlist"), ] def __init__( @@ -641,10 +647,15 @@ class AllowlistSelectionScreen(Screen): yield self.widget yield Footer() - async def action_back(self) -> None: + async def action_go_back(self) -> None: """Go back to previous screen.""" await self.app.pop_screen() + async def action_main_menu(self) -> None: + """Go back to main menu.""" + while len(self.app.screen_stack) > 2: + await self.app.pop_screen() + async def action_refresh(self) -> None: """Refresh the allowlists.""" if hasattr(self, "widget") and self.widget: diff --git a/TUI/Screens/moveagentworkflowscreen.py b/TUI/Screens/moveagentworkflowscreen.py new file mode 100644 index 0000000..16dbfba --- /dev/null +++ b/TUI/Screens/moveagentworkflowscreen.py @@ -0,0 +1,114 @@ +# 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 . +from typing import List, Optional + +from textual.app import ComposeResult +from textual.binding import Binding +from textual.css.query import NoMatches +from textual.screen import Screen + +from models.agent import Agent +from TUI.Widgets.agentmoveoperations import AgentMoveOperations +from TUI.Widgets.multiagentselector import MultiAgentSelector +from TUI.Widgets.resultsdisplay import ResultsDisplay + + +class MoveAgentWorkflowScreen(Screen): + """Screen that handles the agent movement workflow.""" + + BINDINGS = [ + Binding("escape", "go_back", "Back"), + Binding("q", "main_menu", "Main Menu"), + ] + + def __init__(self, all_agents: Optional[List[Agent]]): + super().__init__() + self.all_agents = all_agents + self.selected_agents = None + self.workflow_stage = "select_agents" # Track current stage + + def compose(self) -> ComposeResult: + """Start with the multi-agent selector.""" + yield MultiAgentSelector(self.all_agents) + + def action_go_back(self) -> None: + """Handle escape key to go back one step within the workflow.""" + if self.workflow_stage == "select_agents": + # At first stage, go back to main menu + self.app.pop_screen() + elif self.workflow_stage == "operations": + # Go back to agent selection + try: + ops_widget = self.query_one(AgentMoveOperations) + ops_widget.remove() + except NoMatches: + pass + self.mount(MultiAgentSelector(self.all_agents)) + self.workflow_stage = "select_agents" + elif self.workflow_stage == "results": + # Go back to operations + try: + results_widget = self.query_one(ResultsDisplay) + results_widget.remove() + except NoMatches: + pass + self.mount(AgentMoveOperations(self.selected_agents)) + self.workflow_stage = "operations" + + def action_main_menu(self) -> None: + """Handle q key to go back to main menu.""" + while len(self.app.screen_stack) > 2: + self.app.pop_screen() + + def on_multi_agent_selector_agents_selected( + self, message: MultiAgentSelector.AgentsSelected + ) -> None: + """Handle selected agents - switch to operations screen.""" + self.selected_agents = message.selected_agents + + # Remove the MultiAgentSelector + selector = self.query_one(MultiAgentSelector) + selector.remove() + + # Mount the AgentMoveOperations with the selected Agent objects + self.mount(AgentMoveOperations(self.selected_agents)) + self.workflow_stage = "operations" + + def on_agent_move_operations_operation_complete( + self, message: AgentMoveOperations.OperationComplete + ) -> None: + """Handle completion of move operation - transition to results screen.""" + # Format successful results + success_lines = [] + for agent, result in message.successful: + success_lines.append(f"✔ {agent.hostname}") + + # Format unsuccessful results + failure_lines = [] + for agent, error in message.unsuccessful: + failure_lines.append(f"❌ — {agent.hostname}: {error}") + + successful_text = "\n".join(success_lines) if success_lines else "(none)" + unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)" + + # Remove the operations widget + ops_widget = self.query_one(AgentMoveOperations) + ops_widget.remove() + + # Mount the results display + self.mount( + ResultsDisplay(message.operation, successful_text, unsuccessful_text) + ) + self.workflow_stage = "results" diff --git a/TUI/otpactivityscreen.py b/TUI/Screens/otpactivityscreen.py similarity index 92% rename from TUI/otpactivityscreen.py rename to TUI/Screens/otpactivityscreen.py index 5947cee..3b6714f 100644 --- a/TUI/otpactivityscreen.py +++ b/TUI/Screens/otpactivityscreen.py @@ -1,3 +1,17 @@ +# 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 . from __future__ import annotations from datetime import datetime @@ -11,7 +25,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 TUI.Screens.allowlistselectionscreen import AllowlistSelectionScreen from utils.configmanager import load_env logger = logging.getLogger(__name__) @@ -31,8 +45,7 @@ class OTPActivitiesWidget(Static): """ Reusable widget that contains the sessions table (left) and an Activity Preview (right). The right side shows an Activity Preview that takes ~75% vertical space, and a lower area - with Back and Continue buttons. The Continue button pushes ActivityDetailScreen with the - currently-loaded activities. + with Continue button. """ DEFAULT_CSS = """ @@ -85,15 +98,10 @@ class OTPActivitiesWidget(Static): with Vertical(id="activity_preview_container"): self.activities_table = DataTable(id="activity_preview_table") yield self.activities_table - # Buttons area at the bottom (Back, Continue) + # Button area at the bottom (Continue) with Horizontal(id="activity_buttons"): - # Back takes left side, Continue right side - self.back_btn = Button("Back", id="activity_back_btn") self.continue_btn = Button("Continue", id="activity_continue_btn") - # Stretch buttons nicely - self.back_btn.styles.width = "50%" - self.continue_btn.styles.width = "50%" - yield self.back_btn + self.continue_btn.styles.width = "100%" yield self.continue_btn async def on_mount(self) -> None: @@ -122,7 +130,7 @@ class OTPActivitiesWidget(Static): async def on_button_pressed(self, event) -> None: # type: ignore[override] """ - Handle Back / Continue buttons for the Activity Preview area. + Handle Continue button for the Activity Preview area. """ # Try to resolve the button object from the event btn = ( @@ -136,18 +144,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: logger.info("Continue pressed but no activities loaded.") - await self.post_message( - Static("No activities loaded to continue with.") + self.app.notify( + "No activities loaded to continue with.", severity="warning" ) return # Copy activities DataFrame to pass to new screen @@ -463,7 +465,7 @@ class OTPActivitiesWidget(Static): try: self._activities_df.to_csv(file_path, index=False) logger.info("Exported activities to %s", file_path) - await self.post_message(Static(f"✅ Exported activities to: {file_path}")) + await self.post_message(Static(f"Exported activities to: {file_path}")) except Exception as exc: logger.exception("Failed to export activities to %s: %s", file_path, exc) await self.post_message(Static("Failed to export activities; check logs.")) @@ -472,7 +474,7 @@ class OTPActivitiesWidget(Static): class ActivityDetailWidget(Static): """ Interactive widget for Activity Detail screen. - Shows the provided DataFrame in a DataTable and offers Export + Back buttons. + Shows the provided DataFrame in a DataTable and offers Export button. Now includes Select All/None and Add to Allowlist functionality. """ @@ -529,12 +531,10 @@ class ActivityDetailWidget(Static): # Original buttons at bottom with Horizontal(id="detail_buttons"): - self.detail_back_btn = Button("Back", id="detail_back_btn") self.add_allowlist_btn = Button( - "📋 Add Selected to Allowlist", id="add_allowlist_btn" + "Add Selected to Allowlist", id="add_allowlist_btn" ) yield self.add_allowlist_btn - yield self.detail_back_btn async def on_mount(self) -> None: await self._build_table(rebuild=True) @@ -547,12 +547,12 @@ class ActivityDetailWidget(Static): # Update button labels with count count = len(self.selected_row_ids) - total = len(self.activities_df) + len(self.activities_df) if has_selection: - self.add_allowlist_btn.label = f"📋 Add {count} Selected to Allowlist" + self.add_allowlist_btn.label = f"Add {count} Selected to Allowlist" else: - self.add_allowlist_btn.label = "📋 Add Selected to Allowlist" + 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.""" @@ -591,7 +591,7 @@ class ActivityDetailWidget(Static): 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 "☐" + checkbox = "☑️" if row_id in self.selected_row_ids else "☐" # Add row to table row_key = self.detail_table.add_row(checkbox, *vals) @@ -616,7 +616,7 @@ class ActivityDetailWidget(Static): 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.detail_table.update_cell(row_key, "select", "☑️") # Checked self._update_button_states() @@ -652,19 +652,13 @@ class ActivityDetailWidget(Static): logger.exception("Failed to sort by column %s: %s", column_key, exc) return - # ✅ Only refresh rows, not columns + # 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) - if btn is self.detail_back_btn or btn_id == "detail_back_btn": - while len(self.app.screen_stack) > 2: - self.app.pop_screen() - event.stop() - return - if btn is self.add_allowlist_btn or btn_id == "add_allowlist_btn": await self._open_allowlist_screen() return @@ -676,7 +670,7 @@ class ActivityDetailWidget(Static): # 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.detail_table.update_cell(row_key, "select", "☑️") self._update_button_states() logger.info(f"Selected all {len(self.selected_row_ids)} rows") @@ -688,7 +682,7 @@ class ActivityDetailWidget(Static): # 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.detail_table.update_cell(row_key, "select", "☑️") self._update_button_states() logger.info("Cleared all selections") @@ -730,14 +724,12 @@ class ActivityDetailWidget(Static): async def _export_detail_activities(self) -> None: if self.activities_df is None or self.activities_df.empty: logger.info("No activities to export.") - await self.mount( - Static("❌ No activities to export.", classes="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") + Static("No rows selected for export.", classes="notification") ) return try: @@ -750,7 +742,7 @@ class ActivityDetailWidget(Static): logger.info("Exported selected activities to %s", file_path) await self.mount( Static( - f"✅ Exported selected activities to: {filename}", + f"Exported selected activities to: {filename}", classes="notification", ) ) @@ -758,12 +750,12 @@ class ActivityDetailWidget(Static): logger.exception("Failed to export detail activities: %s", exc) await self.mount( Static( - "❌ Failed to export activities; check logs.", + "¢ Failed to export activities; check logs.", classes="notification", ) ) - # ✅ Helper methods + # Helper methods def get_selected_data(self) -> pd.DataFrame: """Return a DataFrame of the selected rows.""" if not self.selected_row_ids: @@ -795,7 +787,8 @@ class ActivityDetailScreen(Screen): """ BINDINGS = [ - Binding("b", "back", "Back"), + Binding("escape", "go_back", "Back"), + Binding("q", "main_menu", "Main Menu"), Binding("e", "export", "Export"), Binding("a", "select_all", "Select All"), Binding("n", "select_none", "Select None"), @@ -819,11 +812,16 @@ class ActivityDetailScreen(Screen): yield self.widget yield Footer() - async def action_back(self) -> None: + async def action_go_back(self) -> None: try: await self.app.pop_screen() except Exception: - logger.debug("ActivityDetailScreen.action_back pop_screen failed.") + logger.debug("ActivityDetailScreen.action_go_back pop_screen failed.") + + async def action_main_menu(self) -> None: + """Go back to main menu.""" + while len(self.app.screen_stack) > 2: + await self.app.pop_screen() async def action_export(self) -> None: # Delegate to widget export helper @@ -851,9 +849,10 @@ class OTPActivitiesScreen(Screen): """ BINDINGS = [ + Binding("escape", "go_back", "Back"), + Binding("q", "main_menu", "Main Menu"), Binding("r", "refresh_sessions", "Refresh Sessions"), Binding("e", "export_activities", "Export activities"), - Binding("q", "quit", "Quit"), ] def compose(self) -> ComposeResult: @@ -875,6 +874,15 @@ class OTPActivitiesScreen(Screen): else: await self.widget.load_sessions_from_api(api) + async def action_go_back(self) -> None: + """Go back one screen.""" + await self.app.pop_screen() + + async def action_main_menu(self) -> None: + """Go back to main menu.""" + while len(self.app.screen_stack) > 2: + await self.app.pop_screen() + # Simple actions bound to keys async def action_refresh_sessions(self) -> None: api = getattr(self.app, "api", None) @@ -884,10 +892,6 @@ class OTPActivitiesScreen(Screen): logger.info("Refreshing OTP sessions via API.") await self.widget.load_sessions_from_api(api) - async def action_quit(self) -> None: - # Pop the screen or exit app - await self.app.pop_screen() - # If you want an explicit method to fetch activities for a particular otpid from outside: async def fetch_activities_for_otpid(self, otpid, hostname=None) -> None: api = getattr(self.app, "api", None) diff --git a/TUI/Screens/otprevokescreen.py b/TUI/Screens/otprevokescreen.py new file mode 100644 index 0000000..5dc142d --- /dev/null +++ b/TUI/Screens/otprevokescreen.py @@ -0,0 +1,393 @@ +# 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 . + +from __future__ import annotations + +import logging +from typing import List, Optional + +import pandas as pd +from textual.app import ComposeResult +from textual.binding import Binding +from textual.containers import Horizontal, Vertical +from textual.message import Message +from textual.screen import Screen +from textual.widgets import Button, DataTable, Footer, Header, Static + +logger = logging.getLogger(__name__) + + +class OTPRevokeWidget(Static): + """ + Widget for managing OTP session revocation. + Displays active OTP sessions and allows selection for revocation. + """ + + class SessionsRevoked(Message): + """Message sent when sessions are revoked.""" + + def __init__(self, revoked_sessions: List[dict]): + super().__init__() + self.revoked_sessions = revoked_sessions + + DEFAULT_CSS = """ + OTPRevokeWidget { + height: 1fr; + } + #main_container { + width: 100%; + height: 100%; + layout: vertical; + } + #sessions_container { + height: 1fr; + border: none; + padding: 1; + } + #button_container { + height: auto; + padding: 1; + align: center middle; + } + #result_container { + height: auto; + max-height: 10; + border: solid #444444; + padding: 1; + margin: 1; + overflow-y: auto; + } + .panel-title { + text-style: bold; + margin: 0 0 1 0; + } + """ + + def compose(self) -> ComposeResult: + with Vertical(id="main_container"): + # Sessions table + yield Static("OTP Sessions", classes="panel-title") + with Vertical(id="sessions_container"): + self.sessions_table = DataTable(id="sessions_table") + self.sessions_table.styles.width = "100%" + self.sessions_table.styles.height = "1fr" + yield self.sessions_table + + # Action buttons + with Horizontal(id="button_container"): + self.refresh_button = Button("🔄 Refresh", id="refresh_btn") + self.refresh_button.styles.width = "15%" + self.refresh_button.styles.margin = (1, 1, 1, 1) + yield self.refresh_button + + self.select_all_button = Button("☑️ Select All", id="select_all_btn") + self.select_all_button.styles.width = "15%" + self.select_all_button.styles.margin = (1, 1, 1, 1) + yield self.select_all_button + + self.select_none_button = Button( + "❌ Clear Selection", id="select_none_btn" + ) + self.select_none_button.styles.width = "20%" + self.select_none_button.styles.margin = (1, 1, 1, 1) + yield self.select_none_button + + self.revoke_button = Button( + "🛑 Revoke Selected", id="revoke_btn", variant="error" + ) + self.revoke_button.styles.width = "20%" + self.revoke_button.styles.margin = (1, 1, 1, 1) + yield self.revoke_button + + # Results display + with Vertical(id="result_container"): + yield Static("Revocation Results", classes="panel-title") + self.results_display = Static("No actions performed yet.") + yield self.results_display + + async def on_mount(self) -> None: + """Initialize the widget when mounted.""" + # Configure sessions table + self.sessions_table.clear() + self.sessions_table.add_columns( + "☐", "OTP ID", "Hostname", "Status", "Purpose", "Granted" + ) + + # Enable row selection with checkbox column + self.sessions_table.cursor_type = "row" + try: + self.sessions_table.zebra_stripes = True + except Exception: + pass + + # Initialize state + self._sessions_df: Optional[pd.DataFrame] = None + self._filtered_df: Optional[pd.DataFrame] = None + self._selected_otpids: set = set() + + async def load_sessions_from_api(self, api) -> None: + """Load active OTP sessions from the API.""" + try: + # Fetch only active sessions + active_df = api.otp_find_active() + + # Ensure we have a DataFrame + if not isinstance(active_df, pd.DataFrame): + active_df = pd.DataFrame(active_df) + + # Add status column + active_df["status"] = "active" + + # Sort by otpid if column exists + if "otpid" in active_df.columns and not active_df.empty: + active_df = active_df.sort_values(by="otpid", ascending=False) + + # Store the full dataframe + self._sessions_df = active_df + self._filtered_df = active_df.copy() + + # Display in table + await self._refresh_table() + + # Update status + active_count = len(active_df) + + status_msg = f"Loaded {active_count} active sessions" + logger.info(status_msg) + self.results_display.update(status_msg) + + except Exception as e: + logger.exception(f"Failed to load OTP sessions: {e}") + self.results_display.update(f"Error loading sessions: {str(e)}") + + async def _refresh_table(self) -> None: + """Refresh the table display with current filtered data.""" + if self._filtered_df is None or self._filtered_df.empty: + self.sessions_table.clear() + return + + # Ensure expected columns exist + expected_cols = ["otpid", "hostname", "status", "purpose", "granted"] + for col in expected_cols: + if col not in self._filtered_df.columns: + self._filtered_df[col] = "" + + # Clear and repopulate table + self.sessions_table.clear(columns=False) + + for _, row in self._filtered_df.iterrows(): + otpid = str(row.get("otpid", "")) + # Check if this row is selected + checkbox = "☑️" if otpid in self._selected_otpids else "☐" + + self.sessions_table.add_row( + checkbox, + str(otpid), + str(row.get("hostname", "")), + str(row.get("status", "")), + str(row.get("purpose", "")), + str(row.get("granted", "")), + ) + + async def on_button_pressed(self, event) -> None: + """Handle button presses.""" + btn = event.button + + if btn.id == "refresh_btn": + # Refresh sessions + api = getattr(self.app, "api", None) + if api: + await self.load_sessions_from_api(api) + + elif btn.id == "select_all_btn": + # Select all visible rows + if self._filtered_df is not None: + self._selected_otpids = set(str(x) for x in self._filtered_df["otpid"]) + await self._refresh_table() + + elif btn.id == "select_none_btn": + # Clear selection + self._selected_otpids.clear() + await self._refresh_table() + + elif btn.id == "revoke_btn": + # Revoke selected sessions + await self._revoke_selected() + + async def on_data_table_row_selected(self, event) -> None: + """Handle row selection in the table.""" + if event.data_table != self.sessions_table: + return + + try: + # Get the row index from the cursor row + row_index = self.sessions_table.cursor_row + + if self._filtered_df is not None and row_index < len(self._filtered_df): + # Get the OTP ID for this row + otpid = str(self._filtered_df.iloc[row_index]["otpid"]) + + # Toggle selection + if otpid in self._selected_otpids: + self._selected_otpids.remove(otpid) + else: + self._selected_otpids.add(otpid) + + # Refresh table to update checkbox + await self._refresh_table() + + # Restore cursor position + self.sessions_table.move_cursor(row=row_index) + + except Exception as e: + logger.exception(f"Error handling row selection: {e}") + + async def _revoke_selected(self) -> None: + """Revoke the selected OTP sessions.""" + if not self._selected_otpids: + self.results_display.update("❌ No sessions selected for revocation") + return + + api = getattr(self.app, "api", None) + if not api: + self.results_display.update("❌ API not available") + return + + # Collect results + results = [] + success_count = 0 + failure_count = 0 + + for otpid in self._selected_otpids: + try: + # Get hostname for this session + hostname = "Unknown" + if self._sessions_df is not None: + # Convert otpid to same type as in DataFrame for comparison + otpid_compare = otpid + if len(self._sessions_df) > 0: + first_otpid = self._sessions_df["otpid"].iloc[0] + if isinstance(first_otpid, int): + otpid_compare = int(otpid) + + match = self._sessions_df[ + self._sessions_df["otpid"] == otpid_compare + ] + if not match.empty: + hostname = match.iloc[0].get("hostname", "Unknown") + + # Revoke the session + result = api.otp_revoke(otpid) + + if result and result.get("status") != "error": + success_count += 1 + results.append(f"Revoked OTP {otpid} for {hostname}") + logger.info(f"Revoked OTP {otpid} for {hostname}: {result}") + else: + failure_count += 1 + error_msg = ( + result.get("message", "Unknown error") + if result + else "No response" + ) + results.append( + f"❌ Failed to revoke OTP {otpid} for {hostname}: {error_msg}" + ) + logger.error(f"Failed to revoke OTP {otpid}: {error_msg}") + + except Exception as e: + failure_count += 1 + results.append(f"❌ Error revoking OTP {otpid}: {str(e)}") + logger.exception(f"Exception revoking OTP {otpid}: {e}") + + # Update results display + summary = ( + f"Revocation complete: {success_count} succeeded, {failure_count} failed\n" + ) + details = "\n".join(results[-5:]) # Show last 5 results + if len(results) > 5: + details = f"... (showing last 5 of {len(results)} results)\n" + details + + self.results_display.update(summary + details) + + # Clear selection and refresh + self._selected_otpids.clear() + await self.load_sessions_from_api(api) + + # Post message about revoked sessions + if success_count > 0: + self.post_message(self.SessionsRevoked(results)) + + +class OTPRevokeScreen(Screen): + """ + Main screen for OTP session revocation workflow. + This replaces the otp_revoke function from otp.py. + """ + + BINDINGS = [ + Binding("escape", "go_back", "Back"), + Binding("q", "main_menu", "Main Menu"), + Binding("r", "refresh", "Refresh"), + Binding("a", "select_all", "Select All"), + Binding("n", "select_none", "Clear Selection"), + Binding("d", "revoke", "Revoke Selected"), + ] + + def compose(self) -> ComposeResult: + yield Header(show_clock=True) + self.widget = OTPRevokeWidget() + yield self.widget + yield Footer() + + async def on_mount(self) -> None: + """Load sessions when screen mounts.""" + api = getattr(self.app, "api", None) + if api: + await self.widget.load_sessions_from_api(api) + else: + logger.warning("OTPRevokeScreen mounted but no self.app.api found.") + + async def action_refresh(self) -> None: + """Refresh the sessions list.""" + api = getattr(self.app, "api", None) + if api: + await self.widget.load_sessions_from_api(api) + + async def action_select_all(self) -> None: + """Select all visible sessions.""" + if self.widget._filtered_df is not None: + self.widget._selected_otpids = set( + str(x) for x in self.widget._filtered_df["otpid"] + ) + await self.widget._refresh_table() + + async def action_select_none(self) -> None: + """Clear all selections.""" + self.widget._selected_otpids.clear() + await self.widget._refresh_table() + + async def action_revoke(self) -> None: + """Revoke selected sessions.""" + await self.widget._revoke_selected() + + async def action_go_back(self) -> None: + """Go back to previous screen.""" + await self.app.pop_screen() + + async def action_main_menu(self) -> None: + """Go back to main menu.""" + while len(self.app.screen_stack) > 2: + await self.app.pop_screen() diff --git a/TUI/Screens/otpworkflowscreen.py b/TUI/Screens/otpworkflowscreen.py new file mode 100644 index 0000000..124569c --- /dev/null +++ b/TUI/Screens/otpworkflowscreen.py @@ -0,0 +1,52 @@ +# 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 . + +from typing import List, Optional + +from textual.app import ComposeResult +from textual.binding import Binding +from textual.screen import Screen + +from models.agent import Agent +from TUI.Widgets.OTP_generate import OTPGenerator + + +class OTPWorkflowScreen(Screen): + """Screen that handles the OTP generation workflow without agent selection.""" + + BINDINGS = [ + Binding("escape", "go_back", "Back"), + Binding("q", "main_menu", "Main Menu"), + ] + + def __init__(self, selected_agents: Optional[List[Agent]]): + super().__init__() + self.selected_agents = selected_agents + + def compose(self) -> ComposeResult: + """Directly show the OTP generator for the selected agents.""" + yield OTPGenerator(self.selected_agents) + + def action_go_back(self) -> None: + """Handle escape key to go back one screen.""" + self.app.pop_screen() + + def action_main_menu(self) -> None: + """Handle q key to go back to main menu.""" + while len(self.app.screen_stack) > 2: + self.app.pop_screen() + + def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None: + """Handle OTP generation request - pass it up to the app level if needed.""" diff --git a/TUI/Screens/policyprepworkflowscreen.py b/TUI/Screens/policyprepworkflowscreen.py new file mode 100644 index 0000000..7035dc2 --- /dev/null +++ b/TUI/Screens/policyprepworkflowscreen.py @@ -0,0 +1,1822 @@ +# 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 +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: 30%; + margin: 0 1 0 1; + } + + #content_area { + height: 1fr; + } + + 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("r", "refresh", "Refresh"), + 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), + ] + + 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 = "\\\\" # Default path split for Windows + + # 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 + + 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_source_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() + 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 = (1, 2) + checklist_container.styles.padding = 1 + + # Mount the container to the checklist area FIRST + checklist.mount(checklist_container) + + # NOW mount children to the container + checklist_title = Static("Preparation Checklist") + checklist_title.styles.text_style = "bold" + checklist_container.mount(checklist_title) + + # Step 1: Source Policies + step1_status = "✔️" if self.destination_policy else "✖️" + step1_text = f"{step1_status} Source Policies: " + if self.source_policies: + step1_text += ", ".join([p.name for p in self.source_policies[:3]]) + if len(self.source_policies) > 3: + step1_text += f" (+{len(self.source_policies)-3} more)" + else: + step1_text += "Not selected" + step1 = Static(step1_text) + if self.source_policies: + step1.styles.color = "green" + else: + step1.styles.text_style = "dim" + checklist_container.mount(step1) + + # Step 2: Destination Policy + step2_status = "✔️" if self.destination_policy else "✖️" + step2_text = f"{step2_status} Destination Policy: " + step2_text += ( + self.destination_policy.name if self.destination_policy else "Not selected" + ) + step2 = Static(step2_text) + if self.destination_policy: + step2.styles.color = "green" + else: + step2.styles.text_style = "dim" + checklist_container.mount(step2) + + # Step 3: Destination Allowlist + step3_status = "✔️" if self.destination_policy else "✖️" + step3_text = f"{step3_status} Allowlist: " + step3_text += ( + self.destination_allowlist.name + if self.destination_allowlist + else "Not selected" + ) + step3 = Static(step3_text) + if self.destination_allowlist: + step3.styles.color = "green" + else: + step3.styles.text_style = "dim" + checklist_container.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 self.destination_policy else "✖️" + step4_text = f"{step4_status} Data Fetched: " + 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_text += f"{total} executions" + else: + step4_text += "Not fetched" + step4 = Static(step4_text) + if data_fetched: + step4.styles.color = "green" + else: + step4.styles.text_style = "dim" + checklist_container.mount(step4) + + # Step 5: First Review Complete + first_review_path = os.path.join(self.working_dir, "Approved") + 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) + else: + first_review_done = False + + step5_status = "✔️" if self.destination_policy else "✖️" + step5_text = f"{step5_status} First Review: " + step5_text += "Complete" if first_review_done else "Pending" + step5 = Static(step5_text) + if first_review_done: + step5.styles.color = "green" + else: + step5.styles.text_style = "dim" + checklist_container.mount(step5) + + # Step 6: Paths Generated + paths_generated = self.primary_paths_df is not None + step6_status = "✔️" if self.destination_policy else "✖️" + step6_text = f"{step6_status} Paths Generated: " + if paths_generated: + step6_text += f"{len(self.primary_paths_df)} primary paths" + else: + step6_text += "Not generated" + step6 = Static(step6_text) + if paths_generated: + step6.styles.color = "green" + else: + step6.styles.text_style = "dim" + checklist_container.mount(step6) + + 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") + + # Add rows + for policy in self.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 single-select + policy_selector = PolicySelector(self.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 + + # 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() + + 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:" + ) + instruction.styles.margin = (0, 1, 1, 1) + content.mount(instruction) + + # Days input + days_container = Horizontal() + days_container.styles.margin = (1, 1) + content.mount(days_container) + + days_label = Static("History Days (1-150): ") + days_input = Input(value="150", id="history_days_input", type="integer") + days_input.styles.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) + + def _fetch_execution_data(self, history_days: int) -> None: + """Fetch and sort execution data.""" + self.workflow_stage = "fetching" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + status = Static("Fetching execution history...\nThis may take a few moments...") + status.styles.margin = (2, 1) + content.mount(status) + + # Perform fetch in background + self.call_later(lambda: self._perform_fetch(history_days)) + + def _perform_fetch(self, history_days: int) -> None: + """Perform the actual data fetching.""" + try: + # Fetch execution history + policy_executions = ExecutionHistoryRecord.from_policies( + self.api, + self.source_policies, + type_=[1, 2, 6, 7], + history_days=history_days, + ) + + # Enrich with hash data + enriched_executions = ExecutionHistoryRecord.enrich_with_hashes( + self.api, policy_executions + ) + + # Categorize by hash decision + categorized_executions = ( + ExecutionHistoryRecord.categorize_executions_by_hash_decision( + enriched_executions + ) + ) + + # Sort by decision + approved, unapproved, needs_review, unknown = ( + ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions) + ) + + # 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 + self._save_fetched_data() + + # Show 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.""" + self.workflow_stage = "first_review" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + # 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) + + # 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) + + # Show approved table by default + self._show_review_table("approved") + + def _show_review_table(self, table_type: str) -> None: + """Show an editable DataTable for reviewing executions.""" + content = self.query_one("#content_area", Vertical) + + # Remove existing table if any + existing_table = content.query("DataTable") + for table in existing_table: + table.remove() + existing_controls = content.query("#review_controls") + for control in existing_controls: + control.remove() + + # Determine which dataframe 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:" + + 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) + return + + # Instructions + instruction = Static(title) + instruction.styles.margin = (1, 1) + instruction.styles.text_style = "bold" + content.mount(instruction) + + # Help text + help_text = Static( + "Use arrows to navigate, SPACE to select/deselect rows, 'd' to delete selected rows\n" + "Selected rows will be highlighted and removed from the approved list" + ) + 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 = "50vh" + review_table.cursor_type = "row" + review_table.zebra_stripes = True + + # Add columns - checkbox first, then important fields + important_cols = [ + "filename", + "publisher", + "sha256", + "filepath", + "hostname", + "datetime", + ] + 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)) + + content.mount(review_table) + + # Control buttons + control_container = Horizontal(id="review_controls") + control_container.styles.margin = (1, 1) + content.mount(control_container) + + select_all_btn = Button("Select All", id="select_all_rows") + select_all_btn.styles.margin = (0, 1, 0, 0) + + select_none_btn = Button("Clear Selection", id="select_none_rows") + select_none_btn.styles.margin = (0, 1, 0, 0) + + delete_btn = Button( + "Delete Selected", id="delete_selected_rows", variant="error" + ) + delete_btn.styles.margin = (0, 1, 0, 0) + + row_count = Static(f"Total rows: {len(df)}") + row_count.styles.margin = (0, 1, 0, 2) + + control_container.mount(select_all_btn) + control_container.mount(select_none_btn) + control_container.mount(delete_btn) + 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, 0, 1) + continue_container.styles.dock = "bottom" + + # 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 + else: + df = self.needs_review_df + + if df is None: + return + + # Get indices to keep (not in selected rows) + indices_to_delete = [int(idx) for idx in self.selected_rows] + 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 + + # Clear selection + self.selected_rows.clear() + + # Refresh the table + self._show_review_table(self.current_review_type) + + self.app.notify( + f"Deleted {len(indices_to_delete)} rows", severity="information" + ) + + def _build_paths_and_publishers(self) -> None: + """Build path exclusions and publisher lists.""" + self.workflow_stage = "build_paths" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + status = Static("Building path exclusions and publisher lists...") + status.styles.margin = (2, 1) + content.mount(status) + + 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 + remaining = all_approved[ + ~all_approved["sha256"].isin(self.primary_paths_df["sha256"]) + ] + self.secondary_paths_df = self._calculate_paths( + remaining, path_exclusion_const - 1 + ) + + # Remaining hashes + self.remaining_hashes_df = remaining[ + ~remaining["sha256"].isin(self.secondary_paths_df["sha256"]) + ] + + # 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 _calculate_paths(self, df: pd.DataFrame, depth: int) -> pd.DataFrame: + """Calculate path exclusions at specified depth.""" + # Simplified path calculation - in production this would be more complex + if df.empty or "filepath" not in df.columns: + return pd.DataFrame() + + paths = [] + for filepath in df["filepath"].unique(): + if pd.isna(filepath): + continue + parts = filepath.split(self.path_split) + if len(parts) > depth: + truncated = self.path_split.join(parts[:depth]) + paths.append(truncated) + + # Create dataframe with unique paths + if paths: + path_df = pd.DataFrame({"longestcfp": list(set(paths))}) + # Add mock sha256 for compatibility + path_df["sha256"] = path_df.index.astype(str) + return path_df + return pd.DataFrame() + + 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.""" + 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 + ) + + 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) + + # 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) + + # Show paths table by default + self._show_path_review_table("paths") + + def _show_path_review_table(self, table_type: str) -> None: + """Show an editable DataTable for reviewing paths/publishers.""" + content = self.query_one("#content_area", Vertical) + + # Remove existing table if any + existing_table = content.query("DataTable") + for table in existing_table: + table.remove() + existing_controls = content.query("#path_review_controls") + for control in existing_controls: + control.remove() + + # 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) + else: + df = pd.DataFrame() + + table_id = "paths_review_table" + title = "Path Exclusions - Select paths to REMOVE:" + columns = ["longestcfp", "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 [] + ) + + 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 + instruction = Static(title) + instruction.styles.margin = (1, 1) + instruction.styles.text_style = "bold" + content.mount(instruction) + + # Help text (different for remaining hashes) + if table_type != "remaining": + help_text = Static( + "Use arrows to navigate, SPACE to select/deselect rows, 'd' to delete selected rows\n" + "Selected items will be removed from the final approval list" + ) + 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 = "40vh" + 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 = [str(row.get(col, "")) for col in available_cols] + review_table.add_row(checkbox, *row_data, key=str(idx)) + + content.mount(review_table) + + # Control buttons (not for remaining hashes view) + if table_type != "remaining": + control_container = Horizontal(id="path_review_controls") + control_container.styles.margin = (1, 1) + content.mount(control_container) + + select_all_btn = Button("Select All", id="select_all_path_rows") + select_all_btn.styles.margin = (0, 1, 0, 0) + + select_none_btn = Button("Clear Selection", id="select_none_path_rows") + select_none_btn.styles.margin = (0, 1, 0, 0) + + delete_btn = Button( + "Delete Selected", id="delete_selected_path_rows", variant="error" + ) + delete_btn.styles.margin = (0, 1, 0, 0) + + row_count = Static(f"Total items: {len(df)}") + row_count.styles.margin = (0, 1, 0, 2) + + control_container.mount(select_all_btn) + control_container.mount(select_none_btn) + control_container.mount(delete_btn) + control_container.mount(row_count) + + # Continue button (always at bottom) + if not content.query("#path_continue_container"): + continue_container = Horizontal(id="path_continue_container") + continue_container.styles.margin = (2, 1, 0, 1) + continue_container.styles.dock = "bottom" + + 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) + content.mount(continue_container) + + # 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] + + # 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" + ) + + # Clear selection + self.selected_path_rows.clear() + + # Refresh the table + self._show_path_review_table(self.current_path_review_type) + + 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.""" + self.workflow_stage = "test" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + summary = Static( + "Test Mode - Preview Changes\n\n" "The following changes will be applied:\n" + ) + summary.styles.margin = (1, 1) + content.mount(summary) + + # Show what would be changed + changes_text = "" + + if self.destination_policy: + changes_text += f"Policy: {self.destination_policy.name}\n" + + if self.primary_paths_df is not None and not self.primary_paths_df.empty: + changes_text += f"Add {len(self.primary_paths_df)} path exclusions\n" + + if self.publishers_df is not None and not self.publishers_df.empty: + changes_text += f"Add {len(self.publishers_df)} approved publishers\n" + + if self.destination_allowlist: + changes_text += f"\nAllowlist: {self.destination_allowlist.name}\n" + + if self.approved_df is not None and not self.approved_df.empty: + changes_text += f"Add {len(self.approved_df)} approved hashes\n" + + changes = Static(changes_text) + changes.styles.margin = (0, 2) + changes.styles.border = ("round", "cyan") + changes.styles.padding = 1 + content.mount(changes) + + # Warning + warning = Static( + " ⚠️Warning: These changes cannot be easily undone.⚠️\n" + "Please review carefully before proceeding." + ) + warning.styles.margin = (1, 1) + warning.styles.color = "yellow" + content.mount(warning) + + # Buttons + button_container = Horizontal() + button_container.styles.margin = (2, 1) + content.mount(button_container) + + liftoff_btn = Button("Liftoff - Apply Changes", id="liftoff", variant="success") + + button_container.mount(liftoff_btn) + + 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 = [] + + # Apply path exclusions to policy + if self.destination_policy and self.primary_paths_df is not None: + # This would call the actual API methods + results.append("Applied path exclusions to policy") + + # Apply publishers to policy + if self.destination_policy and self.publishers_df is not None: + # This would call the actual API methods + results.append("Applied approved publishers to policy") + + # Apply hashes to allowlist + if self.destination_allowlist and self.approved_df is not None: + # This would call the actual API methods + results.append("Applied approved hashes to allowlist") + + self._show_completion(results) + + except Exception as e: + logger.error(f"Failed to apply changes: {e}", exc_info=True) + self.app.notify(f"Failed to apply changes: {str(e)}", severity="error") + self._show_test_screen() + + def _show_completion(self, results: List[str]) -> None: + """Show completion screen.""" + self.workflow_stage = "complete" + content = self.query_one("#content_area", Vertical) + content.remove_children() + + summary = Static( + "Policy Preparation Complete!\n\n" + "The following changes have been applied:" + ) + summary.styles.margin = (1, 1) + content.mount(summary) + + for result in results: + result_widget = Static(f" {result}") + result_widget.styles.margin = (0, 2) + content.mount(result_widget) + + # Final message + final = Static( + f"\nPolicy '{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(): + # Determine if this row should be checked + checkbox = "☑️" if str(row_key) in selected_keys else "☐" + + # Update the checkbox cell (first column, index 0) + try: + table.update_cell_at((row_index, 0), checkbox) + except Exception as e: + logger.debug(f"Could not update cell at row {row_index}: {e}") + + row_index += 1 + + except Exception as e: + logger.debug(f"Error refreshing table {table_id}: {e}") + + def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None: + """Handle row selection in data tables.""" + table = event.data_table + + # Handle source policy selection + if table.id == "source_policy_table": + # Use cursor_row for reliable row index + row_index = table.cursor_row + # Get the row key from the table + row_keys = list(table.rows.keys()) + if row_index < len(row_keys): + row_key = str(row_keys[row_index]) + 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) + # Refresh checkbox display + 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, + ) + + # Handle allowlist selection + elif 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() + + # Handle review table selections (toggle selection) + elif table.id in ["approved_review_table", "needs_review_table"]: + # Use cursor_row for reliable row index + row_index = table.cursor_row + # Get the row key from the table + row_keys = list(table.rows.keys()) + if row_index < len(row_keys): + row_key = str(row_keys[row_index]) + if row_key in self.selected_rows: + self.selected_rows.remove(row_key) + else: + self.selected_rows.add(row_key) + # Refresh checkbox display + self._refresh_table_checkboxes(table.id, self.selected_rows) + self.app.notify(f"Selected {len(self.selected_rows)} rows", timeout=1) + + # Handle path/publisher review selections + elif table.id in ["paths_review_table", "publishers_review_table"]: + # Use cursor_row for reliable row index + row_index = table.cursor_row + # Get the row key from the table + row_keys = list(table.rows.keys()) + if row_index < len(row_keys): + row_key = str(row_keys[row_index]) + if row_key in self.selected_path_rows: + self.selected_path_rows.remove(row_key) + else: + self.selected_path_rows.add(row_key) + # Refresh checkbox display + self._refresh_table_checkboxes(table.id, self.selected_path_rows) + self.app.notify( + f"Selected {len(self.selected_path_rows)} items", timeout=1 + ) + + def on_button_pressed(self, event: Button.Pressed) -> None: + """Handle button presses.""" + button_id = event.button.id + + # Source policy selection buttons + if 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]}" + ) + 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 <= 150: + self.history_days = history_days + self._fetch_execution_data(history_days) + else: + self.app.notify( + "Please enter a value between 1 and 150", 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": + # 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() + self._build_paths_and_publishers() + + elif button_id == "build_preflight": + # 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 == "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: + self.selected_rows = set(str(i) for i in range(len(df))) + # 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" + 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" + self.selected_path_rows = set( + str(i) for i in range(len(self.publishers_df)) + ) + + # 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 _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_refresh(self) -> None: + """Refresh the current view.""" + self._update_checklist() + + 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 (handled by on_data_table_row_selected).""" + pass # This is handled directly in the DataTable event diff --git a/TUI/policyselectorscreen.py b/TUI/Screens/policyselectorscreen.py similarity index 82% rename from TUI/policyselectorscreen.py rename to TUI/Screens/policyselectorscreen.py index 07ec2ee..7beb704 100644 --- a/TUI/policyselectorscreen.py +++ b/TUI/Screens/policyselectorscreen.py @@ -23,9 +23,11 @@ the policy selection workflow. import logging from textual.app import ComposeResult +from textual.binding import Binding from textual.screen import Screen +from textual.widgets import Footer, Header -from TUI.policyselector import PolicySelector +from TUI.Widgets.policyselector import PolicySelector logger = logging.getLogger(__name__) @@ -42,6 +44,11 @@ class PolicySelectorScreen(Screen): agent_move_operations: Reference to the parent AgentMoveOperations widget. """ + BINDINGS = [ + Binding("escape", "go_back", "Back"), + Binding("q", "main_menu", "Main Menu"), + ] + CSS = """ Screen { layout: vertical; @@ -68,7 +75,18 @@ class PolicySelectorScreen(Screen): def compose(self) -> ComposeResult: """Create the PolicySelector widget.""" + yield Header(show_clock=True) yield PolicySelector(self.policies) + yield Footer() + + def action_go_back(self) -> None: + """Handle escape key to go back one screen.""" + self.app.pop_screen() + + def action_main_menu(self) -> None: + """Handle q key to go back to main menu.""" + while len(self.app.screen_stack) > 2: + self.app.pop_screen() def on_policy_selector_policy_selected( self, message: PolicySelector.PolicySelected diff --git a/TUI/quietagentworkflowscreen.py b/TUI/Screens/quietagentworkflowscreen.py similarity index 91% rename from TUI/quietagentworkflowscreen.py rename to TUI/Screens/quietagentworkflowscreen.py index 552f214..f0d5bbb 100644 --- a/TUI/quietagentworkflowscreen.py +++ b/TUI/Screens/quietagentworkflowscreen.py @@ -39,7 +39,7 @@ 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 TUI.Widgets.policyselector import PolicySelector from utils.configmanager import load_env logger = logging.getLogger(__name__) @@ -69,6 +69,7 @@ class QuietAgentWorkflowScreen(Screen): BINDINGS = [ ("escape", "go_back", "Back"), + ("q", "main_menu", "Main Menu"), ] workflow_stage = reactive("select_policy") # Tracks current workflow stage @@ -96,10 +97,10 @@ class QuietAgentWorkflowScreen(Screen): 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="⚙") + yield Header(show_clock=True, icon="⚙️") # Title area - title = Static("🔒 Quiet Agent Workflow", id="workflow_title") + title = Static("Quiet Agent Workflow", id="workflow_title") title.styles.margin = (0, 0, 0, 1) yield title @@ -130,13 +131,13 @@ class QuietAgentWorkflowScreen(Screen): stage_messages = { "select_policy": "Step 1: Select Policy to Analyze", "select_quiet_days": "Step 2: Select Quiet Time Period", - "analyzing": "📊 Analyzing agent activity...", + "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", + "executing": "Executing agent migrations...", + "complete": "Migration Complete", } status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage")) @@ -208,11 +209,6 @@ class QuietAgentWorkflowScreen(Screen): 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 @@ -226,14 +222,6 @@ class QuietAgentWorkflowScreen(Screen): 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 @@ -278,16 +266,16 @@ class QuietAgentWorkflowScreen(Screen): # Show analyzing message with detailed steps analyzing_msg = Static( - f"📊 Analyzing Agent Activity\n" + 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"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", ) @@ -405,8 +393,8 @@ class QuietAgentWorkflowScreen(Screen): 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() + self.enforce_ready_df = agents[agents["enforce_ready"]].copy() + self.non_enforce_ready_df = agents[not agents["enforce_ready"]].copy() logger.info( f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, " @@ -435,13 +423,13 @@ class QuietAgentWorkflowScreen(Screen): # Build updated message updated_text = ( - f"📊 Analyzing Agent Activity\n" + 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"{status_text}\n\n" f"Please wait - this operation cannot be cancelled." ) @@ -469,9 +457,9 @@ class QuietAgentWorkflowScreen(Screen): 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"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", @@ -500,11 +488,11 @@ class QuietAgentWorkflowScreen(Screen): 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 = 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 = Button("Start Over", id="start_over_btn") start_over_btn.styles.margin = (0, 0, 1, 0) button_container.mount(start_over_btn) @@ -520,7 +508,7 @@ class QuietAgentWorkflowScreen(Screen): enforce_col.styles.margin = (1, 1, 0, 0) tables_container.mount(enforce_col) - enforce_label = Static("✅ Enforce Ready Agents") + enforce_label = Static("Enforce Ready Agents") enforce_label.styles.margin = (0, 0, 1, 0) enforce_col.mount(enforce_label) @@ -548,7 +536,7 @@ class QuietAgentWorkflowScreen(Screen): 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 = Static("Non-Enforce Ready Agents") non_enforce_label.styles.margin = (0, 0, 1, 0) non_enforce_col.mount(non_enforce_label) @@ -589,7 +577,7 @@ class QuietAgentWorkflowScreen(Screen): content.mount(policy_selector) # Skip button - skip_btn = Button("⭕️ Skip - No Migration", id="skip_enforce_target_btn") + 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) @@ -614,7 +602,7 @@ class QuietAgentWorkflowScreen(Screen): content.mount(policy_selector) # Skip button - skip_btn = Button("⭕️ Skip - No Migration", id="skip_non_enforce_target_btn") + 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) @@ -627,29 +615,29 @@ class QuietAgentWorkflowScreen(Screen): # Build confirmation message confirmation_lines = [ - "🔐 Migration Summary\n", + "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" + f"\nEnforce 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" + f"\nNon-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_lines.append("\nNo migrations will be performed.") confirmation = Static("".join(confirmation_lines), id="migration_confirmation") confirmation.styles.margin = (1, 1, 2, 1) @@ -662,11 +650,11 @@ class QuietAgentWorkflowScreen(Screen): 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 = 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") + cancel_btn = Button("Cancel", id="cancel_migration_btn") button_container.mount(cancel_btn) def _execute_migration(self) -> None: @@ -677,7 +665,7 @@ class QuietAgentWorkflowScreen(Screen): # Show executing message executing_msg = Static( - "⏳ Executing agent migrations...\nPlease wait...", + "Executing agent migrations...\nPlease wait...", id="executing_message", ) executing_msg.styles.margin = (2, 1) @@ -696,7 +684,7 @@ class QuietAgentWorkflowScreen(Screen): if self.enforce_ready_target_policy: for idx, row in self.enforce_ready_df.iterrows(): try: - result = self.api.agent_move( + self.api.agent_move( row["agentid"], self.enforce_ready_target_policy.groupid ) successful_migrations.append( @@ -713,7 +701,7 @@ class QuietAgentWorkflowScreen(Screen): if self.non_enforce_ready_target_policy: for idx, row in self.non_enforce_ready_df.iterrows(): try: - result = self.api.agent_move( + self.api.agent_move( row["agentid"], self.non_enforce_ready_target_policy.groupid ) successful_migrations.append( @@ -749,7 +737,7 @@ class QuietAgentWorkflowScreen(Screen): ) results = Static( - f"✅ Migration Complete\n\n" + f"Migration Complete\n\n" f"Total Agents Migrated: {len(successful)}\n" f"Failed Migrations: {len(failed)}\n" f"Success Rate: {success_rate:.1f}%", @@ -764,7 +752,7 @@ class QuietAgentWorkflowScreen(Screen): success_container.styles.margin = (0, 1) content.mount(success_container) - success_label = Static("✅ Successful Migrations") + success_label = Static("Successful Migrations") success_label.styles.margin = (0, 0, 1, 0) success_container.mount(success_label) @@ -785,7 +773,7 @@ class QuietAgentWorkflowScreen(Screen): failed_container.styles.margin = (2, 1, 0, 1) content.mount(failed_container) - failed_label = Static("❌ Failed Migrations") + failed_label = Static("Failed Migrations") failed_label.styles.margin = (0, 0, 1, 0) failed_container.mount(failed_label) @@ -802,7 +790,7 @@ class QuietAgentWorkflowScreen(Screen): failed_container.mount(failed_table) # Action button - done_btn = Button("✔ Done", id="start_over_btn") + 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) @@ -846,3 +834,8 @@ class QuietAgentWorkflowScreen(Screen): self._show_non_enforce_target_selection() 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() diff --git a/TUI/TUI.py b/TUI/TUI.py index 2c5d4fc..1a8ce68 100644 --- a/TUI/TUI.py +++ b/TUI/TUI.py @@ -1,6 +1,20 @@ +# 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 logging import os -import sys from typing import Optional import dotenv @@ -19,24 +33,22 @@ from textual.widgets import ( Tabs, ) -from flows.otp import otp_revoke -from flows.prepPolicy import menu_policy_enforce from models.agent import Agent from models.policy import Policy from services.API import AirlockAPIWrapper -from services.policyhandler import confirmUpdateAfromE -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 TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen +from TUI.Screens.otpactivityscreen import OTPActivitiesScreen +from TUI.Screens.otprevokescreen import OTPRevokeScreen +from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen +from TUI.Screens.policyprepworkflowscreen import PolicyPrepWorkflowScreen +from TUI.Screens.quietagentworkflowscreen import QuietAgentWorkflowScreen +from TUI.Themes.theme_amber_terminal import get_amber_terminal_theme +from TUI.Themes.theme_retro_terminal import get_retro_terminal_theme +from TUI.Themes.themeselector import ThemeSelector +from TUI.Widgets.agentmoveoperations import AgentMoveOperations +from TUI.Widgets.multiagentselector import MultiAgentSelector +from TUI.Widgets.policytreewidget import PolicyTreeWidget +from TUI.Widgets.resultsdisplay import ResultsDisplay from utils.configmanager import get_user_value, load_env, save_user_config from utils.setup import get_base_directory from utils.utils import open_directory @@ -47,7 +59,7 @@ dotenv.load_dotenv() # GLOBAL STASH # --------------------------------------------------------------------------- -_PENDING_JOB = None +_APP_RESTART_REASON = None logger = logging.getLogger(__name__) @@ -83,13 +95,12 @@ class MainMenuScreen(Screen): "🖥️ - Find, Move, or Generate OTP for Agents", "move_agent_workflow_button", ), - ("📊 - Review and appove OTP Activities", "otp_activities_button"), - ("📇 - Find and Move Quiet Hosts to Enforcement", "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"), - ("🔄 - Update Audit Policies", "policy_audit_update_button"), - ("❌ - Revoke OTPs", "otp_revoke_button"), + ("⚖️ - Prepare Policy For Enforcement", "policy_prep_button"), + ("🛑 - Revoke OTPs", "otp_revoke_button"), ], } @@ -128,7 +139,6 @@ class MainMenuScreen(Screen): yield Footer() def on_mount(self) -> None: - api = self.app.api self.switch_tab("agent_actions") # focus helpers @@ -191,42 +201,20 @@ class MainMenuScreen(Screen): self, message: MultiAgentSelector.AgentsSelected ) -> None: """Handle selected agents from AgentSelector.""" - global _PENDING_JOB + global _APP_RESTART_REASON selected_agents = message.selected_agents logger.info("Selected agents: %s", selected_agents) # TODO: Implement actual handling of selected agents - _PENDING_JOB = ("multi_agent_action", selected_agents) + _APP_RESTART_REASON = ("multi_agent_action", selected_agents) self.app.exit() def on_theme_selector_theme_selected( self, message: ThemeSelector.ThemeSelected ) -> None: """Handle theme selection from ThemeSelector.""" - global _PENDING_JOB + global _APP_RESTART_REASON _persist_user_theme(message.theme_name) - _PENDING_JOB = ("restart",) - self.app.exit() - - def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None: - """Handle OTP generation request from the workflow.""" - global _PENDING_JOB - - logger.info( - "OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d", - len(message.devices), - message.requestor, - message.reasoning, - message.duration, - ) - - _PENDING_JOB = ( - "otp_workflow", - message.devices, - message.requestor, - message.reasoning, - message.duration, - ) - + _APP_RESTART_REASON = ("restart",) self.app.exit() def on_agent_move_operations_operation_complete( @@ -282,7 +270,6 @@ class MainMenuScreen(Screen): self.app.bell() def on_button_pressed(self, event: Button.Pressed) -> None: - global _PENDING_JOB button_id = event.button.id logger.debug("Button pressed: %s", button_id) @@ -308,27 +295,23 @@ class MainMenuScreen(Screen): return case "otp_revoke_button": - _PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {}) + self.app.push_screen(OTPRevokeScreen()) + event.stop() + return case "policy_prep_button": - _PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {}) - - case "policy_audit_update_button": - _PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {}) + # Use the new TUI workflow screen instead of legacy + self.app.push_screen( + PolicyPrepWorkflowScreen(self.app.api, self.app.policies) + ) + event.stop() + return case _: self.app.bell() logger.warning("Unknown button pressed: %s", button_id) return - # Only exit the UI loop when we explicitly queued a legacy job. - # The original flow used `self.app.exit()` after setting _PENDING_JOB so - # the outer loop could run legacy code. Keep that behavior only for legacy jobs. - logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB) - if _PENDING_JOB and _PENDING_JOB[0] == "legacy": - # let the main loop pick up the legacy job - self.app.exit() - # --------------------------------------------------------------------------- # 2) APP @@ -353,7 +336,7 @@ class Loxide(App[Message]): ] def __init__(self, api: AirlockAPIWrapper): - self._textual_theme = get_user_value("TEXTUAL_THEME", str, "nord") + self._textual_theme = get_user_value("TEXTUAL_THEME", str, "textual-dark") super().__init__() self.api = api wd = load_env("WORKING_DIR") or os.getcwd() @@ -395,8 +378,8 @@ class Loxide(App[Message]): self.refresh_data() def action_quit(self) -> None: - global _PENDING_JOB - _PENDING_JOB = None + global _APP_RESTART_REASON + _APP_RESTART_REASON = None self.exit() def action_open_fe(self) -> None: @@ -410,45 +393,10 @@ class Loxide(App[Message]): # --------------------------------------------------------------------------- -# 3) TERMINAL + LEGACY -# --------------------------------------------------------------------------- -def _restore_terminal_for_legacy() -> None: - sys.stdout.write("\033[?1049l") - sys.stdout.write("\033[?25h") - sys.stdout.write("\033[0m") - sys.stdout.write("\033[?1000l\033[?1002l\033[?1003l\033[?1006l") - sys.stdout.write("\033[2J\033[H") - sys.stdout.flush() - if os.name == "nt": - try: - import ctypes - - kernel32 = ctypes.windll.kernel32 - handle = kernel32.GetStdHandle(-11) - mode = ctypes.c_ulong() - if kernel32.GetConsoleMode(handle, ctypes.byref(mode)): - kernel32.SetConsoleMode(handle, mode.value | 0x0004) - except Exception as exc: - logger.debug("VT enable on Windows failed: %s", exc) - - -def _run_legacy_job(func, args, kwargs) -> None: - logger.debug("Running legacy job: %s", getattr(func, "__name__", func)) - _restore_terminal_for_legacy() - try: - func(*args, **kwargs) - finally: - try: - input("\nPress Enter to return to the UI...") - except EOFError: - pass - - -# --------------------------------------------------------------------------- -# 4) PUBLIC ENTRYPOINT +# 3) PUBLIC ENTRYPOINT # --------------------------------------------------------------------------- def run_Loxide(api: AirlockAPIWrapper) -> None: - global _PENDING_JOB + global _APP_RESTART_REASON base_dir = get_base_directory() env_path = base_dir / ".env" dotenv.load_dotenv(dotenv_path=env_path, override=True) @@ -458,8 +406,8 @@ def run_Loxide(api: AirlockAPIWrapper) -> None: while attempts < max_attempts: attempts += 1 - logger.debug("Starting job loop iteration (attempt %d)", attempts) - _PENDING_JOB = None + logger.debug("Starting app loop iteration (attempt %d)", attempts) + _APP_RESTART_REASON = None app = Loxide(api) try: @@ -469,42 +417,27 @@ def run_Loxide(api: AirlockAPIWrapper) -> None: logger.debug("Caught SystemExit from Textual: %s", exc) raise - job = _PENDING_JOB - logger.debug("After app.run(), _PENDING_JOB = %r", job) + reason = _APP_RESTART_REASON + logger.debug("After app.run(), _APP_RESTART_REASON = %r", reason) - if not job: - logger.debug("No job pending, exiting loop") + if not reason: + logger.debug("No restart reason, exiting loop") break - if job[0] == "legacy": - _, func, args, kwargs = job - _run_legacy_job(func, args, kwargs) + if reason[0] == "restart": + logger.debug("Restarting app loop") continue - if job[0] == "restart": - logger.debug("Restarting job loop") + if reason[0] == "multi_agent_action": + logger.info("Multi-agent action with selected agents: %s", reason[1]) continue - if job[0] == "multi_agent_action": - logger.info("Multi-agent action with selected agents: %s", job[1]) - continue - - if job[0] == "otp_workflow": - _, devices, requestor, reasoning, duration = job - - def otp_generate_with_params(): - # Your OTP logic here - pass - - _run_legacy_job(otp_generate_with_params, (), {}) - continue - - logger.error("Unknown job type: %r", job) + logger.error("Unknown restart reason: %r", reason) break # --------------------------------------------------------------------------- -# 5) DEV +# 4) DEV # --------------------------------------------------------------------------- if __name__ == "__main__": api = AirlockAPIWrapper() diff --git a/TUI/theme_amber_terminal.py b/TUI/Themes/theme_amber_terminal.py similarity index 55% rename from TUI/theme_amber_terminal.py rename to TUI/Themes/theme_amber_terminal.py index 3662892..738be38 100644 --- a/TUI/theme_amber_terminal.py +++ b/TUI/Themes/theme_amber_terminal.py @@ -1,3 +1,18 @@ +# 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 . + from textual.color import Color from textual.theme import Theme diff --git a/TUI/theme_retro_terminal.py b/TUI/Themes/theme_retro_terminal.py similarity index 50% rename from TUI/theme_retro_terminal.py rename to TUI/Themes/theme_retro_terminal.py index e0cd9c9..1de2669 100644 --- a/TUI/theme_retro_terminal.py +++ b/TUI/Themes/theme_retro_terminal.py @@ -1,3 +1,18 @@ +# 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 . + from textual.color import Color diff --git a/TUI/themeselector.py b/TUI/Themes/themeselector.py similarity index 68% rename from TUI/themeselector.py rename to TUI/Themes/themeselector.py index bcbbb99..b21bb9d 100644 --- a/TUI/themeselector.py +++ b/TUI/Themes/themeselector.py @@ -1,3 +1,18 @@ +# 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 . + from textual.containers import Vertical from textual.message import Message from textual.widget import Widget diff --git a/TUI/OTP_generate.py b/TUI/Widgets/OTP_generate.py similarity index 92% rename from TUI/OTP_generate.py rename to TUI/Widgets/OTP_generate.py index 4ca7977..4dacdf6 100644 --- a/TUI/OTP_generate.py +++ b/TUI/Widgets/OTP_generate.py @@ -1,3 +1,18 @@ +# 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 logging from typing import List, Optional @@ -23,6 +38,8 @@ logger = logging.getLogger(__name__) class OTPGenerator(Widget): + """Widget for generating OTPs for selected devices.""" + # Reactive properties to track form completion requestor_filled = reactive(False) reasoning_filled = reactive(False) @@ -139,14 +156,10 @@ class OTPGenerator(Widget): button_row.styles.height = "auto" button_row.styles.margin = (1, 0, 0, 0) - back_button = Button("← Back", id="back_button") - back_button.styles.width = "1fr" - yield back_button - generate_button = Button( "Generate OTP", id="generate_button", variant="primary" ) - generate_button.styles.width = "2fr" + generate_button.styles.width = "100%" yield generate_button # Right side - Show device list initially, then output after generation @@ -169,7 +182,7 @@ class OTPGenerator(Widget): # Show device list initially device_list_text = "\n".join( - f"• {device.hostname}" for device in self.devices + f"{device.hostname}" for device in self.devices ) device_display = Static(device_list_text, id="device_display") yield device_display @@ -197,14 +210,7 @@ class OTPGenerator(Widget): def on_button_pressed(self, event: Button.Pressed): btn_id = event.button.id - if btn_id == "back_button": - - while len(self.app.screen_stack) > 2: - self.app.pop_screen() - - event.stop() - - elif btn_id == "copy_clipboard_button": + if btn_id == "copy_clipboard_button": try: output_area = self.query_one("#otp_output", TextArea) text_to_copy = output_area.text @@ -213,7 +219,7 @@ class OTPGenerator(Widget): pyperclip.copy(text_to_copy) self.app.notify( - "✅ Copied to clipboard!", severity="information", timeout=2 + "✓ Copied to clipboard!", severity="information", timeout=2 ) except ImportError: self.app.notify( diff --git a/TUI/agentmoveoperations.py b/TUI/Widgets/agentmoveoperations.py similarity index 90% rename from TUI/agentmoveoperations.py rename to TUI/Widgets/agentmoveoperations.py index 26032a0..bf20a82 100644 --- a/TUI/agentmoveoperations.py +++ b/TUI/Widgets/agentmoveoperations.py @@ -1,3 +1,19 @@ +# 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 . + + from dataclasses import asdict from datetime import datetime import logging @@ -10,12 +26,12 @@ from textual.css.query import NoMatches from textual.message import Message from textual.reactive import reactive from textual.widget import Widget -from textual.widgets import Button, DataTable, Header, Static, TextArea +from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea from models.agent import Agent -from TUI.OTP_generate import OTPGenerator -from TUI.otpworkflowscreen import OTPWorkflowScreen -from TUI.policyselectorscreen import PolicySelectorScreen +from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen +from TUI.Screens.policyselectorscreen import PolicySelectorScreen +from TUI.Widgets.OTP_generate import OTPGenerator logger = logging.getLogger(__name__) @@ -183,21 +199,21 @@ class AgentMoveOperations(Widget): f"Operation: {operation_name}", f"{'=' * 50}", "", - f"✅ Successful ({len(successful)}):", + f"✅ Successful ({len(successful)}):", ] if successful: for agent, result in successful: - results_lines.append(f" ✅ {agent.hostname}") + results_lines.append(f" ✅ {agent.hostname}") else: results_lines.append(" (none)") results_lines.append("") - results_lines.append(f"❌ Failed ({len(unsuccessful)}):") + results_lines.append(f"❌ Failed ({len(unsuccessful)}):") if unsuccessful: for agent, error in unsuccessful: - results_lines.append(f" ❌ {agent.hostname}: {error}") + results_lines.append(f" ❌ {agent.hostname}: {error}") else: results_lines.append(" (none)") @@ -232,9 +248,9 @@ class AgentMoveOperations(Widget): - Operations panel: 1/3 width - Results area: Initially hidden, shown after operation completion """ - yield Header(show_clock=True, icon="âš™") + yield Header(show_clock=True, icon="⚙️") title_text = Static( - f"🖥️ Agent Operations - {len(self.agents)} device(s) selected", + f"🖥️ Agent Operations - {len(self.agents)} device(s) selected", id="move_ops_title", ) title_text.styles.margin = (0, 0, 1, 0) @@ -269,32 +285,32 @@ class AgentMoveOperations(Widget): yield operations_label # Operation buttons - export_csv_btn = Button("📈 Export CSV", id="export_csv_btn") + export_csv_btn = Button("📄 Export CSV", id="export_csv_btn") export_csv_btn.styles.width = "100%" export_csv_btn.styles.margin = (0, 0, 1, 0) yield export_csv_btn local_approval_btn = Button( - "✔️ Local Approval Mode", id="local_approval_btn" + "✔️ Local Approval Mode", id="local_approval_btn" ) local_approval_btn.styles.width = "100%" local_approval_btn.styles.margin = (0, 0, 1, 0) yield local_approval_btn - otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn") + otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn") otp_gen_btn.styles.width = "100%" otp_gen_btn.styles.margin = (0, 0, 1, 0) yield otp_gen_btn toggle_enforcement_btn = Button( - "🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn" + "🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn" ) toggle_enforcement_btn.styles.width = "100%" toggle_enforcement_btn.styles.margin = (0, 0, 1, 0) yield toggle_enforcement_btn other_policy_btn = Button( - "🔀 Move to Other Policy", id="other_policy_btn" + "🔀 Move to Other Policy", id="other_policy_btn" ) other_policy_btn.styles.width = "100%" other_policy_btn.styles.margin = (0, 0, 1, 0) @@ -305,10 +321,7 @@ class AgentMoveOperations(Widget): status_label.styles.margin = (2, 0, 0, 0) yield status_label - back_button = Button("← Back", id="back_button") - back_button.styles.width = "50%" - back_button.styles.margin = (0, 1, 1, 0) - yield back_button + yield Footer() def on_mount(self) -> None: """ @@ -343,7 +356,7 @@ class AgentMoveOperations(Widget): Handle button press events from the widget. This Textual event handler routes button presses to appropriate actions: - - back_button: Pop this screen (return to parent) + - copy_results_btn: Copy results text to clipboard (requires pyperclip) - local_approval_btn: Start local approval operation - toggle_enforcement_btn: Start toggle audit/enforcement operation @@ -357,29 +370,24 @@ class AgentMoveOperations(Widget): btn_id = event.button.id - if btn_id == "back_button": - while len(self.app.screen_stack) > 2: - self.app.pop_screen() - event.stop() - - elif btn_id == "copy_results_btn": + if btn_id == "copy_results_btn": try: results_text = self.query_one("#results_text", TextArea) import pyperclip pyperclip.copy(results_text.text) self.app.notify( - "📋✅ Results copied to clipboard!", + "📋✅ Results copied to clipboard!", severity="information", timeout=2, ) except ImportError: self.app.notify( - "❌ pyperclip not installed. Run: pip install pyperclip", + "❌ pyperclip not installed. Run: pip install pyperclip", severity="warning", ) except Exception as e: - self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error") + self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error") event.stop() elif btn_id == "export_csv_btn": self._start_export_csv_operation() @@ -427,7 +435,7 @@ class AgentMoveOperations(Widget): self.operation_in_progress = True status_label = self.query_one("#status_label", Static) - status_label.update("✔️ Moving agents to local approval...") + status_label.update("✔️ Moving agents to local approval...") # Get API from app api = self.app.api @@ -462,12 +470,12 @@ class AgentMoveOperations(Widget): except Exception as e: logger.error(f"Error during local approval operation: {e}") - status_label.update(f"❌ Error: {str(e)}") + status_label.update(f"❌ Error: {str(e)}") self.operation_in_progress = False return self.operation_in_progress = False - status_label.update("✅ Operation complete!") + status_label.update("✅ Operation complete!") # Display results in the widget self._display_results("Local Approval Mode", successful, unsuccessful) @@ -483,7 +491,6 @@ class AgentMoveOperations(Widget): self.selected_operation = "export_csv" self.operation_in_progress = True successful = [] - unsuccessful = [] status_label = self.query_one("#status_label", Static) status_label.update("Exporting CSV...") self.app.refresh_data() @@ -511,9 +518,9 @@ class AgentMoveOperations(Widget): file_path = os.path.join(str(path), filename) df.to_csv(file_path, index=False) successful.append(file_path) - status_label.update(f"✅ Exported to {file_path}") + status_label.update(f"✅ Exported to {file_path}") except Exception: - status_label.update("❌ Failed") + status_label.update("❌ Failed") self.operation_in_progress = False @@ -557,7 +564,7 @@ class AgentMoveOperations(Widget): self.operation_in_progress = True status_label = self.query_one("#status_label", Static) - status_label.update("⏳ Toggling enforcement mode...") + status_label.update("🔄 Toggling enforcement mode...") # Get API from app api = self.app.api @@ -593,12 +600,12 @@ class AgentMoveOperations(Widget): except Exception as e: logger.error(f"Error during toggle enforcement operation: {e}") - status_label.update(f"❌ Error: {str(e)}") + status_label.update(f"❌ Error: {str(e)}") self.operation_in_progress = False return self.operation_in_progress = False - status_label.update("✅ Operation complete!") + status_label.update("✅ Operation complete!") # Display results in the widget self._display_results("Toggle Audit/Enforcement", successful, unsuccessful) @@ -663,7 +670,7 @@ class AgentMoveOperations(Widget): except Exception as e: logger.error(f"Error loading policies: {e}") - status_label.update(f"❌ Error: {str(e)}") + status_label.update(f"❌ Error: {str(e)}") self.operation_in_progress = False self.selected_operation = "" self.app.notify(f"Failed to load policies: {str(e)}", severity="error") @@ -695,7 +702,7 @@ class AgentMoveOperations(Widget): for agent in self.agents: try: # Move agent to target policy - result = api.agent_move(agent.agentid, target_policy.groupid) + api.agent_move(agent.agentid, target_policy.groupid) successful.append((agent, f"Moved to {target_policy.name}")) logger.info( f"Successfully moved {agent.hostname} to policy {target_policy.name}" diff --git a/TUI/multiagentselector.py b/TUI/Widgets/multiagentselector.py similarity index 88% rename from TUI/multiagentselector.py rename to TUI/Widgets/multiagentselector.py index afe8e14..bc77706 100644 --- a/TUI/multiagentselector.py +++ b/TUI/Widgets/multiagentselector.py @@ -1,3 +1,18 @@ +# 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 difflib import re from typing import List, Optional @@ -20,6 +35,8 @@ from models.agent import Agent class MultiAgentSelector(Widget): + """Widget for selecting multiple agents from a list.""" + class AgentsSelected(Message): def __init__(self, selected_agents: List[Agent]): super().__init__() @@ -40,7 +57,7 @@ class MultiAgentSelector(Widget): def compose(self): yield Header(show_clock=True, icon="⚙") - title_text = Static("🖧 Agent Selector", id="selector_title") + title_text = Static("🖥️ Agent Selector", id="selector_title") title_text.styles.margin = (0, 0, 0, 1) yield title_text @@ -60,7 +77,7 @@ class MultiAgentSelector(Widget): text_area.styles.overflow_y = "auto" yield text_area - with Horizontal(id="switch_search_container") as switch_search: + with Horizontal(id="switch_search_container"): switch = Switch(value=False, id="match_switch") switch.styles.width = "auto" switch.styles.margin = (1, 0, 0, 0) @@ -91,11 +108,6 @@ class MultiAgentSelector(Widget): button_row.styles.height = "auto" button_row.styles.margin = (1, 0, 0, 0) - back_button = Button("← Back", id="back_button") - back_button.styles.width = "1fr" - back_button.styles.margin = (0, 0, 0, 1) - yield back_button - submit_button = Button( "▶ Select & Continue", id="submit_selection", variant="primary" ) @@ -123,10 +135,7 @@ class MultiAgentSelector(Widget): match_list = self.query_one("#match_results", SelectionList) except NoMatches: return - if btn_id == "back_button": - self.app.pop_screen() - event.stop() - elif btn_id == "select_all": + if btn_id == "select_all": match_list.select_all() event.stop() elif btn_id == "select_none": diff --git a/TUI/policyselector.py b/TUI/Widgets/policyselector.py similarity index 92% rename from TUI/policyselector.py rename to TUI/Widgets/policyselector.py index d6e4bdf..6e0b71d 100644 --- a/TUI/policyselector.py +++ b/TUI/Widgets/policyselector.py @@ -1,10 +1,17 @@ -""" -Policy Selector Widget Module - -Provides a Textual widget for selecting target policies for bulk agent operations. -Allows users to browse available policies and select one as the destination for -moving agents. Automatically excludes parent/logical policies. -""" +# 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 logging import re @@ -34,7 +41,7 @@ class PolicySelector(Widget): - Wildcard filtering (* and ?) - Interactive table for policy browsing - Explicit confirm button for selection - - Cancel/back button to dismiss + - Use escape key to go back Attributes: policies (list[Policy]): List of available Policy objects to display. @@ -91,10 +98,10 @@ class PolicySelector(Widget): - Clear Filter button - Confirm Selection button - Policy table displaying available policies - - Back buttons for navigation + - Use escape key to go back """ title_text = Static( - "🎯 Select Target Policy", + "Select Target Policy", id="policy_selector_title", ) title_text.styles.margin = (0, 0, 1, 0) @@ -125,12 +132,12 @@ class PolicySelector(Widget): filter_help.styles.margin = (0, 0, 1, 0) yield filter_help - apply_button = Button("✓ Apply Filter", id="filter_button") + apply_button = Button("🔍 Apply Filter", id="filter_button") apply_button.styles.width = "100%" apply_button.styles.margin = (0, 0, 1, 0) yield apply_button - clear_button = Button("Clear Filter", id="clear_filter_button") + clear_button = Button("🧹 Clear Filter", id="clear_filter_button") clear_button.styles.width = "100%" clear_button.styles.margin = (0, 0, 1, 0) yield clear_button @@ -144,11 +151,6 @@ class PolicySelector(Widget): selected_label.styles.margin = (2, 0, 1, 0) yield selected_label - cancel_button = Button("← Back", id="back_button") - cancel_button.styles.width = "100%" - cancel_button.styles.margin = (1, 0, 1, 0) - yield cancel_button - # Right side - Policy table with Vertical() as right_side: right_side.styles.width = "2fr" @@ -222,7 +224,6 @@ class PolicySelector(Widget): Handle button press events from the widget. Routes to: - - back_button (Cancel): Pop screen without selecting - filter_button (Apply Filter): Filter policies with wildcard support - clear_filter_button: Clear filter and show all policies - confirm_button: Confirm selection and post message @@ -232,12 +233,7 @@ class PolicySelector(Widget): """ btn_id = event.button.id - if btn_id == "back_button": - while len(self.app.screen_stack) > 2: - self.app.pop_screen() - event.stop() - - elif btn_id == "filter_button": + if btn_id == "filter_button": self._apply_filter() event.stop() @@ -283,7 +279,7 @@ class PolicySelector(Widget): if self.selected_policy: # Update selection display label = self.query_one("#selected_policy_label", Static) - label.update(f"✓ Selected: {self.selected_policy.name}") + label.update(f"Selected: {self.selected_policy.name}") # Log for debugging logger.debug( @@ -325,7 +321,7 @@ class PolicySelector(Widget): if highlighted_name: label = self.query_one("#selected_policy_label", Static) - label.update(f"→ Highlighting: {highlighted_name}") + label.update(f"Highlighting: {highlighted_name}") except Exception as e: logger.error(f"Error handling row highlight: {e}") @@ -399,7 +395,9 @@ class PolicySelector(Widget): ) displayed_count = len(self._displayed_policies) - status_text = f"📊 Showing {displayed_count} of {len(self._filtered_policies)} policies" + status_text = ( + f"Showing {displayed_count} of {len(self._filtered_policies)} policies" + ) self.app.notify(status_text, severity="information", timeout=2) # Clear selection when filter is applied @@ -474,7 +472,7 @@ class PolicySelector(Widget): """ if self.selected_policy is None: self.app.notify( - "⚠️ Please select a policy first by clicking on a row in the table", + "Please select a policy first by clicking on a row in the table", severity="warning", timeout=3, ) @@ -483,6 +481,6 @@ class PolicySelector(Widget): # Log confirmation for debugging logger.info(f"Confirming selection of policy: {self.selected_policy.name}") self.app.notify( - f"✅ Confirmed: {self.selected_policy.name}", severity="success", timeout=2 + f"Confirmed: {self.selected_policy.name}", severity="success", timeout=2 ) self.post_message(self.PolicySelected(self.selected_policy)) diff --git a/TUI/policytreewidget.py b/TUI/Widgets/policytreewidget.py similarity index 93% rename from TUI/policytreewidget.py rename to TUI/Widgets/policytreewidget.py index 5e01fdf..d51da78 100644 --- a/TUI/policytreewidget.py +++ b/TUI/Widgets/policytreewidget.py @@ -1,3 +1,18 @@ +# 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 . + from collections import defaultdict import logging diff --git a/TUI/Widgets/prepPolicy.py b/TUI/Widgets/prepPolicy.py new file mode 100644 index 0000000..f3448a0 --- /dev/null +++ b/TUI/Widgets/prepPolicy.py @@ -0,0 +1,870 @@ +# 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 logging +import os +import os.path +import re +from typing import List + +import dotenv +import pandas as pd + +from models.execution import ExecutionHistoryRecord +from models.policy import Allowlist, Policy +from services.API import AirlockAPIWrapper +from utils.configmanager import get_system_list, get_system_value, load_env +from utils.selector import Selector +from utils.utils import ( + areYouSure, + clear_screen, + colorText, + formatHTML, + get_sanitized_input, + locked, + open_directory, + print_x_wide, + regulator, +) + +logger = logging.getLogger(__name__) + +dotenv.load_dotenv() + + +def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]: + + policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()] + logger.debug("Prompting for Policies") + print(colorText("Please select policy/policies", "white")) + selected = Selector.select_objects(policies, allow_multiple, prompt_each=True) + + if selected is None: + return [] + + # Normalize to always return a list + logger.debug("Returning {selected.dict}") + return selected if isinstance(selected, list) else [selected] + + +def selectAllowlists( + api: AirlockAPIWrapper, policy=all, allow_multiple=True +) -> List[Allowlist]: + if policy == "all": + allowlists = [ + Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows() + ] + else: + allowlists = [ + Allowlist(**row.to_dict()) + for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows() + ] + logger.debug("Prompting for Allowlist(s)") + print(colorText("Please select allowlist(s)", "white")) + selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True) + + if selected is None: + return [] + + # Normalize to always return a list + logger.debug(f"Returning {selected}") + return selected if isinstance(selected, list) else [selected] + + +def sortHashes( + api: AirlockAPIWrapper, selected_policies: List[Policy], type=[1, 2, 6, 7] +): + working_dir = load_env("WORKING_DIR") + history_days = Selector.select_value( + prompt="Enter how many days of history to pull (1–150): ", + value_type=int, + valid_range=(1, 150), + ) + + logger.debug(f"{history_days} day selected for history") + + if history_days is None: + logging.warning("No history range selected. Aborting.") + return + + policy_executions = ExecutionHistoryRecord.from_policies( + api, selected_policies, type_=type, history_days=history_days + ) + + logger.debug(f"Executions contains {policy_executions}") + + enriched_executions = ExecutionHistoryRecord.enrich_with_hashes( + api, policy_executions + ) + categorized_executions = ( + ExecutionHistoryRecord.categorize_executions_by_hash_decision( + enriched_executions + ) + ) + approved, unapproved, needs_review, unknown = ( + ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions) + ) + + categories = { + "needs_review": needs_review, + "approved": approved, + "unapproved": unapproved, + "leftover": unknown, + } + + for label, records in categories.items(): + if not records: + continue # Skip empty or falsy categories + + csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv" + html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html" + + # Convert ExecutionHistoryRecord objects to dictionaries + df = pd.DataFrame([r.__dict__ for r in records]) + + # Optional: flatten hash_obj if needed + if not df.empty and "hash_obj" in df.columns: + hash_df = df["hash_obj"].apply(lambda h: h.to_dict() if h else {}) + df = pd.concat([df.drop(columns=["hash_obj"]), hash_df], axis=1) + + # Save to CSV + df.to_csv(csv_path, index=False) + logger.info(f"Saved {label} executions to {csv_path}") + + # Generate HTML + formatHTML(df, html_path) + logger.info(f"Generated HTML report at {html_path}") + + +def buildPathsandPublishers(selected_policies: List[Policy], split): + working_dir = load_env("WORKING_DIR") + df1 = pd.DataFrame() + df2 = pd.DataFrame() + all_approved_hashes = pd.DataFrame() + path1 = ( + f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv" + ) + path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv" + path_exclusion_constant = get_system_value("PATH_EXCLUSION_CONST", cast_type=int) + + if os.path.exists(path1): + df1 = pd.read_csv(path1) + else: + logger.warning(f"File not found: {path1}") + + if os.path.exists(path2): + df2 = pd.read_csv(path2) + else: + logger.warning(f"File not found: {path2}") + + if df1.empty and df2.empty: + logger.warning("Both DataFrames are empty. Skipping sort.") + all_approved_hashes = pd.DataFrame() + logger.debug(all_approved_hashes.head) + else: + all_approved_hashes = pd.concat([df1, df2], ignore_index=True) + if "filename" in all_approved_hashes.columns: + all_approved_hashes = all_approved_hashes.sort_values(by="filename") + else: + logger.warning( + "Warning: 'filename' column not found in concatenated DataFrame." + ) + + if not all_approved_hashes.empty and path_exclusion_constant: + + primary_path_exclusions = calculatePath( + all_approved_hashes, + path_exclusion_constant, + split, + ) + remaining_hashes = all_approved_hashes[ + ~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"]) + ] + secondary_path_exclusions = calculatePath( + remaining_hashes, (path_exclusion_constant - 1), split + ) + remaining_hashes = remaining_hashes[ + ~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"]) + ] + dataframes = { + "all_approved_hashes": all_approved_hashes, + "primary_Paths": primary_path_exclusions, + "secondary_Paths": secondary_path_exclusions, + "hashes_not_approvable_by_path": remaining_hashes, + } + logger.debug("Preparing to sort dataframes") + for name, df in dataframes.items(): + logger.debug(f" DataFrame headers: {list(df.columns)}") + if "hashes" in name: + df.sort_values(by="filename", inplace=True) + else: + df.sort_values(by="longestcfp", inplace=True) + + df.to_csv( + f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv", + index=False, + ) + formatHTML( + df, + f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html", + ) + + if not all_approved_hashes.empty: + # Drop all not signed, only keep unique values + publist = all_approved_hashes[ + all_approved_hashes["publisher"] != "Not Signed" + ].drop_duplicates(subset=["publisher"]) + # Remove Bad publisher if somehow they made it this far + pattern = regulator(get_system_list("BAD_PUBLISHERS")) + publist = publist[~publist["publisher"].str.contains(pattern, na=False)] + publist = publist[["publisher"]] + publist.sort_values(by="publisher", inplace=True) + publist.to_csv( + f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv", + index=False, + ) + else: + logger.debug("Approved Hashes list appears empty") + + +def buildPreflights(selected_policies: List[Policy]): + working_dir = load_env("WORKING_DIR") + + df1 = pd.DataFrame() + df2 = pd.DataFrame() + approved_hashes = pd.DataFrame() + approved_publishers = pd.DataFrame() + + hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv" + path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv" + path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv" + publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv" + + # Read in and combine the two path generations + if os.path.exists(path1): + df1 = pd.read_csv(path1) + else: + logger.warning(f"File not found: {path1}") + + if os.path.exists(path2): + df2 = pd.read_csv(path2) + else: + logger.warning(f"File not found: {path2}") + + if df1.empty and df2.empty: + logger.warning("Both DataFrames are empty. Skipping sort.") + approved_paths = pd.DataFrame() + else: + approved_paths = pd.concat([df1, df2], ignore_index=True) + + approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first") + + # We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions. + if os.path.exists(hash): + hashes = pd.read_csv(hash) + approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])] + + approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first") + + else: + logger.warning(f"File not found: {hash}") + + if os.path.exists(publishers): + approved_publishers = pd.read_csv(publishers) + + else: + logger.warning(f"File not found: {publishers}") + + dataframes = { + "approved_paths": approved_paths, + "approved_hashes": approved_hashes, + "approved_publishers": approved_publishers, + } + + for name, df in dataframes.items(): + logger.debug(f" DataFrame headers: {list(df.columns)}") + if name == "approved_paths": + df.sort_values(by="longestcfp", inplace=True) + elif name == "approved_hashes": + df.sort_values(by="filename", inplace=True) + elif name == "approved_publishers": + df.sort_values(by="publisher", inplace=True) + + df.to_csv( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv", + index=False, + ) + formatHTML( + df, + f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html", + ) + + +def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): + min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int) + + def clean_split(path): + 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 + + # Diagnostic: log any non-string entries + non_string_entries = df[ + ~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike))) + ] + if not non_string_entries.empty: + print(f"[WARNING] Non-string entries found in column '{col}':") + print(non_string_entries) + + df = df.copy() + split_paths = df[col].apply(clean_split) + + 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] + + 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): + 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) + + 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) + + return pd.DataFrame(new_rows).drop(columns=["group_key"]) + + +def calculatePath(approved_hashes, path_exclusion_constant, split): + if split: + dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")] + else: + dfs_by_policy = [approved_hashes] + + badpathparts = get_system_list("BAD_PATH_PARTS") + min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int) + + processed_dfs = [] + + for df in dfs_by_policy: + haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename") + haslcp = haslcp.drop_duplicates() + + forbidden = regulator(badpathparts, True) + forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False) + + logger.debug("Removing forbidden filepaths for path exceptions") + print(colorText("Removing forbidden filepaths for path exceptions", "green")) + lcp_not_forbidden = haslcp[~forbidden_lcfp].copy() + + lcp_not_forbidden_review = lcp_not_forbidden[ + [ + "policyname", + "longestcfp", + "middle", + "filename_only", + "file_extension", + "sha256", + ] + ] + + unique_sha_counts = ( + lcp_not_forbidden_review.groupby("longestcfp")["sha256"] + .nunique() + .reset_index() + ) + unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"] + + lcp_not_forbidden_review = lcp_not_forbidden_review.merge( + unique_sha_counts, on="longestcfp", how="left" + ) + lcp_not_forbidden_review = lcp_not_forbidden_review[ + lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path + ] + processed_dfs.append(lcp_not_forbidden_review) + + pathExclusions = pd.concat(processed_dfs, ignore_index=True) + + return pathExclusions + + +def testChange(selected_policies, destination_policy, destination_allowlist): + working_dir = load_env("WORKING_DIR") + + logger.info("These path exclusions would be added to:") + logger.info(destination_policy) + + pathexclusions = pd.read_csv( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv" + ) + hashes = pd.read_csv( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv" + ) + + unique_combinations = pathexclusions[ + ["longestcfp", "file_extension"] + ].drop_duplicates() + + drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\") + processed_paths = [ + (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}" + for path, ext in unique_combinations.itertuples(index=False, name=None) + ] + + for path in processed_paths: + logger.info(path) + + print(colorText("These publishers would added", "yellow")) + processed_publishers = [] + if os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv" + ): + publishers = pd.read_csv( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv" + ) + if publishers.empty: + print(colorText("The publishers list is empty.", "red")) + else: + processed_publishers = ( + publishers[publishers["publisher"] != "Not Signed"]["publisher"] + .drop_duplicates() + .tolist() + ) + for publisher in processed_publishers: + print(publisher) + + print(colorText("These hashes would be added to:", "yellow")) + print(destination_allowlist) + + processed_hashes = hashes["sha256"].unique().tolist() + print_x_wide(processed_hashes, 3) + + return processed_paths, processed_hashes, processed_publishers + + +def menu_policy_enforce( + api: AirlockAPIWrapper, +): # TODO Need to clean up 6 and 7 into functions + selected_policies = [] + destination_policy = [] + destination_allowlist = [] + processed_paths = [] + processed_hashes = [] + processed_publishers = [] + working_dir = load_env("WORKING_DIR") + + while True: + printEnforceChecklist( + selected_policies, destination_policy, destination_allowlist + ) + choice = get_sanitized_input("\nEnter your choice: ") + + if choice == "1": + clear_screen() + selected_policies = selectPolicies(api, True) + + elif choice == "2": + clear_screen() + print( + colorText( + "Please choose destination_name Policy for Path Exclusions", "white" + ) + ) + + destination_policy = selectPolicies(api, False) + + print(colorText("Please choose Allowlist for Hashes", "white")) + + destination_allowlist = selectAllowlists(api, destination_policy, False) + + elif choice == "3": + clear_screen() + sortHashes( + api, + selected_policies, + type=[1, 2, 6, 7], + ) + + elif choice == "4": + clear_screen() + if os.path.exists( + f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv" + ): + buildPathsandPublishers(selected_policies, False) + else: + print( + "File not found. Please make sure it's saved correctly and try again." + ) + + elif choice == "5": + clear_screen() + if os.path.exists( + f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv" + ) and os.path.exists( + f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv" + ): + buildPreflights(selected_policies) + else: + print( + "File not found. Please make sure it's saved correctly and try again." + ) + + elif choice == "6": + clear_screen() + if ( + os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv" + ) + and os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv" + ) + and destination_policy + and destination_allowlist + ): + processed_paths, processed_hashes, processed_publishers = testChange( + selected_policies, destination_policy, destination_allowlist + ) + else: + # Log which condition(s) failed + missing_items = [] + if not os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv" + ): + missing_items.append("approved_paths.csv not found") + if not os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv" + ): + missing_items.append("approved_hashes.csv not found") + if not destination_policy: + missing_items.append("destination_policy is empty or None") + if not destination_allowlist: + missing_items.append("destination_allowlist is empty or None") + + logger.error("Preflight check failed due to the following:") + for item in missing_items: + logger.error(f" - {item}") + + elif choice == "7": + clear_screen() + areYouSure() + confirmation = get_sanitized_input("Type 'I AGREE' to continue: ") + if ( + processed_paths + and processed_hashes + and processed_publishers + and destination_policy + and destination_allowlist + and confirmation.strip() == "I AGREE" + ): + print(colorText("Proceeding with the code...", "yellow")) + api.hash_add_to_allowlist( + destination_allowlist[0].applicationid, processed_hashes + ) + api.policy_add_path_exclusions( + destination_policy[0].groupid, processed_paths + ) + if processed_publishers: + api.policy_add_publishers( + destination_policy[0].groupid, processed_publishers + ) + + locked() + + else: + logger.error("Confirmation block failed. Reasons:") + if not processed_publishers or processed_hashes or processed_paths: + logger.error(" - Test not performed.") + if not destination_policy: + logger.error(" - `destination_policy` is missing or invalid.") + if not destination_allowlist: + logger.error(" - `destination_allowlist` is missing or invalid.") + if confirmation.strip() != "I AGREE": + logger.error( + " - User did not confirm with 'I AGREE'. Received: '%s'", + confirmation.strip(), + ) + + elif choice.upper() == "F": + open_directory(working_dir) + elif choice.upper() == "B": + break + + else: + print(colorText("Invalid choice. Please try again.", "red")) + + +def section_header(title): + print( + colorText( + "\n --------------------------------------------------------------------", + "cyan", + ) + ) + print(colorText(f" ------------- {title} -------------", "cyan")) + print( + colorText( + " --------------------------------------------------------------------", + "cyan", + ) + ) + + +def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist): + working_dir = load_env("WORKING_DIR") + section_header("Prepare to Enforce Policy ") + print( + colorText( + "\nSequentially follow these steps to prepare a policy for enforcement:", + "white", + ) + ) + + # Step 1: Originating Policies + print( + colorText( + "\n1. Choose which policy or policies to gather execution info from", "cyan" + ) + ) + if not selected_policies: + print(colorText(" [✗] No policies have been chosen", "red")) + else: + print(colorText("The following policies have been chosen:", "green")) + for policy in selected_policies: + print(colorText(f" [✓] {policy.name}", "green")) + + # Step 2: Destination Policy and Allowlist + print( + colorText("2. Choose the destination policy and associated allowlist", "cyan") + ) + if destination_policy: + print( + colorText( + f" [✓] {destination_policy[0].name} has been selected as the destination policy", + "green", + ) + ) + else: + print(colorText(" [✗] No destination policy has been chosen", "red")) + + if destination_allowlist: + print( + colorText( + f" [✓] {destination_allowlist[0].name} has been selected as allowlist", + "green", + ) + ) + else: + print(colorText(" [✗] No allowlist has been chosen", "red")) + + # Step 3: Data Preparation + print( + colorText( + f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review", + "cyan", + ) + ) + if selected_policies: + policy_id = selected_policies[0].name + review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv" + print( + colorText( + ( + " [✓] Data has been fetched" + if os.path.exists(review_path) + else " [✗] Data has not been fetched" + ), + "green" if os.path.exists(review_path) else "red", + ) + ) + else: + print( + colorText( + " [✗] No policies selected, cannot check data fetch status", "red" + ) + ) + + # Step 4: Manual Review + print(colorText("4. Manually review the files:", "cyan")) + print( + colorText( + " Remove the rows containing hashes you do not approve of", "cyan" + ) + ) + print( + colorText( + f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.", + "cyan", + ) + ) + print( + colorText( + " This will start the process to generate possible filepath approvals", + "cyan", + ) + ) + + if selected_policies: + policy_id = selected_policies[0].name + approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv" + second_review_path = ( + f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv" + ) + print( + colorText( + ( + " [✓] Reviewed hashes have been loaded" + if os.path.exists(approved_path) + else " [✗] Reviewed hashes have not been loaded" + ), + "green" if os.path.exists(approved_path) else "red", + ) + ) + print( + colorText( + ( + " [✓] Path review list created" + if os.path.exists(second_review_path) + else " [✗] Path review list has not been created" + ), + "green" if os.path.exists(second_review_path) else "red", + ) + ) + else: + print( + colorText( + " [✗] No policies selected, cannot check reviewed hashes or path list", + "red", + ) + ) + + # Step 5: Path Review + print( + colorText( + f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\", + "cyan", + ) + ) + print( + colorText( + " Remove the rows containing path exclusions or publishers you do not approve of.", + "cyan", + ) + ) + print( + colorText( + f" When complete, save the files to {working_dir}\\data\\Approved", + "cyan", + ) + ) + print( + colorText(" Choose this option when done to build your preflights", "cyan") + ) + + if selected_policies: + policy_id = selected_policies[0].name + reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv" + preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv" + preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv" + print( + colorText( + ( + " [✓] Reviewed path list detected" + if os.path.exists(reviewed_path) + else " [✗] Path review list has not been detected" + ), + "green" if os.path.exists(reviewed_path) else "red", + ) + ) + preflight_ready = os.path.exists(preflight_paths) and os.path.exists( + preflight_hashes + ) + print( + colorText( + ( + " [✓] Preflight Path Exclusion List has been generated" + if preflight_ready + else " [✗] Preflight Path Exclusion List has not been generated" + ), + "green" if preflight_ready else "red", + ) + ) + else: + print( + colorText( + " [✗] No policies selected, cannot check preflight status", "red" + ) + ) + + # Final Steps + print( + colorText( + "6. Test ------------------------------------------------------", "cyan" + ) + ) + print( + colorText( + " Prints to console the changes that would be made, must be done to proceed. ", + "cyan", + ) + ) + + print( + colorText( + "7. Liftoff ------------------------------------------------------", "cyan" + ) + ) + print( + colorText( + " Apply path exclusions and approved publishers to selected policy", + "cyan", + ) + ) + print(colorText(" Apply approved hashes to allowlist", "cyan")) + + # Utility Options + print(colorText("F. Open Working Directory", "cyan")) + print(colorText("B. Back", "cyan")) diff --git a/TUI/resultsdisplay.py b/TUI/Widgets/resultsdisplay.py similarity index 74% rename from TUI/resultsdisplay.py rename to TUI/Widgets/resultsdisplay.py index ebabce7..63b204d 100644 --- a/TUI/resultsdisplay.py +++ b/TUI/Widgets/resultsdisplay.py @@ -1,3 +1,18 @@ +# 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 logging from textual.containers import Horizontal, Vertical @@ -59,15 +74,6 @@ class ResultsDisplay(Widget): margin-top: 1; width: 100%; } - - #button_row { - height: auto; - margin: 1 0 0 0; - } - - #back_button { - width: 1fr; - } """ class CopySuccess(Message): @@ -95,9 +101,9 @@ class ResultsDisplay(Widget): def compose(self): with Vertical(id="results_screen"): - yield Header(show_clock=True, icon="⚙") + yield Header(show_clock=True, icon="⚙️") # Title - title = Static(f"📊 {self.operation} - Results", id="results_title") + title = Static(f"{self.operation} - Results", id="results_title") yield title # Two-column layout @@ -107,7 +113,7 @@ class ResultsDisplay(Widget): yield Static("✅ Successful", id="success_label") yield Static(self.successful_results, id="success_results") yield Button( - "📋✅ Copy Success List", + "Copy Success List", id="copy_success", classes="copy_button", ) @@ -117,15 +123,11 @@ class ResultsDisplay(Widget): yield Static("❌ Failed", id="failure_label") yield Static(self.unsuccessful_results, id="failure_results") yield Button( - "📋❌ Copy Failure List", + "Copy Failure List", id="copy_failure", classes="copy_button", ) - # Back Button - with Horizontal(id="button_row"): - back_button = Button("← Back", id="back_button") - yield back_button yield Footer() def on_button_pressed(self, event: Button.Pressed) -> None: @@ -138,18 +140,18 @@ class ResultsDisplay(Widget): pyperclip.copy(str(success_widget.renderable)) self.app.notify( - "✅ Success list copied to clipboard!", + "Success list copied to clipboard!", severity="information", timeout=2, ) self.post_message(self.CopySuccess()) except ImportError: self.app.notify( - "⚠️ pyperclip not installed. Run: pip install pyperclip", + "❌ pyperclip not installed. Run: pip install pyperclip", severity="warning", ) except Exception as e: - self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error") + self.app.notify(f"¢ Failed to copy: {str(e)}", severity="error") event.stop() elif btn_id == "copy_failure": @@ -159,20 +161,16 @@ class ResultsDisplay(Widget): pyperclip.copy(str(failure_widget.renderable)) self.app.notify( - "✅ Failure list copied to clipboard!", + "Failure list copied to clipboard!", severity="information", timeout=2, ) self.post_message(self.CopyFailure()) except ImportError: self.app.notify( - "⚠️ pyperclip not installed. Run: pip install pyperclip", + "❌ pyperclip not installed. Run: pip install pyperclip", severity="warning", ) except Exception as e: - self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error") - event.stop() - - elif btn_id == "back_button": - self.app.pop_screen() + self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error") event.stop() diff --git a/TUI/moveagentworkflowscreen.py b/TUI/moveagentworkflowscreen.py deleted file mode 100644 index 8398ddf..0000000 --- a/TUI/moveagentworkflowscreen.py +++ /dev/null @@ -1,61 +0,0 @@ -from typing import List, Optional - -from textual.app import ComposeResult -from textual.screen import Screen - -from models.agent import Agent -from TUI.agentmoveoperations import AgentMoveOperations -from TUI.multiagentselector import MultiAgentSelector -from TUI.resultsdisplay import ResultsDisplay - - -class MoveAgentWorkflowScreen(Screen): - """Screen that handles the agent movement workflow.""" - - def __init__(self, all_agents: Optional[List[Agent]]): - super().__init__() - self.all_agents = all_agents - self.selected_agents = None - - def compose(self) -> ComposeResult: - """Start with the multi-agent selector.""" - yield MultiAgentSelector(self.all_agents) - - def on_multi_agent_selector_agents_selected( - self, message: MultiAgentSelector.AgentsSelected - ) -> None: - """Handle selected agents - switch to operations screen.""" - self.selected_agents = message.selected_agents - - # Remove the MultiAgentSelector - selector = self.query_one(MultiAgentSelector) - selector.remove() - - # Mount the AgentMoveOperations with the selected Agent objects - self.mount(AgentMoveOperations(self.selected_agents)) - - def on_agent_move_operations_operation_complete( - self, message: AgentMoveOperations.OperationComplete - ) -> None: - """Handle completion of move operation - transition to results screen.""" - # Format successful results - success_lines = [] - for agent, result in message.successful: - success_lines.append(f"✓ {agent.hostname}") - - # Format unsuccessful results - failure_lines = [] - for agent, error in message.unsuccessful: - failure_lines.append(f"✗ {agent.hostname}: {error}") - - successful_text = "\n".join(success_lines) if success_lines else "(none)" - unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)" - - # Remove the operations widget - ops_widget = self.query_one(AgentMoveOperations) - ops_widget.remove() - - # Mount the results display - self.mount( - ResultsDisplay(message.operation, successful_text, unsuccessful_text) - ) diff --git a/TUI/otpworkflowscreen.py b/TUI/otpworkflowscreen.py deleted file mode 100644 index d7d7322..0000000 --- a/TUI/otpworkflowscreen.py +++ /dev/null @@ -1,24 +0,0 @@ -# otp_workflow_screen.py - -from typing import List, Optional - -from textual.app import ComposeResult -from textual.screen import Screen - -from models.agent import Agent -from TUI.OTP_generate import OTPGenerator - - -class OTPWorkflowScreen(Screen): - """Screen that handles the OTP generation workflow without agent selection.""" - - def __init__(self, selected_agents: Optional[List[Agent]]): - super().__init__() - self.selected_agents = selected_agents - - def compose(self) -> ComposeResult: - """Directly show the OTP generator for the selected agents.""" - yield OTPGenerator(self.selected_agents) - - def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None: - """Handle OTP generation request - pass it up to the app level if needed.""" diff --git a/flows/localApproval.py b/flows/localApproval.py index f1054f1..d6bda28 100644 --- a/flows/localApproval.py +++ b/flows/localApproval.py @@ -1,6 +1,17 @@ -""" -This module handles the creation of local approval requests. -""" +# 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 logging import os @@ -51,7 +62,7 @@ class LocalApprovalRequestor: batch_id = int(time.time()) purpose = ( - f"🎫 Local Approval 🎫 - {duration_minutes} mins - " + f" Local Approval - {duration_minutes} mins - " f"batch:{batch_id} Client:{agent_id} User:{self.username}" ) @@ -103,11 +114,9 @@ class LocalApprovalRequestor: success_count = 0 failure_count = 0 - print(colorText(f"\n📦 Processing batch {batch_id}...", "cyan")) - print(colorText(f"👤 Requested by: {self.username}", "cyan")) - print( - colorText(f"📊 Moving {len(agents)} agent(s) to local approval\n", "cyan") - ) + print(colorText(f"\n Processing batch {batch_id}...", "cyan")) + print(colorText(f" Requested by: {self.username}", "cyan")) + print(colorText(f" Moving {len(agents)} agent(s) to local approval\n", "cyan")) for agent in agents: try: @@ -125,11 +134,11 @@ class LocalApprovalRequestor: if not move_success: raise Exception("Failed to move to audit policy") - print(colorText(f"✓ {agent.hostname}", "green")) + print(colorText(f" {agent.hostname}", "green")) success_count += 1 except Exception as e: - print(colorText(f"✗ {agent.hostname}: {e}", "red")) + print(colorText(f" {agent.hostname}: {e}", "red")) logger.error(f"Error processing agent {agent.hostname}: {e}") failure_count += 1 @@ -152,7 +161,7 @@ class LocalApprovalRequestor: ] # Display duration options - print(colorText("\n⏱️ Select Local Approval Duration:", "white")) + print(colorText("\n Select Local Approval Duration:", "white")) print(colorText("=" * 50, "white")) for i, (minutes, label) in enumerate(duration_options, start=1): @@ -166,36 +175,36 @@ class LocalApprovalRequestor: if 1 <= choice <= len(duration_options): duration_minutes, duration_label = duration_options[choice - 1] - print(colorText(f"✓ Selected: {duration_label}", "green")) + print(colorText(f" Selected: {duration_label}", "green")) logger.info(f"User selected duration: {duration_minutes} minutes") else: - print(colorText("❌ Invalid choice.", "red")) + print(colorText("❌ Invalid choice.", "red")) logger.warning("Invalid duration choice") return except ValueError: - print(colorText("❌ Invalid input. Please enter a number.", "red")) + print(colorText("❌ Invalid input. Please enter a number.", "red")) logger.warning("Invalid input for duration selection") return # Select agents - print(colorText("\n🎯 Select Agents for Local Approval:", "white")) + print(colorText("\nSelect Agents for Local Approval:", "white")) agents = selectAgents(self.api) if not agents: - print(colorText("❌ No agents found or error retrieving agents.", "red")) + print(colorText("❌ No agents found or error retrieving agents.", "red")) logger.warning("No agents selected or error retrieving agents") return # Confirm with user - print(colorText("\n📋 Summary:", "cyan")) + print(colorText("\nSummary:", "cyan")) print(colorText(f" Duration: {duration_label}", "white")) print(colorText(f" Agents: {len(agents)}", "white")) confirm = get_sanitized_input("\nProceed? (y/n): ").lower() if confirm != "y": - print(colorText("❌ Operation cancelled.", "yellow")) + print(colorText("❌ Operation cancelled.", "yellow")) return # Process the batch @@ -219,24 +228,24 @@ class LocalApprovalRequestor: failure_count: Number of failed operations """ print(colorText(f"\n{'=' * 60}", "white")) - print(colorText("📊 Local Approval Summary", "cyan")) + print(colorText(" Local Approval Summary", "cyan")) print(colorText("=" * 60, "white")) - print(colorText(f"✓ Successfully processed: {success_count}", "green")) + print(colorText(f" Successfully processed: {success_count}", "green")) if failure_count > 0: - print(colorText(f"✗ Failed: {failure_count}", "red")) + print(colorText(f" Failed: {failure_count}", "red")) - print(colorText(f"\n📦 Batch ID: {batch_id}", "cyan")) - print(colorText(f"⏱️ Duration: {duration_label}", "cyan")) + print(colorText(f"\n Batch ID: {batch_id}", "cyan")) + print(colorText(f" Duration: {duration_label}", "cyan")) print(colorText("=" * 60, "white")) - print(colorText("\n💡 Next Steps:", "yellow")) - print(colorText(" • Agents have been moved to audit policies", "white")) - print(colorText(" • Local approvals are active", "white")) + print(colorText("\n Next Steps:", "yellow")) + print(colorText(" ✅ Agents have been moved to audit policies", "white")) + print(colorText(" ✅ Local approvals are active", "white")) print( colorText( - f" • Agents will return to enforcement after {duration_label}", + f" ✅ Agents will return to enforcement after {duration_label}", "white", ) ) diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py index 53f7ede..e6e6298 100644 --- a/flows/prepPolicy.py +++ b/flows/prepPolicy.py @@ -88,7 +88,7 @@ def sortHashes( ): working_dir = load_env("WORKING_DIR") history_days = Selector.select_value( - prompt="Enter how many days of history to pull (1–150): ", + prompt="Enter how many days of history to pull (1-150): ", value_type=int, valid_range=(1, 150), ) @@ -655,7 +655,7 @@ def section_header(title): def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist): working_dir = load_env("WORKING_DIR") - section_header("🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒") + section_header("Prepare to Enforce Policy") print( colorText( "\nSequentially follow these steps to prepare a policy for enforcement:", @@ -670,11 +670,11 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all ) ) if not selected_policies: - print(colorText(" [✗] No policies have been chosen", "red")) + print(colorText(" [❌] No policies have been chosen", "red")) else: print(colorText("The following policies have been chosen:", "green")) for policy in selected_policies: - print(colorText(f" [✓] {policy.name}", "green")) + print(colorText(f" [✅] {policy.name}", "green")) # Step 2: Destination Policy and Allowlist print( @@ -683,22 +683,22 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all if destination_policy: print( colorText( - f" [✓] {destination_policy[0].name} has been selected as the destination policy", + f" [✅] {destination_policy[0].name} has been selected as the destination policy", "green", ) ) else: - print(colorText(" [✗] No destination policy has been chosen", "red")) + print(colorText(" [❌] No destination policy has been chosen", "red")) if destination_allowlist: print( colorText( - f" [✓] {destination_allowlist[0].name} has been selected as allowlist", + f" [✅] {destination_allowlist[0].name} has been selected as allowlist", "green", ) ) else: - print(colorText(" [✗] No allowlist has been chosen", "red")) + print(colorText(" [❌] No allowlist has been chosen", "red")) # Step 3: Data Preparation print( @@ -713,9 +713,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all print( colorText( ( - " [✓] Data has been fetched" + " [✅] Data has been fetched" if os.path.exists(review_path) - else " [✗] Data has not been fetched" + else " [❌] Data has not been fetched" ), "green" if os.path.exists(review_path) else "red", ) @@ -723,7 +723,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all else: print( colorText( - " [✗] No policies selected, cannot check data fetch status", "red" + " [❌] No policies selected, cannot check data fetch status", "red" ) ) @@ -756,9 +756,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all print( colorText( ( - " [✓] Reviewed hashes have been loaded" + " [✅] Reviewed hashes have been loaded" if os.path.exists(approved_path) - else " [✗] Reviewed hashes have not been loaded" + else " [❌] Reviewed hashes have not been loaded" ), "green" if os.path.exists(approved_path) else "red", ) @@ -766,9 +766,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all print( colorText( ( - " [✓] Path review list created" + " [✅] Path review list created" if os.path.exists(second_review_path) - else " [✗] Path review list has not been created" + else " [❌] Path review list has not been created" ), "green" if os.path.exists(second_review_path) else "red", ) @@ -776,7 +776,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all else: print( colorText( - " [✗] No policies selected, cannot check reviewed hashes or path list", + " [❌] No policies selected, cannot check reviewed hashes or path list", "red", ) ) @@ -812,9 +812,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all print( colorText( ( - " [✓] Reviewed path list detected" + " [✅] Reviewed path list detected" if os.path.exists(reviewed_path) - else " [✗] Path review list has not been detected" + else " [❌] Path review list has not been detected" ), "green" if os.path.exists(reviewed_path) else "red", ) @@ -825,9 +825,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all print( colorText( ( - " [✓] Preflight Path Exclusion List has been generated" + " [✅] Preflight Path Exclusion List has been generated" if preflight_ready - else " [✗] Preflight Path Exclusion List has not been generated" + else " [❌] Preflight Path Exclusion List has not been generated" ), "green" if preflight_ready else "red", ) @@ -835,7 +835,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all else: print( colorText( - " [✗] No policies selected, cannot check preflight status", "red" + " [❌] No policies selected, cannot check preflight status", "red" ) ) @@ -866,5 +866,5 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all print(colorText(" Apply approved hashes to allowlist", "cyan")) # Utility Options - print(colorText("F. 📂 - Open Working Directory", "cyan")) - print(colorText("B. 🔚 - Back", "cyan")) + print(colorText("F. Open Working Directory", "cyan")) + print(colorText("B. Back", "cyan")) diff --git a/models/execution.py b/models/execution.py index 462193d..439515f 100644 --- a/models/execution.py +++ b/models/execution.py @@ -145,13 +145,13 @@ class Hash: approved_count += 1 except (ValueError, TypeError): logger.debug( - "Needs Review: Scannermatch score is missing or invalid. — {e}" + "Needs Review: Scannermatch score is missing or invalid. — {e}" ) hash_obj.at_decision = "needs_review" needs_review_count += 1 logger.debug( - f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}" + f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}" ) return hashes @@ -443,13 +443,13 @@ class ExecutionHistoryRecord: approved_count += 1 except (ValueError, TypeError) as e: logger.debug( - f"Needs Review: Scannermatch score is missing or invalid. — {e}" + f"Needs Review: Scannermatch score is missing or invalid. — {e}" ) hash_obj.at_decision = "needs_review" needs_review_count += 1 logger.debug( - f"Final counts — Needs Review: {needs_review_count}, " + f"Final counts — Needs Review: {needs_review_count}, " f"Approved: {approved_count}, Unapproved: {unapproved_count}" ) diff --git a/services/agenthandler.py b/services/agenthandler.py index 7a788e6..6d6d250 100644 --- a/services/agenthandler.py +++ b/services/agenthandler.py @@ -38,7 +38,7 @@ logger = logging.getLogger(__name__) def devicehistory(api: AirlockAPIWrapper, outputjson: bool): agents = selectAgents(api) history_days = Selector.select_value( - prompt="Enter how many days of history to pull (1–150): ", + prompt="Enter how many days of history to pull (1–150): ", value_type=int, valid_range=(1, 150), ) @@ -60,7 +60,7 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool): except Exception as e: print( colorText( - f"❌ Error retrieving history for {agent.hostname}: {e}", "red" + f"❌ Error retrieving history for {agent.hostname}: {e}", "red" ) ) continue @@ -139,7 +139,7 @@ def findAgents(api, return_dataframe): print( colorText( - f"\n✅ Matched devices exported to: {working_dir}\\{filename}", + f"\n✓ Matched devices exported to: {working_dir}\\{filename}", "green", ) ) @@ -148,7 +148,7 @@ def findAgents(api, return_dataframe): def collect_device_names() -> List[str]: - print(colorText("🔍 Device Search", "cyan")) + print(colorText("🖥��Â Device Search", "cyan")) print( colorText( "Enter the device hostnames you'd like to search for, one per line.", "cyan" @@ -185,7 +185,7 @@ def collect_device_names() -> List[str]: else: print( colorText( - f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", + f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", "yellow", ) ) @@ -235,8 +235,8 @@ def show_unmatched( ] if unmatched: - logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}") - print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow")) + logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}") + print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow")) def enrich_agents(agents: List["Agent"], policies: List["Policy"]): @@ -248,7 +248,7 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]: device_names = collect_device_names() if not device_names: logger.debug("No device names entered") - print(colorText("⚠️ No device names entered.", "red")) + print(colorText("⚠️ No device names entered.", "red")) return [] use_exact = choose_match_type() @@ -261,11 +261,11 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]: show_unmatched(device_names, matched_agents, use_exact) if not matched_agents: - logger.debug("❌ No matching devices found.") - print(colorText("❌ No matching devices found.", "red")) + logger.debug("❌ No matching devices found.") + print(colorText("❌ No matching devices found.", "red")) return [] - print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green")) + print(colorText(f"✓ Found {len(matched_agents)} matching device(s).", "green")) logger.info("Matched agent hostnames:") rows = (len(matched_agents) + 2) // 3 # 3 columns for row in range(rows): @@ -283,8 +283,8 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]: ) if not matched_agents: - logger.debug("❌ No matching devices remain after refinement.") - print(colorText("❌ No matching devices remain after refinement.", "red")) + logger.debug("❌ No matching devices remain after refinement.") + print(colorText("❌ No matching devices remain after refinement.", "red")) return [] enrich_agents(matched_agents, policies) @@ -302,7 +302,7 @@ def moveAgentToRelatedPolicy( Args: api: AirlockAPIWrapper instance. agent: Agent object. - policy_relationship_map: Dict mapping enforcement → audit. + policy_relationship_map: Dict mapping enforcement â–€ –€™ audit. mode: 'audit' to move to audit, 'enforcement' to move to enforcement. """ policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}") diff --git a/services/policyhandler.py b/services/policyhandler.py index 0fa6a11..e36fe23 100644 --- a/services/policyhandler.py +++ b/services/policyhandler.py @@ -27,9 +27,8 @@ import tqdm from models.policy import Policy from services.API import AirlockAPIWrapper -from utils.configmanager import get_system_json from utils.setup import get_base_directory -from utils.utils import areYouSure, colorText, get_sanitized_input +from utils.utils import colorText logger = logging.getLogger(__name__) @@ -233,17 +232,3 @@ def skipback(days): hex_timestamp = format(timestamp, "08x") objectid_hex = hex_timestamp + "0000000000000000" return ObjectId(objectid_hex) - - -def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper): - policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}") - for enforcement_policy, audit_policy in policy_relationship_map.items(): - api.policy_clone(enforcement_policy, audit_policy) - api.policy_set_auditmode(audit_policy, "1") - - -def confirmUpdateAfromE(api: AirlockAPIWrapper): - areYouSure() - confirmation = get_sanitized_input("Type 'I AGREE' to continue: ") - if confirmation.strip() == "I AGREE": - updateAuditPoliciesFromEnforcementPolices(api) diff --git a/utils/selector.py b/utils/selector.py index 2cac574..d43f6d8 100644 --- a/utils/selector.py +++ b/utils/selector.py @@ -184,7 +184,7 @@ class Selector: print(colorText(f"✅ Included {len(selected)} item(s).", "green")) return selected elif mode == "e": - print(colorText(f"🚫 Excluded {len(selected)} item(s).", "yellow")) + print(colorText(f"👫 Excluded {len(selected)} item(s).", "yellow")) return [item for item in items if item not in selected] else: print(colorText("⚠️ Invalid mode. Returning all items.", "yellow")) @@ -279,7 +279,9 @@ class Selector: df = df[columns] items = [row for _, row in df.iterrows()] - label_func = lambda row: str(row.to_dict()) + + def label_func(row): + return str(row.to_dict()) result = Selector._select_from_list( items, @@ -311,7 +313,9 @@ class Selector: df = df[columns] items = df.to_dict("records") - label_func = lambda row: " | ".join(str(row[col]) for col in df.columns) + + def label_func(row): + return " | ".join(str(row[col]) for col in df.columns) # Show rows first print(colorText(header, "cyan")) @@ -346,7 +350,7 @@ class Selector: print(colorText(f"✅ Included {len(selected)} row(s).", "green")) return [pd.Series(row) for row in selected] elif mode == "e": - print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow")) + print(colorText(f"👫 Excluded {len(selected)} row(s).", "yellow")) return [pd.Series(row) for row in items if row not in selected] else: print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow")) diff --git a/utils/setup.py b/utils/setup.py index ee4a384..f28a9a3 100644 --- a/utils/setup.py +++ b/utils/setup.py @@ -118,14 +118,14 @@ def setup(): logging.debug(f"{name.capitalize()} directory ensured at: {path}") # Load system config (immutable) - system_config = load_system_config() + load_system_config() # Configure logging with system-defined log level log_level = get_system_value("LOG_LEVEL", str, "INFO") configure_logging(dirs["logs"], log_level) # Load user config (mutable) - user_config = load_user_config(dirs["config"]) + load_user_config(dirs["config"]) # Set up .env file - ONLY for WORKING_DIR (runtime-configurable value) env_path = base_dir / ".env" @@ -155,6 +155,6 @@ def setup(): for subfolder in subfolders: subfolder_path = folder_path / subfolder subfolder_path.mkdir(parents=True, exist_ok=True) - logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}") + logging.debug(f"'{subfolder}' subfolder created at: {subfolder_path}") logging.info("✅ Setup complete")