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..3f1a64b 100644
--- a/utils/tui.py
+++ b/utils/tui.py
@@ -1,44 +1,45 @@
-# 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
+import sys
-from blessed import Terminal
import dotenv
+from dotenv import set_key
+from textual.app import App, ComposeResult
+from textual.containers import Horizontal, 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_generate, otp_revoke
from flows.prepPolicy import menu_policy_enforce
from flows.quietAgent import findQuietAgents
-from services.agenthandler import (
- findAgents,
- moveAgents,
- toggleEnforcement,
-)
+from services.agenthandler import findAgents, moveAgents, toggleEnforcement
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 # we will actually use this now
+
+dotenv.load_dotenv()
+
+# ---------------------------------------------------------------------------
+# GLOBAL STASH
+# ---------------------------------------------------------------------------
+# ("legacy", func, args, kwargs) OR ("restart",)
+_PENDING_JOB = None
logger = logging.getLogger(__name__)
-dotenv.load_dotenv()
-term = Terminal()
-
-
-logo = r"""
+ASCII_ART = r"""
_____ .__ .__ __ ___________ .__
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
@@ -47,169 +48,421 @@ logo = r"""
\/ \/ \/ \/
"""
-
-footer_keys = ["F: Folder", "S: Settings", "B: Back", "Q: Quit"]
+# ---------------------------------------------------------------------------
+# 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"
-# Box drawing characters
-TOP_LEFT = '╔'
-TOP_RIGHT = '╗'
-BOTTOM_LEFT = '╚'
-BOTTOM_RIGHT = '╝'
-HORIZONTAL = '═'
-VERTICAL = '║'
-SHADOW = '░'
+ # 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")
-logo_lines = logo.splitlines()
-logo_width = max(len(line) for line in logo_lines)
+ # load existing user config
+ user_conf = load_user_config(config_dir)
+ user_conf["TEXTUAL_THEME"] = theme_name
-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
+ # 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)
- 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):
+ # 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 "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
+
+ 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)