feat: Complete Policy Prep Workflow with UX upgrades, Liftoff API, and TUI merge

- Added intro screen with workflow overview, time estimate, and onboarding controls
- Improved visuals: cleaner checkboxes (/), better loading screen layout
- Enforced mandatory tab reviews for critical steps with warnings and blocked navigation
- Optimized logging: INFO for milestones, DEBUG for internals; cleaner production logs
- Implemented Liftoff API integration: paths, publishers, hashes with granular error handling
- Color-coded completion feedback ( success,  failure,  partial) and detailed summaries
- Consolidated architecture: merged TUI.py into Loxide.py (single entry point, no circular imports)
- Fixed race condition in table creation with concurrency locks
This commit is contained in:
2025-12-15 17:01:58 -05:00
parent 0dbc744471
commit 57d0f12000
4 changed files with 820 additions and 512 deletions
+380 -59
View File
@@ -165,6 +165,18 @@ class PolicyPrepWorkflowScreen(Screen):
# Track if we're navigating with keyboard (to prevent selection)
self._keyboard_navigation = False
# Tab review tracking for Step 5 (First Review)
self.approved_tab_reviewed = False
self.needs_review_tab_reviewed = False
# Tab review tracking for Step 6 (Path Review)
self.paths_tab_reviewed = False
self.publishers_tab_reviewed = False
# Lock to prevent concurrent table creation
self._creating_review_table = False
self._creating_path_table = False
def compose(self) -> ComposeResult:
"""Build the UI layout for the workflow screen."""
yield Header(show_clock=True, icon="⚙️")
@@ -191,7 +203,7 @@ class PolicyPrepWorkflowScreen(Screen):
def on_mount(self) -> None:
"""Initialize the screen when mounted."""
self._update_checklist()
self._show_source_policy_selection()
self._show_introduction()
def watch_workflow_stage(self, old_value: str, new_value: str) -> None:
"""React to workflow stage changes."""
@@ -348,6 +360,86 @@ class PolicyPrepWorkflowScreen(Screen):
step8.styles.text_style = "dim"
col2.mount(step8)
def _show_introduction(self) -> None:
"""Show workflow introduction and overview."""
self.workflow_stage = "introduction"
content = self.query_one("#content_area", Vertical)
content.remove_children()
# Title
title = Static("Welcome to Policy Preparation Workflow")
title.styles.margin = (1, 1)
title.styles.text_style = "bold"
title.styles.text_align = "center"
content.mount(title)
# Description
description = Static(
"This workflow will help you:\n"
" • Fetch execution history from selected policies\n"
" • Review and approve safe executions\n"
" • Calculate efficient path exclusions\n"
" • Generate publisher trust rules\n"
" • Apply changes to your destination policy"
)
description.styles.margin = (1, 2)
content.mount(description)
# Process steps
steps_title = Static("The Process:")
steps_title.styles.margin = (1, 2, 0, 2)
steps_title.styles.text_style = "bold"
content.mount(steps_title)
steps = Static(
" 📋 Step 1: Select source policies (data collection)\n"
" 🎯 Step 2: Select destination policy (where changes go)\n"
" 📝 Step 3: Select destination allowlist\n"
" 📊 Step 4: Fetch execution data (may take 1-2 minutes)\n"
" ✅ Step 5: Review approved/needs review executions\n"
" 📁 Step 6: Review path exclusions and publishers\n"
" 🔍 Step 7: Preview changes before applying\n"
" 🚀 Step 8: Liftoff - Apply to production"
)
steps.styles.margin = (0, 2)
content.mount(steps)
# Time estimate
estimate = Static("⏱️ Estimated Time: 15-30 minutes depending on data size")
estimate.styles.margin = (1, 2)
estimate.styles.color = "cyan"
content.mount(estimate)
# Tips
tips_title = Static("💡 Tips:")
tips_title.styles.margin = (1, 2, 0, 2)
tips_title.styles.text_style = "bold"
content.mount(tips_title)
tips = Static(
" • Start with a test policy first\n"
" • Review carefully - changes affect all agents\n"
" • Use path rules when possible (more efficient)\n"
" • Publishers are powerful - use cautiously"
)
tips.styles.margin = (0, 2)
tips.styles.color = "yellow"
content.mount(tips)
# Buttons
button_container = Horizontal()
button_container.styles.margin = (2, 2)
button_container.styles.align = ("center", "middle")
content.mount(button_container)
continue_btn = Button(
"Continue to Policy Selection", id="start_workflow", variant="success"
)
cancel_btn = Button("Cancel", id="cancel_workflow", variant="default")
button_container.mount(continue_btn)
button_container.mount(cancel_btn)
def _show_source_policy_selection(self) -> None:
"""Show the source policy selection screen."""
self.workflow_stage = "select_source"
@@ -366,7 +458,7 @@ class PolicyPrepWorkflowScreen(Screen):
table.zebra_stripes = True
# Add columns - checkbox first, then data columns
table.add_columns("", "Name", "ID", "Parent")
table.add_columns("", "Name", "ID", "Parent")
# Sort policies by name for easier selection
sorted_policies = sorted(self.policies, key=lambda p: p.name.lower())
@@ -376,7 +468,7 @@ class PolicyPrepWorkflowScreen(Screen):
# Skip parent policies
if policy.parent == "global-policy-settings":
continue
checkbox = "" # All start unchecked
checkbox = "" # All start unchecked
table.add_row(
checkbox,
policy.name,
@@ -739,7 +831,7 @@ class PolicyPrepWorkflowScreen(Screen):
def _show_fetch_results(self) -> None:
"""Show the results of data fetching."""
logger.info("=== _show_fetch_results called ===")
logger.debug("=== _show_fetch_results called ===")
self.workflow_stage = "first_review"
content = self.query_one("#content_area", Vertical)
content.remove_children()
@@ -792,15 +884,39 @@ class PolicyPrepWorkflowScreen(Screen):
logger.info("Mounted tab buttons")
# Show approved table by default
logger.info("About to call _show_review_table('approved')")
logger.debug("About to call _show_review_table('approved')")
self._show_review_table("approved")
logger.info("=== _show_fetch_results complete ===")
logger.debug("=== _show_fetch_results complete ===")
def _show_review_table(self, table_type: str) -> None:
"""Show an editable DataTable for reviewing executions."""
logger.info(f"=== _show_review_table called with type: {table_type} ===")
# Prevent concurrent execution
if self._creating_review_table:
logger.warning(
f"Already creating review table, ignoring duplicate call for {table_type}"
)
return
self._creating_review_table = True
try:
self._show_review_table_impl(table_type)
finally:
self._creating_review_table = False
def _show_review_table_impl(self, table_type: str) -> None:
"""Internal implementation of _show_review_table."""
logger.debug(f"_show_review_table called with type: {table_type}")
content = self.query_one("#content_area", Vertical)
# Mark tab as reviewed
if table_type == "approved":
self.approved_tab_reviewed = True
logger.debug("Marked approved tab as reviewed")
else:
self.needs_review_tab_reviewed = True
logger.debug("Marked needs_review tab as reviewed")
# Determine which dataframe and table ID to show
if table_type == "approved":
df = self.approved_df
@@ -841,6 +957,18 @@ class PolicyPrepWorkflowScreen(Screen):
except Exception as e:
logger.debug(f"Error removing existing tables: {e}")
# Remove existing instruction and help text (they accumulate without removal)
# Remove ALL Static widgets - they're just text that needs to be replaced
try:
existing_statics = content.query("Static")
logger.debug(
f"Found {len(existing_statics)} existing Static widgets to remove"
)
for static in existing_statics:
static.remove()
except Exception as e:
logger.debug(f"Error removing Static widgets: {e}")
# Force a refresh to ensure removals are processed
try:
content.refresh()
@@ -850,7 +978,7 @@ class PolicyPrepWorkflowScreen(Screen):
# Note: We no longer remove review_controls or review_continue_container
# They are reused between tabs to avoid DuplicateIds errors
logger.info(
logger.debug(
f"DataFrame for {table_type}: {'empty' if df is None or df.empty else f'{len(df)} rows'}"
)
@@ -861,13 +989,13 @@ class PolicyPrepWorkflowScreen(Screen):
logger.info(f"No data for {table_type}, mounted empty message")
return
# Instructions
# Instructions (no ID needed - we remove all Statics anyway)
instruction = Static(title)
instruction.styles.margin = (1, 1)
instruction.styles.text_style = "bold"
content.mount(instruction)
# Help text
# Help text (no ID needed)
help_text = Static(
"Click to toggle, 'r' for range select (click start, press 'r', click end)\n"
"Space to toggle cursor row, 'd' to delete, 'a' select all, arrows navigate"
@@ -876,6 +1004,19 @@ class PolicyPrepWorkflowScreen(Screen):
help_text.styles.text_style = "dim"
content.mount(help_text)
# CRITICAL: Check if table already exists in content (should not happen after removal above)
try:
existing_check = content.query_one(f"#{table_id}", DataTable)
if existing_check:
logger.error(
f"Table {table_id} STILL EXISTS after removal! This should not happen."
)
# Don't create a new one - just return
return
except Exception:
# Good - table doesn't exist, proceed with creation
pass
# Create the review table
review_table = DataTable(id=table_id)
review_table.styles.height = "40vh" # Increased since we removed button rows
@@ -899,15 +1040,15 @@ class PolicyPrepWorkflowScreen(Screen):
if available_cols:
# Add checkbox column first
review_table.add_columns("", *available_cols)
review_table.add_columns("", *available_cols)
# Add rows with row keys for tracking
for idx, row in df.iterrows():
checkbox = "" # All start unchecked
checkbox = "" # All start unchecked
row_data = [str(row.get(col, "")) for col in available_cols]
review_table.add_row(checkbox, *row_data, key=str(idx))
logger.info(f"About to mount {table_id}")
logger.debug(f"About to mount {table_id}")
# Final safety check - make sure no table with this ID exists before mounting
try:
@@ -923,7 +1064,7 @@ class PolicyPrepWorkflowScreen(Screen):
pass
content.mount(review_table)
logger.info(f"Successfully mounted {table_id} with {len(df)} rows")
logger.debug(f"Successfully mounted {table_id} with {len(df)} rows")
# Row count display only (removed Select All, Clear, Delete buttons)
try:
@@ -1044,13 +1185,19 @@ class PolicyPrepWorkflowScreen(Screen):
content = self.query_one("#content_area", Vertical)
content.remove_children()
# Clear message
# Add spacer to push text to bottom
spacer = Static("")
spacer.styles.height = "1fr"
content.mount(spacer)
# Loading message at bottom (above checklist)
status = Static(
"Building path exclusions and publisher lists...\n\n"
"Building path exclusions and publisher lists...\n"
"This may take a moment for large datasets."
)
status.styles.margin = (2, 1)
status.styles.margin = (1, 1)
status.styles.text_align = "center"
status.styles.color = "cyan"
content.mount(status)
# Force UI refresh to show the loading screen
@@ -1311,8 +1458,8 @@ class PolicyPrepWorkflowScreen(Screen):
logger.info(
f"=== _calculate_paths called with path_exclusion_constant={path_exclusion_constant} ==="
)
logger.info(f"Input DataFrame: {len(df)} rows")
logger.info(f"Columns: {list(df.columns) if not df.empty else 'empty'}")
logger.debug(f"Input DataFrame: {len(df)} rows")
logger.debug(f"Columns: {list(df.columns) if not df.empty else 'empty'}")
if df.empty:
logger.warning("Input DataFrame is empty")
@@ -1327,7 +1474,7 @@ class PolicyPrepWorkflowScreen(Screen):
haslcp = self._split_filepaths_grouped(df, path_exclusion_constant, "filename")
haslcp = haslcp.drop_duplicates()
logger.info(f"After split_filepaths_grouped: {len(haslcp)} rows")
logger.debug(f"After split_filepaths_grouped: {len(haslcp)} rows")
# Filter forbidden paths
badpathparts = get_system_list("BAD_PATH_PARTS")
@@ -1337,7 +1484,7 @@ class PolicyPrepWorkflowScreen(Screen):
forbidden_pattern, case=False, na=False, regex=True
)
logger.info(
logger.debug(
f"Removing forbidden filepaths: {forbidden_lcfp.sum()} paths filtered"
)
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
@@ -1347,7 +1494,7 @@ class PolicyPrepWorkflowScreen(Screen):
)
lcp_not_forbidden = haslcp.copy()
logger.info(f"After forbidden filtering: {len(lcp_not_forbidden)} rows")
logger.debug(f"After forbidden filtering: {len(lcp_not_forbidden)} rows")
# Select relevant columns
if "policyname" in lcp_not_forbidden.columns:
@@ -1399,11 +1546,11 @@ class PolicyPrepWorkflowScreen(Screen):
lcp_not_forbidden_review = lcp_not_forbidden_review[
lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path
]
logger.info(
logger.debug(
f"After MIN_FILES_FOR_PATH filter ({min_files_for_path}): {len(lcp_not_forbidden_review)} rows (removed {before_filter - len(lcp_not_forbidden_review)})"
)
logger.info(
logger.debug(
f"Final result: {len(lcp_not_forbidden_review)} rows with columns: {list(lcp_not_forbidden_review.columns)}"
)
@@ -1441,7 +1588,7 @@ class PolicyPrepWorkflowScreen(Screen):
def _show_path_results(self) -> None:
"""Show the results of path building."""
logger.info("=== _show_path_results called ===")
logger.debug("=== _show_path_results called ===")
self.workflow_stage = "second_review"
content = self.query_one("#content_area", Vertical)
content.remove_children()
@@ -1490,15 +1637,39 @@ class PolicyPrepWorkflowScreen(Screen):
logger.info("Mounted tab buttons")
# Show paths table by default
logger.info("About to call _show_path_review_table('paths')")
logger.debug("About to call _show_path_review_table('paths')")
self._show_path_review_table("paths")
logger.info("=== _show_path_results complete ===")
logger.debug("=== _show_path_results complete ===")
def _show_path_review_table(self, table_type: str) -> None:
"""Show an editable DataTable for reviewing paths/publishers."""
logger.info(f"=== _show_path_review_table called with type: {table_type} ===")
# Prevent concurrent execution
if self._creating_path_table:
logger.warning(
f"Already creating path table, ignoring duplicate call for {table_type}"
)
return
self._creating_path_table = True
try:
self._show_path_review_table_impl(table_type)
finally:
self._creating_path_table = False
def _show_path_review_table_impl(self, table_type: str) -> None:
"""Internal implementation of _show_path_review_table."""
logger.debug(f"_show_path_review_table called with type: {table_type}")
content = self.query_one("#content_area", Vertical)
# Mark tab as reviewed (only paths and publishers, not remaining)
if table_type == "paths":
self.paths_tab_reviewed = True
logger.debug("Marked paths tab as reviewed")
elif table_type == "publishers":
self.publishers_tab_reviewed = True
logger.debug("Marked publishers tab as reviewed")
# Determine which dataframe to show
if table_type == "paths":
# Combine primary and secondary paths for review
@@ -1608,6 +1779,18 @@ class PolicyPrepWorkflowScreen(Screen):
except Exception as e:
logger.debug(f"Error removing existing controls: {e}")
# Remove existing instruction and help text (they accumulate without removal)
# Remove ALL Static widgets - they're just text that needs to be replaced
try:
existing_statics = content.query("Static")
logger.debug(
f"Found {len(existing_statics)} existing Static widgets to remove"
)
for static in existing_statics:
static.remove()
except Exception as e:
logger.debug(f"Error removing Static widgets: {e}")
# Force refresh to ensure removals complete
try:
content.refresh()
@@ -1620,13 +1803,13 @@ class PolicyPrepWorkflowScreen(Screen):
content.mount(empty_msg)
return
# Instructions
# Instructions (no ID needed)
instruction = Static(title)
instruction.styles.margin = (1, 1)
instruction.styles.text_style = "bold"
content.mount(instruction)
# Help text (different for remaining hashes)
# Help text (different for remaining hashes, no ID needed)
if table_type != "remaining":
help_text = Static(
"Click to toggle, 'r' for range select (click start, press 'r', click end)\n"
@@ -1653,11 +1836,11 @@ class PolicyPrepWorkflowScreen(Screen):
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)
review_table.add_columns("", *available_cols)
# Add rows with row keys for tracking
for idx, row in df.iterrows():
checkbox = "" # All start unchecked
checkbox = "" # All start unchecked
row_data = []
for col in available_cols:
value = row.get(col, "")
@@ -1669,7 +1852,7 @@ class PolicyPrepWorkflowScreen(Screen):
row_data.append(str(value))
review_table.add_row(checkbox, *row_data, key=str(idx))
logger.info(f"About to mount {table_id}")
logger.debug(f"About to mount {table_id}")
# Final safety check before mounting table
try:
@@ -1682,7 +1865,7 @@ class PolicyPrepWorkflowScreen(Screen):
pass
content.mount(review_table)
logger.info(f"Successfully mounted {table_id}")
logger.debug(f"Successfully mounted {table_id}")
# Row count display only (removed Select All, Clear, Delete buttons)
if table_type != "remaining":
@@ -1999,53 +2182,166 @@ class PolicyPrepWorkflowScreen(Screen):
"""Perform the actual application of changes."""
try:
results = []
errors = []
# 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 path exclusions to policy (primary + secondary)
if self.destination_policy:
path_rules = []
# Process primary paths
if (
self.primary_paths_df is not None
and not self.primary_paths_df.empty
):
logger.info(
f"Processing {len(self.primary_paths_df)} primary paths"
)
for _, row in self.primary_paths_df.groupby(
["longestcfp", "file_extension"]
):
path = row.iloc[0]["longestcfp"]
ext = row.iloc[0]["file_extension"]
# Format: C:\Path\**.ext
path_rule = f"{path}\\**{ext}"
path_rules.append(path_rule)
# Process secondary paths
if (
self.secondary_paths_df is not None
and not self.secondary_paths_df.empty
):
logger.info(
f"Processing {len(self.secondary_paths_df)} secondary paths"
)
for _, row in self.secondary_paths_df.groupby(
["longestcfp", "file_extension"]
):
path = row.iloc[0]["longestcfp"]
ext = row.iloc[0]["file_extension"]
path_rule = f"{path}\\**{ext}"
path_rules.append(path_rule)
# Apply path rules to policy
if path_rules:
try:
logger.info(
f"Applying {len(path_rules)} path exclusions to policy {self.destination_policy.name}"
)
response = self.api.policy_add_path_exclusions(
str(self.destination_policy.groupid), path_rules
)
results.append(
f"✓ Added {len(path_rules)} path exclusions to policy"
)
logger.info(f"Path exclusions applied successfully: {response}")
except Exception as e:
error_msg = f"✗ Failed to add path exclusions: {str(e)}"
errors.append(error_msg)
logger.error(error_msg, exc_info=True)
# Apply publishers to policy
if self.destination_policy and self.publishers_df is not None:
# This would call the actual API methods
results.append("Applied approved publishers to policy")
if (
self.destination_policy
and self.publishers_df is not None
and not self.publishers_df.empty
):
try:
publishers = self.publishers_df["publisher"].unique().tolist()
logger.info(
f"Applying {len(publishers)} publishers to policy {self.destination_policy.name}"
)
response = self.api.policy_add_publishers(
str(self.destination_policy.groupid), publishers
)
results.append(
f"✓ Added {len(publishers)} trusted publishers to policy"
)
logger.info(f"Publishers applied successfully: {response}")
except Exception as e:
error_msg = f"✗ Failed to add publishers: {str(e)}"
errors.append(error_msg)
logger.error(error_msg, exc_info=True)
# Apply hashes to allowlist
if self.destination_allowlist and self.approved_df is not None:
# This would call the actual API methods
results.append("Applied approved hashes to allowlist")
if (
self.destination_allowlist
and self.approved_df is not None
and not self.approved_df.empty
):
try:
# Get unique hashes
hashes = self.approved_df["sha256"].unique().tolist()
logger.info(
f"Applying {len(hashes)} hashes to allowlist {self.destination_allowlist.name}"
)
response = self.api.hash_add_to_allowlist(
str(self.destination_allowlist.applicationid), hashes
)
results.append(
f"✓ Added {len(hashes):,} approved hashes to allowlist"
)
logger.info(f"Hashes applied successfully: {response}")
except Exception as e:
error_msg = f"✗ Failed to add hashes: {str(e)}"
errors.append(error_msg)
logger.error(error_msg, exc_info=True)
self._show_completion(results)
# Show completion with both results and errors
all_results = results + errors
self._show_completion(all_results, has_errors=len(errors) > 0)
except Exception as e:
logger.error(f"Failed to apply changes: {e}", exc_info=True)
self.app.notify(f"Failed to apply changes: {str(e)}", severity="error")
logger.error(f"Critical failure in _perform_apply: {e}", exc_info=True)
self.app.notify(f"Critical failure: {str(e)}", severity="error")
self._show_test_screen()
def _show_completion(self, results: List[str]) -> None:
def _show_completion(self, results: List[str], has_errors: bool = False) -> None:
"""Show completion screen."""
self.workflow_stage = "complete"
content = self.query_one("#content_area", Vertical)
content.remove_children()
summary = Static(
"Policy Preparation Complete!\n\n"
"The following changes have been applied:"
)
# Title depends on whether there were errors
if has_errors:
title_text = "Policy Preparation Completed with Errors\n\n" "Results:"
title_color = "yellow"
else:
title_text = (
"Policy Preparation Complete!\n\n"
"The following changes have been applied:"
)
title_color = "green"
summary = Static(title_text)
summary.styles.margin = (1, 1)
summary.styles.text_style = "bold"
summary.styles.color = title_color
content.mount(summary)
for result in results:
result_widget = Static(f" {result}")
result_widget.styles.margin = (0, 2)
# Color based on success/failure
if result.startswith(""):
result_widget.styles.color = "green"
elif result.startswith(""):
result_widget.styles.color = "red"
content.mount(result_widget)
# Final message
final = Static(
f"\nPolicy '{self.destination_policy.name}' is now ready for enforcement!"
)
final.styles.margin = (2, 1)
final.styles.color = "green"
if has_errors:
final = Static(
f"\n⚠️ Policy '{self.destination_policy.name}' was partially updated.\n"
"Please review errors above and retry failed operations manually."
)
final.styles.margin = (2, 1)
final.styles.color = "yellow"
else:
final = Static(
f"\n✅ Policy '{self.destination_policy.name}' is now ready for enforcement!"
)
final.styles.margin = (2, 1)
final.styles.color = "green"
content.mount(final)
# Done button
@@ -2078,7 +2374,7 @@ class PolicyPrepWorkflowScreen(Screen):
)
# Determine if this row should be checked
is_selected = row_key_str in selected_keys
checkbox = "☑️" if is_selected else ""
checkbox = "" if is_selected else ""
# Update the checkbox cell (first column, index 0)
try:
@@ -2315,8 +2611,15 @@ class PolicyPrepWorkflowScreen(Screen):
"""Handle button presses."""
button_id = event.button.id
# Introduction screen buttons
if button_id == "start_workflow":
self._show_source_policy_selection()
elif button_id == "cancel_workflow":
self.app.pop_screen()
# Source policy selection buttons
if button_id == "select_none_source":
elif button_id == "select_none_source":
self.selected_source_policy_ids.clear()
# Refresh checkbox display
self._refresh_table_checkboxes(
@@ -2434,6 +2737,15 @@ class PolicyPrepWorkflowScreen(Screen):
)
elif button_id == "continue_from_review":
# Check if both tabs have been reviewed
if not self.approved_tab_reviewed or not self.needs_review_tab_reviewed:
self.app.notify(
"Please review both 'Approved' and 'Needs Review' tabs before continuing.",
severity="warning",
timeout=5,
)
return
# Validate that review is complete
if (self.approved_df is None or self.approved_df.empty) and (
self.needs_review_df is None or self.needs_review_df.empty
@@ -2449,6 +2761,15 @@ class PolicyPrepWorkflowScreen(Screen):
self._show_path_building_screen()
elif button_id == "build_preflight":
# Check if both required tabs have been reviewed
if not self.paths_tab_reviewed or not self.publishers_tab_reviewed:
self.app.notify(
"Please review both 'Paths' and 'Publishers' tabs before continuing.",
severity="warning",
timeout=5,
)
return
# Validate that path review is complete
if (self.primary_paths_df is None or self.primary_paths_df.empty) and (
self.publishers_df is None or self.publishers_df.empty