1823 lines
68 KiB
Python
1823 lines
68 KiB
Python
# 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 <https://www.gnu.org/licenses/>.
|
|
|
|
|
|
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
|