393 lines
13 KiB
Python
393 lines
13 KiB
Python
# 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;
|
|
width: 100%;
|
|
padding: 1;
|
|
align: center middle;
|
|
}
|
|
#button_container Button {
|
|
min-width: 16;
|
|
margin: 0 1;
|
|
}
|
|
#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"):
|
|
yield Button("Refresh", id="refresh_btn")
|
|
yield Button("Select All", id="select_all_btn")
|
|
yield Button("Clear Selection", id="select_none_btn")
|
|
yield Button("Revoke Selected", id="revoke_btn", variant="error")
|
|
|
|
# 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
|
|
and not self._filtered_df.empty
|
|
and "otpid" in self._filtered_df.columns
|
|
):
|
|
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 not self._filtered_df.empty
|
|
and "otpid" in self._filtered_df.columns
|
|
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
|
|
and not self.widget._filtered_df.empty
|
|
and "otpid" in self.widget._filtered_df.columns
|
|
):
|
|
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()
|