513 lines
17 KiB
Python
513 lines
17 KiB
Python
import logging
|
|
import os
|
|
import sys
|
|
import dotenv
|
|
from pathlib import Path
|
|
|
|
from textual.app import App, ComposeResult
|
|
from textual.screen import Screen
|
|
from textual.widgets import (
|
|
Header,
|
|
Tabs,
|
|
Tab,
|
|
Static,
|
|
Footer,
|
|
DirectoryTree,
|
|
Button,
|
|
Tree,
|
|
)
|
|
from textual.containers import Vertical, Horizontal
|
|
from textual.reactive import reactive
|
|
|
|
from dotenv import set_key
|
|
|
|
from utils.configmanager import load_env
|
|
from utils.utils import open_directory
|
|
|
|
from utils.setup import get_base_directory, load_user_config
|
|
|
|
from flows.otp import otp_activities_by_agent, otp_generate, otp_revoke
|
|
from flows.prepPolicy import menu_policy_enforce
|
|
from flows.quietAgent import findQuietAgents
|
|
|
|
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()
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GLOBAL STASH
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_PENDING_JOB = None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
ASCII_ART = r"""
|
|
_____ .__ .__ __ ___________ .__
|
|
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
|
|
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
|
|
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
|
|
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
|
|
\/ \/ \/ \/
|
|
"""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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 = {
|
|
"find": [
|
|
("🔍 - Device Search", "find_device_button"),
|
|
("🔇 - Find Quiet Hosts", "find_quiet_button"),
|
|
],
|
|
"move": [
|
|
("✅ - Move to local approval", "move_local_button"),
|
|
("🔄 - Move to Audit/Enforcement", "move_audit_button"),
|
|
("🔀 - Move - Other", "move_other_button"),
|
|
],
|
|
"otp": [
|
|
("🔐 - Generate OTPs", "otp_generate_button"),
|
|
("📊 - OTP Activities By Agent", "otp_activities_button"),
|
|
("❌ - Revoke OTPs", "otp_revoke_button"),
|
|
],
|
|
"policy": [
|
|
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
|
|
("🔄 - Update Audit Policies", "policy_audit_update_button"),
|
|
],
|
|
}
|
|
|
|
# textual themes to expose
|
|
THEME_BUTTONS = [
|
|
("textual-dark", "textual-dark"),
|
|
("textual-light", "textual-light"),
|
|
("nord", "nord"),
|
|
("gruvbox", "gruvbox"),
|
|
("catppuccin-mocha", "catppuccin-mocha"),
|
|
("dracula", "dracula"),
|
|
("tokyo-night", "tokyo-night"),
|
|
("monokai", "monokai"),
|
|
("flexoki", "flexoki"),
|
|
("catppuccin-latte", "catppuccin-latte"),
|
|
("solarized-light", "solarized-light"),
|
|
]
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
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%" # Make button span full width of parent
|
|
buttons.append(btn)
|
|
return Vertical(*buttons)
|
|
|
|
|
|
|
|
|
|
def compose(self) -> ComposeResult:
|
|
yield Header(show_clock=True, icon="⚙")
|
|
|
|
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"),
|
|
Tab("Directory", id="dir"),
|
|
Tab("Settings", id="settings"),
|
|
]
|
|
|
|
if self.extras == "POLICYPREP":
|
|
tabs.insert(3, Tab("Policy Prep", id="policy"))
|
|
|
|
yield Tabs(*tabs, id="tabs")
|
|
yield Vertical(id="content")
|
|
yield Footer()
|
|
|
|
def on_mount(self) -> None:
|
|
self.switch_tab("find")
|
|
|
|
# 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":
|
|
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")
|
|
horizontal_container.styles.layout = "horizontal"
|
|
horizontal_container.styles.height = "auto"
|
|
content.mount(Static("Theme Options"))
|
|
content.mount(horizontal_container) # Mount the horizontal container first
|
|
|
|
# Create 3 columns
|
|
for i in range(1):
|
|
column = Vertical()
|
|
column.styles.width = "1fr"
|
|
column.styles.height = "auto"
|
|
horizontal_container.mount(column) # Mount each column
|
|
|
|
for j in range(i, len(self.THEME_BUTTONS), 1):
|
|
if j < len(self.THEME_BUTTONS):
|
|
label, btn_id = self.THEME_BUTTONS[j]
|
|
button = Button(label, id=f"set_theme_{btn_id}", compact=True)
|
|
#button.styles.width = "100%"
|
|
column.mount(button) # Mount each button
|
|
|
|
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_tree_node_selected(self, message: Tree.NodeSelected) -> None:
|
|
node = message.node
|
|
data = node.data
|
|
|
|
details_pane = self.query_one("#details-pane", Static)
|
|
|
|
if data:
|
|
details = "\n".join(f"{key}: {value}" for key, value in data.items())
|
|
else:
|
|
details = f"Selected: {node.label}"
|
|
|
|
details_pane.update(details)
|
|
|
|
|
|
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)
|
|
|
|
# theme selection → user config
|
|
if button_id.startswith("set_theme_"):
|
|
theme_name = button_id.replace("set_theme_", "")
|
|
_persist_user_theme(theme_name)
|
|
_PENDING_JOB = ("restart",)
|
|
self.app.exit()
|
|
return
|
|
|
|
match button_id:
|
|
case "find_device_button":
|
|
_PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {})
|
|
case "find_quiet_button":
|
|
_PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {})
|
|
case "move_local_button":
|
|
_PENDING_JOB = (
|
|
"legacy",
|
|
print,
|
|
("Move to local approval (placeholder)",),
|
|
{},
|
|
)
|
|
case "move_audit_button":
|
|
_PENDING_JOB = ("legacy", toggleEnforcement, (self.app.api,), {})
|
|
case "move_other_button":
|
|
_PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {})
|
|
case "otp_generate_button":
|
|
_PENDING_JOB = ("legacy", otp_generate, (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 AirlockTools(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
|
|
self.policies = api.policy_find_all()
|
|
self.devices = api.agent_find_all()
|
|
|
|
def on_mount(self) -> None:
|
|
self.theme = self._textual_theme
|
|
self.push_screen(MainMenuScreen())
|
|
|
|
def action_quit(self) -> None:
|
|
global _PENDING_JOB
|
|
_PENDING_JOB = None
|
|
self.exit()
|
|
|
|
def action_open_dir(self) -> None:
|
|
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_AirlockTools(api: AirlockAPIWrapper) -> None:
|
|
global _PENDING_JOB
|
|
|
|
while True:
|
|
base_dir = get_base_directory()
|
|
env_path = base_dir / ".env"
|
|
dotenv.load_dotenv(dotenv_path=env_path, override=True)
|
|
|
|
_PENDING_JOB = None
|
|
app = AirlockTools(api)
|
|
|
|
try:
|
|
app.run()
|
|
except SystemExit as exc:
|
|
logger.debug("Caught SystemExit from Textual: %s", exc)
|
|
|
|
job = _PENDING_JOB
|
|
logger.debug("After app.run(), _PENDING_JOB = %r", job)
|
|
|
|
if not job:
|
|
break
|
|
|
|
if job[0] == "legacy":
|
|
_, func, args, kwargs = job
|
|
_run_legacy_job(func, args, kwargs)
|
|
continue
|
|
|
|
if job[0] == "restart":
|
|
# just loop again; fresh .env was already loaded at the top
|
|
continue
|
|
|
|
break
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 5) DEV
|
|
# ---------------------------------------------------------------------------
|
|
if __name__ == "__main__":
|
|
api = AirlockAPIWrapper()
|
|
run_AirlockTools(api)
|