617 lines
22 KiB
Python
617 lines
22 KiB
Python
from __future__ import annotations
|
||
|
||
import logging
|
||
from typing import Optional
|
||
|
||
import pandas as pd
|
||
from textual.app import ComposeResult
|
||
from textual.binding import Binding
|
||
from textual.containers import Horizontal, Vertical
|
||
from textual.screen import Screen
|
||
from textual.widgets import (
|
||
Button,
|
||
DataTable,
|
||
Footer,
|
||
Header,
|
||
Static,
|
||
TextArea,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class AllowlistSelectionWidget(Static):
|
||
"""
|
||
Widget for selecting an allowlist and adding hashes to it.
|
||
Can be reused in different workflows.
|
||
"""
|
||
|
||
DEFAULT_CSS = """
|
||
AllowlistSelectionWidget {
|
||
height: 1fr;
|
||
layout: vertical;
|
||
}
|
||
#allowlist_main {
|
||
height: 100%;
|
||
width: 100%;
|
||
}
|
||
#left_panel {
|
||
width: 50%;
|
||
padding: 1;
|
||
border: solid $primary;
|
||
}
|
||
#right_panel {
|
||
width: 50%;
|
||
padding: 1;
|
||
border: solid $primary;
|
||
}
|
||
#allowlist_table {
|
||
height: 70%;
|
||
margin: 1 0;
|
||
}
|
||
#allowlist_table > .datatable--header {
|
||
text-style: bold;
|
||
background: $boost;
|
||
}
|
||
#allowlist_table Row {
|
||
height: 1;
|
||
}
|
||
#preview_area {
|
||
height: 60%;
|
||
margin: 1 0;
|
||
}
|
||
#action_buttons {
|
||
height: 10%;
|
||
padding: 1;
|
||
content-align: center middle;
|
||
}
|
||
.panel-title {
|
||
text-style: bold;
|
||
margin: 0 0 1 0;
|
||
}
|
||
.info-text {
|
||
margin: 1 0;
|
||
}
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
selected_data: pd.DataFrame,
|
||
api=None,
|
||
hostname: Optional[str] = None,
|
||
otpid: Optional[str] = None,
|
||
hash_column: str = "sha256", # Default hash column name
|
||
):
|
||
"""
|
||
Initialize the allowlist selection widget.
|
||
|
||
Args:
|
||
selected_data: DataFrame containing the selected activities
|
||
api: API instance for making allowlist calls
|
||
hostname: Optional hostname for context
|
||
otpid: Optional OTP ID for context
|
||
hash_column: Name of the column containing hashes (default: "sha256")
|
||
"""
|
||
super().__init__()
|
||
self.selected_data = selected_data
|
||
self.api = api
|
||
self.hostname = hostname
|
||
self.otpid = otpid
|
||
self.hash_column = hash_column
|
||
self.allowlists = []
|
||
self.selected_allowlist = None
|
||
self.hashes_to_add = []
|
||
|
||
def compose(self) -> ComposeResult:
|
||
with Horizontal(id="allowlist_main"):
|
||
# Left panel - Allowlist selection
|
||
with Vertical(id="left_panel"):
|
||
yield Static("Select Allowlist", classes="panel-title")
|
||
yield Static(
|
||
f"Choose an allowlist to add {len(self.selected_data)} selected items",
|
||
classes="info-text",
|
||
)
|
||
|
||
# Allowlist table
|
||
self.allowlist_table = DataTable(id="allowlist_table")
|
||
self.allowlist_table.cursor_type = "row"
|
||
yield self.allowlist_table
|
||
|
||
# Refresh button
|
||
self.refresh_btn = Button(
|
||
"🔄 Refresh Allowlists", id="refresh_allowlists_btn"
|
||
)
|
||
yield self.refresh_btn
|
||
|
||
# Right panel - Preview and actions
|
||
with Vertical(id="right_panel"):
|
||
yield Static("Preview", classes="panel-title")
|
||
|
||
# Context information
|
||
context_text = []
|
||
if self.hostname:
|
||
context_text.append(f"Host: {self.hostname}")
|
||
if self.otpid:
|
||
context_text.append(f"OTP: {self.otpid}")
|
||
context_text.append(f"Selected Activities: {len(self.selected_data)}")
|
||
|
||
yield Static(" | ".join(context_text), classes="info-text")
|
||
|
||
# Preview text area
|
||
self.preview_area = TextArea(
|
||
id="preview_area", read_only=True, language="markdown"
|
||
)
|
||
yield self.preview_area
|
||
|
||
# Hash statistics
|
||
self.stats_label = Static("", id="stats_label", classes="info-text")
|
||
yield self.stats_label
|
||
|
||
# Action buttons at bottom
|
||
with Horizontal(id="action_buttons"):
|
||
self.back_btn = Button("← Back", id="back_btn")
|
||
self.add_btn = Button("➕ Add to Allowlist", id="add_to_allowlist_btn")
|
||
|
||
self.back_btn.styles.width = "50%"
|
||
self.add_btn.styles.width = "50%"
|
||
self.add_btn.disabled = True # Disabled until allowlist selected
|
||
|
||
yield self.back_btn
|
||
yield self.add_btn
|
||
|
||
async def on_mount(self) -> None:
|
||
"""Load allowlists when widget mounts."""
|
||
await self.load_allowlists()
|
||
await self.extract_and_preview_hashes()
|
||
|
||
async def load_allowlists(self) -> None:
|
||
"""Load available allowlists from API, grouped by policy association."""
|
||
if not self.api:
|
||
logger.error("No API available")
|
||
self.allowlist_table.add_column("Error")
|
||
self.allowlist_table.add_row("No API available")
|
||
return
|
||
|
||
try:
|
||
# First, try to get the host's policy if hostname is provided
|
||
host_policy_allowlists = []
|
||
host_policy_ids = set()
|
||
policy_name = None
|
||
|
||
if self.hostname:
|
||
try:
|
||
# Get agent info to find its policy
|
||
agents_df = self.api.agent_find_by_hostname(self.hostname)
|
||
if not agents_df.empty:
|
||
# Get the policy group ID for this host
|
||
group_id = agents_df.iloc[0].get("groupid")
|
||
policy_name = agents_df.iloc[0].get(
|
||
"groupname", "Unknown Policy"
|
||
)
|
||
|
||
if group_id:
|
||
# Get allowlists for this policy
|
||
policy_allowlists_df = self.api.policy_list_allowlists(
|
||
group_id
|
||
)
|
||
if not policy_allowlists_df.empty:
|
||
host_policy_allowlists = policy_allowlists_df.to_dict(
|
||
orient="records"
|
||
)
|
||
host_policy_ids = {
|
||
al.get("applicationid")
|
||
for al in host_policy_allowlists
|
||
}
|
||
logger.info(
|
||
f"Found {len(host_policy_allowlists)} allowlists for host's policy"
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Could not get host's policy allowlists: {e}")
|
||
|
||
# Get all allowlists
|
||
all_allowlists_df = self.api.allowlist_find_all()
|
||
|
||
if all_allowlists_df.empty:
|
||
self.allowlist_table.add_column("No Allowlists")
|
||
self.allowlist_table.add_row("No allowlists found")
|
||
return
|
||
|
||
all_allowlists = all_allowlists_df.to_dict(orient="records")
|
||
|
||
# Separate into two groups: policy-associated and others
|
||
other_allowlists = [
|
||
al
|
||
for al in all_allowlists
|
||
if al.get("applicationid") not in host_policy_ids
|
||
]
|
||
|
||
# Sort each group alphabetically by name
|
||
host_policy_allowlists.sort(key=lambda x: x.get("name", "").lower())
|
||
other_allowlists.sort(key=lambda x: x.get("name", "").lower())
|
||
|
||
# Combine lists with policy-associated first
|
||
self.allowlists = host_policy_allowlists + other_allowlists
|
||
|
||
# Setup table columns
|
||
self.allowlist_table.clear()
|
||
self.allowlist_table.add_columns("Name", "Application ID", "Type")
|
||
|
||
# Track which rows are headers vs actual allowlists
|
||
self._row_to_allowlist_map = {}
|
||
current_row = 0
|
||
|
||
# Add policy-associated allowlists if any
|
||
if host_policy_allowlists:
|
||
# Add section header
|
||
header_text = f"━━━ Policy: {policy_name or 'Host Policy'} ━━━"
|
||
self.allowlist_table.add_row(header_text, "", "", key="header_policy")
|
||
current_row += 1
|
||
|
||
# Add policy allowlists
|
||
for idx, allowlist in enumerate(host_policy_allowlists):
|
||
name = allowlist.get("name", "Unknown")
|
||
app_id = allowlist.get("applicationid", "Unknown")
|
||
|
||
self.allowlist_table.add_row(
|
||
f" {name}", # Indent to show grouping
|
||
app_id,
|
||
"Policy",
|
||
key=f"policy_{idx}",
|
||
)
|
||
self._row_to_allowlist_map[current_row] = idx
|
||
current_row += 1
|
||
|
||
# Add other allowlists
|
||
if other_allowlists:
|
||
# Add section header
|
||
if host_policy_allowlists:
|
||
# Add spacer if we have policy allowlists above
|
||
self.allowlist_table.add_row("", "", "", key="spacer")
|
||
current_row += 1
|
||
|
||
self.allowlist_table.add_row(
|
||
"━━━ Other Available Allowlists ━━━", "", "", key="header_other"
|
||
)
|
||
current_row += 1
|
||
|
||
# Add other allowlists
|
||
for idx, allowlist in enumerate(other_allowlists):
|
||
name = allowlist.get("name", "Unknown")
|
||
app_id = allowlist.get("applicationid", "Unknown")
|
||
|
||
self.allowlist_table.add_row(
|
||
f" {name}", # Indent to show grouping
|
||
app_id,
|
||
"General",
|
||
key=f"other_{idx}",
|
||
)
|
||
# Map to the correct index in the combined list
|
||
actual_idx = len(host_policy_allowlists) + idx
|
||
self._row_to_allowlist_map[current_row] = actual_idx
|
||
current_row += 1
|
||
|
||
# Log summary
|
||
logger.info(
|
||
f"Loaded {len(self.allowlists)} total allowlists: "
|
||
f"{len(host_policy_allowlists)} policy-associated, "
|
||
f"{len(other_allowlists)} others"
|
||
)
|
||
|
||
# Update stats label if no allowlists in policy
|
||
if self.hostname and not host_policy_allowlists:
|
||
self.stats_label.update(
|
||
f"Note: No allowlists found for {self.hostname}'s policy | "
|
||
+ self.stats_label.content.plain
|
||
)
|
||
|
||
except Exception as exc:
|
||
logger.exception(f"Failed to load allowlists: {exc}")
|
||
self.allowlist_table.add_column("Error")
|
||
self.allowlist_table.add_row(f"Failed to load: {str(exc)}")
|
||
|
||
async def extract_and_preview_hashes(self) -> None:
|
||
"""Extract hashes from selected data and show preview."""
|
||
preview_lines = ["## Hash Extraction Summary\n"]
|
||
|
||
# Check for hash column
|
||
if self.hash_column not in self.selected_data.columns:
|
||
# Try to find a hash column
|
||
possible_hash_cols = [
|
||
"sha256",
|
||
"SHA256",
|
||
"hash",
|
||
"Hash",
|
||
"sha1",
|
||
"SHA1",
|
||
"md5",
|
||
"MD5",
|
||
"filehash",
|
||
"file_hash",
|
||
]
|
||
found_col = None
|
||
for col in possible_hash_cols:
|
||
if col in self.selected_data.columns:
|
||
found_col = col
|
||
break
|
||
|
||
if found_col:
|
||
self.hash_column = found_col
|
||
preview_lines.append(f"✓ Found hash column: **{found_col}**\n")
|
||
else:
|
||
preview_lines.append("⚠️ **No hash column found**\n")
|
||
preview_lines.append("Available columns:\n")
|
||
for col in self.selected_data.columns:
|
||
if col != "_row_id":
|
||
preview_lines.append(f" - {col}\n")
|
||
|
||
self.preview_area.text = "".join(preview_lines)
|
||
self.stats_label.update("No hashes to add")
|
||
return
|
||
|
||
# Extract unique hashes
|
||
hashes = self.selected_data[self.hash_column].dropna().unique()
|
||
self.hashes_to_add = [h for h in hashes if h and str(h).strip()]
|
||
|
||
# Build preview
|
||
preview_lines.append(f"### Found {len(self.hashes_to_add)} unique hashes\n\n")
|
||
|
||
# Show sample of hashes (first 10)
|
||
preview_lines.append("**Sample hashes to be added:**\n```\n")
|
||
for i, hash_val in enumerate(self.hashes_to_add[:10]):
|
||
preview_lines.append(f"{i+1}. {hash_val}\n")
|
||
if len(self.hashes_to_add) > 10:
|
||
preview_lines.append(f"... and {len(self.hashes_to_add) - 10} more\n")
|
||
preview_lines.append("```\n\n")
|
||
|
||
# Show sample of source data
|
||
preview_lines.append("**Sample source activities:**\n")
|
||
sample_cols = [
|
||
col
|
||
for col in self.selected_data.columns
|
||
if col not in ["_row_id"] and col in ["filename", "path", "action", "user"]
|
||
]
|
||
if not sample_cols:
|
||
sample_cols = [
|
||
col for col in self.selected_data.columns if col != "_row_id"
|
||
][:3]
|
||
|
||
if sample_cols:
|
||
preview_lines.append("```\n")
|
||
for i, row in self.selected_data[sample_cols].head(5).iterrows():
|
||
row_text = " | ".join([f"{col}: {row[col]}" for col in sample_cols])
|
||
preview_lines.append(f"{row_text}\n")
|
||
preview_lines.append("```\n")
|
||
|
||
self.preview_area.text = "".join(preview_lines)
|
||
|
||
# Update statistics
|
||
self.stats_label.update(
|
||
f"Ready to add {len(self.hashes_to_add)} unique hashes | "
|
||
f"From {len(self.selected_data)} selected activities"
|
||
)
|
||
|
||
async def on_data_table_row_selected(self, event) -> None:
|
||
"""Handle allowlist selection."""
|
||
try:
|
||
# Extract row index from event - handle different event structures
|
||
row_index = None
|
||
|
||
# Try to get row index from coordinate
|
||
if hasattr(event, "coordinate") and hasattr(event.coordinate, "row"):
|
||
row_index = event.coordinate.row
|
||
# Try cursor_row as fallback
|
||
elif hasattr(event, "cursor_row"):
|
||
row_index = event.cursor_row
|
||
# Try getting from the table itself
|
||
else:
|
||
table = self.allowlist_table
|
||
if hasattr(table, "cursor_row"):
|
||
row_index = table.cursor_row
|
||
|
||
# Validate row index
|
||
if row_index is not None and isinstance(row_index, int):
|
||
# Account for group headers in the row count
|
||
actual_allowlist_index = self._get_allowlist_index_from_row(row_index)
|
||
|
||
if (
|
||
actual_allowlist_index is not None
|
||
and 0 <= actual_allowlist_index < len(self.allowlists)
|
||
):
|
||
self.selected_allowlist = self.allowlists[actual_allowlist_index]
|
||
self.add_btn.disabled = False
|
||
self.add_btn.label = (
|
||
f"➕ Add to '{self.selected_allowlist.get('name', 'Unknown')}'"
|
||
)
|
||
|
||
# Update preview with selection
|
||
await self._update_preview_with_selection()
|
||
|
||
logger.info(
|
||
f"Selected allowlist: {self.selected_allowlist.get('name')}"
|
||
)
|
||
else:
|
||
logger.debug(f"Row {row_index} is a header or invalid")
|
||
else:
|
||
logger.warning(f"Could not extract valid row index from event: {event}")
|
||
|
||
except Exception as exc:
|
||
logger.exception(f"Failed to select allowlist: {exc}")
|
||
|
||
def _get_allowlist_index_from_row(self, row_index: int) -> Optional[int]:
|
||
"""Convert table row index to allowlist list index, accounting for group headers."""
|
||
# This will be updated when we have group headers
|
||
if hasattr(self, "_row_to_allowlist_map"):
|
||
return self._row_to_allowlist_map.get(row_index)
|
||
return row_index
|
||
|
||
async def _update_preview_with_selection(self) -> None:
|
||
"""Update preview when an allowlist is selected."""
|
||
if not self.selected_allowlist:
|
||
return
|
||
|
||
current_text = self.preview_area.text
|
||
# Remove any existing selection header
|
||
if "### Selected Allowlist:" in current_text:
|
||
lines = current_text.split("\n")
|
||
# Find and remove the selection lines
|
||
new_lines = []
|
||
skip_next = False
|
||
for line in lines:
|
||
if line.startswith("### Selected Allowlist:"):
|
||
skip_next = True
|
||
continue
|
||
if skip_next and line.startswith("Application ID:"):
|
||
skip_next = False
|
||
continue
|
||
if not skip_next:
|
||
new_lines.append(line)
|
||
current_text = "\n".join(new_lines)
|
||
|
||
# Add new selection at the top
|
||
selection_text = (
|
||
f"### Selected Allowlist: **{self.selected_allowlist.get('name')}**\n"
|
||
f"Application ID: {self.selected_allowlist.get('applicationid')}\n\n"
|
||
)
|
||
self.preview_area.text = selection_text + current_text
|
||
|
||
async def on_button_pressed(self, event) -> None:
|
||
"""Handle button presses."""
|
||
btn = getattr(event, "button", None) or getattr(event, "sender", None)
|
||
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None)
|
||
|
||
if btn is self.back_btn or btn_id == "back_btn":
|
||
await self.app.pop_screen()
|
||
event.stop()
|
||
return
|
||
|
||
if btn is self.refresh_btn or btn_id == "refresh_allowlists_btn":
|
||
await self.load_allowlists()
|
||
event.stop()
|
||
return
|
||
|
||
if btn is self.add_btn or btn_id == "add_to_allowlist_btn":
|
||
await self.add_hashes_to_allowlist()
|
||
event.stop()
|
||
return
|
||
|
||
async def add_hashes_to_allowlist(self) -> None:
|
||
"""Add the extracted hashes to the selected allowlist."""
|
||
if not self.selected_allowlist or not self.hashes_to_add:
|
||
self.app.notify(
|
||
"No allowlist selected or no hashes to add", severity="warning"
|
||
)
|
||
return
|
||
|
||
if not self.api:
|
||
self.app.notify("API not available", severity="error")
|
||
return
|
||
|
||
try:
|
||
# Disable button during operation
|
||
self.add_btn.disabled = True
|
||
self.add_btn.label = "⏳ Adding hashes..."
|
||
|
||
# Call API to add hashes
|
||
app_id = self.selected_allowlist.get("applicationid")
|
||
allowlist_name = self.selected_allowlist.get("name", "Unknown")
|
||
|
||
logger.info(
|
||
f"Adding {len(self.hashes_to_add)} hashes to allowlist {allowlist_name} (ID: {app_id})"
|
||
)
|
||
|
||
result = self.api.hash_add_to_allowlist(app_id, self.hashes_to_add)
|
||
|
||
# Success notification
|
||
self.app.notify(
|
||
f"✅ Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'",
|
||
title="Success",
|
||
severity="information",
|
||
timeout=5,
|
||
)
|
||
|
||
# Update preview to show success
|
||
self.preview_area.text = (
|
||
f"## ✅ SUCCESS\n\n"
|
||
f"Added **{len(self.hashes_to_add)} hashes** to allowlist:\n"
|
||
f"**{allowlist_name}** (ID: {app_id})\n\n"
|
||
f"### Operation Details:\n"
|
||
f"- Source: {self.hostname or 'Multiple hosts'}\n"
|
||
f"- OTP ID: {self.otpid or 'N/A'}\n"
|
||
f"- Activities processed: {len(self.selected_data)}\n"
|
||
f"- Unique hashes added: {len(self.hashes_to_add)}\n"
|
||
)
|
||
|
||
# Change button to "Done"
|
||
self.add_btn.label = "✅ Done - Close"
|
||
self.add_btn.disabled = False
|
||
|
||
# When clicked again, close the screen
|
||
self.add_btn_success = True
|
||
|
||
except Exception as exc:
|
||
logger.exception(f"Failed to add hashes to allowlist: {exc}")
|
||
self.app.notify(
|
||
f"❌ Failed to add hashes: {str(exc)}",
|
||
title="Error",
|
||
severity="error",
|
||
timeout=10,
|
||
)
|
||
|
||
# Re-enable button
|
||
self.add_btn.disabled = False
|
||
self.add_btn.label = "➕ Retry Add to Allowlist"
|
||
|
||
|
||
class AllowlistSelectionScreen(Screen):
|
||
"""
|
||
Screen wrapper for the AllowlistSelectionWidget.
|
||
"""
|
||
|
||
BINDINGS = [
|
||
Binding("b", "back", "Back"),
|
||
Binding("r", "refresh", "Refresh Allowlists"),
|
||
Binding("enter", "confirm", "Add to Allowlist"),
|
||
]
|
||
|
||
def __init__(
|
||
self,
|
||
selected_data: pd.DataFrame,
|
||
api=None,
|
||
hostname: Optional[str] = None,
|
||
otpid: Optional[str] = None,
|
||
hash_column: str = "sha256",
|
||
):
|
||
super().__init__()
|
||
self.selected_data = selected_data
|
||
self.api = api
|
||
self.hostname = hostname
|
||
self.otpid = otpid
|
||
self.hash_column = hash_column
|
||
|
||
def compose(self) -> ComposeResult:
|
||
yield Header(show_clock=True)
|
||
self.widget = AllowlistSelectionWidget(
|
||
self.selected_data,
|
||
api=self.api,
|
||
hostname=self.hostname,
|
||
otpid=self.otpid,
|
||
hash_column=self.hash_column,
|
||
)
|
||
yield self.widget
|
||
yield Footer()
|
||
|
||
async def action_back(self) -> None:
|
||
"""Go back to previous screen."""
|
||
await self.app.pop_screen()
|
||
|
||
async def action_refresh(self) -> None:
|
||
"""Refresh the allowlists."""
|
||
if hasattr(self, "widget") and self.widget:
|
||
await self.widget.load_allowlists()
|
||
|
||
async def action_confirm(self) -> None:
|
||
"""Confirm and add to allowlist."""
|
||
if hasattr(self, "widget") and self.widget:
|
||
if self.widget.selected_allowlist and self.widget.hashes_to_add:
|
||
await self.widget.add_hashes_to_allowlist()
|