Files
AirlockTools/screens/otpactivityscreen.py
T

679 lines
26 KiB
Python

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)