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 Horizontal import logging import dotenv import os from Development_Stubs.WIP.localApproval import moveToLocalApproval from flows.prepPolicy import ( buildPathsandPublishers, buildPreflights, selectAllowlists, selectPolicies, sortHashes, ) from flows.quietAgent import findQuietAgents from services.agenthandler import findAgents from services.API import AirlockAPIWrapper import services.policyhandler as policyh from utils.utils import ( areYouSure, colorText, displayIntro, load_env, open_directory, printEnforceChecklist, ) from screens.agent_results import AgentResultsScreen logger = logging.getLogger(__name__) dotenv.load_dotenv() ASCII_ART = displayIntro() class BaseScreen(Screen): def compose(self) -> ComposeResult: yield Header() yield Static("This is a dummy screen.", id="content") yield Footer() class LandingScreen(BaseScreen): def compose(self) -> ComposeResult: yield Header() yield Static(ASCII_ART, id="ascii-art", markup=False) yield Horizontal( Button("Get Started", id="get-started-button"), Button("Settings", id="settings-button"), ) yield Footer() def on_button_pressed(self, event: Button.Pressed) -> None: if event.button.id == "get-started-button": # Go back to the previous screen (MainMenuScreen) self.app.pop_screen() class LocalApprovalScreen(BaseScreen): def compose(self) -> ComposeResult: yield Header() yield Footer() class OTPScreen(BaseScreen): def compose(self) -> ComposeResult: yield Header() yield Static("OTP: Enter your one-time password here.", id="content") yield Footer() class SearchScreen(BaseScreen): def compose(self) -> ComposeResult: yield Header() yield Static("Search: Find items in the system.", id="content") yield Horizontal( Button("Search for Agents", id="searchAgent-button"), ) yield Footer() def on_button_pressed(self, event: Button.Pressed) -> None: if event.button.id == "searchAgent-button": self.app.push_screen(AgentResultsScreen(self.app.api, self.app.working_dir)) class QuietHostsScreen(BaseScreen): def compose(self) -> ComposeResult: yield Header() yield Static("Quiet Hosts: Manage hosts that are not responding.", id="content") yield Footer() class PolicyPrepScreen(BaseScreen): def compose(self) -> ComposeResult: yield Header() yield Static("Policy Prep: Prepare policies for deployment.", id="content") yield Footer() class UpdatePoliciesScreen(BaseScreen): def compose(self) -> ComposeResult: yield Header() yield Static("Update Policies: Update existing policies.", id="content") yield Footer() class SettingsScreen(BaseScreen): def compose(self) -> ComposeResult: yield Header() yield Static("Settings: Configure application settings.", id="content") yield Footer() class DirectoryTreeScreen(Screen): def compose(self) -> ComposeResult: yield Header() yield DirectoryTree("./") yield Footer() class MainMenuScreen(Screen): def compose(self) -> ComposeResult: yield Header() yield Tabs( Tab("Home", id="home_screen"), Tab("Local Approval", id="local_approval"), Tab("OTP", id="otp"), Tab("Search", id="search"), Tab("Quiet Hosts", id="quiet"), Tab("Policy Prep", id="policy"), Tab("Update Policies", id="update"), id="tabs" ) yield Static("Select a tab to begin.", id="content") yield Footer() def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None: match event.tab.id: case "home_screen": self.app.push_screen(LandingScreen()) case "local_approval": self.app.push_screen(LocalApprovalScreen()) case "otp": self.app.push_screen(OTPScreen()) case "search": self.app.push_screen(SearchScreen()) case "quiet": self.app.push_screen(QuietHostsScreen()) case "policy": self.app.push_screen(PolicyPrepScreen()) case "update": self.app.push_screen(UpdatePoliciesScreen()) case "settings": self.app.push_screen(SettingsScreen()) class AirlockTools(App): BINDINGS = [ ("q", "quit", "Quit"), ("d", "open_dir", "Open Directory"), ("b", "back", "Go Back"), ] def __init__(self, api: AirlockAPIWrapper, working_dir: str): super().__init__() self.api = api self.working_dir = working_dir def on_mount(self) -> None: self.push_screen(MainMenuScreen()) def action_quit(self) -> None: self.exit() def action_open_dir(self) -> None: self.push_screen(DirectoryTreeScreen()) def action_back(self) -> None: if len(self.screen_stack) > 2: self.pop_screen() else: self.bell() if __name__ == "__main__": api = AirlockAPIWrapper(base_url=str(os.getenv("URL")),api_key = "b82ab7e2b12430fff0cbf0e798f79e16c1aea2c37ffee5dc884de7ba11d3a3a4") working_dir = "./" # Or a dynamic path AirlockTools(api, working_dir).run()