# Copyright (C) 2025 James Brotosky, Brandon Wickline # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Affero General Public License as published # by the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU Affero General Public License for more details. # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . from __future__ import annotations import logging from typing import List, Optional import pandas as pd from textual.app import ComposeResult from textual.binding import Binding from textual.containers import Horizontal, Vertical from textual.message import Message from textual.screen import Screen from textual.widgets import Button, DataTable, Footer, Header, Static logger = logging.getLogger(__name__) class OTPRevokeWidget(Static): """ Widget for managing OTP session revocation. Displays active OTP sessions and allows selection for revocation. """ class SessionsRevoked(Message): """Message sent when sessions are revoked.""" def __init__(self, revoked_sessions: List[dict]): super().__init__() self.revoked_sessions = revoked_sessions DEFAULT_CSS = """ OTPRevokeWidget { height: 1fr; } #main_container { width: 100%; height: 100%; layout: vertical; } #sessions_container { height: 1fr; border: none; padding: 1; } #button_container { height: auto; padding: 1; align: center middle; } #result_container { height: auto; max-height: 10; border: solid #444444; padding: 1; margin: 1; overflow-y: auto; } .panel-title { text-style: bold; margin: 0 0 1 0; } """ def compose(self) -> ComposeResult: with Vertical(id="main_container"): # Sessions table yield Static("OTP Sessions", classes="panel-title") with Vertical(id="sessions_container"): self.sessions_table = DataTable(id="sessions_table") self.sessions_table.styles.width = "100%" self.sessions_table.styles.height = "1fr" yield self.sessions_table # Action buttons with Horizontal(id="button_container"): self.refresh_button = Button("🔄 Refresh", id="refresh_btn") self.refresh_button.styles.width = "15%" self.refresh_button.styles.margin = (1, 1, 1, 1) yield self.refresh_button self.select_all_button = Button("☑️ Select All", id="select_all_btn") self.select_all_button.styles.width = "15%" self.select_all_button.styles.margin = (1, 1, 1, 1) yield self.select_all_button self.select_none_button = Button( "❌ Clear Selection", id="select_none_btn" ) self.select_none_button.styles.width = "20%" self.select_none_button.styles.margin = (1, 1, 1, 1) yield self.select_none_button self.revoke_button = Button( "🛑 Revoke Selected", id="revoke_btn", variant="error" ) self.revoke_button.styles.width = "20%" self.revoke_button.styles.margin = (1, 1, 1, 1) yield self.revoke_button # Results display with Vertical(id="result_container"): yield Static("Revocation Results", classes="panel-title") self.results_display = Static("No actions performed yet.") yield self.results_display async def on_mount(self) -> None: """Initialize the widget when mounted.""" # Configure sessions table self.sessions_table.clear() self.sessions_table.add_columns( "☐", "OTP ID", "Hostname", "Status", "Purpose", "Granted" ) # Enable row selection with checkbox column self.sessions_table.cursor_type = "row" try: self.sessions_table.zebra_stripes = True except Exception: pass # Initialize state self._sessions_df: Optional[pd.DataFrame] = None self._filtered_df: Optional[pd.DataFrame] = None self._selected_otpids: set = set() async def load_sessions_from_api(self, api) -> None: """Load active OTP sessions from the API.""" try: # Fetch only active sessions active_df = api.otp_find_active() # Ensure we have a DataFrame if not isinstance(active_df, pd.DataFrame): active_df = pd.DataFrame(active_df) # Add status column active_df["status"] = "active" # Sort by otpid if column exists if "otpid" in active_df.columns and not active_df.empty: active_df = active_df.sort_values(by="otpid", ascending=False) # Store the full dataframe self._sessions_df = active_df self._filtered_df = active_df.copy() # Display in table await self._refresh_table() # Update status active_count = len(active_df) status_msg = f"Loaded {active_count} active sessions" logger.info(status_msg) self.results_display.update(status_msg) except Exception as e: logger.exception(f"Failed to load OTP sessions: {e}") self.results_display.update(f"Error loading sessions: {str(e)}") async def _refresh_table(self) -> None: """Refresh the table display with current filtered data.""" if self._filtered_df is None or self._filtered_df.empty: self.sessions_table.clear() return # Ensure expected columns exist expected_cols = ["otpid", "hostname", "status", "purpose", "granted"] for col in expected_cols: if col not in self._filtered_df.columns: self._filtered_df[col] = "" # Clear and repopulate table self.sessions_table.clear(columns=False) for _, row in self._filtered_df.iterrows(): otpid = str(row.get("otpid", "")) # Check if this row is selected checkbox = "☑️" if otpid in self._selected_otpids else "☐" self.sessions_table.add_row( checkbox, str(otpid), str(row.get("hostname", "")), str(row.get("status", "")), str(row.get("purpose", "")), str(row.get("granted", "")), ) async def on_button_pressed(self, event) -> None: """Handle button presses.""" btn = event.button if btn.id == "refresh_btn": # Refresh sessions api = getattr(self.app, "api", None) if api: await self.load_sessions_from_api(api) elif btn.id == "select_all_btn": # Select all visible rows if self._filtered_df is not None: self._selected_otpids = set(str(x) for x in self._filtered_df["otpid"]) await self._refresh_table() elif btn.id == "select_none_btn": # Clear selection self._selected_otpids.clear() await self._refresh_table() elif btn.id == "revoke_btn": # Revoke selected sessions await self._revoke_selected() async def on_data_table_row_selected(self, event) -> None: """Handle row selection in the table.""" if event.data_table != self.sessions_table: return try: # Get the row index from the cursor row row_index = self.sessions_table.cursor_row if self._filtered_df is not None and row_index < len(self._filtered_df): # Get the OTP ID for this row otpid = str(self._filtered_df.iloc[row_index]["otpid"]) # Toggle selection if otpid in self._selected_otpids: self._selected_otpids.remove(otpid) else: self._selected_otpids.add(otpid) # Refresh table to update checkbox await self._refresh_table() # Restore cursor position self.sessions_table.move_cursor(row=row_index) except Exception as e: logger.exception(f"Error handling row selection: {e}") async def _revoke_selected(self) -> None: """Revoke the selected OTP sessions.""" if not self._selected_otpids: self.results_display.update("❌ No sessions selected for revocation") return api = getattr(self.app, "api", None) if not api: self.results_display.update("❌ API not available") return # Collect results results = [] success_count = 0 failure_count = 0 for otpid in self._selected_otpids: try: # Get hostname for this session hostname = "Unknown" if self._sessions_df is not None: # Convert otpid to same type as in DataFrame for comparison otpid_compare = otpid if len(self._sessions_df) > 0: first_otpid = self._sessions_df["otpid"].iloc[0] if isinstance(first_otpid, int): otpid_compare = int(otpid) match = self._sessions_df[ self._sessions_df["otpid"] == otpid_compare ] if not match.empty: hostname = match.iloc[0].get("hostname", "Unknown") # Revoke the session result = api.otp_revoke(otpid) if result and result.get("status") != "error": success_count += 1 results.append(f"Revoked OTP {otpid} for {hostname}") logger.info(f"Revoked OTP {otpid} for {hostname}: {result}") else: failure_count += 1 error_msg = ( result.get("message", "Unknown error") if result else "No response" ) results.append( f"❌ Failed to revoke OTP {otpid} for {hostname}: {error_msg}" ) logger.error(f"Failed to revoke OTP {otpid}: {error_msg}") except Exception as e: failure_count += 1 results.append(f"❌ Error revoking OTP {otpid}: {str(e)}") logger.exception(f"Exception revoking OTP {otpid}: {e}") # Update results display summary = ( f"Revocation complete: {success_count} succeeded, {failure_count} failed\n" ) details = "\n".join(results[-5:]) # Show last 5 results if len(results) > 5: details = f"... (showing last 5 of {len(results)} results)\n" + details self.results_display.update(summary + details) # Clear selection and refresh self._selected_otpids.clear() await self.load_sessions_from_api(api) # Post message about revoked sessions if success_count > 0: self.post_message(self.SessionsRevoked(results)) class OTPRevokeScreen(Screen): """ Main screen for OTP session revocation workflow. This replaces the otp_revoke function from otp.py. """ BINDINGS = [ Binding("escape", "go_back", "Back"), Binding("q", "main_menu", "Main Menu"), Binding("r", "refresh", "Refresh"), Binding("a", "select_all", "Select All"), Binding("n", "select_none", "Clear Selection"), Binding("d", "revoke", "Revoke Selected"), ] def compose(self) -> ComposeResult: yield Header(show_clock=True) self.widget = OTPRevokeWidget() yield self.widget yield Footer() async def on_mount(self) -> None: """Load sessions when screen mounts.""" api = getattr(self.app, "api", None) if api: await self.widget.load_sessions_from_api(api) else: logger.warning("OTPRevokeScreen mounted but no self.app.api found.") async def action_refresh(self) -> None: """Refresh the sessions list.""" api = getattr(self.app, "api", None) if api: await self.widget.load_sessions_from_api(api) async def action_select_all(self) -> None: """Select all visible sessions.""" if self.widget._filtered_df is not None: self.widget._selected_otpids = set( str(x) for x in self.widget._filtered_df["otpid"] ) await self.widget._refresh_table() async def action_select_none(self) -> None: """Clear all selections.""" self.widget._selected_otpids.clear() await self.widget._refresh_table() async def action_revoke(self) -> None: """Revoke selected sessions.""" await self.widget._revoke_selected() async def action_go_back(self) -> None: """Go back to previous screen.""" await self.app.pop_screen() async def action_main_menu(self) -> None: """Go back to main menu.""" while len(self.app.screen_stack) > 2: await self.app.pop_screen()