# 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 . # TODO Continue implementing logger # TODO Add input sanitation and CSV injection prevention # TODO Continue OTP and Local approval rewrites # TODO Explore pywin32 # TODO Fix Requirements.txt # TODO Create Generic system_config.json for gitea from collections import Counter from datetime import UTC, datetime, timedelta import json import logging import os from typing import Optional from bson import ObjectId import dotenv from textual.app import App, ComposeResult from textual.containers import Horizontal, Vertical from textual.message import Message from textual.reactive import reactive from textual.screen import Screen from textual.widgets import ( Button, DirectoryTree, Footer, Header, Select, Static, Tab, Tabs, ) import urllib3 import airlock_libs from models.agent import Agent from models.policy import Policy from services.API import AirlockAPIWrapper from services.security import getAPI from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen 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.policytreewidget import PolicyTreeWidget from TUI.Widgets.resultsdisplay import ResultsDisplay from TUI.Widgets.serverlogwidget import ServerLogWidget 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": [ { "label": "🖥️ - Multi-Agent Operations", "id": "move_agent_workflow_button", "description": "Select agents to: Move policies, Generate OTPs, Toggle audit/enforcement, View history, Export data", }, { "label": "🎫 - Review and approve OTP Activities", "id": "otp_activities_button", }, { "label": "🛑 - Revoke Active OTP Session", "id": "otp_revoke_button", }, ], "policy": [ { "label": "⚖️ - Prepare Policy For Enforcement", "id": "policy_prep_button", }, { "label": "🔕 - Find and Move Quiet Hosts to Enforcement", "id": "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 # Statistics cache: {days: {"data": stats_dict, "timestamp": datetime}} self.statistics_cache = {} self.current_stats_days = 1 # Default to 1 day def _make_buttons_for(self, tab_id: str) -> Vertical: defs = self.BUTTON_DEFS.get(tab_id, []) widgets = [] for item in defs: # Support both old tuple format and new dict format if isinstance(item, dict): label = item["label"] btn_id = item["id"] description = item.get("description") else: # Old tuple format: (label, id) label, btn_id = item description = None btn = Button(label, id=btn_id) btn.styles.width = "100%" widgets.append(btn) # Add description text if provided if description: desc_text = Static(description, classes="button_description") desc_text.styles.width = "100%" desc_text.styles.color = "ansi_bright_black" desc_text.styles.text_align = "center" desc_text.styles.margin = (0, 0, 1, 0) widgets.append(desc_text) return Vertical(*widgets) def compose(self) -> ComposeResult: yield Header(show_clock=True, icon="âš™") tabs = [ Tab("Agents", id="agent_actions"), Tab("Tree View", id="p_tree"), Tab("Server Log", id="server_log"), Tab("Statistics", id="statistics"), 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") def on_key(self, event) -> None: """Handle up/down arrow keys for button navigation.""" if event.key == "down": self._focus_nearby_button(1) event.prevent_default() event.stop() elif event.key == "up": self._focus_nearby_button(-1) event.prevent_default() event.stop() # left/right are handled by Textual's default tab navigation # 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)) elif tab_id == "server_log": content.mount(ServerLogWidget(self.app.api)) 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 == "statistics": content.mount(self._create_statistics_widget()) elif tab_id == "settings": content.mount(ThemeSelector()) else: content.mount(Static(f"Unknown tab: {tab_id}")) def _create_statistics_widget(self) -> Vertical: """Create the enhanced statistics display widget with time period selection.""" stats_container = Vertical(id="statistics_container") # Build all widgets first, then mount them all at once widgets_to_mount = [] # Header header = Static("📊 System Statistics", classes="stats_header") header.styles.text_align = "center" header.styles.text_style = "bold" header.styles.margin = (1, 0, 1, 0) widgets_to_mount.append(header) # Controls row controls = Horizontal() controls.styles.height = "auto" controls.styles.margin = (0, 2, 1, 2) # Time period selector time_label = Static("Time Period: ") time_label.styles.width = "auto" time_label.styles.margin = (0, 1, 0, 0) time_select = Select( options=[ ("1 Day", 1), ("3 Days", 3), ("7 Days", 7), ], value=self.current_stats_days, id="stats_time_select", ) time_select.styles.width = 20 # Refresh button refresh_btn = Button("🔄 Refresh", id="stats_refresh_btn", variant="primary") refresh_btn.styles.margin = (0, 0, 0, 2) # Compose controls controls.compose_add_child(time_label) controls.compose_add_child(time_select) controls.compose_add_child(refresh_btn) widgets_to_mount.append(controls) # Timestamp and status area status_container = Vertical(id="stats_status_area") status_container.styles.margin = (0, 2, 1, 2) widgets_to_mount.append(status_container) # Data display area - using horizontal layout data_container = Horizontal(id="stats_data_area") data_container.styles.margin = (0, 2, 0, 2) data_container.styles.height = "auto" widgets_to_mount.append(data_container) # Compose all widgets into the container for widget in widgets_to_mount: stats_container.compose_add_child(widget) # Schedule data fetch after the widget is mounted if self.current_stats_days not in self.statistics_cache: self.call_after_refresh( lambda: self._fetch_execution_statistics(self.current_stats_days) ) else: self.call_after_refresh(self._update_statistics_display) return stats_container def _skipback(self, days: int) -> ObjectId: """Generate a MongoDB ObjectId for a given number of days ago.""" date_days_ago = datetime.now(UTC) - timedelta(days=days) timestamp = int(date_days_ago.timestamp()) hex_timestamp = format(timestamp, "08x") objectid_hex = hex_timestamp + "0000000000000000" return ObjectId(objectid_hex) def _fetch_execution_statistics(self, days: int) -> None: """Fetch execution history and calculate statistics using airlock_libs.""" try: status_area = self.query_one("#stats_status_area", Vertical) except Exception as e: logger.warning(f"Could not find status area: {e}") return status_area.remove_children() status = Static("⏳ Fetching execution data...") status.styles.text_style = "italic" status_area.mount(status) try: # Non-trusted execution types exec_types = [1, 2, 3, 6, 7, 12, 13, 14, 15, 16] logger.info( f"Fetching {days} days of execution history for all policies..." ) logger.debug(f"Calling with exec_types={exec_types}, days={days}") # Use pull_policy_exec_histories - must pass all 4 args in order: (api, policy_name, type, days) # Empty string for policy_name means all policies execs = airlock_libs.pull_policy_exec_histories( self.app.api, None, str(exec_types), days # Empty string = all policies ) logger.debug(f"Received response, length: {len(execs) if execs else 0}") if not execs: logger.warning("No execution data returned") status_area.remove_children() error = Static("No execution data found for the selected time period.") error.styles.color = "yellow" status_area.mount(error) return # Parse JSON response data = json.loads(execs) logger.debug(f"Parsed JSON, keys: {data.keys()}") all_executions = data.get("response", {}).get("exechistories", []) logger.debug(f"Found {len(all_executions)} executions in response") if not all_executions: logger.warning("No executions in response") status_area.remove_children() error = Static("No execution records found.") error.styles.color = "yellow" status_area.mount(error) return logger.info(f"Retrieved {len(all_executions)} execution records") # Process statistics stats = self._process_execution_data(all_executions) # Cache the results self.statistics_cache[days] = { "data": stats, "timestamp": datetime.now(UTC), } # Update display self._update_statistics_display() except Exception as e: logger.error(f"Failed to fetch execution statistics: {e}", exc_info=True) status_area.remove_children() error = Static(f"❌ Error fetching data: {str(e)}") error.styles.color = "red" status_area.mount(error) def _process_execution_data(self, executions: list) -> dict: """Process raw execution data into statistics.""" # Execution type names type_names = { 0: "Trusted Execution", 1: "Blocked Execution", 2: "Untrusted Execution [Audit]", 3: "Untrusted Execution [OTP]", 4: "Trusted Path Execution", 5: "Trusted Publisher Execution", 6: "Blocklist Execution", 7: "Blocklist Execution [Audit]", 8: "Trusted Process Execution", 9: "Constrained Execution", 10: "Trusted Metadata Execution", 11: "Trusted Browser Execution", 12: "Blocked Browser Execution", 13: "Untrusted Browser Execution [Audit]", 14: "Untrusted Browser Execution [OTP]", 15: "Blocklist Browser Execution [Audit]", 16: "Blocklist Browser Execution", 17: "Trusted Installer Execution", 18: "Trusted Browser Metadata Execution", } # Count execution types type_counter = Counter() policy_counter = Counter() hostname_counter = Counter() for exec_record in executions: exec_type = exec_record.get("type", -1) type_counter[exec_type] += 1 policy_name = exec_record.get("policyname", "Unknown") policy_counter[policy_name] += 1 hostname = exec_record.get("hostname", "Unknown") hostname_counter[hostname] += 1 # Get top 5 policies and hostnames top_policies = policy_counter.most_common(5) top_hostnames = hostname_counter.most_common(5) # Format execution type counts type_counts = [] for exec_type, count in sorted(type_counter.items()): type_name = type_names.get(exec_type, f"Unknown Type {exec_type}") type_counts.append((type_name, count)) return { "total_executions": len(executions), "type_counts": type_counts, "top_policies": top_policies, "top_hostnames": top_hostnames, } def _update_statistics_display(self) -> None: """Update the statistics display with cached data using horizontal layout.""" try: data_area = self.query_one("#stats_data_area", Horizontal) status_area = self.query_one("#stats_status_area", Vertical) except Exception as e: logger.warning(f"Could not find statistics display areas: {e}") return data_area.remove_children() status_area.remove_children() # Get cached data cache_entry = self.statistics_cache.get(self.current_stats_days) if not cache_entry: status = Static("No data available. Click Refresh to fetch.") status.styles.text_style = "italic" status_area.mount(status) return stats = cache_entry["data"] timestamp = cache_entry["timestamp"] # Timestamp display timestamp_str = timestamp.strftime("%Y-%m-%d %H:%M:%S UTC") timestamp_label = Static(f"Last updated: {timestamp_str}") timestamp_label.styles.text_style = "dim" timestamp_label.styles.margin = (0, 0, 1, 0) status_area.mount(timestamp_label) # Create three columns for horizontal layout # Column 1: Summary and Execution Types col1 = Vertical() col1.styles.width = "1fr" col1.styles.padding = (0, 1, 0, 0) # Column 2: Top Policies col2 = Vertical() col2.styles.width = "1fr" col2.styles.padding = (0, 1) # Column 3: Top Machines and System Overview col3 = Vertical() col3.styles.width = "1fr" col3.styles.padding = (0, 0, 0, 1) # === COLUMN 1: Summary and Execution Types === total = Static(f"📊 Total Non-Trusted Executions: {stats['total_executions']}") total.styles.text_style = "bold" total.styles.margin = (0, 0, 1, 0) col1.compose_add_child(total) # Execution type breakdown if stats["type_counts"]: type_header = Static("📋 Execution Types") type_header.styles.text_style = "bold" type_header.styles.margin = (1, 0, 0, 0) col1.compose_add_child(type_header) for type_name, count in stats["type_counts"]: # Abbreviate long type names short_name = type_name.replace("Execution", "Exec").replace( "[Audit]", "[Aud]" ) type_line = Static(f" {short_name}: {count}") type_line.styles.margin = (0, 0, 0, 1) col1.compose_add_child(type_line) # === COLUMN 2: Top Policies === if stats["top_policies"]: policy_header = Static("⚖️ Top 5 Policies") policy_header.styles.text_style = "bold" policy_header.styles.margin = (0, 0, 0, 0) col2.compose_add_child(policy_header) for i, (policy_name, count) in enumerate(stats["top_policies"], 1): policy_line = Static(f"{i}. {policy_name}: {count}") policy_line.styles.margin = (0, 0, 0, 1) col2.compose_add_child(policy_line) # === COLUMN 3: Top Machines and System Overview === if stats["top_hostnames"]: host_header = Static("🖥️ Top 5 Machines") host_header.styles.text_style = "bold" host_header.styles.margin = (0, 0, 0, 0) col3.compose_add_child(host_header) for i, (hostname, count) in enumerate(stats["top_hostnames"], 1): host_line = Static(f"{i}. {hostname}: {count}") host_line.styles.margin = (0, 0, 0, 1) col3.compose_add_child(host_line) # System Overview in column 3 total_agents = len(self.app.devices) if self.app.devices else 0 total_policies = len(self.app.policies) if self.app.policies else 0 if self.app.devices: enforced_agents = sum( 1 for agent in self.app.devices if getattr(agent, "enforcement", False) ) audit_agents = sum( 1 for agent in self.app.devices if not getattr(agent, "enforcement", False) ) online_agents = sum( 1 for agent in self.app.devices if getattr(agent, "online", False) ) offline_agents = total_agents - online_agents else: enforced_agents = audit_agents = online_agents = offline_agents = 0 system_header = Static("ℹ️ System Overview") system_header.styles.text_style = "bold" system_header.styles.margin = (2, 0, 0, 0) col3.compose_add_child(system_header) system_data = [ f"Agents: {total_agents} ({online_agents} online)", f"Policies: {total_policies}", f"Enforced: {enforced_agents}", f"Audit: {audit_agents}", ] for line in system_data: line_widget = Static(f" {line}") line_widget.styles.margin = (0, 0, 0, 1) col3.compose_add_child(line_widget) # Mount all columns data_area.mount(col1) data_area.mount(col2) data_area.mount(col3) def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None: self.switch_tab(event.tab.id) def on_select_changed(self, event: Select.Changed) -> None: """Handle time period selection change.""" if event.select.id == "stats_time_select": new_days = event.value if new_days != self.current_stats_days: self.current_stats_days = new_days # Fetch if not cached if new_days not in self.statistics_cache: self._fetch_execution_statistics(new_days) else: self._update_statistics_display() def on_multi_agent_selector_agents_selected() -> 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_policy_tree_widget_view_execution_history( self, message: PolicyTreeWidget.ViewExecutionHistory ) -> None: """Handle request to view execution history for a device from tree view.""" logger.info("Viewing execution history for device: %s", message.device.hostname) self.app.push_screen(ExecutionHistoryScreen([message.device])) message.stop() def on_policy_tree_widget_generate_otp( self, message: PolicyTreeWidget.GenerateOTP ) -> None: """Handle request to generate OTP for a device from tree view.""" logger.info("Generating OTP for device: %s", message.device.hostname) self.app.push_screen(OTPWorkflowScreen([message.device])) message.stop() def on_policy_tree_widget_toggle_enforcement( self, message: PolicyTreeWidget.ToggleEnforcement ) -> None: """Handle request to toggle enforcement for a device from tree view.""" logger.info("Toggling enforcement for device: %s", message.device.hostname) try: from services.agenthandler import moveAgentToRelatedPolicy from utils.configmanager import get_system_json policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}") # Determine current mode and toggle if message.device.groupid in policy_relationship_map: # Currently in enforcement, move to audit result = moveAgentToRelatedPolicy(self.app.api, message.device, "audit") mode = "audit" else: # Currently in audit, move to enforcement result = moveAgentToRelatedPolicy( self.app.api, message.device, "enforcement" ) mode = "enforcement" logger.info(f"Successfully toggled {message.device.hostname} to {mode}") # Refresh data at the app level self.app.refresh_data() # Refresh the tree widget with new data try: tree_widget = self.query_one(PolicyTreeWidget) tree_widget.refresh_data(self.app.policies, self.app.devices) except: pass except Exception as e: logger.error( f"Failed to toggle enforcement for {message.device.hostname}: {e}" ) self.app.bell() message.stop() 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) # Handle statistics refresh button if button_id == "stats_refresh_btn": self._fetch_execution_statistics(self.current_stats_days) event.stop() return 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 - Updated to accept attach_notification_handler # --------------------------------------------------------------------------- def run_Loxide(api: AirlockAPIWrapper, attach_notification_handler=None) -> 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) # Attach the notification handler if provided if attach_notification_handler: attach_notification_handler(app) 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 - Updated to get and pass attach_notification_handler # --------------------------------------------------------------------------- def main(): irtang() # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored # setup() now returns a function to attach the notification handler attach_notification_handler = setup() logger = logging.getLogger(__name__) try: url = get_system_value("URL") username = os.getenv("USERNAME") if not url: raise ValueError("Missing URL in environment variables.") if not username: raise ValueError("Missing USERNAME in environment variables.") logger.debug(f"Retrieved URL: {url}") logger.debug(f"Retrieved Username: {username}") except ValueError as e: logger.error(f"Configuration error: {e}", exc_info=True) raise api_key = getAPI(username, "Loxide") if api_key is None: raise ValueError("API key for Loxide is missing.") api = AirlockAPIWrapper( base_url=str(url), api_key=api_key, ) run_Loxide(api, attach_notification_handler) if __name__ == "__main__": main()