feat: Complete Policy Prep Workflow with UX upgrades, Liftoff API, and TUI merge
- Added intro screen with workflow overview, time estimate, and onboarding controls - Improved visuals: cleaner checkboxes (/), better loading screen layout - Enforced mandatory tab reviews for critical steps with warnings and blocked navigation - Optimized logging: INFO for milestones, DEBUG for internals; cleaner production logs - Implemented Liftoff API integration: paths, publishers, hashes with granular error handling - Color-coded completion feedback ( success, failure, partial) and detailed summaries - Consolidated architecture: merged TUI.py into Loxide.py (single entry point, no circular imports) - Fixed race condition in table creation with concurrency locks
This commit is contained in:
@@ -23,19 +23,438 @@
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import dotenv
|
||||
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,
|
||||
)
|
||||
import urllib3
|
||||
|
||||
from models.agent import Agent
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.security import getAPI
|
||||
from TUI.TUI import run_Loxide
|
||||
from utils.configmanager import get_system_value
|
||||
from utils.setup import setup
|
||||
from utils.utils import irtang
|
||||
from TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
||||
from TUI.Screens.otpactivityscreen import OTPActivitiesScreen
|
||||
from TUI.Screens.otprevokescreen import OTPRevokeScreen
|
||||
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
|
||||
from TUI.Screens.policyprepworkflowscreen import PolicyPrepWorkflowScreen
|
||||
from TUI.Screens.quietagentworkflowscreen import QuietAgentWorkflowScreen
|
||||
from TUI.Themes.theme_amber_terminal import get_amber_terminal_theme
|
||||
from TUI.Themes.theme_retro_terminal import get_retro_terminal_theme
|
||||
from TUI.Themes.themeselector import ThemeSelector
|
||||
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
|
||||
from TUI.Widgets.multiagentselector import MultiAgentSelector
|
||||
from TUI.Widgets.policytreewidget import PolicyTreeWidget
|
||||
from TUI.Widgets.resultsdisplay import ResultsDisplay
|
||||
from utils.configmanager import (
|
||||
get_system_value,
|
||||
get_user_value,
|
||||
load_env,
|
||||
save_user_config,
|
||||
)
|
||||
from utils.setup import get_base_directory, setup
|
||||
from utils.utils import irtang, open_directory
|
||||
|
||||
dotenv.load_dotenv()
|
||||
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GLOBAL STASH
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_APP_RESTART_REASON = 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 using the config manager.
|
||||
No need to touch .env - config manager handles everything.
|
||||
"""
|
||||
base_dir = get_base_directory()
|
||||
config_dir = base_dir / "config"
|
||||
|
||||
try:
|
||||
save_user_config(config_dir, {"TEXTUAL_THEME": theme_name})
|
||||
logger.debug("Updated user config with TEXTUAL_THEME=%s", theme_name)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to save TEXTUAL_THEME: %s", exc)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1) SCREEN
|
||||
# ---------------------------------------------------------------------------
|
||||
class MainMenuScreen(Screen):
|
||||
api: AirlockAPIWrapper
|
||||
current_tab = reactive("")
|
||||
|
||||
BUTTON_DEFS = {
|
||||
"agent_actions": [
|
||||
(
|
||||
"🖥️ - Find agent, Move agent, or Generate One Time Pass",
|
||||
"move_agent_workflow_button",
|
||||
),
|
||||
("🎫 - Review and approve OTP Activities", "otp_activities_button"),
|
||||
("🛑 - Revoke Active OTP Session", "otp_revoke_button"),
|
||||
],
|
||||
"policy": [
|
||||
("⚖️ - Prepare Policy For Enforcement", "policy_prep_button"),
|
||||
("🔕 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"),
|
||||
],
|
||||
}
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.extras = get_user_value("EXTRAS", str, "NOTTODAY")
|
||||
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:
|
||||
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 _APP_RESTART_REASON
|
||||
selected_agents = message.selected_agents
|
||||
logger.info("Selected agents: %s", selected_agents)
|
||||
# TODO: Implement actual handling of selected agents
|
||||
_APP_RESTART_REASON = ("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 _APP_RESTART_REASON
|
||||
_persist_user_theme(message.theme_name)
|
||||
_APP_RESTART_REASON = ("restart",)
|
||||
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:
|
||||
button_id = event.button.id
|
||||
logger.debug("Button pressed: %s", button_id)
|
||||
|
||||
match button_id:
|
||||
case "move_agent_workflow_button":
|
||||
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
|
||||
event.stop()
|
||||
|
||||
case "otp_generate_button":
|
||||
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
|
||||
event.stop()
|
||||
|
||||
case "find_quiet_button":
|
||||
self.app.push_screen(
|
||||
QuietAgentWorkflowScreen(self.app.api, self.app.policies)
|
||||
)
|
||||
event.stop()
|
||||
return
|
||||
|
||||
case "otp_activities_button":
|
||||
self.app.push_screen(OTPActivitiesScreen())
|
||||
event.stop()
|
||||
return
|
||||
|
||||
case "otp_revoke_button":
|
||||
self.app.push_screen(OTPRevokeScreen())
|
||||
event.stop()
|
||||
return
|
||||
|
||||
case "policy_prep_button":
|
||||
# Use the new TUI workflow screen instead of legacy
|
||||
self.app.push_screen(
|
||||
PolicyPrepWorkflowScreen(self.app.api, self.app.policies)
|
||||
)
|
||||
event.stop()
|
||||
return
|
||||
|
||||
case _:
|
||||
self.app.bell()
|
||||
logger.warning("Unknown button pressed: %s", button_id)
|
||||
return
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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"),
|
||||
("r", "refresh", "Refresh"),
|
||||
]
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper):
|
||||
self._textual_theme = get_user_value("TEXTUAL_THEME", str, "textual-dark")
|
||||
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_refresh(self) -> None:
|
||||
self.refresh_data()
|
||||
|
||||
def action_quit(self) -> None:
|
||||
global _APP_RESTART_REASON
|
||||
_APP_RESTART_REASON = 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) PUBLIC ENTRYPOINT
|
||||
# ---------------------------------------------------------------------------
|
||||
def run_Loxide(api: AirlockAPIWrapper) -> None:
|
||||
global _APP_RESTART_REASON
|
||||
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 app loop iteration (attempt %d)", attempts)
|
||||
_APP_RESTART_REASON = None
|
||||
app = Loxide(api)
|
||||
|
||||
try:
|
||||
app.run()
|
||||
except SystemExit as exc:
|
||||
if exc.code != 0:
|
||||
logger.debug("Caught SystemExit from Textual: %s", exc)
|
||||
raise
|
||||
|
||||
reason = _APP_RESTART_REASON
|
||||
logger.debug("After app.run(), _APP_RESTART_REASON = %r", reason)
|
||||
|
||||
if not reason:
|
||||
logger.debug("No restart reason, exiting loop")
|
||||
break
|
||||
|
||||
if reason[0] == "restart":
|
||||
logger.debug("Restarting app loop")
|
||||
continue
|
||||
|
||||
if reason[0] == "multi_agent_action":
|
||||
logger.info("Multi-agent action with selected agents: %s", reason[1])
|
||||
continue
|
||||
|
||||
logger.error("Unknown restart reason: %r", reason)
|
||||
break
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4) MAIN FUNCTION
|
||||
# ---------------------------------------------------------------------------
|
||||
def main():
|
||||
irtang()
|
||||
# Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
|
||||
|
||||
Reference in New Issue
Block a user