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

Major step towards unification of the UI Implementation of the Back Feature, splitting of TUI files back into subfolders

Browse Source
This commit is contained in:
Zarithas
2025-12-02 16:22:43 -05:00
parent 1f06404a16
commit 1d9caadaf3
29 changed files with 3761 additions and 561 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TUI/Widgets/agentmoveoperations.py
+740
View File
@@ -0,0 +1,740 @@
# 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/>.
from dataclasses import asdict
from datetime import datetime
import logging
import os
from typing import List
import pandas as pd
from textual.containers import Horizontal, Vertical
from textual.css.query import NoMatches
from textual.message import Message
from textual.reactive import reactive
from textual.widget import Widget
from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea
from models.agent import Agent
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
from TUI.Screens.policyselectorscreen import PolicySelectorScreen
from TUI.Widgets.OTP_generate import OTPGenerator
logger = logging.getLogger(__name__)
class AgentMoveOperations(Widget):
"""
A Textual widget for managing bulk agent operations and policy migrations.
This widget provides a comprehensive UI for performing operations on multiple
selected agents. It displays the list of target agents and provides buttons to
trigger various bulk operations like toggling policy modes or enabling local approval.
The widget manages its own state through reactive properties and provides real-time
feedback on operation progress and results. Operations are executed sequentially
per agent with error handling that tracks both successful and failed operations.
Attributes:
operation_in_progress (reactive[bool]): Tracks whether an operation is currently
executing. Used to disable buttons during execution.
selected_operation (reactive[str]): Tracks which operation type is currently
selected or in progress (e.g., "local_approval", "toggle_enforcement").
Example:
```python
agents = [agent1, agent2, agent3]
widget = AgentMoveOperations(agents)
```
"""
# Reactive property to track if an operation is in progress
operation_in_progress = reactive(False)
# Tracks the currently selected operation type
selected_operation = reactive("")
class OperationComplete(Message):
"""
Message posted when a bulk operation completes.
This message is broadcast to parent widgets/screens to notify them of
operation completion along with detailed results. It contains the list
of agents that were processed and the outcome for each.
Attributes:
operation (str): Name of the operation that completed (e.g., "Local Approval Mode").
agents (List[Agent]): List of all agents that were targeted by the operation.
successful (List[tuple]): List of (Agent, result_data) tuples for successfully
processed agents. Result data varies by operation type.
unsuccessful (List[tuple]): List of (Agent, error_message) tuples for agents
where the operation failed. Error message is a string explaining the failure.
"""
def __init__(
self,
operation: str,
agents: List[Agent],
successful: List[tuple],
unsuccessful: List[tuple],
):
super().__init__()
self.operation = operation
self.agents = agents
self.successful = successful # List of (agent, result) tuples
self.unsuccessful = unsuccessful # List of (agent, error) tuples
def __init__(self, agents: List[Agent]):
"""
Initialize the AgentMoveOperations widget.
Args:
agents (List[Agent]): List of Agent objects to perform operations on.
These agents will be displayed in the widget's agent table.
"""
super().__init__()
self.agents = agents
def watch_operation_in_progress(self, old_value: bool, new_value: bool) -> None:
"""
React to changes in the operation_in_progress reactive property.
This is called automatically by Textual when operation_in_progress changes.
It updates the button states to reflect whether an operation is running.
Args:
old_value (bool): Previous value of operation_in_progress.
new_value (bool): New value of operation_in_progress.
"""
self._update_button_states()
def _update_button_states(self) -> None:
"""
Update the enabled/disabled state of operation buttons based on current status.
This method implements the following logic:
- If an operation is in progress: disable all buttons
- If an operation is selected: disable only that operation's button
- If no operation is selected: enable all buttons
The state transitions prevent users from starting multiple operations
simultaneously and provide visual feedback on which operation is active.
Handles NoMatches exceptions gracefully in case buttons are not yet rendered.
"""
try:
export_csv_btn = self.query_one("#export_csv_btn", Button)
local_approval_btn = self.query_one("#local_approval_btn", Button)
toggle_enforcement_btn = self.query_one("#toggle_enforcement_btn", Button)
other_policy_btn = self.query_one("#other_policy_btn", Button)
otp_gen_btn = self.query_one("#otp_gen_btn", Button)
# If operation in progress, disable all
if self.operation_in_progress:
otp_gen_btn = True
export_csv_btn.disabled = True
local_approval_btn.disabled = True
toggle_enforcement_btn.disabled = True
other_policy_btn.disabled = True
else:
# If an operation was selected, disable
if self.selected_operation:
otp_gen_btn.disabled = self.selected_operation == "otp_gen"
export_csv_btn.disabled = self.selected_operation == "export_csv"
local_approval_btn.disabled = (
self.selected_operation == "local_approval"
)
toggle_enforcement_btn.disabled = (
self.selected_operation == "toggle_enforcement"
)
other_policy_btn.disabled = (
self.selected_operation == "other_policy"
)
else:
# Enable all buttons
otp_gen_btn = False
export_csv_btn = False
local_approval_btn.disabled = False
toggle_enforcement_btn.disabled = False
other_policy_btn.disabled = False
except NoMatches:
pass
def _display_results(
self, operation_name: str, successful: list, unsuccessful: list
) -> None:
"""
Display operation results in the results text area.
Formats the results into a human-readable summary including:
- Operation name and separator
- List of successful operations with agent hostnames
- List of failed operations with agent hostnames and error messages
- Summary statistics (total successful/failed count)
The results are displayed in the results_text TextArea widget and the
results container is made visible after being initially hidden.
Args:
operation_name (str): Human-readable name of the operation (e.g., "Local Approval Mode").
successful (list): List of (Agent, result_data) tuples for successful operations.
unsuccessful (list): List of (Agent, error_message) tuples for failed operations.
"""
try:
# Build results text
results_lines = [
f"Operation: {operation_name}",
f"{'=' * 50}",
"",
f"✅ Successful ({len(successful)}):",
]
if successful:
for agent, result in successful:
results_lines.append(f" ✅ {agent.hostname}")
else:
results_lines.append(" (none)")
results_lines.append("")
results_lines.append(f"❌ Failed ({len(unsuccessful)}):")
if unsuccessful:
for agent, error in unsuccessful:
results_lines.append(f" ❌ {agent.hostname}: {error}")
else:
results_lines.append(" (none)")
results_lines.append("")
results_lines.append(f"{'=' * 50}")
results_lines.append(
f"Total: {len(successful)} successful, {len(unsuccessful)} failed"
)
results_text_widget = self.query_one("#results_text", TextArea)
results_text_widget.text = "\n".join(results_lines)
# Show results container
results_container = self.query_one("#results_container", Vertical)
results_container.styles.display = "block"
except Exception as e:
logger.error(f"Error displaying results: {e}")
def compose(self):
"""
Build the UI layout for the AgentMoveOperations widget.
This method is called by Textual to create the widget's UI structure.
It builds a two-column layout with:
- Left side: Agent table showing selected agents and their current policies
- Right side: Operation buttons and results display area
- Bottom: Navigation buttons (Back)
The layout is responsive with:
- Agent table: 2/3 width
- Operations panel: 1/3 width
- Results area: Initially hidden, shown after operation completion
"""
yield Header(show_clock=True, icon="⚙️")
title_text = Static(
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
id="move_ops_title",
)
title_text.styles.margin = (0, 0, 1, 0)
yield title_text
with Horizontal() as main_layout:
main_layout.styles.height = "auto"
# Left side - Agent list
with Vertical() as left_side:
left_side.styles.width = "3fr"
left_side.styles.height = "auto"
agents_label = Static("Selected Agents:")
agents_label.styles.margin = (0, 0, 0, 0)
yield agents_label
# Create a DataTable to show agents with their current policies
agent_table = DataTable(id="agent_table")
agent_table.styles.height = "1fr"
agent_table.styles.margin = (1, 0, 1, 0)
yield agent_table
# Right side - Operation buttons
with Vertical() as right_side:
right_side.styles.width = "2fr"
right_side.styles.margin = (0, 1, 0, 1)
right_side.styles.height = "auto"
operations_label = Static("Operations:")
operations_label.styles.margin = (0, 0, 1, 0)
yield operations_label
# Operation buttons
export_csv_btn = Button("📄 Export CSV", id="export_csv_btn")
export_csv_btn.styles.width = "100%"
export_csv_btn.styles.margin = (0, 0, 1, 0)
yield export_csv_btn
local_approval_btn = Button(
"✔️ Local Approval Mode", id="local_approval_btn"
)
local_approval_btn.styles.width = "100%"
local_approval_btn.styles.margin = (0, 0, 1, 0)
yield local_approval_btn
otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn")
otp_gen_btn.styles.width = "100%"
otp_gen_btn.styles.margin = (0, 0, 1, 0)
yield otp_gen_btn
toggle_enforcement_btn = Button(
"🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn"
)
toggle_enforcement_btn.styles.width = "100%"
toggle_enforcement_btn.styles.margin = (0, 0, 1, 0)
yield toggle_enforcement_btn
other_policy_btn = Button(
"🔀 Move to Other Policy", id="other_policy_btn"
)
other_policy_btn.styles.width = "100%"
other_policy_btn.styles.margin = (0, 0, 1, 0)
yield other_policy_btn
# Status label
status_label = Static("", id="status_label")
status_label.styles.margin = (2, 0, 0, 0)
yield status_label
yield Footer()
def on_mount(self) -> None:
"""
Initialize widget after it has been mounted on the screen.
This Textual lifecycle method is called after the widget is added to the DOM.
It performs initialization tasks:
- Populates the agent table with columns for Hostname, Policy, and Status
- Adds rows to the table for each agent in self.agents
- Initializes button states based on current widget state
The agent table displays agent.hostname, agent.groupname (or "Unknown"),
and agent.status_text (or "Unknown") for each agent.
"""
table = self.query_one("#agent_table", DataTable)
table.add_columns("Hostname", "Current Policy", "Status")
for agent in self.agents:
table.add_row(
agent.hostname,
agent.groupname or "Unknown",
agent.status_text or "Unknown",
)
self._update_button_states()
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
"""Handle OTP generation request - call the actual OTP generation function."""
def on_button_pressed(self, event: Button.Pressed):
"""
Handle button press events from the widget.
This Textual event handler routes button presses to appropriate actions:
- copy_results_btn: Copy results text to clipboard (requires pyperclip)
- local_approval_btn: Start local approval operation
- toggle_enforcement_btn: Start toggle audit/enforcement operation
- other_policy_btn: Start move to other policy operation
After handling, event.stop() is called to prevent event propagation.
Args:
event (Button.Pressed): The button press event containing the button reference.
"""
btn_id = event.button.id
if btn_id == "copy_results_btn":
try:
results_text = self.query_one("#results_text", TextArea)
import pyperclip
pyperclip.copy(results_text.text)
self.app.notify(
"📋✅ Results copied to clipboard!",
severity="information",
timeout=2,
)
except ImportError:
self.app.notify(
"❌ pyperclip not installed. Run: pip install pyperclip",
severity="warning",
)
except Exception as e:
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
event.stop()
elif btn_id == "export_csv_btn":
self._start_export_csv_operation()
event.stop()
elif btn_id == "local_approval_btn":
self._start_local_approval_operation()
event.stop()
elif btn_id == "toggle_enforcement_btn":
self._start_toggle_enforcement_operation()
event.stop()
elif btn_id == "other_policy_btn":
self._start_other_policy_operation()
event.stop()
elif btn_id == "otp_gen_btn":
self._start_OTP_gen_operation()
event.stop()
def _start_local_approval_operation(self) -> None:
"""
Execute the local approval mode operation on all selected agents.
This operation performs the following steps for each agent:
1. Generate a unique batch ID (current Unix timestamp)
2. Create a local approval OTP with default duration of 360 minutes (6 hours)
3. Move the agent to its related audit policy mode
The operation:
- Sets operation state flags (selected_operation, operation_in_progress)
- Updates the status label with progress indicator
- Iterates through all agents, tracking successful and unsuccessful operations
- Displays formatted results via _display_results()
- Posts an OperationComplete message for parent widget handling
Agents that fail are logged and added to the unsuccessful list with error details.
The operation completes and returns to a non-busy state regardless of individual
agent success/failure.
Note: The OTP duration (360 minutes) is currently hardcoded and could be
made configurable in future versions.
"""
self.selected_operation = "local_approval"
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("✔️ Moving agents to local approval...")
# Get API from app
api = self.app.api
successful = []
unsuccessful = []
try:
import time
from services.agenthandler import moveAgentToRelatedPolicy
# Generate batch ID
batch = int(time.time())
duration = 360 # Default 6 hours, could make this configurable
for agent in self.agents:
try:
# Add local approval OTP
addLocalApproval(api, batch, duration, agent.agentid)
# Move to audit mode
result = moveAgentToRelatedPolicy(api, agent, "audit")
successful.append((agent, result))
logger.info(
f"Successfully moved {agent.hostname} to local approval"
)
except Exception as e:
unsuccessful.append((agent, str(e)))
logger.error(
f"Failed to move {agent.hostname} to local approval: {e}"
)
except Exception as e:
logger.error(f"Error during local approval operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
return
self.operation_in_progress = False
status_label.update("✅ Operation complete!")
# Display results in the widget
self._display_results("Local Approval Mode", successful, unsuccessful)
# Also post message for potential parent handling
self.post_message(
self.OperationComplete(
"Local Approval Mode", self.agents, successful, unsuccessful
)
)
def _start_export_csv_operation(self) -> None:
self.selected_operation = "export_csv"
self.operation_in_progress = True
successful = []
status_label = self.query_one("#status_label", Static)
status_label.update("Exporting CSV...")
self.app.refresh_data()
agents = self.agents
policies = self.app.policies
path = self.app.working_dir
try:
# Enrich each agent with policies and status text
for agent in agents:
agent.enrich_with_policies(policies)
# Convert each Agent to a dictionary, including all fields
data = []
for agent in agents:
row = asdict(agent)
# Remove the class-level status_map from the row
row.pop("status_map", None)
data.append(row)
# Create DataFrame
df = pd.DataFrame(data)
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
filename = f"agentsearch_{timestamp}.csv"
file_path = os.path.join(str(path), filename)
df.to_csv(file_path, index=False)
successful.append(file_path)
status_label.update(f"✅ Exported to {file_path}")
except Exception:
status_label.update("❌ Failed")
self.operation_in_progress = False
"""
# Display results in the widget
self._display_results("CSV Export", successful, unsuccessful)
# Also post message for potential parent handling
self.post_message(
self.OperationComplete(
"CSV Export", self.agents, successful, unsuccessful
)
)
"""
def _start_toggle_enforcement_operation(self) -> None:
"""
Toggle agents between enforcement and audit policy modes.
This operation intelligently switches each agent between enforcement and
audit modes based on its current state:
- If agent.groupid is in POLICY_MAP_ENF_AUD: currently enforcing , move to audit
- Otherwise: currently in audit, move to enforcement
The operation:
- Retrieves the enforcement/audit policy relationship map from protected config
- Sets operation state flags and updates status label
- Iterates through agents, determining current mode and toggling to opposite
- Tracks successful toggles with the new mode in the result message
- Logs both successes and failures
- Displays results and posts OperationComplete message
The policy relationship map (POLICY_MAP_ENF_AUD) must be present in protected
configuration and maps enforcement policy IDs to audit policy IDs. If the map
is empty or not found, all agents are assumed to be in audit mode and will
be moved to enforcement.
Returns to a non-busy state after completion regardless of individual results.
"""
self.selected_operation = "toggle_enforcement"
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("🔄 Toggling enforcement mode...")
# Get API from app
api = self.app.api
successful = []
unsuccessful = []
try:
from services.agenthandler import moveAgentToRelatedPolicy
from utils.configmanager import get_system_json
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
for agent in self.agents:
try:
# Determine current mode and toggle
if agent.groupid in policy_relationship_map:
# Currently in enforcement, move to audit
result = moveAgentToRelatedPolicy(api, agent, "audit")
mode = "audit"
else:
# Currently in audit, move to enforcement
result = moveAgentToRelatedPolicy(api, agent, "enforcement")
mode = "enforcement"
successful.append((agent, f"Moved to {mode}: {result}"))
logger.info(f"Successfully toggled {agent.hostname} to {mode}")
self.app.refresh_data()
except Exception as e:
unsuccessful.append((agent, str(e)))
logger.error(f"Failed to toggle {agent.hostname}: {e}")
except Exception as e:
logger.error(f"Error during toggle enforcement operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
return
self.operation_in_progress = False
status_label.update("✅ Operation complete!")
# Display results in the widget
self._display_results("Toggle Audit/Enforcement", successful, unsuccessful)
# Also post message for potential parent handling
self.post_message(
self.OperationComplete(
"Toggle Audit/Enforcement", self.agents, successful, unsuccessful
)
)
def _start_other_policy_operation(self) -> None:
"""
Move agents to a user-selected policy (currently unimplemented).
This operation is intended to allow bulk movement of selected agents to any
alternative policy via a policy selection dialog. Currently, this feature
is not fully implemented.
Planned Implementation:
1. Push a new policy selector screen (TUI modal/overlay)
2. Allow user to choose target policy from available options
3. Move all selected agents to the chosen policy
4. Display results like other operations
Current Behavior:
- Sets selected_operation to "other_policy"
- Displays "Policy selection not yet implemented" status message
- Clears selected_operation without performing any action
TODO: Complete implementation by:
- Creating a policy selector screen component
- Implementing the policy selection logic
- Integrating with moveAgentToPolicy API call
- Adding proper result tracking and display
"""
self.selected_operation = "other_policy"
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("Loading available policies...")
try:
# Fetch all policies from API
api = self.app.api
# Fetch all available policies
all_policies_df = api.policy_find_all()
if all_policies_df.empty:
status_label.update("No policies available")
self.operation_in_progress = False
self.selected_operation = ""
return
# Create and push the policy selector screen
policy_selector_screen = PolicySelectorScreen(
policies=all_policies_df,
agent_move_operations=self,
)
self.app.push_screen(policy_selector_screen)
except Exception as e:
logger.error(f"Error loading policies: {e}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
self.selected_operation = ""
self.app.notify(f"Failed to load policies: {str(e)}", severity="error")
def _start_OTP_gen_operation(self) -> None:
status_label = self.query_one("#status_label", Static)
status_label.update("Generating OTP.")
self.app.push_screen(OTPWorkflowScreen(self.agents))
def _execute_move_to_policy(self, target_policy) -> None:
"""
Execute the actual move of agents to the selected policy.
Moves each agent sequentially to the target policy, tracking success/failure.
Updates the status label and displays results upon completion.
Args:
target_policy: The Policy object selected by the user.
"""
status_label = self.query_one("#status_label", Static)
status_label.update(f"Moving agents to {target_policy.name}...")
api = self.app.api
successful = []
unsuccessful = []
try:
for agent in self.agents:
try:
# Move agent to target policy
api.agent_move(agent.agentid, target_policy.groupid)
successful.append((agent, f"Moved to {target_policy.name}"))
logger.info(
f"Successfully moved {agent.hostname} to policy {target_policy.name}"
)
except Exception as e:
unsuccessful.append((agent, str(e)))
logger.error(
f"Failed to move {agent.hostname} to policy {target_policy.name}: {e}"
)
except Exception as e:
logger.error(f"Error during move to policy operation: {e}")
status_label.update(f"Error: {str(e)}")
self.operation_in_progress = False
return
self.app.refresh_data()
self.operation_in_progress = False
status_label.update("Operation complete!")
# Display results in the widget
self._display_results(
f"Move to {target_policy.name}",
successful,
unsuccessful,
)
# Also post message for potential parent handling
self.post_message(
self.OperationComplete(
f"Move to {target_policy.name}",
self.agents,
successful,
unsuccessful,
)
)