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
utils/tui.py
-542
View File
@@ -1,542 +0,0 @@
import logging
import os
import sys
from typing import Optional
import dotenv
from dotenv import set_key
from textual.app import App, ComposeResult
from textual.containers import Vertical
from textual.message import Message
from textual.reactive import reactive
from textual.screen import Screen
from textual.widgets import (
Button,
DirectoryTree,
Footer,
Header,
Static,
Tab,
Tabs,
)
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 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()
# ---------------------------------------------------------------------------
# GLOBAL STASH
# ---------------------------------------------------------------------------
_PENDING_JOB = None
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
# ---------------------------------------------------------------------------
def _persist_user_theme(theme_name: str) -> None:
"""
Store the chosen Textual theme in the user's config:
<base>/config/user_config.json
and also mirror to <base>/.env so load_env(...) sees it.
"""
base_dir = get_base_directory()
config_dir = base_dir / "config"
user_config_path = config_dir / "user_config.json"
env_path = base_dir / ".env"
# ensure dirs / files exist similarly to setup()
config_dir.mkdir(parents=True, exist_ok=True)
if not user_config_path.exists():
# minimal default like your load_user_config does
user_config_path.write_text(
'{"URL": "", "LOG_LEVEL": "INFO"}\n', encoding="utf-8"
)
# load existing user config
user_conf = load_user_config(config_dir)
user_conf["TEXTUAL_THEME"] = theme_name
# write it back
user_config_path.write_text(
# pretty print so it stays human-readable
__import__("json").dumps(user_conf, indent=4),
encoding="utf-8",
)
logger.debug("Updated user_config.json with TEXTUAL_THEME=%s", theme_name)
# mirror to .env (like setup.write_config_to_env does)
env_path.parent.mkdir(parents=True, exist_ok=True)
if not env_path.exists():
env_path.touch()
try:
set_key(str(env_path), "TEXTUAL_THEME", theme_name)
except Exception as exc: # keep going even if .env write fails
logger.warning("Failed to mirror TEXTUAL_THEME to .env: %s", exc)
# reload so load_env(...) sees the new value right now
dotenv.load_dotenv(dotenv_path=env_path, override=True)
logger.debug("Reloaded .env from %s", env_path)
# ---------------------------------------------------------------------------
# 1) SCREEN
# ---------------------------------------------------------------------------
class MainMenuScreen(Screen):
api: AirlockAPIWrapper
current_tab = reactive("")
BUTTON_DEFS = {
"agent_actions": [
(
"🖥️ - Find, Move, or Generate OTP for Agents",
"move_agent_workflow_button",
),
("📊 - OTP Activities By Agent", "otp_activities_button"),
("🔇 - Find Quiet Hosts", "find_quiet_button"),
],
"policy": [
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
("🔄 - Update Audit Policies", "policy_audit_update_button"),
("❌ - Revoke OTPs", "otp_revoke_button"),
],
}
def __init__(self) -> None:
super().__init__()
self.extras = load_env("EXTRAS")
wd = load_env("WORKING_DIR") or os.getcwd()
if not os.path.isdir(wd):
wd = os.getcwd()
self.working_dir = wd
def _make_buttons_for(self, tab_id: str) -> Vertical:
defs = self.BUTTON_DEFS.get(tab_id, [])
buttons = []
for label, btn_id in defs:
btn = Button(label, id=btn_id)
btn.styles.width = "100%"
buttons.append(btn)
return Vertical(*buttons)
def compose(self) -> ComposeResult:
yield Header(show_clock=True, icon="⚙")
tabs = [
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(2, Tab("Policy Prep", id="policy"))
yield Tabs(*tabs, id="tabs")
yield Vertical(id="content")
yield Footer()
def on_mount(self) -> None:
api = self.app.api
self.switch_tab("agent_actions")
# focus helpers
def _get_content_buttons(self) -> list[Button]:
content = self.query_one("#content", Vertical)
return list(content.query(Button))
def _focus_first_button(self) -> None:
buttons = self._get_content_buttons()
if buttons:
buttons[0].focus()
def _focus_tabs(self) -> None:
tabs = self.query_one("#tabs", Tabs)
tabs.focus()
def _focus_nearby_button(self, direction: int) -> None:
buttons = self._get_content_buttons()
if not buttons:
return
try:
current = next(i for i, b in enumerate(buttons) if b.has_focus)
except StopIteration:
if direction > 0:
buttons[0].focus()
else:
buttons[-1].focus()
return
if direction < 0 and current == 0:
self._focus_tabs()
return
new_index = current + direction
if 0 <= new_index < len(buttons):
buttons[new_index].focus()
def switch_tab(self, tab_id: str) -> None:
self.current_tab = tab_id
content = self.query_one("#content", Vertical)
content.remove_children()
if tab_id in self.BUTTON_DEFS:
content.mount(self._make_buttons_for(tab_id))
self.call_later(self._focus_first_button)
elif tab_id == "dir":
content.mount(DirectoryTree(self.working_dir, id="dir_tree"))
elif tab_id == "p_tree":
content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
elif tab_id == "settings":
content.mount(ThemeSelector())
else:
content.mount(Static(f"Unknown tab: {tab_id}"))
def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
self.switch_tab(event.tab.id)
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
"""Handle selected agents from AgentSelector."""
global _PENDING_JOB
selected_agents = message.selected_agents
logger.info("Selected agents: %s", selected_agents)
# TODO: Implement actual handling of selected agents
_PENDING_JOB = ("multi_agent_action", selected_agents)
self.app.exit()
def on_theme_selector_theme_selected(
self, message: ThemeSelector.ThemeSelected
) -> None:
"""Handle theme selection from ThemeSelector."""
global _PENDING_JOB
_persist_user_theme(message.theme_name)
_PENDING_JOB = ("restart",)
self.app.exit()
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
"""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),
message.requestor,
message.reasoning,
message.duration,
)
# Set up the job to run the OTP generation
_PENDING_JOB = (
"otp_workflow",
message.devices,
message.requestor,
message.reasoning,
message.duration,
)
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:
path = event.path
logger.debug("Directory file selected: %s", path)
try:
open_directory(str(path))
except Exception as exc:
logger.error("Failed to open %s: %s", path, exc)
self.app.bell()
def on_button_pressed(self, event: Button.Pressed) -> None:
global _PENDING_JOB
button_id = event.button.id
logger.debug("Button pressed: %s", button_id)
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,), {})
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
case "otp_revoke_button":
_PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {})
case "policy_prep_button":
_PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {})
case "policy_audit_update_button":
_PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {})
case _:
self.app.bell()
logger.warning("Unknown button pressed: %s", button_id)
return
# Only exit the UI loop when we explicitly queued a legacy job.
# The original flow used `self.app.exit()` after setting _PENDING_JOB so
# the outer loop could run legacy code. Keep that behavior only for legacy jobs.
logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB)
if _PENDING_JOB and _PENDING_JOB[0] == "legacy":
# let the main loop pick up the legacy job
self.app.exit()
# ---------------------------------------------------------------------------
# 2) APP
# ---------------------------------------------------------------------------
class Loxide(App[Message]):
api: AirlockAPIWrapper
working_dir: str
policies: Optional[list[Policy]]
devices: Optional[list[Agent]]
CSS = """
#logo {
width: 100%;
content-align: center middle;
text-align: center;
}
"""
BINDINGS = [
("q", "quit", "Quit"),
("f", "open_fe", "Launch Explorer"),
]
def __init__(self, api: AirlockAPIWrapper):
self._textual_theme = load_env("TEXTUAL_THEME") or "nord"
super().__init__()
self.api = api
wd = load_env("WORKING_DIR") or os.getcwd()
if not os.path.isdir(wd):
wd = os.getcwd()
self.working_dir = wd
# Initial data load
self.refresh_data()
def refresh_data(self) -> None:
"""Public method to refresh policies and devices from the API."""
try:
self.policies = [
Policy(**row.to_dict())
for _, row in self.api.policy_find_all().iterrows()
]
self.devices = [
Agent(**row.to_dict())
for _, row in self.api.agent_find_all().iterrows()
]
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.register_theme(get_amber_terminal_theme())
self.theme = self._textual_theme
self.push_screen(MainMenuScreen())
def action_quit(self) -> None:
global _PENDING_JOB
_PENDING_JOB = None
self.exit()
def action_open_fe(self) -> None:
"""Open the working directory in the OS file manager (footer binding)."""
path_to_open = self.working_dir or os.getcwd()
try:
open_directory(path_to_open)
except Exception as exc:
logger.error("Failed to open directory %s: %s", path_to_open, exc)
self.bell() # optional feedback
# ---------------------------------------------------------------------------
# 3) TERMINAL + LEGACY
# ---------------------------------------------------------------------------
def _restore_terminal_for_legacy() -> None:
sys.stdout.write("\033[?1049l")
sys.stdout.write("\033[?25h")
sys.stdout.write("\033[0m")
sys.stdout.write("\033[?1000l\033[?1002l\033[?1003l\033[?1006l")
sys.stdout.write("\033[2J\033[H")
sys.stdout.flush()
if os.name == "nt":
try:
import ctypes
kernel32 = ctypes.windll.kernel32
handle = kernel32.GetStdHandle(-11)
mode = ctypes.c_ulong()
if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
kernel32.SetConsoleMode(handle, mode.value | 0x0004)
except Exception as exc:
logger.debug("VT enable on Windows failed: %s", exc)
def _run_legacy_job(func, args, kwargs) -> None:
logger.debug("Running legacy job: %s", getattr(func, "__name__", func))
_restore_terminal_for_legacy()
try:
func(*args, **kwargs)
finally:
try:
input("\nPress Enter to return to the UI...")
except EOFError:
pass
# ---------------------------------------------------------------------------
# 4) PUBLIC ENTRYPOINT
# ---------------------------------------------------------------------------
def run_Loxide(api: AirlockAPIWrapper) -> None:
global _PENDING_JOB
base_dir = get_base_directory()
env_path = base_dir / ".env"
dotenv.load_dotenv(dotenv_path=env_path, override=True)
max_attempts = 5
attempts = 0
while attempts < max_attempts:
attempts += 1
logger.debug("Starting job loop iteration (attempt %d)", attempts)
_PENDING_JOB = None
app = Loxide(api)
try:
app.run()
except SystemExit as exc:
if exc.code != 0:
logger.debug("Caught SystemExit from Textual: %s", exc)
raise
job = _PENDING_JOB
logger.debug("After app.run(), _PENDING_JOB = %r", job)
if not job:
logger.debug("No job pending, exiting loop")
break
if job[0] == "legacy":
_, func, args, kwargs = job
_run_legacy_job(func, args, kwargs)
continue
if job[0] == "restart":
logger.debug("Restarting job loop")
continue
if job[0] == "multi_agent_action":
logger.info("Multi-agent action with selected agents: %s", job[1])
continue
if job[0] == "otp_workflow":
_, devices, requestor, reasoning, duration = job
def otp_generate_with_params():
# Your OTP logic here
pass
_run_legacy_job(otp_generate_with_params, (), {})
continue
logger.error("Unknown job type: %r", job)
break
# ---------------------------------------------------------------------------
# 5) DEV
# ---------------------------------------------------------------------------
if __name__ == "__main__":
api = AirlockAPIWrapper()
run_Loxide(api)