Merge remote-tracking branch 'origin/RustImplementation' into RustImplementation
Build Library / Build Library (push) Has been cancelled
Build Library / Build Library (push) Has been cancelled
This commit is contained in:
@@ -1,5 +1,5 @@
|
|||||||
{
|
{
|
||||||
"APPNAME": "AirlockTools",
|
"APPNAME": "Loxide",
|
||||||
"URL": "https://server:3129",
|
"URL": "https://server:3129",
|
||||||
"LOG_LEVEL": "INFO",
|
"LOG_LEVEL": "INFO",
|
||||||
"BAD_PATH_PARTS": ["users","wwwroot","windows\\temp","windows\\task","windows\\system32","startup", "windows\\fonts","Recycle.Bin","AppData","programdata", "Solarwinds","kaseya"],
|
"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,
|
"PATH_EXCLUSION_CONST": 4,
|
||||||
"MIN_FILES_FOR_PATH": 4,
|
"MIN_FILES_FOR_PATH": 4,
|
||||||
"VT_THREAT_TOLERANCE": 4,
|
"VT_THREAT_TOLERANCE": 4,
|
||||||
|
"TELEMETRY": "FALSE",
|
||||||
|
"TELEM_URL": "",
|
||||||
"POLICY_MAP_ENF_AUD": {
|
"POLICY_MAP_ENF_AUD": {
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -29,44 +29,6 @@ from utils.utils import colorText, get_sanitized_input
|
|||||||
logger = logging.getLogger(__name__)
|
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):
|
def otp_activities_by_agent(api: AirlockAPIWrapper):
|
||||||
activeagents = api.otp_find_active()
|
activeagents = api.otp_find_active()
|
||||||
awaitingagents = api.otp_find_awaiting()
|
awaitingagents = api.otp_find_awaiting()
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
from typing import List
|
from typing import List, Optional
|
||||||
|
|
||||||
from textual.app import ComposeResult
|
from textual.app import ComposeResult
|
||||||
from textual.screen import Screen
|
from textual.screen import Screen
|
||||||
@@ -12,7 +12,7 @@ from widgets.resultsdisplay import ResultsDisplay
|
|||||||
class MoveAgentWorkflowScreen(Screen):
|
class MoveAgentWorkflowScreen(Screen):
|
||||||
"""Screen that handles the agent movement workflow."""
|
"""Screen that handles the agent movement workflow."""
|
||||||
|
|
||||||
def __init__(self, all_agents: List[Agent]):
|
def __init__(self, all_agents: Optional[List[Agent]]):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.all_agents = all_agents
|
self.all_agents = all_agents
|
||||||
self.selected_agents = None
|
self.selected_agents = None
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -1,6 +1,6 @@
|
|||||||
# otp_workflow_screen.py
|
# otp_workflow_screen.py
|
||||||
|
|
||||||
from typing import List
|
from typing import List, Optional
|
||||||
|
|
||||||
from textual.app import ComposeResult
|
from textual.app import ComposeResult
|
||||||
from textual.screen import Screen
|
from textual.screen import Screen
|
||||||
@@ -12,7 +12,7 @@ from widgets.OTP_generate import OTPGenerator
|
|||||||
class OTPWorkflowScreen(Screen):
|
class OTPWorkflowScreen(Screen):
|
||||||
"""Screen that handles the OTP generation workflow without agent selection."""
|
"""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__()
|
super().__init__()
|
||||||
self.selected_agents = selected_agents
|
self.selected_agents = selected_agents
|
||||||
|
|
||||||
|
|||||||
@@ -24,6 +24,8 @@ T = TypeVar("T")
|
|||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
PROTECTED_KEYS = [
|
PROTECTED_KEYS = [
|
||||||
|
"URL",
|
||||||
|
"TELEM_URL",
|
||||||
"APPNAME",
|
"APPNAME",
|
||||||
"LOG_LEVEL",
|
"LOG_LEVEL",
|
||||||
"PATH_EXCLUSION_CONST",
|
"PATH_EXCLUSION_CONST",
|
||||||
|
|||||||
+5
-1
@@ -122,7 +122,11 @@ def load_system_config() -> dict:
|
|||||||
def load_user_config(config_dir: Path) -> dict:
|
def load_user_config(config_dir: Path) -> dict:
|
||||||
user_config_path = config_dir / "user_config.json"
|
user_config_path = config_dir / "user_config.json"
|
||||||
if not user_config_path.exists():
|
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:
|
with open(user_config_path, "w") as f:
|
||||||
json.dump(default_user_config, f, indent=4)
|
json.dump(default_user_config, f, indent=4)
|
||||||
logging.debug(f"Created user config at {user_config_path}")
|
logging.debug(f"Created user config at {user_config_path}")
|
||||||
|
|||||||
+43
-17
@@ -1,11 +1,13 @@
|
|||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
import dotenv
|
import dotenv
|
||||||
from dotenv import set_key
|
from dotenv import set_key
|
||||||
from textual.app import App, ComposeResult
|
from textual.app import App, ComposeResult
|
||||||
from textual.containers import Vertical
|
from textual.containers import Vertical
|
||||||
|
from textual.message import Message
|
||||||
from textual.reactive import reactive
|
from textual.reactive import reactive
|
||||||
from textual.screen import Screen
|
from textual.screen import Screen
|
||||||
from textual.widgets import (
|
from textual.widgets import (
|
||||||
@@ -18,25 +20,26 @@ from textual.widgets import (
|
|||||||
Tabs,
|
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.prepPolicy import menu_policy_enforce
|
||||||
from flows.quietAgent import findQuietAgents
|
from flows.quietAgent import findQuietAgents
|
||||||
from models.agent import Agent
|
from models.agent import Agent
|
||||||
from models.policy import Policy
|
from models.policy import Policy
|
||||||
from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
||||||
|
from screens.otpactivityscreen import OTPActivitiesScreen
|
||||||
from screens.otpworkflowscreen import OTPWorkflowScreen
|
from screens.otpworkflowscreen import OTPWorkflowScreen
|
||||||
from services.API import AirlockAPIWrapper
|
from services.API import AirlockAPIWrapper
|
||||||
from services.policyhandler import confirmUpdateAfromE
|
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.configmanager import load_env
|
||||||
from utils.setup import get_base_directory, load_user_config
|
from utils.setup import get_base_directory, load_user_config
|
||||||
from utils.utils import open_directory
|
from utils.utils import open_directory
|
||||||
from widgets.agentmoveoperations import AgentMoveOperations
|
from widgets.agentmoveoperations import AgentMoveOperations
|
||||||
from widgets.amber_terminal_theme import get_amber_terminal_theme
|
|
||||||
from widgets.multiagentselector import MultiAgentSelector
|
from widgets.multiagentselector import MultiAgentSelector
|
||||||
from widgets.OTP_generate import OTPGenerator
|
from widgets.OTP_generate import OTPGenerator
|
||||||
from widgets.policytreewidget import PolicyTreeWidget
|
from widgets.policytreewidget import PolicyTreeWidget
|
||||||
from widgets.resultsdisplay import ResultsDisplay
|
from widgets.resultsdisplay import ResultsDisplay
|
||||||
from widgets.retro_terminal_theme import get_retro_terminal_theme
|
|
||||||
from widgets.themeselector import ThemeSelector
|
from widgets.themeselector import ThemeSelector
|
||||||
|
|
||||||
dotenv.load_dotenv()
|
dotenv.load_dotenv()
|
||||||
@@ -102,6 +105,7 @@ def _persist_user_theme(theme_name: str) -> None:
|
|||||||
# 1) SCREEN
|
# 1) SCREEN
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
class MainMenuScreen(Screen):
|
class MainMenuScreen(Screen):
|
||||||
|
api: AirlockAPIWrapper
|
||||||
current_tab = reactive("")
|
current_tab = reactive("")
|
||||||
|
|
||||||
BUTTON_DEFS = {
|
BUTTON_DEFS = {
|
||||||
@@ -110,19 +114,18 @@ class MainMenuScreen(Screen):
|
|||||||
"🖥️ - Find, Move, or Generate OTP for Agents",
|
"🖥️ - Find, Move, or Generate OTP for Agents",
|
||||||
"move_agent_workflow_button",
|
"move_agent_workflow_button",
|
||||||
),
|
),
|
||||||
|
("📊 - OTP Activities By Agent", "otp_activities_button"),
|
||||||
("🔇 - Find Quiet Hosts", "find_quiet_button"),
|
("🔇 - Find Quiet Hosts", "find_quiet_button"),
|
||||||
],
|
],
|
||||||
"policy": [
|
"policy": [
|
||||||
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
|
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
|
||||||
("🔄 - Update Audit Policies", "policy_audit_update_button"),
|
("🔄 - Update Audit Policies", "policy_audit_update_button"),
|
||||||
("📊 - OTP Activities By Agent", "otp_activities_button"),
|
|
||||||
("❌ - Revoke OTPs", "otp_revoke_button"),
|
("❌ - Revoke OTPs", "otp_revoke_button"),
|
||||||
],
|
],
|
||||||
}
|
}
|
||||||
|
|
||||||
def __init__(self, api: AirlockAPIWrapper) -> None:
|
def __init__(self) -> None:
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.api = api
|
|
||||||
self.extras = load_env("EXTRAS")
|
self.extras = load_env("EXTRAS")
|
||||||
wd = load_env("WORKING_DIR") or os.getcwd()
|
wd = load_env("WORKING_DIR") or os.getcwd()
|
||||||
if not os.path.isdir(wd):
|
if not os.path.isdir(wd):
|
||||||
@@ -156,6 +159,7 @@ class MainMenuScreen(Screen):
|
|||||||
yield Footer()
|
yield Footer()
|
||||||
|
|
||||||
def on_mount(self) -> None:
|
def on_mount(self) -> None:
|
||||||
|
api = self.app.api
|
||||||
self.switch_tab("agent_actions")
|
self.switch_tab("agent_actions")
|
||||||
|
|
||||||
# focus helpers
|
# focus helpers
|
||||||
@@ -321,34 +325,55 @@ class MainMenuScreen(Screen):
|
|||||||
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
|
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
|
||||||
event.stop()
|
event.stop()
|
||||||
return # Don't exit the app
|
return # Don't exit the app
|
||||||
|
|
||||||
case "otp_generate_button":
|
case "otp_generate_button":
|
||||||
# NEW: Push OTP workflow screen instead of legacy function
|
# NEW: Push OTP workflow screen instead of legacy function
|
||||||
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
|
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
|
||||||
event.stop()
|
event.stop()
|
||||||
return # Don't exit the app
|
return # Don't exit the app
|
||||||
|
|
||||||
case "find_quiet_button":
|
case "find_quiet_button":
|
||||||
_PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {})
|
_PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {})
|
||||||
|
|
||||||
case "otp_activities_button":
|
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":
|
case "otp_revoke_button":
|
||||||
_PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {})
|
_PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {})
|
||||||
|
|
||||||
case "policy_prep_button":
|
case "policy_prep_button":
|
||||||
_PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {})
|
_PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {})
|
||||||
|
|
||||||
case "policy_audit_update_button":
|
case "policy_audit_update_button":
|
||||||
_PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {})
|
_PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {})
|
||||||
|
|
||||||
case _:
|
case _:
|
||||||
self.app.bell()
|
self.app.bell()
|
||||||
logger.warning("Unknown button pressed: %s", button_id)
|
logger.warning("Unknown button pressed: %s", button_id)
|
||||||
return
|
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)
|
logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB)
|
||||||
|
if _PENDING_JOB and _PENDING_JOB[0] == "legacy":
|
||||||
|
# let the main loop pick up the legacy job
|
||||||
self.app.exit()
|
self.app.exit()
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# 2) APP
|
# 2) APP
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
class Loxide(App):
|
class Loxide(App[Message]):
|
||||||
|
api: AirlockAPIWrapper
|
||||||
|
working_dir: str
|
||||||
|
policies: Optional[list[Policy]]
|
||||||
|
devices: Optional[list[Agent]]
|
||||||
|
|
||||||
CSS = """
|
CSS = """
|
||||||
#logo {
|
#logo {
|
||||||
width: 100%;
|
width: 100%;
|
||||||
@@ -358,7 +383,7 @@ class Loxide(App):
|
|||||||
"""
|
"""
|
||||||
BINDINGS = [
|
BINDINGS = [
|
||||||
("q", "quit", "Quit"),
|
("q", "quit", "Quit"),
|
||||||
("d", "open_dir", "Open Directory"),
|
("f", "open_fe", "Launch Explorer"),
|
||||||
]
|
]
|
||||||
|
|
||||||
def __init__(self, api: AirlockAPIWrapper):
|
def __init__(self, api: AirlockAPIWrapper):
|
||||||
@@ -398,20 +423,21 @@ class Loxide(App):
|
|||||||
self.register_theme(get_retro_terminal_theme())
|
self.register_theme(get_retro_terminal_theme())
|
||||||
self.register_theme(get_amber_terminal_theme())
|
self.register_theme(get_amber_terminal_theme())
|
||||||
self.theme = self._textual_theme
|
self.theme = self._textual_theme
|
||||||
self.push_screen(MainMenuScreen(api))
|
self.push_screen(MainMenuScreen())
|
||||||
|
|
||||||
def action_quit(self) -> None:
|
def action_quit(self) -> None:
|
||||||
global _PENDING_JOB
|
global _PENDING_JOB
|
||||||
_PENDING_JOB = None
|
_PENDING_JOB = None
|
||||||
self.exit()
|
self.exit()
|
||||||
|
|
||||||
def action_open_dir(self) -> None:
|
def action_open_fe(self) -> None:
|
||||||
# Refresh data before proceeding
|
"""Open the working directory in the OS file manager (footer binding)."""
|
||||||
self.refresh_data()
|
path_to_open = self.working_dir or os.getcwd()
|
||||||
screen = self.screen_stack[-1]
|
try:
|
||||||
if isinstance(screen, MainMenuScreen):
|
open_directory(path_to_open)
|
||||||
if screen.current_tab != "dir":
|
except Exception as exc:
|
||||||
screen.switch_tab("dir")
|
logger.error("Failed to open directory %s: %s", path_to_open, exc)
|
||||||
|
self.bell() # optional feedback
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
import logging
|
import logging
|
||||||
from typing import List
|
from typing import List, Optional
|
||||||
|
|
||||||
from textual.containers import Horizontal, Vertical
|
from textual.containers import Horizontal, Vertical
|
||||||
from textual.css.query import NoMatches
|
from textual.css.query import NoMatches
|
||||||
@@ -31,7 +31,11 @@ class OTPGenerator(Widget):
|
|||||||
|
|
||||||
class OTPInfo(Message):
|
class OTPInfo(Message):
|
||||||
def __init__(
|
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__()
|
super().__init__()
|
||||||
self.devices = devices
|
self.devices = devices
|
||||||
@@ -243,7 +247,7 @@ class OTPGenerator(Widget):
|
|||||||
self.otp_generated = True
|
self.otp_generated = True
|
||||||
|
|
||||||
# Access API from the app - this is the key change!
|
# Access API from the app - this is the key change!
|
||||||
api = self.app.api
|
api = self.app.api # type: ignore
|
||||||
|
|
||||||
output_lines = [
|
output_lines = [
|
||||||
"Requested OTP Codes:",
|
"Requested OTP Codes:",
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
import difflib
|
import difflib
|
||||||
import re
|
import re
|
||||||
from typing import List
|
from typing import List, Optional
|
||||||
|
|
||||||
from textual.containers import Horizontal, Vertical
|
from textual.containers import Horizontal, Vertical
|
||||||
from textual.css.query import NoMatches
|
from textual.css.query import NoMatches
|
||||||
@@ -25,7 +25,7 @@ class MultiAgentSelector(Widget):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
self.selected_agents = selected_agents
|
self.selected_agents = selected_agents
|
||||||
|
|
||||||
def __init__(self, all_agents: List[Agent]):
|
def __init__(self, all_agents: Optional[List[Agent]]):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.all_agents = all_agents
|
self.all_agents = all_agents
|
||||||
self._match_type = "exact"
|
self._match_type = "exact"
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ import logging
|
|||||||
from rich.text import Text
|
from rich.text import Text
|
||||||
from textual.containers import Horizontal, Vertical
|
from textual.containers import Horizontal, Vertical
|
||||||
from textual.widget import Widget
|
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
|
from textual.widgets.option_list import Option
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -19,29 +19,49 @@ class PolicyTreeWidget(Widget):
|
|||||||
self.devices = devices
|
self.devices = devices
|
||||||
self.last_highlighted_node = None
|
self.last_highlighted_node = None
|
||||||
self.leaf_counts = defaultdict(int)
|
self.leaf_counts = defaultdict(int)
|
||||||
|
self.match_type = "Count" # Default to sorting by count
|
||||||
|
|
||||||
def compose(self):
|
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 = Tree("", id="policy_tree") # Label set in on_mount
|
||||||
policy_tree.styles.width = "2fr"
|
policy_tree.styles.width = "2fr"
|
||||||
policy_tree.styles.height = "100%"
|
policy_tree.styles.height = "100%"
|
||||||
|
|
||||||
|
# Create the search box and details pane
|
||||||
label = Static("Device Search:")
|
label = Static("Device Search:")
|
||||||
search_box = Input(
|
search_box = Input(
|
||||||
placeholder="Search policies or devices...", id="tree_search"
|
placeholder="Search policies or devices...", id="tree_search"
|
||||||
)
|
)
|
||||||
details_pane = Static("", id="details_pane")
|
details_pane = Static("", id="details_pane")
|
||||||
|
|
||||||
|
# Layout the UI
|
||||||
with Horizontal():
|
with Horizontal():
|
||||||
yield policy_tree
|
yield policy_tree
|
||||||
with Vertical() as right_pane:
|
with Vertical() as right_pane:
|
||||||
right_pane.styles.width = "3fr"
|
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 label
|
||||||
yield search_box
|
yield search_box
|
||||||
yield details_pane
|
yield details_pane
|
||||||
|
|
||||||
def on_mount(self) -> None:
|
def on_mount(self) -> None:
|
||||||
self._precompute_leaf_counts()
|
self._precompute_leaf_counts()
|
||||||
|
|
||||||
# Update root label with total leaf count
|
# Update root label with total leaf count
|
||||||
total_leaves = sum(
|
total_leaves = sum(
|
||||||
self.leaf_counts.get(policy.groupid, 0)
|
self.leaf_counts.get(policy.groupid, 0)
|
||||||
@@ -50,8 +70,9 @@ class PolicyTreeWidget(Widget):
|
|||||||
)
|
)
|
||||||
policy_tree = self.query_one("#policy_tree", Tree)
|
policy_tree = self.query_one("#policy_tree", Tree)
|
||||||
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
|
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
|
||||||
|
|
||||||
self._build_tree()
|
self._build_tree()
|
||||||
|
# Expand the root node
|
||||||
|
policy_tree.root.expand()
|
||||||
|
|
||||||
def _precompute_leaf_counts(self):
|
def _precompute_leaf_counts(self):
|
||||||
"""Precompute leaf counts for each policy group."""
|
"""Precompute leaf counts for each policy group."""
|
||||||
@@ -76,15 +97,21 @@ class PolicyTreeWidget(Widget):
|
|||||||
|
|
||||||
def _build_tree(self):
|
def _build_tree(self):
|
||||||
policy_tree = self.query_one("#policy_tree", Tree)
|
policy_tree = self.query_one("#policy_tree", Tree)
|
||||||
|
policy_tree.clear() # Clear existing nodes
|
||||||
node_map = {}
|
node_map = {}
|
||||||
|
|
||||||
# Sort top-level policies
|
# Sort top-level policies
|
||||||
top_policies = [
|
top_policies = [
|
||||||
p for p in self.policies if p.parent == "global-policy-settings"
|
p for p in self.policies if p.parent == "global-policy-settings"
|
||||||
]
|
]
|
||||||
|
|
||||||
|
# Sort by count (default) or alphabetically
|
||||||
|
if getattr(self, "match_type", "Count") == "Count":
|
||||||
top_policies.sort(
|
top_policies.sort(
|
||||||
key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
|
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:
|
for policy in top_policies:
|
||||||
label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})"
|
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)
|
children_by_parent[policy.parent].append(policy)
|
||||||
|
|
||||||
for parent_id, children in children_by_parent.items():
|
for parent_id, children in children_by_parent.items():
|
||||||
|
if getattr(self, "match_type", "Count") == "Count":
|
||||||
children.sort(
|
children.sort(
|
||||||
key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
|
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)
|
parent_node = node_map.get(parent_id)
|
||||||
if parent_node:
|
if parent_node:
|
||||||
for policy in children:
|
for policy in children:
|
||||||
@@ -108,11 +139,16 @@ class PolicyTreeWidget(Widget):
|
|||||||
node = parent_node.add(label=label, data=policy)
|
node = parent_node.add(label=label, data=policy)
|
||||||
node_map[policy.groupid] = node
|
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:
|
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)
|
parent_node = node_map.get(group_id)
|
||||||
if parent_node:
|
if parent_node:
|
||||||
|
for device in devices:
|
||||||
parent_node.add(label=device.hostname, data=device)
|
parent_node.add(label=device.hostname, data=device)
|
||||||
|
|
||||||
def _collect_tree_nodes(self, node, all_nodes):
|
def _collect_tree_nodes(self, node, all_nodes):
|
||||||
@@ -157,6 +193,13 @@ class PolicyTreeWidget(Widget):
|
|||||||
|
|
||||||
message.stop()
|
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:
|
def on_input_submitted(self, message: Input.Submitted) -> None:
|
||||||
self._remove_match_selector()
|
self._remove_match_selector()
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user