diff --git a/AirlockTools_Client.py b/Loxide.py similarity index 100% rename from AirlockTools_Client.py rename to Loxide.py diff --git a/default_system_config.json b/default_system_config.json index b09f901..decfd79 100644 --- a/default_system_config.json +++ b/default_system_config.json @@ -1,5 +1,5 @@ { - "APPNAME": "AirlockTools", + "APPNAME": "Loxide", "URL": "https://server:3129", "LOG_LEVEL": "INFO", "BAD_PATH_PARTS": ["users","wwwroot","windows\\temp","windows\\task","windows\\system32","startup", "windows\\fonts","Recycle.Bin","AppData","programdata", "Solarwinds","kaseya"], @@ -8,6 +8,8 @@ "PATH_EXCLUSION_CONST": 4, "MIN_FILES_FOR_PATH": 4, "VT_THREAT_TOLERANCE": 4, + "TELEMETRY": "FALSE", + "TELEM_URL": "", "POLICY_MAP_ENF_AUD": { } diff --git a/flows/otp.py b/flows/otp.py index 8b15fab..720aab3 100644 --- a/flows/otp.py +++ b/flows/otp.py @@ -29,44 +29,6 @@ from utils.utils import colorText, get_sanitized_input logger = logging.getLogger(__name__) -def otp_generate(api: AirlockAPIWrapper): - otp_dict = {} - agents = selectAgents(api) - print(colorText("Would you like to continue with these devices?", "white")) - for agent in agents: - print(agent.hostname) - confirm = Selector.confirm() - if agents and confirm: - requester = get_sanitized_input("Who is requesting the OTP: ") - because = get_sanitized_input("Why/What work are they doing?: ") - - purpose = f"Requester: {requester} - for : {because}" - possible_durations = [15, 60, 360, 1440, 10080] - - print(colorText("Please select a duration in minutes: ", "white")) - print( - colorText( - "15 mins, 60 mins, 360 mins(6 Hours), 1440 mins (24 Hours), 10080 mins (7 Days):", - "white", - ) - ) - duration_selected = Selector.select_int(possible_durations) - - if isinstance(duration_selected, list): - duration_selected = duration_selected[0] if duration_selected else None - - 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.debug(f"Generated OTP for {agent.hostname}: {otp_code}") - otp_dict[agent.hostname] = otp_code - - print(colorText("Requested Codes:", "green")) - for key, value in otp_dict.items(): - print(colorText(f"{key} | {value}", "green")) - - def otp_activities_by_agent(api: AirlockAPIWrapper): activeagents = api.otp_find_active() awaitingagents = api.otp_find_awaiting() diff --git a/screens/moveagentworkflowscreen.py b/screens/moveagentworkflowscreen.py index 7199f55..da96331 100644 --- a/screens/moveagentworkflowscreen.py +++ b/screens/moveagentworkflowscreen.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Optional from textual.app import ComposeResult from textual.screen import Screen @@ -12,7 +12,7 @@ from widgets.resultsdisplay import ResultsDisplay class MoveAgentWorkflowScreen(Screen): """Screen that handles the agent movement workflow.""" - def __init__(self, all_agents: List[Agent]): + def __init__(self, all_agents: Optional[List[Agent]]): super().__init__() self.all_agents = all_agents self.selected_agents = None diff --git a/screens/otpactivityscreen.py b/screens/otpactivityscreen.py new file mode 100644 index 0000000..cc98655 --- /dev/null +++ b/screens/otpactivityscreen.py @@ -0,0 +1,678 @@ +from __future__ import annotations + +from datetime import datetime +import logging +import os + +import pandas as pd +from textual.app import ComposeResult +from textual.binding import Binding +from textual.containers import Horizontal, Vertical +from textual.screen import Screen +from textual.widgets import Button, DataTable, Footer, Header, Static + +from utils.configmanager import load_env + +logger = logging.getLogger(__name__) + + +def _load_working_dir() -> str: + """ + Load the working directory from environment variables or use the current working directory. + """ + wd = os.environ.get("WORKING_DIR") + if wd: + return wd + return os.getcwd() + + +class OTPActivitiesWidget(Static): + """ + Reusable widget that contains the sessions table (left) and an Activity Preview (right). + The right side shows an Activity Preview that takes ~75% vertical space, and a lower area + with Back and Continue buttons. The Continue button pushes ActivityDetailScreen with the + currently-loaded activities. + """ + + DEFAULT_CSS = """ + OTPActivitiesWidget { + height: 1fr; + } + #main_row { + width: 100%; + height: 100%; + layout: horizontal; + } + #left_panel { + width: 60%; + min-width: 60; + border: none; + } + #right_panel { + width: 40%; + min-width: 40; + border: none; + layout: vertical; + } + #activity_preview_container { + height: 75%; + border: none; + padding: 1 1; + } + #activity_buttons { + height: 25%; + padding: 1 1; + content-align: center middle; + } + """ + + def compose(self) -> ComposeResult: + # Layout: horizontal main row with left & right panels + with Horizontal(id="main_row"): + # Left: sessions area + with Vertical(id="left_panel"): + yield Static("OTP Sessions", classes="panel-title") + with Vertical(id="sessions_table_container"): + self.sessions_table = DataTable(id="sessions_table") + self.sessions_table.styles.width = "100%" + yield self.sessions_table + # Right: Activity Preview (top 3/4) + buttons (bottom 1/4) + with Vertical(id="right_panel"): + # Activity preview area (takes ~75% of right panel) + yield Static("Activity Preview", classes="panel-title") + with Vertical(id="activity_preview_container"): + self.activities_table = DataTable(id="activity_preview_table") + yield self.activities_table + # Buttons area at the bottom (Back, Continue) + with Horizontal(id="activity_buttons"): + # Back takes left side, Continue right side + self.back_btn = Button("Back", id="activity_back_btn") + self.continue_btn = Button("Continue", id="activity_continue_btn") + # Stretch buttons nicely + self.back_btn.styles.width = "50%" + self.continue_btn.styles.width = "50%" + yield self.back_btn + yield self.continue_btn + + async def on_mount(self) -> None: + # Configure sessions table and activities preview + self.sessions_table.clear() + self.sessions_table.add_columns( + "otpid", "hostname", "status", "purpose", "granted" + ) + self.activities_table.clear() + # activities_table columns are dynamically added when activities are loaded. + # Selection behavior + self.sessions_table.cursor_type = "row" + try: + self.sessions_table.zebra_stripes = True + except Exception: + pass + self.activities_table.cursor_type = "row" + try: + self.activities_table.zebra_stripes = True + except Exception: + pass + # Store state + self._sessions_df: pd.DataFrame | None = None + self._activities_df: pd.DataFrame | None = None + self._selected_session_otpid: str | int | None = None + + async def on_button_pressed(self, event) -> None: # type: ignore[override] + """ + Handle Back / Continue buttons for the Activity Preview area. + """ + # Try to resolve the button object from the event + btn = ( + getattr(event, "button", None) + or getattr(event, "sender", None) + or getattr(event, "control", None) + or getattr(event, "widget", None) + ) + btn_id = ( + getattr(btn, "id", None) + or getattr(event, "button_id", None) + or getattr(event, "id", None) + ) + + # ---- Back ---- + if btn is self.back_btn or btn_id == getattr(self.back_btn, "id", None): + while len(self.app.screen_stack) > 2: + self.app.pop_screen() + event.stop() + return + + # ---- Continue ---- + if btn is self.continue_btn or btn_id == getattr(self.continue_btn, "id", None): + if self._activities_df is None or self._activities_df.empty: + logger.info("Continue pressed but no activities loaded.") + await self.post_message( + Static("No activities loaded to continue with.") + ) + return + # Copy activities DataFrame to pass to new screen + activities_copy = self._activities_df.copy() + otpid = self._selected_session_otpid + # Optionally include hostname if available + hostname = None + try: + if self._sessions_df is not None: + df = self._sessions_df.reset_index(drop=True) + match = df[df["otpid"] == otpid] + if not match.empty: + hostname = match.iloc[0].get("hostname") + except Exception: + hostname = None + # Create and push ActivityDetailScreen, handing the data + try: + detail_screen = ActivityDetailScreen( + activities_copy, otpid=otpid, hostname=hostname + ) + await self.app.push_screen(detail_screen) + except Exception as exc: + logger.exception("Failed to push ActivityDetailScreen: %s", exc) + return + + # Unknown button on widget + logger.debug( + "Unhandled OTPActivitiesWidget button pressed (resolved btn=%r, id=%r)", + btn, + btn_id, + ) + + async def on_data_table_row_selected(self, event) -> None: # type: ignore[override] + """ + Robust handler for DataTable row-selection across Textual micro-versions. + Tries many attribute names and shapes: + - numeric index (row_key, row_index, index) + - coordinate object or tuple (coordinate.row or (row, col)) + - direct row values (row, values, cells) -> we try to map those back to the sessions DF + - table.cursor_row fallback + """ + # 1) Determine the sending table (best-effort) + sender = None + for attr in ("sender", "table", "data_table", "control"): + sender = getattr(event, attr, None) + if sender is not None: + break + if sender is None: + sender = self.sessions_table # Assume sessions_table if unknown + # Only respond to selections in the sessions table + if sender is not self.sessions_table: + return + + # Helper to log and return + def _bad(msg: str, *args): + logger.warning(msg, *args) + return None + + # 2) Try to extract a numeric index + row_key = None + for attr in ("row_key", "row", "row_index", "index"): + row_key = getattr(event, attr, None) + if row_key is not None: + break + + # If coordinate: try to extract .row or tuple[0] + if row_key is None: + coord = getattr(event, "coordinate", None) or getattr( + event, "cursor_coordinate", None + ) + if coord is not None: + if hasattr(coord, "row"): + row_key = coord.row + elif isinstance(coord, (tuple, list)) and len(coord) >= 1: + row_key = coord[0] + + # If still nothing, maybe the event provides the row's cell values directly + row_values = None + for attr in ("values", "cells", "row", "row_values", "selected_row_values"): + val = getattr(event, attr, None) + if val: + # Prefer actual sequence of cell values + row_values = val + break + + # If we have row_values, try to map them back to the sessions DataFrame + if row_values is not None: + # Normalize into list of strings for comparison + try: + vals = [ + "" if pd.isna(v) else str(v) + for v in ( + list(row_values) + if not isinstance(row_values, str) + else [row_values] + ) + ] + except Exception: + vals = [str(row_values)] + # Try to match against the expected columns order we render + if self._sessions_df is None or self._sessions_df.empty: + logger.warning( + "Sessions DataFrame is empty; cannot map selected row values." + ) + return + df_ordered = self._sessions_df.reset_index(drop=True) + expected_cols = ["otpid", "hostname", "status", "purpose", "granted"] + + # Build stringified candidates for each row in df using the same columns we show + def _row_to_vals(sr): + out = [] + for c in expected_cols: + if c in sr: + v = sr[c] + out.append("" if pd.isna(v) else str(v)) + else: + out.append("") + return out + + match_idx = None + for i, sr in df_ordered.iterrows(): + cand = _row_to_vals(sr) + # Compare prefix: row values might be a subset (e.g. only first 3 cols), so compare prefix only + if len(vals) <= len(cand) and all( + vals[j] == cand[j] for j in range(len(vals)) + ): + match_idx = i + break + if match_idx is None: + # Try looser match: compare first cell only (otpid) + first = vals[0] if vals else None + if first is not None: + for i, sr in df_ordered.iterrows(): + cand0 = "" if pd.isna(sr.get("otpid")) else str(sr.get("otpid")) + if cand0 == first: + match_idx = i + break + if match_idx is None: + logger.warning( + "Unable to locate DataFrame row matching selected row values: %r", + vals, + ) + return + idx = int(match_idx) + else: + # 3) If we have a row_key, try to normalize to an int index + if row_key is not None: + try: + idx = int(row_key) + except Exception: + # Try converting via string + try: + idx = int(str(row_key)) + except Exception: + idx = None + if idx is None: + # Final numeric fallback: use sessions_table.cursor_row if present + try: + idx = getattr(self.sessions_table, "cursor_row") + except Exception: + idx = None + if idx is None: + _bad("Failed to normalize row/key from event: %r", row_key) + return + else: + # 4) Try table cursor_row as last resort + try: + idx = getattr(self.sessions_table, "cursor_row") + except Exception: + logger.warning( + "Could not determine selected row from event: %r", event + ) + # Helpful debug hint for you to paste back if still failing: + logger.debug("Event repr for debugging: %r", event) + return + + # At this point we should have an integer idx + try: + idx = int(idx) + except Exception: + logger.exception( + "Final normalization of selected row index failed: %r", idx + ) + return + + # Validate sessions df + if self._sessions_df is None or self._sessions_df.empty: + logger.warning("Sessions DataFrame empty; nothing to select.") + return + df_ordered = self._sessions_df.reset_index(drop=True) + if idx < 0 or idx >= len(df_ordered): + logger.warning( + "Selected row index %s out of range (0..%d)", idx, len(df_ordered) - 1 + ) + return + row_series = df_ordered.iloc[idx] + otpid = row_series.get("otpid") + hostname = row_series.get("hostname") + # Store selected session and fetch activities + self._selected_session_otpid = otpid + # Obtain api from app (try multiple places) + api = ( + getattr(self.app, "api", None) + or getattr(self, "api", None) + or getattr(self.app, "airlock_api", None) + ) + if api is None: + logger.error("No API available on self.app.api - cannot fetch activities") + return + logger.info( + "Fetching activities for otpid=%s host=%s (selected row=%s)", + otpid, + hostname, + idx, + ) + await self._fetch_activities_for_otpid(api, otpid, hostname=hostname) + + async def load_sessions_from_api(self, api) -> None: + """ + Pulls OTP session lists, adds status column, concatenates and populates the sessions table. + """ + try: + active = api.otp_find_active() + awaiting = api.otp_find_awaiting() + enforced = api.otp_find_enforced() + revoked = api.otp_find_revoked() + except Exception as exc: + logger.exception("Failed to fetch OTP session lists: %s", exc) + # Present empty + active = awaiting = enforced = revoked = pd.DataFrame() + + # Ensure DataFrame objects + def _ensure_df(df): + return df if isinstance(df, pd.DataFrame) else pd.DataFrame(df) + + active = _ensure_df(active) + awaiting = _ensure_df(awaiting) + enforced = _ensure_df(enforced) + revoked = _ensure_df(revoked) + for df, status in [ + (active, "active"), + (awaiting, "awaiting"), + (enforced, "enforced"), + (revoked, "revoked"), + ]: + if "status" not in df.columns: + df["status"] = status + combined = pd.concat([active, awaiting, enforced, revoked], ignore_index=True) + if "otpid" in combined.columns: + combined = combined.sort_values(by="otpid", ascending=False) + self._sessions_df = combined + # Populate DataTable + self.sessions_table.clear() + # Ensure columns exist in DF and when missing add empty column + expected_cols = ["otpid", "hostname", "status", "purpose", "granted"] + for col in expected_cols: + if col not in combined.columns: + combined[col] = "" + self.sessions_table.add_columns(*expected_cols) + # Add rows + for _, row in combined[expected_cols].iterrows(): + # Convert values to str for safe insertion + vals = ["" if pd.isna(v) else v for v in row.to_list()] + self.sessions_table.add_row(*[str(v) for v in vals]) + logger.info("Loaded %d OTP sessions.", len(combined)) + + async def _fetch_activities_for_otpid(self, api, otpid, hostname=None) -> None: + """ + Fetch activities DataFrame for a given otpid and populate activities_table. + """ + try: + result = api.otp_get_activities(otpid) + result_df = ( + result if isinstance(result, pd.DataFrame) else pd.DataFrame(result) + ) + except Exception as exc: + logger.exception("Failed to fetch activities for otpid %s: %s", otpid, exc) + result_df = pd.DataFrame() + # Attach hostname if provided + if hostname is not None: + result_df["hostname"] = hostname + if result_df.empty: + logger.info("No activities found for otpid %s (host: %s)", otpid, hostname) + self._activities_df = pd.DataFrame() + self.activities_table.clear() + return + # Store and render + self._activities_df = result_df.copy() + # Rebuild activities_table columns from result_df + self.activities_table.clear() + # Ensure stable column order + for col in result_df.columns: + self.activities_table.add_column(col) + # Add rows + for _, arow in result_df.iterrows(): + values = ["" if pd.isna(v) else v for v in arow.to_list()] + self.activities_table.add_row(*[str(v) for v in values]) + logger.info( + "Loaded %d activity rows for otpid %s (host: %s)", + len(result_df), + otpid, + hostname, + ) + + async def export_activities(self) -> None: + """ + Export currently-loaded activities DataFrame to CSV. + Can be called directly (programmatically) or from the button handler. + """ + if self._activities_df is None or self._activities_df.empty: + logger.info("No activities loaded to export.") + # On-screen short message + await self.post_message(Static("No activities to export.")) + return + working_dir = _load_working_dir() + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + filename = f"otp_activities_{self._selected_session_otpid}_{timestamp}.csv" + file_path = os.path.join(working_dir, filename) + try: + self._activities_df.to_csv(file_path, index=False) + logger.info("Exported activities to %s", file_path) + await self.post_message(Static(f"✅ Exported activities to: {file_path}")) + except Exception as exc: + logger.exception("Failed to export activities to %s: %s", file_path, exc) + await self.post_message(Static("Failed to export activities; check logs.")) + + +class ActivityDetailWidget(Static): + """ + Interactive widget for Activity Detail screen. + Shows the provided DataFrame in a DataTable and offers Export + Back buttons. + """ + + DEFAULT_CSS = """ + ActivityDetailWidget { + height: 1fr; + layout: vertical; + } + #detail_table_container { + height: 85%; + padding: 1 1; + } + #detail_buttons { + height: 15%; + padding: 1 1; + content-align: center middle; + } + """ + + def __init__(self, activities_df: pd.DataFrame, otpid=None, hostname=None) -> None: + super().__init__() + self.activities_df = ( + activities_df.copy() + if isinstance(activities_df, pd.DataFrame) + else pd.DataFrame(activities_df) + ) + self.otpid = otpid + self.hostname = hostname + + def compose(self) -> ComposeResult: + yield Static( + f"Activity Detail (otpid={self.otpid} host={self.hostname})", + classes="panel-title", + ) + # Table container + with Vertical(id="detail_table_container"): + self.detail_table = DataTable(id="detail_table") + yield self.detail_table + # Buttons at bottom + with Horizontal(id="detail_buttons"): + self.detail_back_btn = Button("Back", id="detail_back_btn") + self.detail_export_btn = Button("Export (CSV)", id="detail_export_btn") + # Make them stretch equally + self.detail_back_btn.styles.width = "50%" + self.detail_export_btn.styles.width = "50%" + yield self.detail_back_btn + yield self.detail_export_btn + + async def on_mount(self) -> None: + # Populate table from activities_df + self.detail_table.clear() + if self.activities_df is None or self.activities_df.empty: + logger.info("ActivityDetailWidget mounted with empty dataframe.") + return + # Add columns + for col in self.activities_df.columns: + self.detail_table.add_column(col) + # Add rows + for _, row in self.activities_df.iterrows(): + vals = ["" if pd.isna(v) else v for v in row.to_list()] + self.detail_table.add_row(*[str(v) for v in vals]) + # Allow sorting / cursor + self.detail_table.cursor_type = "row" + + async def on_button_pressed(self, event) -> None: + btn = getattr(event, "button", None) or getattr(event, "sender", None) + btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None) + # Back button in ActivityDetailWidget + + if btn is self.detail_back_btn or btn_id == "detail_back_btn": + # Pop screens until only the main menu remains + while len(self.app.screen_stack) > 2: + self.app.pop_screen() + event.stop() + return + # Export button + if btn is self.detail_export_btn or btn_id == "detail_export_btn": + await self._export_detail_activities() + return + + async def _export_detail_activities(self) -> None: + + if self.activities_df is None or self.activities_df.empty: + logger.info("No activities to export.") + notification = Static("❌ No activities to export.", classes="notification") + self.mount(notification) + return + try: + working_dir = load_env("WORKING_DIR") or os.getcwd() + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + filename = f"otp_activities_detail_{timestamp}.csv" + file_path = os.path.join(working_dir, filename) + self.activities_df.to_csv(file_path, index=False) + logger.info("Exported detail activities to %s", file_path) + # Show success notification + notification = Static( + f"✅ Exported activities to: {filename}", classes="notification" + ) + self.mount(notification) + except Exception as exc: + logger.exception("Failed to export detail activities: %s", exc) + notification = Static( + "❌ Failed to export activities; check logs.", classes="notification" + ) + self.mount(notification) + + +class ActivityDetailScreen(Screen): + """ + Screen that wraps ActivityDetailWidget. Expects a DataFrame passed on init. + """ + + BINDINGS = [Binding("b", "back", "Back"), Binding("e", "export", "Export")] + + def __init__(self, activities_df: pd.DataFrame, otpid=None, hostname=None) -> None: + super().__init__() + self._activities_df = ( + activities_df.copy() + if isinstance(activities_df, pd.DataFrame) + else pd.DataFrame(activities_df) + ) + self._otpid = otpid + self._hostname = hostname + + def compose(self) -> ComposeResult: + self.widget = ActivityDetailWidget( + self._activities_df, otpid=self._otpid, hostname=self._hostname + ) + yield Header(show_clock=True) + yield self.widget + yield Footer() + + async def action_back(self) -> None: + try: + await self.app.pop_screen() + except Exception: + logger.debug("ActivityDetailScreen.action_back pop_screen failed.") + + async def action_export(self) -> None: + # Delegate to widget export helper + if hasattr(self, "widget") and self.widget is not None: + await self.widget._export_detail_activities() + + +class OTPActivitiesScreen(Screen): + """ + A Screen intended to be pushed into an existing Textual App. + Usage: + app.push_screen(OTPActivitiesScreen()) + or create this screen and call `await screen.load()` inside your app lifecycle. + The screen expects `self.app.api` to exist and be an AirlockAPIWrapper instance. + """ + + BINDINGS = [ + Binding("r", "refresh_sessions", "Refresh Sessions"), + Binding("e", "export_activities", "Export activities"), + Binding("q", "quit", "Quit"), + ] + + def compose(self) -> ComposeResult: + yield Header() + self.widget = OTPActivitiesWidget() + yield self.widget + yield Footer() + + async def on_show(self) -> None: + """Restore focus to the left sessions table when the screen becomes visible.""" + if hasattr(self, "widget") and hasattr(self.widget, "sessions_table"): + self.widget.sessions_table.focus() + + async def on_mount(self) -> None: + # Try to load sessions immediately + api = getattr(self.app, "api", None) + if api is None: + logger.warning("OTPActivitiesScreen mounted but no self.app.api found.") + else: + await self.widget.load_sessions_from_api(api) + + # Simple actions bound to keys + async def action_refresh_sessions(self) -> None: + api = getattr(self.app, "api", None) + if api is None: + logger.error("No API on app; cannot refresh sessions.") + return + logger.info("Refreshing OTP sessions via API.") + await self.widget.load_sessions_from_api(api) + + async def action_quit(self) -> None: + # Pop the screen or exit app + await self.app.pop_screen() + + # If you want an explicit method to fetch activities for a particular otpid from outside: + async def fetch_activities_for_otpid(self, otpid, hostname=None) -> None: + api = getattr(self.app, "api", None) + if api is None: + logger.error("No API on app; cannot fetch activities.") + return + await self.widget._fetch_activities_for_otpid(api, otpid, hostname=hostname) diff --git a/screens/otpworkflowscreen.py b/screens/otpworkflowscreen.py index f3710b3..519ff6b 100644 --- a/screens/otpworkflowscreen.py +++ b/screens/otpworkflowscreen.py @@ -1,6 +1,6 @@ # otp_workflow_screen.py -from typing import List +from typing import List, Optional from textual.app import ComposeResult from textual.screen import Screen @@ -12,7 +12,7 @@ from widgets.OTP_generate import OTPGenerator class OTPWorkflowScreen(Screen): """Screen that handles the OTP generation workflow without agent selection.""" - def __init__(self, selected_agents: List[Agent]): + def __init__(self, selected_agents: Optional[List[Agent]]): super().__init__() self.selected_agents = selected_agents diff --git a/widgets/amber_terminal_theme.py b/themes/amber_terminal_theme.py similarity index 100% rename from widgets/amber_terminal_theme.py rename to themes/amber_terminal_theme.py diff --git a/widgets/retro_terminal_theme.py b/themes/retro_terminal_theme.py similarity index 100% rename from widgets/retro_terminal_theme.py rename to themes/retro_terminal_theme.py diff --git a/utils/configmanager.py b/utils/configmanager.py index b099fa3..7b5aea5 100644 --- a/utils/configmanager.py +++ b/utils/configmanager.py @@ -24,6 +24,8 @@ T = TypeVar("T") logger = logging.getLogger(__name__) PROTECTED_KEYS = [ + "URL", + "TELEM_URL", "APPNAME", "LOG_LEVEL", "PATH_EXCLUSION_CONST", diff --git a/utils/setup.py b/utils/setup.py index 4277389..77bc5a0 100644 --- a/utils/setup.py +++ b/utils/setup.py @@ -122,7 +122,11 @@ def load_system_config() -> dict: def load_user_config(config_dir: Path) -> dict: user_config_path = config_dir / "user_config.json" if not user_config_path.exists(): - default_user_config = {"URL": "", "LOG_LEVEL": "INFO"} + default_user_config = { + "TELEMETRY": "FALSE", + "TEXTUAL_THEME": "gruvbox", + "EXTRAS": "NOTTODAY", + } with open(user_config_path, "w") as f: json.dump(default_user_config, f, indent=4) logging.debug(f"Created user config at {user_config_path}") diff --git a/utils/tui.py b/utils/tui.py index 8cadf33..9f32133 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -1,11 +1,13 @@ import logging import os import sys +from typing import Optional import dotenv from dotenv import set_key from textual.app import App, ComposeResult from textual.containers import Vertical +from textual.message import Message from textual.reactive import reactive from textual.screen import Screen from textual.widgets import ( @@ -18,25 +20,26 @@ from textual.widgets import ( Tabs, ) -from flows.otp import otp_activities_by_agent, otp_revoke +from flows.otp import otp_revoke from flows.prepPolicy import menu_policy_enforce from flows.quietAgent import findQuietAgents from models.agent import Agent from models.policy import Policy from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen +from screens.otpactivityscreen import OTPActivitiesScreen from screens.otpworkflowscreen import OTPWorkflowScreen from services.API import AirlockAPIWrapper from services.policyhandler import confirmUpdateAfromE +from themes.amber_terminal_theme import get_amber_terminal_theme +from themes.retro_terminal_theme import get_retro_terminal_theme from utils.configmanager import load_env from utils.setup import get_base_directory, load_user_config from utils.utils import open_directory from widgets.agentmoveoperations import AgentMoveOperations -from widgets.amber_terminal_theme import get_amber_terminal_theme from widgets.multiagentselector import MultiAgentSelector from widgets.OTP_generate import OTPGenerator from widgets.policytreewidget import PolicyTreeWidget from widgets.resultsdisplay import ResultsDisplay -from widgets.retro_terminal_theme import get_retro_terminal_theme from widgets.themeselector import ThemeSelector dotenv.load_dotenv() @@ -102,6 +105,7 @@ def _persist_user_theme(theme_name: str) -> None: # 1) SCREEN # --------------------------------------------------------------------------- class MainMenuScreen(Screen): + api: AirlockAPIWrapper current_tab = reactive("") BUTTON_DEFS = { @@ -110,19 +114,18 @@ class MainMenuScreen(Screen): "🖥️ - Find, Move, or Generate OTP for Agents", "move_agent_workflow_button", ), + ("📊 - OTP Activities By Agent", "otp_activities_button"), ("🔇 - Find Quiet Hosts", "find_quiet_button"), ], "policy": [ ("🔒 - Prepare Policy For Enforcement", "policy_prep_button"), ("🔄 - Update Audit Policies", "policy_audit_update_button"), - ("📊 - OTP Activities By Agent", "otp_activities_button"), ("❌ - Revoke OTPs", "otp_revoke_button"), ], } - def __init__(self, api: AirlockAPIWrapper) -> None: + def __init__(self) -> None: super().__init__() - self.api = api self.extras = load_env("EXTRAS") wd = load_env("WORKING_DIR") or os.getcwd() if not os.path.isdir(wd): @@ -156,6 +159,7 @@ class MainMenuScreen(Screen): yield Footer() def on_mount(self) -> None: + api = self.app.api self.switch_tab("agent_actions") # focus helpers @@ -321,34 +325,55 @@ class MainMenuScreen(Screen): self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices)) event.stop() return # Don't exit the app + case "otp_generate_button": # NEW: Push OTP workflow screen instead of legacy function self.app.push_screen(OTPWorkflowScreen(self.app.devices)) event.stop() return # Don't exit the app + case "find_quiet_button": _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {}) + case "otp_activities_button": - _PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {}) + # === FIXED: push the Textual OTPActivitiesScreen and return immediately === + # This must return so we don't fall through to the code that exits the app. + self.app.push_screen(OTPActivitiesScreen()) + event.stop() + return + 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 + # Only exit the UI loop when we explicitly queued a legacy job. + # The original flow used `self.app.exit()` after setting _PENDING_JOB so + # the outer loop could run legacy code. Keep that behavior only for legacy jobs. logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB) - self.app.exit() + if _PENDING_JOB and _PENDING_JOB[0] == "legacy": + # let the main loop pick up the legacy job + self.app.exit() # --------------------------------------------------------------------------- # 2) APP # --------------------------------------------------------------------------- -class Loxide(App): +class Loxide(App[Message]): + api: AirlockAPIWrapper + working_dir: str + policies: Optional[list[Policy]] + devices: Optional[list[Agent]] + CSS = """ #logo { width: 100%; @@ -358,7 +383,7 @@ class Loxide(App): """ BINDINGS = [ ("q", "quit", "Quit"), - ("d", "open_dir", "Open Directory"), + ("f", "open_fe", "Launch Explorer"), ] def __init__(self, api: AirlockAPIWrapper): @@ -398,20 +423,21 @@ class Loxide(App): self.register_theme(get_retro_terminal_theme()) self.register_theme(get_amber_terminal_theme()) self.theme = self._textual_theme - self.push_screen(MainMenuScreen(api)) + self.push_screen(MainMenuScreen()) def action_quit(self) -> None: global _PENDING_JOB _PENDING_JOB = None self.exit() - def action_open_dir(self) -> None: - # Refresh data before proceeding - self.refresh_data() - screen = self.screen_stack[-1] - if isinstance(screen, MainMenuScreen): - if screen.current_tab != "dir": - screen.switch_tab("dir") + def action_open_fe(self) -> None: + """Open the working directory in the OS file manager (footer binding).""" + path_to_open = self.working_dir or os.getcwd() + try: + open_directory(path_to_open) + except Exception as exc: + logger.error("Failed to open directory %s: %s", path_to_open, exc) + self.bell() # optional feedback # --------------------------------------------------------------------------- diff --git a/widgets/OTP_generate.py b/widgets/OTP_generate.py index 221394b..4ca7977 100644 --- a/widgets/OTP_generate.py +++ b/widgets/OTP_generate.py @@ -1,5 +1,5 @@ import logging -from typing import List +from typing import List, Optional from textual.containers import Horizontal, Vertical from textual.css.query import NoMatches @@ -31,7 +31,11 @@ class OTPGenerator(Widget): class OTPInfo(Message): def __init__( - self, devices: List[Agent], requestor: str, reasoning: str, duration: int + self, + devices: Optional[List[Agent]], + requestor: str, + reasoning: str, + duration: int, ): super().__init__() self.devices = devices @@ -243,7 +247,7 @@ class OTPGenerator(Widget): self.otp_generated = True # Access API from the app - this is the key change! - api = self.app.api + api = self.app.api # type: ignore output_lines = [ "Requested OTP Codes:", diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py index f4ea371..afe8e14 100644 --- a/widgets/multiagentselector.py +++ b/widgets/multiagentselector.py @@ -1,6 +1,6 @@ import difflib import re -from typing import List +from typing import List, Optional from textual.containers import Horizontal, Vertical from textual.css.query import NoMatches @@ -25,7 +25,7 @@ class MultiAgentSelector(Widget): super().__init__() self.selected_agents = selected_agents - def __init__(self, all_agents: List[Agent]): + def __init__(self, all_agents: Optional[List[Agent]]): super().__init__() self.all_agents = all_agents self._match_type = "exact" diff --git a/widgets/policytreewidget.py b/widgets/policytreewidget.py index f611d5d..5e01fdf 100644 --- a/widgets/policytreewidget.py +++ b/widgets/policytreewidget.py @@ -4,7 +4,7 @@ import logging from rich.text import Text from textual.containers import Horizontal, Vertical from textual.widget import Widget -from textual.widgets import Input, OptionList, Static, Tree +from textual.widgets import Input, OptionList, Static, Switch, Tree from textual.widgets.option_list import Option logger = logging.getLogger(__name__) @@ -19,29 +19,49 @@ class PolicyTreeWidget(Widget): self.devices = devices self.last_highlighted_node = None self.leaf_counts = defaultdict(int) + self.match_type = "Count" # Default to sorting by count def compose(self): + # Create the switch and its label + switch = Switch(value=False, id="match_switch") + switch.styles.margin = (0, 0, 0, 0) # top, right, bottom, left + switch.styles.padding = (0, 0, 0, 0) + + switch_label = Static("Sort: Count", id="match_switch_label") + switch_label.styles.margin = (1, 0, 0, 0) + switch_label.styles.padding = (0, 0, 0, 0) + + # Create the tree policy_tree = Tree("", id="policy_tree") # Label set in on_mount policy_tree.styles.width = "2fr" policy_tree.styles.height = "100%" + # Create the search box and details pane label = Static("Device Search:") search_box = Input( placeholder="Search policies or devices...", id="tree_search" ) details_pane = Static("", id="details_pane") + # Layout the UI with Horizontal(): yield policy_tree with Vertical() as right_pane: right_pane.styles.width = "3fr" + # Use a Horizontal container for the switch and label + with Horizontal() as switch_container: + switch_container.styles.height = 3 + switch_container.styles.margin = (0, 0, 0, 1) + switch_container.styles.padding = (0, 0, 0, 0) + yield switch + yield switch_label + # Add the search box and details pane yield label yield search_box yield details_pane def on_mount(self) -> None: self._precompute_leaf_counts() - # Update root label with total leaf count total_leaves = sum( self.leaf_counts.get(policy.groupid, 0) @@ -50,8 +70,9 @@ class PolicyTreeWidget(Widget): ) policy_tree = self.query_one("#policy_tree", Tree) policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})") - self._build_tree() + # Expand the root node + policy_tree.root.expand() def _precompute_leaf_counts(self): """Precompute leaf counts for each policy group.""" @@ -76,15 +97,21 @@ class PolicyTreeWidget(Widget): def _build_tree(self): policy_tree = self.query_one("#policy_tree", Tree) + policy_tree.clear() # Clear existing nodes node_map = {} # Sort top-level policies top_policies = [ p for p in self.policies if p.parent == "global-policy-settings" ] - top_policies.sort( - key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True - ) + + # Sort by count (default) or alphabetically + if getattr(self, "match_type", "Count") == "Count": + top_policies.sort( + key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True + ) + else: # Alphabetical + top_policies.sort(key=lambda p: p.name.lower()) for policy in top_policies: label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})" @@ -98,9 +125,13 @@ class PolicyTreeWidget(Widget): children_by_parent[policy.parent].append(policy) for parent_id, children in children_by_parent.items(): - children.sort( - key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True - ) + if getattr(self, "match_type", "Count") == "Count": + children.sort( + key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True + ) + else: # Alphabetical + children.sort(key=lambda p: p.name.lower()) + parent_node = node_map.get(parent_id) if parent_node: for policy in children: @@ -108,12 +139,17 @@ class PolicyTreeWidget(Widget): node = parent_node.add(label=label, data=policy) node_map[policy.groupid] = node - # Add devices (leaf nodes) + # Add devices (leaf nodes) - always sort alphabetically + devices_by_group = defaultdict(list) for device in self.devices: - group_id = device.groupid + devices_by_group[device.groupid].append(device) + + for group_id, devices in devices_by_group.items(): + devices.sort(key=lambda d: d.hostname.lower()) # Always sort alphabetically parent_node = node_map.get(group_id) if parent_node: - parent_node.add(label=device.hostname, data=device) + for device in devices: + parent_node.add(label=device.hostname, data=device) def _collect_tree_nodes(self, node, all_nodes): all_nodes.append(node) @@ -157,6 +193,13 @@ class PolicyTreeWidget(Widget): message.stop() + def on_switch_changed(self, event: Switch.Changed): + self.match_type = "Alpha" if event.value else "Count" + self.query_one("#match_switch_label", Static).update( + f"Sort: {self.match_type.capitalize()}" + ) + self._build_tree() + def on_input_submitted(self, message: Input.Submitted) -> None: self._remove_match_selector()