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__":

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"))

utils/tui.py

@@ -1,44 +1,50 @@
-#    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
+import sys
-
-from blessed import Terminal
 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,
+from services.agenthandler import findAgents, moveAgents, toggleEnforcement
-    moveAgents,
-    toggleEnforcement,
-)
-from services.API import AirlockAPIWrapper
 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()
 
 
-
+ASCII_ART = r"""
-
-logo = r"""
    _____  .__       .__                 __     ___________           .__          
   /  _  \ |__|______|  |   ____   ____ |  | __ \__    ___/___   ____ |  |   ______ 
  /  /_\  \|  \_  __ \  |  /  _ \_/ ___\|  |/ /   |    | /  _ \ /  _ \|  |  /  ___/ 
@@ -48,168 +54,420 @@ logo = 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"
 
-footer_keys = ["F: Folder", "S: Settings", "B: Back", "Q: Quit"]
+    # 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")
 
-# Box drawing characters
+    # load existing user config
-TOP_LEFT = '╔'
+    user_conf = load_user_config(config_dir)
-TOP_RIGHT = '╗'
+    user_conf["TEXTUAL_THEME"] = theme_name
-BOTTOM_LEFT = '╚'
-BOTTOM_RIGHT = '╝'
-HORIZONTAL = '═'
-VERTICAL = '║'
-SHADOW = '░'
 
-logo_lines = logo.splitlines()
+    # write it back
-logo_width = max(len(line) for line in logo_lines)
+    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)
 
-def draw_screen(title, items, selected_index):
+    # mirror to .env (like setup.write_config_to_env does)
-    print(term.clear)
+    env_path.parent.mkdir(parents=True, exist_ok=True)
-    center_x = (term.width - logo_width) // 2
+    if not env_path.exists():
-    top = len(logo_lines) + 2
+        env_path.touch()
-    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):
     try:
-        # Exit fullscreen and restore terminal state
+        set_key(str(env_path), "TEXTUAL_THEME", theme_name)
-        print(term.exit_fullscreen, end='', flush=True)
+    except Exception as exc:  # keep going even if .env write fails
-        print(term.normal_cursor, end='', flush=True)
+        logger.warning("Failed to mirror TEXTUAL_THEME to .env: %s", exc)
 
-        # Reset terminal state on Linux
+    # reload so load_env(...) sees the new value right now
-        if platform.system() != 'Windows':
+    dotenv.load_dotenv(dotenv_path=env_path, override=True)
-            os.system('stty sane')
+    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...")
+    finally:
-
+        try:
-    except Exception as e:
+            input("\nPress Enter to return to the UI...")
-        print(f"Error during legacy function: {e}")
+        except EOFError:
-        input("Press Enter to return to the Main Menu...")
+            pass
 
 
-def run_menu(api: AirlockAPIWrapper, start_menu="Main Menu"):
+# ---------------------------------------------------------------------------
-    current_menu = start_menu
+# 4) PUBLIC ENTRYPOINT
-    extras = load_env("EXTRAS")
+# ---------------------------------------------------------------------------
-    selected_index = 0
+def run_AirlockTools(api: AirlockAPIWrapper) -> None:
-    history = []
+    global _PENDING_JOB
-    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])
-    }
 
     while True:
-        func_to_run = None
+        base_dir = get_base_directory()
-        args_to_run = []
+        env_path = base_dir / ".env"
+        dotenv.load_dotenv(dotenv_path=env_path, override=True)
 
-        with term.fullscreen(), term.cbreak(), term.hidden_cursor():
+        _PENDING_JOB = None
-            draw_screen(current_menu, menus[current_menu], selected_index)
+        app = AirlockTools(api)
-            key = term.inkey()
 
-            items = menus[current_menu]
+        try:
-            if key.name == "KEY_UP":
+            app.run()
-                selected_index = (selected_index - 1) % len(items)
+        except SystemExit as exc:
-            elif key.name == "KEY_DOWN":
+            logger.debug("Caught SystemExit from Textual: %s", exc)
-                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":
+        job = _PENDING_JOB
-                history.append(current_menu)  # Add this line
+        logger.debug("After app.run(), _PENDING_JOB = %r", job)
-                current_menu = "Settings"
-                selected_index = 0
 
-            elif key.isdigit():
+        if not job:
-                num = int(key)
+            break
-                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]
 
-        # Run legacy function outside of terminal context
+        if job[0] == "legacy":
-        if func_to_run:
+            _, func, args, kwargs = job
-            run_legacy_function(func_to_run, *args_to_run)
+            _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)