Compare commits
57 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 5555747422 | |||
| 729b45f52a | |||
| 66bb21ed88 | |||
| a7b659c951 | |||
| 53f0b548b0 | |||
| 66eb101c5d | |||
| 22101c1eba | |||
| fc17c869fc | |||
| 1dbbcff5d5 | |||
| f080b0034f | |||
| 57d0f12000 | |||
| 0dbc744471 | |||
| 7a912bddab | |||
| 24211c318b | |||
| 630e0a3cdf | |||
| 797d0f4462 | |||
| 59bb97ec4e | |||
| 0ac3b54d89 | |||
| 154a7efcc8 | |||
| ab5f00d8e7 | |||
| 3ab803c12e | |||
| 98cb23e5ea | |||
| 0aabbfd36e | |||
| b19eeb6c96 | |||
| 76bd3a6087 | |||
| 6ab9de413f | |||
| e5b4b9d959 | |||
| 1d9caadaf3 | |||
| f3c1d97d28 | |||
| 99d7b5f74e | |||
| 6327adeabf | |||
| b289e9324c | |||
| a80c2ca1e1 | |||
| dd206eb272 | |||
| 87ab1e3b28 | |||
| 303ecd8368 | |||
| 697d923172 | |||
| d0fc34fdc7 | |||
| 1f06404a16 | |||
| eb1d710d07 | |||
| fbb5cc396b | |||
| a7fc1b71e1 | |||
| 4a47cbe661 | |||
| 2f41b33dd4 | |||
| 89654d3a8c | |||
| b198362ac8 | |||
| 32e296238b | |||
| 19fab9b703 | |||
| 3ee762a0a1 | |||
| e36e5343d7 | |||
| ecdd991333 | |||
| 6a5a2b5809 | |||
| 2604665247 | |||
| bf5c7d156b | |||
| 0ebb42dcbd | |||
| ac6238873b | |||
| f0e77db414 |
@@ -14,7 +14,7 @@ jobs:
|
|||||||
- name: Install Prerequisites
|
- name: Install Prerequisites
|
||||||
run: |
|
run: |
|
||||||
apt update
|
apt update
|
||||||
apt install curl git python3 pip pkg-config openssl libssl-dev patchelf binutils-mingw-w64-x86-64 mingw-w64 -y
|
apt install curl git python3 pip pkg-config openssl libssl-dev patchelf binutils-mingw-w64-x86-64 mingw-w64 protobuf-compiler -y
|
||||||
curl https://sh.rustup.rs -sSf | sh -s -- -y
|
curl https://sh.rustup.rs -sSf | sh -s -- -y
|
||||||
pip install maturin twine --break-system-packages
|
pip install maturin twine --break-system-packages
|
||||||
|
|
||||||
|
|||||||
@@ -1,85 +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/>.
|
|
||||||
|
|
||||||
|
|
||||||
# TODO Continue implementing logger
|
|
||||||
# TODO Add input sanitation and CSV injection prevention
|
|
||||||
# TODO Continue OTP and Local approval rewrites
|
|
||||||
# TODO Explore pywin32
|
|
||||||
# TODO Fix Requirements.txt
|
|
||||||
# TODO Create Generic system_config.json for gitea
|
|
||||||
|
|
||||||
import logging
|
|
||||||
import os
|
|
||||||
import tempfile
|
|
||||||
|
|
||||||
import dotenv
|
|
||||||
import urllib3
|
|
||||||
|
|
||||||
from services.API import AirlockAPIWrapper
|
|
||||||
from services.security import getAPI
|
|
||||||
from utils.setup import get_base_directory, setup
|
|
||||||
from utils.TUI import run_Loxide
|
|
||||||
from utils.utils import irtang
|
|
||||||
|
|
||||||
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
|
|
||||||
if "NUITKA_ONEFILE_PARENT" in os.environ:
|
|
||||||
splash_filename = os.path.join(
|
|
||||||
tempfile.gettempdir(),
|
|
||||||
f"onefile_{int(os.environ['NUITKA_ONEFILE_PARENT'])}_splash_feedback.tmp",
|
|
||||||
)
|
|
||||||
if os.path.exists(splash_filename):
|
|
||||||
os.unlink(splash_filename)
|
|
||||||
|
|
||||||
irtang()
|
|
||||||
# Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
|
|
||||||
setup()
|
|
||||||
base_dir = get_base_directory()
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
dotenv.load_dotenv(dotenv_path=base_dir / ".env")
|
|
||||||
|
|
||||||
try:
|
|
||||||
url = os.getenv("URL")
|
|
||||||
username = os.getenv("USERNAME")
|
|
||||||
|
|
||||||
if not url:
|
|
||||||
raise ValueError("Missing URL in environment variables.")
|
|
||||||
if not username:
|
|
||||||
raise ValueError("Missing USERNAME in environment variables.")
|
|
||||||
|
|
||||||
logger.debug(f"Retrieved URL: {url}")
|
|
||||||
logger.debug(f"Retrieved Username: {username}")
|
|
||||||
|
|
||||||
except ValueError as e:
|
|
||||||
logger.error(f"Configuration error: {e}", exc_info=True)
|
|
||||||
raise
|
|
||||||
|
|
||||||
api_key = getAPI(username, "Loxide")
|
|
||||||
if api_key is None:
|
|
||||||
raise ValueError("API key for Loxide is missing.")
|
|
||||||
|
|
||||||
api = AirlockAPIWrapper(
|
|
||||||
base_url=str(os.getenv("URL")),
|
|
||||||
api_key=api_key,
|
|
||||||
)
|
|
||||||
run_Loxide(api)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -0,0 +1,606 @@
|
|||||||
|
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||||
|
#
|
||||||
|
# This program is free software: you can redistribute it and/or modify
|
||||||
|
# it under the terms of the GNU Affero General Public License as published
|
||||||
|
# by the Free Software Foundation, either version 3 of the License, or
|
||||||
|
# (at your option) any later version.
|
||||||
|
#
|
||||||
|
# This program is distributed in the hope that it will be useful,
|
||||||
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||||
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||||
|
# GNU Affero General Public License for more details.
|
||||||
|
#
|
||||||
|
# You should have received a copy of the GNU Affero General Public License
|
||||||
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||||
|
|
||||||
|
|
||||||
|
# TODO Continue implementing logger
|
||||||
|
# TODO Add input sanitation and CSV injection prevention
|
||||||
|
# TODO Continue OTP and Local approval rewrites
|
||||||
|
# TODO Explore pywin32
|
||||||
|
# TODO Fix Requirements.txt
|
||||||
|
# TODO Create Generic system_config.json for gitea
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
import dotenv
|
||||||
|
from textual.app import App, ComposeResult
|
||||||
|
from textual.containers import Vertical
|
||||||
|
from textual.message import Message
|
||||||
|
from textual.reactive import reactive
|
||||||
|
from textual.screen import Screen
|
||||||
|
from textual.widgets import (
|
||||||
|
Button,
|
||||||
|
DirectoryTree,
|
||||||
|
Footer,
|
||||||
|
Header,
|
||||||
|
Static,
|
||||||
|
Tab,
|
||||||
|
Tabs,
|
||||||
|
)
|
||||||
|
import urllib3
|
||||||
|
|
||||||
|
from models.agent import Agent
|
||||||
|
from models.policy import Policy
|
||||||
|
from services.API import AirlockAPIWrapper
|
||||||
|
from services.security import getAPI
|
||||||
|
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
|
||||||
|
from TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
||||||
|
from TUI.Screens.otpactivityscreen import OTPActivitiesScreen
|
||||||
|
from TUI.Screens.otprevokescreen import OTPRevokeScreen
|
||||||
|
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
|
||||||
|
from TUI.Screens.policyprepworkflowscreen import PolicyPrepWorkflowScreen
|
||||||
|
from TUI.Screens.quietagentworkflowscreen import QuietAgentWorkflowScreen
|
||||||
|
from TUI.Themes.theme_amber_terminal import get_amber_terminal_theme
|
||||||
|
from TUI.Themes.theme_retro_terminal import get_retro_terminal_theme
|
||||||
|
from TUI.Themes.themeselector import ThemeSelector
|
||||||
|
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
|
||||||
|
from TUI.Widgets.multiagentselector import MultiAgentSelector
|
||||||
|
from TUI.Widgets.policytreewidget import PolicyTreeWidget
|
||||||
|
from TUI.Widgets.resultsdisplay import ResultsDisplay
|
||||||
|
from TUI.Widgets.serverlogwidget import ServerLogWidget
|
||||||
|
from utils.configmanager import (
|
||||||
|
get_system_value,
|
||||||
|
get_user_value,
|
||||||
|
load_env,
|
||||||
|
save_user_config,
|
||||||
|
)
|
||||||
|
from utils.setup import get_base_directory, setup
|
||||||
|
from utils.utils import irtang, open_directory
|
||||||
|
|
||||||
|
dotenv.load_dotenv()
|
||||||
|
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# GLOBAL STASH
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
_APP_RESTART_REASON = None
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
def _persist_user_theme(theme_name: str) -> None:
|
||||||
|
"""
|
||||||
|
Store the chosen Textual theme in the user's config using the config manager.
|
||||||
|
No need to touch .env - config manager handles everything.
|
||||||
|
"""
|
||||||
|
base_dir = get_base_directory()
|
||||||
|
config_dir = base_dir / "config"
|
||||||
|
|
||||||
|
try:
|
||||||
|
save_user_config(config_dir, {"TEXTUAL_THEME": theme_name})
|
||||||
|
logger.debug("Updated user config with TEXTUAL_THEME=%s", theme_name)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("Failed to save TEXTUAL_THEME: %s", exc)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# 1) SCREEN
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
class MainMenuScreen(Screen):
|
||||||
|
api: AirlockAPIWrapper
|
||||||
|
current_tab = reactive("")
|
||||||
|
|
||||||
|
BUTTON_DEFS = {
|
||||||
|
"agent_actions": [
|
||||||
|
{
|
||||||
|
"label": "🖥️ - Multi-Agent Operations",
|
||||||
|
"id": "move_agent_workflow_button",
|
||||||
|
"description": "Select agents to: Move policies, Generate OTPs, Toggle audit/enforcement, View history, Export data",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"label": "🎫 - Review and approve OTP Activities",
|
||||||
|
"id": "otp_activities_button",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"label": "🛑 - Revoke Active OTP Session",
|
||||||
|
"id": "otp_revoke_button",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
"policy": [
|
||||||
|
{
|
||||||
|
"label": "⚖️ - Prepare Policy For Enforcement",
|
||||||
|
"id": "policy_prep_button",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"label": "🔕 - Find and Move Quiet Hosts to Enforcement",
|
||||||
|
"id": "find_quiet_button",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.extras = get_user_value("EXTRAS", str, "NOTTODAY")
|
||||||
|
wd = load_env("WORKING_DIR") or os.getcwd()
|
||||||
|
if not os.path.isdir(wd):
|
||||||
|
wd = os.getcwd()
|
||||||
|
self.working_dir = wd
|
||||||
|
|
||||||
|
def _make_buttons_for(self, tab_id: str) -> Vertical:
|
||||||
|
defs = self.BUTTON_DEFS.get(tab_id, [])
|
||||||
|
widgets = []
|
||||||
|
for item in defs:
|
||||||
|
# Support both old tuple format and new dict format
|
||||||
|
if isinstance(item, dict):
|
||||||
|
label = item["label"]
|
||||||
|
btn_id = item["id"]
|
||||||
|
description = item.get("description")
|
||||||
|
else:
|
||||||
|
# Old tuple format: (label, id)
|
||||||
|
label, btn_id = item
|
||||||
|
description = None
|
||||||
|
|
||||||
|
btn = Button(label, id=btn_id)
|
||||||
|
btn.styles.width = "100%"
|
||||||
|
widgets.append(btn)
|
||||||
|
|
||||||
|
# Add description text if provided
|
||||||
|
if description:
|
||||||
|
desc_text = Static(description, classes="button_description")
|
||||||
|
desc_text.styles.width = "100%"
|
||||||
|
desc_text.styles.color = "ansi_bright_black"
|
||||||
|
desc_text.styles.text_align = "center"
|
||||||
|
desc_text.styles.margin = (0, 0, 1, 0)
|
||||||
|
widgets.append(desc_text)
|
||||||
|
|
||||||
|
return Vertical(*widgets)
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
yield Header(show_clock=True, icon="⚙")
|
||||||
|
|
||||||
|
tabs = [
|
||||||
|
Tab("Agents", id="agent_actions"),
|
||||||
|
Tab("Tree View", id="p_tree"),
|
||||||
|
Tab("Server Log", id="server_log"),
|
||||||
|
Tab("Directory", id="dir"),
|
||||||
|
Tab("Settings", id="settings"),
|
||||||
|
]
|
||||||
|
|
||||||
|
if self.extras == "POLICYPREP":
|
||||||
|
tabs.insert(2, Tab("Policy Prep", id="policy"))
|
||||||
|
|
||||||
|
yield Tabs(*tabs, id="tabs")
|
||||||
|
yield Vertical(id="content")
|
||||||
|
yield Footer()
|
||||||
|
|
||||||
|
def on_mount(self) -> None:
|
||||||
|
self.switch_tab("agent_actions")
|
||||||
|
|
||||||
|
def on_key(self, event) -> None:
|
||||||
|
"""Handle up/down arrow keys for button navigation."""
|
||||||
|
if event.key == "down":
|
||||||
|
self._focus_nearby_button(1)
|
||||||
|
event.prevent_default()
|
||||||
|
event.stop()
|
||||||
|
elif event.key == "up":
|
||||||
|
self._focus_nearby_button(-1)
|
||||||
|
event.prevent_default()
|
||||||
|
event.stop()
|
||||||
|
# left/right are handled by Textual's default tab navigation
|
||||||
|
|
||||||
|
# focus helpers
|
||||||
|
def _get_content_buttons(self) -> list[Button]:
|
||||||
|
content = self.query_one("#content", Vertical)
|
||||||
|
return list(content.query(Button))
|
||||||
|
|
||||||
|
def _focus_first_button(self) -> None:
|
||||||
|
buttons = self._get_content_buttons()
|
||||||
|
if buttons:
|
||||||
|
buttons[0].focus()
|
||||||
|
|
||||||
|
def _focus_tabs(self) -> None:
|
||||||
|
tabs = self.query_one("#tabs", Tabs)
|
||||||
|
tabs.focus()
|
||||||
|
|
||||||
|
def _focus_nearby_button(self, direction: int) -> None:
|
||||||
|
buttons = self._get_content_buttons()
|
||||||
|
if not buttons:
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
current = next(i for i, b in enumerate(buttons) if b.has_focus)
|
||||||
|
except StopIteration:
|
||||||
|
if direction > 0:
|
||||||
|
buttons[0].focus()
|
||||||
|
else:
|
||||||
|
buttons[-1].focus()
|
||||||
|
return
|
||||||
|
|
||||||
|
if direction < 0 and current == 0:
|
||||||
|
self._focus_tabs()
|
||||||
|
return
|
||||||
|
|
||||||
|
new_index = current + direction
|
||||||
|
if 0 <= new_index < len(buttons):
|
||||||
|
buttons[new_index].focus()
|
||||||
|
|
||||||
|
def switch_tab(self, tab_id: str) -> None:
|
||||||
|
self.current_tab = tab_id
|
||||||
|
content = self.query_one("#content", Vertical)
|
||||||
|
content.remove_children()
|
||||||
|
|
||||||
|
if tab_id in self.BUTTON_DEFS:
|
||||||
|
content.mount(self._make_buttons_for(tab_id))
|
||||||
|
elif tab_id == "server_log":
|
||||||
|
content.mount(ServerLogWidget(self.app.api))
|
||||||
|
elif tab_id == "dir":
|
||||||
|
content.mount(DirectoryTree(self.working_dir, id="dir_tree"))
|
||||||
|
elif tab_id == "p_tree":
|
||||||
|
content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
|
||||||
|
elif tab_id == "settings":
|
||||||
|
content.mount(ThemeSelector())
|
||||||
|
else:
|
||||||
|
content.mount(Static(f"Unknown tab: {tab_id}"))
|
||||||
|
|
||||||
|
def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
|
||||||
|
self.switch_tab(event.tab.id)
|
||||||
|
|
||||||
|
def on_multi_agent_selector_agents_selected(
|
||||||
|
self, message: MultiAgentSelector.AgentsSelected
|
||||||
|
) -> None:
|
||||||
|
"""Handle selected agents from AgentSelector."""
|
||||||
|
global _APP_RESTART_REASON
|
||||||
|
selected_agents = message.selected_agents
|
||||||
|
logger.info("Selected agents: %s", selected_agents)
|
||||||
|
# TODO: Implement actual handling of selected agents
|
||||||
|
_APP_RESTART_REASON = ("multi_agent_action", selected_agents)
|
||||||
|
self.app.exit()
|
||||||
|
|
||||||
|
def on_theme_selector_theme_selected(
|
||||||
|
self, message: ThemeSelector.ThemeSelected
|
||||||
|
) -> None:
|
||||||
|
"""Handle theme selection from ThemeSelector."""
|
||||||
|
global _APP_RESTART_REASON
|
||||||
|
_persist_user_theme(message.theme_name)
|
||||||
|
_APP_RESTART_REASON = ("restart",)
|
||||||
|
self.app.exit()
|
||||||
|
|
||||||
|
def on_agent_move_operations_operation_complete(
|
||||||
|
self, message: AgentMoveOperations.OperationComplete
|
||||||
|
) -> None:
|
||||||
|
"""Handle completion of agent move operation - show results."""
|
||||||
|
logger.info(
|
||||||
|
"Agent move operation completed: %s, %d successful, %d unsuccessful",
|
||||||
|
message.operation,
|
||||||
|
len(message.successful),
|
||||||
|
len(message.unsuccessful),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Format results for display
|
||||||
|
successful_text = "\n".join(
|
||||||
|
[f"{agent.hostname}" for agent, _ in message.successful]
|
||||||
|
)
|
||||||
|
unsuccessful_text = "\n".join(
|
||||||
|
[f"{agent.hostname}: {error}" for agent, error in message.unsuccessful]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Remove the operations widget
|
||||||
|
try:
|
||||||
|
ops_widget = self.query_one(AgentMoveOperations)
|
||||||
|
ops_widget.remove()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Show results
|
||||||
|
self.query_one("#content", Vertical).mount(
|
||||||
|
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
|
||||||
|
)
|
||||||
|
|
||||||
|
def on_results_display_go_back(self, message: ResultsDisplay.GoBack) -> None:
|
||||||
|
"""Handle back button from results display."""
|
||||||
|
try:
|
||||||
|
results_widget = self.query_one(ResultsDisplay)
|
||||||
|
results_widget.remove()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
# Return to main menu
|
||||||
|
self.app.pop_screen()
|
||||||
|
|
||||||
|
def on_policy_tree_widget_view_execution_history(
|
||||||
|
self, message: PolicyTreeWidget.ViewExecutionHistory
|
||||||
|
) -> None:
|
||||||
|
"""Handle request to view execution history for a device from tree view."""
|
||||||
|
logger.info("Viewing execution history for device: %s", message.device.hostname)
|
||||||
|
self.app.push_screen(ExecutionHistoryScreen([message.device]))
|
||||||
|
message.stop()
|
||||||
|
|
||||||
|
def on_policy_tree_widget_generate_otp(
|
||||||
|
self, message: PolicyTreeWidget.GenerateOTP
|
||||||
|
) -> None:
|
||||||
|
"""Handle request to generate OTP for a device from tree view."""
|
||||||
|
logger.info("Generating OTP for device: %s", message.device.hostname)
|
||||||
|
self.app.push_screen(OTPWorkflowScreen([message.device]))
|
||||||
|
message.stop()
|
||||||
|
|
||||||
|
def on_policy_tree_widget_toggle_enforcement(
|
||||||
|
self, message: PolicyTreeWidget.ToggleEnforcement
|
||||||
|
) -> None:
|
||||||
|
"""Handle request to toggle enforcement for a device from tree view."""
|
||||||
|
logger.info("Toggling enforcement for device: %s", message.device.hostname)
|
||||||
|
|
||||||
|
try:
|
||||||
|
from services.agenthandler import moveAgentToRelatedPolicy
|
||||||
|
from utils.configmanager import get_system_json
|
||||||
|
|
||||||
|
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||||||
|
|
||||||
|
# Determine current mode and toggle
|
||||||
|
if message.device.groupid in policy_relationship_map:
|
||||||
|
# Currently in enforcement, move to audit
|
||||||
|
result = moveAgentToRelatedPolicy(self.app.api, message.device, "audit")
|
||||||
|
mode = "audit"
|
||||||
|
else:
|
||||||
|
# Currently in audit, move to enforcement
|
||||||
|
result = moveAgentToRelatedPolicy(
|
||||||
|
self.app.api, message.device, "enforcement"
|
||||||
|
)
|
||||||
|
mode = "enforcement"
|
||||||
|
|
||||||
|
logger.info(f"Successfully toggled {message.device.hostname} to {mode}")
|
||||||
|
|
||||||
|
# Refresh data at the app level
|
||||||
|
self.app.refresh_data()
|
||||||
|
|
||||||
|
# Refresh the tree widget with new data
|
||||||
|
try:
|
||||||
|
tree_widget = self.query_one(PolicyTreeWidget)
|
||||||
|
tree_widget.refresh_data(self.app.policies, self.app.devices)
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
f"Failed to toggle enforcement for {message.device.hostname}: {e}"
|
||||||
|
)
|
||||||
|
self.app.bell()
|
||||||
|
|
||||||
|
message.stop()
|
||||||
|
|
||||||
|
def on_directory_tree_file_selected(
|
||||||
|
self, event: DirectoryTree.FileSelected
|
||||||
|
) -> None:
|
||||||
|
path = event.path
|
||||||
|
logger.debug("Directory file selected: %s", path)
|
||||||
|
try:
|
||||||
|
open_directory(str(path))
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("Failed to open %s: %s", path, exc)
|
||||||
|
self.app.bell()
|
||||||
|
|
||||||
|
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||||
|
button_id = event.button.id
|
||||||
|
logger.debug("Button pressed: %s", button_id)
|
||||||
|
|
||||||
|
match button_id:
|
||||||
|
case "move_agent_workflow_button":
|
||||||
|
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
|
||||||
|
event.stop()
|
||||||
|
|
||||||
|
case "otp_generate_button":
|
||||||
|
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
|
||||||
|
event.stop()
|
||||||
|
|
||||||
|
case "find_quiet_button":
|
||||||
|
self.app.push_screen(
|
||||||
|
QuietAgentWorkflowScreen(self.app.api, self.app.policies)
|
||||||
|
)
|
||||||
|
event.stop()
|
||||||
|
return
|
||||||
|
|
||||||
|
case "otp_activities_button":
|
||||||
|
self.app.push_screen(OTPActivitiesScreen())
|
||||||
|
event.stop()
|
||||||
|
return
|
||||||
|
|
||||||
|
case "otp_revoke_button":
|
||||||
|
self.app.push_screen(OTPRevokeScreen())
|
||||||
|
event.stop()
|
||||||
|
return
|
||||||
|
|
||||||
|
case "policy_prep_button":
|
||||||
|
# Use the new TUI workflow screen instead of legacy
|
||||||
|
self.app.push_screen(
|
||||||
|
PolicyPrepWorkflowScreen(self.app.api, self.app.policies)
|
||||||
|
)
|
||||||
|
event.stop()
|
||||||
|
return
|
||||||
|
|
||||||
|
case _:
|
||||||
|
self.app.bell()
|
||||||
|
logger.warning("Unknown button pressed: %s", button_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# 2) APP
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
class Loxide(App[Message]):
|
||||||
|
api: AirlockAPIWrapper
|
||||||
|
working_dir: str
|
||||||
|
policies: Optional[list[Policy]]
|
||||||
|
devices: Optional[list[Agent]]
|
||||||
|
|
||||||
|
CSS = """
|
||||||
|
#logo {
|
||||||
|
width: 100%;
|
||||||
|
content-align: center middle;
|
||||||
|
text-align: center;
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
BINDINGS = [
|
||||||
|
("q", "quit", "Quit"),
|
||||||
|
("f", "open_fe", "Launch Explorer"),
|
||||||
|
("r", "refresh", "Refresh"),
|
||||||
|
]
|
||||||
|
|
||||||
|
def __init__(self, api: AirlockAPIWrapper):
|
||||||
|
self._textual_theme = get_user_value("TEXTUAL_THEME", str, "textual-dark")
|
||||||
|
super().__init__()
|
||||||
|
self.api = api
|
||||||
|
wd = load_env("WORKING_DIR") or os.getcwd()
|
||||||
|
if not os.path.isdir(wd):
|
||||||
|
wd = os.getcwd()
|
||||||
|
self.working_dir = wd
|
||||||
|
# Initial data load
|
||||||
|
self.refresh_data()
|
||||||
|
|
||||||
|
def refresh_data(self) -> None:
|
||||||
|
"""Public method to refresh policies and devices from the API."""
|
||||||
|
try:
|
||||||
|
self.policies = [
|
||||||
|
Policy(**row.to_dict())
|
||||||
|
for _, row in self.api.policy_find_all().iterrows()
|
||||||
|
]
|
||||||
|
self.devices = [
|
||||||
|
Agent(**row.to_dict())
|
||||||
|
for _, row in self.api.agent_find_all().iterrows()
|
||||||
|
]
|
||||||
|
if self.policies and self.devices:
|
||||||
|
for agent in self.devices:
|
||||||
|
agent.enrich_with_policies(self.policies)
|
||||||
|
logger.debug(
|
||||||
|
f"Enriched {len(self.devices)} agents with policy information"
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("Failed to load policies/devices: %s", exc)
|
||||||
|
self.policies = None
|
||||||
|
self.devices = None
|
||||||
|
|
||||||
|
def on_mount(self, api: AirlockAPIWrapper) -> None:
|
||||||
|
self.register_theme(get_retro_terminal_theme())
|
||||||
|
self.register_theme(get_amber_terminal_theme())
|
||||||
|
self.theme = self._textual_theme
|
||||||
|
self.push_screen(MainMenuScreen())
|
||||||
|
|
||||||
|
def action_refresh(self) -> None:
|
||||||
|
self.refresh_data()
|
||||||
|
|
||||||
|
def action_quit(self) -> None:
|
||||||
|
global _APP_RESTART_REASON
|
||||||
|
_APP_RESTART_REASON = None
|
||||||
|
self.exit()
|
||||||
|
|
||||||
|
def action_open_fe(self) -> None:
|
||||||
|
"""Open the working directory in the OS file manager (footer binding)."""
|
||||||
|
path_to_open = self.working_dir or os.getcwd()
|
||||||
|
try:
|
||||||
|
open_directory(path_to_open)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("Failed to open directory %s: %s", path_to_open, exc)
|
||||||
|
self.bell() # optional feedback
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# 3) PUBLIC ENTRYPOINT - Updated to accept attach_notification_handler
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
def run_Loxide(api: AirlockAPIWrapper, attach_notification_handler=None) -> None:
|
||||||
|
global _APP_RESTART_REASON
|
||||||
|
base_dir = get_base_directory()
|
||||||
|
env_path = base_dir / ".env"
|
||||||
|
dotenv.load_dotenv(dotenv_path=env_path, override=True)
|
||||||
|
|
||||||
|
max_attempts = 5
|
||||||
|
attempts = 0
|
||||||
|
|
||||||
|
while attempts < max_attempts:
|
||||||
|
attempts += 1
|
||||||
|
logger.debug("Starting app loop iteration (attempt %d)", attempts)
|
||||||
|
_APP_RESTART_REASON = None
|
||||||
|
app = Loxide(api)
|
||||||
|
|
||||||
|
# Attach the notification handler if provided
|
||||||
|
if attach_notification_handler:
|
||||||
|
attach_notification_handler(app)
|
||||||
|
|
||||||
|
try:
|
||||||
|
app.run()
|
||||||
|
except SystemExit as exc:
|
||||||
|
if exc.code != 0:
|
||||||
|
logger.debug("Caught SystemExit from Textual: %s", exc)
|
||||||
|
raise
|
||||||
|
|
||||||
|
reason = _APP_RESTART_REASON
|
||||||
|
logger.debug("After app.run(), _APP_RESTART_REASON = %r", reason)
|
||||||
|
|
||||||
|
if not reason:
|
||||||
|
logger.debug("No restart reason, exiting loop")
|
||||||
|
break
|
||||||
|
|
||||||
|
if reason[0] == "restart":
|
||||||
|
logger.debug("Restarting app loop")
|
||||||
|
continue
|
||||||
|
|
||||||
|
if reason[0] == "multi_agent_action":
|
||||||
|
logger.info("Multi-agent action with selected agents: %s", reason[1])
|
||||||
|
continue
|
||||||
|
|
||||||
|
logger.error("Unknown restart reason: %r", reason)
|
||||||
|
break
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# 4) MAIN FUNCTION - Updated to get and pass attach_notification_handler
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
def main():
|
||||||
|
irtang()
|
||||||
|
# Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
|
||||||
|
# setup() now returns a function to attach the notification handler
|
||||||
|
attach_notification_handler = setup()
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
try:
|
||||||
|
url = get_system_value("URL")
|
||||||
|
username = os.getenv("USERNAME")
|
||||||
|
|
||||||
|
if not url:
|
||||||
|
raise ValueError("Missing URL in environment variables.")
|
||||||
|
if not username:
|
||||||
|
raise ValueError("Missing USERNAME in environment variables.")
|
||||||
|
|
||||||
|
logger.debug(f"Retrieved URL: {url}")
|
||||||
|
logger.debug(f"Retrieved Username: {username}")
|
||||||
|
|
||||||
|
except ValueError as e:
|
||||||
|
logger.error(f"Configuration error: {e}", exc_info=True)
|
||||||
|
raise
|
||||||
|
|
||||||
|
api_key = getAPI(username, "Loxide")
|
||||||
|
if api_key is None:
|
||||||
|
raise ValueError("API key for Loxide is missing.")
|
||||||
|
|
||||||
|
api = AirlockAPIWrapper(
|
||||||
|
base_url=str(url),
|
||||||
|
api_key=api_key,
|
||||||
|
)
|
||||||
|
run_Loxide(api, attach_notification_handler)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Binary file not shown.
|
After Width: | Height: | Size: 112 KiB |
@@ -0,0 +1,668 @@
|
|||||||
|
# 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 __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
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,
|
||||||
|
Static,
|
||||||
|
TextArea,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class AllowlistSelectionWidget(Static):
|
||||||
|
"""
|
||||||
|
Widget for selecting an allowlist and adding hashes to it.
|
||||||
|
Can be reused in different workflows.
|
||||||
|
"""
|
||||||
|
|
||||||
|
DEFAULT_CSS = """
|
||||||
|
AllowlistSelectionWidget {
|
||||||
|
height: 1fr;
|
||||||
|
layout: vertical;
|
||||||
|
}
|
||||||
|
#allowlist_main {
|
||||||
|
height: 1fr;
|
||||||
|
width: 100%;
|
||||||
|
}
|
||||||
|
#left_panel {
|
||||||
|
width: 50%;
|
||||||
|
padding: 1;
|
||||||
|
border: solid $primary;
|
||||||
|
}
|
||||||
|
#right_panel {
|
||||||
|
width: 50%;
|
||||||
|
padding: 1;
|
||||||
|
border: solid $primary;
|
||||||
|
}
|
||||||
|
#allowlist_table {
|
||||||
|
height: 70%;
|
||||||
|
margin: 1 0;
|
||||||
|
}
|
||||||
|
#allowlist_table > .datatable--header {
|
||||||
|
text-style: bold;
|
||||||
|
background: $boost;
|
||||||
|
}
|
||||||
|
#allowlist_table Row {
|
||||||
|
height: 1;
|
||||||
|
}
|
||||||
|
#preview_area {
|
||||||
|
height: 60%;
|
||||||
|
margin: 1 0;
|
||||||
|
}
|
||||||
|
#action_buttons {
|
||||||
|
height: auto;
|
||||||
|
min-height: 3;
|
||||||
|
padding: 1;
|
||||||
|
content-align: center middle;
|
||||||
|
}
|
||||||
|
.panel-title {
|
||||||
|
text-style: bold;
|
||||||
|
margin: 0 0 1 0;
|
||||||
|
}
|
||||||
|
.info-text {
|
||||||
|
margin: 1 0;
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
selected_data: pd.DataFrame,
|
||||||
|
api=None,
|
||||||
|
hostname: Optional[str] = None,
|
||||||
|
otpid: Optional[str] = None,
|
||||||
|
hash_column: str = "sha256", # Default hash column name
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize the allowlist selection widget.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
selected_data: DataFrame containing the selected activities
|
||||||
|
api: API instance for making allowlist calls
|
||||||
|
hostname: Optional hostname for context
|
||||||
|
otpid: Optional OTP ID for context
|
||||||
|
hash_column: Name of the column containing hashes (default: "sha256")
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
self.selected_data = selected_data
|
||||||
|
self.api = api
|
||||||
|
self.hostname = hostname
|
||||||
|
self.otpid = otpid
|
||||||
|
self.hash_column = hash_column
|
||||||
|
self.allowlists = []
|
||||||
|
self.selected_allowlist = None
|
||||||
|
self.hashes_to_add = []
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
with Horizontal(id="allowlist_main"):
|
||||||
|
# Left panel - Allowlist selection
|
||||||
|
with Vertical(id="left_panel"):
|
||||||
|
yield Static("Select Allowlist", classes="panel-title")
|
||||||
|
yield Static(
|
||||||
|
f"Choose an allowlist to add {len(self.selected_data)} selected items",
|
||||||
|
classes="info-text",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Allowlist table
|
||||||
|
self.allowlist_table = DataTable(id="allowlist_table")
|
||||||
|
self.allowlist_table.cursor_type = "row"
|
||||||
|
yield self.allowlist_table
|
||||||
|
|
||||||
|
# Refresh button
|
||||||
|
self.refresh_btn = Button(
|
||||||
|
"🔄 Refresh Allowlists", id="refresh_allowlists_btn"
|
||||||
|
)
|
||||||
|
yield self.refresh_btn
|
||||||
|
|
||||||
|
# Right panel - Preview and actions
|
||||||
|
with Vertical(id="right_panel"):
|
||||||
|
yield Static("Preview", classes="panel-title")
|
||||||
|
|
||||||
|
# Context information
|
||||||
|
context_text = []
|
||||||
|
if self.hostname:
|
||||||
|
context_text.append(f"Host: {self.hostname}")
|
||||||
|
if self.otpid:
|
||||||
|
context_text.append(f"OTP: {self.otpid}")
|
||||||
|
context_text.append(f"Selected Activities: {len(self.selected_data)}")
|
||||||
|
|
||||||
|
yield Static(" | ".join(context_text), classes="info-text")
|
||||||
|
|
||||||
|
# Preview text area
|
||||||
|
self.preview_area = TextArea(
|
||||||
|
id="preview_area", read_only=True, language="markdown"
|
||||||
|
)
|
||||||
|
yield self.preview_area
|
||||||
|
|
||||||
|
# Hash statistics
|
||||||
|
self.stats_label = Static("", id="stats_label", classes="info-text")
|
||||||
|
yield self.stats_label
|
||||||
|
|
||||||
|
# Action buttons at bottom
|
||||||
|
with Horizontal(id="action_buttons"):
|
||||||
|
self.add_btn = Button("➕ Add to Allowlist", id="add_to_allowlist_btn")
|
||||||
|
|
||||||
|
self.add_btn.styles.width = "100%"
|
||||||
|
self.add_btn.disabled = True # Disabled until allowlist selected
|
||||||
|
|
||||||
|
yield self.add_btn
|
||||||
|
|
||||||
|
async def on_mount(self) -> None:
|
||||||
|
"""Load allowlists when widget mounts."""
|
||||||
|
await self.load_allowlists()
|
||||||
|
await self.extract_and_preview_hashes()
|
||||||
|
|
||||||
|
async def load_allowlists(self) -> None:
|
||||||
|
"""Load available allowlists from API, grouped by policy association."""
|
||||||
|
if not self.api:
|
||||||
|
logger.error("No API available")
|
||||||
|
self.allowlist_table.add_column("Error")
|
||||||
|
self.allowlist_table.add_row("No API available")
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
# First, try to get the host's policy if hostname is provided
|
||||||
|
host_policy_allowlists = []
|
||||||
|
host_policy_ids = set()
|
||||||
|
policy_name = "Unknown Policy" # Default value
|
||||||
|
group_id = None
|
||||||
|
|
||||||
|
if self.hostname:
|
||||||
|
try:
|
||||||
|
# Get agent info to find its policy
|
||||||
|
agents_df = self.api.agent_find_by_hostname(self.hostname)
|
||||||
|
if not agents_df.empty:
|
||||||
|
# Get the policy group ID for this host
|
||||||
|
group_id = agents_df.iloc[0].get("groupid")
|
||||||
|
|
||||||
|
# Look up the policy name from app's cached policies
|
||||||
|
if (
|
||||||
|
group_id
|
||||||
|
and hasattr(self.app, "policies")
|
||||||
|
and self.app.policies
|
||||||
|
):
|
||||||
|
for policy in self.app.policies:
|
||||||
|
if policy.groupid == group_id:
|
||||||
|
policy_name = policy.name
|
||||||
|
logger.info(
|
||||||
|
f"Found policy name: '{policy_name}' for group_id: {group_id}"
|
||||||
|
)
|
||||||
|
break
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
f"Found host '{self.hostname}' in policy '{policy_name}' (group_id: {group_id})"
|
||||||
|
)
|
||||||
|
|
||||||
|
if group_id:
|
||||||
|
# Get allowlists for this policy
|
||||||
|
policy_allowlists_df = self.api.policy_list_allowlists(
|
||||||
|
group_id
|
||||||
|
)
|
||||||
|
if not policy_allowlists_df.empty:
|
||||||
|
host_policy_allowlists = policy_allowlists_df.to_dict(
|
||||||
|
orient="records"
|
||||||
|
)
|
||||||
|
host_policy_ids = {
|
||||||
|
al.get("applicationid")
|
||||||
|
for al in host_policy_allowlists
|
||||||
|
}
|
||||||
|
logger.info(
|
||||||
|
f"Found {len(host_policy_allowlists)} allowlists for host's policy"
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Could not get host's policy allowlists: {e}")
|
||||||
|
|
||||||
|
# If we still don't have a policy name, try to get it from the first allowlist or use a default
|
||||||
|
if not policy_name:
|
||||||
|
# Get all policies and try to find which one has allowlists
|
||||||
|
try:
|
||||||
|
all_policies_df = self.api.policy_find_all()
|
||||||
|
if not all_policies_df.empty:
|
||||||
|
# If we have a group_id from somewhere, use it
|
||||||
|
if group_id:
|
||||||
|
policy_row = all_policies_df[
|
||||||
|
all_policies_df["groupid"] == group_id
|
||||||
|
]
|
||||||
|
if not policy_row.empty:
|
||||||
|
policy_name = policy_row.iloc[0].get(
|
||||||
|
"groupname", "Unknown Policy"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Use the first policy as fallback
|
||||||
|
policy_name = all_policies_df.iloc[0].get(
|
||||||
|
"groupname", "Default Policy"
|
||||||
|
)
|
||||||
|
logger.info(f"Using first available policy: {policy_name}")
|
||||||
|
else:
|
||||||
|
policy_name = "Unknown Policy"
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Could not fetch policies: {e}")
|
||||||
|
policy_name = "Unknown Policy"
|
||||||
|
|
||||||
|
# Get all allowlists
|
||||||
|
all_allowlists_df = self.api.allowlist_find_all()
|
||||||
|
|
||||||
|
if all_allowlists_df.empty:
|
||||||
|
self.allowlist_table.add_column("No Allowlists")
|
||||||
|
self.allowlist_table.add_row("No allowlists found")
|
||||||
|
return
|
||||||
|
|
||||||
|
all_allowlists = all_allowlists_df.to_dict(orient="records")
|
||||||
|
|
||||||
|
# Separate into two groups: policy-associated and others
|
||||||
|
other_allowlists = [
|
||||||
|
al
|
||||||
|
for al in all_allowlists
|
||||||
|
if al.get("applicationid") not in host_policy_ids
|
||||||
|
]
|
||||||
|
|
||||||
|
# Sort each group alphabetically by name
|
||||||
|
host_policy_allowlists.sort(key=lambda x: x.get("name", "").lower())
|
||||||
|
other_allowlists.sort(key=lambda x: x.get("name", "").lower())
|
||||||
|
|
||||||
|
# Combine lists with policy-associated first
|
||||||
|
self.allowlists = host_policy_allowlists + other_allowlists
|
||||||
|
|
||||||
|
# Setup table columns
|
||||||
|
self.allowlist_table.clear()
|
||||||
|
self.allowlist_table.add_columns("Name", "Application ID", "Type")
|
||||||
|
|
||||||
|
# Track which rows are headers vs actual allowlists
|
||||||
|
self._row_to_allowlist_map = {}
|
||||||
|
current_row = 0
|
||||||
|
|
||||||
|
# Add policy-associated allowlists if any
|
||||||
|
if host_policy_allowlists:
|
||||||
|
# Add section header
|
||||||
|
header_text = f"=== Policy: {policy_name or 'Host Policy'} ==="
|
||||||
|
self.allowlist_table.add_row(header_text, "", "", key="header_policy")
|
||||||
|
current_row += 1
|
||||||
|
|
||||||
|
# Add policy allowlists
|
||||||
|
for idx, allowlist in enumerate(host_policy_allowlists):
|
||||||
|
name = allowlist.get("name", "Unknown")
|
||||||
|
app_id = allowlist.get("applicationid", "Unknown")
|
||||||
|
|
||||||
|
self.allowlist_table.add_row(
|
||||||
|
f" {name}", # Indent to show grouping
|
||||||
|
app_id,
|
||||||
|
"Policy",
|
||||||
|
key=f"policy_{idx}",
|
||||||
|
)
|
||||||
|
self._row_to_allowlist_map[current_row] = idx
|
||||||
|
current_row += 1
|
||||||
|
|
||||||
|
# Add other allowlists
|
||||||
|
if other_allowlists:
|
||||||
|
# Add section header
|
||||||
|
if host_policy_allowlists:
|
||||||
|
# Add spacer if we have policy allowlists above
|
||||||
|
self.allowlist_table.add_row("", "", "", key="spacer")
|
||||||
|
current_row += 1
|
||||||
|
|
||||||
|
self.allowlist_table.add_row(
|
||||||
|
"=== Other Available Allowlists ===", "", "", key="header_other"
|
||||||
|
)
|
||||||
|
current_row += 1
|
||||||
|
|
||||||
|
# Add other allowlists
|
||||||
|
for idx, allowlist in enumerate(other_allowlists):
|
||||||
|
name = allowlist.get("name", "Unknown")
|
||||||
|
app_id = allowlist.get("applicationid", "Unknown")
|
||||||
|
|
||||||
|
self.allowlist_table.add_row(
|
||||||
|
f" {name}", # Indent to show grouping
|
||||||
|
app_id,
|
||||||
|
"General",
|
||||||
|
key=f"other_{idx}",
|
||||||
|
)
|
||||||
|
# Map to the correct index in the combined list
|
||||||
|
actual_idx = len(host_policy_allowlists) + idx
|
||||||
|
self._row_to_allowlist_map[current_row] = actual_idx
|
||||||
|
current_row += 1
|
||||||
|
|
||||||
|
# Log summary
|
||||||
|
logger.info(
|
||||||
|
f"Loaded {len(self.allowlists)} total allowlists: "
|
||||||
|
f"{len(host_policy_allowlists)} policy-associated, "
|
||||||
|
f"{len(other_allowlists)} others"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Update stats label if no allowlists in policy
|
||||||
|
if self.hostname and not host_policy_allowlists:
|
||||||
|
self.stats_label.update(
|
||||||
|
f"Note: No allowlists found for {self.hostname}'s policy | "
|
||||||
|
+ self.stats_label.content.plain
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception(f"Failed to load allowlists: {exc}")
|
||||||
|
self.allowlist_table.add_column("Error")
|
||||||
|
self.allowlist_table.add_row(f"Failed to load: {str(exc)}")
|
||||||
|
|
||||||
|
async def extract_and_preview_hashes(self) -> None:
|
||||||
|
"""Extract hashes from selected data and show preview."""
|
||||||
|
preview_lines = ["## Hash Extraction Summary\n"]
|
||||||
|
|
||||||
|
# Check for hash column
|
||||||
|
if self.hash_column not in self.selected_data.columns:
|
||||||
|
# Try to find a hash column
|
||||||
|
possible_hash_cols = [
|
||||||
|
"sha256",
|
||||||
|
"SHA256",
|
||||||
|
"hash",
|
||||||
|
"Hash",
|
||||||
|
"sha1",
|
||||||
|
"SHA1",
|
||||||
|
"md5",
|
||||||
|
"MD5",
|
||||||
|
"filehash",
|
||||||
|
"file_hash",
|
||||||
|
]
|
||||||
|
found_col = None
|
||||||
|
for col in possible_hash_cols:
|
||||||
|
if col in self.selected_data.columns:
|
||||||
|
found_col = col
|
||||||
|
break
|
||||||
|
|
||||||
|
if found_col:
|
||||||
|
self.hash_column = found_col
|
||||||
|
preview_lines.append(f"✅ Found hash column: **{found_col}**\n")
|
||||||
|
else:
|
||||||
|
preview_lines.append("❌ **No hash column found**\n")
|
||||||
|
preview_lines.append("Available columns:\n")
|
||||||
|
for col in self.selected_data.columns:
|
||||||
|
if col != "_row_id":
|
||||||
|
preview_lines.append(f" - {col}\n")
|
||||||
|
|
||||||
|
self.preview_area.text = "".join(preview_lines)
|
||||||
|
self.stats_label.update("No hashes to add")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Extract unique hashes
|
||||||
|
hashes = self.selected_data[self.hash_column].dropna().unique()
|
||||||
|
self.hashes_to_add = [h for h in hashes if h and str(h).strip()]
|
||||||
|
|
||||||
|
# Build preview
|
||||||
|
preview_lines.append(f"### Found {len(self.hashes_to_add)} unique hashes\n\n")
|
||||||
|
|
||||||
|
# Show sample of hashes (first 10)
|
||||||
|
preview_lines.append("**Sample hashes to be added:**\n```\n")
|
||||||
|
for i, hash_val in enumerate(self.hashes_to_add[:10]):
|
||||||
|
preview_lines.append(f"{i+1}. {hash_val}\n")
|
||||||
|
if len(self.hashes_to_add) > 10:
|
||||||
|
preview_lines.append(f"... and {len(self.hashes_to_add) - 10} more\n")
|
||||||
|
preview_lines.append("```\n\n")
|
||||||
|
|
||||||
|
# Show sample of source data
|
||||||
|
preview_lines.append("**Sample source activities:**\n")
|
||||||
|
sample_cols = [
|
||||||
|
col
|
||||||
|
for col in self.selected_data.columns
|
||||||
|
if col not in ["_row_id"] and col in ["filename", "path", "action", "user"]
|
||||||
|
]
|
||||||
|
if not sample_cols:
|
||||||
|
sample_cols = [
|
||||||
|
col for col in self.selected_data.columns if col != "_row_id"
|
||||||
|
][:3]
|
||||||
|
|
||||||
|
if sample_cols:
|
||||||
|
preview_lines.append("```\n")
|
||||||
|
for i, row in self.selected_data[sample_cols].head(5).iterrows():
|
||||||
|
row_text = " | ".join([f"{col}: {row[col]}" for col in sample_cols])
|
||||||
|
preview_lines.append(f"{row_text}\n")
|
||||||
|
preview_lines.append("```\n")
|
||||||
|
|
||||||
|
self.preview_area.text = "".join(preview_lines)
|
||||||
|
|
||||||
|
# Update statistics
|
||||||
|
self.stats_label.update(
|
||||||
|
f"Ready to add {len(self.hashes_to_add)} unique hashes | "
|
||||||
|
f"From {len(self.selected_data)} selected activities"
|
||||||
|
)
|
||||||
|
|
||||||
|
async def on_data_table_row_selected(self, event) -> None:
|
||||||
|
"""Handle allowlist selection."""
|
||||||
|
try:
|
||||||
|
# Extract row index from event - handle different event structures
|
||||||
|
row_index = None
|
||||||
|
|
||||||
|
# Try to get row index from coordinate
|
||||||
|
if hasattr(event, "coordinate") and hasattr(event.coordinate, "row"):
|
||||||
|
row_index = event.coordinate.row
|
||||||
|
# Try cursor_row as fallback
|
||||||
|
elif hasattr(event, "cursor_row"):
|
||||||
|
row_index = event.cursor_row
|
||||||
|
# Try getting from the table itself
|
||||||
|
else:
|
||||||
|
table = self.allowlist_table
|
||||||
|
if hasattr(table, "cursor_row"):
|
||||||
|
row_index = table.cursor_row
|
||||||
|
|
||||||
|
# Validate row index
|
||||||
|
if row_index is not None and isinstance(row_index, int):
|
||||||
|
# Account for group headers in the row count
|
||||||
|
actual_allowlist_index = self._get_allowlist_index_from_row(row_index)
|
||||||
|
|
||||||
|
if (
|
||||||
|
actual_allowlist_index is not None
|
||||||
|
and 0 <= actual_allowlist_index < len(self.allowlists)
|
||||||
|
):
|
||||||
|
self.selected_allowlist = self.allowlists[actual_allowlist_index]
|
||||||
|
self.add_btn.disabled = False
|
||||||
|
self.add_btn.label = (
|
||||||
|
f"➕ Add to '{self.selected_allowlist.get('name', 'Unknown')}'"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Update preview with selection
|
||||||
|
await self._update_preview_with_selection()
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
f"Selected allowlist: {self.selected_allowlist.get('name')}"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.debug(f"Row {row_index} is a header or invalid")
|
||||||
|
else:
|
||||||
|
logger.warning(f"Could not extract valid row index from event: {event}")
|
||||||
|
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception(f"Failed to select allowlist: {exc}")
|
||||||
|
|
||||||
|
def _get_allowlist_index_from_row(self, row_index: int) -> Optional[int]:
|
||||||
|
"""Convert table row index to allowlist list index, accounting for group headers."""
|
||||||
|
# This will be updated when we have group headers
|
||||||
|
if hasattr(self, "_row_to_allowlist_map"):
|
||||||
|
return self._row_to_allowlist_map.get(row_index)
|
||||||
|
return row_index
|
||||||
|
|
||||||
|
async def _update_preview_with_selection(self) -> None:
|
||||||
|
"""Update preview when an allowlist is selected."""
|
||||||
|
if not self.selected_allowlist:
|
||||||
|
return
|
||||||
|
|
||||||
|
current_text = self.preview_area.text
|
||||||
|
# Remove any existing selection header
|
||||||
|
if "### Selected Allowlist:" in current_text:
|
||||||
|
lines = current_text.split("\n")
|
||||||
|
# Find and remove the selection lines
|
||||||
|
new_lines = []
|
||||||
|
skip_next = False
|
||||||
|
for line in lines:
|
||||||
|
if line.startswith("### Selected Allowlist:"):
|
||||||
|
skip_next = True
|
||||||
|
continue
|
||||||
|
if skip_next and line.startswith("Application ID:"):
|
||||||
|
skip_next = False
|
||||||
|
continue
|
||||||
|
if not skip_next:
|
||||||
|
new_lines.append(line)
|
||||||
|
current_text = "\n".join(new_lines)
|
||||||
|
|
||||||
|
# Add new selection at the top
|
||||||
|
selection_text = (
|
||||||
|
f"### Selected Allowlist: **{self.selected_allowlist.get('name')}**\n"
|
||||||
|
f"Application ID: {self.selected_allowlist.get('applicationid')}\n\n"
|
||||||
|
)
|
||||||
|
self.preview_area.text = selection_text + current_text
|
||||||
|
|
||||||
|
async def on_button_pressed(self, event) -> None:
|
||||||
|
"""Handle button presses."""
|
||||||
|
btn = getattr(event, "button", None) or getattr(event, "sender", None)
|
||||||
|
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None)
|
||||||
|
|
||||||
|
if btn is self.refresh_btn or btn_id == "refresh_allowlists_btn":
|
||||||
|
await self.load_allowlists()
|
||||||
|
event.stop()
|
||||||
|
return
|
||||||
|
|
||||||
|
if btn is self.add_btn or btn_id == "add_to_allowlist_btn":
|
||||||
|
await self.add_hashes_to_allowlist()
|
||||||
|
event.stop()
|
||||||
|
return
|
||||||
|
|
||||||
|
async def add_hashes_to_allowlist(self) -> None:
|
||||||
|
"""Add the extracted hashes to the selected allowlist."""
|
||||||
|
if not self.selected_allowlist or not self.hashes_to_add:
|
||||||
|
self.app.notify(
|
||||||
|
"No allowlist selected or no hashes to add", severity="warning"
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
if not self.api:
|
||||||
|
self.app.notify("API not available", severity="error")
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Disable button during operation
|
||||||
|
self.add_btn.disabled = True
|
||||||
|
self.add_btn.label = "Adding hashes..."
|
||||||
|
|
||||||
|
# Call API to add hashes
|
||||||
|
app_id = self.selected_allowlist.get("applicationid")
|
||||||
|
allowlist_name = self.selected_allowlist.get("name", "Unknown")
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
f"Adding {len(self.hashes_to_add)} hashes to allowlist {allowlist_name} (ID: {app_id})"
|
||||||
|
)
|
||||||
|
|
||||||
|
result = self.api.hash_add_to_allowlist(app_id, self.hashes_to_add)
|
||||||
|
logger.debug(f"Hash adding api call: {result}")
|
||||||
|
# Success notification
|
||||||
|
self.app.notify(
|
||||||
|
f"Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'",
|
||||||
|
title="Success",
|
||||||
|
severity="information",
|
||||||
|
timeout=5,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Update preview to show success
|
||||||
|
self.preview_area.text = (
|
||||||
|
f"## SUCCESS\n\n"
|
||||||
|
f"Added **{len(self.hashes_to_add)} hashes** to allowlist:\n"
|
||||||
|
f"**{allowlist_name}** (ID: {app_id})\n\n"
|
||||||
|
f"### Operation Details:\n"
|
||||||
|
f"- Source: {self.hostname or 'Multiple hosts'}\n"
|
||||||
|
f"- OTP ID: {self.otpid or 'N/A'}\n"
|
||||||
|
f"- Activities processed: {len(self.selected_data)}\n"
|
||||||
|
f"- Unique hashes added: {len(self.hashes_to_add)}\n"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Change button to "Done"
|
||||||
|
self.add_btn.label = "Done - Press q to return to main menu"
|
||||||
|
self.add_btn.disabled = True
|
||||||
|
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception(f"Failed to add hashes to allowlist: {exc}")
|
||||||
|
self.app.notify(
|
||||||
|
f"Failed to add hashes: {str(exc)}",
|
||||||
|
title="Error",
|
||||||
|
severity="error",
|
||||||
|
timeout=10,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Re-enable button
|
||||||
|
self.add_btn.disabled = False
|
||||||
|
self.add_btn.label = "Retry Add to Allowlist"
|
||||||
|
|
||||||
|
|
||||||
|
class AllowlistSelectionScreen(Screen):
|
||||||
|
"""
|
||||||
|
Screen wrapper for the AllowlistSelectionWidget.
|
||||||
|
"""
|
||||||
|
|
||||||
|
BINDINGS = [
|
||||||
|
Binding("escape", "go_back", "Back"),
|
||||||
|
Binding("q", "main_menu", "Main Menu"),
|
||||||
|
Binding("r", "refresh", "Refresh Allowlists"),
|
||||||
|
]
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
selected_data: pd.DataFrame,
|
||||||
|
api=None,
|
||||||
|
hostname: Optional[str] = None,
|
||||||
|
otpid: Optional[str] = None,
|
||||||
|
hash_column: str = "sha256",
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.selected_data = selected_data
|
||||||
|
self.api = api
|
||||||
|
self.hostname = hostname
|
||||||
|
self.otpid = otpid
|
||||||
|
self.hash_column = hash_column
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
yield Header(show_clock=True)
|
||||||
|
self.widget = AllowlistSelectionWidget(
|
||||||
|
self.selected_data,
|
||||||
|
api=self.api,
|
||||||
|
hostname=self.hostname,
|
||||||
|
otpid=self.otpid,
|
||||||
|
hash_column=self.hash_column,
|
||||||
|
)
|
||||||
|
yield self.widget
|
||||||
|
yield Footer()
|
||||||
|
|
||||||
|
async def action_go_back(self) -> None:
|
||||||
|
"""Go back to previous screen."""
|
||||||
|
await self.app.pop_screen()
|
||||||
|
|
||||||
|
async def action_main_menu(self) -> None:
|
||||||
|
"""Go back to main menu."""
|
||||||
|
while len(self.app.screen_stack) > 2:
|
||||||
|
await self.app.pop_screen()
|
||||||
|
|
||||||
|
async def action_refresh(self) -> None:
|
||||||
|
"""Refresh the allowlists."""
|
||||||
|
if hasattr(self, "widget") and self.widget:
|
||||||
|
await self.widget.load_allowlists()
|
||||||
|
|
||||||
|
async def action_confirm(self) -> None:
|
||||||
|
"""Confirm and add to allowlist."""
|
||||||
|
if hasattr(self, "widget") and self.widget:
|
||||||
|
if self.widget.selected_allowlist and self.widget.hashes_to_add:
|
||||||
|
await self.widget.add_hashes_to_allowlist()
|
||||||
@@ -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()
|
||||||
@@ -0,0 +1,114 @@
|
|||||||
|
# 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 typing import List, Optional
|
||||||
|
|
||||||
|
from textual.app import ComposeResult
|
||||||
|
from textual.binding import Binding
|
||||||
|
from textual.css.query import NoMatches
|
||||||
|
from textual.screen import Screen
|
||||||
|
|
||||||
|
from models.agent import Agent
|
||||||
|
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
|
||||||
|
from TUI.Widgets.multiagentselector import MultiAgentSelector
|
||||||
|
from TUI.Widgets.resultsdisplay import ResultsDisplay
|
||||||
|
|
||||||
|
|
||||||
|
class MoveAgentWorkflowScreen(Screen):
|
||||||
|
"""Screen that handles the agent movement workflow."""
|
||||||
|
|
||||||
|
BINDINGS = [
|
||||||
|
Binding("escape", "go_back", "Back"),
|
||||||
|
Binding("q", "main_menu", "Main Menu"),
|
||||||
|
]
|
||||||
|
|
||||||
|
def __init__(self, all_agents: Optional[List[Agent]]):
|
||||||
|
super().__init__()
|
||||||
|
self.all_agents = all_agents
|
||||||
|
self.selected_agents = None
|
||||||
|
self.workflow_stage = "select_agents" # Track current stage
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
"""Start with the multi-agent selector."""
|
||||||
|
yield MultiAgentSelector(self.all_agents)
|
||||||
|
|
||||||
|
def action_go_back(self) -> None:
|
||||||
|
"""Handle escape key to go back one step within the workflow."""
|
||||||
|
if self.workflow_stage == "select_agents":
|
||||||
|
# At first stage, go back to main menu
|
||||||
|
self.app.pop_screen()
|
||||||
|
elif self.workflow_stage == "operations":
|
||||||
|
# Go back to agent selection
|
||||||
|
try:
|
||||||
|
ops_widget = self.query_one(AgentMoveOperations)
|
||||||
|
ops_widget.remove()
|
||||||
|
except NoMatches:
|
||||||
|
pass
|
||||||
|
self.mount(MultiAgentSelector(self.all_agents))
|
||||||
|
self.workflow_stage = "select_agents"
|
||||||
|
elif self.workflow_stage == "results":
|
||||||
|
# Go back to operations
|
||||||
|
try:
|
||||||
|
results_widget = self.query_one(ResultsDisplay)
|
||||||
|
results_widget.remove()
|
||||||
|
except NoMatches:
|
||||||
|
pass
|
||||||
|
self.mount(AgentMoveOperations(self.selected_agents))
|
||||||
|
self.workflow_stage = "operations"
|
||||||
|
|
||||||
|
def action_main_menu(self) -> None:
|
||||||
|
"""Handle q key to go back to main menu."""
|
||||||
|
while len(self.app.screen_stack) > 2:
|
||||||
|
self.app.pop_screen()
|
||||||
|
|
||||||
|
def on_multi_agent_selector_agents_selected(
|
||||||
|
self, message: MultiAgentSelector.AgentsSelected
|
||||||
|
) -> None:
|
||||||
|
"""Handle selected agents - switch to operations screen."""
|
||||||
|
self.selected_agents = message.selected_agents
|
||||||
|
|
||||||
|
# Remove the MultiAgentSelector
|
||||||
|
selector = self.query_one(MultiAgentSelector)
|
||||||
|
selector.remove()
|
||||||
|
|
||||||
|
# Mount the AgentMoveOperations with the selected Agent objects
|
||||||
|
self.mount(AgentMoveOperations(self.selected_agents))
|
||||||
|
self.workflow_stage = "operations"
|
||||||
|
|
||||||
|
def on_agent_move_operations_operation_complete(
|
||||||
|
self, message: AgentMoveOperations.OperationComplete
|
||||||
|
) -> None:
|
||||||
|
"""Handle completion of move operation - transition to results screen."""
|
||||||
|
# Format successful results
|
||||||
|
success_lines = []
|
||||||
|
for agent, result in message.successful:
|
||||||
|
success_lines.append(f"✔ {agent.hostname}")
|
||||||
|
|
||||||
|
# Format unsuccessful results
|
||||||
|
failure_lines = []
|
||||||
|
for agent, error in message.unsuccessful:
|
||||||
|
failure_lines.append(f"❌ — {agent.hostname}: {error}")
|
||||||
|
|
||||||
|
successful_text = "\n".join(success_lines) if success_lines else "(none)"
|
||||||
|
unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)"
|
||||||
|
|
||||||
|
# Remove the operations widget
|
||||||
|
ops_widget = self.query_one(AgentMoveOperations)
|
||||||
|
ops_widget.remove()
|
||||||
|
|
||||||
|
# Mount the results display
|
||||||
|
self.mount(
|
||||||
|
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
|
||||||
|
)
|
||||||
|
self.workflow_stage = "results"
|
||||||
@@ -0,0 +1,901 @@
|
|||||||
|
# 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 __future__ import annotations
|
||||||
|
|
||||||
|
from datetime import datetime
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
|
||||||
|
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, Static
|
||||||
|
|
||||||
|
from TUI.Screens.allowlistselectionscreen import AllowlistSelectionScreen
|
||||||
|
from utils.configmanager import load_env
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _load_working_dir() -> str:
|
||||||
|
"""
|
||||||
|
Load the working directory from environment variables or use the current working directory.
|
||||||
|
"""
|
||||||
|
wd = os.environ.get("WORKING_DIR")
|
||||||
|
if wd:
|
||||||
|
return wd
|
||||||
|
return os.getcwd()
|
||||||
|
|
||||||
|
|
||||||
|
class OTPActivitiesWidget(Static):
|
||||||
|
"""
|
||||||
|
Reusable widget that contains the sessions table (left) and an Activity Preview (right).
|
||||||
|
The right side shows an Activity Preview that takes ~75% vertical space, and a lower area
|
||||||
|
with Continue button.
|
||||||
|
"""
|
||||||
|
|
||||||
|
DEFAULT_CSS = """
|
||||||
|
OTPActivitiesWidget {
|
||||||
|
height: 1fr;
|
||||||
|
}
|
||||||
|
#main_row {
|
||||||
|
width: 100%;
|
||||||
|
height: 100%;
|
||||||
|
layout: horizontal;
|
||||||
|
}
|
||||||
|
#left_panel {
|
||||||
|
width: 60%;
|
||||||
|
min-width: 60;
|
||||||
|
border: none;
|
||||||
|
}
|
||||||
|
#right_panel {
|
||||||
|
width: 40%;
|
||||||
|
min-width: 40;
|
||||||
|
border: none;
|
||||||
|
layout: vertical;
|
||||||
|
}
|
||||||
|
#activity_preview_container {
|
||||||
|
height: 1fr;
|
||||||
|
border: none;
|
||||||
|
padding: 1 1;
|
||||||
|
}
|
||||||
|
#activity_buttons {
|
||||||
|
height: auto;
|
||||||
|
min-height: 3;
|
||||||
|
padding: 1 1;
|
||||||
|
content-align: center middle;
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
# Layout: horizontal main row with left & right panels
|
||||||
|
with Horizontal(id="main_row"):
|
||||||
|
# Left: sessions area
|
||||||
|
with Vertical(id="left_panel"):
|
||||||
|
yield Static("OTP Sessions", classes="panel-title")
|
||||||
|
with Vertical(id="sessions_table_container"):
|
||||||
|
self.sessions_table = DataTable(id="sessions_table")
|
||||||
|
self.sessions_table.styles.width = "100%"
|
||||||
|
yield self.sessions_table
|
||||||
|
# Right: Activity Preview (top 3/4) + buttons (bottom 1/4)
|
||||||
|
with Vertical(id="right_panel"):
|
||||||
|
# Activity preview area (takes ~75% of right panel)
|
||||||
|
yield Static("Activity Preview", classes="panel-title")
|
||||||
|
with Vertical(id="activity_preview_container"):
|
||||||
|
self.activities_table = DataTable(id="activity_preview_table")
|
||||||
|
yield self.activities_table
|
||||||
|
# Button area at the bottom (Continue)
|
||||||
|
with Horizontal(id="activity_buttons"):
|
||||||
|
self.continue_btn = Button("Continue", id="activity_continue_btn")
|
||||||
|
self.continue_btn.styles.width = "100%"
|
||||||
|
yield self.continue_btn
|
||||||
|
|
||||||
|
async def on_mount(self) -> None:
|
||||||
|
# Configure sessions table and activities preview
|
||||||
|
self.sessions_table.clear()
|
||||||
|
self.sessions_table.add_columns(
|
||||||
|
"otpid", "hostname", "status", "purpose", "granted"
|
||||||
|
)
|
||||||
|
self.activities_table.clear()
|
||||||
|
# activities_table columns are dynamically added when activities are loaded.
|
||||||
|
# Selection behavior
|
||||||
|
self.sessions_table.cursor_type = "row"
|
||||||
|
try:
|
||||||
|
self.sessions_table.zebra_stripes = True
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
self.activities_table.cursor_type = "row"
|
||||||
|
try:
|
||||||
|
self.activities_table.zebra_stripes = True
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
# Store state
|
||||||
|
self._sessions_df: pd.DataFrame | None = None
|
||||||
|
self._activities_df: pd.DataFrame | None = None
|
||||||
|
self._selected_session_otpid: str | int | None = None
|
||||||
|
|
||||||
|
async def on_button_pressed(self, event) -> None: # type: ignore[override]
|
||||||
|
"""
|
||||||
|
Handle Continue button for the Activity Preview area.
|
||||||
|
"""
|
||||||
|
# Try to resolve the button object from the event
|
||||||
|
btn = (
|
||||||
|
getattr(event, "button", None)
|
||||||
|
or getattr(event, "sender", None)
|
||||||
|
or getattr(event, "control", None)
|
||||||
|
or getattr(event, "widget", None)
|
||||||
|
)
|
||||||
|
btn_id = (
|
||||||
|
getattr(btn, "id", None)
|
||||||
|
or getattr(event, "button_id", None)
|
||||||
|
or getattr(event, "id", None)
|
||||||
|
)
|
||||||
|
# ---- Continue ----
|
||||||
|
if btn is self.continue_btn or btn_id == getattr(self.continue_btn, "id", None):
|
||||||
|
if self._activities_df is None or self._activities_df.empty:
|
||||||
|
logger.info("Continue pressed but no activities loaded.")
|
||||||
|
self.app.notify(
|
||||||
|
"No activities loaded to continue with.", severity="warning"
|
||||||
|
)
|
||||||
|
return
|
||||||
|
# Copy activities DataFrame to pass to new screen
|
||||||
|
activities_copy = self._activities_df.copy()
|
||||||
|
otpid = self._selected_session_otpid
|
||||||
|
# Optionally include hostname if available
|
||||||
|
hostname = None
|
||||||
|
try:
|
||||||
|
if self._sessions_df is not None:
|
||||||
|
df = self._sessions_df.reset_index(drop=True)
|
||||||
|
match = df[df["otpid"] == otpid]
|
||||||
|
if not match.empty:
|
||||||
|
hostname = match.iloc[0].get("hostname")
|
||||||
|
except Exception:
|
||||||
|
hostname = None
|
||||||
|
# Create and push ActivityDetailScreen, handing the data
|
||||||
|
try:
|
||||||
|
detail_screen = ActivityDetailScreen(
|
||||||
|
activities_copy, otpid=otpid, hostname=hostname
|
||||||
|
)
|
||||||
|
await self.app.push_screen(detail_screen)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception("Failed to push ActivityDetailScreen: %s", exc)
|
||||||
|
return
|
||||||
|
# Unknown button on widget
|
||||||
|
logger.debug(
|
||||||
|
"Unhandled OTPActivitiesWidget button pressed (resolved btn=%r, id=%r)",
|
||||||
|
btn,
|
||||||
|
btn_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def on_data_table_row_selected(self, event) -> None: # type: ignore[override]
|
||||||
|
"""
|
||||||
|
Robust handler for DataTable row-selection across Textual micro-versions.
|
||||||
|
Tries many attribute names and shapes:
|
||||||
|
- numeric index (row_key, row_index, index)
|
||||||
|
- coordinate object or tuple (coordinate.row or (row, col))
|
||||||
|
- direct row values (row, values, cells) -> we try to map those back to the sessions DF
|
||||||
|
- table.cursor_row fallback
|
||||||
|
"""
|
||||||
|
# 1) Determine the sending table (best-effort)
|
||||||
|
sender = None
|
||||||
|
for attr in ("sender", "table", "data_table", "control"):
|
||||||
|
sender = getattr(event, attr, None)
|
||||||
|
if sender is not None:
|
||||||
|
break
|
||||||
|
if sender is None:
|
||||||
|
sender = self.sessions_table # Assume sessions_table if unknown
|
||||||
|
# Only respond to selections in the sessions table
|
||||||
|
if sender is not self.sessions_table:
|
||||||
|
return
|
||||||
|
|
||||||
|
# Helper to log and return
|
||||||
|
def _bad(msg: str, *args):
|
||||||
|
logger.warning(msg, *args)
|
||||||
|
return None
|
||||||
|
|
||||||
|
# 2) Try to extract a numeric index
|
||||||
|
row_key = None
|
||||||
|
for attr in ("row_key", "row", "row_index", "index"):
|
||||||
|
row_key = getattr(event, attr, None)
|
||||||
|
if row_key is not None:
|
||||||
|
break
|
||||||
|
# If coordinate: try to extract .row or tuple[0]
|
||||||
|
if row_key is None:
|
||||||
|
coord = getattr(event, "coordinate", None) or getattr(
|
||||||
|
event, "cursor_coordinate", None
|
||||||
|
)
|
||||||
|
if coord is not None:
|
||||||
|
if hasattr(coord, "row"):
|
||||||
|
row_key = coord.row
|
||||||
|
elif isinstance(coord, (tuple, list)) and len(coord) >= 1:
|
||||||
|
row_key = coord[0]
|
||||||
|
# If still nothing, maybe the event provides the row's cell values directly
|
||||||
|
row_values = None
|
||||||
|
for attr in ("values", "cells", "row", "row_values", "selected_row_values"):
|
||||||
|
val = getattr(event, attr, None)
|
||||||
|
if val:
|
||||||
|
# Prefer actual sequence of cell values
|
||||||
|
row_values = val
|
||||||
|
break
|
||||||
|
# If we have row_values, try to map them back to the sessions DataFrame
|
||||||
|
if row_values is not None:
|
||||||
|
# Normalize into list of strings for comparison
|
||||||
|
try:
|
||||||
|
vals = [
|
||||||
|
"" if pd.isna(v) else str(v)
|
||||||
|
for v in (
|
||||||
|
list(row_values)
|
||||||
|
if not isinstance(row_values, str)
|
||||||
|
else [row_values]
|
||||||
|
)
|
||||||
|
]
|
||||||
|
except Exception:
|
||||||
|
vals = [str(row_values)]
|
||||||
|
# Try to match against the expected columns order we render
|
||||||
|
if self._sessions_df is None or self._sessions_df.empty:
|
||||||
|
logger.warning(
|
||||||
|
"Sessions DataFrame is empty; cannot map selected row values."
|
||||||
|
)
|
||||||
|
return
|
||||||
|
df_ordered = self._sessions_df.reset_index(drop=True)
|
||||||
|
expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]
|
||||||
|
|
||||||
|
# Build stringified candidates for each row in df using the same columns we show
|
||||||
|
def _row_to_vals(sr):
|
||||||
|
out = []
|
||||||
|
for c in expected_cols:
|
||||||
|
if c in sr:
|
||||||
|
v = sr[c]
|
||||||
|
out.append("" if pd.isna(v) else str(v))
|
||||||
|
else:
|
||||||
|
out.append("")
|
||||||
|
return out
|
||||||
|
|
||||||
|
match_idx = None
|
||||||
|
for i, sr in df_ordered.iterrows():
|
||||||
|
cand = _row_to_vals(sr)
|
||||||
|
# Compare prefix: row values might be a subset (e.g. only first 3 cols), so compare prefix only
|
||||||
|
if len(vals) <= len(cand) and all(
|
||||||
|
vals[j] == cand[j] for j in range(len(vals))
|
||||||
|
):
|
||||||
|
match_idx = i
|
||||||
|
break
|
||||||
|
if match_idx is None:
|
||||||
|
# Try looser match: compare first cell only (otpid)
|
||||||
|
first = vals[0] if vals else None
|
||||||
|
if first is not None:
|
||||||
|
for i, sr in df_ordered.iterrows():
|
||||||
|
cand0 = "" if pd.isna(sr.get("otpid")) else str(sr.get("otpid"))
|
||||||
|
if cand0 == first:
|
||||||
|
match_idx = i
|
||||||
|
break
|
||||||
|
if match_idx is None:
|
||||||
|
logger.warning(
|
||||||
|
"Unable to locate DataFrame row matching selected row values: %r",
|
||||||
|
vals,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
idx = int(match_idx)
|
||||||
|
else:
|
||||||
|
# 3) If we have a row_key, try to normalize to an int index
|
||||||
|
if row_key is not None:
|
||||||
|
try:
|
||||||
|
idx = int(row_key)
|
||||||
|
except Exception:
|
||||||
|
# Try converting via string
|
||||||
|
try:
|
||||||
|
idx = int(str(row_key))
|
||||||
|
except Exception:
|
||||||
|
idx = None
|
||||||
|
if idx is None:
|
||||||
|
# Final numeric fallback: use sessions_table.cursor_row if present
|
||||||
|
try:
|
||||||
|
idx = getattr(self.sessions_table, "cursor_row")
|
||||||
|
except Exception:
|
||||||
|
idx = None
|
||||||
|
if idx is None:
|
||||||
|
_bad("Failed to normalize row/key from event: %r", row_key)
|
||||||
|
return
|
||||||
|
else:
|
||||||
|
# 4) Try table cursor_row as last resort
|
||||||
|
try:
|
||||||
|
idx = getattr(self.sessions_table, "cursor_row")
|
||||||
|
except Exception:
|
||||||
|
logger.warning(
|
||||||
|
"Could not determine selected row from event: %r", event
|
||||||
|
)
|
||||||
|
# Helpful debug hint for you to paste back if still failing:
|
||||||
|
logger.debug("Event repr for debugging: %r", event)
|
||||||
|
return
|
||||||
|
# At this point we should have an integer idx
|
||||||
|
try:
|
||||||
|
idx = int(idx)
|
||||||
|
except Exception:
|
||||||
|
logger.exception(
|
||||||
|
"Final normalization of selected row index failed: %r", idx
|
||||||
|
)
|
||||||
|
return
|
||||||
|
# Validate sessions df
|
||||||
|
if self._sessions_df is None or self._sessions_df.empty:
|
||||||
|
logger.warning("Sessions DataFrame empty; nothing to select.")
|
||||||
|
return
|
||||||
|
df_ordered = self._sessions_df.reset_index(drop=True)
|
||||||
|
if idx < 0 or idx >= len(df_ordered):
|
||||||
|
logger.warning(
|
||||||
|
"Selected row index %s out of range (0..%d)", idx, len(df_ordered) - 1
|
||||||
|
)
|
||||||
|
return
|
||||||
|
row_series = df_ordered.iloc[idx]
|
||||||
|
otpid = row_series.get("otpid")
|
||||||
|
hostname = row_series.get("hostname")
|
||||||
|
# Store selected session and fetch activities
|
||||||
|
self._selected_session_otpid = otpid
|
||||||
|
# Obtain api from app (try multiple places)
|
||||||
|
api = (
|
||||||
|
getattr(self.app, "api", None)
|
||||||
|
or getattr(self, "api", None)
|
||||||
|
or getattr(self.app, "airlock_api", None)
|
||||||
|
)
|
||||||
|
if api is None:
|
||||||
|
logger.error("No API available on self.app.api - cannot fetch activities")
|
||||||
|
return
|
||||||
|
logger.info(
|
||||||
|
"Fetching activities for otpid=%s host=%s (selected row=%s)",
|
||||||
|
otpid,
|
||||||
|
hostname,
|
||||||
|
idx,
|
||||||
|
)
|
||||||
|
await self._fetch_activities_for_otpid(api, otpid, hostname=hostname)
|
||||||
|
|
||||||
|
async def load_sessions_from_api(self, api) -> None:
|
||||||
|
"""
|
||||||
|
Pulls OTP session lists, adds status column, concatenates and populates the sessions table.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
active = api.otp_find_active()
|
||||||
|
awaiting = api.otp_find_awaiting()
|
||||||
|
enforced = api.otp_find_enforced()
|
||||||
|
revoked = api.otp_find_revoked()
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception("Failed to fetch OTP session lists: %s", exc)
|
||||||
|
# Present empty
|
||||||
|
active = awaiting = enforced = revoked = pd.DataFrame()
|
||||||
|
|
||||||
|
# Ensure DataFrame objects
|
||||||
|
def _ensure_df(df):
|
||||||
|
return df if isinstance(df, pd.DataFrame) else pd.DataFrame(df)
|
||||||
|
|
||||||
|
active = _ensure_df(active)
|
||||||
|
awaiting = _ensure_df(awaiting)
|
||||||
|
enforced = _ensure_df(enforced)
|
||||||
|
revoked = _ensure_df(revoked)
|
||||||
|
for df, status in [
|
||||||
|
(active, "active"),
|
||||||
|
(awaiting, "awaiting"),
|
||||||
|
(enforced, "enforced"),
|
||||||
|
(revoked, "revoked"),
|
||||||
|
]:
|
||||||
|
if "status" not in df.columns:
|
||||||
|
df["status"] = status
|
||||||
|
combined = pd.concat([active, awaiting, enforced, revoked], ignore_index=True)
|
||||||
|
if "otpid" in combined.columns:
|
||||||
|
combined = combined.sort_values(by="otpid", ascending=False)
|
||||||
|
self._sessions_df = combined
|
||||||
|
# Populate DataTable
|
||||||
|
self.sessions_table.clear()
|
||||||
|
# Ensure columns exist in DF and when missing add empty column
|
||||||
|
expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]
|
||||||
|
for col in expected_cols:
|
||||||
|
if col not in combined.columns:
|
||||||
|
combined[col] = ""
|
||||||
|
self.sessions_table.add_columns(*expected_cols)
|
||||||
|
# Add rows
|
||||||
|
for _, row in combined[expected_cols].iterrows():
|
||||||
|
# Convert values to str for safe insertion
|
||||||
|
vals = ["" if pd.isna(v) else v for v in row.to_list()]
|
||||||
|
self.sessions_table.add_row(*[str(v) for v in vals])
|
||||||
|
logger.info("Loaded %d OTP sessions.", len(combined))
|
||||||
|
|
||||||
|
async def _fetch_activities_for_otpid(self, api, otpid, hostname=None) -> None:
|
||||||
|
"""
|
||||||
|
Fetch activities DataFrame for a given otpid and populate activities_table.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
result = api.otp_get_activities(otpid)
|
||||||
|
result_df = (
|
||||||
|
result if isinstance(result, pd.DataFrame) else pd.DataFrame(result)
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception("Failed to fetch activities for otpid %s: %s", otpid, exc)
|
||||||
|
result_df = pd.DataFrame()
|
||||||
|
# Attach hostname if provided
|
||||||
|
if hostname is not None:
|
||||||
|
result_df["hostname"] = hostname
|
||||||
|
if result_df.empty:
|
||||||
|
logger.info("No activities found for otpid %s (host: %s)", otpid, hostname)
|
||||||
|
self._activities_df = pd.DataFrame()
|
||||||
|
self.activities_table.clear()
|
||||||
|
return
|
||||||
|
# Store and render
|
||||||
|
self._activities_df = result_df.copy()
|
||||||
|
# Rebuild activities_table columns from result_df
|
||||||
|
self.activities_table.clear()
|
||||||
|
# Ensure stable column order
|
||||||
|
for col in result_df.columns:
|
||||||
|
self.activities_table.add_column(col)
|
||||||
|
# Add rows
|
||||||
|
for _, arow in result_df.iterrows():
|
||||||
|
values = ["" if pd.isna(v) else v for v in arow.to_list()]
|
||||||
|
self.activities_table.add_row(*[str(v) for v in values])
|
||||||
|
logger.info(
|
||||||
|
"Loaded %d activity rows for otpid %s (host: %s)",
|
||||||
|
len(result_df),
|
||||||
|
otpid,
|
||||||
|
hostname,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def export_activities(self) -> None:
|
||||||
|
"""
|
||||||
|
Export currently-loaded activities DataFrame to CSV.
|
||||||
|
Can be called directly (programmatically) or from the button handler.
|
||||||
|
"""
|
||||||
|
if self._activities_df is None or self._activities_df.empty:
|
||||||
|
logger.info("No activities loaded to export.")
|
||||||
|
# On-screen short message
|
||||||
|
await self.post_message(Static("No activities to export."))
|
||||||
|
return
|
||||||
|
working_dir = _load_working_dir()
|
||||||
|
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||||
|
filename = f"otp_activities_{self._selected_session_otpid}_{timestamp}.csv"
|
||||||
|
file_path = os.path.join(working_dir, filename)
|
||||||
|
try:
|
||||||
|
self._activities_df.to_csv(file_path, index=False)
|
||||||
|
logger.info("Exported activities to %s", file_path)
|
||||||
|
await self.post_message(Static(f"Exported activities to: {file_path}"))
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception("Failed to export activities to %s: %s", file_path, exc)
|
||||||
|
await self.post_message(Static("Failed to export activities; check logs."))
|
||||||
|
|
||||||
|
|
||||||
|
class ActivityDetailWidget(Static):
|
||||||
|
"""
|
||||||
|
Interactive widget for Activity Detail screen.
|
||||||
|
Shows the provided DataFrame in a DataTable and offers Export button.
|
||||||
|
Now includes Select All/None and Add to Allowlist functionality.
|
||||||
|
"""
|
||||||
|
|
||||||
|
DEFAULT_CSS = """
|
||||||
|
ActivityDetailWidget {
|
||||||
|
height: 1fr;
|
||||||
|
layout: vertical;
|
||||||
|
}
|
||||||
|
#detail_table_container {
|
||||||
|
height: 1fr;
|
||||||
|
padding: 1 1;
|
||||||
|
}
|
||||||
|
#selection_buttons {
|
||||||
|
height: auto;
|
||||||
|
min-height: 3;
|
||||||
|
padding: 1 1;
|
||||||
|
content-align: center middle;
|
||||||
|
}
|
||||||
|
#detail_buttons {
|
||||||
|
height: auto;
|
||||||
|
min-height: 3;
|
||||||
|
padding: 1 1;
|
||||||
|
content-align: center middle;
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, activities_df: pd.DataFrame, otpid=None, hostname=None) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.activities_df = (
|
||||||
|
activities_df.copy()
|
||||||
|
if isinstance(activities_df, pd.DataFrame)
|
||||||
|
else pd.DataFrame(activities_df)
|
||||||
|
)
|
||||||
|
# Add a unique identifier column if not present
|
||||||
|
if "_row_id" not in self.activities_df.columns:
|
||||||
|
self.activities_df["_row_id"] = range(len(self.activities_df))
|
||||||
|
|
||||||
|
self.otpid = otpid
|
||||||
|
self.hostname = hostname
|
||||||
|
self.selected_row_ids = set() # Track selected rows by unique ID
|
||||||
|
self.row_key_to_id = {} # Map DataTable row keys to unique row IDs
|
||||||
|
self.table_row_to_id = {} # Map table row indices to unique row IDs
|
||||||
|
self._last_sort = None # Track last sort column and order
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
yield Static(
|
||||||
|
f"Activity Detail (otpid={self.otpid} host={self.hostname})",
|
||||||
|
classes="panel-title",
|
||||||
|
)
|
||||||
|
# Table container
|
||||||
|
with Vertical(id="detail_table_container"):
|
||||||
|
self.detail_table = DataTable(id="detail_table")
|
||||||
|
yield self.detail_table
|
||||||
|
|
||||||
|
# Original buttons at bottom
|
||||||
|
with Horizontal(id="detail_buttons"):
|
||||||
|
self.add_allowlist_btn = Button(
|
||||||
|
"Add Selected to Allowlist", id="add_allowlist_btn"
|
||||||
|
)
|
||||||
|
yield self.add_allowlist_btn
|
||||||
|
|
||||||
|
async def on_mount(self) -> None:
|
||||||
|
await self._build_table(rebuild=True)
|
||||||
|
self._update_button_states()
|
||||||
|
|
||||||
|
def _update_button_states(self) -> None:
|
||||||
|
"""Update button states based on selection."""
|
||||||
|
has_selection = len(self.selected_row_ids) > 0
|
||||||
|
self.add_allowlist_btn.disabled = not has_selection
|
||||||
|
|
||||||
|
# Update button labels with count
|
||||||
|
count = len(self.selected_row_ids)
|
||||||
|
len(self.activities_df)
|
||||||
|
|
||||||
|
if has_selection:
|
||||||
|
self.add_allowlist_btn.label = f"Add {count} Selected to Allowlist"
|
||||||
|
else:
|
||||||
|
self.add_allowlist_btn.label = "Add Selected to Allowlist"
|
||||||
|
|
||||||
|
async def _build_table(self, rebuild: bool = True) -> None:
|
||||||
|
"""Rebuild the DataTable. If rebuild=False, only refresh rows."""
|
||||||
|
if rebuild:
|
||||||
|
# Full rebuild: clear columns and rows
|
||||||
|
self.detail_table.clear()
|
||||||
|
self.detail_table.columns.clear()
|
||||||
|
self.row_key_to_id.clear()
|
||||||
|
self.table_row_to_id.clear()
|
||||||
|
|
||||||
|
if self.activities_df is None or self.activities_df.empty:
|
||||||
|
logger.info("ActivityDetailWidget mounted with empty dataframe.")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Add columns (checkbox + data columns, excluding internal _row_id)
|
||||||
|
self.detail_table.add_column("Select", key="select")
|
||||||
|
for col in self.activities_df.columns:
|
||||||
|
if col != "_row_id": # Don't display the internal ID column
|
||||||
|
self.detail_table.add_column(col)
|
||||||
|
else:
|
||||||
|
# Partial rebuild: clear rows only
|
||||||
|
self.detail_table.clear()
|
||||||
|
self.row_key_to_id.clear()
|
||||||
|
self.table_row_to_id.clear()
|
||||||
|
|
||||||
|
# Add rows
|
||||||
|
for table_idx, (df_idx, row) in enumerate(self.activities_df.iterrows()):
|
||||||
|
# Get the unique row ID
|
||||||
|
row_id = row["_row_id"]
|
||||||
|
|
||||||
|
# Build values list (excluding _row_id column)
|
||||||
|
vals = []
|
||||||
|
for col in self.activities_df.columns:
|
||||||
|
if col != "_row_id":
|
||||||
|
v = row[col]
|
||||||
|
vals.append("" if pd.isna(v) else str(v))
|
||||||
|
|
||||||
|
# Check if this row is selected
|
||||||
|
checkbox = "☑️" if row_id in self.selected_row_ids else "☐"
|
||||||
|
|
||||||
|
# Add row to table
|
||||||
|
row_key = self.detail_table.add_row(checkbox, *vals)
|
||||||
|
|
||||||
|
# Map the row key and table index to the unique row ID
|
||||||
|
self.row_key_to_id[row_key] = row_id
|
||||||
|
self.table_row_to_id[table_idx] = row_id
|
||||||
|
|
||||||
|
async def on_data_table_cell_selected(self, event: DataTable.CellSelected) -> None:
|
||||||
|
# Toggle selection when the "Select" column is clicked
|
||||||
|
if event.cell_key.column_key.value == "select":
|
||||||
|
table_row_index = event.coordinate.row
|
||||||
|
|
||||||
|
# Get the unique row ID for this table row
|
||||||
|
row_id = self.table_row_to_id.get(table_row_index)
|
||||||
|
if row_id is not None:
|
||||||
|
# Get the row key for updating the cell
|
||||||
|
row_key = event.cell_key.row_key
|
||||||
|
|
||||||
|
if row_id in self.selected_row_ids:
|
||||||
|
self.selected_row_ids.remove(row_id)
|
||||||
|
self.detail_table.update_cell(row_key, "select", "☐") # Unchecked
|
||||||
|
else:
|
||||||
|
self.selected_row_ids.add(row_id)
|
||||||
|
self.detail_table.update_cell(row_key, "select", "☑️") # Checked
|
||||||
|
|
||||||
|
self._update_button_states()
|
||||||
|
|
||||||
|
async def on_data_table_header_selected(
|
||||||
|
self, event: DataTable.HeaderSelected
|
||||||
|
) -> None:
|
||||||
|
column_key = event.column_key.value if event.column_key else None
|
||||||
|
if not column_key:
|
||||||
|
col_index = event.column_index
|
||||||
|
if col_index == 0: # First column is "Select"
|
||||||
|
return
|
||||||
|
# Adjust for hidden _row_id column
|
||||||
|
visible_cols = [
|
||||||
|
col for col in self.activities_df.columns if col != "_row_id"
|
||||||
|
]
|
||||||
|
if col_index - 1 < len(visible_cols):
|
||||||
|
column_key = visible_cols[col_index - 1]
|
||||||
|
else:
|
||||||
|
return
|
||||||
|
if column_key == "select" or column_key == "_row_id":
|
||||||
|
return
|
||||||
|
|
||||||
|
ascending = True
|
||||||
|
if self._last_sort == (column_key, True):
|
||||||
|
ascending = False
|
||||||
|
self._last_sort = (column_key, ascending)
|
||||||
|
|
||||||
|
try:
|
||||||
|
self.activities_df.sort_values(
|
||||||
|
by=column_key, ascending=ascending, inplace=True
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception("Failed to sort by column %s: %s", column_key, exc)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Only refresh rows, not columns
|
||||||
|
await self._build_table(rebuild=False)
|
||||||
|
|
||||||
|
async def on_button_pressed(self, event) -> None:
|
||||||
|
btn = getattr(event, "button", None) or getattr(event, "sender", None)
|
||||||
|
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None)
|
||||||
|
|
||||||
|
if btn is self.add_allowlist_btn or btn_id == "add_allowlist_btn":
|
||||||
|
await self._open_allowlist_screen()
|
||||||
|
return
|
||||||
|
|
||||||
|
async def _select_all(self) -> None:
|
||||||
|
"""Select all rows in the table."""
|
||||||
|
# Add all row IDs to selected set
|
||||||
|
self.selected_row_ids = set(self.activities_df["_row_id"].tolist())
|
||||||
|
|
||||||
|
# Update all checkboxes in the table
|
||||||
|
for row_key, row_id in self.row_key_to_id.items():
|
||||||
|
self.detail_table.update_cell(row_key, "select", "☑️")
|
||||||
|
|
||||||
|
self._update_button_states()
|
||||||
|
logger.info(f"Selected all {len(self.selected_row_ids)} rows")
|
||||||
|
|
||||||
|
async def _select_none(self) -> None:
|
||||||
|
"""Deselect all rows in the table."""
|
||||||
|
# Clear selected set
|
||||||
|
self.selected_row_ids.clear()
|
||||||
|
|
||||||
|
# Update all checkboxes in the table
|
||||||
|
for row_key, row_id in self.row_key_to_id.items():
|
||||||
|
self.detail_table.update_cell(row_key, "select", "☑️")
|
||||||
|
|
||||||
|
self._update_button_states()
|
||||||
|
logger.info("Cleared all selections")
|
||||||
|
|
||||||
|
async def _open_allowlist_screen(self) -> None:
|
||||||
|
"""Open the allowlist selection screen with selected activities."""
|
||||||
|
if not self.selected_row_ids:
|
||||||
|
self.app.notify("No rows selected", severity="warning")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Get selected data
|
||||||
|
selected_df = self.get_selected_data()
|
||||||
|
|
||||||
|
# Get API from app
|
||||||
|
api = getattr(self.app, "api", None)
|
||||||
|
if api is None:
|
||||||
|
logger.error("No API available on self.app.api")
|
||||||
|
self.app.notify("API not available", severity="error")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Create and push AllowlistSelectionScreen
|
||||||
|
try:
|
||||||
|
allowlist_screen = AllowlistSelectionScreen(
|
||||||
|
selected_df, api=api, hostname=self.hostname, otpid=self.otpid
|
||||||
|
)
|
||||||
|
await self.app.push_screen(allowlist_screen)
|
||||||
|
logger.info(
|
||||||
|
f"Opened allowlist screen with {len(selected_df)} selected activities"
|
||||||
|
)
|
||||||
|
except ImportError as e:
|
||||||
|
logger.error(f"Failed to import AllowlistSelectionScreen: {e}")
|
||||||
|
self.app.notify("Allowlist screen module not found", severity="error")
|
||||||
|
except Exception as e:
|
||||||
|
logger.exception(f"Failed to open allowlist screen: {e}")
|
||||||
|
self.app.notify(
|
||||||
|
f"Error opening allowlist screen: {str(e)}", severity="error"
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _export_detail_activities(self) -> None:
|
||||||
|
if self.activities_df is None or self.activities_df.empty:
|
||||||
|
logger.info("No activities to export.")
|
||||||
|
await self.mount(Static("No activities to export.", classes="notification"))
|
||||||
|
return
|
||||||
|
if not self.selected_row_ids:
|
||||||
|
logger.info("No rows selected for export.")
|
||||||
|
await self.mount(
|
||||||
|
Static("No rows selected for export.", classes="notification")
|
||||||
|
)
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
working_dir = load_env("WORKING_DIR") or os.getcwd()
|
||||||
|
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||||
|
filename = f"otp_activities_detail_{timestamp}.csv"
|
||||||
|
file_path = os.path.join(working_dir, filename)
|
||||||
|
selected_df = self.get_selected_data()
|
||||||
|
selected_df.to_csv(file_path, index=False)
|
||||||
|
logger.info("Exported selected activities to %s", file_path)
|
||||||
|
await self.mount(
|
||||||
|
Static(
|
||||||
|
f"Exported selected activities to: {filename}",
|
||||||
|
classes="notification",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception("Failed to export detail activities: %s", exc)
|
||||||
|
await self.mount(
|
||||||
|
Static(
|
||||||
|
"¢ Failed to export activities; check logs.",
|
||||||
|
classes="notification",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Helper methods
|
||||||
|
def get_selected_data(self) -> pd.DataFrame:
|
||||||
|
"""Return a DataFrame of the selected rows."""
|
||||||
|
if not self.selected_row_ids:
|
||||||
|
return pd.DataFrame()
|
||||||
|
# Filter by selected row IDs and drop the internal _row_id column
|
||||||
|
selected_df = self.activities_df[
|
||||||
|
self.activities_df["_row_id"].isin(self.selected_row_ids)
|
||||||
|
].copy()
|
||||||
|
if "_row_id" in selected_df.columns:
|
||||||
|
selected_df = selected_df.drop(columns=["_row_id"])
|
||||||
|
return selected_df
|
||||||
|
|
||||||
|
def get_selected_records(self) -> list[dict]:
|
||||||
|
"""Return selected rows as a list of dicts."""
|
||||||
|
if not self.selected_row_ids:
|
||||||
|
return []
|
||||||
|
# Filter by selected row IDs and drop the internal _row_id column
|
||||||
|
selected_df = self.activities_df[
|
||||||
|
self.activities_df["_row_id"].isin(self.selected_row_ids)
|
||||||
|
].copy()
|
||||||
|
if "_row_id" in selected_df.columns:
|
||||||
|
selected_df = selected_df.drop(columns=["_row_id"])
|
||||||
|
return selected_df.to_dict(orient="records")
|
||||||
|
|
||||||
|
|
||||||
|
class ActivityDetailScreen(Screen):
|
||||||
|
"""
|
||||||
|
Screen that wraps ActivityDetailWidget. Expects a DataFrame passed on init.
|
||||||
|
"""
|
||||||
|
|
||||||
|
BINDINGS = [
|
||||||
|
Binding("escape", "go_back", "Back"),
|
||||||
|
Binding("q", "main_menu", "Main Menu"),
|
||||||
|
Binding("e", "export", "Export"),
|
||||||
|
Binding("a", "select_all", "Select All"),
|
||||||
|
Binding("n", "select_none", "Select None"),
|
||||||
|
]
|
||||||
|
|
||||||
|
def __init__(self, activities_df: pd.DataFrame, otpid=None, hostname=None) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self._activities_df = (
|
||||||
|
activities_df.copy()
|
||||||
|
if isinstance(activities_df, pd.DataFrame)
|
||||||
|
else pd.DataFrame(activities_df)
|
||||||
|
)
|
||||||
|
self._otpid = otpid
|
||||||
|
self._hostname = hostname
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
self.widget = ActivityDetailWidget(
|
||||||
|
self._activities_df, otpid=self._otpid, hostname=self._hostname
|
||||||
|
)
|
||||||
|
yield Header(show_clock=True)
|
||||||
|
yield self.widget
|
||||||
|
yield Footer()
|
||||||
|
|
||||||
|
async def action_go_back(self) -> None:
|
||||||
|
try:
|
||||||
|
await self.app.pop_screen()
|
||||||
|
except Exception:
|
||||||
|
logger.debug("ActivityDetailScreen.action_go_back pop_screen failed.")
|
||||||
|
|
||||||
|
async def action_main_menu(self) -> None:
|
||||||
|
"""Go back to main menu."""
|
||||||
|
while len(self.app.screen_stack) > 2:
|
||||||
|
await self.app.pop_screen()
|
||||||
|
|
||||||
|
async def action_export(self) -> None:
|
||||||
|
# Delegate to widget export helper
|
||||||
|
if hasattr(self, "widget") and self.widget is not None:
|
||||||
|
await self.widget._export_detail_activities()
|
||||||
|
|
||||||
|
async def action_select_all(self) -> None:
|
||||||
|
"""Handle 'a' key for select all."""
|
||||||
|
if hasattr(self, "widget") and self.widget is not None:
|
||||||
|
await self.widget._select_all()
|
||||||
|
|
||||||
|
async def action_select_none(self) -> None:
|
||||||
|
"""Handle 'n' key for select none."""
|
||||||
|
if hasattr(self, "widget") and self.widget is not None:
|
||||||
|
await self.widget._select_none()
|
||||||
|
|
||||||
|
|
||||||
|
class OTPActivitiesScreen(Screen):
|
||||||
|
"""
|
||||||
|
A Screen intended to be pushed into an existing Textual App.
|
||||||
|
Usage:
|
||||||
|
app.push_screen(OTPActivitiesScreen())
|
||||||
|
or create this screen and call `await screen.load()` inside your app lifecycle.
|
||||||
|
The screen expects `self.app.api` to exist and be an AirlockAPIWrapper instance.
|
||||||
|
"""
|
||||||
|
|
||||||
|
BINDINGS = [
|
||||||
|
Binding("escape", "go_back", "Back"),
|
||||||
|
Binding("q", "main_menu", "Main Menu"),
|
||||||
|
Binding("r", "refresh_sessions", "Refresh Sessions"),
|
||||||
|
Binding("e", "export_activities", "Export activities"),
|
||||||
|
]
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
yield Header()
|
||||||
|
self.widget = OTPActivitiesWidget()
|
||||||
|
yield self.widget
|
||||||
|
yield Footer()
|
||||||
|
|
||||||
|
async def on_show(self) -> None:
|
||||||
|
"""Restore focus to the left sessions table when the screen becomes visible."""
|
||||||
|
if hasattr(self, "widget") and hasattr(self.widget, "sessions_table"):
|
||||||
|
self.widget.sessions_table.focus()
|
||||||
|
|
||||||
|
async def on_mount(self) -> None:
|
||||||
|
# Try to load sessions immediately
|
||||||
|
api = getattr(self.app, "api", None)
|
||||||
|
if api is None:
|
||||||
|
logger.warning("OTPActivitiesScreen mounted but no self.app.api found.")
|
||||||
|
else:
|
||||||
|
await self.widget.load_sessions_from_api(api)
|
||||||
|
|
||||||
|
async def action_go_back(self) -> None:
|
||||||
|
"""Go back one screen."""
|
||||||
|
await self.app.pop_screen()
|
||||||
|
|
||||||
|
async def action_main_menu(self) -> None:
|
||||||
|
"""Go back to main menu."""
|
||||||
|
while len(self.app.screen_stack) > 2:
|
||||||
|
await self.app.pop_screen()
|
||||||
|
|
||||||
|
# Simple actions bound to keys
|
||||||
|
async def action_refresh_sessions(self) -> None:
|
||||||
|
api = getattr(self.app, "api", None)
|
||||||
|
if api is None:
|
||||||
|
logger.error("No API on app; cannot refresh sessions.")
|
||||||
|
return
|
||||||
|
logger.info("Refreshing OTP sessions via API.")
|
||||||
|
await self.widget.load_sessions_from_api(api)
|
||||||
|
|
||||||
|
# If you want an explicit method to fetch activities for a particular otpid from outside:
|
||||||
|
async def fetch_activities_for_otpid(self, otpid, hostname=None) -> None:
|
||||||
|
api = getattr(self.app, "api", None)
|
||||||
|
if api is None:
|
||||||
|
logger.error("No API on app; cannot fetch activities.")
|
||||||
|
return
|
||||||
|
await self.widget._fetch_activities_for_otpid(api, otpid, hostname=hostname)
|
||||||
@@ -0,0 +1,392 @@
|
|||||||
|
# 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 __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from typing import List, Optional
|
||||||
|
|
||||||
|
import pandas as pd
|
||||||
|
from textual.app import ComposeResult
|
||||||
|
from textual.binding import Binding
|
||||||
|
from textual.containers import Horizontal, Vertical
|
||||||
|
from textual.message import Message
|
||||||
|
from textual.screen import Screen
|
||||||
|
from textual.widgets import Button, DataTable, Footer, Header, Static
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class OTPRevokeWidget(Static):
|
||||||
|
"""
|
||||||
|
Widget for managing OTP session revocation.
|
||||||
|
Displays active OTP sessions and allows selection for revocation.
|
||||||
|
"""
|
||||||
|
|
||||||
|
class SessionsRevoked(Message):
|
||||||
|
"""Message sent when sessions are revoked."""
|
||||||
|
|
||||||
|
def __init__(self, revoked_sessions: List[dict]):
|
||||||
|
super().__init__()
|
||||||
|
self.revoked_sessions = revoked_sessions
|
||||||
|
|
||||||
|
DEFAULT_CSS = """
|
||||||
|
OTPRevokeWidget {
|
||||||
|
height: 1fr;
|
||||||
|
}
|
||||||
|
#main_container {
|
||||||
|
width: 100%;
|
||||||
|
height: 100%;
|
||||||
|
layout: vertical;
|
||||||
|
}
|
||||||
|
#sessions_container {
|
||||||
|
height: 1fr;
|
||||||
|
border: none;
|
||||||
|
padding: 1;
|
||||||
|
}
|
||||||
|
#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;
|
||||||
|
border: solid #444444;
|
||||||
|
padding: 1;
|
||||||
|
margin: 1;
|
||||||
|
overflow-y: auto;
|
||||||
|
}
|
||||||
|
.panel-title {
|
||||||
|
text-style: bold;
|
||||||
|
margin: 0 0 1 0;
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
with Vertical(id="main_container"):
|
||||||
|
# Sessions table
|
||||||
|
yield Static("OTP Sessions", classes="panel-title")
|
||||||
|
with Vertical(id="sessions_container"):
|
||||||
|
self.sessions_table = DataTable(id="sessions_table")
|
||||||
|
self.sessions_table.styles.width = "100%"
|
||||||
|
self.sessions_table.styles.height = "1fr"
|
||||||
|
yield self.sessions_table
|
||||||
|
|
||||||
|
# Action buttons
|
||||||
|
with Horizontal(id="button_container"):
|
||||||
|
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"):
|
||||||
|
yield Static("Revocation Results", classes="panel-title")
|
||||||
|
self.results_display = Static("No actions performed yet.")
|
||||||
|
yield self.results_display
|
||||||
|
|
||||||
|
async def on_mount(self) -> None:
|
||||||
|
"""Initialize the widget when mounted."""
|
||||||
|
# Configure sessions table
|
||||||
|
self.sessions_table.clear()
|
||||||
|
self.sessions_table.add_columns(
|
||||||
|
"", "OTP ID", "Hostname", "Status", "Purpose", "Granted"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Enable row selection with checkbox column
|
||||||
|
self.sessions_table.cursor_type = "row"
|
||||||
|
try:
|
||||||
|
self.sessions_table.zebra_stripes = True
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Initialize state
|
||||||
|
self._sessions_df: Optional[pd.DataFrame] = None
|
||||||
|
self._filtered_df: Optional[pd.DataFrame] = None
|
||||||
|
self._selected_otpids: set = set()
|
||||||
|
|
||||||
|
async def load_sessions_from_api(self, api) -> None:
|
||||||
|
"""Load active OTP sessions from the API."""
|
||||||
|
try:
|
||||||
|
# Fetch only active sessions
|
||||||
|
active_df = api.otp_find_active()
|
||||||
|
|
||||||
|
# Ensure we have a DataFrame
|
||||||
|
if not isinstance(active_df, pd.DataFrame):
|
||||||
|
active_df = pd.DataFrame(active_df)
|
||||||
|
|
||||||
|
# Add status column
|
||||||
|
active_df["status"] = "active"
|
||||||
|
|
||||||
|
# Sort by otpid if column exists
|
||||||
|
if "otpid" in active_df.columns and not active_df.empty:
|
||||||
|
active_df = active_df.sort_values(by="otpid", ascending=False)
|
||||||
|
|
||||||
|
# Store the full dataframe
|
||||||
|
self._sessions_df = active_df
|
||||||
|
self._filtered_df = active_df.copy()
|
||||||
|
|
||||||
|
# Display in table
|
||||||
|
await self._refresh_table()
|
||||||
|
|
||||||
|
# Update status
|
||||||
|
active_count = len(active_df)
|
||||||
|
|
||||||
|
status_msg = f"Loaded {active_count} active sessions"
|
||||||
|
logger.info(status_msg)
|
||||||
|
self.results_display.update(status_msg)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.exception(f"Failed to load OTP sessions: {e}")
|
||||||
|
self.results_display.update(f"Error loading sessions: {str(e)}")
|
||||||
|
|
||||||
|
async def _refresh_table(self) -> None:
|
||||||
|
"""Refresh the table display with current filtered data."""
|
||||||
|
if self._filtered_df is None or self._filtered_df.empty:
|
||||||
|
self.sessions_table.clear()
|
||||||
|
return
|
||||||
|
|
||||||
|
# Ensure expected columns exist
|
||||||
|
expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]
|
||||||
|
for col in expected_cols:
|
||||||
|
if col not in self._filtered_df.columns:
|
||||||
|
self._filtered_df[col] = ""
|
||||||
|
|
||||||
|
# Clear and repopulate table
|
||||||
|
self.sessions_table.clear(columns=False)
|
||||||
|
|
||||||
|
for _, row in self._filtered_df.iterrows():
|
||||||
|
otpid = str(row.get("otpid", ""))
|
||||||
|
# Check if this row is selected
|
||||||
|
checkbox = "☑️" if otpid in self._selected_otpids else "☐"
|
||||||
|
|
||||||
|
self.sessions_table.add_row(
|
||||||
|
checkbox,
|
||||||
|
str(otpid),
|
||||||
|
str(row.get("hostname", "")),
|
||||||
|
str(row.get("status", "")),
|
||||||
|
str(row.get("purpose", "")),
|
||||||
|
str(row.get("granted", "")),
|
||||||
|
)
|
||||||
|
|
||||||
|
async def on_button_pressed(self, event) -> None:
|
||||||
|
"""Handle button presses."""
|
||||||
|
btn = event.button
|
||||||
|
|
||||||
|
if btn.id == "refresh_btn":
|
||||||
|
# Refresh sessions
|
||||||
|
api = getattr(self.app, "api", None)
|
||||||
|
if api:
|
||||||
|
await self.load_sessions_from_api(api)
|
||||||
|
|
||||||
|
elif btn.id == "select_all_btn":
|
||||||
|
# Select all visible rows
|
||||||
|
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()
|
||||||
|
|
||||||
|
elif btn.id == "select_none_btn":
|
||||||
|
# Clear selection
|
||||||
|
self._selected_otpids.clear()
|
||||||
|
await self._refresh_table()
|
||||||
|
|
||||||
|
elif btn.id == "revoke_btn":
|
||||||
|
# Revoke selected sessions
|
||||||
|
await self._revoke_selected()
|
||||||
|
|
||||||
|
async def on_data_table_row_selected(self, event) -> None:
|
||||||
|
"""Handle row selection in the table."""
|
||||||
|
if event.data_table != self.sessions_table:
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Get the row index from the cursor row
|
||||||
|
row_index = self.sessions_table.cursor_row
|
||||||
|
|
||||||
|
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"])
|
||||||
|
|
||||||
|
# Toggle selection
|
||||||
|
if otpid in self._selected_otpids:
|
||||||
|
self._selected_otpids.remove(otpid)
|
||||||
|
else:
|
||||||
|
self._selected_otpids.add(otpid)
|
||||||
|
|
||||||
|
# Refresh table to update checkbox
|
||||||
|
await self._refresh_table()
|
||||||
|
|
||||||
|
# Restore cursor position
|
||||||
|
self.sessions_table.move_cursor(row=row_index)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.exception(f"Error handling row selection: {e}")
|
||||||
|
|
||||||
|
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")
|
||||||
|
return
|
||||||
|
|
||||||
|
api = getattr(self.app, "api", None)
|
||||||
|
if not api:
|
||||||
|
self.results_display.update("API not available")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Collect results
|
||||||
|
results = []
|
||||||
|
success_count = 0
|
||||||
|
failure_count = 0
|
||||||
|
|
||||||
|
for otpid in self._selected_otpids:
|
||||||
|
try:
|
||||||
|
# Get hostname for this session
|
||||||
|
hostname = "Unknown"
|
||||||
|
if self._sessions_df is not None:
|
||||||
|
# Convert otpid to same type as in DataFrame for comparison
|
||||||
|
otpid_compare = otpid
|
||||||
|
if len(self._sessions_df) > 0:
|
||||||
|
first_otpid = self._sessions_df["otpid"].iloc[0]
|
||||||
|
if isinstance(first_otpid, int):
|
||||||
|
otpid_compare = int(otpid)
|
||||||
|
|
||||||
|
match = self._sessions_df[
|
||||||
|
self._sessions_df["otpid"] == otpid_compare
|
||||||
|
]
|
||||||
|
if not match.empty:
|
||||||
|
hostname = match.iloc[0].get("hostname", "Unknown")
|
||||||
|
|
||||||
|
# Revoke the session
|
||||||
|
result = api.otp_revoke(otpid)
|
||||||
|
|
||||||
|
if result and result.get("status") != "error":
|
||||||
|
success_count += 1
|
||||||
|
results.append(f"Revoked OTP {otpid} for {hostname}")
|
||||||
|
logger.info(f"Revoked OTP {otpid} for {hostname}: {result}")
|
||||||
|
else:
|
||||||
|
failure_count += 1
|
||||||
|
error_msg = (
|
||||||
|
result.get("message", "Unknown error")
|
||||||
|
if result
|
||||||
|
else "No response"
|
||||||
|
)
|
||||||
|
results.append(
|
||||||
|
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)}")
|
||||||
|
logger.exception(f"Exception revoking OTP {otpid}: {e}")
|
||||||
|
|
||||||
|
# Update results display
|
||||||
|
summary = (
|
||||||
|
f"Revocation complete: {success_count} succeeded, {failure_count} failed\n"
|
||||||
|
)
|
||||||
|
details = "\n".join(results[-5:]) # Show last 5 results
|
||||||
|
if len(results) > 5:
|
||||||
|
details = f"... (showing last 5 of {len(results)} results)\n" + details
|
||||||
|
|
||||||
|
self.results_display.update(summary + details)
|
||||||
|
|
||||||
|
# Clear selection and refresh
|
||||||
|
self._selected_otpids.clear()
|
||||||
|
await self.load_sessions_from_api(api)
|
||||||
|
|
||||||
|
# Post message about revoked sessions
|
||||||
|
if success_count > 0:
|
||||||
|
self.post_message(self.SessionsRevoked(results))
|
||||||
|
|
||||||
|
|
||||||
|
class OTPRevokeScreen(Screen):
|
||||||
|
"""
|
||||||
|
Main screen for OTP session revocation workflow.
|
||||||
|
This replaces the otp_revoke function from otp.py.
|
||||||
|
"""
|
||||||
|
|
||||||
|
BINDINGS = [
|
||||||
|
Binding("escape", "go_back", "Back"),
|
||||||
|
Binding("q", "main_menu", "Main Menu"),
|
||||||
|
Binding("r", "refresh", "Refresh"),
|
||||||
|
Binding("a", "select_all", "Select All"),
|
||||||
|
Binding("n", "select_none", "Clear Selection"),
|
||||||
|
Binding("d", "revoke", "Revoke Selected"),
|
||||||
|
]
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
yield Header(show_clock=True)
|
||||||
|
self.widget = OTPRevokeWidget()
|
||||||
|
yield self.widget
|
||||||
|
yield Footer()
|
||||||
|
|
||||||
|
async def on_mount(self) -> None:
|
||||||
|
"""Load sessions when screen mounts."""
|
||||||
|
api = getattr(self.app, "api", None)
|
||||||
|
if api:
|
||||||
|
await self.widget.load_sessions_from_api(api)
|
||||||
|
else:
|
||||||
|
logger.warning("OTPRevokeScreen mounted but no self.app.api found.")
|
||||||
|
|
||||||
|
async def action_refresh(self) -> None:
|
||||||
|
"""Refresh the sessions list."""
|
||||||
|
api = getattr(self.app, "api", None)
|
||||||
|
if api:
|
||||||
|
await self.widget.load_sessions_from_api(api)
|
||||||
|
|
||||||
|
async def action_select_all(self) -> None:
|
||||||
|
"""Select all visible sessions."""
|
||||||
|
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"]
|
||||||
|
)
|
||||||
|
await self.widget._refresh_table()
|
||||||
|
|
||||||
|
async def action_select_none(self) -> None:
|
||||||
|
"""Clear all selections."""
|
||||||
|
self.widget._selected_otpids.clear()
|
||||||
|
await self.widget._refresh_table()
|
||||||
|
|
||||||
|
async def action_revoke(self) -> None:
|
||||||
|
"""Revoke selected sessions."""
|
||||||
|
await self.widget._revoke_selected()
|
||||||
|
|
||||||
|
async def action_go_back(self) -> None:
|
||||||
|
"""Go back to previous screen."""
|
||||||
|
await self.app.pop_screen()
|
||||||
|
|
||||||
|
async def action_main_menu(self) -> None:
|
||||||
|
"""Go back to main menu."""
|
||||||
|
while len(self.app.screen_stack) > 2:
|
||||||
|
await self.app.pop_screen()
|
||||||
@@ -0,0 +1,52 @@
|
|||||||
|
# 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 typing import List, Optional
|
||||||
|
|
||||||
|
from textual.app import ComposeResult
|
||||||
|
from textual.binding import Binding
|
||||||
|
from textual.screen import Screen
|
||||||
|
|
||||||
|
from models.agent import Agent
|
||||||
|
from TUI.Widgets.OTP_generate import OTPGenerator
|
||||||
|
|
||||||
|
|
||||||
|
class OTPWorkflowScreen(Screen):
|
||||||
|
"""Screen that handles the OTP generation workflow without agent selection."""
|
||||||
|
|
||||||
|
BINDINGS = [
|
||||||
|
Binding("escape", "go_back", "Back"),
|
||||||
|
Binding("q", "main_menu", "Main Menu"),
|
||||||
|
]
|
||||||
|
|
||||||
|
def __init__(self, selected_agents: Optional[List[Agent]]):
|
||||||
|
super().__init__()
|
||||||
|
self.selected_agents = selected_agents
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
"""Directly show the OTP generator for the selected agents."""
|
||||||
|
yield OTPGenerator(self.selected_agents)
|
||||||
|
|
||||||
|
def action_go_back(self) -> None:
|
||||||
|
"""Handle escape key to go back one screen."""
|
||||||
|
self.app.pop_screen()
|
||||||
|
|
||||||
|
def action_main_menu(self) -> None:
|
||||||
|
"""Handle q key to go back to main menu."""
|
||||||
|
while len(self.app.screen_stack) > 2:
|
||||||
|
self.app.pop_screen()
|
||||||
|
|
||||||
|
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
|
||||||
|
"""Handle OTP generation request - pass it up to the app level if needed."""
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -23,9 +23,11 @@ the policy selection workflow.
|
|||||||
import logging
|
import logging
|
||||||
|
|
||||||
from textual.app import ComposeResult
|
from textual.app import ComposeResult
|
||||||
|
from textual.binding import Binding
|
||||||
from textual.screen import Screen
|
from textual.screen import Screen
|
||||||
|
from textual.widgets import Footer, Header
|
||||||
|
|
||||||
from widgets.policyselector import PolicySelector
|
from TUI.Widgets.policyselector import PolicySelector
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -42,6 +44,11 @@ class PolicySelectorScreen(Screen):
|
|||||||
agent_move_operations: Reference to the parent AgentMoveOperations widget.
|
agent_move_operations: Reference to the parent AgentMoveOperations widget.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
BINDINGS = [
|
||||||
|
Binding("escape", "go_back", "Back"),
|
||||||
|
Binding("q", "main_menu", "Main Menu"),
|
||||||
|
]
|
||||||
|
|
||||||
CSS = """
|
CSS = """
|
||||||
Screen {
|
Screen {
|
||||||
layout: vertical;
|
layout: vertical;
|
||||||
@@ -68,7 +75,18 @@ class PolicySelectorScreen(Screen):
|
|||||||
|
|
||||||
def compose(self) -> ComposeResult:
|
def compose(self) -> ComposeResult:
|
||||||
"""Create the PolicySelector widget."""
|
"""Create the PolicySelector widget."""
|
||||||
|
yield Header(show_clock=True)
|
||||||
yield PolicySelector(self.policies)
|
yield PolicySelector(self.policies)
|
||||||
|
yield Footer()
|
||||||
|
|
||||||
|
def action_go_back(self) -> None:
|
||||||
|
"""Handle escape key to go back one screen."""
|
||||||
|
self.app.pop_screen()
|
||||||
|
|
||||||
|
def action_main_menu(self) -> None:
|
||||||
|
"""Handle q key to go back to main menu."""
|
||||||
|
while len(self.app.screen_stack) > 2:
|
||||||
|
self.app.pop_screen()
|
||||||
|
|
||||||
def on_policy_selector_policy_selected(
|
def on_policy_selector_policy_selected(
|
||||||
self, message: PolicySelector.PolicySelected
|
self, message: PolicySelector.PolicySelected
|
||||||
@@ -0,0 +1,914 @@
|
|||||||
|
# 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/>.
|
||||||
|
|
||||||
|
"""
|
||||||
|
Quiet Agent Workflow Screen Module
|
||||||
|
|
||||||
|
Provides a TUI workflow for identifying quiet agents and moving them to target policies.
|
||||||
|
This screen replaces the legacy quietAgent.py with a comprehensive TUI interface that:
|
||||||
|
1. Allows selection of an initial policy to analyze
|
||||||
|
2. Categorizes devices into "Enforce Ready" and "Non-Enforce Ready" based on activity
|
||||||
|
3. Allows users to select target policies for each category
|
||||||
|
4. Uses the API to move devices to their target policies
|
||||||
|
"""
|
||||||
|
|
||||||
|
import datetime
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
from typing import List, Optional
|
||||||
|
|
||||||
|
import pandas as pd
|
||||||
|
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, Input, Static
|
||||||
|
|
||||||
|
from models.policy import Policy
|
||||||
|
from services.API import AirlockAPIWrapper
|
||||||
|
from services.policyhandler import getPolicyInfo
|
||||||
|
from TUI.Widgets.policyselector import PolicySelector
|
||||||
|
from utils.configmanager import load_env
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class QuietAgentWorkflowScreen(Screen):
|
||||||
|
"""
|
||||||
|
A Textual screen for the Quiet Agent analysis and migration workflow.
|
||||||
|
|
||||||
|
This screen provides a multi-step workflow:
|
||||||
|
1. Select initial policy to analyze
|
||||||
|
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, 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
|
||||||
|
workflow_stage (str): Current stage of the workflow
|
||||||
|
"""
|
||||||
|
|
||||||
|
BINDINGS = [
|
||||||
|
("escape", "go_back", "Back"),
|
||||||
|
("q", "main_menu", "Main Menu"),
|
||||||
|
]
|
||||||
|
|
||||||
|
workflow_stage = reactive("select_policy") # Tracks current workflow stage
|
||||||
|
|
||||||
|
def __init__(self, api: AirlockAPIWrapper, policies: List[Policy]):
|
||||||
|
"""
|
||||||
|
Initialize the QuietAgentWorkflowScreen.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
api (AirlockAPIWrapper): API wrapper for Airlock operations
|
||||||
|
policies (List[Policy]): List of all available policies
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
self.api = api
|
||||||
|
self.policies = policies
|
||||||
|
self.selected_policy: Optional[Policy] = None
|
||||||
|
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
|
||||||
|
self.non_enforce_ready_df: Optional[pd.DataFrame] = None
|
||||||
|
self.enforce_ready_target_policy: Optional[Policy] = None
|
||||||
|
self.non_enforce_ready_target_policy: Optional[Policy] = None
|
||||||
|
|
||||||
|
def compose(self) -> ComposeResult:
|
||||||
|
"""Build the UI layout for the workflow screen."""
|
||||||
|
# Include Header and Footer like other standalone screens
|
||||||
|
yield Header(show_clock=True, icon="⚙️")
|
||||||
|
|
||||||
|
# Title area
|
||||||
|
title = Static("Quiet Agent Workflow", id="workflow_title")
|
||||||
|
title.styles.margin = (0, 0, 0, 1)
|
||||||
|
yield title
|
||||||
|
|
||||||
|
# Status area
|
||||||
|
status = Static("Step 1: Select Policy to Analyze", id="workflow_status")
|
||||||
|
status.styles.margin = (0, 0, 1, 1)
|
||||||
|
yield status
|
||||||
|
|
||||||
|
# Content area - dynamically populated based on workflow stage
|
||||||
|
yield Vertical(id="content_area")
|
||||||
|
|
||||||
|
yield Footer()
|
||||||
|
|
||||||
|
def on_mount(self) -> None:
|
||||||
|
"""Initialize the screen when mounted."""
|
||||||
|
# Show initial policy selection
|
||||||
|
self._show_policy_selection()
|
||||||
|
|
||||||
|
def watch_workflow_stage(self, old_value: str, new_value: str) -> None:
|
||||||
|
"""React to workflow stage changes."""
|
||||||
|
logger.debug(f"Workflow stage changed from {old_value} to {new_value}")
|
||||||
|
self._update_status_message()
|
||||||
|
|
||||||
|
def _update_status_message(self) -> None:
|
||||||
|
"""Update the status message based on current workflow stage."""
|
||||||
|
status_widget = self.query_one("#workflow_status", Static)
|
||||||
|
|
||||||
|
stage_messages = {
|
||||||
|
"select_policy": "Step 1: Select Policy to Analyze",
|
||||||
|
"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",
|
||||||
|
"select_non_enforce_target": "Step 5: Select Target Policy for Non-Enforce Ready Agents",
|
||||||
|
"confirm_migration": "Step 6: Confirm and Execute Migration",
|
||||||
|
"executing": "Executing agent migrations...",
|
||||||
|
"complete": "Migration Complete",
|
||||||
|
}
|
||||||
|
|
||||||
|
status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage"))
|
||||||
|
|
||||||
|
def _show_policy_selection(self) -> None:
|
||||||
|
"""Show the initial policy selection screen."""
|
||||||
|
self.workflow_stage = "select_policy"
|
||||||
|
content = self.query_one("#content_area", Vertical)
|
||||||
|
content.remove_children()
|
||||||
|
|
||||||
|
# Create policy selector widget
|
||||||
|
policy_selector = PolicySelector(self.policies)
|
||||||
|
content.mount(policy_selector)
|
||||||
|
|
||||||
|
def on_policy_selector_policy_selected(
|
||||||
|
self, message: PolicySelector.PolicySelected
|
||||||
|
) -> None:
|
||||||
|
"""Handle policy selection from PolicySelector widget."""
|
||||||
|
# Handle based on current workflow stage
|
||||||
|
if self.workflow_stage == "select_policy":
|
||||||
|
# Initial policy selection for analysis
|
||||||
|
self.selected_policy = message.policy
|
||||||
|
logger.info(f"Selected policy for analysis: {self.selected_policy.name}")
|
||||||
|
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
|
||||||
|
logger.info(
|
||||||
|
f"Selected target policy for enforce ready: {message.policy.name}"
|
||||||
|
)
|
||||||
|
self._show_non_enforce_target_selection()
|
||||||
|
elif self.workflow_stage == "select_non_enforce_target":
|
||||||
|
# Target policy selection for non-enforce ready agents
|
||||||
|
self.non_enforce_ready_target_policy = message.policy
|
||||||
|
logger.info(
|
||||||
|
f"Selected target policy for non-enforce ready: {message.policy.name}"
|
||||||
|
)
|
||||||
|
self._show_migration_confirmation()
|
||||||
|
|
||||||
|
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"
|
||||||
|
"Configure Analysis Parameters:",
|
||||||
|
id="analysis_params_info",
|
||||||
|
)
|
||||||
|
info_widget.styles.margin = (0, 0, 2, 0)
|
||||||
|
content.mount(info_widget)
|
||||||
|
|
||||||
|
# Create input container
|
||||||
|
input_container = Vertical(id="analysis_params_input_container")
|
||||||
|
input_container.styles.height = "auto"
|
||||||
|
content.mount(input_container)
|
||||||
|
|
||||||
|
# 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}"
|
||||||
|
)
|
||||||
|
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
|
||||||
|
|
||||||
|
# Analysis parameters submit button
|
||||||
|
if button_id == "analysis_params_submit":
|
||||||
|
self._validate_and_submit_history_days()
|
||||||
|
return
|
||||||
|
|
||||||
|
# Navigation buttons
|
||||||
|
if button_id == "select_enforce_target_btn":
|
||||||
|
self._show_enforce_target_selection()
|
||||||
|
return
|
||||||
|
|
||||||
|
if button_id == "select_non_enforce_target_btn":
|
||||||
|
self._show_non_enforce_target_selection()
|
||||||
|
return
|
||||||
|
|
||||||
|
if button_id == "skip_enforce_target_btn":
|
||||||
|
# Skip enforce ready target selection
|
||||||
|
self.enforce_ready_target_policy = None
|
||||||
|
self._show_non_enforce_target_selection()
|
||||||
|
return
|
||||||
|
|
||||||
|
if button_id == "skip_non_enforce_target_btn":
|
||||||
|
# Skip non-enforce ready target selection
|
||||||
|
self.non_enforce_ready_target_policy = None
|
||||||
|
self._show_migration_confirmation()
|
||||||
|
return
|
||||||
|
|
||||||
|
if button_id == "confirm_migration_btn":
|
||||||
|
self._execute_migration()
|
||||||
|
return
|
||||||
|
|
||||||
|
if button_id == "cancel_migration_btn":
|
||||||
|
self._show_results()
|
||||||
|
return
|
||||||
|
|
||||||
|
if button_id == "export_results_btn":
|
||||||
|
self._export_results()
|
||||||
|
return
|
||||||
|
|
||||||
|
if button_id == "start_over_btn":
|
||||||
|
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."""
|
||||||
|
# Show notification that analysis is starting
|
||||||
|
self.app.notify(
|
||||||
|
"Starting analysis - this may take several minutes for large policies",
|
||||||
|
severity="information",
|
||||||
|
timeout=5,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 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("")
|
||||||
|
|
||||||
|
# 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:
|
||||||
|
# Get agents in the selected policy
|
||||||
|
agents = self.api.agents_find_by_group(self.selected_policy.groupid)
|
||||||
|
|
||||||
|
if agents.empty:
|
||||||
|
self.app.notify(
|
||||||
|
f"No agents found in policy: {self.selected_policy.name}",
|
||||||
|
severity="warning",
|
||||||
|
timeout=5,
|
||||||
|
)
|
||||||
|
self._show_policy_selection()
|
||||||
|
return
|
||||||
|
|
||||||
|
# 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
|
||||||
|
)
|
||||||
|
|
||||||
|
if policy_exec_history.empty:
|
||||||
|
logger.info(
|
||||||
|
"No execution history found for the selected policy and time range."
|
||||||
|
)
|
||||||
|
# All agents are quiet (no executions)
|
||||||
|
agents["execution_count"] = 0
|
||||||
|
agents["days_since"] = None
|
||||||
|
agents["required_quiet"] = self.quiet_days
|
||||||
|
agents["enforce_ready"] = True
|
||||||
|
else:
|
||||||
|
# Convert datetime column
|
||||||
|
policy_exec_history["datetime"] = pd.to_datetime(
|
||||||
|
policy_exec_history["datetime"],
|
||||||
|
format="%Y-%m-%dT%H:%M:%SZ",
|
||||||
|
utc=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Calculate days ago
|
||||||
|
now = datetime.datetime.now(datetime.timezone.utc)
|
||||||
|
policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
|
||||||
|
lambda dt: (now - dt).days
|
||||||
|
)
|
||||||
|
|
||||||
|
# Count total executions per hostname
|
||||||
|
hostname_counts = policy_exec_history["hostname"].value_counts()
|
||||||
|
agents["execution_count"] = (
|
||||||
|
agents["hostname"].map(hostname_counts).fillna(0).astype(int)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Find most recent execution per hostname
|
||||||
|
most_recent_exec = policy_exec_history.sort_values(
|
||||||
|
by="days_ago"
|
||||||
|
).drop_duplicates(subset="hostname", keep="first")
|
||||||
|
|
||||||
|
# Map most recent execution age to agents
|
||||||
|
agents["days_since"] = agents["hostname"].map(
|
||||||
|
most_recent_exec.set_index("hostname")["days_ago"]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check for enforcement readiness
|
||||||
|
agents["required_quiet"] = self.quiet_days
|
||||||
|
agents["enforce_ready"] = agents["days_since"].apply(
|
||||||
|
lambda x: True if pd.isna(x) or x > self.quiet_days else False
|
||||||
|
)
|
||||||
|
|
||||||
|
# Sort agents
|
||||||
|
agents = agents.sort_values(
|
||||||
|
by=["execution_count", "hostname"], ascending=[True, True]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Store the results
|
||||||
|
self.agents_df = agents
|
||||||
|
|
||||||
|
# Categorize agents into DataFrames
|
||||||
|
self.enforce_ready_df = agents[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, "
|
||||||
|
f"{len(self.non_enforce_ready_df)} non-enforce ready"
|
||||||
|
)
|
||||||
|
|
||||||
|
self.app.notify(
|
||||||
|
f"Analysis complete! Found {len(self.enforce_ready_df)} enforce ready, "
|
||||||
|
f"{len(self.non_enforce_ready_df)} not ready",
|
||||||
|
severity="success",
|
||||||
|
timeout=5,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Show results
|
||||||
|
self._show_results()
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error during analysis: {e}", exc_info=True)
|
||||||
|
self.app.notify(f"Analysis failed: {str(e)}", severity="error", timeout=5)
|
||||||
|
self._show_policy_selection()
|
||||||
|
|
||||||
|
def _show_results(self) -> None:
|
||||||
|
"""Show the categorized results."""
|
||||||
|
self.workflow_stage = "view_results"
|
||||||
|
content = self.query_one("#content_area", Vertical)
|
||||||
|
content.remove_children()
|
||||||
|
|
||||||
|
# Create results display container and mount it first
|
||||||
|
results_container = Vertical(id="results_container")
|
||||||
|
results_container.styles.height = "auto"
|
||||||
|
results_container.styles.margin = (1, 1)
|
||||||
|
content.mount(results_container)
|
||||||
|
|
||||||
|
# Summary statistics
|
||||||
|
total_agents = len(self.enforce_ready_df) + len(self.non_enforce_ready_df)
|
||||||
|
ready_count = len(self.enforce_ready_df)
|
||||||
|
not_ready_count = len(self.non_enforce_ready_df)
|
||||||
|
ready_percentage = (ready_count / total_agents * 100) if total_agents > 0 else 0
|
||||||
|
|
||||||
|
summary = Static(
|
||||||
|
f"Analysis Results for: {self.selected_policy.name}\n\n"
|
||||||
|
f"Total Agents: {total_agents}\n"
|
||||||
|
f"Enforce Ready: {ready_count} ({ready_percentage:.1f}%)\n"
|
||||||
|
f"Not Ready: {not_ready_count} ({100 - ready_percentage:.1f}%)\n\n"
|
||||||
|
f"Quiet Threshold: {self.quiet_days} days\n"
|
||||||
|
f"History Period: {self.history_days} days",
|
||||||
|
id="results_summary",
|
||||||
|
)
|
||||||
|
summary.styles.margin = (0, 0, 2, 0)
|
||||||
|
results_container.mount(summary)
|
||||||
|
|
||||||
|
# Action buttons
|
||||||
|
button_container = Horizontal(id="results_buttons")
|
||||||
|
button_container.styles.height = "auto"
|
||||||
|
results_container.mount(button_container)
|
||||||
|
|
||||||
|
if ready_count > 0:
|
||||||
|
enforce_btn = Button(
|
||||||
|
f"Select Target for Enforce Ready ({ready_count})",
|
||||||
|
id="select_enforce_target_btn",
|
||||||
|
)
|
||||||
|
enforce_btn.styles.margin = (0, 1, 1, 0)
|
||||||
|
button_container.mount(enforce_btn)
|
||||||
|
|
||||||
|
if not_ready_count > 0:
|
||||||
|
non_enforce_btn = Button(
|
||||||
|
f"Select Target for Non-Enforce Ready ({not_ready_count})",
|
||||||
|
id="select_non_enforce_target_btn",
|
||||||
|
)
|
||||||
|
non_enforce_btn.styles.margin = (0, 1, 1, 0)
|
||||||
|
button_container.mount(non_enforce_btn)
|
||||||
|
|
||||||
|
export_btn = Button("Export Results", id="export_results_btn")
|
||||||
|
export_btn.styles.margin = (0, 1, 1, 0)
|
||||||
|
button_container.mount(export_btn)
|
||||||
|
|
||||||
|
start_over_btn = Button("Start Over", id="start_over_btn")
|
||||||
|
start_over_btn.styles.margin = (0, 0, 1, 0)
|
||||||
|
button_container.mount(start_over_btn)
|
||||||
|
|
||||||
|
# Tables showing agents
|
||||||
|
tables_container = Horizontal()
|
||||||
|
tables_container.styles.height = "1fr"
|
||||||
|
results_container.mount(tables_container)
|
||||||
|
|
||||||
|
# Enforce Ready table
|
||||||
|
if ready_count > 0:
|
||||||
|
enforce_col = Vertical()
|
||||||
|
enforce_col.styles.width = "1fr"
|
||||||
|
enforce_col.styles.margin = (1, 1, 0, 0)
|
||||||
|
tables_container.mount(enforce_col)
|
||||||
|
|
||||||
|
enforce_label = Static("Enforce Ready Agents")
|
||||||
|
enforce_label.styles.margin = (0, 0, 1, 0)
|
||||||
|
enforce_col.mount(enforce_label)
|
||||||
|
|
||||||
|
enforce_table = DataTable(id="enforce_ready_table")
|
||||||
|
enforce_table.styles.height = "1fr"
|
||||||
|
enforce_table.add_columns("Hostname", "Last Exec (days)")
|
||||||
|
|
||||||
|
# Display first 50 agents
|
||||||
|
for idx, row in self.enforce_ready_df.head(50).iterrows():
|
||||||
|
days_since = row["days_since"]
|
||||||
|
days_str = f"{int(days_since)}" if not pd.isna(days_since) else "Never"
|
||||||
|
enforce_table.add_row(row["hostname"], days_str)
|
||||||
|
|
||||||
|
if len(self.enforce_ready_df) > 50:
|
||||||
|
enforce_table.add_row(
|
||||||
|
f"... and {len(self.enforce_ready_df) - 50} more", ""
|
||||||
|
)
|
||||||
|
|
||||||
|
enforce_col.mount(enforce_table)
|
||||||
|
|
||||||
|
# Non-Enforce Ready table
|
||||||
|
if not_ready_count > 0:
|
||||||
|
non_enforce_col = Vertical()
|
||||||
|
non_enforce_col.styles.width = "1fr"
|
||||||
|
non_enforce_col.styles.margin = (1, 0, 0, 1)
|
||||||
|
tables_container.mount(non_enforce_col)
|
||||||
|
|
||||||
|
non_enforce_label = Static("Non-Enforce Ready Agents")
|
||||||
|
non_enforce_label.styles.margin = (0, 0, 1, 0)
|
||||||
|
non_enforce_col.mount(non_enforce_label)
|
||||||
|
|
||||||
|
non_enforce_table = DataTable(id="non_enforce_ready_table")
|
||||||
|
non_enforce_table.styles.height = "1fr"
|
||||||
|
non_enforce_table.add_columns("Hostname", "Last Exec (days)")
|
||||||
|
|
||||||
|
# Display first 50 agents
|
||||||
|
for idx, row in self.non_enforce_ready_df.head(50).iterrows():
|
||||||
|
days_since = row["days_since"]
|
||||||
|
days_str = f"{int(days_since)}" if not pd.isna(days_since) else "N/A"
|
||||||
|
non_enforce_table.add_row(row["hostname"], days_str)
|
||||||
|
|
||||||
|
if len(self.non_enforce_ready_df) > 50:
|
||||||
|
non_enforce_table.add_row(
|
||||||
|
f"... and {len(self.non_enforce_ready_df) - 50} more", ""
|
||||||
|
)
|
||||||
|
|
||||||
|
non_enforce_col.mount(non_enforce_table)
|
||||||
|
|
||||||
|
def _show_enforce_target_selection(self) -> None:
|
||||||
|
"""Show policy selection for enforce ready agents."""
|
||||||
|
self.workflow_stage = "select_enforce_target"
|
||||||
|
content = self.query_one("#content_area", Vertical)
|
||||||
|
content.remove_children()
|
||||||
|
|
||||||
|
# Info message
|
||||||
|
info = Static(
|
||||||
|
f"Select target policy for {len(self.enforce_ready_df)} Enforce Ready agents\n"
|
||||||
|
f"Source Policy: {self.selected_policy.name}",
|
||||||
|
id="enforce_target_info",
|
||||||
|
)
|
||||||
|
info.styles.margin = (0, 0, 2, 0)
|
||||||
|
content.mount(info)
|
||||||
|
|
||||||
|
# Policy selector
|
||||||
|
policy_selector = PolicySelector(self.policies)
|
||||||
|
content.mount(policy_selector)
|
||||||
|
|
||||||
|
# Skip button
|
||||||
|
skip_btn = Button("Skip - No Migration", id="skip_enforce_target_btn")
|
||||||
|
skip_btn.styles.width = "50%"
|
||||||
|
skip_btn.styles.margin = (2, 0, 0, 0)
|
||||||
|
content.mount(skip_btn)
|
||||||
|
|
||||||
|
def _show_non_enforce_target_selection(self) -> None:
|
||||||
|
"""Show policy selection for non-enforce ready agents."""
|
||||||
|
self.workflow_stage = "select_non_enforce_target"
|
||||||
|
content = self.query_one("#content_area", Vertical)
|
||||||
|
content.remove_children()
|
||||||
|
|
||||||
|
# Info message
|
||||||
|
info = Static(
|
||||||
|
f"Select target policy for {len(self.non_enforce_ready_df)} Non-Enforce Ready agents\n"
|
||||||
|
f"Source Policy: {self.selected_policy.name}",
|
||||||
|
id="non_enforce_target_info",
|
||||||
|
)
|
||||||
|
info.styles.margin = (0, 0, 2, 0)
|
||||||
|
content.mount(info)
|
||||||
|
|
||||||
|
# Policy selector
|
||||||
|
policy_selector = PolicySelector(self.policies)
|
||||||
|
content.mount(policy_selector)
|
||||||
|
|
||||||
|
# Skip button
|
||||||
|
skip_btn = Button("Skip - No Migration", id="skip_non_enforce_target_btn")
|
||||||
|
skip_btn.styles.width = "50%"
|
||||||
|
skip_btn.styles.margin = (2, 0, 0, 0)
|
||||||
|
content.mount(skip_btn)
|
||||||
|
|
||||||
|
def _show_migration_confirmation(self) -> None:
|
||||||
|
"""Show migration confirmation screen."""
|
||||||
|
self.workflow_stage = "confirm_migration"
|
||||||
|
content = self.query_one("#content_area", Vertical)
|
||||||
|
content.remove_children()
|
||||||
|
|
||||||
|
# Build confirmation message
|
||||||
|
confirmation_lines = [
|
||||||
|
"Migration Summary\n",
|
||||||
|
f"Source Policy: {self.selected_policy.name}\n",
|
||||||
|
]
|
||||||
|
|
||||||
|
if self.enforce_ready_target_policy:
|
||||||
|
confirmation_lines.append(
|
||||||
|
f"\nEnforce Ready Migration:\n"
|
||||||
|
f"Agents: {len(self.enforce_ready_df)}\n"
|
||||||
|
f"Target: {self.enforce_ready_target_policy.name}\n"
|
||||||
|
)
|
||||||
|
|
||||||
|
if self.non_enforce_ready_target_policy:
|
||||||
|
confirmation_lines.append(
|
||||||
|
f"\nNon-Enforce Ready Migration:\n"
|
||||||
|
f"Agents: {len(self.non_enforce_ready_df)}\n"
|
||||||
|
f"Target: {self.non_enforce_ready_target_policy.name}\n"
|
||||||
|
)
|
||||||
|
|
||||||
|
if (
|
||||||
|
not self.enforce_ready_target_policy
|
||||||
|
and not self.non_enforce_ready_target_policy
|
||||||
|
):
|
||||||
|
confirmation_lines.append("\nNo migrations will be performed.")
|
||||||
|
|
||||||
|
confirmation = Static("".join(confirmation_lines), id="migration_confirmation")
|
||||||
|
confirmation.styles.margin = (1, 1, 2, 1)
|
||||||
|
content.mount(confirmation)
|
||||||
|
|
||||||
|
# Action buttons - mount container first, then add buttons
|
||||||
|
button_container = Horizontal(id="confirmation_buttons")
|
||||||
|
button_container.styles.height = "auto"
|
||||||
|
button_container.styles.margin = (1, 1)
|
||||||
|
content.mount(button_container)
|
||||||
|
|
||||||
|
if self.enforce_ready_target_policy or self.non_enforce_ready_target_policy:
|
||||||
|
confirm_btn = Button("Confirm Migration", id="confirm_migration_btn")
|
||||||
|
confirm_btn.styles.margin = (0, 1, 0, 0)
|
||||||
|
button_container.mount(confirm_btn)
|
||||||
|
|
||||||
|
cancel_btn = Button("Cancel", id="cancel_migration_btn")
|
||||||
|
button_container.mount(cancel_btn)
|
||||||
|
|
||||||
|
def _execute_migration(self) -> None:
|
||||||
|
"""Execute the agent migrations."""
|
||||||
|
self.workflow_stage = "executing"
|
||||||
|
content = self.query_one("#content_area", Vertical)
|
||||||
|
content.remove_children()
|
||||||
|
|
||||||
|
# Show executing message
|
||||||
|
executing_msg = Static(
|
||||||
|
"Executing agent migrations...\nPlease wait...",
|
||||||
|
id="executing_message",
|
||||||
|
)
|
||||||
|
executing_msg.styles.margin = (2, 1)
|
||||||
|
content.mount(executing_msg)
|
||||||
|
|
||||||
|
# Perform migrations asynchronously
|
||||||
|
self.call_later(self._perform_migrations)
|
||||||
|
|
||||||
|
def _perform_migrations(self) -> None:
|
||||||
|
"""Perform the actual agent migrations."""
|
||||||
|
successful_migrations = []
|
||||||
|
failed_migrations = []
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Migrate enforce ready agents
|
||||||
|
if self.enforce_ready_target_policy:
|
||||||
|
for idx, row in self.enforce_ready_df.iterrows():
|
||||||
|
try:
|
||||||
|
self.api.agent_move(
|
||||||
|
row["agentid"], self.enforce_ready_target_policy.groupid
|
||||||
|
)
|
||||||
|
successful_migrations.append(
|
||||||
|
(row["hostname"], self.enforce_ready_target_policy.name)
|
||||||
|
)
|
||||||
|
logger.debug(
|
||||||
|
f"Moved {row['hostname']} to {self.enforce_ready_target_policy.name}"
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
failed_migrations.append((row["hostname"], str(e)))
|
||||||
|
logger.error(f"Failed to move {row['hostname']}: {e}")
|
||||||
|
|
||||||
|
# Migrate non-enforce ready agents
|
||||||
|
if self.non_enforce_ready_target_policy:
|
||||||
|
for idx, row in self.non_enforce_ready_df.iterrows():
|
||||||
|
try:
|
||||||
|
self.api.agent_move(
|
||||||
|
row["agentid"], self.non_enforce_ready_target_policy.groupid
|
||||||
|
)
|
||||||
|
successful_migrations.append(
|
||||||
|
(row["hostname"], self.non_enforce_ready_target_policy.name)
|
||||||
|
)
|
||||||
|
logger.debug(
|
||||||
|
f"Moved {row['hostname']} to {self.non_enforce_ready_target_policy.name}"
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
failed_migrations.append((row["hostname"], str(e)))
|
||||||
|
logger.error(f"Failed to move {row['hostname']}: {e}")
|
||||||
|
|
||||||
|
# Show completion results
|
||||||
|
self._show_completion_results(successful_migrations, failed_migrations)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error during migration execution: {e}", exc_info=True)
|
||||||
|
self.app.notify(f"Migration failed: {str(e)}", severity="error", timeout=5)
|
||||||
|
self._show_results()
|
||||||
|
|
||||||
|
def _show_completion_results(
|
||||||
|
self, successful: List[tuple], failed: List[tuple]
|
||||||
|
) -> None:
|
||||||
|
"""Show migration completion results."""
|
||||||
|
self.workflow_stage = "complete"
|
||||||
|
content = self.query_one("#content_area", Vertical)
|
||||||
|
content.remove_children()
|
||||||
|
|
||||||
|
# Results summary
|
||||||
|
total_attempted = len(successful) + len(failed)
|
||||||
|
success_rate = (
|
||||||
|
(len(successful) / total_attempted * 100) if total_attempted > 0 else 0
|
||||||
|
)
|
||||||
|
|
||||||
|
results = Static(
|
||||||
|
f"Migration Complete\n\n"
|
||||||
|
f"Total Agents Migrated: {len(successful)}\n"
|
||||||
|
f"Failed Migrations: {len(failed)}\n"
|
||||||
|
f"Success Rate: {success_rate:.1f}%",
|
||||||
|
id="completion_summary",
|
||||||
|
)
|
||||||
|
results.styles.margin = (1, 1, 2, 1)
|
||||||
|
content.mount(results)
|
||||||
|
|
||||||
|
# Details tables
|
||||||
|
if successful:
|
||||||
|
success_container = Vertical()
|
||||||
|
success_container.styles.margin = (0, 1)
|
||||||
|
content.mount(success_container)
|
||||||
|
|
||||||
|
success_label = Static("Successful Migrations")
|
||||||
|
success_label.styles.margin = (0, 0, 1, 0)
|
||||||
|
success_container.mount(success_label)
|
||||||
|
|
||||||
|
success_table = DataTable(id="success_table")
|
||||||
|
success_table.styles.height = "auto"
|
||||||
|
success_table.add_columns("Hostname", "Target Policy")
|
||||||
|
|
||||||
|
for hostname, target_policy in successful[:25]: # Show first 25
|
||||||
|
success_table.add_row(hostname, target_policy)
|
||||||
|
|
||||||
|
if len(successful) > 25:
|
||||||
|
success_table.add_row(f"... and {len(successful) - 25} more", "")
|
||||||
|
|
||||||
|
success_container.mount(success_table)
|
||||||
|
|
||||||
|
if failed:
|
||||||
|
failed_container = Vertical()
|
||||||
|
failed_container.styles.margin = (2, 1, 0, 1)
|
||||||
|
content.mount(failed_container)
|
||||||
|
|
||||||
|
failed_label = Static("Failed Migrations")
|
||||||
|
failed_label.styles.margin = (0, 0, 1, 0)
|
||||||
|
failed_container.mount(failed_label)
|
||||||
|
|
||||||
|
failed_table = DataTable(id="failed_table")
|
||||||
|
failed_table.styles.height = "auto"
|
||||||
|
failed_table.add_columns("Hostname", "Error")
|
||||||
|
|
||||||
|
for hostname, error in failed[:25]: # Show first 25
|
||||||
|
failed_table.add_row(hostname, error[:50]) # Truncate error
|
||||||
|
|
||||||
|
if len(failed) > 25:
|
||||||
|
failed_table.add_row(f"... and {len(failed) - 25} more", "")
|
||||||
|
|
||||||
|
failed_container.mount(failed_table)
|
||||||
|
|
||||||
|
# Action button
|
||||||
|
done_btn = Button("Done", id="start_over_btn")
|
||||||
|
done_btn.styles.width = "50%"
|
||||||
|
done_btn.styles.margin = (2, 0, 0, 0)
|
||||||
|
content.mount(done_btn)
|
||||||
|
|
||||||
|
def _export_results(self) -> None:
|
||||||
|
"""Export analysis results to CSV."""
|
||||||
|
try:
|
||||||
|
working_dir = load_env("WORKING_DIR") or os.getcwd()
|
||||||
|
filename = os.path.join(
|
||||||
|
working_dir,
|
||||||
|
f"{self.selected_policy.name}_quiet_analysis_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
|
||||||
|
)
|
||||||
|
|
||||||
|
self.agents_df.to_csv(filename, index=False)
|
||||||
|
logger.info(f"Exported results to {filename}")
|
||||||
|
self.app.notify(
|
||||||
|
f"Results exported to:\n{filename}",
|
||||||
|
severity="information",
|
||||||
|
timeout=5,
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to export results: {e}")
|
||||||
|
self.app.notify(f"Export failed: {str(e)}", severity="error", timeout=5)
|
||||||
|
|
||||||
|
def action_go_back(self) -> None:
|
||||||
|
"""Handle back/escape action."""
|
||||||
|
# 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_history_days":
|
||||||
|
self._show_policy_selection()
|
||||||
|
elif self.workflow_stage == "select_enforce_target":
|
||||||
|
self._show_results()
|
||||||
|
elif self.workflow_stage == "select_non_enforce_target":
|
||||||
|
if self.enforce_ready_target_policy:
|
||||||
|
self._show_enforce_target_selection()
|
||||||
|
else:
|
||||||
|
self._show_results()
|
||||||
|
elif self.workflow_stage == "confirm_migration":
|
||||||
|
self._show_non_enforce_target_selection()
|
||||||
|
else:
|
||||||
|
self.app.pop_screen()
|
||||||
|
|
||||||
|
def action_main_menu(self) -> None:
|
||||||
|
"""Go back to main menu."""
|
||||||
|
while len(self.app.screen_stack) > 2:
|
||||||
|
self.app.pop_screen()
|
||||||
@@ -1,3 +1,18 @@
|
|||||||
|
# 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 textual.color import Color
|
from textual.color import Color
|
||||||
from textual.theme import Theme
|
from textual.theme import Theme
|
||||||
|
|
||||||
@@ -12,7 +27,7 @@ def get_amber_terminal_theme():
|
|||||||
success=Color.parse("#ffb733"),
|
success=Color.parse("#ffb733"),
|
||||||
warning=Color.parse("#ffff66"),
|
warning=Color.parse("#ffff66"),
|
||||||
error=Color.parse("#ff3300"),
|
error=Color.parse("#ff3300"),
|
||||||
surface=Color.parse("#3a1f00"), # brighter brown for blending
|
surface=Color.parse("#49331a"), # brighter brown for blending
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -1,3 +1,18 @@
|
|||||||
|
# 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 textual.color import Color
|
from textual.color import Color
|
||||||
|
|
||||||
|
|
||||||
@@ -1,3 +1,18 @@
|
|||||||
|
# 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 textual.containers import Vertical
|
from textual.containers import Vertical
|
||||||
from textual.message import Message
|
from textual.message import Message
|
||||||
from textual.widget import Widget
|
from textual.widget import Widget
|
||||||
@@ -1,5 +1,20 @@
|
|||||||
|
# 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 logging
|
||||||
from typing import List
|
from typing import List, Optional
|
||||||
|
|
||||||
from textual.containers import Horizontal, Vertical
|
from textual.containers import Horizontal, Vertical
|
||||||
from textual.css.query import NoMatches
|
from textual.css.query import NoMatches
|
||||||
@@ -23,6 +38,8 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
|
|
||||||
class OTPGenerator(Widget):
|
class OTPGenerator(Widget):
|
||||||
|
"""Widget for generating OTPs for selected devices."""
|
||||||
|
|
||||||
# Reactive properties to track form completion
|
# Reactive properties to track form completion
|
||||||
requestor_filled = reactive(False)
|
requestor_filled = reactive(False)
|
||||||
reasoning_filled = reactive(False)
|
reasoning_filled = reactive(False)
|
||||||
@@ -31,7 +48,11 @@ class OTPGenerator(Widget):
|
|||||||
|
|
||||||
class OTPInfo(Message):
|
class OTPInfo(Message):
|
||||||
def __init__(
|
def __init__(
|
||||||
self, devices: List[Agent], requestor: str, reasoning: str, duration: int
|
self,
|
||||||
|
devices: Optional[List[Agent]],
|
||||||
|
requestor: str,
|
||||||
|
reasoning: str,
|
||||||
|
duration: int,
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.devices = devices
|
self.devices = devices
|
||||||
@@ -135,14 +156,10 @@ class OTPGenerator(Widget):
|
|||||||
button_row.styles.height = "auto"
|
button_row.styles.height = "auto"
|
||||||
button_row.styles.margin = (1, 0, 0, 0)
|
button_row.styles.margin = (1, 0, 0, 0)
|
||||||
|
|
||||||
back_button = Button("← Back", id="back_button")
|
|
||||||
back_button.styles.width = "1fr"
|
|
||||||
yield back_button
|
|
||||||
|
|
||||||
generate_button = Button(
|
generate_button = Button(
|
||||||
"Generate OTP", id="generate_button", variant="primary"
|
"Generate OTP", id="generate_button", variant="primary"
|
||||||
)
|
)
|
||||||
generate_button.styles.width = "2fr"
|
generate_button.styles.width = "100%"
|
||||||
yield generate_button
|
yield generate_button
|
||||||
|
|
||||||
# Right side - Show device list initially, then output after generation
|
# Right side - Show device list initially, then output after generation
|
||||||
@@ -165,7 +182,7 @@ class OTPGenerator(Widget):
|
|||||||
|
|
||||||
# Show device list initially
|
# Show device list initially
|
||||||
device_list_text = "\n".join(
|
device_list_text = "\n".join(
|
||||||
f"• {device.hostname}" for device in self.devices
|
f"{device.hostname}" for device in self.devices
|
||||||
)
|
)
|
||||||
device_display = Static(device_list_text, id="device_display")
|
device_display = Static(device_list_text, id="device_display")
|
||||||
yield device_display
|
yield device_display
|
||||||
@@ -193,14 +210,7 @@ class OTPGenerator(Widget):
|
|||||||
def on_button_pressed(self, event: Button.Pressed):
|
def on_button_pressed(self, event: Button.Pressed):
|
||||||
btn_id = event.button.id
|
btn_id = event.button.id
|
||||||
|
|
||||||
if btn_id == "back_button":
|
if btn_id == "copy_clipboard_button":
|
||||||
|
|
||||||
while len(self.app.screen_stack) > 2:
|
|
||||||
self.app.pop_screen()
|
|
||||||
|
|
||||||
event.stop()
|
|
||||||
|
|
||||||
elif btn_id == "copy_clipboard_button":
|
|
||||||
try:
|
try:
|
||||||
output_area = self.query_one("#otp_output", TextArea)
|
output_area = self.query_one("#otp_output", TextArea)
|
||||||
text_to_copy = output_area.text
|
text_to_copy = output_area.text
|
||||||
@@ -209,7 +219,7 @@ class OTPGenerator(Widget):
|
|||||||
|
|
||||||
pyperclip.copy(text_to_copy)
|
pyperclip.copy(text_to_copy)
|
||||||
self.app.notify(
|
self.app.notify(
|
||||||
"✅ Copied to clipboard!", severity="information", timeout=2
|
"✓ Copied to clipboard!", severity="information", timeout=2
|
||||||
)
|
)
|
||||||
except ImportError:
|
except ImportError:
|
||||||
self.app.notify(
|
self.app.notify(
|
||||||
@@ -243,7 +253,7 @@ class OTPGenerator(Widget):
|
|||||||
self.otp_generated = True
|
self.otp_generated = True
|
||||||
|
|
||||||
# Access API from the app - this is the key change!
|
# Access API from the app - this is the key change!
|
||||||
api = self.app.api
|
api = self.app.api # type: ignore
|
||||||
|
|
||||||
output_lines = [
|
output_lines = [
|
||||||
"Requested OTP Codes:",
|
"Requested OTP Codes:",
|
||||||
@@ -1,22 +1,18 @@
|
|||||||
"""
|
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||||
Agent Move Operations Widget Module
|
#
|
||||||
|
# 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/>.
|
||||||
|
|
||||||
This module provides a Textual-based UI widget for performing bulk operations on
|
|
||||||
agent devices in the Airlock system. It allows users to:
|
|
||||||
- View selected agents and their current policy assignments
|
|
||||||
- Move agents to local approval mode with OTP enforcement
|
|
||||||
- Toggle agents between audit and enforcement policy modes
|
|
||||||
- Select and move agents to alternate policies (future implementation)
|
|
||||||
|
|
||||||
The widget tracks operation state, manages button availability, and displays
|
|
||||||
results with success/failure summaries that can be copied to clipboard.
|
|
||||||
|
|
||||||
Dependencies:
|
|
||||||
- textual: TUI framework for building the widget and UI components
|
|
||||||
- models.agent: Agent model class
|
|
||||||
- services.agenthandler: Core agent operation functions
|
|
||||||
- flows.localApproval: Local approval workflow handling
|
|
||||||
"""
|
|
||||||
|
|
||||||
from dataclasses import asdict
|
from dataclasses import asdict
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
@@ -30,12 +26,13 @@ from textual.css.query import NoMatches
|
|||||||
from textual.message import Message
|
from textual.message import Message
|
||||||
from textual.reactive import reactive
|
from textual.reactive import reactive
|
||||||
from textual.widget import Widget
|
from textual.widget import Widget
|
||||||
from textual.widgets import Button, DataTable, Header, Static, TextArea
|
from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea
|
||||||
|
|
||||||
from models.agent import Agent
|
from models.agent import Agent
|
||||||
from screens.otpworkflowscreen import OTPWorkflowScreen
|
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
|
||||||
from screens.policyselectorscreen import PolicySelectorScreen
|
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
|
||||||
from widgets.OTP_generate import OTPGenerator
|
from TUI.Screens.policyselectorscreen import PolicySelectorScreen
|
||||||
|
from TUI.Widgets.OTP_generate import OTPGenerator
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -144,6 +141,7 @@ class AgentMoveOperations(Widget):
|
|||||||
toggle_enforcement_btn = self.query_one("#toggle_enforcement_btn", Button)
|
toggle_enforcement_btn = self.query_one("#toggle_enforcement_btn", Button)
|
||||||
other_policy_btn = self.query_one("#other_policy_btn", Button)
|
other_policy_btn = self.query_one("#other_policy_btn", Button)
|
||||||
otp_gen_btn = self.query_one("#otp_gen_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 operation in progress, disable all
|
||||||
if self.operation_in_progress:
|
if self.operation_in_progress:
|
||||||
@@ -152,6 +150,7 @@ class AgentMoveOperations(Widget):
|
|||||||
local_approval_btn.disabled = True
|
local_approval_btn.disabled = True
|
||||||
toggle_enforcement_btn.disabled = True
|
toggle_enforcement_btn.disabled = True
|
||||||
other_policy_btn.disabled = True
|
other_policy_btn.disabled = True
|
||||||
|
exec_history_btn.disabled = True
|
||||||
else:
|
else:
|
||||||
# If an operation was selected, disable
|
# If an operation was selected, disable
|
||||||
if self.selected_operation:
|
if self.selected_operation:
|
||||||
@@ -166,6 +165,9 @@ class AgentMoveOperations(Widget):
|
|||||||
other_policy_btn.disabled = (
|
other_policy_btn.disabled = (
|
||||||
self.selected_operation == "other_policy"
|
self.selected_operation == "other_policy"
|
||||||
)
|
)
|
||||||
|
exec_history_btn.disabled = (
|
||||||
|
self.selected_operation == "exec_history"
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
# Enable all buttons
|
# Enable all buttons
|
||||||
otp_gen_btn = False
|
otp_gen_btn = False
|
||||||
@@ -173,6 +175,7 @@ class AgentMoveOperations(Widget):
|
|||||||
local_approval_btn.disabled = False
|
local_approval_btn.disabled = False
|
||||||
toggle_enforcement_btn.disabled = False
|
toggle_enforcement_btn.disabled = False
|
||||||
other_policy_btn.disabled = False
|
other_policy_btn.disabled = False
|
||||||
|
exec_history_btn.disabled = False
|
||||||
|
|
||||||
except NoMatches:
|
except NoMatches:
|
||||||
pass
|
pass
|
||||||
@@ -252,7 +255,7 @@ class AgentMoveOperations(Widget):
|
|||||||
- Operations panel: 1/3 width
|
- Operations panel: 1/3 width
|
||||||
- Results area: Initially hidden, shown after operation completion
|
- Results area: Initially hidden, shown after operation completion
|
||||||
"""
|
"""
|
||||||
yield Header(show_clock=True, icon="⚙")
|
yield Header(show_clock=True, icon="⚙️")
|
||||||
title_text = Static(
|
title_text = Static(
|
||||||
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
|
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
|
||||||
id="move_ops_title",
|
id="move_ops_title",
|
||||||
@@ -289,7 +292,7 @@ class AgentMoveOperations(Widget):
|
|||||||
yield operations_label
|
yield operations_label
|
||||||
|
|
||||||
# Operation buttons
|
# Operation buttons
|
||||||
export_csv_btn = Button("📈 Export CSV", id="export_csv_btn")
|
export_csv_btn = Button("📄 Export CSV", id="export_csv_btn")
|
||||||
export_csv_btn.styles.width = "100%"
|
export_csv_btn.styles.width = "100%"
|
||||||
export_csv_btn.styles.margin = (0, 0, 1, 0)
|
export_csv_btn.styles.margin = (0, 0, 1, 0)
|
||||||
yield export_csv_btn
|
yield export_csv_btn
|
||||||
@@ -314,21 +317,25 @@ class AgentMoveOperations(Widget):
|
|||||||
yield toggle_enforcement_btn
|
yield toggle_enforcement_btn
|
||||||
|
|
||||||
other_policy_btn = Button(
|
other_policy_btn = Button(
|
||||||
"🔀 Move to Other Policy", id="other_policy_btn"
|
"🔀 Move to Other Policy", id="other_policy_btn"
|
||||||
)
|
)
|
||||||
other_policy_btn.styles.width = "100%"
|
other_policy_btn.styles.width = "100%"
|
||||||
other_policy_btn.styles.margin = (0, 0, 1, 0)
|
other_policy_btn.styles.margin = (0, 0, 1, 0)
|
||||||
yield other_policy_btn
|
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
|
||||||
status_label = Static("", id="status_label")
|
status_label = Static("", id="status_label")
|
||||||
status_label.styles.margin = (2, 0, 0, 0)
|
status_label.styles.margin = (2, 0, 0, 0)
|
||||||
yield status_label
|
yield status_label
|
||||||
|
|
||||||
back_button = Button("← Back", id="back_button")
|
yield Footer()
|
||||||
back_button.styles.width = "50%"
|
|
||||||
back_button.styles.margin = (0, 1, 1, 0)
|
|
||||||
yield back_button
|
|
||||||
|
|
||||||
def on_mount(self) -> None:
|
def on_mount(self) -> None:
|
||||||
"""
|
"""
|
||||||
@@ -363,7 +370,7 @@ class AgentMoveOperations(Widget):
|
|||||||
Handle button press events from the widget.
|
Handle button press events from the widget.
|
||||||
|
|
||||||
This Textual event handler routes button presses to appropriate actions:
|
This Textual event handler routes button presses to appropriate actions:
|
||||||
- back_button: Pop this screen (return to parent)
|
|
||||||
- copy_results_btn: Copy results text to clipboard (requires pyperclip)
|
- copy_results_btn: Copy results text to clipboard (requires pyperclip)
|
||||||
- local_approval_btn: Start local approval operation
|
- local_approval_btn: Start local approval operation
|
||||||
- toggle_enforcement_btn: Start toggle audit/enforcement operation
|
- toggle_enforcement_btn: Start toggle audit/enforcement operation
|
||||||
@@ -377,12 +384,7 @@ class AgentMoveOperations(Widget):
|
|||||||
|
|
||||||
btn_id = event.button.id
|
btn_id = event.button.id
|
||||||
|
|
||||||
if btn_id == "back_button":
|
if btn_id == "copy_results_btn":
|
||||||
while len(self.app.screen_stack) > 2:
|
|
||||||
self.app.pop_screen()
|
|
||||||
event.stop()
|
|
||||||
|
|
||||||
elif btn_id == "copy_results_btn":
|
|
||||||
try:
|
try:
|
||||||
results_text = self.query_one("#results_text", TextArea)
|
results_text = self.query_one("#results_text", TextArea)
|
||||||
import pyperclip
|
import pyperclip
|
||||||
@@ -399,7 +401,7 @@ class AgentMoveOperations(Widget):
|
|||||||
severity="warning",
|
severity="warning",
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.app.notify(f"⌠Failed to copy: {str(e)}", severity="error")
|
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
|
||||||
event.stop()
|
event.stop()
|
||||||
elif btn_id == "export_csv_btn":
|
elif btn_id == "export_csv_btn":
|
||||||
self._start_export_csv_operation()
|
self._start_export_csv_operation()
|
||||||
@@ -419,6 +421,9 @@ class AgentMoveOperations(Widget):
|
|||||||
elif btn_id == "otp_gen_btn":
|
elif btn_id == "otp_gen_btn":
|
||||||
self._start_OTP_gen_operation()
|
self._start_OTP_gen_operation()
|
||||||
event.stop()
|
event.stop()
|
||||||
|
elif btn_id == "exec_history_btn":
|
||||||
|
self._start_execution_history_operation()
|
||||||
|
event.stop()
|
||||||
|
|
||||||
def _start_local_approval_operation(self) -> None:
|
def _start_local_approval_operation(self) -> None:
|
||||||
"""
|
"""
|
||||||
@@ -503,7 +508,6 @@ class AgentMoveOperations(Widget):
|
|||||||
self.selected_operation = "export_csv"
|
self.selected_operation = "export_csv"
|
||||||
self.operation_in_progress = True
|
self.operation_in_progress = True
|
||||||
successful = []
|
successful = []
|
||||||
unsuccessful = []
|
|
||||||
status_label = self.query_one("#status_label", Static)
|
status_label = self.query_one("#status_label", Static)
|
||||||
status_label.update("Exporting CSV...")
|
status_label.update("Exporting CSV...")
|
||||||
self.app.refresh_data()
|
self.app.refresh_data()
|
||||||
@@ -577,7 +581,7 @@ class AgentMoveOperations(Widget):
|
|||||||
self.operation_in_progress = True
|
self.operation_in_progress = True
|
||||||
|
|
||||||
status_label = self.query_one("#status_label", Static)
|
status_label = self.query_one("#status_label", Static)
|
||||||
status_label.update("â³ Toggling enforcement mode...")
|
status_label.update("🔄 Toggling enforcement mode...")
|
||||||
|
|
||||||
# Get API from app
|
# Get API from app
|
||||||
api = self.app.api
|
api = self.app.api
|
||||||
@@ -587,9 +591,9 @@ class AgentMoveOperations(Widget):
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
from services.agenthandler import moveAgentToRelatedPolicy
|
from services.agenthandler import moveAgentToRelatedPolicy
|
||||||
from utils.configmanager import get_protected_json
|
from utils.configmanager import get_system_json
|
||||||
|
|
||||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||||||
|
|
||||||
for agent in self.agents:
|
for agent in self.agents:
|
||||||
try:
|
try:
|
||||||
@@ -618,7 +622,7 @@ class AgentMoveOperations(Widget):
|
|||||||
return
|
return
|
||||||
|
|
||||||
self.operation_in_progress = False
|
self.operation_in_progress = False
|
||||||
status_label.update("✅ Operation complete!")
|
status_label.update("✅ Operation complete!")
|
||||||
|
|
||||||
# Display results in the widget
|
# Display results in the widget
|
||||||
self._display_results("Toggle Audit/Enforcement", successful, unsuccessful)
|
self._display_results("Toggle Audit/Enforcement", successful, unsuccessful)
|
||||||
@@ -694,6 +698,35 @@ class AgentMoveOperations(Widget):
|
|||||||
|
|
||||||
self.app.push_screen(OTPWorkflowScreen(self.agents))
|
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:
|
def _execute_move_to_policy(self, target_policy) -> None:
|
||||||
"""
|
"""
|
||||||
Execute the actual move of agents to the selected policy.
|
Execute the actual move of agents to the selected policy.
|
||||||
@@ -715,7 +748,7 @@ class AgentMoveOperations(Widget):
|
|||||||
for agent in self.agents:
|
for agent in self.agents:
|
||||||
try:
|
try:
|
||||||
# Move agent to target policy
|
# Move agent to target policy
|
||||||
result = api.agent_move(agent.agentid, target_policy.groupid)
|
api.agent_move(agent.agentid, target_policy.groupid)
|
||||||
successful.append((agent, f"Moved to {target_policy.name}"))
|
successful.append((agent, f"Moved to {target_policy.name}"))
|
||||||
logger.info(
|
logger.info(
|
||||||
f"Successfully moved {agent.hostname} to policy {target_policy.name}"
|
f"Successfully moved {agent.hostname} to policy {target_policy.name}"
|
||||||
@@ -1,6 +1,22 @@
|
|||||||
|
# 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 difflib
|
import difflib
|
||||||
|
from pathlib import Path
|
||||||
import re
|
import re
|
||||||
from typing import List
|
from typing import List, Optional
|
||||||
|
|
||||||
from textual.containers import Horizontal, Vertical
|
from textual.containers import Horizontal, Vertical
|
||||||
from textual.css.query import NoMatches
|
from textual.css.query import NoMatches
|
||||||
@@ -20,12 +36,14 @@ from models.agent import Agent
|
|||||||
|
|
||||||
|
|
||||||
class MultiAgentSelector(Widget):
|
class MultiAgentSelector(Widget):
|
||||||
|
"""Widget for selecting multiple agents from a list."""
|
||||||
|
|
||||||
class AgentsSelected(Message):
|
class AgentsSelected(Message):
|
||||||
def __init__(self, selected_agents: List[Agent]):
|
def __init__(self, selected_agents: List[Agent]):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.selected_agents = selected_agents
|
self.selected_agents = selected_agents
|
||||||
|
|
||||||
def __init__(self, all_agents: List[Agent]):
|
def __init__(self, all_agents: Optional[List[Agent]]):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.all_agents = all_agents
|
self.all_agents = all_agents
|
||||||
self._match_type = "exact"
|
self._match_type = "exact"
|
||||||
@@ -40,7 +58,7 @@ class MultiAgentSelector(Widget):
|
|||||||
|
|
||||||
def compose(self):
|
def compose(self):
|
||||||
yield Header(show_clock=True, icon="⚙")
|
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)
|
title_text.styles.margin = (0, 0, 0, 1)
|
||||||
yield title_text
|
yield title_text
|
||||||
|
|
||||||
@@ -60,7 +78,7 @@ class MultiAgentSelector(Widget):
|
|||||||
text_area.styles.overflow_y = "auto"
|
text_area.styles.overflow_y = "auto"
|
||||||
yield text_area
|
yield text_area
|
||||||
|
|
||||||
with Horizontal(id="switch_search_container") as switch_search:
|
with Horizontal(id="switch_container"):
|
||||||
switch = Switch(value=False, id="match_switch")
|
switch = Switch(value=False, id="match_switch")
|
||||||
switch.styles.width = "auto"
|
switch.styles.width = "auto"
|
||||||
switch.styles.margin = (1, 0, 0, 0)
|
switch.styles.margin = (1, 0, 0, 0)
|
||||||
@@ -72,8 +90,13 @@ class MultiAgentSelector(Widget):
|
|||||||
switch_label.styles.margin = (2, 1, 0, 0)
|
switch_label.styles.margin = (2, 1, 0, 0)
|
||||||
yield switch_label
|
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 = Button("🔍 Search", id="search_button")
|
||||||
search.styles.margin = (1, 0, 0, 0)
|
search.styles.margin = (1, 0, 0, 1)
|
||||||
yield search
|
yield search
|
||||||
|
|
||||||
with Horizontal() as select_buttons:
|
with Horizontal() as select_buttons:
|
||||||
@@ -91,11 +114,6 @@ class MultiAgentSelector(Widget):
|
|||||||
button_row.styles.height = "auto"
|
button_row.styles.height = "auto"
|
||||||
button_row.styles.margin = (1, 0, 0, 0)
|
button_row.styles.margin = (1, 0, 0, 0)
|
||||||
|
|
||||||
back_button = Button("← Back", id="back_button")
|
|
||||||
back_button.styles.width = "1fr"
|
|
||||||
back_button.styles.margin = (0, 0, 0, 1)
|
|
||||||
yield back_button
|
|
||||||
|
|
||||||
submit_button = Button(
|
submit_button = Button(
|
||||||
"▶ Select & Continue", id="submit_selection", variant="primary"
|
"▶ Select & Continue", id="submit_selection", variant="primary"
|
||||||
)
|
)
|
||||||
@@ -123,10 +141,7 @@ class MultiAgentSelector(Widget):
|
|||||||
match_list = self.query_one("#match_results", SelectionList)
|
match_list = self.query_one("#match_results", SelectionList)
|
||||||
except NoMatches:
|
except NoMatches:
|
||||||
return
|
return
|
||||||
if btn_id == "back_button":
|
if btn_id == "select_all":
|
||||||
self.app.pop_screen()
|
|
||||||
event.stop()
|
|
||||||
elif btn_id == "select_all":
|
|
||||||
match_list.select_all()
|
match_list.select_all()
|
||||||
event.stop()
|
event.stop()
|
||||||
elif btn_id == "select_none":
|
elif btn_id == "select_none":
|
||||||
@@ -143,6 +158,9 @@ class MultiAgentSelector(Widget):
|
|||||||
]
|
]
|
||||||
self.post_message(self.AgentsSelected(selected_agents))
|
self.post_message(self.AgentsSelected(selected_agents))
|
||||||
event.stop()
|
event.stop()
|
||||||
|
elif btn_id == "load_file_button":
|
||||||
|
self._load_from_file()
|
||||||
|
event.stop()
|
||||||
elif btn_id == "search_button":
|
elif btn_id == "search_button":
|
||||||
self.update_matches()
|
self.update_matches()
|
||||||
event.stop()
|
event.stop()
|
||||||
@@ -157,7 +175,7 @@ class MultiAgentSelector(Widget):
|
|||||||
match_list.add_option((name, name))
|
match_list.add_option((name, name))
|
||||||
unmatched_label = self.query_one("#unmatched_label", Static)
|
unmatched_label = self.query_one("#unmatched_label", Static)
|
||||||
if unmatched:
|
if unmatched:
|
||||||
unmatched_label.update(f"⚠️ No matches for: {', '.join(unmatched)}")
|
unmatched_label.update(f"âš ï¸ No matches for: {', '.join(unmatched)}")
|
||||||
else:
|
else:
|
||||||
unmatched_label.update("")
|
unmatched_label.update("")
|
||||||
|
|
||||||
@@ -205,3 +223,118 @@ class MultiAgentSelector(Widget):
|
|||||||
else:
|
else:
|
||||||
unmatched.append(name)
|
unmatched.append(name)
|
||||||
return sorted(matched), unmatched
|
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
|
||||||
|
)
|
||||||
@@ -1,10 +1,17 @@
|
|||||||
"""
|
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||||
Policy Selector Widget Module
|
#
|
||||||
|
# This program is free software: you can redistribute it and/or modify
|
||||||
Provides a Textual widget for selecting target policies for bulk agent operations.
|
# it under the terms of the GNU Affero General Public License as published
|
||||||
Allows users to browse available policies and select one as the destination for
|
# by the Free Software Foundation, either version 3 of the License, or
|
||||||
moving agents. Automatically excludes parent/logical policies.
|
# (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 logging
|
||||||
import re
|
import re
|
||||||
@@ -14,7 +21,7 @@ import pandas as pd
|
|||||||
from textual.containers import Horizontal, Vertical
|
from textual.containers import Horizontal, Vertical
|
||||||
from textual.message import Message
|
from textual.message import Message
|
||||||
from textual.widget import Widget
|
from textual.widget import Widget
|
||||||
from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea
|
from textual.widgets import Button, DataTable, Static, TextArea
|
||||||
|
|
||||||
from models.policy import Policy
|
from models.policy import Policy
|
||||||
|
|
||||||
@@ -34,7 +41,7 @@ class PolicySelector(Widget):
|
|||||||
- Wildcard filtering (* and ?)
|
- Wildcard filtering (* and ?)
|
||||||
- Interactive table for policy browsing
|
- Interactive table for policy browsing
|
||||||
- Explicit confirm button for selection
|
- Explicit confirm button for selection
|
||||||
- Cancel/back button to dismiss
|
- Use escape key to go back
|
||||||
|
|
||||||
Attributes:
|
Attributes:
|
||||||
policies (list[Policy]): List of available Policy objects to display.
|
policies (list[Policy]): List of available Policy objects to display.
|
||||||
@@ -91,11 +98,10 @@ class PolicySelector(Widget):
|
|||||||
- Clear Filter button
|
- Clear Filter button
|
||||||
- Confirm Selection button
|
- Confirm Selection button
|
||||||
- Policy table displaying available policies
|
- Policy table displaying available policies
|
||||||
- Back buttons for navigation
|
- Use escape key to go back
|
||||||
"""
|
"""
|
||||||
yield Header(show_clock=True, icon="⚙")
|
|
||||||
title_text = Static(
|
title_text = Static(
|
||||||
"🎯 Select Target Policy",
|
"Select Target Policy",
|
||||||
id="policy_selector_title",
|
id="policy_selector_title",
|
||||||
)
|
)
|
||||||
title_text.styles.margin = (0, 0, 1, 0)
|
title_text.styles.margin = (0, 0, 1, 0)
|
||||||
@@ -126,12 +132,12 @@ class PolicySelector(Widget):
|
|||||||
filter_help.styles.margin = (0, 0, 1, 0)
|
filter_help.styles.margin = (0, 0, 1, 0)
|
||||||
yield filter_help
|
yield filter_help
|
||||||
|
|
||||||
apply_button = Button("✓ Apply Filter", id="filter_button")
|
apply_button = Button("🔍 Apply Filter", id="filter_button")
|
||||||
apply_button.styles.width = "100%"
|
apply_button.styles.width = "100%"
|
||||||
apply_button.styles.margin = (0, 0, 1, 0)
|
apply_button.styles.margin = (0, 0, 1, 0)
|
||||||
yield apply_button
|
yield apply_button
|
||||||
|
|
||||||
clear_button = Button("Clear Filter", id="clear_filter_button")
|
clear_button = Button("🧹 Clear Filter", id="clear_filter_button")
|
||||||
clear_button.styles.width = "100%"
|
clear_button.styles.width = "100%"
|
||||||
clear_button.styles.margin = (0, 0, 1, 0)
|
clear_button.styles.margin = (0, 0, 1, 0)
|
||||||
yield clear_button
|
yield clear_button
|
||||||
@@ -145,11 +151,6 @@ class PolicySelector(Widget):
|
|||||||
selected_label.styles.margin = (2, 0, 1, 0)
|
selected_label.styles.margin = (2, 0, 1, 0)
|
||||||
yield selected_label
|
yield selected_label
|
||||||
|
|
||||||
cancel_button = Button("← Back", id="back_button")
|
|
||||||
cancel_button.styles.width = "100%"
|
|
||||||
cancel_button.styles.margin = (1, 0, 1, 0)
|
|
||||||
yield cancel_button
|
|
||||||
|
|
||||||
# Right side - Policy table
|
# Right side - Policy table
|
||||||
with Vertical() as right_side:
|
with Vertical() as right_side:
|
||||||
right_side.styles.width = "2fr"
|
right_side.styles.width = "2fr"
|
||||||
@@ -164,8 +165,6 @@ class PolicySelector(Widget):
|
|||||||
policy_table.styles.margin = (1, 0, 1, 0)
|
policy_table.styles.margin = (1, 0, 1, 0)
|
||||||
yield policy_table
|
yield policy_table
|
||||||
|
|
||||||
yield Footer()
|
|
||||||
|
|
||||||
def on_mount(self) -> None:
|
def on_mount(self) -> None:
|
||||||
"""
|
"""
|
||||||
Initialize the policy table when the widget is mounted.
|
Initialize the policy table when the widget is mounted.
|
||||||
@@ -225,7 +224,6 @@ class PolicySelector(Widget):
|
|||||||
Handle button press events from the widget.
|
Handle button press events from the widget.
|
||||||
|
|
||||||
Routes to:
|
Routes to:
|
||||||
- back_button (Cancel): Pop screen without selecting
|
|
||||||
- filter_button (Apply Filter): Filter policies with wildcard support
|
- filter_button (Apply Filter): Filter policies with wildcard support
|
||||||
- clear_filter_button: Clear filter and show all policies
|
- clear_filter_button: Clear filter and show all policies
|
||||||
- confirm_button: Confirm selection and post message
|
- confirm_button: Confirm selection and post message
|
||||||
@@ -235,12 +233,7 @@ class PolicySelector(Widget):
|
|||||||
"""
|
"""
|
||||||
btn_id = event.button.id
|
btn_id = event.button.id
|
||||||
|
|
||||||
if btn_id == "back_button":
|
if btn_id == "filter_button":
|
||||||
while len(self.app.screen_stack) > 2:
|
|
||||||
self.app.pop_screen()
|
|
||||||
event.stop()
|
|
||||||
|
|
||||||
elif btn_id == "filter_button":
|
|
||||||
self._apply_filter()
|
self._apply_filter()
|
||||||
event.stop()
|
event.stop()
|
||||||
|
|
||||||
@@ -286,7 +279,7 @@ class PolicySelector(Widget):
|
|||||||
if self.selected_policy:
|
if self.selected_policy:
|
||||||
# Update selection display
|
# Update selection display
|
||||||
label = self.query_one("#selected_policy_label", Static)
|
label = self.query_one("#selected_policy_label", Static)
|
||||||
label.update(f"✓ Selected: {self.selected_policy.name}")
|
label.update(f"Selected: {self.selected_policy.name}")
|
||||||
|
|
||||||
# Log for debugging
|
# Log for debugging
|
||||||
logger.debug(
|
logger.debug(
|
||||||
@@ -328,7 +321,7 @@ class PolicySelector(Widget):
|
|||||||
|
|
||||||
if highlighted_name:
|
if highlighted_name:
|
||||||
label = self.query_one("#selected_policy_label", Static)
|
label = self.query_one("#selected_policy_label", Static)
|
||||||
label.update(f"→ Highlighting: {highlighted_name}")
|
label.update(f"Highlighting: {highlighted_name}")
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error handling row highlight: {e}")
|
logger.error(f"Error handling row highlight: {e}")
|
||||||
@@ -402,7 +395,9 @@ class PolicySelector(Widget):
|
|||||||
)
|
)
|
||||||
|
|
||||||
displayed_count = len(self._displayed_policies)
|
displayed_count = len(self._displayed_policies)
|
||||||
status_text = f"📊 Showing {displayed_count} of {len(self._filtered_policies)} policies"
|
status_text = (
|
||||||
|
f"Showing {displayed_count} of {len(self._filtered_policies)} policies"
|
||||||
|
)
|
||||||
self.app.notify(status_text, severity="information", timeout=2)
|
self.app.notify(status_text, severity="information", timeout=2)
|
||||||
|
|
||||||
# Clear selection when filter is applied
|
# Clear selection when filter is applied
|
||||||
@@ -477,7 +472,7 @@ class PolicySelector(Widget):
|
|||||||
"""
|
"""
|
||||||
if self.selected_policy is None:
|
if self.selected_policy is None:
|
||||||
self.app.notify(
|
self.app.notify(
|
||||||
"⚠️ Please select a policy first by clicking on a row in the table",
|
"Please select a policy first by clicking on a row in the table",
|
||||||
severity="warning",
|
severity="warning",
|
||||||
timeout=3,
|
timeout=3,
|
||||||
)
|
)
|
||||||
@@ -486,6 +481,6 @@ class PolicySelector(Widget):
|
|||||||
# Log confirmation for debugging
|
# Log confirmation for debugging
|
||||||
logger.info(f"Confirming selection of policy: {self.selected_policy.name}")
|
logger.info(f"Confirming selection of policy: {self.selected_policy.name}")
|
||||||
self.app.notify(
|
self.app.notify(
|
||||||
f"✅ Confirmed: {self.selected_policy.name}", severity="success", timeout=2
|
f"Confirmed: {self.selected_policy.name}", severity="success", timeout=2
|
||||||
)
|
)
|
||||||
self.post_message(self.PolicySelected(self.selected_policy))
|
self.post_message(self.PolicySelected(self.selected_policy))
|
||||||
@@ -1,10 +1,26 @@
|
|||||||
|
# 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 collections import defaultdict
|
from collections import defaultdict
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
from rich.text import Text
|
from rich.text import Text
|
||||||
from textual.containers import Horizontal, Vertical
|
from textual.containers import Horizontal, Vertical
|
||||||
|
from textual.message import Message
|
||||||
from textual.widget import Widget
|
from textual.widget import Widget
|
||||||
from textual.widgets import Input, OptionList, Static, Tree
|
from textual.widgets import Button, Input, OptionList, Static, Switch, Tree
|
||||||
from textual.widgets.option_list import Option
|
from textual.widgets.option_list import Option
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -13,35 +29,95 @@ logger = logging.getLogger(__name__)
|
|||||||
class PolicyTreeWidget(Widget):
|
class PolicyTreeWidget(Widget):
|
||||||
"""Widget for displaying and searching a hierarchical policy tree."""
|
"""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):
|
def __init__(self, policies, devices):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.policies = policies
|
self.policies = policies
|
||||||
self.devices = devices
|
self.devices = devices
|
||||||
self.last_highlighted_node = None
|
self.last_highlighted_node = None
|
||||||
self.leaf_counts = defaultdict(int)
|
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):
|
def compose(self):
|
||||||
|
# Create the switch and its label
|
||||||
|
switch = Switch(value=False, id="match_switch")
|
||||||
|
switch.styles.margin = (0, 0, 0, 0) # top, right, bottom, left
|
||||||
|
switch.styles.padding = (0, 0, 0, 0)
|
||||||
|
|
||||||
|
switch_label = Static("Sort: Count", id="match_switch_label")
|
||||||
|
switch_label.styles.margin = (1, 0, 0, 0)
|
||||||
|
switch_label.styles.padding = (0, 0, 0, 0)
|
||||||
|
|
||||||
|
# Create the tree
|
||||||
policy_tree = Tree("", id="policy_tree") # Label set in on_mount
|
policy_tree = Tree("", id="policy_tree") # Label set in on_mount
|
||||||
policy_tree.styles.width = "2fr"
|
policy_tree.styles.width = "2fr"
|
||||||
policy_tree.styles.height = "100%"
|
policy_tree.styles.height = "100%"
|
||||||
|
|
||||||
|
# Create the search box and details pane
|
||||||
label = Static("Device Search:")
|
label = Static("Device Search:")
|
||||||
search_box = Input(
|
search_box = Input(
|
||||||
placeholder="Search policies or devices...", id="tree_search"
|
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")
|
details_pane = Static("", id="details_pane")
|
||||||
|
|
||||||
|
# Layout the UI
|
||||||
with Horizontal():
|
with Horizontal():
|
||||||
yield policy_tree
|
yield policy_tree
|
||||||
with Vertical() as right_pane:
|
with Vertical() as right_pane:
|
||||||
right_pane.styles.width = "3fr"
|
right_pane.styles.width = "3fr"
|
||||||
|
# Use a Horizontal container for the switch and label
|
||||||
|
with Horizontal() as switch_container:
|
||||||
|
switch_container.styles.height = 3
|
||||||
|
switch_container.styles.margin = (0, 0, 0, 1)
|
||||||
|
switch_container.styles.padding = (0, 0, 0, 0)
|
||||||
|
yield switch
|
||||||
|
yield switch_label
|
||||||
|
# Add the search box and details pane
|
||||||
yield label
|
yield label
|
||||||
yield search_box
|
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
|
yield details_pane
|
||||||
|
|
||||||
def on_mount(self) -> None:
|
def on_mount(self) -> None:
|
||||||
self._precompute_leaf_counts()
|
self._precompute_leaf_counts()
|
||||||
|
|
||||||
# Update root label with total leaf count
|
# Update root label with total leaf count
|
||||||
total_leaves = sum(
|
total_leaves = sum(
|
||||||
self.leaf_counts.get(policy.groupid, 0)
|
self.leaf_counts.get(policy.groupid, 0)
|
||||||
@@ -50,8 +126,35 @@ class PolicyTreeWidget(Widget):
|
|||||||
)
|
)
|
||||||
policy_tree = self.query_one("#policy_tree", Tree)
|
policy_tree = self.query_one("#policy_tree", Tree)
|
||||||
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
|
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
|
||||||
|
|
||||||
self._build_tree()
|
self._build_tree()
|
||||||
|
# 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):
|
def _precompute_leaf_counts(self):
|
||||||
"""Precompute leaf counts for each policy group."""
|
"""Precompute leaf counts for each policy group."""
|
||||||
@@ -76,15 +179,21 @@ class PolicyTreeWidget(Widget):
|
|||||||
|
|
||||||
def _build_tree(self):
|
def _build_tree(self):
|
||||||
policy_tree = self.query_one("#policy_tree", Tree)
|
policy_tree = self.query_one("#policy_tree", Tree)
|
||||||
|
policy_tree.clear() # Clear existing nodes
|
||||||
node_map = {}
|
node_map = {}
|
||||||
|
|
||||||
# Sort top-level policies
|
# Sort top-level policies
|
||||||
top_policies = [
|
top_policies = [
|
||||||
p for p in self.policies if p.parent == "global-policy-settings"
|
p for p in self.policies if p.parent == "global-policy-settings"
|
||||||
]
|
]
|
||||||
top_policies.sort(
|
|
||||||
key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
|
# Sort by count (default) or alphabetically
|
||||||
)
|
if getattr(self, "match_type", "Count") == "Count":
|
||||||
|
top_policies.sort(
|
||||||
|
key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
|
||||||
|
)
|
||||||
|
else: # Alphabetical
|
||||||
|
top_policies.sort(key=lambda p: p.name.lower())
|
||||||
|
|
||||||
for policy in top_policies:
|
for policy in top_policies:
|
||||||
label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})"
|
label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})"
|
||||||
@@ -98,9 +207,13 @@ class PolicyTreeWidget(Widget):
|
|||||||
children_by_parent[policy.parent].append(policy)
|
children_by_parent[policy.parent].append(policy)
|
||||||
|
|
||||||
for parent_id, children in children_by_parent.items():
|
for parent_id, children in children_by_parent.items():
|
||||||
children.sort(
|
if getattr(self, "match_type", "Count") == "Count":
|
||||||
key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
|
children.sort(
|
||||||
)
|
key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
|
||||||
|
)
|
||||||
|
else: # Alphabetical
|
||||||
|
children.sort(key=lambda p: p.name.lower())
|
||||||
|
|
||||||
parent_node = node_map.get(parent_id)
|
parent_node = node_map.get(parent_id)
|
||||||
if parent_node:
|
if parent_node:
|
||||||
for policy in children:
|
for policy in children:
|
||||||
@@ -108,12 +221,17 @@ class PolicyTreeWidget(Widget):
|
|||||||
node = parent_node.add(label=label, data=policy)
|
node = parent_node.add(label=label, data=policy)
|
||||||
node_map[policy.groupid] = node
|
node_map[policy.groupid] = node
|
||||||
|
|
||||||
# Add devices (leaf nodes)
|
# Add devices (leaf nodes) - always sort alphabetically
|
||||||
|
devices_by_group = defaultdict(list)
|
||||||
for device in self.devices:
|
for device in self.devices:
|
||||||
group_id = device.groupid
|
devices_by_group[device.groupid].append(device)
|
||||||
|
|
||||||
|
for group_id, devices in devices_by_group.items():
|
||||||
|
devices.sort(key=lambda d: d.hostname.lower()) # Always sort alphabetically
|
||||||
parent_node = node_map.get(group_id)
|
parent_node = node_map.get(group_id)
|
||||||
if parent_node:
|
if parent_node:
|
||||||
parent_node.add(label=device.hostname, data=device)
|
for device in devices:
|
||||||
|
parent_node.add(label=device.hostname, data=device)
|
||||||
|
|
||||||
def _collect_tree_nodes(self, node, all_nodes):
|
def _collect_tree_nodes(self, node, all_nodes):
|
||||||
all_nodes.append(node)
|
all_nodes.append(node)
|
||||||
@@ -133,6 +251,9 @@ class PolicyTreeWidget(Widget):
|
|||||||
node = message.node
|
node = message.node
|
||||||
data = node.data
|
data = node.data
|
||||||
details_pane = self.query_one("#details_pane", Static)
|
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:
|
if self.last_highlighted_node is not None:
|
||||||
original_label = str(self.last_highlighted_node.label).strip()
|
original_label = str(self.last_highlighted_node.label).strip()
|
||||||
@@ -147,6 +268,20 @@ class PolicyTreeWidget(Widget):
|
|||||||
node.set_label(highlighted_label)
|
node.set_label(highlighted_label)
|
||||||
self.last_highlighted_node = node
|
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:
|
if data:
|
||||||
details = "\n".join(
|
details = "\n".join(
|
||||||
f"{key}: {value}" for key, value in data.__dict__.items()
|
f"{key}: {value}" for key, value in data.__dict__.items()
|
||||||
@@ -157,6 +292,13 @@ class PolicyTreeWidget(Widget):
|
|||||||
|
|
||||||
message.stop()
|
message.stop()
|
||||||
|
|
||||||
|
def on_switch_changed(self, event: Switch.Changed):
|
||||||
|
self.match_type = "Alpha" if event.value else "Count"
|
||||||
|
self.query_one("#match_switch_label", Static).update(
|
||||||
|
f"Sort: {self.match_type.capitalize()}"
|
||||||
|
)
|
||||||
|
self._build_tree()
|
||||||
|
|
||||||
def on_input_submitted(self, message: Input.Submitted) -> None:
|
def on_input_submitted(self, message: Input.Submitted) -> None:
|
||||||
self._remove_match_selector()
|
self._remove_match_selector()
|
||||||
|
|
||||||
@@ -229,3 +371,18 @@ class PolicyTreeWidget(Widget):
|
|||||||
option_list.remove()
|
option_list.remove()
|
||||||
except:
|
except:
|
||||||
pass
|
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()
|
||||||
@@ -0,0 +1,870 @@
|
|||||||
|
# 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
|
||||||
|
import os.path
|
||||||
|
import re
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
import dotenv
|
||||||
|
import pandas as pd
|
||||||
|
|
||||||
|
from models.execution import ExecutionHistoryRecord
|
||||||
|
from models.policy import Allowlist, Policy
|
||||||
|
from services.API import AirlockAPIWrapper
|
||||||
|
from utils.configmanager import get_system_list, get_system_value, load_env
|
||||||
|
from utils.selector import Selector
|
||||||
|
from utils.utils import (
|
||||||
|
areYouSure,
|
||||||
|
clear_screen,
|
||||||
|
colorText,
|
||||||
|
formatHTML,
|
||||||
|
get_sanitized_input,
|
||||||
|
locked,
|
||||||
|
open_directory,
|
||||||
|
print_x_wide,
|
||||||
|
regulator,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
dotenv.load_dotenv()
|
||||||
|
|
||||||
|
|
||||||
|
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
|
||||||
|
|
||||||
|
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||||||
|
logger.debug("Prompting for Policies")
|
||||||
|
print(colorText("Please select policy/policies", "white"))
|
||||||
|
selected = Selector.select_objects(policies, allow_multiple, prompt_each=True)
|
||||||
|
|
||||||
|
if selected is None:
|
||||||
|
return []
|
||||||
|
|
||||||
|
# Normalize to always return a list
|
||||||
|
logger.debug("Returning {selected.dict}")
|
||||||
|
return selected if isinstance(selected, list) else [selected]
|
||||||
|
|
||||||
|
|
||||||
|
def selectAllowlists(
|
||||||
|
api: AirlockAPIWrapper, policy=all, allow_multiple=True
|
||||||
|
) -> List[Allowlist]:
|
||||||
|
if policy == "all":
|
||||||
|
allowlists = [
|
||||||
|
Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()
|
||||||
|
]
|
||||||
|
else:
|
||||||
|
allowlists = [
|
||||||
|
Allowlist(**row.to_dict())
|
||||||
|
for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()
|
||||||
|
]
|
||||||
|
logger.debug("Prompting for Allowlist(s)")
|
||||||
|
print(colorText("Please select allowlist(s)", "white"))
|
||||||
|
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
|
||||||
|
|
||||||
|
if selected is None:
|
||||||
|
return []
|
||||||
|
|
||||||
|
# Normalize to always return a list
|
||||||
|
logger.debug(f"Returning {selected}")
|
||||||
|
return selected if isinstance(selected, list) else [selected]
|
||||||
|
|
||||||
|
|
||||||
|
def sortHashes(
|
||||||
|
api: AirlockAPIWrapper, selected_policies: List[Policy], type=[1, 2, 6, 7]
|
||||||
|
):
|
||||||
|
working_dir = load_env("WORKING_DIR")
|
||||||
|
history_days = Selector.select_value(
|
||||||
|
prompt="Enter how many days of history to pull (1-365): ",
|
||||||
|
value_type=int,
|
||||||
|
valid_range=(1, 365),
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(f"{history_days} day selected for history")
|
||||||
|
|
||||||
|
if history_days is None:
|
||||||
|
logging.warning("No history range selected. Aborting.")
|
||||||
|
return
|
||||||
|
|
||||||
|
policy_executions = ExecutionHistoryRecord.from_policies(
|
||||||
|
api, selected_policies, type_=type, history_days=history_days
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(f"Executions contains {policy_executions}")
|
||||||
|
|
||||||
|
enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(
|
||||||
|
api, policy_executions
|
||||||
|
)
|
||||||
|
categorized_executions = (
|
||||||
|
ExecutionHistoryRecord.categorize_executions_by_hash_decision(
|
||||||
|
enriched_executions
|
||||||
|
)
|
||||||
|
)
|
||||||
|
approved, unapproved, needs_review, unknown = (
|
||||||
|
ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions)
|
||||||
|
)
|
||||||
|
|
||||||
|
categories = {
|
||||||
|
"needs_review": needs_review,
|
||||||
|
"approved": approved,
|
||||||
|
"unapproved": unapproved,
|
||||||
|
"leftover": unknown,
|
||||||
|
}
|
||||||
|
|
||||||
|
for label, records in categories.items():
|
||||||
|
if not records:
|
||||||
|
continue # Skip empty or falsy categories
|
||||||
|
|
||||||
|
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
|
||||||
|
html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html"
|
||||||
|
|
||||||
|
# Convert ExecutionHistoryRecord objects to dictionaries
|
||||||
|
df = pd.DataFrame([r.__dict__ for r in records])
|
||||||
|
|
||||||
|
# Optional: flatten hash_obj if needed
|
||||||
|
if not df.empty and "hash_obj" in df.columns:
|
||||||
|
hash_df = df["hash_obj"].apply(lambda h: h.to_dict() if h else {})
|
||||||
|
df = pd.concat([df.drop(columns=["hash_obj"]), hash_df], axis=1)
|
||||||
|
|
||||||
|
# Save to CSV
|
||||||
|
df.to_csv(csv_path, index=False)
|
||||||
|
logger.info(f"Saved {label} executions to {csv_path}")
|
||||||
|
|
||||||
|
# Generate HTML
|
||||||
|
formatHTML(df, html_path)
|
||||||
|
logger.info(f"Generated HTML report at {html_path}")
|
||||||
|
|
||||||
|
|
||||||
|
def buildPathsandPublishers(selected_policies: List[Policy], split):
|
||||||
|
working_dir = load_env("WORKING_DIR")
|
||||||
|
df1 = pd.DataFrame()
|
||||||
|
df2 = pd.DataFrame()
|
||||||
|
all_approved_hashes = pd.DataFrame()
|
||||||
|
path1 = (
|
||||||
|
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
|
||||||
|
)
|
||||||
|
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
|
||||||
|
path_exclusion_constant = get_system_value("PATH_EXCLUSION_CONST", cast_type=int)
|
||||||
|
|
||||||
|
if os.path.exists(path1):
|
||||||
|
df1 = pd.read_csv(path1)
|
||||||
|
else:
|
||||||
|
logger.warning(f"File not found: {path1}")
|
||||||
|
|
||||||
|
if os.path.exists(path2):
|
||||||
|
df2 = pd.read_csv(path2)
|
||||||
|
else:
|
||||||
|
logger.warning(f"File not found: {path2}")
|
||||||
|
|
||||||
|
if df1.empty and df2.empty:
|
||||||
|
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||||
|
all_approved_hashes = pd.DataFrame()
|
||||||
|
logger.debug(all_approved_hashes.head)
|
||||||
|
else:
|
||||||
|
all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
|
||||||
|
if "filename" in all_approved_hashes.columns:
|
||||||
|
all_approved_hashes = all_approved_hashes.sort_values(by="filename")
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
"Warning: 'filename' column not found in concatenated DataFrame."
|
||||||
|
)
|
||||||
|
|
||||||
|
if not all_approved_hashes.empty and path_exclusion_constant:
|
||||||
|
|
||||||
|
primary_path_exclusions = calculatePath(
|
||||||
|
all_approved_hashes,
|
||||||
|
path_exclusion_constant,
|
||||||
|
split,
|
||||||
|
)
|
||||||
|
remaining_hashes = all_approved_hashes[
|
||||||
|
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
|
||||||
|
]
|
||||||
|
secondary_path_exclusions = calculatePath(
|
||||||
|
remaining_hashes, (path_exclusion_constant - 1), split
|
||||||
|
)
|
||||||
|
remaining_hashes = remaining_hashes[
|
||||||
|
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
|
||||||
|
]
|
||||||
|
dataframes = {
|
||||||
|
"all_approved_hashes": all_approved_hashes,
|
||||||
|
"primary_Paths": primary_path_exclusions,
|
||||||
|
"secondary_Paths": secondary_path_exclusions,
|
||||||
|
"hashes_not_approvable_by_path": remaining_hashes,
|
||||||
|
}
|
||||||
|
logger.debug("Preparing to sort dataframes")
|
||||||
|
for name, df in dataframes.items():
|
||||||
|
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||||||
|
if "hashes" in name:
|
||||||
|
df.sort_values(by="filename", inplace=True)
|
||||||
|
else:
|
||||||
|
df.sort_values(by="longestcfp", inplace=True)
|
||||||
|
|
||||||
|
df.to_csv(
|
||||||
|
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv",
|
||||||
|
index=False,
|
||||||
|
)
|
||||||
|
formatHTML(
|
||||||
|
df,
|
||||||
|
f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html",
|
||||||
|
)
|
||||||
|
|
||||||
|
if not all_approved_hashes.empty:
|
||||||
|
# Drop all not signed, only keep unique values
|
||||||
|
publist = all_approved_hashes[
|
||||||
|
all_approved_hashes["publisher"] != "Not Signed"
|
||||||
|
].drop_duplicates(subset=["publisher"])
|
||||||
|
# Remove Bad publisher if somehow they made it this far
|
||||||
|
pattern = regulator(get_system_list("BAD_PUBLISHERS"))
|
||||||
|
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
|
||||||
|
publist = publist[["publisher"]]
|
||||||
|
publist.sort_values(by="publisher", inplace=True)
|
||||||
|
publist.to_csv(
|
||||||
|
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv",
|
||||||
|
index=False,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.debug("Approved Hashes list appears empty")
|
||||||
|
|
||||||
|
|
||||||
|
def buildPreflights(selected_policies: List[Policy]):
|
||||||
|
working_dir = load_env("WORKING_DIR")
|
||||||
|
|
||||||
|
df1 = pd.DataFrame()
|
||||||
|
df2 = pd.DataFrame()
|
||||||
|
approved_hashes = pd.DataFrame()
|
||||||
|
approved_publishers = pd.DataFrame()
|
||||||
|
|
||||||
|
hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv"
|
||||||
|
path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
|
||||||
|
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv"
|
||||||
|
publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv"
|
||||||
|
|
||||||
|
# Read in and combine the two path generations
|
||||||
|
if os.path.exists(path1):
|
||||||
|
df1 = pd.read_csv(path1)
|
||||||
|
else:
|
||||||
|
logger.warning(f"File not found: {path1}")
|
||||||
|
|
||||||
|
if os.path.exists(path2):
|
||||||
|
df2 = pd.read_csv(path2)
|
||||||
|
else:
|
||||||
|
logger.warning(f"File not found: {path2}")
|
||||||
|
|
||||||
|
if df1.empty and df2.empty:
|
||||||
|
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||||
|
approved_paths = pd.DataFrame()
|
||||||
|
else:
|
||||||
|
approved_paths = pd.concat([df1, df2], ignore_index=True)
|
||||||
|
|
||||||
|
approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first")
|
||||||
|
|
||||||
|
# We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions.
|
||||||
|
if os.path.exists(hash):
|
||||||
|
hashes = pd.read_csv(hash)
|
||||||
|
approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])]
|
||||||
|
|
||||||
|
approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first")
|
||||||
|
|
||||||
|
else:
|
||||||
|
logger.warning(f"File not found: {hash}")
|
||||||
|
|
||||||
|
if os.path.exists(publishers):
|
||||||
|
approved_publishers = pd.read_csv(publishers)
|
||||||
|
|
||||||
|
else:
|
||||||
|
logger.warning(f"File not found: {publishers}")
|
||||||
|
|
||||||
|
dataframes = {
|
||||||
|
"approved_paths": approved_paths,
|
||||||
|
"approved_hashes": approved_hashes,
|
||||||
|
"approved_publishers": approved_publishers,
|
||||||
|
}
|
||||||
|
|
||||||
|
for name, df in dataframes.items():
|
||||||
|
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||||||
|
if name == "approved_paths":
|
||||||
|
df.sort_values(by="longestcfp", inplace=True)
|
||||||
|
elif name == "approved_hashes":
|
||||||
|
df.sort_values(by="filename", inplace=True)
|
||||||
|
elif name == "approved_publishers":
|
||||||
|
df.sort_values(by="publisher", inplace=True)
|
||||||
|
|
||||||
|
df.to_csv(
|
||||||
|
f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv",
|
||||||
|
index=False,
|
||||||
|
)
|
||||||
|
formatHTML(
|
||||||
|
df,
|
||||||
|
f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
|
||||||
|
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
|
||||||
|
|
||||||
|
def clean_split(path):
|
||||||
|
if not isinstance(path, (str, bytes, os.PathLike)):
|
||||||
|
return []
|
||||||
|
parts = str(os.path.normpath(path)).split(os.sep)
|
||||||
|
parts = [p for p in parts if p] # Remove empty strings
|
||||||
|
return parts
|
||||||
|
|
||||||
|
# Diagnostic: log any non-string entries
|
||||||
|
non_string_entries = df[
|
||||||
|
~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))
|
||||||
|
]
|
||||||
|
if not non_string_entries.empty:
|
||||||
|
print(f"[WARNING] Non-string entries found in column '{col}':")
|
||||||
|
print(non_string_entries)
|
||||||
|
|
||||||
|
df = df.copy()
|
||||||
|
split_paths = df[col].apply(clean_split)
|
||||||
|
|
||||||
|
if min_files_for_path is not None:
|
||||||
|
df = df[
|
||||||
|
split_paths.apply(lambda parts: len(parts) >= min_files_for_path)
|
||||||
|
].copy()
|
||||||
|
split_paths = split_paths[df.index]
|
||||||
|
|
||||||
|
df["group_key"] = split_paths.apply(
|
||||||
|
lambda parts: os.sep.join(parts[:path_exclusion_constant])
|
||||||
|
)
|
||||||
|
grouped = df.groupby("group_key")
|
||||||
|
new_rows = []
|
||||||
|
|
||||||
|
for _, group_df in grouped:
|
||||||
|
paths = group_df[col].tolist()
|
||||||
|
split_parts = [clean_split(p) for p in paths]
|
||||||
|
|
||||||
|
def longest_common_prefix(paths):
|
||||||
|
if not paths:
|
||||||
|
return []
|
||||||
|
prefix = paths[0]
|
||||||
|
for path in paths[1:]:
|
||||||
|
prefix = [a for a, b in zip(prefix, path) if a == b]
|
||||||
|
if not prefix:
|
||||||
|
break
|
||||||
|
return prefix
|
||||||
|
|
||||||
|
common_prefix = longest_common_prefix(split_parts)
|
||||||
|
prefix_str = os.sep.join(common_prefix)
|
||||||
|
|
||||||
|
for i, parts in enumerate(split_parts):
|
||||||
|
filename = parts[-1]
|
||||||
|
middle = (
|
||||||
|
os.sep.join(parts[len(common_prefix) : -1])
|
||||||
|
if len(parts) > len(common_prefix) + 1
|
||||||
|
else ""
|
||||||
|
)
|
||||||
|
row = group_df.iloc[i].copy()
|
||||||
|
row["longestcfp"] = prefix_str
|
||||||
|
row["middle"] = middle
|
||||||
|
row["filename_only"] = filename
|
||||||
|
row["file_extension"] = os.path.splitext(filename)[1].lower()
|
||||||
|
new_rows.append(row)
|
||||||
|
|
||||||
|
return pd.DataFrame(new_rows).drop(columns=["group_key"])
|
||||||
|
|
||||||
|
|
||||||
|
def calculatePath(approved_hashes, path_exclusion_constant, split):
|
||||||
|
if split:
|
||||||
|
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
|
||||||
|
else:
|
||||||
|
dfs_by_policy = [approved_hashes]
|
||||||
|
|
||||||
|
badpathparts = get_system_list("BAD_PATH_PARTS")
|
||||||
|
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
|
||||||
|
|
||||||
|
processed_dfs = []
|
||||||
|
|
||||||
|
for df in dfs_by_policy:
|
||||||
|
haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename")
|
||||||
|
haslcp = haslcp.drop_duplicates()
|
||||||
|
|
||||||
|
forbidden = regulator(badpathparts, True)
|
||||||
|
forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
|
||||||
|
|
||||||
|
logger.debug("Removing forbidden filepaths for path exceptions")
|
||||||
|
print(colorText("Removing forbidden filepaths for path exceptions", "green"))
|
||||||
|
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
|
||||||
|
|
||||||
|
lcp_not_forbidden_review = lcp_not_forbidden[
|
||||||
|
[
|
||||||
|
"policyname",
|
||||||
|
"longestcfp",
|
||||||
|
"middle",
|
||||||
|
"filename_only",
|
||||||
|
"file_extension",
|
||||||
|
"sha256",
|
||||||
|
]
|
||||||
|
]
|
||||||
|
|
||||||
|
unique_sha_counts = (
|
||||||
|
lcp_not_forbidden_review.groupby("longestcfp")["sha256"]
|
||||||
|
.nunique()
|
||||||
|
.reset_index()
|
||||||
|
)
|
||||||
|
unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]
|
||||||
|
|
||||||
|
lcp_not_forbidden_review = lcp_not_forbidden_review.merge(
|
||||||
|
unique_sha_counts, on="longestcfp", how="left"
|
||||||
|
)
|
||||||
|
lcp_not_forbidden_review = lcp_not_forbidden_review[
|
||||||
|
lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path
|
||||||
|
]
|
||||||
|
processed_dfs.append(lcp_not_forbidden_review)
|
||||||
|
|
||||||
|
pathExclusions = pd.concat(processed_dfs, ignore_index=True)
|
||||||
|
|
||||||
|
return pathExclusions
|
||||||
|
|
||||||
|
|
||||||
|
def testChange(selected_policies, destination_policy, destination_allowlist):
|
||||||
|
working_dir = load_env("WORKING_DIR")
|
||||||
|
|
||||||
|
logger.info("These path exclusions would be added to:")
|
||||||
|
logger.info(destination_policy)
|
||||||
|
|
||||||
|
pathexclusions = pd.read_csv(
|
||||||
|
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
|
||||||
|
)
|
||||||
|
hashes = pd.read_csv(
|
||||||
|
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
|
||||||
|
)
|
||||||
|
|
||||||
|
unique_combinations = pathexclusions[
|
||||||
|
["longestcfp", "file_extension"]
|
||||||
|
].drop_duplicates()
|
||||||
|
|
||||||
|
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
|
||||||
|
processed_paths = [
|
||||||
|
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
|
||||||
|
for path, ext in unique_combinations.itertuples(index=False, name=None)
|
||||||
|
]
|
||||||
|
|
||||||
|
for path in processed_paths:
|
||||||
|
logger.info(path)
|
||||||
|
|
||||||
|
print(colorText("These publishers would added", "yellow"))
|
||||||
|
processed_publishers = []
|
||||||
|
if os.path.exists(
|
||||||
|
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
|
||||||
|
):
|
||||||
|
publishers = pd.read_csv(
|
||||||
|
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
|
||||||
|
)
|
||||||
|
if publishers.empty:
|
||||||
|
print(colorText("The publishers list is empty.", "red"))
|
||||||
|
else:
|
||||||
|
processed_publishers = (
|
||||||
|
publishers[publishers["publisher"] != "Not Signed"]["publisher"]
|
||||||
|
.drop_duplicates()
|
||||||
|
.tolist()
|
||||||
|
)
|
||||||
|
for publisher in processed_publishers:
|
||||||
|
print(publisher)
|
||||||
|
|
||||||
|
print(colorText("These hashes would be added to:", "yellow"))
|
||||||
|
print(destination_allowlist)
|
||||||
|
|
||||||
|
processed_hashes = hashes["sha256"].unique().tolist()
|
||||||
|
print_x_wide(processed_hashes, 3)
|
||||||
|
|
||||||
|
return processed_paths, processed_hashes, processed_publishers
|
||||||
|
|
||||||
|
|
||||||
|
def menu_policy_enforce(
|
||||||
|
api: AirlockAPIWrapper,
|
||||||
|
): # TODO Need to clean up 6 and 7 into functions
|
||||||
|
selected_policies = []
|
||||||
|
destination_policy = []
|
||||||
|
destination_allowlist = []
|
||||||
|
processed_paths = []
|
||||||
|
processed_hashes = []
|
||||||
|
processed_publishers = []
|
||||||
|
working_dir = load_env("WORKING_DIR")
|
||||||
|
|
||||||
|
while True:
|
||||||
|
printEnforceChecklist(
|
||||||
|
selected_policies, destination_policy, destination_allowlist
|
||||||
|
)
|
||||||
|
choice = get_sanitized_input("\nEnter your choice: ")
|
||||||
|
|
||||||
|
if choice == "1":
|
||||||
|
clear_screen()
|
||||||
|
selected_policies = selectPolicies(api, True)
|
||||||
|
|
||||||
|
elif choice == "2":
|
||||||
|
clear_screen()
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
"Please choose destination_name Policy for Path Exclusions", "white"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
destination_policy = selectPolicies(api, False)
|
||||||
|
|
||||||
|
print(colorText("Please choose Allowlist for Hashes", "white"))
|
||||||
|
|
||||||
|
destination_allowlist = selectAllowlists(api, destination_policy, False)
|
||||||
|
|
||||||
|
elif choice == "3":
|
||||||
|
clear_screen()
|
||||||
|
sortHashes(
|
||||||
|
api,
|
||||||
|
selected_policies,
|
||||||
|
type=[1, 2, 6, 7],
|
||||||
|
)
|
||||||
|
|
||||||
|
elif choice == "4":
|
||||||
|
clear_screen()
|
||||||
|
if os.path.exists(
|
||||||
|
f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"
|
||||||
|
):
|
||||||
|
buildPathsandPublishers(selected_policies, False)
|
||||||
|
else:
|
||||||
|
print(
|
||||||
|
"File not found. Please make sure it's saved correctly and try again."
|
||||||
|
)
|
||||||
|
|
||||||
|
elif choice == "5":
|
||||||
|
clear_screen()
|
||||||
|
if os.path.exists(
|
||||||
|
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
|
||||||
|
) and os.path.exists(
|
||||||
|
f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
|
||||||
|
):
|
||||||
|
buildPreflights(selected_policies)
|
||||||
|
else:
|
||||||
|
print(
|
||||||
|
"File not found. Please make sure it's saved correctly and try again."
|
||||||
|
)
|
||||||
|
|
||||||
|
elif choice == "6":
|
||||||
|
clear_screen()
|
||||||
|
if (
|
||||||
|
os.path.exists(
|
||||||
|
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
|
||||||
|
)
|
||||||
|
and os.path.exists(
|
||||||
|
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
|
||||||
|
)
|
||||||
|
and destination_policy
|
||||||
|
and destination_allowlist
|
||||||
|
):
|
||||||
|
processed_paths, processed_hashes, processed_publishers = testChange(
|
||||||
|
selected_policies, destination_policy, destination_allowlist
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Log which condition(s) failed
|
||||||
|
missing_items = []
|
||||||
|
if not os.path.exists(
|
||||||
|
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
|
||||||
|
):
|
||||||
|
missing_items.append("approved_paths.csv not found")
|
||||||
|
if not os.path.exists(
|
||||||
|
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
|
||||||
|
):
|
||||||
|
missing_items.append("approved_hashes.csv not found")
|
||||||
|
if not destination_policy:
|
||||||
|
missing_items.append("destination_policy is empty or None")
|
||||||
|
if not destination_allowlist:
|
||||||
|
missing_items.append("destination_allowlist is empty or None")
|
||||||
|
|
||||||
|
logger.error("Preflight check failed due to the following:")
|
||||||
|
for item in missing_items:
|
||||||
|
logger.error(f" - {item}")
|
||||||
|
|
||||||
|
elif choice == "7":
|
||||||
|
clear_screen()
|
||||||
|
areYouSure()
|
||||||
|
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||||
|
if (
|
||||||
|
processed_paths
|
||||||
|
and processed_hashes
|
||||||
|
and processed_publishers
|
||||||
|
and destination_policy
|
||||||
|
and destination_allowlist
|
||||||
|
and confirmation.strip() == "I AGREE"
|
||||||
|
):
|
||||||
|
print(colorText("Proceeding with the code...", "yellow"))
|
||||||
|
api.hash_add_to_allowlist(
|
||||||
|
destination_allowlist[0].applicationid, processed_hashes
|
||||||
|
)
|
||||||
|
api.policy_add_path_exclusions(
|
||||||
|
destination_policy[0].groupid, processed_paths
|
||||||
|
)
|
||||||
|
if processed_publishers:
|
||||||
|
api.policy_add_publishers(
|
||||||
|
destination_policy[0].groupid, processed_publishers
|
||||||
|
)
|
||||||
|
|
||||||
|
locked()
|
||||||
|
|
||||||
|
else:
|
||||||
|
logger.error("Confirmation block failed. Reasons:")
|
||||||
|
if not processed_publishers or processed_hashes or processed_paths:
|
||||||
|
logger.error(" - Test not performed.")
|
||||||
|
if not destination_policy:
|
||||||
|
logger.error(" - `destination_policy` is missing or invalid.")
|
||||||
|
if not destination_allowlist:
|
||||||
|
logger.error(" - `destination_allowlist` is missing or invalid.")
|
||||||
|
if confirmation.strip() != "I AGREE":
|
||||||
|
logger.error(
|
||||||
|
" - User did not confirm with 'I AGREE'. Received: '%s'",
|
||||||
|
confirmation.strip(),
|
||||||
|
)
|
||||||
|
|
||||||
|
elif choice.upper() == "F":
|
||||||
|
open_directory(working_dir)
|
||||||
|
elif choice.upper() == "B":
|
||||||
|
break
|
||||||
|
|
||||||
|
else:
|
||||||
|
print(colorText("Invalid choice. Please try again.", "red"))
|
||||||
|
|
||||||
|
|
||||||
|
def section_header(title):
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
"\n --------------------------------------------------------------------",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(colorText(f" ------------- {title} -------------", "cyan"))
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
" --------------------------------------------------------------------",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
|
||||||
|
working_dir = load_env("WORKING_DIR")
|
||||||
|
section_header("Prepare to Enforce Policy")
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
"\nSequentially follow these steps to prepare a policy for enforcement:",
|
||||||
|
"white",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Step 1: Originating Policies
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
"\n1. Choose which policy or policies to gather execution info from", "cyan"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if not selected_policies:
|
||||||
|
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"))
|
||||||
|
|
||||||
|
# Step 2: Destination Policy and Allowlist
|
||||||
|
print(
|
||||||
|
colorText("2. Choose the destination policy and associated allowlist", "cyan")
|
||||||
|
)
|
||||||
|
if destination_policy:
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
f" [✅] {destination_policy[0].name} has been selected as the destination policy",
|
||||||
|
"green",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
print(colorText(" [âŒ] No destination policy has been chosen", "red"))
|
||||||
|
|
||||||
|
if destination_allowlist:
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
f" [✅] {destination_allowlist[0].name} has been selected as allowlist",
|
||||||
|
"green",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
print(colorText(" [âŒ] No allowlist has been chosen", "red"))
|
||||||
|
|
||||||
|
# Step 3: Data Preparation
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if selected_policies:
|
||||||
|
policy_id = selected_policies[0].name
|
||||||
|
review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv"
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
(
|
||||||
|
" [✅] Data has been fetched"
|
||||||
|
if os.path.exists(review_path)
|
||||||
|
else " [âŒ] Data has not been fetched"
|
||||||
|
),
|
||||||
|
"green" if os.path.exists(review_path) else "red",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
" [âŒ] No policies selected, cannot check data fetch status", "red"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Step 4: Manual Review
|
||||||
|
print(colorText("4. Manually review the files:", "cyan"))
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
" Remove the rows containing hashes you do not approve of", "cyan"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
" This will start the process to generate possible filepath approvals",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
if selected_policies:
|
||||||
|
policy_id = selected_policies[0].name
|
||||||
|
approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv"
|
||||||
|
second_review_path = (
|
||||||
|
f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv"
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
(
|
||||||
|
" [✅] Reviewed hashes have been loaded"
|
||||||
|
if os.path.exists(approved_path)
|
||||||
|
else " [âŒ] Reviewed hashes have not been loaded"
|
||||||
|
),
|
||||||
|
"green" if os.path.exists(approved_path) else "red",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
(
|
||||||
|
" [✅] Path review list created"
|
||||||
|
if os.path.exists(second_review_path)
|
||||||
|
else " [âŒ] Path review list has not been created"
|
||||||
|
),
|
||||||
|
"green" if os.path.exists(second_review_path) else "red",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
" [âŒ] No policies selected, cannot check reviewed hashes or path list",
|
||||||
|
"red",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Step 5: Path Review
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
" Remove the rows containing path exclusions or publishers you do not approve of.",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
f" When complete, save the files to {working_dir}\\data\\Approved",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(" Choose this option when done to build your preflights", "cyan")
|
||||||
|
)
|
||||||
|
|
||||||
|
if selected_policies:
|
||||||
|
policy_id = selected_policies[0].name
|
||||||
|
reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv"
|
||||||
|
preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv"
|
||||||
|
preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv"
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
(
|
||||||
|
" [✅] Reviewed path list detected"
|
||||||
|
if os.path.exists(reviewed_path)
|
||||||
|
else " [âŒ] Path review list has not been detected"
|
||||||
|
),
|
||||||
|
"green" if os.path.exists(reviewed_path) else "red",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
preflight_ready = os.path.exists(preflight_paths) and os.path.exists(
|
||||||
|
preflight_hashes
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
(
|
||||||
|
" [✅] Preflight Path Exclusion List has been generated"
|
||||||
|
if preflight_ready
|
||||||
|
else " [âŒ] Preflight Path Exclusion List has not been generated"
|
||||||
|
),
|
||||||
|
"green" if preflight_ready else "red",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
" [âŒ] No policies selected, cannot check preflight status", "red"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Final Steps
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
"6. Test ------------------------------------------------------", "cyan"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
" Prints to console the changes that would be made, must be done to proceed. ",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
"7. Liftoff ------------------------------------------------------", "cyan"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
colorText(
|
||||||
|
" Apply path exclusions and approved publishers to selected policy",
|
||||||
|
"cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(colorText(" Apply approved hashes to allowlist", "cyan"))
|
||||||
|
|
||||||
|
# Utility Options
|
||||||
|
print(colorText("F. Open Working Directory", "cyan"))
|
||||||
|
print(colorText("B. Back", "cyan"))
|
||||||
@@ -1,3 +1,18 @@
|
|||||||
|
# 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 logging
|
||||||
|
|
||||||
from textual.containers import Horizontal, Vertical
|
from textual.containers import Horizontal, Vertical
|
||||||
@@ -59,15 +74,6 @@ class ResultsDisplay(Widget):
|
|||||||
margin-top: 1;
|
margin-top: 1;
|
||||||
width: 100%;
|
width: 100%;
|
||||||
}
|
}
|
||||||
|
|
||||||
#button_row {
|
|
||||||
height: auto;
|
|
||||||
margin: 1 0 0 0;
|
|
||||||
}
|
|
||||||
|
|
||||||
#back_button {
|
|
||||||
width: 1fr;
|
|
||||||
}
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
class CopySuccess(Message):
|
class CopySuccess(Message):
|
||||||
@@ -95,9 +101,9 @@ class ResultsDisplay(Widget):
|
|||||||
|
|
||||||
def compose(self):
|
def compose(self):
|
||||||
with Vertical(id="results_screen"):
|
with Vertical(id="results_screen"):
|
||||||
yield Header(show_clock=True, icon="⚙")
|
yield Header(show_clock=True, icon="⚙️")
|
||||||
# Title
|
# Title
|
||||||
title = Static(f"📊 {self.operation} - Results", id="results_title")
|
title = Static(f"{self.operation} - Results", id="results_title")
|
||||||
yield title
|
yield title
|
||||||
|
|
||||||
# Two-column layout
|
# Two-column layout
|
||||||
@@ -107,7 +113,7 @@ class ResultsDisplay(Widget):
|
|||||||
yield Static("✅ Successful", id="success_label")
|
yield Static("✅ Successful", id="success_label")
|
||||||
yield Static(self.successful_results, id="success_results")
|
yield Static(self.successful_results, id="success_results")
|
||||||
yield Button(
|
yield Button(
|
||||||
"📋✅ Copy Success List",
|
"Copy Success List",
|
||||||
id="copy_success",
|
id="copy_success",
|
||||||
classes="copy_button",
|
classes="copy_button",
|
||||||
)
|
)
|
||||||
@@ -117,15 +123,11 @@ class ResultsDisplay(Widget):
|
|||||||
yield Static("❌ Failed", id="failure_label")
|
yield Static("❌ Failed", id="failure_label")
|
||||||
yield Static(self.unsuccessful_results, id="failure_results")
|
yield Static(self.unsuccessful_results, id="failure_results")
|
||||||
yield Button(
|
yield Button(
|
||||||
"📋❌ Copy Failure List",
|
"Copy Failure List",
|
||||||
id="copy_failure",
|
id="copy_failure",
|
||||||
classes="copy_button",
|
classes="copy_button",
|
||||||
)
|
)
|
||||||
|
|
||||||
# Back Button
|
|
||||||
with Horizontal(id="button_row"):
|
|
||||||
back_button = Button("← Back", id="back_button")
|
|
||||||
yield back_button
|
|
||||||
yield Footer()
|
yield Footer()
|
||||||
|
|
||||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||||
@@ -138,18 +140,18 @@ class ResultsDisplay(Widget):
|
|||||||
|
|
||||||
pyperclip.copy(str(success_widget.renderable))
|
pyperclip.copy(str(success_widget.renderable))
|
||||||
self.app.notify(
|
self.app.notify(
|
||||||
"✅ Success list copied to clipboard!",
|
"Success list copied to clipboard!",
|
||||||
severity="information",
|
severity="information",
|
||||||
timeout=2,
|
timeout=2,
|
||||||
)
|
)
|
||||||
self.post_message(self.CopySuccess())
|
self.post_message(self.CopySuccess())
|
||||||
except ImportError:
|
except ImportError:
|
||||||
self.app.notify(
|
self.app.notify(
|
||||||
"âš ï¸ pyperclip not installed. Run: pip install pyperclip",
|
"❌ pyperclip not installed. Run: pip install pyperclip",
|
||||||
severity="warning",
|
severity="warning",
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.app.notify(f"⌠Failed to copy: {str(e)}", severity="error")
|
self.app.notify(f"¢ Failed to copy: {str(e)}", severity="error")
|
||||||
event.stop()
|
event.stop()
|
||||||
|
|
||||||
elif btn_id == "copy_failure":
|
elif btn_id == "copy_failure":
|
||||||
@@ -159,20 +161,16 @@ class ResultsDisplay(Widget):
|
|||||||
|
|
||||||
pyperclip.copy(str(failure_widget.renderable))
|
pyperclip.copy(str(failure_widget.renderable))
|
||||||
self.app.notify(
|
self.app.notify(
|
||||||
"✅ Failure list copied to clipboard!",
|
"Failure list copied to clipboard!",
|
||||||
severity="information",
|
severity="information",
|
||||||
timeout=2,
|
timeout=2,
|
||||||
)
|
)
|
||||||
self.post_message(self.CopyFailure())
|
self.post_message(self.CopyFailure())
|
||||||
except ImportError:
|
except ImportError:
|
||||||
self.app.notify(
|
self.app.notify(
|
||||||
"âš ï¸ pyperclip not installed. Run: pip install pyperclip",
|
"❌ pyperclip not installed. Run: pip install pyperclip",
|
||||||
severity="warning",
|
severity="warning",
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.app.notify(f"⌠Failed to copy: {str(e)}", severity="error")
|
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
|
||||||
event.stop()
|
|
||||||
|
|
||||||
elif btn_id == "back_button":
|
|
||||||
self.app.pop_screen()
|
|
||||||
event.stop()
|
event.stop()
|
||||||
@@ -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()
|
||||||
@@ -1,3 +1,4 @@
|
|||||||
/target
|
/target
|
||||||
build.sh
|
build.sh
|
||||||
pythontest.py
|
pythontest.py
|
||||||
|
changelog.md
|
||||||
Generated
+1239
-336
File diff suppressed because it is too large
Load Diff
+16
-6
@@ -1,21 +1,31 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "airlock_libs"
|
name = "airlock_libs"
|
||||||
version = "2.0.0"
|
version = "7.1.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
|
||||||
[lib]
|
|
||||||
crate-type = ["cdylib"]
|
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
chrono = "0.4.42"
|
chrono = "0.4.42"
|
||||||
indicatif = "0.18.2"
|
indicatif = "0.18.2"
|
||||||
mongodb = "3.3.0"
|
mongodb = "3.3.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"] }
|
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 = "1.0.228"
|
||||||
serde-pyobject = "0.8.0"
|
serde-pyobject = "0.8.0"
|
||||||
serde_json = "1.0.145"
|
serde_json = "1.0.145"
|
||||||
tokio = { version = "1.48.0", features = ["full"] }
|
tokio = { version = "1.48.0", features = ["full"] }
|
||||||
|
tonic = { version = "0.12.3", features = ["tls-roots"] }
|
||||||
|
tracing = "0.1.41"
|
||||||
|
tracing-subscriber = "0.3.20"
|
||||||
|
tracing-opentelemetry = "0.32.0"
|
||||||
|
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]
|
[package.metadata.maturin]
|
||||||
generate-abi-stubs = true
|
generate-abi-stubs = true
|
||||||
@@ -28,4 +38,4 @@ codegen-units = 1
|
|||||||
panic = 'abort'
|
panic = 'abort'
|
||||||
strip = true
|
strip = true
|
||||||
debug-assertions = false
|
debug-assertions = false
|
||||||
overflow-checks = false
|
overflow-checks = true
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ build-backend = "maturin"
|
|||||||
|
|
||||||
[project]
|
[project]
|
||||||
name = "airlock_libs"
|
name = "airlock_libs"
|
||||||
version = "2.0.0"
|
version = "7.1.0"
|
||||||
description = "Airlock Digital API Wrapper"
|
description = "Airlock Digital API Wrapper"
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
license = { text = "AGPL-3.0-only" }
|
license = { text = "AGPL-3.0-only" }
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
use pyo3::prelude::*;
|
use pyo3::prelude::*;
|
||||||
mod services;
|
pub mod modules;
|
||||||
|
pub mod prelude;
|
||||||
|
pub mod services;
|
||||||
#[pymodule]
|
#[pymodule]
|
||||||
fn airlock_libs(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
|
fn airlock_libs(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
|
||||||
m.add_function(wrap_pyfunction!(services::pull_policy_exec_histories, py)?)?;
|
m.add_function(wrap_pyfunction!(services::pull_policy_exec_histories, py)?)?;
|
||||||
|
|||||||
@@ -0,0 +1,131 @@
|
|||||||
|
use crate::prelude::*;
|
||||||
|
use crate::services::get_base_directory;
|
||||||
|
#[allow(non_snake_case)]
|
||||||
|
#[derive(Deserialize, Debug)]
|
||||||
|
pub struct TelemetryConfig {
|
||||||
|
pub TELEMETRY: bool,
|
||||||
|
pub TELEM_URL: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TelemetryConfig {
|
||||||
|
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 {
|
||||||
|
TELEMETRY: false,
|
||||||
|
TELEM_URL: None,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
match fs::read_to_string(&cfg_path) {
|
||||||
|
Ok(contents) => serde_json::from_str::<Self>(&contents).unwrap_or(Self {
|
||||||
|
TELEMETRY: false,
|
||||||
|
TELEM_URL: None,
|
||||||
|
}),
|
||||||
|
Err(_) => Self {
|
||||||
|
TELEMETRY: false,
|
||||||
|
TELEM_URL: None,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize, Serialize)]
|
||||||
|
pub struct ApiResponse {
|
||||||
|
pub(crate) error: String,
|
||||||
|
pub(crate) response: ExecHistories,
|
||||||
|
}
|
||||||
|
#[derive(Debug, Deserialize, Serialize)]
|
||||||
|
pub struct ExecHistories {
|
||||||
|
pub(crate) exechistories: Vec<Group>,
|
||||||
|
}
|
||||||
|
#[derive(Debug, Deserialize, Serialize, Clone)]
|
||||||
|
pub struct Group {
|
||||||
|
pub(crate) checkpoint: String,
|
||||||
|
#[serde(rename = "type")]
|
||||||
|
pub(crate) exectype: u8,
|
||||||
|
pub(crate) username: String,
|
||||||
|
pub(crate) hostname: String,
|
||||||
|
pub(crate) netdomain: String,
|
||||||
|
pub(crate) filename: String,
|
||||||
|
pub(crate) ppolicy: String,
|
||||||
|
pub(crate) policyname: String,
|
||||||
|
pub(crate) policyver: String,
|
||||||
|
pub(crate) commandline: String,
|
||||||
|
pub(crate) publisher: String,
|
||||||
|
pub(crate) pprocess: String,
|
||||||
|
pub(crate) gprocess: String,
|
||||||
|
pub(crate) sha256: String,
|
||||||
|
pub(crate) datetime: String,
|
||||||
|
pub(crate) md5: String,
|
||||||
|
pub(crate) sha128: String,
|
||||||
|
pub(crate) sha384: String,
|
||||||
|
pub(crate) sha512: String,
|
||||||
|
pub(crate) ip: String,
|
||||||
|
pub(crate) localip: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct PyData {
|
||||||
|
pub headers: reqwest::header::HeaderMap,
|
||||||
|
pub base_url: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let base_url = obj.getattr(py, "base_url").unwrap().to_string();
|
||||||
|
Self {
|
||||||
|
headers: header_map,
|
||||||
|
base_url,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct SkipBack;
|
||||||
|
|
||||||
|
impl SkipBack {
|
||||||
|
pub fn find_checkpoint(days: i64) -> ObjectId {
|
||||||
|
let date_days_ago = Local::now() - Duration::days(days);
|
||||||
|
let timestamp = date_days_ago.timestamp() as u32;
|
||||||
|
let mut hex_timestamp = String::new();
|
||||||
|
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap();
|
||||||
|
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
|
||||||
|
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
pub mod datatypes;
|
||||||
@@ -0,0 +1,32 @@
|
|||||||
|
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::GlobalTracerProvider;
|
||||||
|
pub use opentelemetry::trace::noop::NoopTracerProvider;
|
||||||
|
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 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,
|
||||||
|
fmt::Write,
|
||||||
|
fs::{self, File},
|
||||||
|
io::{Read, Seek, SeekFrom},
|
||||||
|
path::PathBuf,
|
||||||
|
str::FromStr,
|
||||||
|
};
|
||||||
|
pub use tonic::transport::{Channel, ClientTlsConfig};
|
||||||
+255
-205
@@ -1,234 +1,289 @@
|
|||||||
use chrono::{Duration, Local, NaiveDate};
|
use crate::modules::datatypes::*;
|
||||||
use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
|
use crate::prelude::*;
|
||||||
use mongodb::bson::oid::ObjectId;
|
|
||||||
use pyo3::{prelude::*, types::PyString};
|
|
||||||
use reqwest::{
|
|
||||||
Client,
|
|
||||||
header::{HeaderMap, HeaderName, HeaderValue},
|
|
||||||
};
|
|
||||||
use serde::{Deserialize, Serialize};
|
|
||||||
use serde_json::Value;
|
|
||||||
use std::{
|
|
||||||
collections::HashMap,
|
|
||||||
env,
|
|
||||||
fmt::Write,
|
|
||||||
fs::{self, File},
|
|
||||||
io::{Read, Seek, SeekFrom},
|
|
||||||
path::PathBuf,
|
|
||||||
str::FromStr,
|
|
||||||
};
|
|
||||||
|
|
||||||
#[derive(Debug, Deserialize, Serialize)]
|
|
||||||
struct ApiResponse {
|
|
||||||
error: String,
|
|
||||||
response: ExecHistories,
|
|
||||||
}
|
|
||||||
#[derive(Debug, Deserialize, Serialize)]
|
|
||||||
struct ExecHistories {
|
|
||||||
exechistories: Vec<Group>,
|
|
||||||
}
|
|
||||||
#[derive(Debug, Deserialize, Serialize, Clone)]
|
|
||||||
struct Group {
|
|
||||||
checkpoint: String,
|
|
||||||
#[serde(rename = "type")]
|
|
||||||
exectype: u8,
|
|
||||||
username: String,
|
|
||||||
hostname: String,
|
|
||||||
netdomain: String,
|
|
||||||
filename: String,
|
|
||||||
ppolicy: String,
|
|
||||||
policyname: String,
|
|
||||||
policyver: String,
|
|
||||||
commandline: String,
|
|
||||||
publisher: String,
|
|
||||||
pprocess: String,
|
|
||||||
gprocess: String,
|
|
||||||
sha256: String,
|
|
||||||
datetime: String,
|
|
||||||
md5: String,
|
|
||||||
sha128: String,
|
|
||||||
sha384: String,
|
|
||||||
sha512: String,
|
|
||||||
ip: String,
|
|
||||||
localip: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[pyfunction]
|
#[pyfunction]
|
||||||
pub fn pull_policy_exec_histories(
|
pub fn pull_policy_exec_histories(
|
||||||
py: Python<'_>,
|
py: Python<'_>,
|
||||||
py_self: Py<PyAny>,
|
py_self: Py<PyAny>,
|
||||||
policy_names: String,
|
policy_names: Option<String>,
|
||||||
exec_types: String,
|
exec_types: String,
|
||||||
days: i64,
|
days: i64,
|
||||||
) -> Py<PyString> {
|
) -> Py<PyString> {
|
||||||
let file_path: PathBuf = format!(
|
println!();
|
||||||
"{}\\cache\\chunkinator.json",
|
let data: PyData = PyData::extract_data(py, &py_self);
|
||||||
get_base_directory().display()
|
let headers: HeaderMap = data.headers;
|
||||||
)
|
let base_url: String = data.base_url;
|
||||||
.into();
|
let handle: thread::JoinHandle<String> = std::thread::spawn(move || {
|
||||||
let writeable_filepath = file_path.clone();
|
let rt: tokio::runtime::Runtime = match tokio::runtime::Runtime::new() {
|
||||||
if !file_path.exists() {
|
Ok(rt) => rt,
|
||||||
if let Some(parent_dir) = file_path.parent()
|
Err(e) => {
|
||||||
&& !parent_dir.exists()
|
println!("Failed to build Tokio Runtime: {:?}", e);
|
||||||
{
|
std::process::abort();
|
||||||
fs::create_dir_all(parent_dir).unwrap();
|
}
|
||||||
}
|
|
||||||
fs::File::create(file_path).unwrap();
|
|
||||||
}
|
|
||||||
let data = ApiResponse {
|
|
||||||
error: "Success".to_string(),
|
|
||||||
response: ExecHistories {
|
|
||||||
exechistories: vec![],
|
|
||||||
},
|
|
||||||
};
|
|
||||||
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
|
|
||||||
fs::write(writeable_filepath.clone(), data_write).unwrap();
|
|
||||||
let mut checkpoint_number: String = skipback(days).to_string();
|
|
||||||
let multi_progress = MultiProgress::new();
|
|
||||||
multi_progress.set_draw_target(ProgressDrawTarget::stdout());
|
|
||||||
let progress_bar = multi_progress.add(ProgressBar::new(100));
|
|
||||||
progress_bar.set_style(
|
|
||||||
ProgressStyle::default_bar()
|
|
||||||
.template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
|
|
||||||
.unwrap(),
|
|
||||||
);
|
|
||||||
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
|
|
||||||
let client = build_client(py, &py_self);
|
|
||||||
let api: Py<PyAny> = py_self;
|
|
||||||
let cutoff = Local::now().naive_local() - Duration::days(days);
|
|
||||||
let mut f = File::open(&writeable_filepath).unwrap();
|
|
||||||
loop {
|
|
||||||
f.seek(SeekFrom::Start(0)).unwrap();
|
|
||||||
let execution_histories = history_logging(
|
|
||||||
py,
|
|
||||||
&api,
|
|
||||||
&exec_types,
|
|
||||||
&checkpoint_number,
|
|
||||||
&policy_names,
|
|
||||||
&client,
|
|
||||||
);
|
|
||||||
let parsed_responses = execution_histories.response.exechistories;
|
|
||||||
if parsed_responses.is_empty() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() {
|
|
||||||
let mut contents = String::new();
|
|
||||||
f.read_to_string(&mut contents).unwrap();
|
|
||||||
let existing_data: ApiResponse =
|
|
||||||
serde_json::from_str(&contents).unwrap_or(ApiResponse {
|
|
||||||
error: "Success".to_string(),
|
|
||||||
response: ExecHistories {
|
|
||||||
exechistories: vec![],
|
|
||||||
},
|
|
||||||
});
|
|
||||||
existing_data
|
|
||||||
.response
|
|
||||||
.exechistories
|
|
||||||
.into_iter()
|
|
||||||
.map(|entry| {
|
|
||||||
(
|
|
||||||
(
|
|
||||||
entry.sha256.clone(),
|
|
||||||
entry.filename.clone(),
|
|
||||||
entry.hostname.clone(),
|
|
||||||
),
|
|
||||||
entry,
|
|
||||||
)
|
|
||||||
})
|
|
||||||
.collect()
|
|
||||||
} else {
|
|
||||||
HashMap::new()
|
|
||||||
};
|
};
|
||||||
for (index, executions) in parsed_responses.iter().enumerate() {
|
let tracer_provider = rt.block_on(async { TelemetryConfig::init_tracer() });
|
||||||
if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
|
global::set_tracer_provider(tracer_provider.clone());
|
||||||
continue;
|
let tracer: global::BoxedTracer = global::tracer("tracer");
|
||||||
|
let _cx: Context = Context::new();
|
||||||
|
let file_path: PathBuf = format!(
|
||||||
|
"{}\\cache\\chunkinator.json",
|
||||||
|
get_base_directory().display()
|
||||||
|
)
|
||||||
|
.into();
|
||||||
|
if !&file_path.exists() {
|
||||||
|
if let Some(parent_dir) = &file_path.parent()
|
||||||
|
&& !parent_dir.exists()
|
||||||
|
{
|
||||||
|
match fs::create_dir_all(parent_dir) {
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(e) => {
|
||||||
|
println!("Failed to Create Directory {:?}: {}", parent_dir, e);
|
||||||
|
std::process::abort();
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
if index == parsed_responses.len() - 1 {
|
match fs::File::create(&file_path) {
|
||||||
checkpoint_number = executions.checkpoint.clone();
|
Ok(_) => {}
|
||||||
break;
|
Err(e) => {
|
||||||
}
|
println!("Failed to Create Directory {:?}: {}", &file_path, e);
|
||||||
let history_date = match NaiveDate::parse_from_str(
|
std::process::abort();
|
||||||
&executions.datetime.replace(" +0000 UTC", ""),
|
}
|
||||||
"%Y-%m-%dT%H:%M:%SZ",
|
|
||||||
) {
|
|
||||||
Ok(date) => date,
|
|
||||||
Err(_) => continue,
|
|
||||||
};
|
|
||||||
if history_date >= cutoff.into() {
|
|
||||||
let key = (
|
|
||||||
executions.sha256.clone(),
|
|
||||||
executions.filename.clone(),
|
|
||||||
executions.hostname.clone(),
|
|
||||||
);
|
|
||||||
seen.entry(key).or_insert(executions.clone());
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let final_response = ApiResponse {
|
let data: ApiResponse = ApiResponse {
|
||||||
error: "Success".to_string(),
|
error: "Success".to_string(),
|
||||||
response: ExecHistories {
|
response: ExecHistories {
|
||||||
exechistories: seen.values().cloned().collect(),
|
exechistories: vec![],
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
let data_write = serde_json::to_string_pretty(&final_response).unwrap();
|
let writeable_filepath: PathBuf = file_path.clone();
|
||||||
fs::write(&writeable_filepath, data_write).unwrap();
|
let data_write: String = serde_json::to_string_pretty(&data).expect("Failed to serialize");
|
||||||
if let Some(last_item) = &final_response.response.exechistories.last()
|
match fs::write(writeable_filepath.clone(), data_write) {
|
||||||
&& let Ok(last_date) = NaiveDate::parse_from_str(
|
Ok(_) => {}
|
||||||
&last_item.datetime.replace(" +0000 UTC", ""),
|
Err(e) => {
|
||||||
"%Y-%m-%dT%H:%M:%SZ",
|
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
|
||||||
)
|
std::process::abort();
|
||||||
{
|
|
||||||
let date_diff = Local::now().naive_local().date() - last_date;
|
|
||||||
let percentage_diff = (days - date_diff.num_days()) as f64 / days as f64 * 100.0;
|
|
||||||
progress_bar.set_position(percentage_diff.round() as u64);
|
|
||||||
progress_bar.set_message("Total Percent Complete");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
progress_bar.finish_with_message("All Checkpoints Complete");
|
|
||||||
let return_data = fs::read_to_string(&writeable_filepath).unwrap();
|
|
||||||
PyString::new(py, &return_data).into()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Client {
|
|
||||||
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);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
let mut checkpoint_number: String = SkipBack::find_checkpoint(days).to_string();
|
||||||
|
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 - Policy Name: {msg}: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
|
||||||
|
.unwrap().progress_chars("⣿⣦⣀")
|
||||||
|
);
|
||||||
|
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(
|
||||||
|
"info",
|
||||||
|
vec![KeyValue::new(
|
||||||
|
"Client Built Successfully",
|
||||||
|
format!("{:?}", client_result),
|
||||||
|
)],
|
||||||
|
);
|
||||||
|
client_result
|
||||||
|
}
|
||||||
|
Err(client_result) => {
|
||||||
|
cx.span().add_event(
|
||||||
|
"warn",
|
||||||
|
vec![KeyValue::new(
|
||||||
|
"Client Failed to Build",
|
||||||
|
format!("{:?}", &client_result),
|
||||||
|
)],
|
||||||
|
);
|
||||||
|
cx.span()
|
||||||
|
.set_status(Status::error("Client Failed to Build"));
|
||||||
|
println!("Failed to Build Client: {:?}", client_result);
|
||||||
|
std::process::abort();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let cutoff: chrono::NaiveDateTime =
|
||||||
|
Local::now().naive_local() - chrono::Duration::days(days);
|
||||||
|
let (tx, rx) = unbounded::<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: 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(),
|
||||||
|
response: ExecHistories {
|
||||||
|
exechistories: vec![],
|
||||||
|
},
|
||||||
|
});
|
||||||
|
existing
|
||||||
|
.response
|
||||||
|
.exechistories
|
||||||
|
.into_iter()
|
||||||
|
.map(|entry: Group| {
|
||||||
|
(
|
||||||
|
(
|
||||||
|
entry.sha256.clone(),
|
||||||
|
entry.filename.clone(),
|
||||||
|
entry.hostname.clone(),
|
||||||
|
),
|
||||||
|
entry,
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
} else {
|
||||||
|
HashMap::new()
|
||||||
|
};
|
||||||
|
while let Ok(parsed_responses) = rx.recv() {
|
||||||
|
for executions in parsed_responses {
|
||||||
|
if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let history_date: NaiveDate = match NaiveDate::parse_from_str(
|
||||||
|
&executions.datetime.replace(" +0000 UTC", ""),
|
||||||
|
"%Y-%m-%dT%H:%M:%SZ",
|
||||||
|
) {
|
||||||
|
Ok(date) => date,
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
if history_date >= cutoff.into() {
|
||||||
|
let key: (String, String, String) = (
|
||||||
|
executions.sha256.clone(),
|
||||||
|
executions.filename.clone(),
|
||||||
|
executions.hostname.clone(),
|
||||||
|
);
|
||||||
|
seen.entry(key).or_insert(executions.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let final_response: ApiResponse = ApiResponse {
|
||||||
|
error: "Success".to_string(),
|
||||||
|
response: ExecHistories {
|
||||||
|
exechistories: seen.values().cloned().collect(),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
let data_write: String = serde_json::to_string_pretty(&final_response).unwrap();
|
||||||
|
match fs::write(&writeable_filepath, data_write) {
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(e) => {
|
||||||
|
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let mut first_date: Option<NaiveDate> = None;
|
||||||
|
tracer.in_span("Airlock Data Retreival", |cx| {
|
||||||
|
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(
|
||||||
|
&base_url,
|
||||||
|
&exec_types,
|
||||||
|
&checkpoint_number,
|
||||||
|
&policy_names,
|
||||||
|
&client,
|
||||||
|
);
|
||||||
|
cx.span().set_attribute(KeyValue::new(
|
||||||
|
"items_in_response",
|
||||||
|
results.response.exechistories.len().to_string(),
|
||||||
|
));
|
||||||
|
results
|
||||||
|
});
|
||||||
|
let parsed_responses: Vec<Group> = execution_histories.response.exechistories;
|
||||||
|
if parsed_responses.is_empty() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
tx.send(parsed_responses.clone()).unwrap();
|
||||||
|
checkpoint_number = parsed_responses.last().unwrap().checkpoint.clone();
|
||||||
|
if let Some(last_item) = parsed_responses.last()
|
||||||
|
&& let Ok(last_date) = NaiveDate::parse_from_str(
|
||||||
|
&last_item.datetime.replace(" +0000 UTC", ""),
|
||||||
|
"%Y-%m-%dT%H:%M:%SZ",
|
||||||
|
)
|
||||||
|
{
|
||||||
|
if first_date.is_none() {
|
||||||
|
first_date = Some(last_date);
|
||||||
|
}
|
||||||
|
if let Some(base_date) = first_date {
|
||||||
|
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;
|
||||||
|
pb_clone.lock().unwrap().set_position(percentage);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
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);
|
||||||
|
return_data.to_string()
|
||||||
|
});
|
||||||
|
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> {
|
||||||
Client::builder()
|
Client::builder()
|
||||||
.danger_accept_invalid_certs(true)
|
.danger_accept_invalid_certs(true)
|
||||||
.default_headers(header_map)
|
.default_headers(headers)
|
||||||
.timeout(std::time::Duration::from_secs(300))
|
.timeout(std::time::Duration::from_secs(300))
|
||||||
.build()
|
.build()
|
||||||
.unwrap()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
async fn history_logging(
|
async fn history_logging(
|
||||||
py: Python<'_>,
|
base_url: &String,
|
||||||
py_self: &Py<PyAny>,
|
|
||||||
exec_types: &String,
|
exec_types: &String,
|
||||||
checkpoint_number: &String,
|
checkpoint_number: &String,
|
||||||
policy_names: &String,
|
policy_names: &Option<String>,
|
||||||
client: &Client,
|
client: &Client,
|
||||||
) -> ApiResponse {
|
) -> ApiResponse {
|
||||||
let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
|
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!(
|
let payload = format!(
|
||||||
r#"{{
|
r#"{{
|
||||||
"type": {},
|
"type": {},
|
||||||
"checkpoint": "{}",
|
"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))
|
.post(format!("{}/v1/logging/exechistories", base_url))
|
||||||
.body(payload)
|
.body(payload)
|
||||||
.send()
|
.send()
|
||||||
@@ -251,7 +306,7 @@ async fn history_logging(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn get_base_directory() -> PathBuf {
|
pub fn get_base_directory() -> PathBuf {
|
||||||
let home = env::var_os("HOME")
|
let home = env::var_os("HOME")
|
||||||
.map(PathBuf::from)
|
.map(PathBuf::from)
|
||||||
.or_else(|| env::var_os("USERPROFILE").map(PathBuf::from))
|
.or_else(|| env::var_os("USERPROFILE").map(PathBuf::from))
|
||||||
@@ -264,15 +319,10 @@ fn get_base_directory() -> PathBuf {
|
|||||||
.unwrap_or_else(|| home.join("AppData").join("Roaming"));
|
.unwrap_or_else(|| home.join("AppData").join("Roaming"));
|
||||||
appdata.join("Loxide")
|
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 skipback(days: i64) -> ObjectId {
|
|
||||||
let date_days_ago = Local::now() - Duration::days(days);
|
|
||||||
let timestamp = date_days_ago.timestamp() as u32;
|
|
||||||
let mut hex_timestamp = String::new();
|
|
||||||
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap();
|
|
||||||
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
|
|
||||||
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -1,14 +1,39 @@
|
|||||||
{
|
{
|
||||||
"APPNAME": "AirlockTools",
|
"APPNAME": "Loxide",
|
||||||
"URL": "https://server:3129",
|
"URL": "https://server:3129",
|
||||||
"LOG_LEVEL": "INFO",
|
"LOG_LEVEL": "INFO",
|
||||||
"BAD_PATH_PARTS": ["users","wwwroot","windows\\temp","windows\\task","windows\\system32","startup", "windows\\fonts","Recycle.Bin","AppData","programdata", "Solarwinds","kaseya"],
|
"BAD_PATH_PARTS": [
|
||||||
"BAD_PUBLISHERS": ["Brave", "Zoom", "GlavSoft", "VNC"],
|
"users",
|
||||||
"PUPS":["logmein","invalid","nmap","LTSvc","VNC","Kaseya","Solarwinds","mRemoteNG"],
|
"wwwroot",
|
||||||
|
"windows\\temp",
|
||||||
|
"windows\\task",
|
||||||
|
"windows\\system32",
|
||||||
|
"startup",
|
||||||
|
"windows\\fonts",
|
||||||
|
"Recycle.Bin",
|
||||||
|
"AppData",
|
||||||
|
"programdata",
|
||||||
|
"Solarwinds",
|
||||||
|
"kaseya"
|
||||||
|
],
|
||||||
|
"BAD_PUBLISHERS": [
|
||||||
|
"Brave",
|
||||||
|
"Zoom",
|
||||||
|
"GlavSoft",
|
||||||
|
"VNC"
|
||||||
|
],
|
||||||
|
"PUPS": [
|
||||||
|
"logmein",
|
||||||
|
"invalid",
|
||||||
|
"nmap",
|
||||||
|
"LTSvc",
|
||||||
|
"VNC",
|
||||||
|
"Kaseya",
|
||||||
|
"Solarwinds",
|
||||||
|
"mRemoteNG"
|
||||||
|
],
|
||||||
"PATH_EXCLUSION_CONST": 4,
|
"PATH_EXCLUSION_CONST": 4,
|
||||||
"MIN_FILES_FOR_PATH": 4,
|
"MIN_FILES_FOR_PATH": 4,
|
||||||
"VT_THREAT_TOLERANCE": 4,
|
"VT_THREAT_TOLERANCE": 4,
|
||||||
"POLICY_MAP_ENF_AUD": {
|
"POLICY_MAP_ENF_AUD": {}
|
||||||
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
+36
-26
@@ -1,6 +1,17 @@
|
|||||||
"""
|
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||||
This module handles the creation of local approval requests.
|
#
|
||||||
"""
|
# 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 logging
|
||||||
import os
|
import os
|
||||||
@@ -10,7 +21,7 @@ from typing import List, Optional
|
|||||||
from models.agent import Agent
|
from models.agent import Agent
|
||||||
from services.agenthandler import moveAgentToRelatedPolicy, selectAgents
|
from services.agenthandler import moveAgentToRelatedPolicy, selectAgents
|
||||||
from services.API import AirlockAPIWrapper
|
from services.API import AirlockAPIWrapper
|
||||||
from utils.configmanager import get_protected_json
|
from utils.configmanager import get_system_json
|
||||||
from utils.utils import colorText, get_sanitized_input
|
from utils.utils import colorText, get_sanitized_input
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -28,7 +39,7 @@ class LocalApprovalRequestor:
|
|||||||
username: Username creating the approvals (for tracking)
|
username: Username creating the approvals (for tracking)
|
||||||
"""
|
"""
|
||||||
self.api = api
|
self.api = api
|
||||||
self.policy_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
self.policy_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||||||
self.username = (
|
self.username = (
|
||||||
username or os.getenv("USERNAME") or os.getenv("USER") or "unknown"
|
username or os.getenv("USERNAME") or os.getenv("USER") or "unknown"
|
||||||
)
|
)
|
||||||
@@ -51,7 +62,7 @@ class LocalApprovalRequestor:
|
|||||||
batch_id = int(time.time())
|
batch_id = int(time.time())
|
||||||
|
|
||||||
purpose = (
|
purpose = (
|
||||||
f"🎫 Local Approval 🎫 - {duration_minutes} mins - "
|
f" Local Approval - {duration_minutes} mins - "
|
||||||
f"batch:{batch_id} Client:{agent_id} User:{self.username}"
|
f"batch:{batch_id} Client:{agent_id} User:{self.username}"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -103,11 +114,9 @@ class LocalApprovalRequestor:
|
|||||||
success_count = 0
|
success_count = 0
|
||||||
failure_count = 0
|
failure_count = 0
|
||||||
|
|
||||||
print(colorText(f"\n📦 Processing batch {batch_id}...", "cyan"))
|
print(colorText(f"\n Processing batch {batch_id}...", "cyan"))
|
||||||
print(colorText(f"👤 Requested by: {self.username}", "cyan"))
|
print(colorText(f" Requested by: {self.username}", "cyan"))
|
||||||
print(
|
print(colorText(f" Moving {len(agents)} agent(s) to local approval\n", "cyan"))
|
||||||
colorText(f"📊 Moving {len(agents)} agent(s) to local approval\n", "cyan")
|
|
||||||
)
|
|
||||||
|
|
||||||
for agent in agents:
|
for agent in agents:
|
||||||
try:
|
try:
|
||||||
@@ -125,11 +134,11 @@ class LocalApprovalRequestor:
|
|||||||
if not move_success:
|
if not move_success:
|
||||||
raise Exception("Failed to move to audit policy")
|
raise Exception("Failed to move to audit policy")
|
||||||
|
|
||||||
print(colorText(f"✓ {agent.hostname}", "green"))
|
print(colorText(f" {agent.hostname}", "green"))
|
||||||
success_count += 1
|
success_count += 1
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(colorText(f"✗ {agent.hostname}: {e}", "red"))
|
print(colorText(f" {agent.hostname}: {e}", "red"))
|
||||||
logger.error(f"Error processing agent {agent.hostname}: {e}")
|
logger.error(f"Error processing agent {agent.hostname}: {e}")
|
||||||
failure_count += 1
|
failure_count += 1
|
||||||
|
|
||||||
@@ -152,7 +161,7 @@ class LocalApprovalRequestor:
|
|||||||
]
|
]
|
||||||
|
|
||||||
# Display duration options
|
# Display duration options
|
||||||
print(colorText("\n⏱️ Select Local Approval Duration:", "white"))
|
print(colorText("\n Select Local Approval Duration:", "white"))
|
||||||
print(colorText("=" * 50, "white"))
|
print(colorText("=" * 50, "white"))
|
||||||
|
|
||||||
for i, (minutes, label) in enumerate(duration_options, start=1):
|
for i, (minutes, label) in enumerate(duration_options, start=1):
|
||||||
@@ -166,7 +175,7 @@ class LocalApprovalRequestor:
|
|||||||
|
|
||||||
if 1 <= choice <= len(duration_options):
|
if 1 <= choice <= len(duration_options):
|
||||||
duration_minutes, duration_label = duration_options[choice - 1]
|
duration_minutes, duration_label = duration_options[choice - 1]
|
||||||
print(colorText(f"✓ Selected: {duration_label}", "green"))
|
print(colorText(f" Selected: {duration_label}", "green"))
|
||||||
logger.info(f"User selected duration: {duration_minutes} minutes")
|
logger.info(f"User selected duration: {duration_minutes} minutes")
|
||||||
else:
|
else:
|
||||||
print(colorText("❌ Invalid choice.", "red"))
|
print(colorText("❌ Invalid choice.", "red"))
|
||||||
@@ -179,7 +188,7 @@ class LocalApprovalRequestor:
|
|||||||
return
|
return
|
||||||
|
|
||||||
# Select agents
|
# Select agents
|
||||||
print(colorText("\n🎯 Select Agents for Local Approval:", "white"))
|
print(colorText("\nSelect Agents for Local Approval:", "white"))
|
||||||
agents = selectAgents(self.api)
|
agents = selectAgents(self.api)
|
||||||
|
|
||||||
if not agents:
|
if not agents:
|
||||||
@@ -188,7 +197,7 @@ class LocalApprovalRequestor:
|
|||||||
return
|
return
|
||||||
|
|
||||||
# Confirm with user
|
# Confirm with user
|
||||||
print(colorText("\n📋 Summary:", "cyan"))
|
print(colorText("\nSummary:", "cyan"))
|
||||||
print(colorText(f" Duration: {duration_label}", "white"))
|
print(colorText(f" Duration: {duration_label}", "white"))
|
||||||
print(colorText(f" Agents: {len(agents)}", "white"))
|
print(colorText(f" Agents: {len(agents)}", "white"))
|
||||||
|
|
||||||
@@ -219,24 +228,25 @@ class LocalApprovalRequestor:
|
|||||||
failure_count: Number of failed operations
|
failure_count: Number of failed operations
|
||||||
"""
|
"""
|
||||||
print(colorText(f"\n{'=' * 60}", "white"))
|
print(colorText(f"\n{'=' * 60}", "white"))
|
||||||
print(colorText("📊 Local Approval Summary", "cyan"))
|
print(colorText(" Local Approval Summary", "cyan"))
|
||||||
print(colorText("=" * 60, "white"))
|
print(colorText("=" * 60, "white"))
|
||||||
|
|
||||||
print(colorText(f"✓ Successfully processed: {success_count}", "green"))
|
print(colorText(f" Successfully processed: {success_count}", "green"))
|
||||||
|
|
||||||
if failure_count > 0:
|
if failure_count > 0:
|
||||||
print(colorText(f"✗ Failed: {failure_count}", "red"))
|
print(colorText(f" Failed: {failure_count}", "red"))
|
||||||
|
|
||||||
print(colorText(f"\n📦 Batch ID: {batch_id}", "cyan"))
|
print(colorText(f"\n Batch ID: {batch_id}", "cyan"))
|
||||||
print(colorText(f"⏱️ Duration: {duration_label}", "cyan"))
|
print(colorText(f" Duration: {duration_label}", "cyan"))
|
||||||
|
|
||||||
print(colorText("=" * 60, "white"))
|
print(colorText("=" * 60, "white"))
|
||||||
print(colorText("\n💡 Next Steps:", "yellow"))
|
print(colorText("\n Next Steps:", "yellow"))
|
||||||
print(colorText(" • Agents have been moved to audit policies", "white"))
|
print(colorText(" ✅ Agents have been moved to audit policies", "white"))
|
||||||
print(colorText(" • Local approvals are active", "white"))
|
print(colorText(" ✅ Local approvals are active", "white"))
|
||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
f" • Agents will return to enforcement after {duration_label}", "white"
|
f" ✅ Agents will return to enforcement after {duration_label}",
|
||||||
|
"white",
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
print(colorText("=" * 60 + "\n", "white"))
|
print(colorText("=" * 60 + "\n", "white"))
|
||||||
|
|||||||
+1
-119
@@ -14,136 +14,18 @@
|
|||||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||||
|
|
||||||
|
|
||||||
from datetime import datetime
|
|
||||||
import logging
|
import logging
|
||||||
import os
|
|
||||||
|
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
|
|
||||||
from services.agenthandler import selectAgents
|
from services.agenthandler import selectAgents
|
||||||
from services.API import AirlockAPIWrapper
|
from services.API import AirlockAPIWrapper
|
||||||
from utils.configmanager import load_env
|
|
||||||
from utils.selector import Selector
|
from utils.selector import Selector
|
||||||
from utils.utils import colorText, get_sanitized_input
|
from utils.utils import get_sanitized_input
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
def otp_generate(api: AirlockAPIWrapper):
|
|
||||||
otp_dict = {}
|
|
||||||
agents = selectAgents(api)
|
|
||||||
print(colorText("Would you like to continue with these devices?", "white"))
|
|
||||||
for agent in agents:
|
|
||||||
print(agent.hostname)
|
|
||||||
confirm = Selector.confirm()
|
|
||||||
if agents and confirm:
|
|
||||||
requester = get_sanitized_input("Who is requesting the OTP: ")
|
|
||||||
because = get_sanitized_input("Why/What work are they doing?: ")
|
|
||||||
|
|
||||||
purpose = f"Requester: {requester} - for : {because}"
|
|
||||||
possible_durations = [15, 60, 360, 1440, 10080]
|
|
||||||
|
|
||||||
print(colorText("Please select a duration in minutes: ", "white"))
|
|
||||||
print(
|
|
||||||
colorText(
|
|
||||||
"15 mins, 60 mins, 360 mins(6 Hours), 1440 mins (24 Hours), 10080 mins (7 Days):",
|
|
||||||
"white",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
duration_selected = Selector.select_int(possible_durations)
|
|
||||||
|
|
||||||
if isinstance(duration_selected, list):
|
|
||||||
duration_selected = duration_selected[0] if duration_selected else None
|
|
||||||
|
|
||||||
if duration_selected is not None:
|
|
||||||
for agent in agents:
|
|
||||||
logging.info(f"Querying API for {agent.hostname}")
|
|
||||||
otp_code = api.otp_generate(agent.agentid, duration_selected, purpose)
|
|
||||||
logger.debug(f"Generated OTP for {agent.hostname}: {otp_code}")
|
|
||||||
otp_dict[agent.hostname] = otp_code
|
|
||||||
|
|
||||||
print(colorText("Requested Codes:", "green"))
|
|
||||||
for key, value in otp_dict.items():
|
|
||||||
print(colorText(f"{key} | {value}", "green"))
|
|
||||||
|
|
||||||
|
|
||||||
def otp_activities_by_agent(api: AirlockAPIWrapper):
|
|
||||||
activeagents = api.otp_find_active()
|
|
||||||
awaitingagents = api.otp_find_awaiting()
|
|
||||||
enforcedagents = api.otp_find_enforced()
|
|
||||||
revokedagents = api.otp_find_revoked()
|
|
||||||
|
|
||||||
# Add a 'status' column to each DataFrame
|
|
||||||
activeagents["status"] = "active"
|
|
||||||
awaitingagents["status"] = "awaiting"
|
|
||||||
enforcedagents["status"] = "enforced"
|
|
||||||
revokedagents["status"] = "revoked"
|
|
||||||
|
|
||||||
# Combine all into one DataFrame
|
|
||||||
combined_agents = pd.concat(
|
|
||||||
[activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True
|
|
||||||
)
|
|
||||||
combined_agents = combined_agents.sort_values(by="otpid", ascending=False)
|
|
||||||
|
|
||||||
# Optionally, select specific hosts
|
|
||||||
user_input = (
|
|
||||||
get_sanitized_input("\nWould you like to search for a specific device? (y/n): ")
|
|
||||||
.strip()
|
|
||||||
.lower()
|
|
||||||
)
|
|
||||||
if user_input == "y":
|
|
||||||
agentnames = []
|
|
||||||
agents = selectAgents(api)
|
|
||||||
for agent in agents:
|
|
||||||
agentnames.append(agent.hostname)
|
|
||||||
|
|
||||||
combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)]
|
|
||||||
|
|
||||||
# Present and select rows
|
|
||||||
selected_rows = Selector.select_dataframe_with_mode(
|
|
||||||
combined_agents,
|
|
||||||
columns=["otpid", "hostname", "status", "purpose", "granted"],
|
|
||||||
header="OTP Sessions",
|
|
||||||
)
|
|
||||||
combined_df = pd.DataFrame()
|
|
||||||
|
|
||||||
for row in selected_rows:
|
|
||||||
otpid = row["otpid"]
|
|
||||||
hostname = row["hostname"]
|
|
||||||
result = api.otp_get_activities(otpid)
|
|
||||||
result["hostname"] = hostname
|
|
||||||
if not result.empty:
|
|
||||||
logger.info(f"Activities for {hostname} (otpid: {otpid}):\n{result}")
|
|
||||||
combined_df = pd.concat([combined_df, result], ignore_index=True)
|
|
||||||
else:
|
|
||||||
logger.info(f"No activities found for {hostname} (otpid: {otpid})")
|
|
||||||
|
|
||||||
user_input = (
|
|
||||||
get_sanitized_input(
|
|
||||||
"\nWould you like to export the results to a CSV file? (y/n): "
|
|
||||||
)
|
|
||||||
.strip()
|
|
||||||
.lower()
|
|
||||||
)
|
|
||||||
if user_input == "y":
|
|
||||||
working_dir = load_env("WORKING_DIR")
|
|
||||||
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
|
||||||
filename = f"otp_activities_{timestamp}.csv"
|
|
||||||
file_path = os.path.join(str(working_dir), filename)
|
|
||||||
|
|
||||||
combined_df.to_csv(file_path, index=False)
|
|
||||||
logging.info(f"Exported Data to {file_path}")
|
|
||||||
|
|
||||||
print(
|
|
||||||
colorText(
|
|
||||||
f"\n✅ OTP Activity exported to: {working_dir}\\{filename}",
|
|
||||||
"green",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
logging.debug("User declined to export the DataFrame.")
|
|
||||||
|
|
||||||
|
|
||||||
def otp_revoke(api: AirlockAPIWrapper):
|
def otp_revoke(api: AirlockAPIWrapper):
|
||||||
|
|
||||||
activeagents = api.otp_find_active()
|
activeagents = api.otp_find_active()
|
||||||
|
|||||||
+29
-29
@@ -25,7 +25,7 @@ import pandas as pd
|
|||||||
from models.execution import ExecutionHistoryRecord
|
from models.execution import ExecutionHistoryRecord
|
||||||
from models.policy import Allowlist, Policy
|
from models.policy import Allowlist, Policy
|
||||||
from services.API import AirlockAPIWrapper
|
from services.API import AirlockAPIWrapper
|
||||||
from utils.configmanager import get_protected_value, load_env, load_env_json
|
from utils.configmanager import get_system_list, get_system_value, load_env
|
||||||
from utils.selector import Selector
|
from utils.selector import Selector
|
||||||
from utils.utils import (
|
from utils.utils import (
|
||||||
areYouSure,
|
areYouSure,
|
||||||
@@ -88,7 +88,7 @@ def sortHashes(
|
|||||||
):
|
):
|
||||||
working_dir = load_env("WORKING_DIR")
|
working_dir = load_env("WORKING_DIR")
|
||||||
history_days = Selector.select_value(
|
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-150): ",
|
||||||
value_type=int,
|
value_type=int,
|
||||||
valid_range=(1, 150),
|
valid_range=(1, 150),
|
||||||
)
|
)
|
||||||
@@ -157,7 +157,7 @@ def buildPathsandPublishers(selected_policies: List[Policy], split):
|
|||||||
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
|
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
|
||||||
)
|
)
|
||||||
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
|
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
|
||||||
path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type=int)
|
path_exclusion_constant = get_system_value("PATH_EXCLUSION_CONST", cast_type=int)
|
||||||
|
|
||||||
if os.path.exists(path1):
|
if os.path.exists(path1):
|
||||||
df1 = pd.read_csv(path1)
|
df1 = pd.read_csv(path1)
|
||||||
@@ -227,7 +227,7 @@ def buildPathsandPublishers(selected_policies: List[Policy], split):
|
|||||||
all_approved_hashes["publisher"] != "Not Signed"
|
all_approved_hashes["publisher"] != "Not Signed"
|
||||||
].drop_duplicates(subset=["publisher"])
|
].drop_duplicates(subset=["publisher"])
|
||||||
# Remove Bad publisher if somehow they made it this far
|
# Remove Bad publisher if somehow they made it this far
|
||||||
pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
|
pattern = regulator(get_system_list("BAD_PUBLISHERS"))
|
||||||
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
|
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
|
||||||
publist = publist[["publisher"]]
|
publist = publist[["publisher"]]
|
||||||
publist.sort_values(by="publisher", inplace=True)
|
publist.sort_values(by="publisher", inplace=True)
|
||||||
@@ -313,7 +313,7 @@ def buildPreflights(selected_policies: List[Policy]):
|
|||||||
|
|
||||||
|
|
||||||
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
|
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
|
||||||
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type=int)
|
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
|
||||||
|
|
||||||
def clean_split(path):
|
def clean_split(path):
|
||||||
if not isinstance(path, (str, bytes, os.PathLike)):
|
if not isinstance(path, (str, bytes, os.PathLike)):
|
||||||
@@ -385,8 +385,8 @@ def calculatePath(approved_hashes, path_exclusion_constant, split):
|
|||||||
else:
|
else:
|
||||||
dfs_by_policy = [approved_hashes]
|
dfs_by_policy = [approved_hashes]
|
||||||
|
|
||||||
badpathparts = load_env_json("BAD_PATH_PARTS", "[]")
|
badpathparts = get_system_list("BAD_PATH_PARTS")
|
||||||
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type=int)
|
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
|
||||||
|
|
||||||
processed_dfs = []
|
processed_dfs = []
|
||||||
|
|
||||||
@@ -655,7 +655,7 @@ def section_header(title):
|
|||||||
|
|
||||||
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
|
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
|
||||||
working_dir = load_env("WORKING_DIR")
|
working_dir = load_env("WORKING_DIR")
|
||||||
section_header("🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒")
|
section_header("Prepare to Enforce Policy")
|
||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
"\nSequentially follow these steps to prepare a policy for enforcement:",
|
"\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:
|
if not selected_policies:
|
||||||
print(colorText(" [✗] No policies have been chosen", "red"))
|
print(colorText(" [❌] No policies have been chosen", "red"))
|
||||||
else:
|
else:
|
||||||
print(colorText("The following policies have been chosen:", "green"))
|
print(colorText("The following policies have been chosen:", "green"))
|
||||||
for policy in selected_policies:
|
for policy in selected_policies:
|
||||||
print(colorText(f" [✓] {policy.name}", "green"))
|
print(colorText(f" [✅] {policy.name}", "green"))
|
||||||
|
|
||||||
# Step 2: Destination Policy and Allowlist
|
# Step 2: Destination Policy and Allowlist
|
||||||
print(
|
print(
|
||||||
@@ -683,22 +683,22 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
|||||||
if destination_policy:
|
if destination_policy:
|
||||||
print(
|
print(
|
||||||
colorText(
|
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",
|
"green",
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
print(colorText(" [✗] No destination policy has been chosen", "red"))
|
print(colorText(" [❌] No destination policy has been chosen", "red"))
|
||||||
|
|
||||||
if destination_allowlist:
|
if destination_allowlist:
|
||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
|
f" [✅] {destination_allowlist[0].name} has been selected as allowlist",
|
||||||
"green",
|
"green",
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
print(colorText(" [✗] No allowlist has been chosen", "red"))
|
print(colorText(" [❌] No allowlist has been chosen", "red"))
|
||||||
|
|
||||||
# Step 3: Data Preparation
|
# Step 3: Data Preparation
|
||||||
print(
|
print(
|
||||||
@@ -713,9 +713,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
|||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
(
|
(
|
||||||
" [✓] Data has been fetched"
|
" [✅] Data has been fetched"
|
||||||
if os.path.exists(review_path)
|
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",
|
"green" if os.path.exists(review_path) else "red",
|
||||||
)
|
)
|
||||||
@@ -723,7 +723,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
|||||||
else:
|
else:
|
||||||
print(
|
print(
|
||||||
colorText(
|
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(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
(
|
(
|
||||||
" [✓] Reviewed hashes have been loaded"
|
" [✅] Reviewed hashes have been loaded"
|
||||||
if os.path.exists(approved_path)
|
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",
|
"green" if os.path.exists(approved_path) else "red",
|
||||||
)
|
)
|
||||||
@@ -766,9 +766,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
|||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
(
|
(
|
||||||
" [✓] Path review list created"
|
" [✅] Path review list created"
|
||||||
if os.path.exists(second_review_path)
|
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",
|
"green" if os.path.exists(second_review_path) else "red",
|
||||||
)
|
)
|
||||||
@@ -776,7 +776,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
|||||||
else:
|
else:
|
||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
" [✗] No policies selected, cannot check reviewed hashes or path list",
|
" [❌] No policies selected, cannot check reviewed hashes or path list",
|
||||||
"red",
|
"red",
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
@@ -812,9 +812,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
|||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
(
|
(
|
||||||
" [✓] Reviewed path list detected"
|
" [✅] Reviewed path list detected"
|
||||||
if os.path.exists(reviewed_path)
|
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",
|
"green" if os.path.exists(reviewed_path) else "red",
|
||||||
)
|
)
|
||||||
@@ -825,9 +825,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
|||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
(
|
(
|
||||||
" [✓] Preflight Path Exclusion List has been generated"
|
" [✅] Preflight Path Exclusion List has been generated"
|
||||||
if preflight_ready
|
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",
|
"green" if preflight_ready else "red",
|
||||||
)
|
)
|
||||||
@@ -835,7 +835,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
|||||||
else:
|
else:
|
||||||
print(
|
print(
|
||||||
colorText(
|
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"))
|
print(colorText(" Apply approved hashes to allowlist", "cyan"))
|
||||||
|
|
||||||
# Utility Options
|
# Utility Options
|
||||||
print(colorText("F. 📂 - Open Working Directory", "cyan"))
|
print(colorText("F. Open Working Directory", "cyan"))
|
||||||
print(colorText("B. 🔚 - Back", "cyan"))
|
print(colorText("B. Back", "cyan"))
|
||||||
|
|||||||
@@ -1,134 +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 datetime
|
|
||||||
import logging
|
|
||||||
|
|
||||||
import dotenv
|
|
||||||
import pandas as pd
|
|
||||||
|
|
||||||
from flows.prepPolicy import selectPolicies
|
|
||||||
from services.API import AirlockAPIWrapper
|
|
||||||
from services.policyhandler import getPolicyInfo
|
|
||||||
from utils.configmanager import load_env
|
|
||||||
from utils.selector import Selector
|
|
||||||
from utils.utils import colorText, get_sanitized_input
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
|
|
||||||
dotenv.load_dotenv()
|
|
||||||
|
|
||||||
|
|
||||||
def findQuietAgents(api: AirlockAPIWrapper):
|
|
||||||
working_dir = load_env("WORKING_DIR")
|
|
||||||
# Get policy selection and agent list
|
|
||||||
selected_policy = selectPolicies(api, False)
|
|
||||||
if selected_policy:
|
|
||||||
agents = api.agents_find_by_group(selected_policy[0].groupid)
|
|
||||||
|
|
||||||
# Prompt user for history range
|
|
||||||
history_days = Selector.select_value(
|
|
||||||
prompt="Enter how many days of history to pull (1–150): ",
|
|
||||||
value_type=int,
|
|
||||||
valid_range=(1, 150),
|
|
||||||
)
|
|
||||||
required_quiet = Selector.select_value(
|
|
||||||
prompt="Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–365): ",
|
|
||||||
value_type=int,
|
|
||||||
valid_range=(1, 150),
|
|
||||||
)
|
|
||||||
|
|
||||||
confirm = Selector.confirm(
|
|
||||||
f"Do you wish to proceed to pull history for {selected_policy[0].name}? Y/N : "
|
|
||||||
)
|
|
||||||
# Get execution history as a DataFrame
|
|
||||||
if confirm:
|
|
||||||
policy_exec_history = getPolicyInfo(
|
|
||||||
api, selected_policy[0], [1, 2, 6, 7], history_days
|
|
||||||
)
|
|
||||||
|
|
||||||
if policy_exec_history.empty:
|
|
||||||
logging.info(
|
|
||||||
"No execution history found for the selected policy and time range."
|
|
||||||
)
|
|
||||||
get_sanitized_input("Press enter to continue")
|
|
||||||
return
|
|
||||||
|
|
||||||
# Convert 'datetime' column to timezone-aware datetime objects
|
|
||||||
policy_exec_history["datetime"] = pd.to_datetime(
|
|
||||||
policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True
|
|
||||||
)
|
|
||||||
|
|
||||||
# Get current UTC time
|
|
||||||
now = datetime.datetime.now(datetime.timezone.utc)
|
|
||||||
|
|
||||||
# Calculate days ago
|
|
||||||
policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
|
|
||||||
lambda dt: (now - dt).days
|
|
||||||
)
|
|
||||||
|
|
||||||
# Count total executions per hostname
|
|
||||||
hostname_counts = policy_exec_history["hostname"].value_counts()
|
|
||||||
|
|
||||||
# Map execution counts to agents
|
|
||||||
agents["execution_count"] = (
|
|
||||||
agents["hostname"].map(hostname_counts).fillna(0).astype(int)
|
|
||||||
)
|
|
||||||
|
|
||||||
# Find most recent execution per hostname
|
|
||||||
most_recent_exec = policy_exec_history.sort_values(
|
|
||||||
by="days_ago"
|
|
||||||
).drop_duplicates(subset="hostname", keep="first")
|
|
||||||
|
|
||||||
# Map most recent execution age to agents
|
|
||||||
agents["days_since"] = agents["hostname"].map(
|
|
||||||
most_recent_exec.set_index("hostname")["days_ago"]
|
|
||||||
)
|
|
||||||
|
|
||||||
# Check for enforcement readiness
|
|
||||||
agents["required_quiet"] = required_quiet
|
|
||||||
agents["enforce_ready"] = agents["days_since"].apply(
|
|
||||||
lambda x: True if pd.isna(x) or x > required_quiet else False
|
|
||||||
)
|
|
||||||
|
|
||||||
# Sort agents by execution count and hostname
|
|
||||||
agents = agents.sort_values(
|
|
||||||
by=["execution_count", "hostname"], ascending=[True, True]
|
|
||||||
)
|
|
||||||
|
|
||||||
# Save to CSV
|
|
||||||
filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv"
|
|
||||||
logging.debug(f"Saving CSV to {filename}")
|
|
||||||
print(colorText(f"Saving CSV to {filename}", "green"))
|
|
||||||
agents.to_csv(filename, index=False)
|
|
||||||
|
|
||||||
# Summary statistics
|
|
||||||
total_agents = len(agents)
|
|
||||||
ready_agents = agents["enforce_ready"].sum()
|
|
||||||
not_ready_agents = total_agents - ready_agents
|
|
||||||
ready_percentage = (ready_agents / total_agents) * 100
|
|
||||||
|
|
||||||
# Print results
|
|
||||||
|
|
||||||
message = (
|
|
||||||
f"Total agents: {total_agents}\n"
|
|
||||||
f"Agents marked as 'enforce_ready': {ready_agents}\n"
|
|
||||||
f"Agents not ready: {not_ready_agents}\n"
|
|
||||||
f"Percentage ready for enforcement: {ready_percentage:.2f}%"
|
|
||||||
)
|
|
||||||
logger.debug(message)
|
|
||||||
colorText(message, "green")
|
|
||||||
get_sanitized_input("Press enter to continue")
|
|
||||||
+7
-7
@@ -28,7 +28,7 @@ import pandas as pd
|
|||||||
|
|
||||||
import airlock_libs
|
import airlock_libs
|
||||||
from services.API import AirlockAPIWrapper
|
from services.API import AirlockAPIWrapper
|
||||||
from utils.configmanager import get_protected_value, load_env_json
|
from utils.configmanager import get_system_list, get_system_value
|
||||||
from utils.utils import colorText, regulator
|
from utils.utils import colorText, regulator
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -90,9 +90,9 @@ class Hash:
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def categorize_hashes(cls, hashes):
|
def categorize_hashes(cls, hashes):
|
||||||
threat_tolerance = get_protected_value("VT_THREAT_TOLERANCE", cast_type=int)
|
threat_tolerance = get_system_value("VT_THREAT_TOLERANCE", cast_type=int)
|
||||||
bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
|
bad_publishers_pattern = regulator(get_system_list("BAD_PUBLISHERS"))
|
||||||
pups_pattern = regulator(load_env_json("PUPS", "[]"))
|
pups_pattern = regulator(get_system_list("PUPS"))
|
||||||
|
|
||||||
approved_count = 0
|
approved_count = 0
|
||||||
unapproved_count = 0
|
unapproved_count = 0
|
||||||
@@ -384,9 +384,9 @@ class ExecutionHistoryRecord:
|
|||||||
Returns:
|
Returns:
|
||||||
List[ExecutionHistoryRecord]: The same list, with hash_obj.at_decision updated.
|
List[ExecutionHistoryRecord]: The same list, with hash_obj.at_decision updated.
|
||||||
"""
|
"""
|
||||||
threat_tolerance = get_protected_value("VT_THREAT_TOLERANCE", cast_type=int)
|
threat_tolerance = get_system_value("VT_THREAT_TOLERANCE", cast_type=int)
|
||||||
bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
|
bad_publishers_pattern = regulator(get_system_list("BAD_PUBLISHERS"))
|
||||||
pups_pattern = regulator(load_env_json("PUPS", "[]"))
|
pups_pattern = regulator(get_system_list("PUPS"))
|
||||||
|
|
||||||
approved_count = 0
|
approved_count = 0
|
||||||
unapproved_count = 0
|
unapproved_count = 0
|
||||||
|
|||||||
+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
|
cryptography==46.0.3
|
||||||
keyring==25.6.0
|
keyring==25.6.0
|
||||||
numpy==2.3.4
|
|
||||||
pandas==2.3.3
|
# Environment management
|
||||||
pymongo==4.15.3
|
|
||||||
python-dotenv==1.2.1
|
python-dotenv==1.2.1
|
||||||
Requests==2.32.5
|
|
||||||
textual==6.5.0
|
# Utilities
|
||||||
tqdm==4.67.1
|
tqdm==4.67.1
|
||||||
urllib3==2.5.0
|
urllib3==2.5.0
|
||||||
pyperclip==1.11.0
|
pyperclip==1.11.0
|
||||||
|
|
||||||
|
# Custom/Private packages
|
||||||
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
|
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
|
||||||
airlock_libs==2.0.0
|
airlock_libs==7.1.0
|
||||||
@@ -1,61 +0,0 @@
|
|||||||
from typing import List
|
|
||||||
|
|
||||||
from textual.app import ComposeResult
|
|
||||||
from textual.screen import Screen
|
|
||||||
|
|
||||||
from models.agent import Agent
|
|
||||||
from widgets.agentmoveoperations import AgentMoveOperations
|
|
||||||
from widgets.multiagentselector import MultiAgentSelector
|
|
||||||
from widgets.resultsdisplay import ResultsDisplay
|
|
||||||
|
|
||||||
|
|
||||||
class MoveAgentWorkflowScreen(Screen):
|
|
||||||
"""Screen that handles the agent movement workflow."""
|
|
||||||
|
|
||||||
def __init__(self, all_agents: List[Agent]):
|
|
||||||
super().__init__()
|
|
||||||
self.all_agents = all_agents
|
|
||||||
self.selected_agents = None
|
|
||||||
|
|
||||||
def compose(self) -> ComposeResult:
|
|
||||||
"""Start with the multi-agent selector."""
|
|
||||||
yield MultiAgentSelector(self.all_agents)
|
|
||||||
|
|
||||||
def on_multi_agent_selector_agents_selected(
|
|
||||||
self, message: MultiAgentSelector.AgentsSelected
|
|
||||||
) -> None:
|
|
||||||
"""Handle selected agents - switch to operations screen."""
|
|
||||||
self.selected_agents = message.selected_agents
|
|
||||||
|
|
||||||
# Remove the MultiAgentSelector
|
|
||||||
selector = self.query_one(MultiAgentSelector)
|
|
||||||
selector.remove()
|
|
||||||
|
|
||||||
# Mount the AgentMoveOperations with the selected Agent objects
|
|
||||||
self.mount(AgentMoveOperations(self.selected_agents))
|
|
||||||
|
|
||||||
def on_agent_move_operations_operation_complete(
|
|
||||||
self, message: AgentMoveOperations.OperationComplete
|
|
||||||
) -> None:
|
|
||||||
"""Handle completion of move operation - transition to results screen."""
|
|
||||||
# Format successful results
|
|
||||||
success_lines = []
|
|
||||||
for agent, result in message.successful:
|
|
||||||
success_lines.append(f"✓ {agent.hostname}")
|
|
||||||
|
|
||||||
# Format unsuccessful results
|
|
||||||
failure_lines = []
|
|
||||||
for agent, error in message.unsuccessful:
|
|
||||||
failure_lines.append(f"✗ {agent.hostname}: {error}")
|
|
||||||
|
|
||||||
successful_text = "\n".join(success_lines) if success_lines else "(none)"
|
|
||||||
unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)"
|
|
||||||
|
|
||||||
# Remove the operations widget
|
|
||||||
ops_widget = self.query_one(AgentMoveOperations)
|
|
||||||
ops_widget.remove()
|
|
||||||
|
|
||||||
# Mount the results display
|
|
||||||
self.mount(
|
|
||||||
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
|
|
||||||
)
|
|
||||||
@@ -1,24 +0,0 @@
|
|||||||
# otp_workflow_screen.py
|
|
||||||
|
|
||||||
from typing import List
|
|
||||||
|
|
||||||
from textual.app import ComposeResult
|
|
||||||
from textual.screen import Screen
|
|
||||||
|
|
||||||
from models.agent import Agent
|
|
||||||
from widgets.OTP_generate import OTPGenerator
|
|
||||||
|
|
||||||
|
|
||||||
class OTPWorkflowScreen(Screen):
|
|
||||||
"""Screen that handles the OTP generation workflow without agent selection."""
|
|
||||||
|
|
||||||
def __init__(self, selected_agents: List[Agent]):
|
|
||||||
super().__init__()
|
|
||||||
self.selected_agents = selected_agents
|
|
||||||
|
|
||||||
def compose(self) -> ComposeResult:
|
|
||||||
"""Directly show the OTP generator for the selected agents."""
|
|
||||||
yield OTPGenerator(self.selected_agents)
|
|
||||||
|
|
||||||
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
|
|
||||||
"""Handle OTP generation request - pass it up to the app level if needed."""
|
|
||||||
@@ -64,6 +64,23 @@ class AirlockAPIWrapper:
|
|||||||
logger.error(f"API request failed: {e}")
|
logger.error(f"API request failed: {e}")
|
||||||
raise
|
raise
|
||||||
|
|
||||||
|
def _post_raw(self, endpoint: str, payload: Optional[dict] = None) -> bytes:
|
||||||
|
"""
|
||||||
|
Send POST request and return raw response content (bytes).
|
||||||
|
Useful for XML endpoints.
|
||||||
|
"""
|
||||||
|
url = f"{self.base_url}{endpoint}"
|
||||||
|
data = json.dumps(payload or {})
|
||||||
|
try:
|
||||||
|
logger.debug(f"POST Request to {url} with payload: {payload}")
|
||||||
|
response = requests.post(url, headers=self.headers, data=data, verify=False)
|
||||||
|
response.raise_for_status()
|
||||||
|
logger.debug(f"Raw response received from {url}")
|
||||||
|
return response.content # bytes
|
||||||
|
except requests.exceptions.RequestException as e:
|
||||||
|
logger.error(f"API request failed: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
# Allowlist Management
|
# Allowlist Management
|
||||||
def allowlist_find_all(self) -> pd.DataFrame:
|
def allowlist_find_all(self) -> pd.DataFrame:
|
||||||
"""
|
"""
|
||||||
@@ -75,6 +92,12 @@ class AirlockAPIWrapper:
|
|||||||
result = self._post("/v1/application", {})
|
result = self._post("/v1/application", {})
|
||||||
return pd.DataFrame(result["response"]["applications"])
|
return pd.DataFrame(result["response"]["applications"])
|
||||||
|
|
||||||
|
def allowlist_export(self, applicationid) -> bytes:
|
||||||
|
"""Return Allowlist XML as bytes to save to file"""
|
||||||
|
payload = {"applicationid": applicationid}
|
||||||
|
result = self._post_raw("/v1/application/export", payload)
|
||||||
|
return result # should be bytes
|
||||||
|
|
||||||
# Agent Management
|
# Agent Management
|
||||||
def agent_find_all(self) -> pd.DataFrame:
|
def agent_find_all(self) -> pd.DataFrame:
|
||||||
"""Retrieve all agents."""
|
"""Retrieve all agents."""
|
||||||
@@ -116,6 +139,30 @@ class AirlockAPIWrapper:
|
|||||||
result = self._post("/v1/agent/find", payload)
|
result = self._post("/v1/agent/find", payload)
|
||||||
return pd.DataFrame(result["response"]["agents"])
|
return pd.DataFrame(result["response"]["agents"])
|
||||||
|
|
||||||
|
# Baseline Managment
|
||||||
|
def baseline_find_all(self) -> pd.DataFrame:
|
||||||
|
"""Retrieve all Baselines."""
|
||||||
|
result = self._post("/v1/baseline", {})
|
||||||
|
return pd.DataFrame(result["response"]["baselines"])
|
||||||
|
|
||||||
|
def baseline_export(self, baselineid) -> bytes:
|
||||||
|
"""Return Baseline XML as bytes to save to file"""
|
||||||
|
payload = {"baselineid": baselineid}
|
||||||
|
result = self._post_raw("/v1/baseline/export", payload)
|
||||||
|
return result
|
||||||
|
|
||||||
|
# Blocklist Managment
|
||||||
|
def blocklist_find_all(self) -> pd.DataFrame:
|
||||||
|
"""Retrieve all Baselines."""
|
||||||
|
result = self._post("/v1/blocklist", {})
|
||||||
|
return pd.DataFrame(result["response"]["blocklists"])
|
||||||
|
|
||||||
|
def blocklist_export(self, blocklistid) -> bytes:
|
||||||
|
"""Return Blocklist XML as bytes to save to file"""
|
||||||
|
payload = {"blocklistid": blocklistid}
|
||||||
|
result = self._post_raw("/v1/blocklist/export", payload)
|
||||||
|
return result
|
||||||
|
|
||||||
# Hash Management
|
# Hash Management
|
||||||
def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict:
|
def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict:
|
||||||
"""Add hashes to the allowlist for a specific application."""
|
"""Add hashes to the allowlist for a specific application."""
|
||||||
@@ -253,6 +300,19 @@ class AirlockAPIWrapper:
|
|||||||
}
|
}
|
||||||
return self._post("/v1/group/settings/script_custom", payload)
|
return self._post("/v1/group/settings/script_custom", payload)
|
||||||
|
|
||||||
|
def policy_set_upgradetarget(
|
||||||
|
self,
|
||||||
|
groupid: str,
|
||||||
|
windows: str,
|
||||||
|
macos: str,
|
||||||
|
) -> dict:
|
||||||
|
payload = {
|
||||||
|
"groupid": groupid,
|
||||||
|
"windows": windows,
|
||||||
|
"macos": macos,
|
||||||
|
}
|
||||||
|
return self._post("/v1/group/settings/selfupgrade/target", payload)
|
||||||
|
|
||||||
# Execution History
|
# Execution History
|
||||||
def history_logging(
|
def history_logging(
|
||||||
self, type: List[str], checkpoint: str, policy: List[str]
|
self, type: List[str], checkpoint: str, policy: List[str]
|
||||||
@@ -291,6 +351,14 @@ class AirlockAPIWrapper:
|
|||||||
result = self._post("/v1/getexechistory", payload)
|
result = self._post("/v1/getexechistory", payload)
|
||||||
return result["response"]["exechistory"]
|
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
|
from services.API import AirlockAPIWrapper
|
||||||
|
|||||||
+17
-17
@@ -28,7 +28,7 @@ from flows.prepPolicy import selectPolicies
|
|||||||
from models.agent import Agent
|
from models.agent import Agent
|
||||||
from models.policy import Policy
|
from models.policy import Policy
|
||||||
from services.API import AirlockAPIWrapper
|
from services.API import AirlockAPIWrapper
|
||||||
from utils.configmanager import get_protected_json, load_env
|
from utils.configmanager import get_system_json, load_env
|
||||||
from utils.selector import Selector
|
from utils.selector import Selector
|
||||||
from utils.utils import colorText, get_sanitized_input
|
from utils.utils import colorText, get_sanitized_input
|
||||||
|
|
||||||
@@ -38,9 +38,9 @@ logger = logging.getLogger(__name__)
|
|||||||
def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||||||
agents = selectAgents(api)
|
agents = selectAgents(api)
|
||||||
history_days = Selector.select_value(
|
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,
|
value_type=int,
|
||||||
valid_range=(1, 150),
|
valid_range=(1, 365),
|
||||||
)
|
)
|
||||||
|
|
||||||
if not agents or not history_days:
|
if not agents or not history_days:
|
||||||
@@ -60,7 +60,7 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
|
f"⌠Error retrieving history for {agent.hostname}: {e}", "red"
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
continue
|
continue
|
||||||
@@ -139,7 +139,7 @@ def findAgents(api, return_dataframe):
|
|||||||
|
|
||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
|
f"\n✓ Matched devices exported to: {working_dir}\\{filename}",
|
||||||
"green",
|
"green",
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
@@ -148,7 +148,7 @@ def findAgents(api, return_dataframe):
|
|||||||
|
|
||||||
|
|
||||||
def collect_device_names() -> List[str]:
|
def collect_device_names() -> List[str]:
|
||||||
print(colorText("🔍 Device Search", "cyan"))
|
print(colorText("🖥�� Device Search", "cyan"))
|
||||||
print(
|
print(
|
||||||
colorText(
|
colorText(
|
||||||
"Enter the device hostnames you'd like to search for, one per line.", "cyan"
|
"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:
|
else:
|
||||||
print(
|
print(
|
||||||
colorText(
|
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",
|
"yellow",
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
@@ -235,8 +235,8 @@ def show_unmatched(
|
|||||||
]
|
]
|
||||||
|
|
||||||
if unmatched:
|
if unmatched:
|
||||||
logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
|
logger.debug(f"âš ï¸ No matches for: {', '.join(unmatched)}")
|
||||||
print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
|
print(colorText(f"âš ï¸ No matches for: {', '.join(unmatched)}", "yellow"))
|
||||||
|
|
||||||
|
|
||||||
def enrich_agents(agents: List["Agent"], policies: List["Policy"]):
|
def enrich_agents(agents: List["Agent"], policies: List["Policy"]):
|
||||||
@@ -248,7 +248,7 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
|||||||
device_names = collect_device_names()
|
device_names = collect_device_names()
|
||||||
if not device_names:
|
if not device_names:
|
||||||
logger.debug("No device names entered")
|
logger.debug("No device names entered")
|
||||||
print(colorText("⚠️ No device names entered.", "red"))
|
print(colorText("âš ï¸ No device names entered.", "red"))
|
||||||
return []
|
return []
|
||||||
|
|
||||||
use_exact = choose_match_type()
|
use_exact = choose_match_type()
|
||||||
@@ -261,11 +261,11 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
|||||||
show_unmatched(device_names, matched_agents, use_exact)
|
show_unmatched(device_names, matched_agents, use_exact)
|
||||||
|
|
||||||
if not matched_agents:
|
if not matched_agents:
|
||||||
logger.debug("❌ No matching devices found.")
|
logger.debug("⌠No matching devices found.")
|
||||||
print(colorText("❌ No matching devices found.", "red"))
|
print(colorText("⌠No matching devices found.", "red"))
|
||||||
return []
|
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:")
|
logger.info("Matched agent hostnames:")
|
||||||
rows = (len(matched_agents) + 2) // 3 # 3 columns
|
rows = (len(matched_agents) + 2) // 3 # 3 columns
|
||||||
for row in range(rows):
|
for row in range(rows):
|
||||||
@@ -283,8 +283,8 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
|||||||
)
|
)
|
||||||
|
|
||||||
if not matched_agents:
|
if not matched_agents:
|
||||||
logger.debug("❌ No matching devices remain after refinement.")
|
logger.debug("⌠No matching devices remain after refinement.")
|
||||||
print(colorText("❌ No matching devices remain after refinement.", "red"))
|
print(colorText("⌠No matching devices remain after refinement.", "red"))
|
||||||
return []
|
return []
|
||||||
|
|
||||||
enrich_agents(matched_agents, policies)
|
enrich_agents(matched_agents, policies)
|
||||||
@@ -302,10 +302,10 @@ def moveAgentToRelatedPolicy(
|
|||||||
Args:
|
Args:
|
||||||
api: AirlockAPIWrapper instance.
|
api: AirlockAPIWrapper instance.
|
||||||
agent: Agent object.
|
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.
|
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
|
||||||
"""
|
"""
|
||||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||||||
|
|
||||||
if mode == "audit":
|
if mode == "audit":
|
||||||
if agent.groupid in policy_relationship_map:
|
if agent.groupid in policy_relationship_map:
|
||||||
|
|||||||
@@ -27,9 +27,8 @@ import tqdm
|
|||||||
|
|
||||||
from models.policy import Policy
|
from models.policy import Policy
|
||||||
from services.API import AirlockAPIWrapper
|
from services.API import AirlockAPIWrapper
|
||||||
from utils.configmanager import get_protected_json
|
|
||||||
from utils.setup import get_base_directory
|
from utils.setup import get_base_directory
|
||||||
from utils.utils import areYouSure, colorText, get_sanitized_input
|
from utils.utils import colorText
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -233,17 +232,3 @@ def skipback(days):
|
|||||||
hex_timestamp = format(timestamp, "08x")
|
hex_timestamp = format(timestamp, "08x")
|
||||||
objectid_hex = hex_timestamp + "0000000000000000"
|
objectid_hex = hex_timestamp + "0000000000000000"
|
||||||
return ObjectId(objectid_hex)
|
return ObjectId(objectid_hex)
|
||||||
|
|
||||||
|
|
||||||
def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
|
|
||||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
|
||||||
for enforcement_policy, audit_policy in policy_relationship_map.items():
|
|
||||||
api.policy_clone(enforcement_policy, audit_policy)
|
|
||||||
api.policy_set_auditmode(audit_policy, "1")
|
|
||||||
|
|
||||||
|
|
||||||
def confirmUpdateAfromE(api: AirlockAPIWrapper):
|
|
||||||
areYouSure()
|
|
||||||
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
|
||||||
if confirmation.strip() == "I AGREE":
|
|
||||||
updateAuditPoliciesFromEnforcementPolices(api)
|
|
||||||
|
|||||||
+249
-43
@@ -18,117 +18,280 @@ import logging
|
|||||||
import os
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import sys
|
import sys
|
||||||
from typing import Callable, Optional, TypeVar
|
from typing import Any, Callable, Optional, TypeVar
|
||||||
|
|
||||||
T = TypeVar("T")
|
T = TypeVar("T")
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
PROTECTED_KEYS = [
|
# System config keys - these are immutable and come from system_config.json (bundled in exe)
|
||||||
|
SYSTEM_CONFIG_KEYS = [
|
||||||
|
"URL",
|
||||||
"APPNAME",
|
"APPNAME",
|
||||||
"LOG_LEVEL",
|
"LOG_LEVEL",
|
||||||
|
"BAD_PATH_PARTS",
|
||||||
|
"BAD_PUBLISHERS",
|
||||||
|
"PUPS",
|
||||||
"PATH_EXCLUSION_CONST",
|
"PATH_EXCLUSION_CONST",
|
||||||
"MIN_FILES_FOR_PATH",
|
"MIN_FILES_FOR_PATH",
|
||||||
"VT_THREAT_TOLERANCE",
|
"VT_THREAT_TOLERANCE",
|
||||||
"POLICY_MAP_ENF_AUD",
|
"POLICY_MAP_ENF_AUD",
|
||||||
]
|
]
|
||||||
|
|
||||||
_protected_config = {}
|
# User config keys - these can be changed by the end user
|
||||||
|
USER_CONFIG_KEYS = [
|
||||||
|
"TELEMETRY", # User opt-in/out for telemetry
|
||||||
|
"TELEM_URL",
|
||||||
|
"TEXTUAL_THEME", # UI theme preference
|
||||||
|
"EXTRAS", # Feature flags
|
||||||
|
]
|
||||||
|
|
||||||
|
# In-memory config storage
|
||||||
|
_system_config = {}
|
||||||
|
_user_config = {}
|
||||||
|
|
||||||
|
|
||||||
def get_system_config_path() -> Path:
|
def get_system_config_path() -> Path:
|
||||||
|
"""
|
||||||
|
Get path to system_config.json.
|
||||||
|
Priority:
|
||||||
|
1. Bundled in exe (_MEIPASS)
|
||||||
|
2. Next to this file (development)
|
||||||
|
"""
|
||||||
# Check inside bundled EXE directory first
|
# Check inside bundled EXE directory first
|
||||||
bundled_dir = Path(getattr(sys, "_MEIPASS", ""))
|
bundled_dir = Path(getattr(sys, "_MEIPASS", ""))
|
||||||
bundled_path = bundled_dir / "system_config.json"
|
bundled_path = bundled_dir / "system_config.json"
|
||||||
if bundled_path.exists():
|
if bundled_path.exists():
|
||||||
return bundled_path
|
return bundled_path
|
||||||
|
|
||||||
# Fallback to external location
|
# Fallback to development location (next to this file)
|
||||||
return Path(__file__).parent.parent / "system_config.json"
|
return Path(__file__).parent.parent / "system_config.json"
|
||||||
|
|
||||||
|
|
||||||
def load_protected_config() -> dict:
|
def load_system_config() -> dict:
|
||||||
global _protected_config
|
"""
|
||||||
|
Load system configuration from system_config.json.
|
||||||
|
This should only be called once at startup.
|
||||||
|
Returns the full system config dict.
|
||||||
|
"""
|
||||||
|
global _system_config
|
||||||
|
|
||||||
try:
|
try:
|
||||||
with open(get_system_config_path(), "r") as f:
|
config_path = get_system_config_path()
|
||||||
system_config = json.load(f)
|
with open(config_path, "r") as f:
|
||||||
|
_system_config = json.load(f)
|
||||||
|
logger.debug(f"✅ Loaded system config from {config_path}")
|
||||||
except FileNotFoundError:
|
except FileNotFoundError:
|
||||||
logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
|
logger.warning("⚠️ system_config.json not found. Using minimal defaults.")
|
||||||
system_config = {
|
# Minimal defaults for development without system_config.json
|
||||||
|
_system_config = {
|
||||||
"APPNAME": "Loxide",
|
"APPNAME": "Loxide",
|
||||||
|
"LOG_LEVEL": "INFO",
|
||||||
"PATH_EXCLUSION_CONST": 4,
|
"PATH_EXCLUSION_CONST": 4,
|
||||||
"MIN_FILES_FOR_PATH": 4,
|
"MIN_FILES_FOR_PATH": 4,
|
||||||
"VT_THREAT_TOLERANCE": 4,
|
"VT_THREAT_TOLERANCE": 4,
|
||||||
"POLICY_MAP_ENF_AUD": {"enforced_id": "audit_id"},
|
"POLICY_MAP_ENF_AUD": {},
|
||||||
}
|
}
|
||||||
|
|
||||||
_protected_config = {key: system_config[key] for key in PROTECTED_KEYS}
|
return _system_config
|
||||||
return _protected_config
|
|
||||||
|
|
||||||
|
|
||||||
def get_protected_value(
|
def get_system_value(
|
||||||
key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None
|
key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None
|
||||||
) -> Optional[T]:
|
) -> Optional[T]:
|
||||||
value = _protected_config.get(key)
|
"""
|
||||||
|
Get a value from system config (immutable).
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
key: The config key to retrieve
|
||||||
|
cast_type: Function to cast the value to desired type
|
||||||
|
default: Default value if key not found
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
The config value cast to the desired type, or default
|
||||||
|
"""
|
||||||
|
value = _system_config.get(key)
|
||||||
if value is None:
|
if value is None:
|
||||||
logging.warning(f"Protected config key '{key}' not found.")
|
logger.warning(f"System config key '{key}' not found.")
|
||||||
return default
|
return default
|
||||||
|
|
||||||
try:
|
try:
|
||||||
if isinstance(value, str):
|
if isinstance(value, str):
|
||||||
value = value.strip("'\"")
|
value = value.strip("'\"")
|
||||||
return cast_type(value)
|
return cast_type(value)
|
||||||
except (ValueError, TypeError):
|
except (ValueError, TypeError):
|
||||||
logging.warning(
|
logger.warning(
|
||||||
f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}."
|
f"Invalid value for system key '{key}': {value}. Expected type {cast_type.__name__}."
|
||||||
)
|
)
|
||||||
return default
|
return default
|
||||||
|
|
||||||
|
|
||||||
def get_protected_json(key: str, default: str = "{}") -> dict:
|
def get_system_json(key: str, default: Optional[dict] = None) -> dict:
|
||||||
raw = _protected_config.get(key, default)
|
"""
|
||||||
|
Get a JSON/dict value from system config.
|
||||||
|
Handles both dict values and JSON strings.
|
||||||
|
"""
|
||||||
|
if default is None:
|
||||||
|
default = {}
|
||||||
|
|
||||||
|
raw = _system_config.get(key, default)
|
||||||
if isinstance(raw, dict):
|
if isinstance(raw, dict):
|
||||||
return raw
|
return raw
|
||||||
|
|
||||||
try:
|
try:
|
||||||
return json.loads(raw)
|
return json.loads(raw)
|
||||||
except json.JSONDecodeError:
|
except (json.JSONDecodeError, TypeError) as e:
|
||||||
try:
|
logger.error(f"Failed to parse system JSON key '{key}': {e}")
|
||||||
escaped = raw.encode("unicode_escape").decode("utf-8")
|
return default
|
||||||
return json.loads(escaped)
|
|
||||||
except Exception as e:
|
|
||||||
logging.error(f"Failed to parse protected JSON key '{key}': {e}")
|
|
||||||
return json.loads(default)
|
|
||||||
|
|
||||||
|
|
||||||
def load_env_json(key: str, default: str):
|
def get_system_list(key: str, default: Optional[list] = None) -> list:
|
||||||
raw = os.getenv(key, default)
|
"""
|
||||||
|
Get a list value from system config.
|
||||||
|
Handles both list values and JSON strings.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
key: The config key to retrieve
|
||||||
|
default: Default value if key not found or parsing fails
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
The list value or default
|
||||||
|
"""
|
||||||
|
if default is None:
|
||||||
|
default = []
|
||||||
|
|
||||||
|
raw = _system_config.get(key, default)
|
||||||
|
if isinstance(raw, list):
|
||||||
|
return raw
|
||||||
|
|
||||||
try:
|
try:
|
||||||
return json.loads(raw)
|
result = json.loads(raw) if isinstance(raw, str) else raw
|
||||||
except json.JSONDecodeError:
|
if isinstance(result, list):
|
||||||
try:
|
return result
|
||||||
escaped = raw.encode("unicode_escape").decode("utf-8")
|
logger.warning(f"System config key '{key}' is not a list: {type(result)}")
|
||||||
return json.loads(escaped)
|
return default
|
||||||
except Exception as e:
|
except (json.JSONDecodeError, TypeError) as e:
|
||||||
logging.error(f"Failed to parse {key}: {e}")
|
logger.error(f"Failed to parse system list key '{key}': {e}")
|
||||||
return json.loads(default)
|
return default
|
||||||
|
|
||||||
|
|
||||||
|
def load_user_config(config_dir: Path) -> dict:
|
||||||
|
"""
|
||||||
|
Load user configuration from user_config.json.
|
||||||
|
Creates the file with defaults if it doesn't exist.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
config_dir: Directory containing user_config.json
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
The user config dict
|
||||||
|
"""
|
||||||
|
global _user_config
|
||||||
|
|
||||||
|
user_config_path = config_dir / "user_config.json"
|
||||||
|
|
||||||
|
if not user_config_path.exists():
|
||||||
|
# Create default user config
|
||||||
|
default_user_config = {
|
||||||
|
"TELEMETRY": False,
|
||||||
|
"TELEM_URL": "",
|
||||||
|
"TEXTUAL_THEME": "gruvbox",
|
||||||
|
"EXTRAS": "NOTTODAY",
|
||||||
|
}
|
||||||
|
user_config_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
with open(user_config_path, "w") as f:
|
||||||
|
json.dump(default_user_config, f, indent=4)
|
||||||
|
logger.debug(f"Created default user config at {user_config_path}")
|
||||||
|
_user_config = default_user_config
|
||||||
|
else:
|
||||||
|
with open(user_config_path, "r") as f:
|
||||||
|
_user_config = json.load(f)
|
||||||
|
logger.debug(f"✅ Loaded user config from {user_config_path}")
|
||||||
|
|
||||||
|
return _user_config
|
||||||
|
|
||||||
|
|
||||||
|
def save_user_config(config_dir: Path, updates: dict) -> None:
|
||||||
|
"""
|
||||||
|
Save updates to user configuration.
|
||||||
|
Only keys in USER_CONFIG_KEYS are allowed.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
config_dir: Directory containing user_config.json
|
||||||
|
updates: Dict of key-value pairs to update
|
||||||
|
"""
|
||||||
|
global _user_config
|
||||||
|
|
||||||
|
# Validate that only user-configurable keys are being updated
|
||||||
|
invalid_keys = [k for k in updates.keys() if k not in USER_CONFIG_KEYS]
|
||||||
|
if invalid_keys:
|
||||||
|
logger.error(f"Attempted to save invalid user config keys: {invalid_keys}")
|
||||||
|
raise ValueError(f"Cannot modify system config keys: {invalid_keys}")
|
||||||
|
|
||||||
|
# Update in-memory config
|
||||||
|
_user_config.update(updates)
|
||||||
|
|
||||||
|
# Write to file
|
||||||
|
user_config_path = config_dir / "user_config.json"
|
||||||
|
user_config_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
with open(user_config_path, "w") as f:
|
||||||
|
json.dump(_user_config, f, indent=4)
|
||||||
|
|
||||||
|
logger.debug(f"✅ Saved user config to {user_config_path}: {updates}")
|
||||||
|
|
||||||
|
|
||||||
|
def get_user_value(
|
||||||
|
key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None
|
||||||
|
) -> Optional[T]:
|
||||||
|
"""
|
||||||
|
Get a value from user config (mutable).
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
key: The config key to retrieve
|
||||||
|
cast_type: Function to cast the value to desired type
|
||||||
|
default: Default value if key not found
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
The config value cast to the desired type, or default
|
||||||
|
"""
|
||||||
|
value = _user_config.get(key)
|
||||||
|
if value is None:
|
||||||
|
logger.warning(f"User config key '{key}' not found.")
|
||||||
|
return default
|
||||||
|
|
||||||
|
try:
|
||||||
|
if isinstance(value, str):
|
||||||
|
value = value.strip("'\"")
|
||||||
|
return cast_type(value)
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
logger.warning(
|
||||||
|
f"Invalid value for user key '{key}': {value}. Expected type {cast_type.__name__}."
|
||||||
|
)
|
||||||
|
return default
|
||||||
|
|
||||||
|
|
||||||
def load_env(
|
def load_env(
|
||||||
key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None
|
key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None
|
||||||
) -> Optional[T]:
|
) -> Optional[T]:
|
||||||
"""
|
"""
|
||||||
Safely retrieves an environment variable and casts it to the desired type.
|
Safely retrieves an environment variable from .env and casts it to the desired type.
|
||||||
|
This should ONLY be used for runtime/dynamic values like WORKING_DIR.
|
||||||
|
|
||||||
|
For system config, use get_system_value().
|
||||||
|
For user config, use get_user_value().
|
||||||
|
|
||||||
Parameters:
|
Parameters:
|
||||||
key (str): The name of the environment variable.
|
key: The name of the environment variable
|
||||||
cast_type (Callable[[str], T], optional): Function to cast the value. Defaults to str.
|
cast_type: Function to cast the value. Defaults to str
|
||||||
default (Optional[T], optional): Default value if the variable is not set or invalid.
|
default: Default value if the variable is not set or invalid
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Optional[T]: The casted value or the default.
|
The casted value or the default
|
||||||
"""
|
"""
|
||||||
value = os.getenv(key)
|
value = os.getenv(key)
|
||||||
if value is None:
|
if value is None:
|
||||||
logger.warning(f"Environment variable '{key}' not set.")
|
logger.debug(f"Environment variable '{key}' not set, using default.")
|
||||||
return default
|
return default
|
||||||
|
|
||||||
try:
|
try:
|
||||||
value = value.strip("'\"") # Strip surrounding quotes
|
value = value.strip("'\"") # Strip surrounding quotes
|
||||||
return cast_type(value)
|
return cast_type(value)
|
||||||
@@ -137,3 +300,46 @@ def load_env(
|
|||||||
f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}."
|
f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}."
|
||||||
)
|
)
|
||||||
return default
|
return default
|
||||||
|
|
||||||
|
|
||||||
|
def load_env_json(key: str, default: str = "[]") -> Any:
|
||||||
|
"""
|
||||||
|
Load a JSON value from environment or system config.
|
||||||
|
|
||||||
|
DEPRECATED: This function is kept for backward compatibility.
|
||||||
|
- For system config lists (BAD_PUBLISHERS, PUPS, BAD_PATH_PARTS), use get_system_list()
|
||||||
|
- For system config dicts, use get_system_json()
|
||||||
|
- For actual .env JSON values, parse manually
|
||||||
|
|
||||||
|
This function automatically redirects known system config keys to system config.
|
||||||
|
"""
|
||||||
|
# Known system config list keys - redirect to system config
|
||||||
|
system_list_keys = ["BAD_PUBLISHERS", "PUPS", "BAD_PATH_PARTS"]
|
||||||
|
if key in system_list_keys:
|
||||||
|
logger.debug(f"Redirecting load_env_json('{key}') to get_system_list()")
|
||||||
|
return get_system_list(key, json.loads(default) if default else [])
|
||||||
|
|
||||||
|
# Known system config dict keys - redirect to system config
|
||||||
|
system_dict_keys = ["POLICY_MAP_ENF_AUD"]
|
||||||
|
if key in system_dict_keys:
|
||||||
|
logger.debug(f"Redirecting load_env_json('{key}') to get_system_json()")
|
||||||
|
return get_system_json(key, json.loads(default) if default else {})
|
||||||
|
|
||||||
|
# Fall back to reading from .env (backward compatibility for unknown keys)
|
||||||
|
raw = os.getenv(key, default)
|
||||||
|
try:
|
||||||
|
return json.loads(raw)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
try:
|
||||||
|
escaped = raw.encode("unicode_escape").decode("utf-8")
|
||||||
|
return json.loads(escaped)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to parse {key}: {e}")
|
||||||
|
return json.loads(default)
|
||||||
|
|
||||||
|
|
||||||
|
# Backwards compatibility aliases (deprecated - use get_system_value instead)
|
||||||
|
get_protected_value = get_system_value
|
||||||
|
get_protected_json = get_system_json
|
||||||
|
load_protected_config = load_system_config
|
||||||
|
PROTECTED_KEYS = SYSTEM_CONFIG_KEYS # For backwards compatibility
|
||||||
|
|||||||
+8
-4
@@ -184,7 +184,7 @@ class Selector:
|
|||||||
print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
|
print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
|
||||||
return selected
|
return selected
|
||||||
elif mode == "e":
|
elif mode == "e":
|
||||||
print(colorText(f"🚫 Excluded {len(selected)} item(s).", "yellow"))
|
print(colorText(f"👫 Excluded {len(selected)} item(s).", "yellow"))
|
||||||
return [item for item in items if item not in selected]
|
return [item for item in items if item not in selected]
|
||||||
else:
|
else:
|
||||||
print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
|
print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
|
||||||
@@ -279,7 +279,9 @@ class Selector:
|
|||||||
df = df[columns]
|
df = df[columns]
|
||||||
|
|
||||||
items = [row for _, row in df.iterrows()]
|
items = [row for _, row in df.iterrows()]
|
||||||
label_func = lambda row: str(row.to_dict())
|
|
||||||
|
def label_func(row):
|
||||||
|
return str(row.to_dict())
|
||||||
|
|
||||||
result = Selector._select_from_list(
|
result = Selector._select_from_list(
|
||||||
items,
|
items,
|
||||||
@@ -311,7 +313,9 @@ class Selector:
|
|||||||
df = df[columns]
|
df = df[columns]
|
||||||
|
|
||||||
items = df.to_dict("records")
|
items = df.to_dict("records")
|
||||||
label_func = lambda row: " | ".join(str(row[col]) for col in df.columns)
|
|
||||||
|
def label_func(row):
|
||||||
|
return " | ".join(str(row[col]) for col in df.columns)
|
||||||
|
|
||||||
# Show rows first
|
# Show rows first
|
||||||
print(colorText(header, "cyan"))
|
print(colorText(header, "cyan"))
|
||||||
@@ -346,7 +350,7 @@ class Selector:
|
|||||||
print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
|
print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
|
||||||
return [pd.Series(row) for row in selected]
|
return [pd.Series(row) for row in selected]
|
||||||
elif mode == "e":
|
elif mode == "e":
|
||||||
print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow"))
|
print(colorText(f"👫 Excluded {len(selected)} row(s).", "yellow"))
|
||||||
return [pd.Series(row) for row in items if row not in selected]
|
return [pd.Series(row) for row in items if row not in selected]
|
||||||
else:
|
else:
|
||||||
print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
|
print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
|
||||||
|
|||||||
+100
-94
@@ -13,18 +13,63 @@
|
|||||||
# You should have received a copy of the GNU Affero General Public License
|
# 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/>.
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||||
|
|
||||||
import json
|
|
||||||
import logging
|
import logging
|
||||||
import logging.config
|
import logging.config
|
||||||
import logging.handlers
|
import logging.handlers
|
||||||
import os
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import platform
|
import platform
|
||||||
import sys
|
|
||||||
|
|
||||||
from dotenv import load_dotenv, set_key
|
from dotenv import load_dotenv, set_key
|
||||||
|
|
||||||
from utils.configmanager import PROTECTED_KEYS, load_protected_config
|
from utils.configmanager import (
|
||||||
|
get_system_value,
|
||||||
|
load_system_config,
|
||||||
|
load_user_config,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
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:
|
def get_base_directory() -> Path:
|
||||||
@@ -38,42 +83,35 @@ def get_base_directory() -> Path:
|
|||||||
return home / ".local" / "share" / "Loxide"
|
return home / ".local" / "share" / "Loxide"
|
||||||
|
|
||||||
|
|
||||||
def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
|
def configure_logging(log_dir: Path, log_level: str = "INFO"):
|
||||||
log_file = log_dir / "Loxide.log"
|
log_file = log_dir / "Loxide.log"
|
||||||
|
|
||||||
config = {
|
config = {
|
||||||
"version": 1, # Required key for dictConfig format version
|
"version": 1,
|
||||||
"disable_existing_loggers": False, # Keeps existing loggers active
|
"disable_existing_loggers": False,
|
||||||
"formatters": {
|
"formatters": {
|
||||||
"detailed": {
|
"detailed": {
|
||||||
"format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
"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": {
|
"handlers": {
|
||||||
"file": {
|
"file": {
|
||||||
"class": "logging.handlers.TimedRotatingFileHandler",
|
"class": "logging.handlers.TimedRotatingFileHandler",
|
||||||
"filename": str(log_file),
|
"filename": str(log_file),
|
||||||
"when": "midnight", # Rotate logs at midnight
|
"when": "midnight",
|
||||||
"interval": 1, # Every 1 day
|
"interval": 1,
|
||||||
"backupCount": 7, # Keep 7 days of logs
|
"backupCount": 7,
|
||||||
"encoding": "utf-8", # Ensure UTF-8 encoding
|
"encoding": "utf-8",
|
||||||
"level": "DEBUG", # Always log DEBUG and above
|
"level": "DEBUG",
|
||||||
"formatter": "detailed", # Use detailed format
|
"formatter": "detailed",
|
||||||
},
|
|
||||||
"console": {
|
|
||||||
"class": "logging.StreamHandler",
|
|
||||||
"level": log_level.upper(), # Configurable log level
|
|
||||||
"formatter": "simple", # Use simple format
|
|
||||||
},
|
},
|
||||||
|
# REMOVED console handler - it interferes with Textual TUI
|
||||||
},
|
},
|
||||||
"root": {
|
"root": {
|
||||||
"level": "DEBUG", # Root logger level
|
"level": "DEBUG",
|
||||||
"handlers": ["file", "console"], # Attach both handlers
|
"handlers": ["file"], # Only use file handler, not console
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -94,56 +132,31 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
|
|||||||
logging.config.dictConfig(config)
|
logging.config.dictConfig(config)
|
||||||
logging.getLogger().debug("✅ Logging configured.")
|
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.")
|
||||||
|
|
||||||
def get_system_config_path() -> Path:
|
return attach_notification_handler
|
||||||
base_path = Path(
|
|
||||||
getattr(sys, "_MEIPASS", os.path.dirname(os.path.abspath(__file__)))
|
|
||||||
)
|
|
||||||
return base_path.parent / "system_config.json"
|
|
||||||
|
|
||||||
|
|
||||||
def load_system_config() -> dict:
|
|
||||||
try:
|
|
||||||
config_path = get_system_config_path()
|
|
||||||
with open(config_path, "r") as f:
|
|
||||||
return json.load(f)
|
|
||||||
except FileNotFoundError:
|
|
||||||
logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
|
|
||||||
return {
|
|
||||||
"APPNAME": "Loxide",
|
|
||||||
"LOG_LEVEL": "DEBUG",
|
|
||||||
"PATH_EXCLUSION_CONST": 4,
|
|
||||||
"MIN_FILES_FOR_PATH": 4,
|
|
||||||
"VT_THREAT_TOLERANCE": 4,
|
|
||||||
"POLICY_MAP_ENF_AUD": {"enforced_id": "audit_id"},
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def load_user_config(config_dir: Path) -> dict:
|
|
||||||
user_config_path = config_dir / "user_config.json"
|
|
||||||
if not user_config_path.exists():
|
|
||||||
default_user_config = {"URL": "", "LOG_LEVEL": "INFO"}
|
|
||||||
with open(user_config_path, "w") as f:
|
|
||||||
json.dump(default_user_config, f, indent=4)
|
|
||||||
logging.debug(f"Created user config at {user_config_path}")
|
|
||||||
with open(user_config_path, "r") as f:
|
|
||||||
return json.load(f)
|
|
||||||
|
|
||||||
|
|
||||||
def write_config_to_env(config: dict, env_path: Path):
|
|
||||||
for key, value in config.items():
|
|
||||||
if key in PROTECTED_KEYS:
|
|
||||||
continue # Skip protected keys
|
|
||||||
try:
|
|
||||||
serialized = (
|
|
||||||
json.dumps(value) if isinstance(value, (list, dict)) else str(value)
|
|
||||||
)
|
|
||||||
set_key(env_path, key, serialized)
|
|
||||||
except Exception as e:
|
|
||||||
logging.warning(f"Failed to write {key} to .env: {e}")
|
|
||||||
|
|
||||||
|
|
||||||
def setup():
|
def setup():
|
||||||
|
"""
|
||||||
|
Initialize the application environment:
|
||||||
|
1. Create directory structure
|
||||||
|
2. Load system config (immutable, from system_config.json)
|
||||||
|
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()
|
base_dir = get_base_directory()
|
||||||
dirs = {
|
dirs = {
|
||||||
"config": base_dir / "config",
|
"config": base_dir / "config",
|
||||||
@@ -155,25 +168,34 @@ def setup():
|
|||||||
path.mkdir(parents=True, exist_ok=True)
|
path.mkdir(parents=True, exist_ok=True)
|
||||||
logging.debug(f"{name.capitalize()} directory ensured at: {path}")
|
logging.debug(f"{name.capitalize()} directory ensured at: {path}")
|
||||||
|
|
||||||
system_config = load_system_config()
|
# Load system config (immutable)
|
||||||
configure_logging(dirs["logs"], system_config.get("LOG_LEVEL", "DEBUG"))
|
load_system_config()
|
||||||
|
|
||||||
|
# Configure logging with system-defined log level
|
||||||
|
log_level = get_system_value("LOG_LEVEL", str, "INFO")
|
||||||
|
attach_handler = configure_logging(dirs["logs"], log_level)
|
||||||
|
|
||||||
|
# Load user config (mutable)
|
||||||
|
load_user_config(dirs["config"])
|
||||||
|
|
||||||
|
# Set up .env file - ONLY for WORKING_DIR (runtime-configurable value)
|
||||||
env_path = base_dir / ".env"
|
env_path = base_dir / ".env"
|
||||||
if not env_path.exists():
|
if not env_path.exists():
|
||||||
env_path.touch()
|
env_path.touch()
|
||||||
load_dotenv(dotenv_path=env_path, override=True)
|
load_dotenv(dotenv_path=env_path, override=True)
|
||||||
|
|
||||||
|
# Set up working directory (only dynamic value in .env)
|
||||||
working_dir = Path(os.getenv("WORKING_DIR") or (base_dir / "data"))
|
working_dir = Path(os.getenv("WORKING_DIR") or (base_dir / "data"))
|
||||||
working_dir.mkdir(parents=True, exist_ok=True)
|
working_dir.mkdir(parents=True, exist_ok=True)
|
||||||
set_key(env_path, "WORKING_DIR", str(working_dir))
|
set_key(str(env_path), "WORKING_DIR", str(working_dir))
|
||||||
os.environ["WORKING_DIR"] = str(working_dir)
|
os.environ["WORKING_DIR"] = str(working_dir)
|
||||||
logging.debug(f"Working directory set to: {working_dir}")
|
logging.debug(f"Working directory set to: {working_dir}")
|
||||||
|
|
||||||
|
# Create folder structure in working directory
|
||||||
folders_structure = {
|
folders_structure = {
|
||||||
"Approved": [],
|
"Approved": [],
|
||||||
"Needs_Review": ["Review_First", "Review_Second", "HTML"],
|
"Needs_Review": ["Review_First", "Review_Second", "HTML"],
|
||||||
"Preflight": ["HTML"],
|
"Preflight": ["HTML"],
|
||||||
"Archived": [],
|
|
||||||
}
|
}
|
||||||
|
|
||||||
for folder_name, subfolders in folders_structure.items():
|
for folder_name, subfolders in folders_structure.items():
|
||||||
@@ -183,25 +205,9 @@ def setup():
|
|||||||
for subfolder in subfolders:
|
for subfolder in subfolders:
|
||||||
subfolder_path = folder_path / subfolder
|
subfolder_path = folder_path / subfolder
|
||||||
subfolder_path.mkdir(parents=True, exist_ok=True)
|
subfolder_path.mkdir(parents=True, exist_ok=True)
|
||||||
logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
|
logging.debug(f"'{subfolder}' subfolder created at: {subfolder_path}")
|
||||||
|
|
||||||
user_config = load_user_config(dirs["config"])
|
logging.info("✅ Setup complete")
|
||||||
merged_config = {**system_config, **user_config}
|
|
||||||
|
|
||||||
protected_config = load_protected_config()
|
# Return the attach handler function
|
||||||
merged_config.update(protected_config)
|
return attach_handler
|
||||||
|
|
||||||
# ✅ URL resolution order: system_config → .env → user prompt
|
|
||||||
url = system_config.get("URL")
|
|
||||||
if not url:
|
|
||||||
url = os.getenv("URL")
|
|
||||||
if not url:
|
|
||||||
url = input(
|
|
||||||
"🌐 Enter the service URL (e.g., https://example.com/api): "
|
|
||||||
).strip()
|
|
||||||
merged_config["URL"] = url
|
|
||||||
set_key(env_path, "URL", url)
|
|
||||||
os.environ["URL"] = url
|
|
||||||
logging.debug(f"Service URL set to: {url}")
|
|
||||||
|
|
||||||
write_config_to_env(merged_config, env_path)
|
|
||||||
|
|||||||
-516
@@ -1,516 +0,0 @@
|
|||||||
import logging
|
|
||||||
import os
|
|
||||||
import sys
|
|
||||||
|
|
||||||
import dotenv
|
|
||||||
from dotenv import set_key
|
|
||||||
from textual.app import App, ComposeResult
|
|
||||||
from textual.containers import Vertical
|
|
||||||
from textual.reactive import reactive
|
|
||||||
from textual.screen import Screen
|
|
||||||
from textual.widgets import (
|
|
||||||
Button,
|
|
||||||
DirectoryTree,
|
|
||||||
Footer,
|
|
||||||
Header,
|
|
||||||
Static,
|
|
||||||
Tab,
|
|
||||||
Tabs,
|
|
||||||
)
|
|
||||||
|
|
||||||
from flows.otp import otp_activities_by_agent, otp_revoke
|
|
||||||
from flows.prepPolicy import menu_policy_enforce
|
|
||||||
from flows.quietAgent import findQuietAgents
|
|
||||||
from models.agent import Agent
|
|
||||||
from models.policy import Policy
|
|
||||||
from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
|
||||||
from screens.otpworkflowscreen import OTPWorkflowScreen
|
|
||||||
from services.API import AirlockAPIWrapper
|
|
||||||
from services.policyhandler import confirmUpdateAfromE
|
|
||||||
from utils.configmanager import load_env
|
|
||||||
from utils.setup import get_base_directory, load_user_config
|
|
||||||
from utils.utils import open_directory
|
|
||||||
from widgets.agentmoveoperations import AgentMoveOperations
|
|
||||||
from widgets.amber_terminal_theme import get_amber_terminal_theme
|
|
||||||
from widgets.multiagentselector import MultiAgentSelector
|
|
||||||
from widgets.OTP_generate import OTPGenerator
|
|
||||||
from widgets.policytreewidget import PolicyTreeWidget
|
|
||||||
from widgets.resultsdisplay import ResultsDisplay
|
|
||||||
from widgets.retro_terminal_theme import get_retro_terminal_theme
|
|
||||||
from widgets.themeselector import ThemeSelector
|
|
||||||
|
|
||||||
dotenv.load_dotenv()
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# GLOBAL STASH
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
_PENDING_JOB = None
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
def _persist_user_theme(theme_name: str) -> None:
|
|
||||||
"""
|
|
||||||
Store the chosen Textual theme in the user's config:
|
|
||||||
<base>/config/user_config.json
|
|
||||||
and also mirror to <base>/.env so load_env(...) sees it.
|
|
||||||
"""
|
|
||||||
base_dir = get_base_directory()
|
|
||||||
config_dir = base_dir / "config"
|
|
||||||
user_config_path = config_dir / "user_config.json"
|
|
||||||
env_path = base_dir / ".env"
|
|
||||||
|
|
||||||
# ensure dirs / files exist similarly to setup()
|
|
||||||
config_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
if not user_config_path.exists():
|
|
||||||
# minimal default like your load_user_config does
|
|
||||||
user_config_path.write_text(
|
|
||||||
'{"URL": "", "LOG_LEVEL": "INFO"}\n', encoding="utf-8"
|
|
||||||
)
|
|
||||||
|
|
||||||
# load existing user config
|
|
||||||
user_conf = load_user_config(config_dir)
|
|
||||||
user_conf["TEXTUAL_THEME"] = theme_name
|
|
||||||
|
|
||||||
# write it back
|
|
||||||
user_config_path.write_text(
|
|
||||||
# pretty print so it stays human-readable
|
|
||||||
__import__("json").dumps(user_conf, indent=4),
|
|
||||||
encoding="utf-8",
|
|
||||||
)
|
|
||||||
logger.debug("Updated user_config.json with TEXTUAL_THEME=%s", theme_name)
|
|
||||||
|
|
||||||
# mirror to .env (like setup.write_config_to_env does)
|
|
||||||
env_path.parent.mkdir(parents=True, exist_ok=True)
|
|
||||||
if not env_path.exists():
|
|
||||||
env_path.touch()
|
|
||||||
try:
|
|
||||||
set_key(str(env_path), "TEXTUAL_THEME", theme_name)
|
|
||||||
except Exception as exc: # keep going even if .env write fails
|
|
||||||
logger.warning("Failed to mirror TEXTUAL_THEME to .env: %s", exc)
|
|
||||||
|
|
||||||
# reload so load_env(...) sees the new value right now
|
|
||||||
dotenv.load_dotenv(dotenv_path=env_path, override=True)
|
|
||||||
logger.debug("Reloaded .env from %s", env_path)
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# 1) SCREEN
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
class MainMenuScreen(Screen):
|
|
||||||
current_tab = reactive("")
|
|
||||||
|
|
||||||
BUTTON_DEFS = {
|
|
||||||
"agent_actions": [
|
|
||||||
(
|
|
||||||
"🖥️ - Find, Move, or Generate OTP for Agents",
|
|
||||||
"move_agent_workflow_button",
|
|
||||||
),
|
|
||||||
("🔇 - Find Quiet Hosts", "find_quiet_button"),
|
|
||||||
],
|
|
||||||
"policy": [
|
|
||||||
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
|
|
||||||
("🔄 - Update Audit Policies", "policy_audit_update_button"),
|
|
||||||
("📊 - OTP Activities By Agent", "otp_activities_button"),
|
|
||||||
("❌ - Revoke OTPs", "otp_revoke_button"),
|
|
||||||
],
|
|
||||||
}
|
|
||||||
|
|
||||||
def __init__(self, api: AirlockAPIWrapper) -> None:
|
|
||||||
super().__init__()
|
|
||||||
self.api = api
|
|
||||||
self.extras = load_env("EXTRAS")
|
|
||||||
wd = load_env("WORKING_DIR") or os.getcwd()
|
|
||||||
if not os.path.isdir(wd):
|
|
||||||
wd = os.getcwd()
|
|
||||||
self.working_dir = wd
|
|
||||||
|
|
||||||
def _make_buttons_for(self, tab_id: str) -> Vertical:
|
|
||||||
defs = self.BUTTON_DEFS.get(tab_id, [])
|
|
||||||
buttons = []
|
|
||||||
for label, btn_id in defs:
|
|
||||||
btn = Button(label, id=btn_id)
|
|
||||||
btn.styles.width = "100%"
|
|
||||||
buttons.append(btn)
|
|
||||||
return Vertical(*buttons)
|
|
||||||
|
|
||||||
def compose(self) -> ComposeResult:
|
|
||||||
yield Header(show_clock=True, icon="⚙")
|
|
||||||
|
|
||||||
tabs = [
|
|
||||||
Tab("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 _PENDING_JOB
|
|
||||||
selected_agents = message.selected_agents
|
|
||||||
logger.info("Selected agents: %s", selected_agents)
|
|
||||||
# TODO: Implement actual handling of selected agents
|
|
||||||
_PENDING_JOB = ("multi_agent_action", selected_agents)
|
|
||||||
self.app.exit()
|
|
||||||
|
|
||||||
def on_theme_selector_theme_selected(
|
|
||||||
self, message: ThemeSelector.ThemeSelected
|
|
||||||
) -> None:
|
|
||||||
"""Handle theme selection from ThemeSelector."""
|
|
||||||
global _PENDING_JOB
|
|
||||||
_persist_user_theme(message.theme_name)
|
|
||||||
_PENDING_JOB = ("restart",)
|
|
||||||
self.app.exit()
|
|
||||||
|
|
||||||
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
|
|
||||||
"""Handle OTP generation request from the workflow."""
|
|
||||||
global _PENDING_JOB
|
|
||||||
|
|
||||||
# Log what we received
|
|
||||||
logger.info(
|
|
||||||
"OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d",
|
|
||||||
len(message.devices),
|
|
||||||
message.requestor,
|
|
||||||
message.reasoning,
|
|
||||||
message.duration,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Set up the job to run the OTP generation
|
|
||||||
_PENDING_JOB = (
|
|
||||||
"otp_workflow",
|
|
||||||
message.devices,
|
|
||||||
message.requestor,
|
|
||||||
message.reasoning,
|
|
||||||
message.duration,
|
|
||||||
)
|
|
||||||
|
|
||||||
self.app.exit()
|
|
||||||
|
|
||||||
def on_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:
|
|
||||||
global _PENDING_JOB
|
|
||||||
button_id = event.button.id
|
|
||||||
logger.debug("Button pressed: %s", button_id)
|
|
||||||
|
|
||||||
match button_id:
|
|
||||||
case "move_agent_workflow_button":
|
|
||||||
# Push Move Agent workflow screen
|
|
||||||
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
|
|
||||||
event.stop()
|
|
||||||
return # Don't exit the app
|
|
||||||
case "otp_generate_button":
|
|
||||||
# NEW: Push OTP workflow screen instead of legacy function
|
|
||||||
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
|
|
||||||
event.stop()
|
|
||||||
return # Don't exit the app
|
|
||||||
case "find_quiet_button":
|
|
||||||
_PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {})
|
|
||||||
case "otp_activities_button":
|
|
||||||
_PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {})
|
|
||||||
case "otp_revoke_button":
|
|
||||||
_PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {})
|
|
||||||
case "policy_prep_button":
|
|
||||||
_PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {})
|
|
||||||
case "policy_audit_update_button":
|
|
||||||
_PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {})
|
|
||||||
case _:
|
|
||||||
self.app.bell()
|
|
||||||
logger.warning("Unknown button pressed: %s", button_id)
|
|
||||||
return
|
|
||||||
|
|
||||||
logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB)
|
|
||||||
self.app.exit()
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# 2) APP
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
class Loxide(App):
|
|
||||||
CSS = """
|
|
||||||
#logo {
|
|
||||||
width: 100%;
|
|
||||||
content-align: center middle;
|
|
||||||
text-align: center;
|
|
||||||
}
|
|
||||||
"""
|
|
||||||
BINDINGS = [
|
|
||||||
("q", "quit", "Quit"),
|
|
||||||
("d", "open_dir", "Open Directory"),
|
|
||||||
]
|
|
||||||
|
|
||||||
def __init__(self, api: AirlockAPIWrapper):
|
|
||||||
self._textual_theme = load_env("TEXTUAL_THEME") or "nord"
|
|
||||||
super().__init__()
|
|
||||||
self.api = api
|
|
||||||
wd = load_env("WORKING_DIR") or os.getcwd()
|
|
||||||
if not os.path.isdir(wd):
|
|
||||||
wd = os.getcwd()
|
|
||||||
self.working_dir = wd
|
|
||||||
# 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(api))
|
|
||||||
|
|
||||||
def action_quit(self) -> None:
|
|
||||||
global _PENDING_JOB
|
|
||||||
_PENDING_JOB = None
|
|
||||||
self.exit()
|
|
||||||
|
|
||||||
def action_open_dir(self) -> None:
|
|
||||||
# Refresh data before proceeding
|
|
||||||
self.refresh_data()
|
|
||||||
screen = self.screen_stack[-1]
|
|
||||||
if isinstance(screen, MainMenuScreen):
|
|
||||||
if screen.current_tab != "dir":
|
|
||||||
screen.switch_tab("dir")
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# 3) TERMINAL + LEGACY
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
def _restore_terminal_for_legacy() -> None:
|
|
||||||
sys.stdout.write("\033[?1049l")
|
|
||||||
sys.stdout.write("\033[?25h")
|
|
||||||
sys.stdout.write("\033[0m")
|
|
||||||
sys.stdout.write("\033[?1000l\033[?1002l\033[?1003l\033[?1006l")
|
|
||||||
sys.stdout.write("\033[2J\033[H")
|
|
||||||
sys.stdout.flush()
|
|
||||||
if os.name == "nt":
|
|
||||||
try:
|
|
||||||
import ctypes
|
|
||||||
|
|
||||||
kernel32 = ctypes.windll.kernel32
|
|
||||||
handle = kernel32.GetStdHandle(-11)
|
|
||||||
mode = ctypes.c_ulong()
|
|
||||||
if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
|
|
||||||
kernel32.SetConsoleMode(handle, mode.value | 0x0004)
|
|
||||||
except Exception as exc:
|
|
||||||
logger.debug("VT enable on Windows failed: %s", exc)
|
|
||||||
|
|
||||||
|
|
||||||
def _run_legacy_job(func, args, kwargs) -> None:
|
|
||||||
logger.debug("Running legacy job: %s", getattr(func, "__name__", func))
|
|
||||||
_restore_terminal_for_legacy()
|
|
||||||
try:
|
|
||||||
func(*args, **kwargs)
|
|
||||||
finally:
|
|
||||||
try:
|
|
||||||
input("\nPress Enter to return to the UI...")
|
|
||||||
except EOFError:
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# 4) PUBLIC ENTRYPOINT
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
def run_Loxide(api: AirlockAPIWrapper) -> None:
|
|
||||||
global _PENDING_JOB
|
|
||||||
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 job loop iteration (attempt %d)", attempts)
|
|
||||||
_PENDING_JOB = None
|
|
||||||
app = Loxide(api)
|
|
||||||
|
|
||||||
try:
|
|
||||||
app.run()
|
|
||||||
except SystemExit as exc:
|
|
||||||
if exc.code != 0:
|
|
||||||
logger.debug("Caught SystemExit from Textual: %s", exc)
|
|
||||||
raise
|
|
||||||
|
|
||||||
job = _PENDING_JOB
|
|
||||||
logger.debug("After app.run(), _PENDING_JOB = %r", job)
|
|
||||||
|
|
||||||
if not job:
|
|
||||||
logger.debug("No job pending, exiting loop")
|
|
||||||
break
|
|
||||||
|
|
||||||
if job[0] == "legacy":
|
|
||||||
_, func, args, kwargs = job
|
|
||||||
_run_legacy_job(func, args, kwargs)
|
|
||||||
continue
|
|
||||||
|
|
||||||
if job[0] == "restart":
|
|
||||||
logger.debug("Restarting job loop")
|
|
||||||
continue
|
|
||||||
|
|
||||||
if job[0] == "multi_agent_action":
|
|
||||||
logger.info("Multi-agent action with selected agents: %s", job[1])
|
|
||||||
continue
|
|
||||||
|
|
||||||
if job[0] == "otp_workflow":
|
|
||||||
_, devices, requestor, reasoning, duration = job
|
|
||||||
|
|
||||||
def otp_generate_with_params():
|
|
||||||
# Your OTP logic here
|
|
||||||
pass
|
|
||||||
|
|
||||||
_run_legacy_job(otp_generate_with_params, (), {})
|
|
||||||
continue
|
|
||||||
|
|
||||||
logger.error("Unknown job type: %r", job)
|
|
||||||
break
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# 5) DEV
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
if __name__ == "__main__":
|
|
||||||
api = AirlockAPIWrapper()
|
|
||||||
run_Loxide(api)
|
|
||||||
Reference in New Issue
Block a user