mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

TUI Upgrade

Browse Source
This commit is contained in:
Zarithas
2025-11-03 10:29:13 -05:00
parent b3f48eb4df
commit 5aa0b5a88b
3 changed files with 448 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AirlockTools_Client.py
+2 -2
View File
@@ -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__":
flows/otp.py
+2 -1
View File
@@ -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"))
utils/tui.py
+444 -186
View File
@@ -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 <https://www.gnu.org/licenses/>.
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:
<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:
# 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)