Files
AirlockTools/Loxide.py
T
Zarithas fc17c869fc feat: Multiple UI improvements and new server log functionality
- Add Server Log tab with DataTable display of server activity logs

- Fix keyboard navigation bug in agents tab

- Add execution history viewer for selected agents

- Improve policy tree widget functionality by adding single device operations

- Integrate logging notifications into TUI
  - Add TextualNotificationHandler to setup.py
  - Display ERROR/WARNING/CRITICAL logs as toast notifications
  - Remove terminal output to prevent interference with TUI
  - Logs still written to Loxide.log file

closes #45
2025-12-16 16:16:49 -05:00

607 lines
21 KiB
Python

# 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/>.
# 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
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.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.multiagentselector import MultiAgentSelector
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
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("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 == "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_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)
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()