Major step towards unification of the UI Implementation of the Back Feature, splitting of TUI files back into subfolders
This commit is contained in:
+63
-130
@@ -1,6 +1,20 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
import dotenv
|
||||
@@ -19,24 +33,22 @@ from textual.widgets import (
|
||||
Tabs,
|
||||
)
|
||||
|
||||
from flows.otp import otp_revoke
|
||||
from flows.prepPolicy import menu_policy_enforce
|
||||
from models.agent import Agent
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.policyhandler import confirmUpdateAfromE
|
||||
from TUI.agentmoveoperations import AgentMoveOperations
|
||||
from TUI.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
||||
from TUI.multiagentselector import MultiAgentSelector
|
||||
from TUI.OTP_generate import OTPGenerator
|
||||
from TUI.otpactivityscreen import OTPActivitiesScreen
|
||||
from TUI.otpworkflowscreen import OTPWorkflowScreen
|
||||
from TUI.policytreewidget import PolicyTreeWidget
|
||||
from TUI.quietagentworkflowscreen import QuietAgentWorkflowScreen
|
||||
from TUI.resultsdisplay import ResultsDisplay
|
||||
from TUI.theme_amber_terminal import get_amber_terminal_theme
|
||||
from TUI.theme_retro_terminal import get_retro_terminal_theme
|
||||
from TUI.themeselector import ThemeSelector
|
||||
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_user_value, load_env, save_user_config
|
||||
from utils.setup import get_base_directory
|
||||
from utils.utils import open_directory
|
||||
@@ -47,7 +59,7 @@ dotenv.load_dotenv()
|
||||
# GLOBAL STASH
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PENDING_JOB = None
|
||||
_APP_RESTART_REASON = None
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -83,13 +95,12 @@ class MainMenuScreen(Screen):
|
||||
"🖥️ - Find, Move, or Generate OTP for Agents",
|
||||
"move_agent_workflow_button",
|
||||
),
|
||||
("📊 - Review and appove OTP Activities", "otp_activities_button"),
|
||||
("📇 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"),
|
||||
("🎫 - Review and appove OTP Activities", "otp_activities_button"),
|
||||
("🔕 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"),
|
||||
],
|
||||
"policy": [
|
||||
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
|
||||
("🔄 - Update Audit Policies", "policy_audit_update_button"),
|
||||
("❌ - Revoke OTPs", "otp_revoke_button"),
|
||||
("⚖️ - Prepare Policy For Enforcement", "policy_prep_button"),
|
||||
("🛑 - Revoke OTPs", "otp_revoke_button"),
|
||||
],
|
||||
}
|
||||
|
||||
@@ -128,7 +139,6 @@ class MainMenuScreen(Screen):
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
api = self.app.api
|
||||
self.switch_tab("agent_actions")
|
||||
|
||||
# focus helpers
|
||||
@@ -191,42 +201,20 @@ class MainMenuScreen(Screen):
|
||||
self, message: MultiAgentSelector.AgentsSelected
|
||||
) -> None:
|
||||
"""Handle selected agents from AgentSelector."""
|
||||
global _PENDING_JOB
|
||||
global _APP_RESTART_REASON
|
||||
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)
|
||||
_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 _PENDING_JOB
|
||||
global _APP_RESTART_REASON
|
||||
_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
|
||||
|
||||
logger.info(
|
||||
"OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d",
|
||||
len(message.devices),
|
||||
message.requestor,
|
||||
message.reasoning,
|
||||
message.duration,
|
||||
)
|
||||
|
||||
_PENDING_JOB = (
|
||||
"otp_workflow",
|
||||
message.devices,
|
||||
message.requestor,
|
||||
message.reasoning,
|
||||
message.duration,
|
||||
)
|
||||
|
||||
_APP_RESTART_REASON = ("restart",)
|
||||
self.app.exit()
|
||||
|
||||
def on_agent_move_operations_operation_complete(
|
||||
@@ -282,7 +270,6 @@ class MainMenuScreen(Screen):
|
||||
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)
|
||||
|
||||
@@ -308,27 +295,23 @@ class MainMenuScreen(Screen):
|
||||
return
|
||||
|
||||
case "otp_revoke_button":
|
||||
_PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {})
|
||||
self.app.push_screen(OTPRevokeScreen())
|
||||
event.stop()
|
||||
return
|
||||
|
||||
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,), {})
|
||||
# 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
|
||||
|
||||
# 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
|
||||
@@ -353,7 +336,7 @@ class Loxide(App[Message]):
|
||||
]
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper):
|
||||
self._textual_theme = get_user_value("TEXTUAL_THEME", str, "nord")
|
||||
self._textual_theme = get_user_value("TEXTUAL_THEME", str, "textual-dark")
|
||||
super().__init__()
|
||||
self.api = api
|
||||
wd = load_env("WORKING_DIR") or os.getcwd()
|
||||
@@ -395,8 +378,8 @@ class Loxide(App[Message]):
|
||||
self.refresh_data()
|
||||
|
||||
def action_quit(self) -> None:
|
||||
global _PENDING_JOB
|
||||
_PENDING_JOB = None
|
||||
global _APP_RESTART_REASON
|
||||
_APP_RESTART_REASON = None
|
||||
self.exit()
|
||||
|
||||
def action_open_fe(self) -> None:
|
||||
@@ -410,45 +393,10 @@ class Loxide(App[Message]):
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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
|
||||
# 3) PUBLIC ENTRYPOINT
|
||||
# ---------------------------------------------------------------------------
|
||||
def run_Loxide(api: AirlockAPIWrapper) -> None:
|
||||
global _PENDING_JOB
|
||||
global _APP_RESTART_REASON
|
||||
base_dir = get_base_directory()
|
||||
env_path = base_dir / ".env"
|
||||
dotenv.load_dotenv(dotenv_path=env_path, override=True)
|
||||
@@ -458,8 +406,8 @@ def run_Loxide(api: AirlockAPIWrapper) -> None:
|
||||
|
||||
while attempts < max_attempts:
|
||||
attempts += 1
|
||||
logger.debug("Starting job loop iteration (attempt %d)", attempts)
|
||||
_PENDING_JOB = None
|
||||
logger.debug("Starting app loop iteration (attempt %d)", attempts)
|
||||
_APP_RESTART_REASON = None
|
||||
app = Loxide(api)
|
||||
|
||||
try:
|
||||
@@ -469,42 +417,27 @@ def run_Loxide(api: AirlockAPIWrapper) -> None:
|
||||
logger.debug("Caught SystemExit from Textual: %s", exc)
|
||||
raise
|
||||
|
||||
job = _PENDING_JOB
|
||||
logger.debug("After app.run(), _PENDING_JOB = %r", job)
|
||||
reason = _APP_RESTART_REASON
|
||||
logger.debug("After app.run(), _APP_RESTART_REASON = %r", reason)
|
||||
|
||||
if not job:
|
||||
logger.debug("No job pending, exiting loop")
|
||||
if not reason:
|
||||
logger.debug("No restart reason, exiting loop")
|
||||
break
|
||||
|
||||
if job[0] == "legacy":
|
||||
_, func, args, kwargs = job
|
||||
_run_legacy_job(func, args, kwargs)
|
||||
if reason[0] == "restart":
|
||||
logger.debug("Restarting app loop")
|
||||
continue
|
||||
|
||||
if job[0] == "restart":
|
||||
logger.debug("Restarting job loop")
|
||||
if reason[0] == "multi_agent_action":
|
||||
logger.info("Multi-agent action with selected agents: %s", reason[1])
|
||||
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)
|
||||
logger.error("Unknown restart reason: %r", reason)
|
||||
break
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5) DEV
|
||||
# 4) DEV
|
||||
# ---------------------------------------------------------------------------
|
||||
if __name__ == "__main__":
|
||||
api = AirlockAPIWrapper()
|
||||
|
||||
Reference in New Issue
Block a user