From 5aa0b5a88b6c32dd36262761b7cfe884e9b7ceb9 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Mon, 3 Nov 2025 10:29:13 -0500 Subject: [PATCH] TUI Upgrade --- AirlockTools_Client.py | 4 +- flows/otp.py | 3 +- utils/tui.py | 630 +++++++++++++++++++++++++++++------------ 3 files changed, 448 insertions(+), 189 deletions(-) diff --git a/AirlockTools_Client.py b/AirlockTools_Client.py index 8f55d97..0554d43 100644 --- a/AirlockTools_Client.py +++ b/AirlockTools_Client.py @@ -30,7 +30,7 @@ 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_menu +from utils.TUI import run_AirlockTools from utils.utils import irtang urllib3.disable_warnings( @@ -70,8 +70,8 @@ def main(): base_url=str(os.getenv("URL")), api_key=api_key, ) + run_AirlockTools(api) - run_menu(api) if __name__ == "__main__": diff --git a/flows/otp.py b/flows/otp.py index 61b44fc..6c09db2 100644 --- a/flows/otp.py +++ b/flows/otp.py @@ -53,8 +53,9 @@ def otp_generate(api: AirlockAPIWrapper): 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.info(f"Generated OTP for {agent.hostname}: {otp_code}") + logger.debug(f"Generated OTP for {agent.hostname}: {otp_code}") otp_dict[agent.hostname] = otp_code print(colorText("Requested Codes:", "green")) diff --git a/utils/tui.py b/utils/tui.py index 5643b0f..906ac5a 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -1,215 +1,473 @@ -# 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 . import logging import os -import platform - -from blessed import Terminal +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, +) +from textual.containers import Vertical, Horizontal, Grid +from textual.reactive import reactive + +from dotenv import set_key + +from utils.configmanager import load_env +from utils.utils import open_directory # we will actually use this now + +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.API import AirlockAPIWrapper + +from services.agenthandler import findAgents, moveAgents, toggleEnforcement from services.policyhandler import confirmUpdateAfromE -from utils.configmanager import load_env +from services.API import AirlockAPIWrapper + +dotenv.load_dotenv() + +# --------------------------------------------------------------------------- +# GLOBAL STASH +# --------------------------------------------------------------------------- + +_PENDING_JOB = None logger = logging.getLogger(__name__) -dotenv.load_dotenv() -term = Terminal() - - -logo = r""" +ASCII_ART = r""" _____ .__ .__ __ ___________ .__ - / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______ - / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/ -/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \ + / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______ + / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/ +/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \ \____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ > - \/ \/ \/ \/ -""" - - - -footer_keys = ["F: Folder", "S: Settings", "B: Back", "Q: Quit"] - -# Box drawing characters -TOP_LEFT = '╔' -TOP_RIGHT = '╗' -BOTTOM_LEFT = '╚' -BOTTOM_RIGHT = '╝' -HORIZONTAL = '═' -VERTICAL = '║' -SHADOW = '░' - -logo_lines = logo.splitlines() -logo_width = max(len(line) for line in logo_lines) - -def draw_screen(title, items, selected_index): - print(term.clear) - center_x = (term.width - logo_width) // 2 - top = len(logo_lines) + 2 - height = len(items) + 4 - - for i, line in enumerate(logo_lines): - print(term.move_xy(center_x, i) + term.cyan(line)) - - box_color = term.bright_magenta # or term.cyan, term.green, etc. - - print(term.move_xy(center_x, top) + box_color(TOP_LEFT + HORIZONTAL * logo_width + TOP_RIGHT)) - for i in range(height): - print(term.move_xy(center_x, top + 1 + i) + box_color(VERTICAL) + ' ' * logo_width + box_color(VERTICAL)) - print(term.move_xy(center_x, top + 1 + height) + box_color(BOTTOM_LEFT + HORIZONTAL * logo_width + BOTTOM_RIGHT)) - - for i in range(1,height + 2): - print(term.move_xy(center_x + logo_width + 2, top + i) + term.darkgray(SHADOW)) - print(term.move_xy(center_x + 1, top + height + 2) + term.darkgray(SHADOW * (logo_width + 2))) - - print(term.move_xy(center_x + 4, top) + term.bold_magenta(title)) - for i, item in enumerate(items): - style = term.reverse if i == selected_index else term.bold_yellow - print(term.move_xy(center_x + 4, top + 2 + i) + style(f"{i+1}. {item}")) - - footer_text = " ".join([term.bold_cyan(k) for k in footer_keys]) - footer_x = (term.width - len(footer_text)) // 2 - print(term.move_xy(footer_x, term.height - 2) + footer_text) - print(term.move_xy(footer_x, term.height - 4) + term.bold("Use ↑/↓ or number keys. Press Enter to select."), end='', flush=True) + \/ \/ \/ \/ + """ -def run_legacy_function(func, *args, **kwargs): +# --------------------------------------------------------------------------- +# 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: + /config/user_config.json + and also mirror to /.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: - # Exit fullscreen and restore terminal state - print(term.exit_fullscreen, end='', flush=True) - print(term.normal_cursor, end='', flush=True) + 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) - # Reset terminal state on Linux - if platform.system() != 'Windows': - os.system('stty sane') + # 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) - # Clear screen - os.system('cls' if os.name == 'nt' else 'clear') - # Run the legacy function +# --------------------------------------------------------------------------- +# 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() + yield Static(ASCII_ART, id="logo") + + tabs = [ + 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 == "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_key(self, event) -> None: + key = event.key + logger.debug("KEY: %r", key) + + if self.current_tab == "dir": + return + + 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 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 + + 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 + + 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) - input("Press Enter to return to the previous menu...") - - except Exception as e: - print(f"Error during legacy function: {e}") - input("Press Enter to return to the Main Menu...") + finally: + try: + input("\nPress Enter to return to the UI...") + except EOFError: + pass -def run_menu(api: AirlockAPIWrapper, start_menu="Main Menu"): - current_menu = start_menu - extras = load_env("EXTRAS") - selected_index = 0 - history = [] - menus = { - "Main Menu": [ - "🔍 - Device Search", - "🔀 - Move Device(s)", - "🎫 - One Time Pass (OTP)", - "🔇 - Find Quiet Hosts", - ], - "🔀 - Move Device(s)": [ - "✅ - Move to local approval (Placeholder)", - "🔄 - Move to Audit/Enforcement", - "🔀 - Move - Other" - ], - "🎫 - One Time Pass (OTP)": [ - "🔐 - Generate OTPs", - "📊 - OTP Activities By Agent", - "❌ - Revoke OTPs" - ], - "Settings": [ - "(Placeholder) Change Working Directory" - ] - } - - if extras == "POLICYPREP": - menus["Main Menu"].insert(4, "🛡️ - Policy Enforcement Tools") - menus["🛡️ - Policy Enforcement Tools"] = [ - "🔒 - Prepare Policy For Enforcement", - "🔄 - Update Audit Policies" - ] - - actions = { - "🔍 - Device Search": (findAgents, [api, False]), - "🔇 - Find Quiet Hosts": (findQuietAgents, [api]), - "🔄 - Move to Audit/Enforcement": (toggleEnforcement, [api]), - "🔀 - Move - Other": (moveAgents, [api]), - "🔐 - Generate OTPs": (otp_generate, [api]), - "📊 - OTP Activities By Agent": (otp_activities_by_agent, [api]), - "❌ - Revoke OTPs": (otp_revoke, [api]), - "🔄 - Update Audit Policies": (confirmUpdateAfromE, [api]), - "🔒 - Prepare Policy For Enforcement": (menu_policy_enforce, [api]) - } +# --------------------------------------------------------------------------- +# 4) PUBLIC ENTRYPOINT +# --------------------------------------------------------------------------- +def run_AirlockTools(api: AirlockAPIWrapper) -> None: + global _PENDING_JOB while True: - func_to_run = None - args_to_run = [] + base_dir = get_base_directory() + env_path = base_dir / ".env" + dotenv.load_dotenv(dotenv_path=env_path, override=True) - with term.fullscreen(), term.cbreak(), term.hidden_cursor(): - draw_screen(current_menu, menus[current_menu], selected_index) - key = term.inkey() + _PENDING_JOB = None + app = AirlockTools(api) - items = menus[current_menu] - if key.name == "KEY_UP": - selected_index = (selected_index - 1) % len(items) - elif key.name == "KEY_DOWN": - selected_index = (selected_index + 1) % len(items) - elif key.name == "KEY_ENTER" or key == "\n": - selected_item = items[selected_index] - if selected_item in menus: - history.append(current_menu) - current_menu = selected_item - selected_index = 0 - elif selected_item in actions: - func_to_run, args_to_run = actions[selected_item] - elif key.upper() == "Q": - return - elif key.upper() == "B": - if history: - current_menu = history.pop() - selected_index = 0 - elif key.upper() == "F": - print(term.move_xy(4, term.height - 6) + term.bold_green("Folder selected")) - term.inkey(timeout=2) - - elif key.upper() == "S": - history.append(current_menu) # Add this line - current_menu = "Settings" - selected_index = 0 + try: + app.run() + except SystemExit as exc: + logger.debug("Caught SystemExit from Textual: %s", exc) - elif key.isdigit(): - num = int(key) - if 1 <= num <= len(items): - selected_index = num - 1 - selected_item = items[selected_index] - if selected_item in menus: - history.append(current_menu) - current_menu = selected_item - selected_index = 0 - elif selected_item in actions: - func_to_run, args_to_run = actions[selected_item] + job = _PENDING_JOB + logger.debug("After app.run(), _PENDING_JOB = %r", job) - # Run legacy function outside of terminal context - if func_to_run: - run_legacy_function(func_to_run, *args_to_run) + 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)