AirlockTools/Development/WIP/gu2i.py

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