Major step towards unification of the UI Implementation of the Back Feature, splitting of TUI files back into subfolders
This commit is contained in:
@@ -0,0 +1,393 @@
|
||||
# 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 <https://www.gnu.org/licenses/>.
|
||||
|
||||
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()
|
||||
Reference in New Issue
Block a user