RustImplementation #28

Merged
mysticmomba merged 37 commits from RustImplementation into master 2025-11-18 14:51:42 -05:00
10 changed files with 188 additions and 159 deletions
Showing only changes of commit d9cbce6175 - Show all commits
+1 -1
View File
@@ -30,7 +30,7 @@ def history_logging(
checkpoint_number: str,
policy_names: str,
) -> List[Dict[str, Any]]:
"""
"""
Query execution history logs from the Airlock API.
Parameters
+8 -25
View File
@@ -1,41 +1,24 @@
# otp_workflow_screen.py
from typing import List
from textual.app import ComposeResult
from textual.screen import Screen
from models.agent import Agent
from widgets.multiagentselector import MultiAgentSelector
from widgets.OTP_generate import OTPGenerator
class OTPWorkflowScreen(Screen):
"""Screen that handles the OTP generation workflow."""
"""Screen that handles the OTP generation workflow without agent selection."""
def __init__(self, all_agents: List[Agent]):
def __init__(self, selected_agents: List[Agent]):
super().__init__()
self.all_agents = all_agents
self.selected_devices = None
self.selected_agents = selected_agents
def compose(self) -> ComposeResult:
"""Start with the multi-agent selector."""
yield MultiAgentSelector(self.all_agents)
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
"""Handle selected agents - switch to OTP generator."""
self.selected_devices = message.selected_agents
# Remove the MultiAgentSelector
selector = self.query_one(MultiAgentSelector)
selector.remove()
# Mount the OTPGenerator with the selected Agent objects
# No need to pass API - it will access self.app.api directly
self.mount(OTPGenerator(self.selected_devices))
"""Directly show the OTP generator for the selected agents."""
yield OTPGenerator(self.selected_agents)
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
"""Handle OTP generation request - call the actual OTP generation function."""
# This will be handled by the main app, but we can also do it here
# For now, just pass it up to the app level
pass
"""Handle OTP generation request - pass it up to the app level if needed."""
+14 -37
View File
@@ -20,18 +20,17 @@ from textual.widgets import (
from flows.otp import otp_activities_by_agent, 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.otpworkflowscreen import OTPWorkflowScreen
from services.agenthandler import findAgents, moveAgents, toggleEnforcement
from services.API import AirlockAPIWrapper
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.amber_terminal_theme import get_amber_terminal_theme
from widgets.multiagentselector import MultiAgentSelector
from widgets.OTP_generate import OTPGenerator
from widgets.policytreewidget import PolicyTreeWidget
@@ -105,24 +104,18 @@ class MainMenuScreen(Screen):
current_tab = reactive("")
BUTTON_DEFS = {
"find": [
("🔍 - Device Search", "find_device_button"),
"agent_actions": [
(
"🖥️ - Find, Move, or Generate OTP for Agents",
"move_agent_workflow_button",
),
("🔇 - 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"),
],
"otp": [
("🎫 - Generate OTPs", "otp_generate_button"),
("📊 - OTP Activities By Agent", "otp_activities_button"),
("❌ - Revoke OTPs", "otp_revoke_button"),
],
"policy": [
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
("🔄 - Update Audit Policies", "policy_audit_update_button"),
("📊 - OTP Activities By Agent", "otp_activities_button"),
("❌ - Revoke OTPs", "otp_revoke_button"),
],
}
@@ -148,23 +141,21 @@ class MainMenuScreen(Screen):
yield Header(show_clock=True, icon="")
tabs = [
Tab("Policy Tree", id="p_tree"),
Tab("Device Search", id="find"),
Tab("Move Agent", id="move"),
Tab("OTP", id="otp"),
Tab("Tree View", id="p_tree"),
Tab("Agents", id="agent_actions"),
Tab("Directory", id="dir"),
Tab("Settings", id="settings"),
]
if self.extras == "POLICYPREP":
tabs.insert(3, Tab("Policy Prep", id="policy"))
tabs.insert(2, Tab("Policy Prep", id="policy"))
yield Tabs(*tabs, id="tabs")
yield Vertical(id="content")
yield Footer()
def on_mount(self) -> None:
self.switch_tab("find")
self.switch_tab("agent_actions")
# focus helpers
def _get_content_buttons(self) -> list[Button]:
@@ -225,7 +216,7 @@ class MainMenuScreen(Screen):
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
"""Handle selected agents from MultiAgentSelector."""
"""Handle selected agents from AgentSelector."""
global _PENDING_JOB
selected_agents = message.selected_agents
logger.info("Selected agents: %s", selected_agents)
@@ -324,26 +315,11 @@ class MainMenuScreen(Screen):
logger.debug("Button pressed: %s", button_id)
match button_id:
case "find_device_button":
_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",
print,
("Move to local approval (placeholder)",),
{},
)
case "move_audit_button":
_PENDING_JOB = ("legacy", toggleEnforcement, (self.app.api,), {})
case "move_other_button":
_PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {})
case "otp_generate_button":
# NEW: Push OTP workflow screen instead of legacy function
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
@@ -415,6 +391,7 @@ class Loxide(App):
def on_mount(self, api: AirlockAPIWrapper) -> None:
self.register_theme(get_retro_terminal_theme())
self.register_theme(get_amber_terminal_theme())
self.theme = self._textual_theme
self.push_screen(MainMenuScreen(api))
+4 -1
View File
@@ -194,7 +194,10 @@ class OTPGenerator(Widget):
btn_id = event.button.id
if btn_id == "back_button":
self.app.pop_screen()
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "copy_clipboard_button":
+114 -68
View File
@@ -18,9 +18,13 @@ Dependencies:
- 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
@@ -29,7 +33,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
logger = logging.getLogger(__name__)
@@ -133,18 +139,24 @@ class AgentMoveOperations(Widget):
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, keep it disabled, enable others
# 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"
)
@@ -156,6 +168,8 @@ class AgentMoveOperations(Widget):
)
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
@@ -189,21 +203,21 @@ class AgentMoveOperations(Widget):
f"Operation: {operation_name}",
f"{'=' * 50}",
"",
f"✅ Successful ({len(successful)}):",
f" Successful ({len(successful)}):",
]
if successful:
for agent, result in successful:
results_lines.append(f" • {agent.hostname}")
results_lines.append(f" {agent.hostname}")
else:
results_lines.append(" (none)")
results_lines.append("")
results_lines.append(f"❌ Failed ({len(unsuccessful)}):")
results_lines.append(f" Failed ({len(unsuccessful)}):")
if unsuccessful:
for agent, error in unsuccessful:
results_lines.append(f" • {agent.hostname}: {error}")
results_lines.append(f" {agent.hostname}: {error}")
else:
results_lines.append(" (none)")
@@ -231,7 +245,7 @@ class AgentMoveOperations(Widget):
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, Reset)
- Bottom: Navigation buttons (Back)
The layout is responsive with:
- Agent table: 2/3 width
@@ -240,7 +254,7 @@ class AgentMoveOperations(Widget):
"""
yield Header(show_clock=True, icon="")
title_text = Static(
f"↔️ Move Agent Operations - {len(self.agents)} device(s) selected",
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
id="move_ops_title",
)
title_text.styles.margin = (0, 0, 1, 0)
@@ -251,7 +265,7 @@ class AgentMoveOperations(Widget):
# Left side - Agent list
with Vertical() as left_side:
left_side.styles.width = "2fr"
left_side.styles.width = "3fr"
left_side.styles.height = "auto"
agents_label = Static("Selected Agents:")
@@ -266,7 +280,8 @@ class AgentMoveOperations(Widget):
# Right side - Operation buttons
with Vertical() as right_side:
right_side.styles.width = "1fr"
right_side.styles.width = "2fr"
right_side.styles.margin = (0, 1, 0, 1)
right_side.styles.height = "auto"
operations_label = Static("Operations:")
@@ -274,13 +289,23 @@ class AgentMoveOperations(Widget):
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 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"
)
@@ -300,39 +325,10 @@ class AgentMoveOperations(Widget):
status_label.styles.margin = (2, 0, 0, 0)
yield status_label
# Results display area (initially hidden)
with Vertical(id="results_container") as results_container:
results_container.styles.height = "auto"
results_container.styles.margin = (1, 0, 0, 0)
results_container.styles.display = "none"
results_label = Static("📊 Results:", id="results_label")
results_label.styles.margin = (0, 0, 0, 0)
yield results_label
results_text = TextArea(id="results_text", read_only=True)
results_text.styles.height = 15
results_text.styles.margin = (0, 0, 1, 0)
yield results_text
copy_results_btn = Button(
"📋 Copy Results to Clipboard", id="copy_results_btn"
)
copy_results_btn.styles.width = "100%"
yield copy_results_btn
# Bottom buttons
with Horizontal() as button_row:
button_row.styles.height = "auto"
button_row.styles.margin = (1, 0, 0, 0)
back_button = Button("← Back", id="back_button")
back_button.styles.width = "1fr"
yield back_button
reset_button = Button("🔄 Reset Selection", id="reset_button")
reset_button.styles.width = "1fr"
yield reset_button
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:
"""
@@ -359,13 +355,15 @@ class AgentMoveOperations(Widget):
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)
- reset_button: Clear operation state and hide results
- 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
@@ -380,21 +378,8 @@ class AgentMoveOperations(Widget):
btn_id = event.button.id
if btn_id == "back_button":
self.app.pop_screen()
event.stop()
elif btn_id == "reset_button":
# Reset operation selection
self.selected_operation = ""
self.operation_in_progress = False
status_label = self.query_one("#status_label", Static)
status_label.update("")
# Hide results
try:
results_container = self.query_one("#results_container", Vertical)
results_container.styles.display = "none"
except NoMatches:
pass
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "copy_results_btn":
@@ -404,18 +389,21 @@ class AgentMoveOperations(Widget):
pyperclip.copy(results_text.text)
self.app.notify(
"✅ Results copied to clipboard!",
"📋✅ Results copied to clipboard!",
severity="information",
timeout=2,
)
except ImportError:
self.app.notify(
"⚠️ pyperclip not installed. Run: pip install pyperclip",
" 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()
@@ -428,6 +416,9 @@ class AgentMoveOperations(Widget):
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:
"""
@@ -456,7 +447,7 @@ class AgentMoveOperations(Widget):
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("⏳ Moving agents to local approval...")
status_label.update("✔️ Moving agents to local approval...")
# Get API from app
api = self.app.api
@@ -491,12 +482,12 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error during local approval operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f" Error: {str(e)}")
self.operation_in_progress = False
return
self.operation_in_progress = False
status_label.update("✅ Operation complete!")
status_label.update(" Operation complete!")
# Display results in the widget
self._display_results("Local Approval Mode", successful, unsuccessful)
@@ -508,14 +499,63 @@ class AgentMoveOperations(Widget):
)
)
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
- 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
@@ -570,7 +610,7 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error during toggle enforcement operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f" Error: {str(e)}")
self.operation_in_progress = False
return
@@ -640,11 +680,17 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error loading policies: {e}")
status_label.update(f"Error: {str(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.
+35
View File
@@ -0,0 +1,35 @@
from textual.color import Color
from textual.theme import Theme
def get_amber_terminal_theme():
"""Amber CRT theme with compensated brightness for blending."""
return Theme(
name="amber-terminal",
background=Color.parse("#000000"), # pure black
primary=Color.parse("#ffb733"), # bright amber
secondary=Color.parse("#e69500"), # strong amber
success=Color.parse("#ffb733"),
warning=Color.parse("#ffff66"),
error=Color.parse("#ff3300"),
surface=Color.parse("#3a1f00"), # brighter brown for blending
)
AMBER_TERMINAL_CSS = """
Screen {
align: center middle;
background: #000000; /* force black */
color: #ffb733; /* force amber text */
}
.widget {
border: tall #ffb733; /* force amber border */
background: #3a1f00; /* compensated surface */
width: 80%;
}
* {
font-family: "Courier New", monospace;
}
"""
+1 -1
View File
@@ -40,7 +40,7 @@ class MultiAgentSelector(Widget):
def compose(self):
yield Header(show_clock=True, icon="")
title_text = Static("🖧 Multi-Agent Selector", id="selector_title")
title_text = Static("🖧 Agent Selector", id="selector_title")
title_text.styles.margin = (0, 0, 0, 1)
yield title_text
+8 -24
View File
@@ -91,7 +91,7 @@ class PolicySelector(Widget):
- Clear Filter button
- Confirm Selection button
- Policy table displaying available policies
- Back and Continue buttons for navigation
- Back buttons for navigation
"""
yield Header(show_clock=True, icon="")
title_text = Static(
@@ -144,6 +144,11 @@ class PolicySelector(Widget):
selected_label.styles.margin = (2, 0, 1, 0)
yield selected_label
cancel_button = Button("← Back", id="back_button")
cancel_button.styles.width = "100%"
cancel_button.styles.margin = (1, 0, 1, 0)
yield cancel_button
# Right side - Policy table
with Vertical() as right_side:
right_side.styles.width = "2fr"
@@ -158,23 +163,6 @@ class PolicySelector(Widget):
policy_table.styles.margin = (1, 0, 1, 0)
yield policy_table
# Bottom buttons
with Horizontal() as button_row:
button_row.styles.height = "auto"
button_row.styles.margin = (1, 0, 0, 0)
cancel_button = Button("✕ Cancel", id="back_button", variant="error")
cancel_button.styles.width = "1fr"
yield cancel_button
continue_button = Button(
"▶ Continue",
id="continue_button",
variant="primary",
)
continue_button.styles.width = "1fr"
continue_button.styles.margin = (0, 0, 0, 1)
yield continue_button
yield Footer()
def on_mount(self) -> None:
@@ -239,7 +227,6 @@ class PolicySelector(Widget):
- filter_button (Apply Filter): Filter policies with wildcard support
- clear_filter_button: Clear filter and show all policies
- confirm_button: Confirm selection and post message
- continue_button: Continue without posting message
Args:
event (Button.Pressed): The button press event.
@@ -247,7 +234,8 @@ class PolicySelector(Widget):
btn_id = event.button.id
if btn_id == "back_button":
self.app.pop_screen()
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "filter_button":
@@ -262,10 +250,6 @@ class PolicySelector(Widget):
self._confirm_selection()
event.stop()
elif btn_id == "continue_button":
self.app.pop_screen()
event.stop()
def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
"""
Handle row selection in the policy table.
+1 -1
View File
@@ -12,7 +12,7 @@ def get_retro_terminal_theme():
success=Color.parse("#00ff00"),
warning=Color.parse("#ffff00"),
error=Color.parse("#ff0000"),
surface=Color.parse("#111111"),
surface=Color.parse("#071802"),
)
+2 -1
View File
@@ -26,7 +26,8 @@ class ThemeSelector(Widget):
("Flexoki", "flexoki"),
("Catppuccin Latte", "catppuccin-latte"),
("Solarized Light", "solarized-light"),
("Retro Terminal", "retro-terminal"), # your custom theme
("Retro Terminal", "retro-terminal"),
("Amber Terminal", "amber-terminal"), # your custom theme
]
def compose(self):