From 70477f3885bafa3c76f975a4c4cb537660d05301 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Mon, 3 Nov 2025 15:59:08 -0500 Subject: [PATCH] Policy Tree Tab added --- services/API.py | 16 +++++++++ utils/tui.py | 86 +++++++++++++++++++++++++++++++++++-------------- 2 files changed, 78 insertions(+), 24 deletions(-) diff --git a/services/API.py b/services/API.py index c618a79..e5a3c76 100644 --- a/services/API.py +++ b/services/API.py @@ -236,6 +236,22 @@ class AirlockAPIWrapper: """Set audit mode for a policy group. 1=Audit, 0=Enforcement""" payload = {"groupid": groupid, "auditmode": auditmode} return self._post("/v1/group/settings/auditmode", payload) + + def policy_set_script_custom(self, + groupid: str, + script_custom: int, + scripts_audit: List[str], + scripts_disabled: List[str], + scripts_respect: List[str], + ) -> dict: + """Set audit mode for a policy group. 1=Audit, 0=Enforcement""" + payload = {"groupid": groupid, + "script_custom": script_custom, + "scripts_audit": scripts_audit, + "scripts_disabled": scripts_disabled, + "scripts_respect": scripts_respect + } + return self._post("/v1/group/settings/script_custom", payload) # Execution History def history_logging(self, type: List[str], checkpoint: str, policy: List[str]) -> str: diff --git a/utils/tui.py b/utils/tui.py index 823c0c8..9dfa6e3 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -14,6 +14,7 @@ from textual.widgets import ( Footer, DirectoryTree, Button, + Tree, ) from textual.containers import Vertical, Horizontal from textual.reactive import reactive @@ -33,6 +34,12 @@ from services.agenthandler import findAgents, moveAgents, toggleEnforcement from services.policyhandler import confirmUpdateAfromE from services.API import AirlockAPIWrapper + +from rich.text import Text +from rich.style import Style +from rich.color import Color +from textual.widgets import Static + dotenv.load_dotenv() # --------------------------------------------------------------------------- @@ -100,6 +107,8 @@ def _persist_user_theme(theme_name: str) -> None: logger.debug("Reloaded .env from %s", env_path) + + # --------------------------------------------------------------------------- # 1) SCREEN # --------------------------------------------------------------------------- @@ -169,6 +178,7 @@ class MainMenuScreen(Screen): yield Static(ASCII_ART, id="logo") tabs = [ + Tab("Policy Tree", id="p_tree"), Tab("Device Search", id="find"), Tab("Move Agent", id="move"), Tab("OTP", id="otp"), @@ -233,6 +243,46 @@ class MainMenuScreen(Screen): 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": + layout = Horizontal() + content.mount(layout) + + # Left: Policy Tree + policy_tree = Tree("Policies", id="policy_tree") + policy_tree.styles.width = "2fr" + layout.mount(policy_tree) + + # Right: Details pane + details_pane = Static("Select a policy or device to view details", id="details-pane") + details_pane.styles.width = "3fr" + layout.mount(details_pane) + + # Build the tree + node_map = {} + + # Top-level policies + for _, policy in self.app.policies.iterrows(): + if policy["parent"] == "global-policy-settings": + node = policy_tree.root.add(label=policy["name"], data=policy.to_dict()) + node_map[policy["groupid"]] = node + + # Child policies + for _, policy in self.app.policies.iterrows(): + parent_id = policy["parent"] + if parent_id in node_map: + parent_node = node_map[parent_id] + node = parent_node.add(label=policy["name"], data=policy.to_dict()) + node_map[policy["groupid"]] = node + + # Devices under policies + for _, device in self.app.devices.iterrows(): + group_id = device["groupid"] + if group_id in node_map: + parent_node = node_map[group_id] + label = device["hostname"] # Keep tree clean + parent_node.add(label=label, data=device.to_dict()) + + elif tab_id == "settings": # Create and mount the horizontal container horizontal_container = Horizontal(id="settings_grid") @@ -261,33 +311,19 @@ class MainMenuScreen(Screen): def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None: self.switch_tab(event.tab.id) - def on_key(self, event) -> None: - key = event.key - logger.debug("KEY: %r", key) + def on_tree_node_selected(self, message: Tree.NodeSelected) -> None: + node = message.node + data = node.data - if self.current_tab == "dir": - return + details_pane = self.query_one("#details-pane", Static) - if key in ("down", "j"): - self._focus_nearby_button(+1) - event.stop() - return - if key in ("up", "k"): - self._focus_nearby_button(-1) - event.stop() - return + if data: + details = "\n".join(f"{key}: {value}" for key, value in data.items()) + else: + details = f"Selected: {node.label}" - if key in ("left", "right"): - tabs = self.query_one("#tabs", Tabs) - if not tabs.has_focus: - tabs.focus() - if key == "left": - tabs.action_previous_tab() - else: - tabs.action_next_tab() - event.stop() - return - return + details_pane.update(details) + def on_directory_tree_file_selected(self, event: DirectoryTree.FileSelected) -> None: path = event.path @@ -375,6 +411,8 @@ class AirlockTools(App): if not os.path.isdir(wd): wd = os.getcwd() self.working_dir = wd + self.policies = api.policy_find_all() + self.devices = api.agent_find_all() def on_mount(self) -> None: self.theme = self._textual_theme