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

Checking WIP Agent Movement Workflow

Browse Source
This commit is contained in:
Zarithas
2025-11-09 21:06:07 -05:00
parent ed07c79b74
commit 5510d22cbd
17 changed files with 1868 additions and 626 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/selector.py
+3 -3
View File
@@ -33,7 +33,7 @@ class Selector:
def _display_choices(
items: List[Any],
label_func: Callable[[Any], str],
num_columns: int = 4,
num_columns: int = 3,
header: str = "Available Choices:",
) -> None:
# Force single column if items are DataFrame rows
@@ -54,7 +54,7 @@ class Selector:
@staticmethod
def _display_selected_items(
selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 4
selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 3
) -> None:
print(colorText("\nCurrent selections:", "cyan"))
if not selected:
@@ -93,7 +93,7 @@ class Selector:
allow_multiple: bool = False,
prompt_each: bool = False,
header: str = "Available Choices:",
num_columns: int = 4,
num_columns: int = 3,
) -> Union[Optional[Any], List[Any]]:
if not items:
logger.warning("No items available for selection.")
utils/tui.py
+60
View File
@@ -23,6 +23,7 @@ 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.otpworkflowscreen import OTPWorkflowScreen
from services.agenthandler import findAgents, moveAgents, toggleEnforcement
from services.API import AirlockAPIWrapper
@@ -30,9 +31,12 @@ from services.policyhandler import confirmUpdateAfromE
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.retro_terminal_theme import get_retro_terminal_theme
from widgets.themeselector import ThemeSelector
dotenv.load_dotenv()
@@ -106,6 +110,7 @@ class MainMenuScreen(Screen):
("🔇 - Find Quiet Hosts", "find_quiet_button"),
],
"move": [
("🔄 - Move Agent Workflow", "move_agent_workflow_button"),
("✅ - Move to local approval", "move_local_button"),
("🔄 - Move to Audit/Enforcement", "move_audit_button"),
("🔀 - Move - Other", "move_other_button"),
@@ -261,6 +266,47 @@ class MainMenuScreen(Screen):
self.app.exit()
def on_agent_move_operations_operation_complete(
self, message: AgentMoveOperations.OperationComplete
) -> None:
"""Handle completion of agent move operation - show results."""
logger.info(
"Agent move operation completed: %s, %d successful, %d unsuccessful",
message.operation,
len(message.successful),
len(message.unsuccessful),
)
# Format results for display
successful_text = "\n".join(
[f"{agent.hostname}" for agent, _ in message.successful]
)
unsuccessful_text = "\n".join(
[f"{agent.hostname}: {error}" for agent, error in message.unsuccessful]
)
# Remove the operations widget
try:
ops_widget = self.query_one(AgentMoveOperations)
ops_widget.remove()
except Exception:
pass
# Show results
self.query_one("#content", Vertical).mount(
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
)
def on_results_display_go_back(self, message: ResultsDisplay.GoBack) -> None:
"""Handle back button from results display."""
try:
results_widget = self.query_one(ResultsDisplay)
results_widget.remove()
except Exception:
pass
# Return to main menu
self.app.pop_screen()
def on_directory_tree_file_selected(
self, event: DirectoryTree.FileSelected
) -> None:
@@ -282,6 +328,11 @@ class MainMenuScreen(Screen):
_PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {})
case "find_quiet_button":
_PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {})
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 "move_local_button":
_PENDING_JOB = (
"legacy",
@@ -349,12 +400,21 @@ class Loxide(App):
self.devices = [
Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()
]
# Enrich agents with policy information
if self.policies and self.devices:
for agent in self.devices:
agent.enrich_with_policies(self.policies)
logger.debug(
f"Enriched {len(self.devices)} agents with policy information"
)
except Exception as exc:
logger.error("Failed to load policies/devices: %s", exc)
self.policies = None
self.devices = None
def on_mount(self, api: AirlockAPIWrapper) -> None:
self.register_theme(get_retro_terminal_theme())
self.theme = self._textual_theme
self.push_screen(MainMenuScreen(api))
utils/utils.py
-38
View File
@@ -151,44 +151,6 @@ def irtang():
)
def displayIntro():
print(
colorText(
r"""
_____ .__ .__ __ ___________ .__
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
\/ \/ \/ \/
""",
"cyan",
)
)
def welcome():
print(
colorText(
"=================================================================================",
"cyan",
)
)
print(
colorText(
"======================== Welcome to the Airlock API Tool ========================",
"cyan",
)
)
print(
colorText(
"=================================================================================",
"cyan",
)
)
def section_header(title):
print(
colorText(