Compare commits
24 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 72681218e7 | |||
| 5555747422 | |||
| 729b45f52a | |||
| 66bb21ed88 | |||
| a7b659c951 | |||
| 53f0b548b0 | |||
| 66eb101c5d | |||
| 22101c1eba | |||
| fc17c869fc | |||
| 1dbbcff5d5 | |||
| f080b0034f | |||
| 57d0f12000 | |||
| 0dbc744471 | |||
| 7a912bddab | |||
| 24211c318b | |||
| 630e0a3cdf | |||
| 797d0f4462 | |||
| 59bb97ec4e | |||
| 0ac3b54d89 | |||
| 154a7efcc8 | |||
| ab5f00d8e7 | |||
| 3ab803c12e | |||
| 98cb23e5ea | |||
| 0aabbfd36e |
@@ -23,23 +23,556 @@
|
||||
|
||||
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.TUI import run_Loxide
|
||||
from utils.configmanager import get_system_value
|
||||
from utils.setup import setup
|
||||
from utils.utils import irtang
|
||||
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()
|
||||
# setup() now returns a function to attach the notification handler
|
||||
attach_notification_handler = setup()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
@@ -66,7 +599,7 @@ def main():
|
||||
base_url=str(url),
|
||||
api_key=api_key,
|
||||
)
|
||||
run_Loxide(api)
|
||||
run_Loxide(api, attach_notification_handler)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,639 @@
|
||||
# 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/>.
|
||||
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
import logging
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
import pandas as pd
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import (
|
||||
Button,
|
||||
DataTable,
|
||||
Footer,
|
||||
Header,
|
||||
Label,
|
||||
Select,
|
||||
Static,
|
||||
)
|
||||
|
||||
from models.agent import Agent
|
||||
from models.execution import ExecutionHistoryRecord
|
||||
from utils.configmanager import load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ExecutionHistoryScreen(Screen):
|
||||
"""
|
||||
A screen for viewing and exporting execution history for selected agents.
|
||||
|
||||
This screen allows users to:
|
||||
1. Select a start date and end date using dropdown selects
|
||||
2. Fetch execution history for all selected agents
|
||||
3. View the results in a DataTable
|
||||
4. Export the results to CSV using a keybinding
|
||||
|
||||
Attributes:
|
||||
agents (List[Agent]): List of agents to fetch execution history for
|
||||
execution_data (pd.DataFrame): Combined execution history data
|
||||
working_dir (str): Directory for CSV exports
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
ExecutionHistoryScreen {
|
||||
align: center top;
|
||||
}
|
||||
|
||||
#main_container {
|
||||
width: 95%;
|
||||
height: 1fr;
|
||||
border: solid $primary;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
#title {
|
||||
text-style: bold;
|
||||
color: $text;
|
||||
text-align: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#date_container {
|
||||
height: auto;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#start_date_row, #end_date_row {
|
||||
height: auto;
|
||||
align-horizontal: left;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
.date_label {
|
||||
width: 8;
|
||||
margin-right: 1;
|
||||
}
|
||||
|
||||
.date_selector {
|
||||
width: 18;
|
||||
margin: 0 1;
|
||||
}
|
||||
|
||||
#quick_buttons_row {
|
||||
height: auto;
|
||||
align-horizontal: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
.quick_select_btn {
|
||||
margin: 0 1;
|
||||
}
|
||||
|
||||
#button_row {
|
||||
height: auto;
|
||||
align-horizontal: center;
|
||||
margin-top: 1;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
Button {
|
||||
margin: 0 1;
|
||||
}
|
||||
|
||||
#status_label {
|
||||
text-align: center;
|
||||
color: $accent;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#results_container {
|
||||
height: 1fr;
|
||||
display: none;
|
||||
}
|
||||
|
||||
#results_button_row {
|
||||
height: auto;
|
||||
align-horizontal: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#history_table {
|
||||
height: 1fr;
|
||||
border: solid $primary;
|
||||
}
|
||||
|
||||
DataTable > .datatable--header {
|
||||
text-style: bold;
|
||||
background: $primary 20%;
|
||||
}
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "close_screen", "Close"),
|
||||
Binding("e", "export_csv", "Export CSV"),
|
||||
Binding("q", "close_screen", "Quit"),
|
||||
]
|
||||
|
||||
def __init__(self, agents: List[Agent]):
|
||||
"""
|
||||
Initialize the ExecutionHistoryScreen.
|
||||
|
||||
Args:
|
||||
agents (List[Agent]): List of agents to fetch execution history for
|
||||
"""
|
||||
super().__init__()
|
||||
self.agents = agents
|
||||
self.execution_data = pd.DataFrame()
|
||||
self.working_dir = load_env("WORKING_DIR") or os.getcwd()
|
||||
|
||||
# Generate dropdown options
|
||||
today = datetime.now().date()
|
||||
|
||||
# Month options - format is (display_text, value)
|
||||
self.month_options = [
|
||||
("January", "01"),
|
||||
("February", "02"),
|
||||
("March", "03"),
|
||||
("April", "04"),
|
||||
("May", "05"),
|
||||
("June", "06"),
|
||||
("July", "07"),
|
||||
("August", "08"),
|
||||
("September", "09"),
|
||||
("October", "10"),
|
||||
("November", "11"),
|
||||
("December", "12"),
|
||||
]
|
||||
|
||||
# Day options (1-31) - format is (display_text, value)
|
||||
self.day_options = [(f"{i}", f"{i:02d}") for i in range(1, 32)]
|
||||
|
||||
# Year options (current year back 5 years) - format is (display_text, value)
|
||||
current_year = today.year
|
||||
self.year_options = [
|
||||
(str(year), str(year)) for year in range(current_year, current_year - 6, -1)
|
||||
]
|
||||
|
||||
# Default dates: last 30 days
|
||||
start_date = today - timedelta(days=30)
|
||||
self.start_month = f"{start_date.month:02d}"
|
||||
self.start_day = f"{start_date.day:02d}"
|
||||
self.start_year = str(start_date.year)
|
||||
|
||||
self.end_month = f"{today.month:02d}"
|
||||
self.end_day = f"{today.day:02d}"
|
||||
self.end_year = str(today.year)
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Build the UI layout."""
|
||||
yield Header(show_clock=True, icon="📊")
|
||||
|
||||
with Vertical(id="main_container"):
|
||||
title_text = f"Execution History - {len(self.agents)} Agent(s)"
|
||||
yield Static(title_text, id="title")
|
||||
|
||||
# Date selection area
|
||||
with Vertical(id="date_container"):
|
||||
yield Label("Select Date Range:")
|
||||
|
||||
# Start date row
|
||||
with Horizontal(id="start_date_row"):
|
||||
yield Label("From:", classes="date_label")
|
||||
yield Select(
|
||||
options=self.month_options,
|
||||
value=self.start_month,
|
||||
id="start_month_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
yield Select(
|
||||
options=self.day_options,
|
||||
value=self.start_day,
|
||||
id="start_day_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
yield Select(
|
||||
options=self.year_options,
|
||||
value=self.start_year,
|
||||
id="start_year_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
|
||||
# End date row
|
||||
with Horizontal(id="end_date_row"):
|
||||
yield Label("To:", classes="date_label")
|
||||
yield Select(
|
||||
options=self.month_options,
|
||||
value=self.end_month,
|
||||
id="end_month_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
yield Select(
|
||||
options=self.day_options,
|
||||
value=self.end_day,
|
||||
id="end_day_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
yield Select(
|
||||
options=self.year_options,
|
||||
value=self.end_year,
|
||||
id="end_year_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
|
||||
# Quick select buttons
|
||||
with Horizontal(id="quick_buttons_row"):
|
||||
yield Button(
|
||||
"1 Day",
|
||||
id="quick_1day",
|
||||
classes="quick_select_btn",
|
||||
variant="default",
|
||||
)
|
||||
yield Button(
|
||||
"1 Week",
|
||||
id="quick_1week",
|
||||
classes="quick_select_btn",
|
||||
variant="default",
|
||||
)
|
||||
yield Button(
|
||||
"30 Days",
|
||||
id="quick_30days",
|
||||
classes="quick_select_btn",
|
||||
variant="default",
|
||||
)
|
||||
|
||||
# Buttons
|
||||
with Horizontal(id="button_row"):
|
||||
yield Button("Fetch History", id="fetch_btn", variant="primary")
|
||||
yield Button("Close", id="close_btn", variant="error")
|
||||
|
||||
# Status
|
||||
yield Static(
|
||||
"Select date range and click 'Fetch History'", id="status_label"
|
||||
)
|
||||
|
||||
# Results container (hidden initially, shown after fetch)
|
||||
with Vertical(id="results_container"):
|
||||
with Horizontal(id="results_button_row"):
|
||||
yield Button("Export CSV", id="export_btn", variant="success")
|
||||
yield Button("Back", id="back_btn", variant="default")
|
||||
yield DataTable(id="history_table")
|
||||
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Initialize the table when screen is mounted."""
|
||||
table = self.query_one("#history_table", DataTable)
|
||||
table.cursor_type = "row"
|
||||
table.zebra_stripes = True
|
||||
|
||||
# Initially empty - will populate after fetch
|
||||
logger.info(f"ExecutionHistoryScreen mounted with {len(self.agents)} agents")
|
||||
|
||||
def on_select_changed(self, event: Select.Changed) -> None:
|
||||
"""Handle date selection changes."""
|
||||
select_id = event.select.id
|
||||
|
||||
if select_id == "start_month_select":
|
||||
self.start_month = event.value
|
||||
logger.debug(f"Start month changed to: {self.start_month}")
|
||||
elif select_id == "start_day_select":
|
||||
self.start_day = event.value
|
||||
logger.debug(f"Start day changed to: {self.start_day}")
|
||||
elif select_id == "start_year_select":
|
||||
self.start_year = event.value
|
||||
logger.debug(f"Start year changed to: {self.start_year}")
|
||||
elif select_id == "end_month_select":
|
||||
self.end_month = event.value
|
||||
logger.debug(f"End month changed to: {self.end_month}")
|
||||
elif select_id == "end_day_select":
|
||||
self.end_day = event.value
|
||||
logger.debug(f"End day changed to: {self.end_day}")
|
||||
elif select_id == "end_year_select":
|
||||
self.end_year = event.value
|
||||
logger.debug(f"End year changed to: {self.end_year}")
|
||||
|
||||
def _set_quick_date_range(self, days: int) -> None:
|
||||
"""Set the date range based on quick select button."""
|
||||
today = datetime.now().date()
|
||||
start_date = today - timedelta(days=days)
|
||||
|
||||
# Update internal values
|
||||
self.start_month = f"{start_date.month:02d}"
|
||||
self.start_day = f"{start_date.day:02d}"
|
||||
self.start_year = str(start_date.year)
|
||||
|
||||
self.end_month = f"{today.month:02d}"
|
||||
self.end_day = f"{today.day:02d}"
|
||||
self.end_year = str(today.year)
|
||||
|
||||
# Update the Select widgets
|
||||
try:
|
||||
self.query_one("#start_month_select", Select).value = self.start_month
|
||||
self.query_one("#start_day_select", Select).value = self.start_day
|
||||
self.query_one("#start_year_select", Select).value = self.start_year
|
||||
|
||||
self.query_one("#end_month_select", Select).value = self.end_month
|
||||
self.query_one("#end_day_select", Select).value = self.end_day
|
||||
self.query_one("#end_year_select", Select).value = self.end_year
|
||||
|
||||
self.app.notify(
|
||||
f"Date range set to last {days} day(s)",
|
||||
severity="information",
|
||||
timeout=2,
|
||||
)
|
||||
logger.info(f"Quick select: Set date range to last {days} days")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update date selects: {e}")
|
||||
|
||||
def _show_date_selection(self) -> None:
|
||||
"""Show the date selection view and hide results."""
|
||||
try:
|
||||
self.query_one("#date_container").styles.display = "block"
|
||||
self.query_one("#button_row").styles.display = "block"
|
||||
self.query_one("#status_label").styles.display = "block"
|
||||
self.query_one("#results_container").styles.display = "none"
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to show date selection: {e}")
|
||||
|
||||
def _show_results(self) -> None:
|
||||
"""Hide date selection view and show results."""
|
||||
try:
|
||||
self.query_one("#date_container").styles.display = "none"
|
||||
self.query_one("#button_row").styles.display = "none"
|
||||
self.query_one("#status_label").styles.display = "none"
|
||||
self.query_one("#results_container").styles.display = "block"
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to show results: {e}")
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button clicks."""
|
||||
if event.button.id == "fetch_btn":
|
||||
self._fetch_execution_history()
|
||||
elif event.button.id == "export_btn":
|
||||
self._export_to_csv()
|
||||
elif event.button.id == "close_btn":
|
||||
self.app.pop_screen()
|
||||
elif event.button.id == "back_btn":
|
||||
self._show_date_selection()
|
||||
elif event.button.id == "quick_1day":
|
||||
self._set_quick_date_range(days=1)
|
||||
elif event.button.id == "quick_1week":
|
||||
self._set_quick_date_range(days=7)
|
||||
elif event.button.id == "quick_30days":
|
||||
self._set_quick_date_range(days=30)
|
||||
|
||||
def _fetch_execution_history(self) -> None:
|
||||
"""Fetch execution history for all selected agents."""
|
||||
status_label = self.query_one("#status_label", Static)
|
||||
status_label.update("â³ Fetching execution history...")
|
||||
|
||||
# Disable buttons during fetch
|
||||
fetch_btn = self.query_one("#fetch_btn", Button)
|
||||
export_btn = self.query_one("#export_btn", Button)
|
||||
fetch_btn.disabled = True
|
||||
export_btn.disabled = True
|
||||
|
||||
api = self.app.api
|
||||
all_history = []
|
||||
|
||||
try:
|
||||
# Construct dates from dropdowns
|
||||
start_date_str = f"{self.start_year}-{self.start_month}-{self.start_day}"
|
||||
end_date_str = f"{self.end_year}-{self.end_month}-{self.end_day}"
|
||||
|
||||
# Validate dates
|
||||
try:
|
||||
start_dt = datetime.strptime(start_date_str, "%Y-%m-%d")
|
||||
end_dt = datetime.strptime(end_date_str, "%Y-%m-%d")
|
||||
except ValueError as e:
|
||||
status_label.update(f"⌠Invalid date: {str(e)}")
|
||||
fetch_btn.disabled = False
|
||||
export_btn.disabled = False
|
||||
self.app.notify(f"Invalid date selected: {str(e)}", severity="error")
|
||||
return
|
||||
|
||||
if start_dt > end_dt:
|
||||
status_label.update("⌠Error: Start date must be before end date")
|
||||
fetch_btn.disabled = False
|
||||
export_btn.disabled = False
|
||||
return
|
||||
|
||||
# Fetch history for each agent
|
||||
for i, agent in enumerate(self.agents):
|
||||
try:
|
||||
status_label.update(
|
||||
f"â³ Fetching history for {agent.hostname} ({i+1}/{len(self.agents)})..."
|
||||
)
|
||||
|
||||
# Call API - note the API expects 'dateto' first, then 'datefrom'
|
||||
history = api.history_execution(
|
||||
today=end_date_str,
|
||||
date_selected=start_date_str,
|
||||
agent_name=agent.hostname,
|
||||
)
|
||||
|
||||
if history:
|
||||
# Add agent hostname to each record for identification
|
||||
for record in history:
|
||||
record["agent_hostname"] = agent.hostname
|
||||
all_history.extend(history)
|
||||
logger.info(
|
||||
f"Fetched {len(history)} records for {agent.hostname}"
|
||||
)
|
||||
else:
|
||||
logger.info(f"No history found for {agent.hostname}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch history for {agent.hostname}: {e}")
|
||||
self.app.notify(
|
||||
f"Warning: Failed to fetch history for {agent.hostname}",
|
||||
severity="warning",
|
||||
)
|
||||
|
||||
# Convert to DataFrame
|
||||
if all_history:
|
||||
status_label.update(
|
||||
"â³ Enriching execution data with hash information..."
|
||||
)
|
||||
|
||||
# Normalize field names (handle API typos)
|
||||
for record in all_history:
|
||||
if "policver" in record and "policyver" not in record:
|
||||
record["policyver"] = record.pop("policver")
|
||||
|
||||
# Convert dict records to ExecutionHistoryRecord objects
|
||||
execution_records = []
|
||||
for record in all_history:
|
||||
try:
|
||||
execution_records.append(ExecutionHistoryRecord(**record))
|
||||
except TypeError as e:
|
||||
logger.warning(f"Failed to create ExecutionHistoryRecord: {e}")
|
||||
# If it fails, just keep the dict
|
||||
continue
|
||||
|
||||
# Enrich with hash data if we have ExecutionHistoryRecord objects
|
||||
if execution_records:
|
||||
try:
|
||||
enriched_records = ExecutionHistoryRecord.enrich_with_hashes(
|
||||
api, execution_records
|
||||
)
|
||||
logger.info(
|
||||
f"Enriched {len(enriched_records)} records with hash data"
|
||||
)
|
||||
|
||||
# Convert back to DataFrame
|
||||
self.execution_data = pd.DataFrame(
|
||||
[r.__dict__ for r in enriched_records]
|
||||
)
|
||||
|
||||
# Flatten hash_obj if present
|
||||
if (
|
||||
not self.execution_data.empty
|
||||
and "hash_obj" in self.execution_data.columns
|
||||
):
|
||||
hash_df = self.execution_data["hash_obj"].apply(
|
||||
lambda h: (
|
||||
h.to_dict() if h and hasattr(h, "to_dict") else {}
|
||||
)
|
||||
)
|
||||
self.execution_data = pd.concat(
|
||||
[
|
||||
self.execution_data.drop(columns=["hash_obj"]),
|
||||
hash_df,
|
||||
],
|
||||
axis=1,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to enrich with hashes: {e}")
|
||||
# Fall back to plain DataFrame
|
||||
self.execution_data = pd.DataFrame(all_history)
|
||||
else:
|
||||
# If we couldn't create any ExecutionHistoryRecord objects, just use raw data
|
||||
self.execution_data = pd.DataFrame(all_history)
|
||||
|
||||
self._populate_table()
|
||||
self._show_results() # Switch to results view
|
||||
self.app.notify(
|
||||
f"Successfully loaded {len(self.execution_data)} records",
|
||||
severity="information",
|
||||
)
|
||||
else:
|
||||
status_label.update(
|
||||
"â„¹ï¸ No execution history found for selected agents/dates"
|
||||
)
|
||||
self.app.notify("No execution history found", severity="information")
|
||||
self.execution_data = pd.DataFrame()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error fetching execution history: {e}")
|
||||
status_label.update(f"⌠Error: {str(e)}")
|
||||
self.app.notify(f"Failed to fetch history: {str(e)}", severity="error")
|
||||
|
||||
finally:
|
||||
# Re-enable buttons
|
||||
fetch_btn.disabled = False
|
||||
export_btn.disabled = False
|
||||
|
||||
def _populate_table(self) -> None:
|
||||
"""Populate the DataTable with execution history data."""
|
||||
table = self.query_one("#history_table", DataTable)
|
||||
table.clear(columns=True)
|
||||
|
||||
if self.execution_data.empty:
|
||||
return
|
||||
|
||||
# Define preferred column order (your specified order)
|
||||
preferred_order = [
|
||||
"policyname",
|
||||
"policyver",
|
||||
"hostname",
|
||||
"username",
|
||||
"publisher",
|
||||
"filename",
|
||||
"pprocess",
|
||||
"gprocess",
|
||||
"sha256",
|
||||
"commandline",
|
||||
"agent_hostname", # Our custom field
|
||||
]
|
||||
|
||||
# Get available columns in preferred order, then add any remaining columns
|
||||
available_cols = []
|
||||
for col in preferred_order:
|
||||
if col in self.execution_data.columns:
|
||||
available_cols.append(col)
|
||||
|
||||
# Add any remaining columns not in preferred order
|
||||
for col in self.execution_data.columns:
|
||||
if col not in available_cols:
|
||||
available_cols.append(col)
|
||||
|
||||
# Add columns to table
|
||||
for col in available_cols:
|
||||
table.add_column(col, key=col)
|
||||
|
||||
# Add rows
|
||||
for idx, row in self.execution_data.iterrows():
|
||||
row_data = []
|
||||
for col in available_cols:
|
||||
value = row[col]
|
||||
# Convert to string, handle None/NaN
|
||||
if pd.isna(value):
|
||||
row_data.append("")
|
||||
else:
|
||||
row_data.append(str(value))
|
||||
table.add_row(*row_data, key=str(idx))
|
||||
|
||||
logger.info(f"Populated table with {len(self.execution_data)} rows")
|
||||
|
||||
def _export_to_csv(self) -> None:
|
||||
"""Export the current execution data to CSV."""
|
||||
if self.execution_data.empty:
|
||||
self.app.notify("No data to export", severity="warning")
|
||||
return
|
||||
|
||||
try:
|
||||
# Create filename with timestamp
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"execution_history_{timestamp}.csv"
|
||||
filepath = os.path.join(self.working_dir, filename)
|
||||
|
||||
# Export to CSV
|
||||
self.execution_data.to_csv(filepath, index=False, encoding="utf-8-sig")
|
||||
|
||||
self.app.notify(
|
||||
f"✅ Exported {len(self.execution_data)} records to: {filepath}",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
logger.info(f"Exported execution history to: {filepath}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to export CSV: {e}")
|
||||
self.app.notify(f"Failed to export CSV: {str(e)}", severity="error")
|
||||
|
||||
def action_export_csv(self) -> None:
|
||||
"""Keybinding action to export CSV."""
|
||||
self._export_to_csv()
|
||||
|
||||
def action_close_screen(self) -> None:
|
||||
"""Close this screen and return to previous."""
|
||||
self.app.pop_screen()
|
||||
@@ -58,9 +58,14 @@ class OTPRevokeWidget(Static):
|
||||
}
|
||||
#button_container {
|
||||
height: auto;
|
||||
width: 100%;
|
||||
padding: 1;
|
||||
align: center middle;
|
||||
}
|
||||
#button_container Button {
|
||||
min-width: 16;
|
||||
margin: 0 1;
|
||||
}
|
||||
#result_container {
|
||||
height: auto;
|
||||
max-height: 10;
|
||||
@@ -87,29 +92,10 @@ class OTPRevokeWidget(Static):
|
||||
|
||||
# Action buttons
|
||||
with Horizontal(id="button_container"):
|
||||
self.refresh_button = Button("🔄 Refresh", id="refresh_btn")
|
||||
self.refresh_button.styles.width = "15%"
|
||||
self.refresh_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.refresh_button
|
||||
|
||||
self.select_all_button = Button("☑️ Select All", id="select_all_btn")
|
||||
self.select_all_button.styles.width = "15%"
|
||||
self.select_all_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.select_all_button
|
||||
|
||||
self.select_none_button = Button(
|
||||
"❌ Clear Selection", id="select_none_btn"
|
||||
)
|
||||
self.select_none_button.styles.width = "20%"
|
||||
self.select_none_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.select_none_button
|
||||
|
||||
self.revoke_button = Button(
|
||||
"🛑 Revoke Selected", id="revoke_btn", variant="error"
|
||||
)
|
||||
self.revoke_button.styles.width = "20%"
|
||||
self.revoke_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.revoke_button
|
||||
yield Button("Refresh", id="refresh_btn")
|
||||
yield Button("Select All", id="select_all_btn")
|
||||
yield Button("Clear Selection", id="select_none_btn")
|
||||
yield Button("Revoke Selected", id="revoke_btn", variant="error")
|
||||
|
||||
# Results display
|
||||
with Vertical(id="result_container"):
|
||||
@@ -122,7 +108,7 @@ class OTPRevokeWidget(Static):
|
||||
# Configure sessions table
|
||||
self.sessions_table.clear()
|
||||
self.sessions_table.add_columns(
|
||||
"☐", "OTP ID", "Hostname", "Status", "Purpose", "Granted"
|
||||
"", "OTP ID", "Hostname", "Status", "Purpose", "Granted"
|
||||
)
|
||||
|
||||
# Enable row selection with checkbox column
|
||||
@@ -213,7 +199,11 @@ class OTPRevokeWidget(Static):
|
||||
|
||||
elif btn.id == "select_all_btn":
|
||||
# Select all visible rows
|
||||
if self._filtered_df is not None:
|
||||
if (
|
||||
self._filtered_df is not None
|
||||
and not self._filtered_df.empty
|
||||
and "otpid" in self._filtered_df.columns
|
||||
):
|
||||
self._selected_otpids = set(str(x) for x in self._filtered_df["otpid"])
|
||||
await self._refresh_table()
|
||||
|
||||
@@ -235,7 +225,12 @@ class OTPRevokeWidget(Static):
|
||||
# Get the row index from the cursor row
|
||||
row_index = self.sessions_table.cursor_row
|
||||
|
||||
if self._filtered_df is not None and row_index < len(self._filtered_df):
|
||||
if (
|
||||
self._filtered_df is not None
|
||||
and not self._filtered_df.empty
|
||||
and "otpid" in self._filtered_df.columns
|
||||
and row_index < len(self._filtered_df)
|
||||
):
|
||||
# Get the OTP ID for this row
|
||||
otpid = str(self._filtered_df.iloc[row_index]["otpid"])
|
||||
|
||||
@@ -257,12 +252,12 @@ class OTPRevokeWidget(Static):
|
||||
async def _revoke_selected(self) -> None:
|
||||
"""Revoke the selected OTP sessions."""
|
||||
if not self._selected_otpids:
|
||||
self.results_display.update("❌ No sessions selected for revocation")
|
||||
self.results_display.update("No sessions selected for revocation")
|
||||
return
|
||||
|
||||
api = getattr(self.app, "api", None)
|
||||
if not api:
|
||||
self.results_display.update("❌ API not available")
|
||||
self.results_display.update("API not available")
|
||||
return
|
||||
|
||||
# Collect results
|
||||
@@ -303,13 +298,13 @@ class OTPRevokeWidget(Static):
|
||||
else "No response"
|
||||
)
|
||||
results.append(
|
||||
f"❌ Failed to revoke OTP {otpid} for {hostname}: {error_msg}"
|
||||
f"Failed to revoke OTP {otpid} for {hostname}: {error_msg}"
|
||||
)
|
||||
logger.error(f"Failed to revoke OTP {otpid}: {error_msg}")
|
||||
|
||||
except Exception as e:
|
||||
failure_count += 1
|
||||
results.append(f"❌ Error revoking OTP {otpid}: {str(e)}")
|
||||
results.append(f"Error revoking OTP {otpid}: {str(e)}")
|
||||
logger.exception(f"Exception revoking OTP {otpid}: {e}")
|
||||
|
||||
# Update results display
|
||||
@@ -368,7 +363,11 @@ class OTPRevokeScreen(Screen):
|
||||
|
||||
async def action_select_all(self) -> None:
|
||||
"""Select all visible sessions."""
|
||||
if self.widget._filtered_df is not None:
|
||||
if (
|
||||
self.widget._filtered_df is not None
|
||||
and not self.widget._filtered_df.empty
|
||||
and "otpid" in self.widget._filtered_df.columns
|
||||
):
|
||||
self.widget._selected_otpids = set(
|
||||
str(x) for x in self.widget._filtered_df["otpid"]
|
||||
)
|
||||
|
||||
+2237
-362
File diff suppressed because it is too large
Load Diff
@@ -34,7 +34,7 @@ from textual.app import ComposeResult
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.reactive import reactive
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Static
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Input, Static
|
||||
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
@@ -51,16 +51,17 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
|
||||
This screen provides a multi-step workflow:
|
||||
1. Select initial policy to analyze
|
||||
2. View categorized agents (enforce ready vs. non-enforce ready)
|
||||
3. Select target policies for each category
|
||||
4. Execute agent migrations
|
||||
2. Configure analysis parameters (history period and quiet time period)
|
||||
3. View categorized agents (enforce ready vs. non-enforce ready)
|
||||
4. Select target policies for each category
|
||||
5. Execute agent migrations
|
||||
|
||||
Attributes:
|
||||
api (AirlockAPIWrapper): API wrapper for Airlock operations
|
||||
policies (List[Policy]): List of all available policies
|
||||
selected_policy (Optional[Policy]): The initially selected policy to analyze
|
||||
history_days (int): Number of days of history to pull (default: 150)
|
||||
quiet_days (int): Number of days without execution to be considered quiet (default: 45)
|
||||
history_days (int): Number of days of history to pull (default: 150, range: 1-365)
|
||||
quiet_days (int): Number of days without execution to be considered quiet (default: 45, range: 1-365)
|
||||
agents_df (Optional[pd.DataFrame]): DataFrame of all agents with analysis results
|
||||
enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents ready for enforcement
|
||||
non_enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents not ready for enforcement
|
||||
@@ -86,7 +87,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
self.api = api
|
||||
self.policies = policies
|
||||
self.selected_policy: Optional[Policy] = None
|
||||
self.history_days = 150 # Fixed as per requirements
|
||||
self.history_days = 150 # Default value, user-selectable
|
||||
self.quiet_days = 45 # Default value
|
||||
self.agents_df: Optional[pd.DataFrame] = None
|
||||
self.enforce_ready_df: Optional[pd.DataFrame] = None
|
||||
@@ -130,7 +131,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
|
||||
stage_messages = {
|
||||
"select_policy": "Step 1: Select Policy to Analyze",
|
||||
"select_quiet_days": "Step 2: Select Quiet Time Period",
|
||||
"select_history_days": "Step 2: Configure Analysis Parameters",
|
||||
"analyzing": "Analyzing agent activity...",
|
||||
"view_results": "Step 3: Review Categorized Agents",
|
||||
"select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents",
|
||||
@@ -161,7 +162,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
# Initial policy selection for analysis
|
||||
self.selected_policy = message.policy
|
||||
logger.info(f"Selected policy for analysis: {self.selected_policy.name}")
|
||||
self._show_quiet_days_selection()
|
||||
self._show_history_days_selection()
|
||||
elif self.workflow_stage == "select_enforce_target":
|
||||
# Target policy selection for enforce ready agents
|
||||
self.enforce_ready_target_policy = message.policy
|
||||
@@ -177,48 +178,167 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
)
|
||||
self._show_migration_confirmation()
|
||||
|
||||
def _show_quiet_days_selection(self) -> None:
|
||||
"""Show the quiet days selection screen."""
|
||||
self.workflow_stage = "select_quiet_days"
|
||||
def _show_history_days_selection(self) -> None:
|
||||
"""Show the history days and quiet days selection screen."""
|
||||
self.workflow_stage = "select_history_days"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Create info text
|
||||
info_widget = Static(
|
||||
f"Policy Selected: {self.selected_policy.name}\n\n"
|
||||
f"History Period: {self.history_days} days\n\n"
|
||||
"Select quiet time period (days without untrusted execution):",
|
||||
id="quiet_days_info",
|
||||
"Configure Analysis Parameters:",
|
||||
id="analysis_params_info",
|
||||
)
|
||||
info_widget.styles.margin = (0, 0, 2, 0)
|
||||
content.mount(info_widget)
|
||||
|
||||
# Create button container and mount it first
|
||||
button_container = Vertical(id="quiet_days_buttons")
|
||||
button_container.styles.height = "auto"
|
||||
content.mount(button_container)
|
||||
# Create input container
|
||||
input_container = Vertical(id="analysis_params_input_container")
|
||||
input_container.styles.height = "auto"
|
||||
content.mount(input_container)
|
||||
|
||||
# Now add buttons to the mounted container
|
||||
for days in [15, 30, 45, 60]:
|
||||
btn = Button(
|
||||
f"{days} days {'(Default)' if days == 45 else ''}",
|
||||
id=f"quiet_days_{days}",
|
||||
classes="quiet_day_btn",
|
||||
# History days label
|
||||
history_label = Static("History Period (days of execution history to pull):")
|
||||
history_label.styles.margin = (0, 0, 1, 0)
|
||||
input_container.mount(history_label)
|
||||
|
||||
# Add history days input field
|
||||
history_input = Input(
|
||||
placeholder="Enter days (1-365, default: 150)",
|
||||
value="150",
|
||||
id="history_days_input",
|
||||
)
|
||||
history_input.styles.width = "50"
|
||||
history_input.styles.margin = (0, 0, 2, 0)
|
||||
input_container.mount(history_input)
|
||||
|
||||
# Quiet days label
|
||||
quiet_label = Static(
|
||||
"Quiet Time Period (days without execution to be considered quiet):"
|
||||
)
|
||||
quiet_label.styles.margin = (0, 0, 1, 0)
|
||||
input_container.mount(quiet_label)
|
||||
|
||||
# Add quiet days input field
|
||||
quiet_input = Input(
|
||||
placeholder="Enter days (1-365, default: 45)",
|
||||
value="45",
|
||||
id="quiet_days_input",
|
||||
)
|
||||
quiet_input.styles.width = "50"
|
||||
quiet_input.styles.margin = (0, 0, 2, 0)
|
||||
input_container.mount(quiet_input)
|
||||
|
||||
# Add submit button
|
||||
submit_btn = Button(
|
||||
"Continue",
|
||||
id="analysis_params_submit",
|
||||
variant="primary",
|
||||
)
|
||||
submit_btn.styles.width = "50"
|
||||
submit_btn.styles.margin = (1, 0, 0, 0)
|
||||
input_container.mount(submit_btn)
|
||||
|
||||
# Focus the first input field
|
||||
history_input.focus()
|
||||
|
||||
def _validate_and_submit_history_days(self) -> None:
|
||||
"""Validate and submit the history days and quiet days inputs."""
|
||||
try:
|
||||
history_input = self.query_one("#history_days_input", Input)
|
||||
quiet_input = self.query_one("#quiet_days_input", Input)
|
||||
|
||||
history_value = history_input.value.strip()
|
||||
quiet_value = quiet_input.value.strip()
|
||||
|
||||
# Validate history days
|
||||
if not history_value:
|
||||
self.app.notify(
|
||||
"Please enter a history period value", severity="error", timeout=3
|
||||
)
|
||||
history_input.focus()
|
||||
return
|
||||
|
||||
try:
|
||||
history_days = int(history_value)
|
||||
except ValueError:
|
||||
self.app.notify(
|
||||
"Please enter a valid number for history period",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
history_input.focus()
|
||||
return
|
||||
|
||||
if history_days < 1 or history_days > 365:
|
||||
self.app.notify(
|
||||
"History period must be between 1 and 365 days",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
history_input.focus()
|
||||
return
|
||||
|
||||
# Validate quiet days
|
||||
if not quiet_value:
|
||||
self.app.notify(
|
||||
"Please enter a quiet time period value",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
quiet_input.focus()
|
||||
return
|
||||
|
||||
try:
|
||||
quiet_days = int(quiet_value)
|
||||
except ValueError:
|
||||
self.app.notify(
|
||||
"Please enter a valid number for quiet time period",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
quiet_input.focus()
|
||||
return
|
||||
|
||||
if quiet_days < 1 or quiet_days > 365:
|
||||
self.app.notify(
|
||||
"Quiet time period must be between 1 and 365 days",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
quiet_input.focus()
|
||||
return
|
||||
|
||||
# Check that quiet days doesn't exceed history days
|
||||
if quiet_days > history_days:
|
||||
self.app.notify(
|
||||
"Quiet time period cannot exceed history period",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
quiet_input.focus()
|
||||
return
|
||||
|
||||
# All validation passed
|
||||
self.history_days = history_days
|
||||
self.quiet_days = quiet_days
|
||||
logger.info(
|
||||
f"Selected history days: {history_days}, quiet days: {quiet_days}"
|
||||
)
|
||||
btn.styles.width = "100%"
|
||||
btn.styles.margin = (0, 0, 1, 0)
|
||||
button_container.mount(btn)
|
||||
self._start_analysis()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error validating analysis parameters: {e}")
|
||||
self.app.notify(f"Error: {str(e)}", severity="error", timeout=3)
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button press events."""
|
||||
button_id = event.button.id
|
||||
|
||||
# Quiet days selection buttons
|
||||
if button_id and button_id.startswith("quiet_days_"):
|
||||
days = int(button_id.split("_")[-1])
|
||||
self.quiet_days = days
|
||||
logger.info(f"Selected quiet days: {days}")
|
||||
self._start_analysis()
|
||||
# Analysis parameters submit button
|
||||
if button_id == "analysis_params_submit":
|
||||
self._validate_and_submit_history_days()
|
||||
return
|
||||
|
||||
# Navigation buttons
|
||||
@@ -258,46 +378,44 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
self._show_policy_selection()
|
||||
return
|
||||
|
||||
def on_input_submitted(self, event: Input.Submitted) -> None:
|
||||
"""Handle input submission (Enter key pressed)."""
|
||||
if event.input.id in ["history_days_input", "quiet_days_input"]:
|
||||
self._validate_and_submit_history_days()
|
||||
|
||||
def _start_analysis(self) -> None:
|
||||
"""Start the agent activity analysis."""
|
||||
self.workflow_stage = "analyzing"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Show analyzing message with detailed steps
|
||||
analyzing_msg = Static(
|
||||
f"Analyzing Agent Activity\n"
|
||||
f"{'=' * 50}\n\n"
|
||||
f"Policy: {self.selected_policy.name}\n"
|
||||
f"History Period: {self.history_days} days\n"
|
||||
f"Quiet Threshold: {self.quiet_days} days\n\n"
|
||||
f"Progress:\n"
|
||||
f"Step 1/4: Fetching agents from policy...\n"
|
||||
f"Step 2/4: Pulling execution history (this may take a moment)...\n"
|
||||
f"Step 3/4: Analyzing activity patterns...\n"
|
||||
f"Step 4/4: Categorizing agents...\n\n"
|
||||
f"Please wait - this operation cannot be cancelled.",
|
||||
id="analyzing_message",
|
||||
)
|
||||
analyzing_msg.styles.margin = (2, 1)
|
||||
content.mount(analyzing_msg)
|
||||
|
||||
# Show notification
|
||||
# Show notification that analysis is starting
|
||||
self.app.notify(
|
||||
"Starting analysis - this may take several minutes for large policies",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
# Perform the analysis asynchronously
|
||||
self.call_later(self._perform_analysis)
|
||||
# Clear the screen to provide a blank canvas for Rust progress output
|
||||
# (Rust output displays over the TUI, so we clear everything except header/footer)
|
||||
try:
|
||||
# Clear title
|
||||
title_widget = self.query_one("#workflow_title", Static)
|
||||
title_widget.update("")
|
||||
|
||||
def _perform_analysis(self) -> None:
|
||||
# Clear status
|
||||
status_widget = self.query_one("#workflow_status", Static)
|
||||
status_widget.update("")
|
||||
|
||||
# Clear content area
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not clear screen for analysis: {e}")
|
||||
|
||||
# Delay the analysis start to ensure UI refresh completes first
|
||||
# This prevents Rust output from starting before the screen is cleared
|
||||
self.set_timer(0.5, self._perform_analysis_worker)
|
||||
|
||||
def _perform_analysis_worker(self) -> None:
|
||||
"""Perform the actual agent activity analysis."""
|
||||
try:
|
||||
# Update status: Fetching agents
|
||||
self._update_analysis_status("Step 1/4: Fetching agents from policy...")
|
||||
|
||||
# Get agents in the selected policy
|
||||
agents = self.api.agents_find_by_group(self.selected_policy.groupid)
|
||||
|
||||
@@ -310,32 +428,11 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
self._show_policy_selection()
|
||||
return
|
||||
|
||||
agent_count = len(agents)
|
||||
self.app.notify(
|
||||
f"Found {agent_count} agents - fetching execution history...",
|
||||
severity="information",
|
||||
timeout=3,
|
||||
)
|
||||
|
||||
# Update status: Pulling execution history
|
||||
self._update_analysis_status(
|
||||
f"Step 2/4: Pulling execution history for {agent_count} agents...\n"
|
||||
f"(This may take several minutes - progress shown in terminal)"
|
||||
)
|
||||
|
||||
# Get execution history (this shows progress bars in terminal via airlock_libs)
|
||||
policy_exec_history = getPolicyInfo(
|
||||
self.api, self.selected_policy, [1, 2, 6, 7], self.history_days
|
||||
)
|
||||
|
||||
# Update status: Analyzing patterns
|
||||
self._update_analysis_status("Step 3/4: Analyzing activity patterns...")
|
||||
self.app.notify(
|
||||
"History retrieved - analyzing patterns...",
|
||||
severity="information",
|
||||
timeout=2,
|
||||
)
|
||||
|
||||
if policy_exec_history.empty:
|
||||
logger.info(
|
||||
"No execution history found for the selected policy and time range."
|
||||
@@ -381,9 +478,6 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
lambda x: True if pd.isna(x) or x > self.quiet_days else False
|
||||
)
|
||||
|
||||
# Update status: Categorizing
|
||||
self._update_analysis_status("Step 4/4: Categorizing agents...")
|
||||
|
||||
# Sort agents
|
||||
agents = agents.sort_values(
|
||||
by=["execution_count", "hostname"], ascending=[True, True]
|
||||
@@ -394,7 +488,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
|
||||
# Categorize agents into DataFrames
|
||||
self.enforce_ready_df = agents[agents["enforce_ready"]].copy()
|
||||
self.non_enforce_ready_df = agents[not agents["enforce_ready"]].copy()
|
||||
self.non_enforce_ready_df = agents[~agents["enforce_ready"]].copy()
|
||||
|
||||
logger.info(
|
||||
f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, "
|
||||
@@ -416,27 +510,6 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
self.app.notify(f"Analysis failed: {str(e)}", severity="error", timeout=5)
|
||||
self._show_policy_selection()
|
||||
|
||||
def _update_analysis_status(self, status_text: str) -> None:
|
||||
"""Update the analysis status message."""
|
||||
try:
|
||||
analyzing_msg = self.query_one("#analyzing_message", Static)
|
||||
|
||||
# Build updated message
|
||||
updated_text = (
|
||||
f"Analyzing Agent Activity\n"
|
||||
f"{'=' * 50}\n\n"
|
||||
f"Policy: {self.selected_policy.name}\n"
|
||||
f"History Period: {self.history_days} days\n"
|
||||
f"Quiet Threshold: {self.quiet_days} days\n\n"
|
||||
f"Progress:\n"
|
||||
f"{status_text}\n\n"
|
||||
f"Please wait - this operation cannot be cancelled."
|
||||
)
|
||||
|
||||
analyzing_msg.update(updated_text)
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not update analysis status: {e}")
|
||||
|
||||
def _show_results(self) -> None:
|
||||
"""Show the categorized results."""
|
||||
self.workflow_stage = "view_results"
|
||||
@@ -821,7 +894,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
# Depending on stage, go back to previous stage or exit
|
||||
if self.workflow_stage in ["select_policy", "view_results", "complete"]:
|
||||
self.app.pop_screen()
|
||||
elif self.workflow_stage == "select_quiet_days":
|
||||
elif self.workflow_stage == "select_history_days":
|
||||
self._show_policy_selection()
|
||||
elif self.workflow_stage == "select_enforce_target":
|
||||
self._show_results()
|
||||
|
||||
-444
@@ -1,444 +0,0 @@
|
||||
# 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
|
||||
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,
|
||||
)
|
||||
|
||||
from models.agent import Agent
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
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
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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": [
|
||||
(
|
||||
"🖥️ - 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"),
|
||||
],
|
||||
"policy": [
|
||||
("⚖️ - Prepare Policy For Enforcement", "policy_prep_button"),
|
||||
("🛑 - Revoke OTPs", "otp_revoke_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, [])
|
||||
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("Tree View", id="p_tree"),
|
||||
Tab("Agents", id="agent_actions"),
|
||||
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")
|
||||
|
||||
# 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 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_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
|
||||
# ---------------------------------------------------------------------------
|
||||
def run_Loxide(api: AirlockAPIWrapper) -> 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)
|
||||
|
||||
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) DEV
|
||||
# ---------------------------------------------------------------------------
|
||||
if __name__ == "__main__":
|
||||
api = AirlockAPIWrapper()
|
||||
run_Loxide(api)
|
||||
@@ -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
|
||||
@@ -316,6 +323,13 @@ class AgentMoveOperations(Widget):
|
||||
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)
|
||||
@@ -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:
|
||||
"""
|
||||
@@ -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.
|
||||
|
||||
@@ -14,6 +14,7 @@
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import difflib
|
||||
from pathlib import Path
|
||||
import re
|
||||
from typing import List, Optional
|
||||
|
||||
@@ -57,7 +58,7 @@ class MultiAgentSelector(Widget):
|
||||
|
||||
def compose(self):
|
||||
yield Header(show_clock=True, icon="⚙")
|
||||
title_text = Static("🖥️ Agent Selector", id="selector_title")
|
||||
title_text = Static("🖥️ Agent Selector", id="selector_title")
|
||||
title_text.styles.margin = (0, 0, 0, 1)
|
||||
yield title_text
|
||||
|
||||
@@ -77,7 +78,7 @@ class MultiAgentSelector(Widget):
|
||||
text_area.styles.overflow_y = "auto"
|
||||
yield text_area
|
||||
|
||||
with Horizontal(id="switch_search_container"):
|
||||
with Horizontal(id="switch_container"):
|
||||
switch = Switch(value=False, id="match_switch")
|
||||
switch.styles.width = "auto"
|
||||
switch.styles.margin = (1, 0, 0, 0)
|
||||
@@ -89,8 +90,13 @@ class MultiAgentSelector(Widget):
|
||||
switch_label.styles.margin = (2, 1, 0, 0)
|
||||
yield switch_label
|
||||
|
||||
with Horizontal(id="action_buttons_container"):
|
||||
load_file = Button("📂 Load File", id="load_file_button")
|
||||
load_file.styles.margin = (1, 1, 0, 1)
|
||||
yield load_file
|
||||
|
||||
search = Button("🔍 Search", id="search_button")
|
||||
search.styles.margin = (1, 0, 0, 0)
|
||||
search.styles.margin = (1, 0, 0, 1)
|
||||
yield search
|
||||
|
||||
with Horizontal() as select_buttons:
|
||||
@@ -152,6 +158,9 @@ class MultiAgentSelector(Widget):
|
||||
]
|
||||
self.post_message(self.AgentsSelected(selected_agents))
|
||||
event.stop()
|
||||
elif btn_id == "load_file_button":
|
||||
self._load_from_file()
|
||||
event.stop()
|
||||
elif btn_id == "search_button":
|
||||
self.update_matches()
|
||||
event.stop()
|
||||
@@ -166,7 +175,7 @@ class MultiAgentSelector(Widget):
|
||||
match_list.add_option((name, name))
|
||||
unmatched_label = self.query_one("#unmatched_label", Static)
|
||||
if unmatched:
|
||||
unmatched_label.update(f"⚠️ No matches for: {', '.join(unmatched)}")
|
||||
unmatched_label.update(f"âš ï¸ No matches for: {', '.join(unmatched)}")
|
||||
else:
|
||||
unmatched_label.update("")
|
||||
|
||||
@@ -214,3 +223,118 @@ class MultiAgentSelector(Widget):
|
||||
else:
|
||||
unmatched.append(name)
|
||||
return sorted(matched), unmatched
|
||||
|
||||
def _load_from_file(self):
|
||||
"""Safely load device names from a text file."""
|
||||
try:
|
||||
# Import here to avoid issues if tkinter isn't available
|
||||
import tkinter as tk
|
||||
from tkinter import filedialog
|
||||
|
||||
# Create file dialog
|
||||
root = tk.Tk()
|
||||
root.withdraw()
|
||||
|
||||
file_path = filedialog.askopenfilename(
|
||||
title="Select device list file",
|
||||
filetypes=[
|
||||
("Text files", "*.txt"),
|
||||
("CSV files", "*.csv"),
|
||||
("All files", "*.*"),
|
||||
],
|
||||
)
|
||||
|
||||
if not file_path:
|
||||
# User cancelled
|
||||
return
|
||||
|
||||
# Validate file path
|
||||
path_obj = Path(file_path)
|
||||
if not path_obj.exists():
|
||||
self.app.notify("File does not exist", severity="error", timeout=3)
|
||||
return
|
||||
|
||||
if not path_obj.is_file():
|
||||
self.app.notify(
|
||||
"Selected path is not a file", severity="error", timeout=3
|
||||
)
|
||||
return
|
||||
|
||||
# Check file size (limit to 1 MB for safety)
|
||||
file_size = path_obj.stat().st_size
|
||||
if file_size > 1_000_000: # 1 MB
|
||||
self.app.notify(
|
||||
f"File too large ({file_size:,} bytes). Maximum 1 MB.",
|
||||
severity="error",
|
||||
timeout=5,
|
||||
)
|
||||
return
|
||||
|
||||
# Read file with proper encoding to preserve emojis
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
except UnicodeDecodeError:
|
||||
# Try with different encoding if UTF-8 fails
|
||||
try:
|
||||
with open(file_path, "r", encoding="latin-1") as f:
|
||||
content = f.read()
|
||||
self.app.notify(
|
||||
"File loaded with Latin-1 encoding (UTF-8 failed)",
|
||||
severity="warning",
|
||||
timeout=3,
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(
|
||||
f"Error reading file: {str(e)}", severity="error", timeout=5
|
||||
)
|
||||
return
|
||||
|
||||
# Validate and sanitize content
|
||||
lines = content.split("\n")
|
||||
valid_lines = []
|
||||
invalid_count = 0
|
||||
|
||||
# Pattern for valid hostnames/device names
|
||||
# Allows: letters, numbers, hyphens, underscores, periods, and Unicode chars
|
||||
hostname_pattern = re.compile(r"^[\w\-\.\u0080-\uFFFF]+$", re.UNICODE)
|
||||
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue # Skip empty lines
|
||||
|
||||
# Check if line looks like a valid hostname/device name
|
||||
if hostname_pattern.match(line):
|
||||
valid_lines.append(line)
|
||||
else:
|
||||
invalid_count += 1
|
||||
# Log but don't add invalid entries
|
||||
|
||||
if not valid_lines:
|
||||
self.app.notify(
|
||||
"No valid device names found in file", severity="warning", timeout=3
|
||||
)
|
||||
return
|
||||
|
||||
# Update text area with validated content
|
||||
text_area = self.query_one("#device_input", TextArea)
|
||||
text_area.text = "\n".join(valid_lines)
|
||||
|
||||
# Show notification
|
||||
msg = f"✅ Loaded {len(valid_lines)} devices from file"
|
||||
if invalid_count > 0:
|
||||
msg += f" ({invalid_count} invalid entries skipped)"
|
||||
|
||||
self.app.notify(msg, severity="information", timeout=5)
|
||||
|
||||
except ImportError:
|
||||
self.app.notify(
|
||||
"tkinter not available - cannot open file dialog",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(
|
||||
f"Error loading file: {str(e)}", severity="error", timeout=5
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
|
||||
+24
-24
@@ -88,9 +88,9 @@ def sortHashes(
|
||||
):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
prompt="Enter how many days of history to pull (1-365): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
valid_range=(1, 365),
|
||||
)
|
||||
|
||||
logger.debug(f"{history_days} day selected for history")
|
||||
@@ -655,7 +655,7 @@ def section_header(title):
|
||||
|
||||
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
section_header("Prepare to Enforce Policy ")
|
||||
section_header("Prepare to Enforce Policy")
|
||||
print(
|
||||
colorText(
|
||||
"\nSequentially follow these steps to prepare a policy for enforcement:",
|
||||
@@ -670,11 +670,11 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
)
|
||||
)
|
||||
if not selected_policies:
|
||||
print(colorText(" [✗] No policies have been chosen", "red"))
|
||||
print(colorText(" [âŒ] No policies have been chosen", "red"))
|
||||
else:
|
||||
print(colorText("The following policies have been chosen:", "green"))
|
||||
for policy in selected_policies:
|
||||
print(colorText(f" [✓] {policy.name}", "green"))
|
||||
print(colorText(f" [✅] {policy.name}", "green"))
|
||||
|
||||
# Step 2: Destination Policy and Allowlist
|
||||
print(
|
||||
@@ -683,22 +683,22 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
if destination_policy:
|
||||
print(
|
||||
colorText(
|
||||
f" [✓] {destination_policy[0].name} has been selected as the destination policy",
|
||||
f" [✅] {destination_policy[0].name} has been selected as the destination policy",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(colorText(" [✗] No destination policy has been chosen", "red"))
|
||||
print(colorText(" [âŒ] No destination policy has been chosen", "red"))
|
||||
|
||||
if destination_allowlist:
|
||||
print(
|
||||
colorText(
|
||||
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
|
||||
f" [✅] {destination_allowlist[0].name} has been selected as allowlist",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(colorText(" [✗] No allowlist has been chosen", "red"))
|
||||
print(colorText(" [âŒ] No allowlist has been chosen", "red"))
|
||||
|
||||
# Step 3: Data Preparation
|
||||
print(
|
||||
@@ -713,9 +713,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Data has been fetched"
|
||||
" [✅] Data has been fetched"
|
||||
if os.path.exists(review_path)
|
||||
else " [✗] Data has not been fetched"
|
||||
else " [âŒ] Data has not been fetched"
|
||||
),
|
||||
"green" if os.path.exists(review_path) else "red",
|
||||
)
|
||||
@@ -723,7 +723,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check data fetch status", "red"
|
||||
" [âŒ] No policies selected, cannot check data fetch status", "red"
|
||||
)
|
||||
)
|
||||
|
||||
@@ -756,9 +756,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Reviewed hashes have been loaded"
|
||||
" [✅] Reviewed hashes have been loaded"
|
||||
if os.path.exists(approved_path)
|
||||
else " [✗] Reviewed hashes have not been loaded"
|
||||
else " [âŒ] Reviewed hashes have not been loaded"
|
||||
),
|
||||
"green" if os.path.exists(approved_path) else "red",
|
||||
)
|
||||
@@ -766,9 +766,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Path review list created"
|
||||
" [✅] Path review list created"
|
||||
if os.path.exists(second_review_path)
|
||||
else " [✗] Path review list has not been created"
|
||||
else " [âŒ] Path review list has not been created"
|
||||
),
|
||||
"green" if os.path.exists(second_review_path) else "red",
|
||||
)
|
||||
@@ -776,7 +776,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check reviewed hashes or path list",
|
||||
" [âŒ] No policies selected, cannot check reviewed hashes or path list",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
@@ -812,9 +812,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Reviewed path list detected"
|
||||
" [✅] Reviewed path list detected"
|
||||
if os.path.exists(reviewed_path)
|
||||
else " [✗] Path review list has not been detected"
|
||||
else " [âŒ] Path review list has not been detected"
|
||||
),
|
||||
"green" if os.path.exists(reviewed_path) else "red",
|
||||
)
|
||||
@@ -825,9 +825,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Preflight Path Exclusion List has been generated"
|
||||
" [✅] Preflight Path Exclusion List has been generated"
|
||||
if preflight_ready
|
||||
else " [✗] Preflight Path Exclusion List has not been generated"
|
||||
else " [âŒ] Preflight Path Exclusion List has not been generated"
|
||||
),
|
||||
"green" if preflight_ready else "red",
|
||||
)
|
||||
@@ -835,7 +835,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check preflight status", "red"
|
||||
" [âŒ] No policies selected, cannot check preflight status", "red"
|
||||
)
|
||||
)
|
||||
|
||||
@@ -866,5 +866,5 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(colorText(" Apply approved hashes to allowlist", "cyan"))
|
||||
|
||||
# Utility Options
|
||||
print(colorText("F. Open Working Directory", "cyan"))
|
||||
print(colorText("B. Back", "cyan"))
|
||||
print(colorText("F. Open Working Directory", "cyan"))
|
||||
print(colorText("B. Back", "cyan"))
|
||||
|
||||
@@ -0,0 +1,261 @@
|
||||
# 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, Input, 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;
|
||||
align: left middle;
|
||||
}
|
||||
|
||||
ServerLogWidget .filter_label {
|
||||
width: auto;
|
||||
height: 3;
|
||||
content-align: left middle;
|
||||
padding-right: 1;
|
||||
}
|
||||
|
||||
ServerLogWidget #filter_input {
|
||||
width: 40;
|
||||
margin-right: 1;
|
||||
}
|
||||
|
||||
ServerLogWidget Button {
|
||||
margin-right: 1;
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper):
|
||||
super().__init__()
|
||||
self.api = api
|
||||
self.all_logs = [] # Store all logs for filtering
|
||||
self.columns = [] # Store column names
|
||||
|
||||
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 Static("Filter:", classes="filter_label")
|
||||
yield Input(placeholder="Filter (use * and ? wildcards)", id="filter_input")
|
||||
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)
|
||||
self.all_logs = []
|
||||
self.columns = []
|
||||
return
|
||||
|
||||
# Store all logs for filtering
|
||||
self.all_logs = logs
|
||||
|
||||
# Clear existing data
|
||||
table.clear(columns=True)
|
||||
|
||||
# Add columns based on the first log entry
|
||||
if logs:
|
||||
first_log = logs[0]
|
||||
self.columns = [col for col in first_log.keys() if col != "checkpoint"]
|
||||
|
||||
for col in self.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 self.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.")
|
||||
self.all_logs = []
|
||||
self.columns = []
|
||||
|
||||
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)
|
||||
self.all_logs = []
|
||||
self.columns = []
|
||||
|
||||
def filter_logs(self, filter_text: str) -> None:
|
||||
"""Filter the logs based on the filter text with wildcard support."""
|
||||
import fnmatch
|
||||
|
||||
table = self.query_one("#server_log_table", DataTable)
|
||||
status = self.query_one("#status_bar", Static)
|
||||
|
||||
if not self.all_logs:
|
||||
return
|
||||
|
||||
# Clear existing data
|
||||
table.clear(columns=True)
|
||||
|
||||
# Re-add columns
|
||||
for col in self.columns:
|
||||
table.add_column(col, key=col)
|
||||
|
||||
# Filter logs
|
||||
filtered_logs = []
|
||||
if filter_text.strip():
|
||||
filter_pattern = filter_text.strip().lower()
|
||||
for log_entry in self.all_logs:
|
||||
# Check if any field matches the filter pattern
|
||||
match = False
|
||||
for col in self.columns:
|
||||
value = str(log_entry.get(col, "")).lower()
|
||||
if fnmatch.fnmatch(value, filter_pattern):
|
||||
match = True
|
||||
break
|
||||
if match:
|
||||
filtered_logs.append(log_entry)
|
||||
else:
|
||||
# No filter, show all logs
|
||||
filtered_logs = self.all_logs
|
||||
|
||||
# Add filtered rows in reverse order
|
||||
for log_entry in reversed(filtered_logs):
|
||||
row_data = []
|
||||
for col in self.columns:
|
||||
value = log_entry.get(col, "")
|
||||
# Format datetime column to be more readable
|
||||
if col == "datetime" and value:
|
||||
try:
|
||||
dt = datetime.datetime.fromisoformat(
|
||||
str(value).replace("Z", "+00:00")
|
||||
)
|
||||
value = dt.strftime("%Y-%m-%d %H:%M:%S")
|
||||
except Exception:
|
||||
pass
|
||||
row_data.append(str(value))
|
||||
table.add_row(*row_data)
|
||||
|
||||
if filter_text.strip():
|
||||
status.update(
|
||||
f"✅ Showing {len(filtered_logs)} of {len(self.all_logs)} log entries (filtered)"
|
||||
)
|
||||
else:
|
||||
status.update(
|
||||
f"✅ Loaded {len(self.all_logs)} log entries from the last 72 hours"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Filtered to {len(filtered_logs)} entries with pattern: {filter_text}"
|
||||
)
|
||||
|
||||
def on_input_changed(self, event: Input.Changed) -> None:
|
||||
"""Handle filter input changes."""
|
||||
if event.input.id == "filter_input":
|
||||
self.filter_logs(event.value)
|
||||
|
||||
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()
|
||||
# Clear the filter input when refreshing
|
||||
try:
|
||||
filter_input = self.query_one("#filter_input", Input)
|
||||
filter_input.value = ""
|
||||
except Exception:
|
||||
pass
|
||||
event.stop()
|
||||
Generated
+353
-587
File diff suppressed because it is too large
Load Diff
+12
-12
@@ -1,31 +1,31 @@
|
||||
[package]
|
||||
name = "airlock_libs"
|
||||
version = "5.1.1"
|
||||
version = "7.3.0"
|
||||
edition = "2024"
|
||||
|
||||
[lib]
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
chrono = "0.4.42"
|
||||
indicatif = "0.18.2"
|
||||
mongodb = "3.3.0"
|
||||
opentelemetry = { version = "0.18.0", features = ["rt-tokio", "metrics", "trace"] }
|
||||
opentelemetry-otlp = { version = "0.11.0", features = ["trace", "metrics"] }
|
||||
opentelemetry-semantic-conventions = { version = "0.10.0" }
|
||||
opentelemetry-proto = { version = "0.1.0"}
|
||||
opentelemetry = { version = "0.27.0", features = ["logs", "metrics", "trace"] }
|
||||
opentelemetry-otlp = { version = "0.27.0", features = ["trace", "metrics", "grpc-tonic", "http-proto", "tls", "reqwest-client", "reqwest-rustls"] }
|
||||
opentelemetry-semantic-conventions = { version = "0.27.0" }
|
||||
opentelemetry-proto = { version = "0.27.0"}
|
||||
pyo3 = { version = "0.27.0", features = ["extension-module", "generate-import-lib"] }
|
||||
reqwest = { version = "0.12.24", features = ["json", "native-tls"] }
|
||||
reqwest = { version = "0.12.24", features = ["json", "native-tls", "rustls-tls"] }
|
||||
serde = "1.0.228"
|
||||
serde-pyobject = "0.8.0"
|
||||
serde_json = "1.0.145"
|
||||
tokio = { version = "1.48.0", features = ["full"] }
|
||||
tonic = { version = "0.8.2", features = ["tls-roots"] }
|
||||
tonic = { version = "0.12.3", features = ["tls-roots"] }
|
||||
tracing = "0.1.41"
|
||||
tracing-subscriber = "0.3.20"
|
||||
tracing-opentelemetry = "0.32.0"
|
||||
pyo3-async-runtimes = { version = "0.27.0", features = ["async-std", "tokio"] }
|
||||
crossbeam = "0.8.4"
|
||||
log = "0.4.29"
|
||||
flexi_logger = "0.31.7"
|
||||
opentelemetry-appender-log = "0.27.0"
|
||||
opentelemetry_sdk = { version = "0.27.0", features = ["rt-tokio", "testing", "trace"] }
|
||||
|
||||
[package.metadata.maturin]
|
||||
generate-abi-stubs = true
|
||||
@@ -38,4 +38,4 @@ codegen-units = 1
|
||||
panic = 'abort'
|
||||
strip = true
|
||||
debug-assertions = false
|
||||
overflow-checks = false
|
||||
overflow-checks = true
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "airlock_libs"
|
||||
version = "5.1.1"
|
||||
version = "7.3.0"
|
||||
description = "Airlock Digital API Wrapper"
|
||||
readme = "README.md"
|
||||
license = { text = "AGPL-3.0-only" }
|
||||
|
||||
@@ -8,7 +8,32 @@ pub struct TelemetryConfig {
|
||||
}
|
||||
|
||||
impl TelemetryConfig {
|
||||
pub fn load() -> Self {
|
||||
pub fn init_tracer() -> opentelemetry_sdk::trace::TracerProvider {
|
||||
let cfg: TelemetryConfig = TelemetryConfig::load();
|
||||
if !cfg.TELEMETRY {
|
||||
return TracerProvider::builder().build();
|
||||
}
|
||||
let endpoint = cfg.TELEM_URL.unwrap_or_default();
|
||||
let channel = Channel::from_shared(endpoint.clone())
|
||||
.unwrap()
|
||||
.tls_config(ClientTlsConfig::new().with_native_roots())
|
||||
.unwrap()
|
||||
.connect_lazy();
|
||||
let exporter = opentelemetry_otlp::SpanExporter::builder()
|
||||
.with_tonic()
|
||||
.with_endpoint(endpoint.clone())
|
||||
.with_channel(channel)
|
||||
.build()
|
||||
.expect("Failed to build exporter");
|
||||
opentelemetry_sdk::trace::TracerProvider::builder()
|
||||
.with_simple_exporter(exporter)
|
||||
.with_resource(Resource::new(vec![KeyValue::new(
|
||||
"service.name",
|
||||
"LoxideLibs",
|
||||
)]))
|
||||
.build()
|
||||
}
|
||||
fn load() -> Self {
|
||||
let cfg_path = get_base_directory().join("config\\user_config.json");
|
||||
if !cfg_path.exists() {
|
||||
return Self {
|
||||
@@ -64,36 +89,30 @@ pub struct Group {
|
||||
pub(crate) localip: String,
|
||||
}
|
||||
|
||||
pub enum ExtractedValues {
|
||||
Headers(reqwest::header::HeaderMap),
|
||||
BaseUrl(String),
|
||||
pub struct PyData {
|
||||
pub headers: reqwest::header::HeaderMap,
|
||||
pub base_url: String,
|
||||
}
|
||||
|
||||
pub trait Converter {
|
||||
fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues;
|
||||
}
|
||||
|
||||
pub struct PyData;
|
||||
|
||||
impl Converter for PyData {
|
||||
fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues {
|
||||
if extract_headers {
|
||||
let headers = py_self.getattr(py, "headers").unwrap().to_string();
|
||||
let headers_replace = headers.replace('\'', "\"");
|
||||
let parsed: Value = serde_json::from_str(headers_replace.as_str()).unwrap();
|
||||
let mut header_map = HeaderMap::new();
|
||||
if let Some(obj) = parsed.as_object() {
|
||||
for (_key, value) in obj {
|
||||
if let Some(v) = value.as_str() {
|
||||
let val = HeaderValue::from_str(v).unwrap();
|
||||
header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val);
|
||||
}
|
||||
impl PyData {
|
||||
pub fn extract_data(py: Python<'_>, obj: &Py<PyAny>) -> Self {
|
||||
let headers_raw = obj.getattr(py, "headers").unwrap().to_string();
|
||||
let headers_json = headers_raw.replace('\'', "\"");
|
||||
let parsed: Value = serde_json::from_str(&headers_json).unwrap();
|
||||
let mut header_map = HeaderMap::new();
|
||||
if let Some(obj) = parsed.as_object() {
|
||||
for (key, val) in obj {
|
||||
if let Some(v) = val.as_str() {
|
||||
let header_name = HeaderName::from_str(key).unwrap();
|
||||
let header_value: HeaderValue = HeaderValue::from_str(v).unwrap();
|
||||
header_map.insert(header_name, header_value);
|
||||
}
|
||||
}
|
||||
ExtractedValues::Headers(header_map)
|
||||
} else {
|
||||
let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
|
||||
ExtractedValues::BaseUrl(base_url)
|
||||
}
|
||||
let base_url = obj.getattr(py, "base_url").unwrap().to_string();
|
||||
Self {
|
||||
headers: header_map,
|
||||
base_url,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,21 +1,25 @@
|
||||
pub use chrono::{Duration, Local, NaiveDate};
|
||||
pub use crossbeam::channel::unbounded;
|
||||
pub use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
|
||||
pub use mongodb::bson::oid::ObjectId;
|
||||
pub use opentelemetry::global::shutdown_tracer_provider;
|
||||
pub use opentelemetry::sdk::Resource;
|
||||
pub use opentelemetry::global::GlobalTracerProvider;
|
||||
pub use opentelemetry::trace::noop::NoopTracerProvider;
|
||||
pub use opentelemetry::trace::{Status, TraceContextExt, TraceError};
|
||||
pub use opentelemetry::{Context, KeyValue, sdk::trace as sdktrace, trace::Tracer};
|
||||
pub use opentelemetry::{Key, global};
|
||||
pub use opentelemetry::trace::{Status, TraceContextExt, Tracer};
|
||||
pub use opentelemetry::*;
|
||||
pub use opentelemetry_otlp::ExportConfig;
|
||||
pub use opentelemetry_otlp::WithExportConfig;
|
||||
pub use opentelemetry_otlp::WithTonicConfig;
|
||||
pub use opentelemetry_sdk::Resource;
|
||||
pub use opentelemetry_sdk::trace::{Config, TracerProvider};
|
||||
pub use pyo3::{prelude::*, types::PyString};
|
||||
pub use pyo3_async_runtimes::async_std;
|
||||
pub use reqwest::{
|
||||
Client,
|
||||
header::{HeaderMap, HeaderName, HeaderValue},
|
||||
};
|
||||
pub use serde::{Deserialize, Serialize};
|
||||
pub use serde_json::Value;
|
||||
pub use std::sync::{Arc, Mutex};
|
||||
pub use std::thread;
|
||||
pub use std::{
|
||||
collections::HashMap,
|
||||
env,
|
||||
@@ -25,3 +29,4 @@ pub use std::{
|
||||
path::PathBuf,
|
||||
str::FromStr,
|
||||
};
|
||||
pub use tonic::transport::{Channel, ClientTlsConfig};
|
||||
|
||||
+137
-89
@@ -1,38 +1,30 @@
|
||||
use std::thread;
|
||||
|
||||
use crossbeam::channel::unbounded;
|
||||
|
||||
use crate::modules::datatypes::*;
|
||||
use crate::prelude::*;
|
||||
|
||||
#[pyfunction]
|
||||
pub fn pull_policy_exec_histories(
|
||||
py: Python<'_>,
|
||||
py_self: Py<PyAny>,
|
||||
policy_names: String,
|
||||
policy_names: Option<String>,
|
||||
exec_types: String,
|
||||
days: i64,
|
||||
) -> Py<PyString> {
|
||||
let headers: HeaderMap = match PyData::convert(py, &py_self, true) {
|
||||
ExtractedValues::Headers(h) => h,
|
||||
ExtractedValues::BaseUrl(_) => std::process::abort(),
|
||||
};
|
||||
let base_url = match PyData::convert(py, &py_self, false) {
|
||||
ExtractedValues::Headers(_) => std::process::abort(),
|
||||
ExtractedValues::BaseUrl(b) => b,
|
||||
};
|
||||
let handle = std::thread::spawn(move || {
|
||||
let rt = match tokio::runtime::Runtime::new() {
|
||||
println!();
|
||||
let data: PyData = PyData::extract_data(py, &py_self);
|
||||
let headers: HeaderMap = data.headers;
|
||||
let base_url: String = data.base_url;
|
||||
let handle: thread::JoinHandle<String> = std::thread::spawn(move || {
|
||||
let rt: tokio::runtime::Runtime = match tokio::runtime::Runtime::new() {
|
||||
Ok(rt) => rt,
|
||||
Err(e) => {
|
||||
println!("Failed to build Tokio Runtime: {:?}", e);
|
||||
std::process::abort();
|
||||
}
|
||||
};
|
||||
rt.block_on(async {
|
||||
let _ = init_tracer();
|
||||
});
|
||||
let tracer = global::tracer("global_tracer");
|
||||
let _cx = Context::new();
|
||||
let tracer_provider = rt.block_on(async { TelemetryConfig::init_tracer() });
|
||||
global::set_tracer_provider(tracer_provider.clone());
|
||||
let tracer: global::BoxedTracer = global::tracer("tracer");
|
||||
let _cx: Context = Context::new();
|
||||
let file_path: PathBuf = format!(
|
||||
"{}\\cache\\chunkinator.json",
|
||||
get_base_directory().display()
|
||||
@@ -58,14 +50,14 @@ pub fn pull_policy_exec_histories(
|
||||
}
|
||||
}
|
||||
}
|
||||
let data = ApiResponse {
|
||||
let data: ApiResponse = ApiResponse {
|
||||
error: "Success".to_string(),
|
||||
response: ExecHistories {
|
||||
exechistories: vec![],
|
||||
},
|
||||
};
|
||||
let writeable_filepath = file_path.clone();
|
||||
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
|
||||
let writeable_filepath: PathBuf = file_path.clone();
|
||||
let data_write: String = serde_json::to_string_pretty(&data).expect("Failed to serialize");
|
||||
match fs::write(writeable_filepath.clone(), data_write) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
@@ -74,17 +66,18 @@ pub fn pull_policy_exec_histories(
|
||||
}
|
||||
}
|
||||
let mut checkpoint_number: String = SkipBack::find_checkpoint(days).to_string();
|
||||
let multi_progress = MultiProgress::new();
|
||||
multi_progress.set_draw_target(ProgressDrawTarget::stderr());
|
||||
let progress_bar = multi_progress.add(ProgressBar::new(100));
|
||||
progress_bar.set_style(
|
||||
let progress_bar = Arc::new(Mutex::new(ProgressBar::new(100)));
|
||||
progress_bar
|
||||
.lock()
|
||||
.unwrap()
|
||||
.set_draw_target(ProgressDrawTarget::stderr());
|
||||
progress_bar.lock().unwrap().set_style(
|
||||
ProgressStyle::default_bar()
|
||||
.template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len} {message}")
|
||||
.unwrap(),
|
||||
.template("Total - Policy Name: {msg}: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
|
||||
.unwrap().progress_chars("⣿⣦⣀")
|
||||
);
|
||||
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
|
||||
let client = tracer.in_span("Building HTTP Client", |cx| {
|
||||
let client_result = build_client(headers);
|
||||
let client: Client = tracer.in_span("Building HTTP Client", |cx| {
|
||||
let client_result: Result<Client, reqwest::Error> = build_client(headers);
|
||||
match client_result {
|
||||
Ok(client_result) => {
|
||||
cx.span().add_event(
|
||||
@@ -111,12 +104,14 @@ pub fn pull_policy_exec_histories(
|
||||
}
|
||||
}
|
||||
});
|
||||
let cutoff = Local::now().naive_local() - Duration::days(days);
|
||||
let (tx, rx) = unbounded::<Vec<Group>>();
|
||||
let cutoff: chrono::NaiveDateTime =
|
||||
Local::now().naive_local() - chrono::Duration::days(days);
|
||||
let (tx, rx) = unbounded::<(Context, Vec<Group>)>();
|
||||
let pb_clone = progress_bar.clone();
|
||||
thread::spawn(move || {
|
||||
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
|
||||
{
|
||||
let contents = fs::read_to_string(&writeable_filepath).unwrap_or_default();
|
||||
let contents: String = fs::read_to_string(&writeable_filepath).unwrap_or_default();
|
||||
let existing: ApiResponse =
|
||||
serde_json::from_str(&contents).unwrap_or(ApiResponse {
|
||||
error: "Success".to_string(),
|
||||
@@ -128,7 +123,7 @@ pub fn pull_policy_exec_histories(
|
||||
.response
|
||||
.exechistories
|
||||
.into_iter()
|
||||
.map(|entry| {
|
||||
.map(|entry: Group| {
|
||||
(
|
||||
(
|
||||
entry.sha256.clone(),
|
||||
@@ -142,12 +137,19 @@ pub fn pull_policy_exec_histories(
|
||||
} else {
|
||||
HashMap::new()
|
||||
};
|
||||
while let Ok(parsed_responses) = rx.recv() {
|
||||
while let Ok((cx, parsed_responses)) = rx.recv() {
|
||||
cx.span().add_event(
|
||||
"Received Data from Producer",
|
||||
vec![KeyValue::new(
|
||||
"Items to Process",
|
||||
parsed_responses.len().to_string(),
|
||||
)],
|
||||
);
|
||||
for executions in parsed_responses {
|
||||
if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let history_date = match NaiveDate::parse_from_str(
|
||||
let history_date: NaiveDate = match NaiveDate::parse_from_str(
|
||||
&executions.datetime.replace(" +0000 UTC", ""),
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
) {
|
||||
@@ -155,7 +157,7 @@ pub fn pull_policy_exec_histories(
|
||||
Err(_) => continue,
|
||||
};
|
||||
if history_date >= cutoff.into() {
|
||||
let key = (
|
||||
let key: (String, String, String) = (
|
||||
executions.sha256.clone(),
|
||||
executions.filename.clone(),
|
||||
executions.hostname.clone(),
|
||||
@@ -163,46 +165,100 @@ pub fn pull_policy_exec_histories(
|
||||
seen.entry(key).or_insert(executions.clone());
|
||||
}
|
||||
}
|
||||
let final_response = ApiResponse {
|
||||
let final_response: ApiResponse = ApiResponse {
|
||||
error: "Success".to_string(),
|
||||
response: ExecHistories {
|
||||
exechistories: seen.values().cloned().collect(),
|
||||
},
|
||||
};
|
||||
let data_write = serde_json::to_string_pretty(&final_response).unwrap();
|
||||
let data_write: String = serde_json::to_string_pretty(&final_response).unwrap();
|
||||
match fs::write(&writeable_filepath, data_write) {
|
||||
Ok(_) => {}
|
||||
Ok(_) => {
|
||||
cx.span().add_event(
|
||||
"Writing Data to File",
|
||||
vec![KeyValue::new("Success", "Ok".to_string())],
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
|
||||
cx.span().add_event(
|
||||
"Writing Data to File",
|
||||
vec![KeyValue::new("Failed", e.to_string())],
|
||||
);
|
||||
cx.span()
|
||||
.set_status(Status::error("Failed to Write to File"));
|
||||
}
|
||||
}
|
||||
cx.span().add_event(
|
||||
"Finished Deduplicating Data",
|
||||
vec![KeyValue::new(
|
||||
"Items Successfully Processed",
|
||||
seen.len().to_string(),
|
||||
)],
|
||||
);
|
||||
}
|
||||
});
|
||||
let mut first_date: Option<NaiveDate> = None;
|
||||
tracer.in_span("Airlock Data Retreival", |cx| {
|
||||
let span = cx.span();
|
||||
span.set_attribute(Key::new("Days").string(days.to_string()));
|
||||
span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
|
||||
pb_clone
|
||||
.lock()
|
||||
.unwrap()
|
||||
.enable_steady_tick(std::time::Duration::from_millis(100));
|
||||
pb_clone
|
||||
.lock()
|
||||
.unwrap()
|
||||
.set_message(policy_names.clone().unwrap_or("Statistics".to_string()));
|
||||
let span: opentelemetry::trace::SpanRef<'_> = cx.span();
|
||||
span.set_attribute(KeyValue::new("Days", days.to_string()));
|
||||
span.set_attribute(KeyValue::new(
|
||||
"Policy Name",
|
||||
policy_names
|
||||
.clone()
|
||||
.unwrap_or("Statistics Monitoring".to_string()),
|
||||
));
|
||||
loop {
|
||||
let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
|
||||
let results: ApiResponse = history_logging(
|
||||
cx.span().add_event(
|
||||
"Retrieving Responses from API",
|
||||
vec![KeyValue::new(
|
||||
"Checkpoint Number",
|
||||
checkpoint_number.to_string(),
|
||||
)],
|
||||
);
|
||||
let results: ApiResponse = rt.block_on(history_logging(
|
||||
&base_url,
|
||||
&exec_types,
|
||||
&checkpoint_number,
|
||||
&policy_names,
|
||||
&client,
|
||||
));
|
||||
cx.span().add_event(
|
||||
"Got Responses from API",
|
||||
vec![KeyValue::new(
|
||||
"Items in Response",
|
||||
results.response.exechistories.len().to_string(),
|
||||
)],
|
||||
);
|
||||
cx.span().set_status(Status::Ok);
|
||||
cx.span().set_attribute(KeyValue::new(
|
||||
"items_in_response",
|
||||
results.response.exechistories.len().to_string(),
|
||||
));
|
||||
results
|
||||
});
|
||||
let parsed_responses = execution_histories.response.exechistories;
|
||||
let parsed_responses: Vec<Group> = execution_histories.response.exechistories;
|
||||
if parsed_responses.is_empty() {
|
||||
break;
|
||||
}
|
||||
tx.send(parsed_responses.clone()).unwrap();
|
||||
match tx.send((cx.clone(), parsed_responses.clone())) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
cx.span().add_event(
|
||||
"Failed to Send Items to Processor",
|
||||
vec![KeyValue::new("Response from Processor", e.to_string())],
|
||||
);
|
||||
cx.span().set_status(Status::error("Processor Failed"))
|
||||
}
|
||||
}
|
||||
checkpoint_number = parsed_responses.last().unwrap().checkpoint.clone();
|
||||
if let Some(last_item) = parsed_responses.last()
|
||||
&& let Ok(last_date) = NaiveDate::parse_from_str(
|
||||
@@ -214,30 +270,37 @@ pub fn pull_policy_exec_histories(
|
||||
first_date = Some(last_date);
|
||||
}
|
||||
if let Some(base_date) = first_date {
|
||||
let date_diff = last_date - base_date;
|
||||
let total_span = (Local::now().naive_local().date() - base_date).num_days();
|
||||
let percentage = ((date_diff.num_days() as f64 / total_span as f64) * 100.0)
|
||||
let date_diff: chrono::TimeDelta = last_date - base_date;
|
||||
let total_span: i64 =
|
||||
(Local::now().naive_local().date() - base_date).num_days();
|
||||
let percentage: u64 = ((date_diff.num_days() as f64 / total_span as f64)
|
||||
* 100.0)
|
||||
.clamp(0.0, 100.0)
|
||||
.round() as u64;
|
||||
progress_bar.set_position(percentage);
|
||||
pb_clone.lock().unwrap().set_position(percentage);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
progress_bar.finish_with_message("All Checkpoints Complete");
|
||||
let return_data = match fs::read_to_string(file_path.clone()) {
|
||||
progress_bar
|
||||
.lock()
|
||||
.unwrap()
|
||||
.finish_with_message("All Checkpoints Complete");
|
||||
let return_data: String = match fs::read_to_string(file_path.clone()) {
|
||||
Ok(return_data) => return_data,
|
||||
Err(e) => {
|
||||
println!("Failed to read data from: {:?}: {}", &file_path, e);
|
||||
std::process::abort();
|
||||
}
|
||||
};
|
||||
tracer_provider
|
||||
.shutdown()
|
||||
.expect("Failed to Shutdown Tracer Provdier");
|
||||
drop(tx);
|
||||
shutdown_tracer_provider();
|
||||
return_data.to_string()
|
||||
});
|
||||
let gil_value = handle.join().unwrap();
|
||||
Python::attach(|py| PyString::new(py, &gil_value).into())
|
||||
let gil_value: String = handle.join().unwrap();
|
||||
Python::attach(|py: Python<'_>| PyString::new(py, &gil_value).into())
|
||||
}
|
||||
|
||||
fn build_client(headers: HeaderMap) -> Result<reqwest::Client, reqwest::Error> {
|
||||
@@ -248,23 +311,27 @@ fn build_client(headers: HeaderMap) -> Result<reqwest::Client, reqwest::Error> {
|
||||
.build()
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
#[tracing::instrument(name = "history_logging")]
|
||||
async fn history_logging(
|
||||
base_url: &String,
|
||||
exec_types: &String,
|
||||
checkpoint_number: &String,
|
||||
policy_names: &String,
|
||||
policy_names: &Option<String>,
|
||||
client: &Client,
|
||||
) -> ApiResponse {
|
||||
let policy_json = match policy_names {
|
||||
Some(name) => format!(r#"[ "{}" ]"#, name), // JSON array with one element
|
||||
None => "[]".to_string(), // Empty JSON array
|
||||
};
|
||||
let payload = format!(
|
||||
r#"{{
|
||||
"type": {},
|
||||
"checkpoint": "{}",
|
||||
"policy": ["{}"]
|
||||
"policy": {}
|
||||
}}"#,
|
||||
exec_types, checkpoint_number, policy_names
|
||||
exec_types, checkpoint_number, policy_json
|
||||
);
|
||||
let res = client
|
||||
let res: Result<reqwest::Response, reqwest::Error> = client
|
||||
.post(format!("{}/v1/logging/exechistories", base_url))
|
||||
.body(payload)
|
||||
.send()
|
||||
@@ -273,7 +340,7 @@ async fn history_logging(
|
||||
Ok(res) => {
|
||||
let first_response: ApiResponse = serde_json::from_str(&res.text().await.unwrap())
|
||||
.expect("Failed to retrieve response from API");
|
||||
return first_response;
|
||||
first_response
|
||||
}
|
||||
Err(_res) => {
|
||||
let failed_response: ApiResponse = ApiResponse {
|
||||
@@ -282,7 +349,7 @@ async fn history_logging(
|
||||
exechistories: vec![],
|
||||
},
|
||||
};
|
||||
return failed_response;
|
||||
failed_response
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -300,29 +367,10 @@ pub fn get_base_directory() -> PathBuf {
|
||||
.unwrap_or_else(|| home.join("AppData").join("Roaming"));
|
||||
appdata.join("Loxide")
|
||||
}
|
||||
_ => home.join(".local").join("share").join("Loxide"),
|
||||
"linux" => home.join(".local").join("share").join("Loxide"),
|
||||
_ => {
|
||||
println!("{} is currently not compatible with LoxideLibs", os);
|
||||
std::process::abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn init_tracer() -> Result<Option<sdktrace::Tracer>, TraceError> {
|
||||
let cfg = TelemetryConfig::load();
|
||||
if !cfg.TELEMETRY {
|
||||
global::set_tracer_provider(NoopTracerProvider::new());
|
||||
return Ok(None);
|
||||
}
|
||||
let endpoint = cfg.TELEM_URL.unwrap_or_default();
|
||||
let tracer =
|
||||
opentelemetry_otlp::new_pipeline()
|
||||
.tracing()
|
||||
.with_exporter(
|
||||
opentelemetry_otlp::new_exporter()
|
||||
.tonic()
|
||||
.with_endpoint(endpoint),
|
||||
)
|
||||
.with_trace_config(sdktrace::config().with_resource(Resource::new(vec![
|
||||
KeyValue::new("service.name", "LoxideLibs"),
|
||||
])))
|
||||
.install_simple()
|
||||
.unwrap();
|
||||
Ok(Some(tracer))
|
||||
}
|
||||
|
||||
+18
-6
@@ -1,14 +1,26 @@
|
||||
# Core TUI dependencies
|
||||
textual==6.5.0
|
||||
|
||||
# API and data handling
|
||||
Requests==2.32.5
|
||||
pandas==2.3.3
|
||||
numpy==2.3.4
|
||||
|
||||
# Database
|
||||
pymongo==4.15.3
|
||||
|
||||
# Security and encryption
|
||||
cryptography==46.0.3
|
||||
keyring==25.6.0
|
||||
numpy==2.3.4
|
||||
pandas==2.3.3
|
||||
pymongo==4.15.3
|
||||
|
||||
# Environment management
|
||||
python-dotenv==1.2.1
|
||||
Requests==2.32.5
|
||||
textual==6.5.0
|
||||
|
||||
# Utilities
|
||||
tqdm==4.67.1
|
||||
urllib3==2.5.0
|
||||
pyperclip==1.11.0
|
||||
|
||||
# Custom/Private packages
|
||||
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
|
||||
airlock_libs==5.1.1
|
||||
airlock_libs==7.3.0
|
||||
@@ -351,6 +351,14 @@ class AirlockAPIWrapper:
|
||||
result = self._post("/v1/getexechistory", payload)
|
||||
return result["response"]["exechistory"]
|
||||
|
||||
def server_logs(self, checkpoint: str | None = None) -> str:
|
||||
"""Retrieves Server Activity History Logs."""
|
||||
payload = {}
|
||||
if checkpoint is not None:
|
||||
payload["checkpoint"] = checkpoint
|
||||
result = self._post("/v1/logging/svractivities?checkpoint", payload)
|
||||
return result["response"]["svractivities"]
|
||||
|
||||
|
||||
"""
|
||||
from services.API import AirlockAPIWrapper
|
||||
|
||||
+15
-15
@@ -38,9 +38,9 @@ logger = logging.getLogger(__name__)
|
||||
def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||||
agents = selectAgents(api)
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
prompt="Enter how many days of history to pull (1–365): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
valid_range=(1, 365),
|
||||
)
|
||||
|
||||
if not agents or not history_days:
|
||||
@@ -60,7 +60,7 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||||
except Exception as e:
|
||||
print(
|
||||
colorText(
|
||||
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
|
||||
f"⌠Error retrieving history for {agent.hostname}: {e}", "red"
|
||||
)
|
||||
)
|
||||
continue
|
||||
@@ -139,7 +139,7 @@ def findAgents(api, return_dataframe):
|
||||
|
||||
print(
|
||||
colorText(
|
||||
f"\n✓ Matched devices exported to: {working_dir}\\{filename}",
|
||||
f"\n✓ Matched devices exported to: {working_dir}\\{filename}",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
@@ -148,7 +148,7 @@ def findAgents(api, return_dataframe):
|
||||
|
||||
|
||||
def collect_device_names() -> List[str]:
|
||||
print(colorText("🖥�� Device Search", "cyan"))
|
||||
print(colorText("🖥�� Device Search", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
"Enter the device hostnames you'd like to search for, one per line.", "cyan"
|
||||
@@ -185,7 +185,7 @@ def collect_device_names() -> List[str]:
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
|
||||
f"âš ï¸ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
@@ -235,8 +235,8 @@ def show_unmatched(
|
||||
]
|
||||
|
||||
if unmatched:
|
||||
logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
|
||||
print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
|
||||
logger.debug(f"âš ï¸ No matches for: {', '.join(unmatched)}")
|
||||
print(colorText(f"âš ï¸ No matches for: {', '.join(unmatched)}", "yellow"))
|
||||
|
||||
|
||||
def enrich_agents(agents: List["Agent"], policies: List["Policy"]):
|
||||
@@ -248,7 +248,7 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
||||
device_names = collect_device_names()
|
||||
if not device_names:
|
||||
logger.debug("No device names entered")
|
||||
print(colorText("⚠️ No device names entered.", "red"))
|
||||
print(colorText("âš ï¸ No device names entered.", "red"))
|
||||
return []
|
||||
|
||||
use_exact = choose_match_type()
|
||||
@@ -261,11 +261,11 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
||||
show_unmatched(device_names, matched_agents, use_exact)
|
||||
|
||||
if not matched_agents:
|
||||
logger.debug("❌ No matching devices found.")
|
||||
print(colorText("❌ No matching devices found.", "red"))
|
||||
logger.debug("⌠No matching devices found.")
|
||||
print(colorText("⌠No matching devices found.", "red"))
|
||||
return []
|
||||
|
||||
print(colorText(f"✓ Found {len(matched_agents)} matching device(s).", "green"))
|
||||
print(colorText(f"✓ Found {len(matched_agents)} matching device(s).", "green"))
|
||||
logger.info("Matched agent hostnames:")
|
||||
rows = (len(matched_agents) + 2) // 3 # 3 columns
|
||||
for row in range(rows):
|
||||
@@ -283,8 +283,8 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
||||
)
|
||||
|
||||
if not matched_agents:
|
||||
logger.debug("❌ No matching devices remain after refinement.")
|
||||
print(colorText("❌ No matching devices remain after refinement.", "red"))
|
||||
logger.debug("⌠No matching devices remain after refinement.")
|
||||
print(colorText("⌠No matching devices remain after refinement.", "red"))
|
||||
return []
|
||||
|
||||
enrich_agents(matched_agents, policies)
|
||||
@@ -302,7 +302,7 @@ def moveAgentToRelatedPolicy(
|
||||
Args:
|
||||
api: AirlockAPIWrapper instance.
|
||||
agent: Agent object.
|
||||
policy_relationship_map: Dict mapping enforcement â–€ –€™ audit.
|
||||
policy_relationship_map: Dict mapping enforcement â–€ –€™ audit.
|
||||
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
|
||||
"""
|
||||
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
|
||||
+75
-22
@@ -29,6 +29,49 @@ from utils.configmanager import (
|
||||
)
|
||||
|
||||
|
||||
class TextualNotificationHandler(logging.Handler):
|
||||
"""
|
||||
Custom logging handler that sends ERROR, WARNING, and CRITICAL logs
|
||||
to Textual toast notifications.
|
||||
"""
|
||||
|
||||
def __init__(self, app):
|
||||
super().__init__()
|
||||
self.app = app
|
||||
|
||||
def emit(self, record):
|
||||
try:
|
||||
# Only handle ERROR, WARNING, and CRITICAL
|
||||
if record.levelno >= logging.WARNING:
|
||||
# Format the message
|
||||
msg = self.format(record)
|
||||
|
||||
# Map log levels to Textual severity
|
||||
severity_map = {
|
||||
logging.WARNING: "warning",
|
||||
logging.ERROR: "error",
|
||||
logging.CRITICAL: "error",
|
||||
}
|
||||
severity = severity_map.get(record.levelno, "information")
|
||||
|
||||
# Send to Textual notification
|
||||
# Use call_from_thread if logging from non-main thread
|
||||
try:
|
||||
self.app.notify(msg, severity=severity, timeout=5)
|
||||
except Exception:
|
||||
# If we're not on the main thread, schedule it
|
||||
try:
|
||||
self.app.call_from_thread(
|
||||
self.app.notify, msg, severity=severity, timeout=5
|
||||
)
|
||||
except Exception:
|
||||
# Silently fail to avoid breaking the logging system
|
||||
pass
|
||||
except Exception:
|
||||
# Silently fail to avoid breaking the logging system
|
||||
pass
|
||||
|
||||
|
||||
def get_base_directory() -> Path:
|
||||
system = platform.system()
|
||||
home = Path.home()
|
||||
@@ -44,38 +87,31 @@ def configure_logging(log_dir: Path, log_level: str = "INFO"):
|
||||
log_file = log_dir / "Loxide.log"
|
||||
|
||||
config = {
|
||||
"version": 1, # Required key for dictConfig format version
|
||||
"disable_existing_loggers": False, # Keeps existing loggers active
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
"formatters": {
|
||||
"detailed": {
|
||||
"format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
# Includes timestamp, logger name, level, and message
|
||||
},
|
||||
"simple": {
|
||||
"format": "%(levelname)s - %(message)s"
|
||||
# Minimal format for console output
|
||||
},
|
||||
"simple": {"format": "%(levelname)s - %(message)s"},
|
||||
"toast": {"format": "%(name)s: %(message)s"}, # Simpler format for toasts
|
||||
},
|
||||
"handlers": {
|
||||
"file": {
|
||||
"class": "logging.handlers.TimedRotatingFileHandler",
|
||||
"filename": str(log_file),
|
||||
"when": "midnight", # Rotate logs at midnight
|
||||
"interval": 1, # Every 1 day
|
||||
"backupCount": 7, # Keep 7 days of logs
|
||||
"encoding": "utf-8", # Ensure UTF-8 encoding
|
||||
"level": "DEBUG", # Always log DEBUG and above to file
|
||||
"formatter": "detailed", # Use detailed format
|
||||
},
|
||||
"console": {
|
||||
"class": "logging.StreamHandler",
|
||||
"level": log_level.upper(), # System-configured level for console
|
||||
"formatter": "simple", # Use simple format
|
||||
"when": "midnight",
|
||||
"interval": 1,
|
||||
"backupCount": 7,
|
||||
"encoding": "utf-8",
|
||||
"level": "DEBUG",
|
||||
"formatter": "detailed",
|
||||
},
|
||||
# REMOVED console handler - it interferes with Textual TUI
|
||||
},
|
||||
"root": {
|
||||
"level": "DEBUG", # Root logger level
|
||||
"handlers": ["file", "console"], # Attach both handlers
|
||||
"level": "DEBUG",
|
||||
"handlers": ["file"], # Only use file handler, not console
|
||||
},
|
||||
}
|
||||
|
||||
@@ -96,6 +132,18 @@ def configure_logging(log_dir: Path, log_level: str = "INFO"):
|
||||
logging.config.dictConfig(config)
|
||||
logging.getLogger().debug("✅ Logging configured.")
|
||||
|
||||
# Return a function to attach the notification handler once the app is created
|
||||
def attach_notification_handler(app):
|
||||
"""Attach the Textual notification handler to the root logger."""
|
||||
handler = TextualNotificationHandler(app)
|
||||
handler.setLevel(logging.WARNING) # Only WARNING and above
|
||||
formatter = logging.Formatter("%(name)s: %(message)s")
|
||||
handler.setFormatter(formatter)
|
||||
logging.getLogger().addHandler(handler)
|
||||
logging.getLogger().debug("✅ Textual notification handler attached.")
|
||||
|
||||
return attach_notification_handler
|
||||
|
||||
|
||||
def setup():
|
||||
"""
|
||||
@@ -105,6 +153,9 @@ def setup():
|
||||
3. Load user config (mutable, from user_config.json)
|
||||
4. Configure logging
|
||||
5. Set up .env with WORKING_DIR only
|
||||
|
||||
Returns:
|
||||
attach_notification_handler: Function to attach notification handler to TUI app
|
||||
"""
|
||||
base_dir = get_base_directory()
|
||||
dirs = {
|
||||
@@ -122,7 +173,7 @@ def setup():
|
||||
|
||||
# Configure logging with system-defined log level
|
||||
log_level = get_system_value("LOG_LEVEL", str, "INFO")
|
||||
configure_logging(dirs["logs"], log_level)
|
||||
attach_handler = configure_logging(dirs["logs"], log_level)
|
||||
|
||||
# Load user config (mutable)
|
||||
load_user_config(dirs["config"])
|
||||
@@ -145,7 +196,6 @@ def setup():
|
||||
"Approved": [],
|
||||
"Needs_Review": ["Review_First", "Review_Second", "HTML"],
|
||||
"Preflight": ["HTML"],
|
||||
"Archived": [],
|
||||
}
|
||||
|
||||
for folder_name, subfolders in folders_structure.items():
|
||||
@@ -158,3 +208,6 @@ def setup():
|
||||
logging.debug(f"'{subfolder}' subfolder created at: {subfolder_path}")
|
||||
|
||||
logging.info("✅ Setup complete")
|
||||
|
||||
# Return the attach handler function
|
||||
return attach_handler
|
||||
|
||||
Reference in New Issue
Block a user