mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Restructured TUI, expanded quietagent workflow

Browse Source
This commit is contained in:
Zarithas
2025-11-14 17:07:49 -05:00
parent 0ebb42dcbd
commit 3ee762a0a1
20 changed files with 1770 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Loxide.py
+1 -1
View File
@@ -30,8 +30,8 @@ import urllib3
from services.API import AirlockAPIWrapper
from services.security import getAPI
from TUI.TUI import run_Loxide
from utils.setup import get_base_directory, setup
from utils.TUI import run_Loxide
from utils.utils import irtang
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
Loxide_Icon.ico
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 112 KiB

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