mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

feat: Multiple UI improvements and new server log functionality

Browse Source 
- 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
This commit is contained in:
Zarithas
2025-12-16 16:16:49 -05:00
parent 57d0f12000
commit fc17c869fc
7 changed files with 1183 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TUI/Widgets/agentmoveoperations.py
+69 -23
View File
@@ -29,6 +29,7 @@ from textual.widget import Widget
from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea
from models.agent import Agent
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
from TUI.Screens.policyselectorscreen import PolicySelectorScreen
from TUI.Widgets.OTP_generate import OTPGenerator
@@ -140,6 +141,7 @@ class AgentMoveOperations(Widget):
toggle_enforcement_btn = self.query_one("#toggle_enforcement_btn", Button)
other_policy_btn = self.query_one("#other_policy_btn", Button)
otp_gen_btn = self.query_one("#otp_gen_btn", Button)
exec_history_btn = self.query_one("#exec_history_btn", Button)
# If operation in progress, disable all
if self.operation_in_progress:
@@ -148,6 +150,7 @@ class AgentMoveOperations(Widget):
local_approval_btn.disabled = True
toggle_enforcement_btn.disabled = True
other_policy_btn.disabled = True
exec_history_btn.disabled = True
else:
# If an operation was selected, disable
if self.selected_operation:
@@ -162,6 +165,9 @@ class AgentMoveOperations(Widget):
other_policy_btn.disabled = (
self.selected_operation == "other_policy"
)
exec_history_btn.disabled = (
self.selected_operation == "exec_history"
)
else:
# Enable all buttons
otp_gen_btn = False
@@ -169,6 +175,7 @@ class AgentMoveOperations(Widget):
local_approval_btn.disabled = False
toggle_enforcement_btn.disabled = False
other_policy_btn.disabled = False
exec_history_btn.disabled = False
except NoMatches:
pass
@@ -199,21 +206,21 @@ class AgentMoveOperations(Widget):
f"Operation: {operation_name}",
f"{'=' * 50}",
"",
f"✅ Successful ({len(successful)}):",
f"✅ Successful ({len(successful)}):",
]
if successful:
for agent, result in successful:
results_lines.append(f" ✅ {agent.hostname}")
results_lines.append(f" ✅ {agent.hostname}")
else:
results_lines.append(" (none)")
results_lines.append("")
results_lines.append(f"❌ Failed ({len(unsuccessful)}):")
results_lines.append(f"❌ Failed ({len(unsuccessful)}):")
if unsuccessful:
for agent, error in unsuccessful:
results_lines.append(f" ❌ {agent.hostname}: {error}")
results_lines.append(f" ❌ {agent.hostname}: {error}")
else:
results_lines.append(" (none)")
@@ -248,9 +255,9 @@ class AgentMoveOperations(Widget):
- Operations panel: 1/3 width
- Results area: Initially hidden, shown after operation completion
"""
yield Header(show_clock=True, icon="⚙️")
yield Header(show_clock=True, icon="⚙️")
title_text = Static(
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
id="move_ops_title",
)
title_text.styles.margin = (0, 0, 1, 0)
@@ -285,37 +292,44 @@ class AgentMoveOperations(Widget):
yield operations_label
# Operation buttons
export_csv_btn = Button("📄 Export CSV", id="export_csv_btn")
export_csv_btn = Button("📄 Export CSV", id="export_csv_btn")
export_csv_btn.styles.width = "100%"
export_csv_btn.styles.margin = (0, 0, 1, 0)
yield export_csv_btn
local_approval_btn = Button(
"✔️ Local Approval Mode", id="local_approval_btn"
"✔️ Local Approval Mode", id="local_approval_btn"
)
local_approval_btn.styles.width = "100%"
local_approval_btn.styles.margin = (0, 0, 1, 0)
yield local_approval_btn
otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn")
otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn")
otp_gen_btn.styles.width = "100%"
otp_gen_btn.styles.margin = (0, 0, 1, 0)
yield otp_gen_btn
toggle_enforcement_btn = Button(
"🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn"
"🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn"
)
toggle_enforcement_btn.styles.width = "100%"
toggle_enforcement_btn.styles.margin = (0, 0, 1, 0)
yield toggle_enforcement_btn
other_policy_btn = Button(
"🔀 Move to Other Policy", id="other_policy_btn"
"🔀 Move to Other Policy", id="other_policy_btn"
)
other_policy_btn.styles.width = "100%"
other_policy_btn.styles.margin = (0, 0, 1, 0)
yield other_policy_btn
exec_history_btn = Button(
"📊 View Execution History", id="exec_history_btn"
)
exec_history_btn.styles.width = "100%"
exec_history_btn.styles.margin = (0, 0, 1, 0)
yield exec_history_btn
# Status label
status_label = Static("", id="status_label")
status_label.styles.margin = (2, 0, 0, 0)
@@ -377,17 +391,17 @@ class AgentMoveOperations(Widget):
pyperclip.copy(results_text.text)
self.app.notify(
"📋✅ Results copied to clipboard!",
"📋✅ Results copied to clipboard!",
severity="information",
timeout=2,
)
except ImportError:
self.app.notify(
"❌ pyperclip not installed. Run: pip install pyperclip",
"❌ pyperclip not installed. Run: pip install pyperclip",
severity="warning",
)
except Exception as e:
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
event.stop()
elif btn_id == "export_csv_btn":
self._start_export_csv_operation()
@@ -407,6 +421,9 @@ class AgentMoveOperations(Widget):
elif btn_id == "otp_gen_btn":
self._start_OTP_gen_operation()
event.stop()
elif btn_id == "exec_history_btn":
self._start_execution_history_operation()
event.stop()
def _start_local_approval_operation(self) -> None:
"""
@@ -435,7 +452,7 @@ class AgentMoveOperations(Widget):
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("✔️ Moving agents to local approval...")
status_label.update("✔️ Moving agents to local approval...")
# Get API from app
api = self.app.api
@@ -470,12 +487,12 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error during local approval operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
return
self.operation_in_progress = False
status_label.update("✅ Operation complete!")
status_label.update("✅ Operation complete!")
# Display results in the widget
self._display_results("Local Approval Mode", successful, unsuccessful)
@@ -518,9 +535,9 @@ class AgentMoveOperations(Widget):
file_path = os.path.join(str(path), filename)
df.to_csv(file_path, index=False)
successful.append(file_path)
status_label.update(f"✅ Exported to {file_path}")
status_label.update(f"✅ Exported to {file_path}")
except Exception:
status_label.update("❌ Failed")
status_label.update("❌ Failed")
self.operation_in_progress = False
@@ -564,7 +581,7 @@ class AgentMoveOperations(Widget):
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("🔄 Toggling enforcement mode...")
status_label.update("🔄 Toggling enforcement mode...")
# Get API from app
api = self.app.api
@@ -600,12 +617,12 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error during toggle enforcement operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
return
self.operation_in_progress = False
status_label.update("✅ Operation complete!")
status_label.update("✅ Operation complete!")
# Display results in the widget
self._display_results("Toggle Audit/Enforcement", successful, unsuccessful)
@@ -670,7 +687,7 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error loading policies: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
self.selected_operation = ""
self.app.notify(f"Failed to load policies: {str(e)}", severity="error")
@@ -681,6 +698,35 @@ class AgentMoveOperations(Widget):
self.app.push_screen(OTPWorkflowScreen(self.agents))
def _start_execution_history_operation(self) -> None:
"""
Launch the execution history viewer for selected agents.
This operation opens a new screen that allows the user to:
1. Select a date range for execution history
2. Fetch execution logs for all selected agents
3. View the results in a table
4. Export the results to CSV
The screen is pushed onto the screen stack, allowing the user to return
to this screen when done.
"""
status_label = self.query_one("#status_label", Static)
status_label.update("Opening execution history viewer...")
try:
# Push the execution history screen
self.app.push_screen(ExecutionHistoryScreen(self.agents))
logger.info(
f"Opened execution history viewer for {len(self.agents)} agents"
)
except Exception as e:
logger.error(f"Failed to open execution history viewer: {e}")
status_label.update(f"❌ Error: {str(e)}")
self.app.notify(
f"Failed to open execution history: {str(e)}", severity="error"
)
def _execute_move_to_policy(self, target_policy) -> None:
"""
Execute the actual move of agents to the selected policy.
TUI/Widgets/policytreewidget.py
+100 -1
View File
@@ -18,8 +18,9 @@ import logging
from rich.text import Text
from textual.containers import Horizontal, Vertical
from textual.message import Message
from textual.widget import Widget
from textual.widgets import Input, OptionList, Static, Switch, Tree
from textual.widgets import Button, Input, OptionList, Static, Switch, Tree
from textual.widgets.option_list import Option
logger = logging.getLogger(__name__)
@@ -28,6 +29,27 @@ logger = logging.getLogger(__name__)
class PolicyTreeWidget(Widget):
"""Widget for displaying and searching a hierarchical policy tree."""
class ViewExecutionHistory(Message):
"""Message sent when user wants to view execution history for a device."""
def __init__(self, device):
super().__init__()
self.device = device
class GenerateOTP(Message):
"""Message sent when user wants to generate OTP for a device."""
def __init__(self, device):
super().__init__()
self.device = device
class ToggleEnforcement(Message):
"""Message sent when user wants to toggle audit/enforcement for a device."""
def __init__(self, device):
super().__init__()
self.device = device
def __init__(self, policies, devices):
super().__init__()
self.policies = policies
@@ -35,6 +57,7 @@ class PolicyTreeWidget(Widget):
self.last_highlighted_node = None
self.leaf_counts = defaultdict(int)
self.match_type = "Count" # Default to sorting by count
self.selected_device = None # Track currently selected device
def compose(self):
# Create the switch and its label
@@ -56,6 +79,18 @@ class PolicyTreeWidget(Widget):
search_box = Input(
placeholder="Search policies or devices...", id="tree_search"
)
exec_history_button = Button(
"📊 Execution History", id="view_exec_history_button", disabled=True
)
exec_history_button.styles.margin = (0, 1, 0, 0) # Right margin
otp_button = Button("🎫 Generate OTP", id="generate_otp_button", disabled=True)
otp_button.styles.margin = (0, 1, 0, 0) # Right margin
toggle_enforcement_button = Button(
"🔄 Toggle Enforcement/Audit", id="toggle_enforcement_button", disabled=True
)
# No right margin on last button
details_pane = Static("", id="details_pane")
# Layout the UI
@@ -73,6 +108,12 @@ class PolicyTreeWidget(Widget):
# Add the search box and details pane
yield label
yield search_box
# Action buttons in a horizontal row
with Horizontal() as button_row:
button_row.styles.height = "auto"
yield exec_history_button
yield otp_button
yield toggle_enforcement_button
yield details_pane
def on_mount(self) -> None:
@@ -89,6 +130,32 @@ class PolicyTreeWidget(Widget):
# Expand the root node
policy_tree.root.expand()
def refresh_data(self, policies, devices):
"""Refresh the widget with new data and rebuild the tree."""
self.policies = policies
self.devices = devices
self.selected_device = None
# Disable all buttons since selection is lost
try:
self.query_one("#view_exec_history_button", Button).disabled = True
self.query_one("#generate_otp_button", Button).disabled = True
self.query_one("#toggle_enforcement_button", Button).disabled = True
except:
pass
# Rebuild tree with new data
self._precompute_leaf_counts()
total_leaves = sum(
self.leaf_counts.get(policy.groupid, 0)
for policy in self.policies
if policy.parent == "global-policy-settings"
)
policy_tree = self.query_one("#policy_tree", Tree)
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
self._build_tree()
policy_tree.root.expand()
def _precompute_leaf_counts(self):
"""Precompute leaf counts for each policy group."""
device_counts = defaultdict(int)
@@ -184,6 +251,9 @@ class PolicyTreeWidget(Widget):
node = message.node
data = node.data
details_pane = self.query_one("#details_pane", Static)
exec_history_button = self.query_one("#view_exec_history_button", Button)
otp_button = self.query_one("#generate_otp_button", Button)
toggle_enforcement_button = self.query_one("#toggle_enforcement_button", Button)
if self.last_highlighted_node is not None:
original_label = str(self.last_highlighted_node.label).strip()
@@ -198,6 +268,20 @@ class PolicyTreeWidget(Widget):
node.set_label(highlighted_label)
self.last_highlighted_node = node
# Check if selected node is a device (has Agent data)
from models.agent import Agent
if data and isinstance(data, Agent):
self.selected_device = data
exec_history_button.disabled = False
otp_button.disabled = False
toggle_enforcement_button.disabled = False
else:
self.selected_device = None
exec_history_button.disabled = True
otp_button.disabled = True
toggle_enforcement_button.disabled = True
if data:
details = "\n".join(
f"{key}: {value}" for key, value in data.__dict__.items()
@@ -287,3 +371,18 @@ class PolicyTreeWidget(Widget):
option_list.remove()
except:
pass
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button presses."""
if event.button.id == "view_exec_history_button":
if self.selected_device:
self.post_message(self.ViewExecutionHistory(self.selected_device))
event.stop()
elif event.button.id == "generate_otp_button":
if self.selected_device:
self.post_message(self.GenerateOTP(self.selected_device))
event.stop()
elif event.button.id == "toggle_enforcement_button":
if self.selected_device:
self.post_message(self.ToggleEnforcement(self.selected_device))
event.stop()
TUI/Widgets/serverlogwidget.py
+159
View File
@@ -0,0 +1,159 @@
# 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 datetime
import logging
from bson import ObjectId
from textual.app import ComposeResult
from textual.containers import Container, Vertical
from textual.widgets import Button, DataTable, Static
from services.API import AirlockAPIWrapper
logger = logging.getLogger(__name__)
def skipback(days):
"""
Generate a MongoDB ObjectId for a given number of days ago from today.
"""
adjusted_days = days
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
days=adjusted_days
)
timestamp = int(date_days_ago.timestamp())
hex_timestamp = format(timestamp, "08x")
objectid_hex = hex_timestamp + "0000000000000000"
return ObjectId(objectid_hex)
class ServerLogWidget(Vertical):
"""Widget for displaying server activity logs in a DataTable."""
DEFAULT_CSS = """
ServerLogWidget {
width: 100%;
height: 100%;
}
ServerLogWidget #status_bar {
width: 100%;
height: auto;
background: $surface;
padding: 1;
margin-bottom: 1;
}
ServerLogWidget DataTable {
height: 1fr;
border: solid $primary;
}
ServerLogWidget #button_container {
width: 100%;
height: auto;
layout: horizontal;
padding: 1;
}
ServerLogWidget Button {
margin-right: 1;
}
"""
def __init__(self, api: AirlockAPIWrapper):
super().__init__()
self.api = api
def compose(self) -> ComposeResult:
yield Static("Loading server logs (last 72 hours)...", id="status_bar")
yield DataTable(id="server_log_table")
with Container(id="button_container"):
yield Button("Refresh", id="refresh_button", variant="primary")
def on_mount(self) -> None:
"""Initialize the DataTable and load server logs."""
self.load_logs()
def load_logs(self) -> None:
"""Load server logs from the API and populate the DataTable."""
table = self.query_one("#server_log_table", DataTable)
status = self.query_one("#status_bar", Static)
try:
status.update("⏳ Loading server logs (last 72 hours)...")
# Create a fake checkpoint for 3 days ago (72 hours)
checkpoint = str(skipback(3))
# Get server logs from API
logs = self.api.server_logs(checkpoint=checkpoint)
if not logs:
status.update("ℹï¸ No server logs found in the last 72 hours.")
table.clear(columns=True)
return
# Clear existing data
table.clear(columns=True)
# Add columns based on the first log entry
if logs:
first_log = logs[0]
columns = [col for col in first_log.keys() if col != "checkpoint"]
for col in columns:
table.add_column(col, key=col)
# Add rows in reverse order so newest entries are at the top
for log_entry in reversed(logs):
row_data = []
for col in columns:
value = log_entry.get(col, "")
# Format datetime column to be more readable
if col == "datetime" and value:
try:
# Parse ISO format and convert to readable format
dt = datetime.datetime.fromisoformat(
str(value).replace("Z", "+00:00")
)
value = dt.strftime("%Y-%m-%d %H:%M:%S")
except Exception:
# If parsing fails, just use the original value
pass
row_data.append(str(value))
table.add_row(*row_data)
status.update(
f"✅ Loaded {len(logs)} log entries from the last 72 hours"
)
logger.info(f"Loaded {len(logs)} server log entries")
else:
status.update("ℹï¸ No log entries found.")
except Exception as exc:
error_msg = f"❌ Error loading server logs: {exc}"
status.update(error_msg)
logger.error(f"Failed to load server logs: {exc}", exc_info=True)
table.clear(columns=True)
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button presses."""
button_id = event.button.id
if button_id == "refresh_button":
self.load_logs()
event.stop()