import logging import os import sys import dotenv from dotenv import set_key from textual.app import App, ComposeResult from textual.containers import Vertical 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_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.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.multiagentselector import MultiAgentSelector from widgets.OTP_generate import OTPGenerator from widgets.policytreewidget import PolicyTreeWidget 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: /config/user_config.json and also mirror to /.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): current_tab = reactive("") BUTTON_DEFS = { "find": [ ("🔍 - Device Search", "find_device_button"), ("🔇 - Find Quiet Hosts", "find_quiet_button"), ], "move": [ ("✅ - 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"), ], } def __init__(self, api: AirlockAPIWrapper) -> None: super().__init__() self.api = api 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("Policy Tree", id="p_tree"), Tab("Device Search", id="find"), Tab("Move Agent", id="move"), Tab("OTP", id="otp"), Tab("Directory", id="dir"), Tab("Settings", id="settings"), ] if self.extras == "POLICYPREP": tabs.insert(3, 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") # 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 MultiAgentSelector.""" 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_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 "find_device_button": _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {}) case "find_quiet_button": _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {}) 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)) event.stop() return # Don't exit the app case "otp_activities_button": _PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {}) 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 logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB) self.app.exit() # --------------------------------------------------------------------------- # 2) APP # --------------------------------------------------------------------------- class Loxide(App): CSS = """ #logo { width: 100%; content-align: center middle; text-align: center; } """ BINDINGS = [ ("q", "quit", "Quit"), ("d", "open_dir", "Open Directory"), ] 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 # Add error handling for API calls try: self.policies = [ Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows() ] self.devices = [ Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows() ] 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.theme = self._textual_theme self.push_screen(MainMenuScreen(api)) def action_quit(self) -> None: global _PENDING_JOB _PENDING_JOB = None self.exit() def action_open_dir(self) -> None: screen = self.screen_stack[-1] if isinstance(screen, MainMenuScreen): if screen.current_tab != "dir": screen.switch_tab("dir") # --------------------------------------------------------------------------- # 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 while True: base_dir = get_base_directory() env_path = base_dir / ".env" dotenv.load_dotenv(dotenv_path=env_path, override=True) _PENDING_JOB = None app = Loxide(api) try: app.run() except SystemExit as exc: logger.debug("Caught SystemExit from Textual: %s", exc) job = _PENDING_JOB logger.debug("After app.run(), _PENDING_JOB = %r", job) if not job: break if job[0] == "legacy": _, func, args, kwargs = job _run_legacy_job(func, args, kwargs) continue if job[0] == "restart": # just loop again; fresh .env was already loaded at the top continue if job[0] == "multi_agent_action": # Handle multi-agent selection logger.info("Multi-agent action with selected agents: %s", job[1]) continue # NEW: Handle OTP workflow if job[0] == "otp_workflow": _, devices, requestor, reasoning, duration = job # Call your OTP generation with the parameters def otp_generate_with_params(): print(f"\n{'='*60}") print("OTP GENERATION") print(f"{'='*60}") print(f"Requestor: {requestor}") print(f"Reasoning: {reasoning}") print(f"Duration: {duration} minutes") print(f"\nGenerating OTPs for {len(devices)} devices:") print(f"{'='*60}\n") # Call your actual OTP generation function # You'll need to adapt otp_generate to accept these parameters # For now, this is a placeholder showing the structure for device in devices: print(f"Device: {device}") print(f" Requestor: {requestor}") print(f" Reason: {reasoning}") print(f" Duration: {duration} minutes") # TODO: Actually call your API to generate OTP # result = api.generate_otp(device, requestor, reasoning, duration) print() print(f"{'='*60}") print("OTP Generation Complete!") print(f"{'='*60}") _run_legacy_job(otp_generate_with_params, (), {}) continue break # --------------------------------------------------------------------------- # 5) DEV # --------------------------------------------------------------------------- if __name__ == "__main__": api = AirlockAPIWrapper() run_Loxide(api)