751 lines
29 KiB
Python
751 lines
29 KiB
Python
"""
|
|
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
|
|
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, Header, Static, TextArea
|
|
|
|
from models.agent import Agent
|
|
from screens.otpworkflowscreen import OTPWorkflowScreen
|
|
from screens.policyselectorscreen import PolicySelectorScreen
|
|
from 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
|
|
|
|
back_button = Button("← Back", id="back_button")
|
|
back_button.styles.width = "50%"
|
|
back_button.styles.margin = (0, 1, 1, 0)
|
|
yield back_button
|
|
|
|
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:
|
|
- back_button: Pop this screen (return to parent)
|
|
- 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 == "back_button":
|
|
while len(self.app.screen_stack) > 2:
|
|
self.app.pop_screen()
|
|
event.stop()
|
|
|
|
elif 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 = []
|
|
unsuccessful = []
|
|
status_label = self.query_one("#status_label", Static)
|
|
status_label.update("Exporting CSV...")
|
|
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_protected_json
|
|
|
|
policy_relationship_map = get_protected_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}")
|
|
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
|
|
result = 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.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,
|
|
)
|
|
)
|