Compare commits
82 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 23529f1f94 | |||
| 963aa3bcfb | |||
| a609a088d5 | |||
| ccfdf4e7ab | |||
| a860dce421 | |||
| e19b748bac | |||
| 01b80429f8 | |||
| 6a76c2ded7 | |||
| f1c5080c97 | |||
| 72681218e7 | |||
| ededef6d30 | |||
| 1aae27b0f4 | |||
| 663dbc3cc2 | |||
| d2d181de4a | |||
| 31d45ca1db | |||
| a086956b48 | |||
| 76cfe62f08 | |||
| 5b53612644 | |||
| 2ed1d2e4cd | |||
| 040a513bfb | |||
| 423e9e8208 | |||
| 7f4edcba46 | |||
| e1e0cb7ac7 | |||
| 5555747422 | |||
| 729b45f52a | |||
| 09d2c125cd | |||
| fbd5b8b4b8 | |||
| 66bb21ed88 | |||
| a7b659c951 | |||
| 53f0b548b0 | |||
| 66eb101c5d | |||
| 22101c1eba | |||
| fc17c869fc | |||
| 1dbbcff5d5 | |||
| f080b0034f | |||
| 57d0f12000 | |||
| 0dbc744471 | |||
| 7a912bddab | |||
| 24211c318b | |||
| 630e0a3cdf | |||
| 797d0f4462 | |||
| 59bb97ec4e | |||
| 0ac3b54d89 | |||
| 154a7efcc8 | |||
| ab5f00d8e7 | |||
| 3ab803c12e | |||
| 98cb23e5ea | |||
| 0aabbfd36e | |||
| b19eeb6c96 | |||
| 76bd3a6087 | |||
| 6ab9de413f | |||
| e5b4b9d959 | |||
| 1d9caadaf3 | |||
| f3c1d97d28 | |||
| 99d7b5f74e | |||
| 6327adeabf | |||
| b289e9324c | |||
| a80c2ca1e1 | |||
| dd206eb272 | |||
| 87ab1e3b28 | |||
| 303ecd8368 | |||
| 697d923172 | |||
| d0fc34fdc7 | |||
| 1f06404a16 | |||
| eb1d710d07 | |||
| fbb5cc396b | |||
| a7fc1b71e1 | |||
| 4a47cbe661 | |||
| 2f41b33dd4 | |||
| 89654d3a8c | |||
| b198362ac8 | |||
| 32e296238b | |||
| 19fab9b703 | |||
| 3ee762a0a1 | |||
| e36e5343d7 | |||
| ecdd991333 | |||
| 6a5a2b5809 | |||
| 2604665247 | |||
| bf5c7d156b | |||
| 0ebb42dcbd | |||
| ac6238873b | |||
| f0e77db414 |
@@ -2,8 +2,8 @@ name: Build Library
|
||||
run-name: ${{ gitea.actor }}
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- RustImplementation
|
||||
branches-ignore:
|
||||
- master
|
||||
paths:
|
||||
- airlock_libs/**
|
||||
|
||||
@@ -14,7 +14,7 @@ jobs:
|
||||
- name: Install Prerequisites
|
||||
run: |
|
||||
apt update
|
||||
apt install curl git python3 pip pkg-config openssl libssl-dev patchelf binutils-mingw-w64-x86-64 mingw-w64 -y
|
||||
apt install curl git python3 pip pkg-config openssl libssl-dev patchelf binutils-mingw-w64-x86-64 mingw-w64 protobuf-compiler -y
|
||||
curl https://sh.rustup.rs -sSf | sh -s -- -y
|
||||
pip install maturin twine --break-system-packages
|
||||
|
||||
|
||||
@@ -14,3 +14,6 @@ securitytest.py
|
||||
system_config.json
|
||||
Development/
|
||||
AirlockTools_client*/
|
||||
*.build/
|
||||
*.dist/
|
||||
*.onefile-build/
|
||||
@@ -1,85 +0,0 @@
|
||||
# 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/>.
|
||||
|
||||
|
||||
# TODO Continue implementing logger
|
||||
# TODO Add input sanitation and CSV injection prevention
|
||||
# TODO Continue OTP and Local approval rewrites
|
||||
# TODO Explore pywin32
|
||||
# TODO Fix Requirements.txt
|
||||
# TODO Create Generic system_config.json for gitea
|
||||
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import dotenv
|
||||
import urllib3
|
||||
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.security import getAPI
|
||||
from utils.setup import get_base_directory, setup
|
||||
from utils.TUI import run_Loxide
|
||||
from utils.utils import irtang
|
||||
|
||||
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
if "NUITKA_ONEFILE_PARENT" in os.environ:
|
||||
splash_filename = os.path.join(
|
||||
tempfile.gettempdir(),
|
||||
f"onefile_{int(os.environ['NUITKA_ONEFILE_PARENT'])}_splash_feedback.tmp",
|
||||
)
|
||||
if os.path.exists(splash_filename):
|
||||
os.unlink(splash_filename)
|
||||
|
||||
irtang()
|
||||
# Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
|
||||
setup()
|
||||
base_dir = get_base_directory()
|
||||
logger = logging.getLogger(__name__)
|
||||
dotenv.load_dotenv(dotenv_path=base_dir / ".env")
|
||||
|
||||
try:
|
||||
url = os.getenv("URL")
|
||||
username = os.getenv("USERNAME")
|
||||
|
||||
if not url:
|
||||
raise ValueError("Missing URL in environment variables.")
|
||||
if not username:
|
||||
raise ValueError("Missing USERNAME in environment variables.")
|
||||
|
||||
logger.debug(f"Retrieved URL: {url}")
|
||||
logger.debug(f"Retrieved Username: {username}")
|
||||
|
||||
except ValueError as e:
|
||||
logger.error(f"Configuration error: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
api_key = getAPI(username, "Loxide")
|
||||
if api_key is None:
|
||||
raise ValueError("API key for Loxide is missing.")
|
||||
|
||||
api = AirlockAPIWrapper(
|
||||
base_url=str(os.getenv("URL")),
|
||||
api_key=api_key,
|
||||
)
|
||||
run_Loxide(api)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 112 KiB |
@@ -0,0 +1,668 @@
|
||||
# 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 Optional
|
||||
|
||||
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,
|
||||
TextArea,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AllowlistSelectionWidget(Static):
|
||||
"""
|
||||
Widget for selecting an allowlist and adding hashes to it.
|
||||
Can be reused in different workflows.
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
AllowlistSelectionWidget {
|
||||
height: 1fr;
|
||||
layout: vertical;
|
||||
}
|
||||
#allowlist_main {
|
||||
height: 1fr;
|
||||
width: 100%;
|
||||
}
|
||||
#left_panel {
|
||||
width: 50%;
|
||||
padding: 1;
|
||||
border: solid $primary;
|
||||
}
|
||||
#right_panel {
|
||||
width: 50%;
|
||||
padding: 1;
|
||||
border: solid $primary;
|
||||
}
|
||||
#allowlist_table {
|
||||
height: 70%;
|
||||
margin: 1 0;
|
||||
}
|
||||
#allowlist_table > .datatable--header {
|
||||
text-style: bold;
|
||||
background: $boost;
|
||||
}
|
||||
#allowlist_table Row {
|
||||
height: 1;
|
||||
}
|
||||
#preview_area {
|
||||
height: 60%;
|
||||
margin: 1 0;
|
||||
}
|
||||
#action_buttons {
|
||||
height: auto;
|
||||
min-height: 3;
|
||||
padding: 1;
|
||||
content-align: center middle;
|
||||
}
|
||||
.panel-title {
|
||||
text-style: bold;
|
||||
margin: 0 0 1 0;
|
||||
}
|
||||
.info-text {
|
||||
margin: 1 0;
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
selected_data: pd.DataFrame,
|
||||
api=None,
|
||||
hostname: Optional[str] = None,
|
||||
otpid: Optional[str] = None,
|
||||
hash_column: str = "sha256", # Default hash column name
|
||||
):
|
||||
"""
|
||||
Initialize the allowlist selection widget.
|
||||
|
||||
Args:
|
||||
selected_data: DataFrame containing the selected activities
|
||||
api: API instance for making allowlist calls
|
||||
hostname: Optional hostname for context
|
||||
otpid: Optional OTP ID for context
|
||||
hash_column: Name of the column containing hashes (default: "sha256")
|
||||
"""
|
||||
super().__init__()
|
||||
self.selected_data = selected_data
|
||||
self.api = api
|
||||
self.hostname = hostname
|
||||
self.otpid = otpid
|
||||
self.hash_column = hash_column
|
||||
self.allowlists = []
|
||||
self.selected_allowlist = None
|
||||
self.hashes_to_add = []
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Horizontal(id="allowlist_main"):
|
||||
# Left panel - Allowlist selection
|
||||
with Vertical(id="left_panel"):
|
||||
yield Static("Select Allowlist", classes="panel-title")
|
||||
yield Static(
|
||||
f"Choose an allowlist to add {len(self.selected_data)} selected items",
|
||||
classes="info-text",
|
||||
)
|
||||
|
||||
# Allowlist table
|
||||
self.allowlist_table = DataTable(id="allowlist_table")
|
||||
self.allowlist_table.cursor_type = "row"
|
||||
yield self.allowlist_table
|
||||
|
||||
# Refresh button
|
||||
self.refresh_btn = Button(
|
||||
"🔄 Refresh Allowlists", id="refresh_allowlists_btn"
|
||||
)
|
||||
yield self.refresh_btn
|
||||
|
||||
# Right panel - Preview and actions
|
||||
with Vertical(id="right_panel"):
|
||||
yield Static("Preview", classes="panel-title")
|
||||
|
||||
# Context information
|
||||
context_text = []
|
||||
if self.hostname:
|
||||
context_text.append(f"Host: {self.hostname}")
|
||||
if self.otpid:
|
||||
context_text.append(f"OTP: {self.otpid}")
|
||||
context_text.append(f"Selected Activities: {len(self.selected_data)}")
|
||||
|
||||
yield Static(" | ".join(context_text), classes="info-text")
|
||||
|
||||
# Preview text area
|
||||
self.preview_area = TextArea(
|
||||
id="preview_area", read_only=True, language="markdown"
|
||||
)
|
||||
yield self.preview_area
|
||||
|
||||
# Hash statistics
|
||||
self.stats_label = Static("", id="stats_label", classes="info-text")
|
||||
yield self.stats_label
|
||||
|
||||
# Action buttons at bottom
|
||||
with Horizontal(id="action_buttons"):
|
||||
self.add_btn = Button("➕ Add to Allowlist", id="add_to_allowlist_btn")
|
||||
|
||||
self.add_btn.styles.width = "100%"
|
||||
self.add_btn.disabled = True # Disabled until allowlist selected
|
||||
|
||||
yield self.add_btn
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
"""Load allowlists when widget mounts."""
|
||||
await self.load_allowlists()
|
||||
await self.extract_and_preview_hashes()
|
||||
|
||||
async def load_allowlists(self) -> None:
|
||||
"""Load available allowlists from API, grouped by policy association."""
|
||||
if not self.api:
|
||||
logger.error("No API available")
|
||||
self.allowlist_table.add_column("Error")
|
||||
self.allowlist_table.add_row("No API available")
|
||||
return
|
||||
|
||||
try:
|
||||
# First, try to get the host's policy if hostname is provided
|
||||
host_policy_allowlists = []
|
||||
host_policy_ids = set()
|
||||
policy_name = "Unknown Policy" # Default value
|
||||
group_id = None
|
||||
|
||||
if self.hostname:
|
||||
try:
|
||||
# Get agent info to find its policy
|
||||
agents_df = self.api.agent_find_by_hostname(self.hostname)
|
||||
if not agents_df.empty:
|
||||
# Get the policy group ID for this host
|
||||
group_id = agents_df.iloc[0].get("groupid")
|
||||
|
||||
# Look up the policy name from app's cached policies
|
||||
if (
|
||||
group_id
|
||||
and hasattr(self.app, "policies")
|
||||
and self.app.policies
|
||||
):
|
||||
for policy in self.app.policies:
|
||||
if policy.groupid == group_id:
|
||||
policy_name = policy.name
|
||||
logger.info(
|
||||
f"Found policy name: '{policy_name}' for group_id: {group_id}"
|
||||
)
|
||||
break
|
||||
|
||||
logger.info(
|
||||
f"Found host '{self.hostname}' in policy '{policy_name}' (group_id: {group_id})"
|
||||
)
|
||||
|
||||
if group_id:
|
||||
# Get allowlists for this policy
|
||||
policy_allowlists_df = self.api.policy_list_allowlists(
|
||||
group_id
|
||||
)
|
||||
if not policy_allowlists_df.empty:
|
||||
host_policy_allowlists = policy_allowlists_df.to_dict(
|
||||
orient="records"
|
||||
)
|
||||
host_policy_ids = {
|
||||
al.get("applicationid")
|
||||
for al in host_policy_allowlists
|
||||
}
|
||||
logger.info(
|
||||
f"Found {len(host_policy_allowlists)} allowlists for host's policy"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not get host's policy allowlists: {e}")
|
||||
|
||||
# If we still don't have a policy name, try to get it from the first allowlist or use a default
|
||||
if not policy_name:
|
||||
# Get all policies and try to find which one has allowlists
|
||||
try:
|
||||
all_policies_df = self.api.policy_find_all()
|
||||
if not all_policies_df.empty:
|
||||
# If we have a group_id from somewhere, use it
|
||||
if group_id:
|
||||
policy_row = all_policies_df[
|
||||
all_policies_df["groupid"] == group_id
|
||||
]
|
||||
if not policy_row.empty:
|
||||
policy_name = policy_row.iloc[0].get(
|
||||
"groupname", "Unknown Policy"
|
||||
)
|
||||
else:
|
||||
# Use the first policy as fallback
|
||||
policy_name = all_policies_df.iloc[0].get(
|
||||
"groupname", "Default Policy"
|
||||
)
|
||||
logger.info(f"Using first available policy: {policy_name}")
|
||||
else:
|
||||
policy_name = "Unknown Policy"
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not fetch policies: {e}")
|
||||
policy_name = "Unknown Policy"
|
||||
|
||||
# Get all allowlists
|
||||
all_allowlists_df = self.api.allowlist_find_all()
|
||||
|
||||
if all_allowlists_df.empty:
|
||||
self.allowlist_table.add_column("No Allowlists")
|
||||
self.allowlist_table.add_row("No allowlists found")
|
||||
return
|
||||
|
||||
all_allowlists = all_allowlists_df.to_dict(orient="records")
|
||||
|
||||
# Separate into two groups: policy-associated and others
|
||||
other_allowlists = [
|
||||
al
|
||||
for al in all_allowlists
|
||||
if al.get("applicationid") not in host_policy_ids
|
||||
]
|
||||
|
||||
# Sort each group alphabetically by name
|
||||
host_policy_allowlists.sort(key=lambda x: x.get("name", "").lower())
|
||||
other_allowlists.sort(key=lambda x: x.get("name", "").lower())
|
||||
|
||||
# Combine lists with policy-associated first
|
||||
self.allowlists = host_policy_allowlists + other_allowlists
|
||||
|
||||
# Setup table columns
|
||||
self.allowlist_table.clear()
|
||||
self.allowlist_table.add_columns("Name", "Application ID", "Type")
|
||||
|
||||
# Track which rows are headers vs actual allowlists
|
||||
self._row_to_allowlist_map = {}
|
||||
current_row = 0
|
||||
|
||||
# Add policy-associated allowlists if any
|
||||
if host_policy_allowlists:
|
||||
# Add section header
|
||||
header_text = f"=== Policy: {policy_name or 'Host Policy'} ==="
|
||||
self.allowlist_table.add_row(header_text, "", "", key="header_policy")
|
||||
current_row += 1
|
||||
|
||||
# Add policy allowlists
|
||||
for idx, allowlist in enumerate(host_policy_allowlists):
|
||||
name = allowlist.get("name", "Unknown")
|
||||
app_id = allowlist.get("applicationid", "Unknown")
|
||||
|
||||
self.allowlist_table.add_row(
|
||||
f" {name}", # Indent to show grouping
|
||||
app_id,
|
||||
"Policy",
|
||||
key=f"policy_{idx}",
|
||||
)
|
||||
self._row_to_allowlist_map[current_row] = idx
|
||||
current_row += 1
|
||||
|
||||
# Add other allowlists
|
||||
if other_allowlists:
|
||||
# Add section header
|
||||
if host_policy_allowlists:
|
||||
# Add spacer if we have policy allowlists above
|
||||
self.allowlist_table.add_row("", "", "", key="spacer")
|
||||
current_row += 1
|
||||
|
||||
self.allowlist_table.add_row(
|
||||
"=== Other Available Allowlists ===", "", "", key="header_other"
|
||||
)
|
||||
current_row += 1
|
||||
|
||||
# Add other allowlists
|
||||
for idx, allowlist in enumerate(other_allowlists):
|
||||
name = allowlist.get("name", "Unknown")
|
||||
app_id = allowlist.get("applicationid", "Unknown")
|
||||
|
||||
self.allowlist_table.add_row(
|
||||
f" {name}", # Indent to show grouping
|
||||
app_id,
|
||||
"General",
|
||||
key=f"other_{idx}",
|
||||
)
|
||||
# Map to the correct index in the combined list
|
||||
actual_idx = len(host_policy_allowlists) + idx
|
||||
self._row_to_allowlist_map[current_row] = actual_idx
|
||||
current_row += 1
|
||||
|
||||
# Log summary
|
||||
logger.info(
|
||||
f"Loaded {len(self.allowlists)} total allowlists: "
|
||||
f"{len(host_policy_allowlists)} policy-associated, "
|
||||
f"{len(other_allowlists)} others"
|
||||
)
|
||||
|
||||
# Update stats label if no allowlists in policy
|
||||
if self.hostname and not host_policy_allowlists:
|
||||
self.stats_label.update(
|
||||
f"Note: No allowlists found for {self.hostname}'s policy | "
|
||||
+ self.stats_label.content.plain
|
||||
)
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception(f"Failed to load allowlists: {exc}")
|
||||
self.allowlist_table.add_column("Error")
|
||||
self.allowlist_table.add_row(f"Failed to load: {str(exc)}")
|
||||
|
||||
async def extract_and_preview_hashes(self) -> None:
|
||||
"""Extract hashes from selected data and show preview."""
|
||||
preview_lines = ["## Hash Extraction Summary\n"]
|
||||
|
||||
# Check for hash column
|
||||
if self.hash_column not in self.selected_data.columns:
|
||||
# Try to find a hash column
|
||||
possible_hash_cols = [
|
||||
"sha256",
|
||||
"SHA256",
|
||||
"hash",
|
||||
"Hash",
|
||||
"sha1",
|
||||
"SHA1",
|
||||
"md5",
|
||||
"MD5",
|
||||
"filehash",
|
||||
"file_hash",
|
||||
]
|
||||
found_col = None
|
||||
for col in possible_hash_cols:
|
||||
if col in self.selected_data.columns:
|
||||
found_col = col
|
||||
break
|
||||
|
||||
if found_col:
|
||||
self.hash_column = found_col
|
||||
preview_lines.append(f"✅ Found hash column: **{found_col}**\n")
|
||||
else:
|
||||
preview_lines.append("❌ **No hash column found**\n")
|
||||
preview_lines.append("Available columns:\n")
|
||||
for col in self.selected_data.columns:
|
||||
if col != "_row_id":
|
||||
preview_lines.append(f" - {col}\n")
|
||||
|
||||
self.preview_area.text = "".join(preview_lines)
|
||||
self.stats_label.update("No hashes to add")
|
||||
return
|
||||
|
||||
# Extract unique hashes
|
||||
hashes = self.selected_data[self.hash_column].dropna().unique()
|
||||
self.hashes_to_add = [h for h in hashes if h and str(h).strip()]
|
||||
|
||||
# Build preview
|
||||
preview_lines.append(f"### Found {len(self.hashes_to_add)} unique hashes\n\n")
|
||||
|
||||
# Show sample of hashes (first 10)
|
||||
preview_lines.append("**Sample hashes to be added:**\n```\n")
|
||||
for i, hash_val in enumerate(self.hashes_to_add[:10]):
|
||||
preview_lines.append(f"{i+1}. {hash_val}\n")
|
||||
if len(self.hashes_to_add) > 10:
|
||||
preview_lines.append(f"... and {len(self.hashes_to_add) - 10} more\n")
|
||||
preview_lines.append("```\n\n")
|
||||
|
||||
# Show sample of source data
|
||||
preview_lines.append("**Sample source activities:**\n")
|
||||
sample_cols = [
|
||||
col
|
||||
for col in self.selected_data.columns
|
||||
if col not in ["_row_id"] and col in ["filename", "path", "action", "user"]
|
||||
]
|
||||
if not sample_cols:
|
||||
sample_cols = [
|
||||
col for col in self.selected_data.columns if col != "_row_id"
|
||||
][:3]
|
||||
|
||||
if sample_cols:
|
||||
preview_lines.append("```\n")
|
||||
for i, row in self.selected_data[sample_cols].head(5).iterrows():
|
||||
row_text = " | ".join([f"{col}: {row[col]}" for col in sample_cols])
|
||||
preview_lines.append(f"{row_text}\n")
|
||||
preview_lines.append("```\n")
|
||||
|
||||
self.preview_area.text = "".join(preview_lines)
|
||||
|
||||
# Update statistics
|
||||
self.stats_label.update(
|
||||
f"Ready to add {len(self.hashes_to_add)} unique hashes | "
|
||||
f"From {len(self.selected_data)} selected activities"
|
||||
)
|
||||
|
||||
async def on_data_table_row_selected(self, event) -> None:
|
||||
"""Handle allowlist selection."""
|
||||
try:
|
||||
# Extract row index from event - handle different event structures
|
||||
row_index = None
|
||||
|
||||
# Try to get row index from coordinate
|
||||
if hasattr(event, "coordinate") and hasattr(event.coordinate, "row"):
|
||||
row_index = event.coordinate.row
|
||||
# Try cursor_row as fallback
|
||||
elif hasattr(event, "cursor_row"):
|
||||
row_index = event.cursor_row
|
||||
# Try getting from the table itself
|
||||
else:
|
||||
table = self.allowlist_table
|
||||
if hasattr(table, "cursor_row"):
|
||||
row_index = table.cursor_row
|
||||
|
||||
# Validate row index
|
||||
if row_index is not None and isinstance(row_index, int):
|
||||
# Account for group headers in the row count
|
||||
actual_allowlist_index = self._get_allowlist_index_from_row(row_index)
|
||||
|
||||
if (
|
||||
actual_allowlist_index is not None
|
||||
and 0 <= actual_allowlist_index < len(self.allowlists)
|
||||
):
|
||||
self.selected_allowlist = self.allowlists[actual_allowlist_index]
|
||||
self.add_btn.disabled = False
|
||||
self.add_btn.label = (
|
||||
f"➕ Add to '{self.selected_allowlist.get('name', 'Unknown')}'"
|
||||
)
|
||||
|
||||
# Update preview with selection
|
||||
await self._update_preview_with_selection()
|
||||
|
||||
logger.info(
|
||||
f"Selected allowlist: {self.selected_allowlist.get('name')}"
|
||||
)
|
||||
else:
|
||||
logger.debug(f"Row {row_index} is a header or invalid")
|
||||
else:
|
||||
logger.warning(f"Could not extract valid row index from event: {event}")
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception(f"Failed to select allowlist: {exc}")
|
||||
|
||||
def _get_allowlist_index_from_row(self, row_index: int) -> Optional[int]:
|
||||
"""Convert table row index to allowlist list index, accounting for group headers."""
|
||||
# This will be updated when we have group headers
|
||||
if hasattr(self, "_row_to_allowlist_map"):
|
||||
return self._row_to_allowlist_map.get(row_index)
|
||||
return row_index
|
||||
|
||||
async def _update_preview_with_selection(self) -> None:
|
||||
"""Update preview when an allowlist is selected."""
|
||||
if not self.selected_allowlist:
|
||||
return
|
||||
|
||||
current_text = self.preview_area.text
|
||||
# Remove any existing selection header
|
||||
if "### Selected Allowlist:" in current_text:
|
||||
lines = current_text.split("\n")
|
||||
# Find and remove the selection lines
|
||||
new_lines = []
|
||||
skip_next = False
|
||||
for line in lines:
|
||||
if line.startswith("### Selected Allowlist:"):
|
||||
skip_next = True
|
||||
continue
|
||||
if skip_next and line.startswith("Application ID:"):
|
||||
skip_next = False
|
||||
continue
|
||||
if not skip_next:
|
||||
new_lines.append(line)
|
||||
current_text = "\n".join(new_lines)
|
||||
|
||||
# Add new selection at the top
|
||||
selection_text = (
|
||||
f"### Selected Allowlist: **{self.selected_allowlist.get('name')}**\n"
|
||||
f"Application ID: {self.selected_allowlist.get('applicationid')}\n\n"
|
||||
)
|
||||
self.preview_area.text = selection_text + current_text
|
||||
|
||||
async def on_button_pressed(self, event) -> None:
|
||||
"""Handle button presses."""
|
||||
btn = getattr(event, "button", None) or getattr(event, "sender", None)
|
||||
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None)
|
||||
|
||||
if btn is self.refresh_btn or btn_id == "refresh_allowlists_btn":
|
||||
await self.load_allowlists()
|
||||
event.stop()
|
||||
return
|
||||
|
||||
if btn is self.add_btn or btn_id == "add_to_allowlist_btn":
|
||||
await self.add_hashes_to_allowlist()
|
||||
event.stop()
|
||||
return
|
||||
|
||||
async def add_hashes_to_allowlist(self) -> None:
|
||||
"""Add the extracted hashes to the selected allowlist."""
|
||||
if not self.selected_allowlist or not self.hashes_to_add:
|
||||
self.app.notify(
|
||||
"No allowlist selected or no hashes to add", severity="warning"
|
||||
)
|
||||
return
|
||||
|
||||
if not self.api:
|
||||
self.app.notify("API not available", severity="error")
|
||||
return
|
||||
|
||||
try:
|
||||
# Disable button during operation
|
||||
self.add_btn.disabled = True
|
||||
self.add_btn.label = "Adding hashes..."
|
||||
|
||||
# Call API to add hashes
|
||||
app_id = self.selected_allowlist.get("applicationid")
|
||||
allowlist_name = self.selected_allowlist.get("name", "Unknown")
|
||||
|
||||
logger.info(
|
||||
f"Adding {len(self.hashes_to_add)} hashes to allowlist {allowlist_name} (ID: {app_id})"
|
||||
)
|
||||
|
||||
result = self.api.hash_add_to_allowlist(app_id, self.hashes_to_add)
|
||||
logger.debug(f"Hash adding api call: {result}")
|
||||
# Success notification
|
||||
self.app.notify(
|
||||
f"Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'",
|
||||
title="Success",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
# Update preview to show success
|
||||
self.preview_area.text = (
|
||||
f"## SUCCESS\n\n"
|
||||
f"Added **{len(self.hashes_to_add)} hashes** to allowlist:\n"
|
||||
f"**{allowlist_name}** (ID: {app_id})\n\n"
|
||||
f"### Operation Details:\n"
|
||||
f"- Source: {self.hostname or 'Multiple hosts'}\n"
|
||||
f"- OTP ID: {self.otpid or 'N/A'}\n"
|
||||
f"- Activities processed: {len(self.selected_data)}\n"
|
||||
f"- Unique hashes added: {len(self.hashes_to_add)}\n"
|
||||
)
|
||||
|
||||
# Change button to "Done"
|
||||
self.add_btn.label = "Done - Press q to return to main menu"
|
||||
self.add_btn.disabled = True
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception(f"Failed to add hashes to allowlist: {exc}")
|
||||
self.app.notify(
|
||||
f"Failed to add hashes: {str(exc)}",
|
||||
title="Error",
|
||||
severity="error",
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
# Re-enable button
|
||||
self.add_btn.disabled = False
|
||||
self.add_btn.label = "Retry Add to Allowlist"
|
||||
|
||||
|
||||
class AllowlistSelectionScreen(Screen):
|
||||
"""
|
||||
Screen wrapper for the AllowlistSelectionWidget.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("r", "refresh", "Refresh Allowlists"),
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
selected_data: pd.DataFrame,
|
||||
api=None,
|
||||
hostname: Optional[str] = None,
|
||||
otpid: Optional[str] = None,
|
||||
hash_column: str = "sha256",
|
||||
):
|
||||
super().__init__()
|
||||
self.selected_data = selected_data
|
||||
self.api = api
|
||||
self.hostname = hostname
|
||||
self.otpid = otpid
|
||||
self.hash_column = hash_column
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header(show_clock=True)
|
||||
self.widget = AllowlistSelectionWidget(
|
||||
self.selected_data,
|
||||
api=self.api,
|
||||
hostname=self.hostname,
|
||||
otpid=self.otpid,
|
||||
hash_column=self.hash_column,
|
||||
)
|
||||
yield self.widget
|
||||
yield Footer()
|
||||
|
||||
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()
|
||||
|
||||
async def action_refresh(self) -> None:
|
||||
"""Refresh the allowlists."""
|
||||
if hasattr(self, "widget") and self.widget:
|
||||
await self.widget.load_allowlists()
|
||||
|
||||
async def action_confirm(self) -> None:
|
||||
"""Confirm and add to allowlist."""
|
||||
if hasattr(self, "widget") and self.widget:
|
||||
if self.widget.selected_allowlist and self.widget.hashes_to_add:
|
||||
await self.widget.add_hashes_to_allowlist()
|
||||
@@ -0,0 +1,639 @@
|
||||
# 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 datetime import datetime, timedelta
|
||||
import logging
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
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,
|
||||
Label,
|
||||
Select,
|
||||
Static,
|
||||
)
|
||||
|
||||
from models.agent import Agent
|
||||
from models.execution import ExecutionHistoryRecord
|
||||
from utils.configmanager import load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ExecutionHistoryScreen(Screen):
|
||||
"""
|
||||
A screen for viewing and exporting execution history for selected agents.
|
||||
|
||||
This screen allows users to:
|
||||
1. Select a start date and end date using dropdown selects
|
||||
2. Fetch execution history for all selected agents
|
||||
3. View the results in a DataTable
|
||||
4. Export the results to CSV using a keybinding
|
||||
|
||||
Attributes:
|
||||
agents (List[Agent]): List of agents to fetch execution history for
|
||||
execution_data (pd.DataFrame): Combined execution history data
|
||||
working_dir (str): Directory for CSV exports
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
ExecutionHistoryScreen {
|
||||
align: center top;
|
||||
}
|
||||
|
||||
#main_container {
|
||||
width: 95%;
|
||||
height: 1fr;
|
||||
border: solid $primary;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
#title {
|
||||
text-style: bold;
|
||||
color: $text;
|
||||
text-align: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#date_container {
|
||||
height: auto;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#start_date_row, #end_date_row {
|
||||
height: auto;
|
||||
align-horizontal: left;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
.date_label {
|
||||
width: 8;
|
||||
margin-right: 1;
|
||||
}
|
||||
|
||||
.date_selector {
|
||||
width: 18;
|
||||
margin: 0 1;
|
||||
}
|
||||
|
||||
#quick_buttons_row {
|
||||
height: auto;
|
||||
align-horizontal: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
.quick_select_btn {
|
||||
margin: 0 1;
|
||||
}
|
||||
|
||||
#button_row {
|
||||
height: auto;
|
||||
align-horizontal: center;
|
||||
margin-top: 1;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
Button {
|
||||
margin: 0 1;
|
||||
}
|
||||
|
||||
#status_label {
|
||||
text-align: center;
|
||||
color: $accent;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#results_container {
|
||||
height: 1fr;
|
||||
display: none;
|
||||
}
|
||||
|
||||
#results_button_row {
|
||||
height: auto;
|
||||
align-horizontal: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#history_table {
|
||||
height: 1fr;
|
||||
border: solid $primary;
|
||||
}
|
||||
|
||||
DataTable > .datatable--header {
|
||||
text-style: bold;
|
||||
background: $primary 20%;
|
||||
}
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "close_screen", "Close"),
|
||||
Binding("e", "export_csv", "Export CSV"),
|
||||
Binding("q", "close_screen", "Quit"),
|
||||
]
|
||||
|
||||
def __init__(self, agents: List[Agent]):
|
||||
"""
|
||||
Initialize the ExecutionHistoryScreen.
|
||||
|
||||
Args:
|
||||
agents (List[Agent]): List of agents to fetch execution history for
|
||||
"""
|
||||
super().__init__()
|
||||
self.agents = agents
|
||||
self.execution_data = pd.DataFrame()
|
||||
self.working_dir = load_env("WORKING_DIR") or os.getcwd()
|
||||
|
||||
# Generate dropdown options
|
||||
today = datetime.now().date()
|
||||
|
||||
# Month options - format is (display_text, value)
|
||||
self.month_options = [
|
||||
("January", "01"),
|
||||
("February", "02"),
|
||||
("March", "03"),
|
||||
("April", "04"),
|
||||
("May", "05"),
|
||||
("June", "06"),
|
||||
("July", "07"),
|
||||
("August", "08"),
|
||||
("September", "09"),
|
||||
("October", "10"),
|
||||
("November", "11"),
|
||||
("December", "12"),
|
||||
]
|
||||
|
||||
# Day options (1-31) - format is (display_text, value)
|
||||
self.day_options = [(f"{i}", f"{i:02d}") for i in range(1, 32)]
|
||||
|
||||
# Year options (current year back 5 years) - format is (display_text, value)
|
||||
current_year = today.year
|
||||
self.year_options = [
|
||||
(str(year), str(year)) for year in range(current_year, current_year - 6, -1)
|
||||
]
|
||||
|
||||
# Default dates: last 30 days
|
||||
start_date = today - timedelta(days=30)
|
||||
self.start_month = f"{start_date.month:02d}"
|
||||
self.start_day = f"{start_date.day:02d}"
|
||||
self.start_year = str(start_date.year)
|
||||
|
||||
self.end_month = f"{today.month:02d}"
|
||||
self.end_day = f"{today.day:02d}"
|
||||
self.end_year = str(today.year)
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Build the UI layout."""
|
||||
yield Header(show_clock=True, icon="📊")
|
||||
|
||||
with Vertical(id="main_container"):
|
||||
title_text = f"Execution History - {len(self.agents)} Agent(s)"
|
||||
yield Static(title_text, id="title")
|
||||
|
||||
# Date selection area
|
||||
with Vertical(id="date_container"):
|
||||
yield Label("Select Date Range:")
|
||||
|
||||
# Start date row
|
||||
with Horizontal(id="start_date_row"):
|
||||
yield Label("From:", classes="date_label")
|
||||
yield Select(
|
||||
options=self.month_options,
|
||||
value=self.start_month,
|
||||
id="start_month_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
yield Select(
|
||||
options=self.day_options,
|
||||
value=self.start_day,
|
||||
id="start_day_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
yield Select(
|
||||
options=self.year_options,
|
||||
value=self.start_year,
|
||||
id="start_year_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
|
||||
# End date row
|
||||
with Horizontal(id="end_date_row"):
|
||||
yield Label("To:", classes="date_label")
|
||||
yield Select(
|
||||
options=self.month_options,
|
||||
value=self.end_month,
|
||||
id="end_month_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
yield Select(
|
||||
options=self.day_options,
|
||||
value=self.end_day,
|
||||
id="end_day_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
yield Select(
|
||||
options=self.year_options,
|
||||
value=self.end_year,
|
||||
id="end_year_select",
|
||||
classes="date_selector",
|
||||
)
|
||||
|
||||
# Quick select buttons
|
||||
with Horizontal(id="quick_buttons_row"):
|
||||
yield Button(
|
||||
"1 Day",
|
||||
id="quick_1day",
|
||||
classes="quick_select_btn",
|
||||
variant="default",
|
||||
)
|
||||
yield Button(
|
||||
"1 Week",
|
||||
id="quick_1week",
|
||||
classes="quick_select_btn",
|
||||
variant="default",
|
||||
)
|
||||
yield Button(
|
||||
"30 Days",
|
||||
id="quick_30days",
|
||||
classes="quick_select_btn",
|
||||
variant="default",
|
||||
)
|
||||
|
||||
# Buttons
|
||||
with Horizontal(id="button_row"):
|
||||
yield Button("Fetch History", id="fetch_btn", variant="primary")
|
||||
yield Button("Close", id="close_btn", variant="error")
|
||||
|
||||
# Status
|
||||
yield Static(
|
||||
"Select date range and click 'Fetch History'", id="status_label"
|
||||
)
|
||||
|
||||
# Results container (hidden initially, shown after fetch)
|
||||
with Vertical(id="results_container"):
|
||||
with Horizontal(id="results_button_row"):
|
||||
yield Button("Export CSV", id="export_btn", variant="success")
|
||||
yield Button("Back", id="back_btn", variant="default")
|
||||
yield DataTable(id="history_table")
|
||||
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Initialize the table when screen is mounted."""
|
||||
table = self.query_one("#history_table", DataTable)
|
||||
table.cursor_type = "row"
|
||||
table.zebra_stripes = True
|
||||
|
||||
# Initially empty - will populate after fetch
|
||||
logger.info(f"ExecutionHistoryScreen mounted with {len(self.agents)} agents")
|
||||
|
||||
def on_select_changed(self, event: Select.Changed) -> None:
|
||||
"""Handle date selection changes."""
|
||||
select_id = event.select.id
|
||||
|
||||
if select_id == "start_month_select":
|
||||
self.start_month = event.value
|
||||
logger.debug(f"Start month changed to: {self.start_month}")
|
||||
elif select_id == "start_day_select":
|
||||
self.start_day = event.value
|
||||
logger.debug(f"Start day changed to: {self.start_day}")
|
||||
elif select_id == "start_year_select":
|
||||
self.start_year = event.value
|
||||
logger.debug(f"Start year changed to: {self.start_year}")
|
||||
elif select_id == "end_month_select":
|
||||
self.end_month = event.value
|
||||
logger.debug(f"End month changed to: {self.end_month}")
|
||||
elif select_id == "end_day_select":
|
||||
self.end_day = event.value
|
||||
logger.debug(f"End day changed to: {self.end_day}")
|
||||
elif select_id == "end_year_select":
|
||||
self.end_year = event.value
|
||||
logger.debug(f"End year changed to: {self.end_year}")
|
||||
|
||||
def _set_quick_date_range(self, days: int) -> None:
|
||||
"""Set the date range based on quick select button."""
|
||||
today = datetime.now().date()
|
||||
start_date = today - timedelta(days=days)
|
||||
|
||||
# Update internal values
|
||||
self.start_month = f"{start_date.month:02d}"
|
||||
self.start_day = f"{start_date.day:02d}"
|
||||
self.start_year = str(start_date.year)
|
||||
|
||||
self.end_month = f"{today.month:02d}"
|
||||
self.end_day = f"{today.day:02d}"
|
||||
self.end_year = str(today.year)
|
||||
|
||||
# Update the Select widgets
|
||||
try:
|
||||
self.query_one("#start_month_select", Select).value = self.start_month
|
||||
self.query_one("#start_day_select", Select).value = self.start_day
|
||||
self.query_one("#start_year_select", Select).value = self.start_year
|
||||
|
||||
self.query_one("#end_month_select", Select).value = self.end_month
|
||||
self.query_one("#end_day_select", Select).value = self.end_day
|
||||
self.query_one("#end_year_select", Select).value = self.end_year
|
||||
|
||||
self.app.notify(
|
||||
f"Date range set to last {days} day(s)",
|
||||
severity="information",
|
||||
timeout=2,
|
||||
)
|
||||
logger.info(f"Quick select: Set date range to last {days} days")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update date selects: {e}")
|
||||
|
||||
def _show_date_selection(self) -> None:
|
||||
"""Show the date selection view and hide results."""
|
||||
try:
|
||||
self.query_one("#date_container").styles.display = "block"
|
||||
self.query_one("#button_row").styles.display = "block"
|
||||
self.query_one("#status_label").styles.display = "block"
|
||||
self.query_one("#results_container").styles.display = "none"
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to show date selection: {e}")
|
||||
|
||||
def _show_results(self) -> None:
|
||||
"""Hide date selection view and show results."""
|
||||
try:
|
||||
self.query_one("#date_container").styles.display = "none"
|
||||
self.query_one("#button_row").styles.display = "none"
|
||||
self.query_one("#status_label").styles.display = "none"
|
||||
self.query_one("#results_container").styles.display = "block"
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to show results: {e}")
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button clicks."""
|
||||
if event.button.id == "fetch_btn":
|
||||
self._fetch_execution_history()
|
||||
elif event.button.id == "export_btn":
|
||||
self._export_to_csv()
|
||||
elif event.button.id == "close_btn":
|
||||
self.app.pop_screen()
|
||||
elif event.button.id == "back_btn":
|
||||
self._show_date_selection()
|
||||
elif event.button.id == "quick_1day":
|
||||
self._set_quick_date_range(days=1)
|
||||
elif event.button.id == "quick_1week":
|
||||
self._set_quick_date_range(days=7)
|
||||
elif event.button.id == "quick_30days":
|
||||
self._set_quick_date_range(days=30)
|
||||
|
||||
def _fetch_execution_history(self) -> None:
|
||||
"""Fetch execution history for all selected agents."""
|
||||
status_label = self.query_one("#status_label", Static)
|
||||
status_label.update("â³ Fetching execution history...")
|
||||
|
||||
# Disable buttons during fetch
|
||||
fetch_btn = self.query_one("#fetch_btn", Button)
|
||||
export_btn = self.query_one("#export_btn", Button)
|
||||
fetch_btn.disabled = True
|
||||
export_btn.disabled = True
|
||||
|
||||
api = self.app.api
|
||||
all_history = []
|
||||
|
||||
try:
|
||||
# Construct dates from dropdowns
|
||||
start_date_str = f"{self.start_year}-{self.start_month}-{self.start_day}"
|
||||
end_date_str = f"{self.end_year}-{self.end_month}-{self.end_day}"
|
||||
|
||||
# Validate dates
|
||||
try:
|
||||
start_dt = datetime.strptime(start_date_str, "%Y-%m-%d")
|
||||
end_dt = datetime.strptime(end_date_str, "%Y-%m-%d")
|
||||
except ValueError as e:
|
||||
status_label.update(f"⌠Invalid date: {str(e)}")
|
||||
fetch_btn.disabled = False
|
||||
export_btn.disabled = False
|
||||
self.app.notify(f"Invalid date selected: {str(e)}", severity="error")
|
||||
return
|
||||
|
||||
if start_dt > end_dt:
|
||||
status_label.update("⌠Error: Start date must be before end date")
|
||||
fetch_btn.disabled = False
|
||||
export_btn.disabled = False
|
||||
return
|
||||
|
||||
# Fetch history for each agent
|
||||
for i, agent in enumerate(self.agents):
|
||||
try:
|
||||
status_label.update(
|
||||
f"â³ Fetching history for {agent.hostname} ({i+1}/{len(self.agents)})..."
|
||||
)
|
||||
|
||||
# Call API - note the API expects 'dateto' first, then 'datefrom'
|
||||
history = api.history_execution(
|
||||
today=end_date_str,
|
||||
date_selected=start_date_str,
|
||||
agent_name=agent.hostname,
|
||||
)
|
||||
|
||||
if history:
|
||||
# Add agent hostname to each record for identification
|
||||
for record in history:
|
||||
record["agent_hostname"] = agent.hostname
|
||||
all_history.extend(history)
|
||||
logger.info(
|
||||
f"Fetched {len(history)} records for {agent.hostname}"
|
||||
)
|
||||
else:
|
||||
logger.info(f"No history found for {agent.hostname}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch history for {agent.hostname}: {e}")
|
||||
self.app.notify(
|
||||
f"Warning: Failed to fetch history for {agent.hostname}",
|
||||
severity="warning",
|
||||
)
|
||||
|
||||
# Convert to DataFrame
|
||||
if all_history:
|
||||
status_label.update(
|
||||
"â³ Enriching execution data with hash information..."
|
||||
)
|
||||
|
||||
# Normalize field names (handle API typos)
|
||||
for record in all_history:
|
||||
if "policver" in record and "policyver" not in record:
|
||||
record["policyver"] = record.pop("policver")
|
||||
|
||||
# Convert dict records to ExecutionHistoryRecord objects
|
||||
execution_records = []
|
||||
for record in all_history:
|
||||
try:
|
||||
execution_records.append(ExecutionHistoryRecord(**record))
|
||||
except TypeError as e:
|
||||
logger.warning(f"Failed to create ExecutionHistoryRecord: {e}")
|
||||
# If it fails, just keep the dict
|
||||
continue
|
||||
|
||||
# Enrich with hash data if we have ExecutionHistoryRecord objects
|
||||
if execution_records:
|
||||
try:
|
||||
enriched_records = ExecutionHistoryRecord.enrich_with_hashes(
|
||||
api, execution_records
|
||||
)
|
||||
logger.info(
|
||||
f"Enriched {len(enriched_records)} records with hash data"
|
||||
)
|
||||
|
||||
# Convert back to DataFrame
|
||||
self.execution_data = pd.DataFrame(
|
||||
[r.__dict__ for r in enriched_records]
|
||||
)
|
||||
|
||||
# Flatten hash_obj if present
|
||||
if (
|
||||
not self.execution_data.empty
|
||||
and "hash_obj" in self.execution_data.columns
|
||||
):
|
||||
hash_df = self.execution_data["hash_obj"].apply(
|
||||
lambda h: (
|
||||
h.to_dict() if h and hasattr(h, "to_dict") else {}
|
||||
)
|
||||
)
|
||||
self.execution_data = pd.concat(
|
||||
[
|
||||
self.execution_data.drop(columns=["hash_obj"]),
|
||||
hash_df,
|
||||
],
|
||||
axis=1,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to enrich with hashes: {e}")
|
||||
# Fall back to plain DataFrame
|
||||
self.execution_data = pd.DataFrame(all_history)
|
||||
else:
|
||||
# If we couldn't create any ExecutionHistoryRecord objects, just use raw data
|
||||
self.execution_data = pd.DataFrame(all_history)
|
||||
|
||||
self._populate_table()
|
||||
self._show_results() # Switch to results view
|
||||
self.app.notify(
|
||||
f"Successfully loaded {len(self.execution_data)} records",
|
||||
severity="information",
|
||||
)
|
||||
else:
|
||||
status_label.update(
|
||||
"â„¹ï¸ No execution history found for selected agents/dates"
|
||||
)
|
||||
self.app.notify("No execution history found", severity="information")
|
||||
self.execution_data = pd.DataFrame()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error fetching execution history: {e}")
|
||||
status_label.update(f"⌠Error: {str(e)}")
|
||||
self.app.notify(f"Failed to fetch history: {str(e)}", severity="error")
|
||||
|
||||
finally:
|
||||
# Re-enable buttons
|
||||
fetch_btn.disabled = False
|
||||
export_btn.disabled = False
|
||||
|
||||
def _populate_table(self) -> None:
|
||||
"""Populate the DataTable with execution history data."""
|
||||
table = self.query_one("#history_table", DataTable)
|
||||
table.clear(columns=True)
|
||||
|
||||
if self.execution_data.empty:
|
||||
return
|
||||
|
||||
# Define preferred column order (your specified order)
|
||||
preferred_order = [
|
||||
"policyname",
|
||||
"policyver",
|
||||
"hostname",
|
||||
"username",
|
||||
"publisher",
|
||||
"filename",
|
||||
"pprocess",
|
||||
"gprocess",
|
||||
"sha256",
|
||||
"commandline",
|
||||
"agent_hostname", # Our custom field
|
||||
]
|
||||
|
||||
# Get available columns in preferred order, then add any remaining columns
|
||||
available_cols = []
|
||||
for col in preferred_order:
|
||||
if col in self.execution_data.columns:
|
||||
available_cols.append(col)
|
||||
|
||||
# Add any remaining columns not in preferred order
|
||||
for col in self.execution_data.columns:
|
||||
if col not in available_cols:
|
||||
available_cols.append(col)
|
||||
|
||||
# Add columns to table
|
||||
for col in available_cols:
|
||||
table.add_column(col, key=col)
|
||||
|
||||
# Add rows
|
||||
for idx, row in self.execution_data.iterrows():
|
||||
row_data = []
|
||||
for col in available_cols:
|
||||
value = row[col]
|
||||
# Convert to string, handle None/NaN
|
||||
if pd.isna(value):
|
||||
row_data.append("")
|
||||
else:
|
||||
row_data.append(str(value))
|
||||
table.add_row(*row_data, key=str(idx))
|
||||
|
||||
logger.info(f"Populated table with {len(self.execution_data)} rows")
|
||||
|
||||
def _export_to_csv(self) -> None:
|
||||
"""Export the current execution data to CSV."""
|
||||
if self.execution_data.empty:
|
||||
self.app.notify("No data to export", severity="warning")
|
||||
return
|
||||
|
||||
try:
|
||||
# Create filename with timestamp
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"execution_history_{timestamp}.csv"
|
||||
filepath = os.path.join(self.working_dir, filename)
|
||||
|
||||
# Export to CSV
|
||||
self.execution_data.to_csv(filepath, index=False, encoding="utf-8-sig")
|
||||
|
||||
self.app.notify(
|
||||
f"✅ Exported {len(self.execution_data)} records to: {filepath}",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
logger.info(f"Exported execution history to: {filepath}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to export CSV: {e}")
|
||||
self.app.notify(f"Failed to export CSV: {str(e)}", severity="error")
|
||||
|
||||
def action_export_csv(self) -> None:
|
||||
"""Keybinding action to export CSV."""
|
||||
self._export_to_csv()
|
||||
|
||||
def action_close_screen(self) -> None:
|
||||
"""Close this screen and return to previous."""
|
||||
self.app.pop_screen()
|
||||
@@ -0,0 +1,114 @@
|
||||
# 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 typing import List, Optional
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.css.query import NoMatches
|
||||
from textual.screen import Screen
|
||||
|
||||
from models.agent import Agent
|
||||
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
|
||||
from TUI.Widgets.multiagentselector import MultiAgentSelector
|
||||
from TUI.Widgets.resultsdisplay import ResultsDisplay
|
||||
|
||||
|
||||
class MoveAgentWorkflowScreen(Screen):
|
||||
"""Screen that handles the agent movement workflow."""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
def __init__(self, all_agents: Optional[List[Agent]]):
|
||||
super().__init__()
|
||||
self.all_agents = all_agents
|
||||
self.selected_agents = None
|
||||
self.workflow_stage = "select_agents" # Track current stage
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Start with the multi-agent selector."""
|
||||
yield MultiAgentSelector(self.all_agents)
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle escape key to go back one step within the workflow."""
|
||||
if self.workflow_stage == "select_agents":
|
||||
# At first stage, go back to main menu
|
||||
self.app.pop_screen()
|
||||
elif self.workflow_stage == "operations":
|
||||
# Go back to agent selection
|
||||
try:
|
||||
ops_widget = self.query_one(AgentMoveOperations)
|
||||
ops_widget.remove()
|
||||
except NoMatches:
|
||||
pass
|
||||
self.mount(MultiAgentSelector(self.all_agents))
|
||||
self.workflow_stage = "select_agents"
|
||||
elif self.workflow_stage == "results":
|
||||
# Go back to operations
|
||||
try:
|
||||
results_widget = self.query_one(ResultsDisplay)
|
||||
results_widget.remove()
|
||||
except NoMatches:
|
||||
pass
|
||||
self.mount(AgentMoveOperations(self.selected_agents))
|
||||
self.workflow_stage = "operations"
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Handle q key to go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_multi_agent_selector_agents_selected(
|
||||
self, message: MultiAgentSelector.AgentsSelected
|
||||
) -> None:
|
||||
"""Handle selected agents - switch to operations screen."""
|
||||
self.selected_agents = message.selected_agents
|
||||
|
||||
# Remove the MultiAgentSelector
|
||||
selector = self.query_one(MultiAgentSelector)
|
||||
selector.remove()
|
||||
|
||||
# Mount the AgentMoveOperations with the selected Agent objects
|
||||
self.mount(AgentMoveOperations(self.selected_agents))
|
||||
self.workflow_stage = "operations"
|
||||
|
||||
def on_agent_move_operations_operation_complete(
|
||||
self, message: AgentMoveOperations.OperationComplete
|
||||
) -> None:
|
||||
"""Handle completion of move operation - transition to results screen."""
|
||||
# Format successful results
|
||||
success_lines = []
|
||||
for agent, result in message.successful:
|
||||
success_lines.append(f"✔ {agent.hostname}")
|
||||
|
||||
# Format unsuccessful results
|
||||
failure_lines = []
|
||||
for agent, error in message.unsuccessful:
|
||||
failure_lines.append(f"❌ — {agent.hostname}: {error}")
|
||||
|
||||
successful_text = "\n".join(success_lines) if success_lines else "(none)"
|
||||
unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)"
|
||||
|
||||
# Remove the operations widget
|
||||
ops_widget = self.query_one(AgentMoveOperations)
|
||||
ops_widget.remove()
|
||||
|
||||
# Mount the results display
|
||||
self.mount(
|
||||
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
|
||||
)
|
||||
self.workflow_stage = "results"
|
||||
@@ -0,0 +1,901 @@
|
||||
# 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
|
||||
|
||||
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 TUI.Screens.allowlistselectionscreen import AllowlistSelectionScreen
|
||||
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 Continue button.
|
||||
"""
|
||||
|
||||
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: 1fr;
|
||||
border: none;
|
||||
padding: 1 1;
|
||||
}
|
||||
#activity_buttons {
|
||||
height: auto;
|
||||
min-height: 3;
|
||||
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
|
||||
# Button area at the bottom (Continue)
|
||||
with Horizontal(id="activity_buttons"):
|
||||
self.continue_btn = Button("Continue", id="activity_continue_btn")
|
||||
self.continue_btn.styles.width = "100%"
|
||||
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 Continue button 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)
|
||||
)
|
||||
# ---- 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.")
|
||||
self.app.notify(
|
||||
"No activities loaded to continue with.", severity="warning"
|
||||
)
|
||||
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 button.
|
||||
Now includes Select All/None and Add to Allowlist functionality.
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
ActivityDetailWidget {
|
||||
height: 1fr;
|
||||
layout: vertical;
|
||||
}
|
||||
#detail_table_container {
|
||||
height: 1fr;
|
||||
padding: 1 1;
|
||||
}
|
||||
#selection_buttons {
|
||||
height: auto;
|
||||
min-height: 3;
|
||||
padding: 1 1;
|
||||
content-align: center middle;
|
||||
}
|
||||
#detail_buttons {
|
||||
height: auto;
|
||||
min-height: 3;
|
||||
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)
|
||||
)
|
||||
# Add a unique identifier column if not present
|
||||
if "_row_id" not in self.activities_df.columns:
|
||||
self.activities_df["_row_id"] = range(len(self.activities_df))
|
||||
|
||||
self.otpid = otpid
|
||||
self.hostname = hostname
|
||||
self.selected_row_ids = set() # Track selected rows by unique ID
|
||||
self.row_key_to_id = {} # Map DataTable row keys to unique row IDs
|
||||
self.table_row_to_id = {} # Map table row indices to unique row IDs
|
||||
self._last_sort = None # Track last sort column and order
|
||||
|
||||
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
|
||||
|
||||
# Original buttons at bottom
|
||||
with Horizontal(id="detail_buttons"):
|
||||
self.add_allowlist_btn = Button(
|
||||
"Add Selected to Allowlist", id="add_allowlist_btn"
|
||||
)
|
||||
yield self.add_allowlist_btn
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
await self._build_table(rebuild=True)
|
||||
self._update_button_states()
|
||||
|
||||
def _update_button_states(self) -> None:
|
||||
"""Update button states based on selection."""
|
||||
has_selection = len(self.selected_row_ids) > 0
|
||||
self.add_allowlist_btn.disabled = not has_selection
|
||||
|
||||
# Update button labels with count
|
||||
count = len(self.selected_row_ids)
|
||||
len(self.activities_df)
|
||||
|
||||
if has_selection:
|
||||
self.add_allowlist_btn.label = f"Add {count} Selected to Allowlist"
|
||||
else:
|
||||
self.add_allowlist_btn.label = "Add Selected to Allowlist"
|
||||
|
||||
async def _build_table(self, rebuild: bool = True) -> None:
|
||||
"""Rebuild the DataTable. If rebuild=False, only refresh rows."""
|
||||
if rebuild:
|
||||
# Full rebuild: clear columns and rows
|
||||
self.detail_table.clear()
|
||||
self.detail_table.columns.clear()
|
||||
self.row_key_to_id.clear()
|
||||
self.table_row_to_id.clear()
|
||||
|
||||
if self.activities_df is None or self.activities_df.empty:
|
||||
logger.info("ActivityDetailWidget mounted with empty dataframe.")
|
||||
return
|
||||
|
||||
# Add columns (checkbox + data columns, excluding internal _row_id)
|
||||
self.detail_table.add_column("Select", key="select")
|
||||
for col in self.activities_df.columns:
|
||||
if col != "_row_id": # Don't display the internal ID column
|
||||
self.detail_table.add_column(col)
|
||||
else:
|
||||
# Partial rebuild: clear rows only
|
||||
self.detail_table.clear()
|
||||
self.row_key_to_id.clear()
|
||||
self.table_row_to_id.clear()
|
||||
|
||||
# Add rows
|
||||
for table_idx, (df_idx, row) in enumerate(self.activities_df.iterrows()):
|
||||
# Get the unique row ID
|
||||
row_id = row["_row_id"]
|
||||
|
||||
# Build values list (excluding _row_id column)
|
||||
vals = []
|
||||
for col in self.activities_df.columns:
|
||||
if col != "_row_id":
|
||||
v = row[col]
|
||||
vals.append("" if pd.isna(v) else str(v))
|
||||
|
||||
# Check if this row is selected
|
||||
checkbox = "☑️" if row_id in self.selected_row_ids else "☐"
|
||||
|
||||
# Add row to table
|
||||
row_key = self.detail_table.add_row(checkbox, *vals)
|
||||
|
||||
# Map the row key and table index to the unique row ID
|
||||
self.row_key_to_id[row_key] = row_id
|
||||
self.table_row_to_id[table_idx] = row_id
|
||||
|
||||
async def on_data_table_cell_selected(self, event: DataTable.CellSelected) -> None:
|
||||
# Toggle selection when the "Select" column is clicked
|
||||
if event.cell_key.column_key.value == "select":
|
||||
table_row_index = event.coordinate.row
|
||||
|
||||
# Get the unique row ID for this table row
|
||||
row_id = self.table_row_to_id.get(table_row_index)
|
||||
if row_id is not None:
|
||||
# Get the row key for updating the cell
|
||||
row_key = event.cell_key.row_key
|
||||
|
||||
if row_id in self.selected_row_ids:
|
||||
self.selected_row_ids.remove(row_id)
|
||||
self.detail_table.update_cell(row_key, "select", "☐") # Unchecked
|
||||
else:
|
||||
self.selected_row_ids.add(row_id)
|
||||
self.detail_table.update_cell(row_key, "select", "☑️") # Checked
|
||||
|
||||
self._update_button_states()
|
||||
|
||||
async def on_data_table_header_selected(
|
||||
self, event: DataTable.HeaderSelected
|
||||
) -> None:
|
||||
column_key = event.column_key.value if event.column_key else None
|
||||
if not column_key:
|
||||
col_index = event.column_index
|
||||
if col_index == 0: # First column is "Select"
|
||||
return
|
||||
# Adjust for hidden _row_id column
|
||||
visible_cols = [
|
||||
col for col in self.activities_df.columns if col != "_row_id"
|
||||
]
|
||||
if col_index - 1 < len(visible_cols):
|
||||
column_key = visible_cols[col_index - 1]
|
||||
else:
|
||||
return
|
||||
if column_key == "select" or column_key == "_row_id":
|
||||
return
|
||||
|
||||
ascending = True
|
||||
if self._last_sort == (column_key, True):
|
||||
ascending = False
|
||||
self._last_sort = (column_key, ascending)
|
||||
|
||||
try:
|
||||
self.activities_df.sort_values(
|
||||
by=column_key, ascending=ascending, inplace=True
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to sort by column %s: %s", column_key, exc)
|
||||
return
|
||||
|
||||
# Only refresh rows, not columns
|
||||
await self._build_table(rebuild=False)
|
||||
|
||||
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)
|
||||
|
||||
if btn is self.add_allowlist_btn or btn_id == "add_allowlist_btn":
|
||||
await self._open_allowlist_screen()
|
||||
return
|
||||
|
||||
async def _select_all(self) -> None:
|
||||
"""Select all rows in the table."""
|
||||
# Add all row IDs to selected set
|
||||
self.selected_row_ids = set(self.activities_df["_row_id"].tolist())
|
||||
|
||||
# Update all checkboxes in the table
|
||||
for row_key, row_id in self.row_key_to_id.items():
|
||||
self.detail_table.update_cell(row_key, "select", "☑️")
|
||||
|
||||
self._update_button_states()
|
||||
logger.info(f"Selected all {len(self.selected_row_ids)} rows")
|
||||
|
||||
async def _select_none(self) -> None:
|
||||
"""Deselect all rows in the table."""
|
||||
# Clear selected set
|
||||
self.selected_row_ids.clear()
|
||||
|
||||
# Update all checkboxes in the table
|
||||
for row_key, row_id in self.row_key_to_id.items():
|
||||
self.detail_table.update_cell(row_key, "select", "☑️")
|
||||
|
||||
self._update_button_states()
|
||||
logger.info("Cleared all selections")
|
||||
|
||||
async def _open_allowlist_screen(self) -> None:
|
||||
"""Open the allowlist selection screen with selected activities."""
|
||||
if not self.selected_row_ids:
|
||||
self.app.notify("No rows selected", severity="warning")
|
||||
return
|
||||
|
||||
# Get selected data
|
||||
selected_df = self.get_selected_data()
|
||||
|
||||
# Get API from app
|
||||
api = getattr(self.app, "api", None)
|
||||
if api is None:
|
||||
logger.error("No API available on self.app.api")
|
||||
self.app.notify("API not available", severity="error")
|
||||
return
|
||||
|
||||
# Create and push AllowlistSelectionScreen
|
||||
try:
|
||||
allowlist_screen = AllowlistSelectionScreen(
|
||||
selected_df, api=api, hostname=self.hostname, otpid=self.otpid
|
||||
)
|
||||
await self.app.push_screen(allowlist_screen)
|
||||
logger.info(
|
||||
f"Opened allowlist screen with {len(selected_df)} selected activities"
|
||||
)
|
||||
except ImportError as e:
|
||||
logger.error(f"Failed to import AllowlistSelectionScreen: {e}")
|
||||
self.app.notify("Allowlist screen module not found", severity="error")
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to open allowlist screen: {e}")
|
||||
self.app.notify(
|
||||
f"Error opening allowlist screen: {str(e)}", severity="error"
|
||||
)
|
||||
|
||||
async def _export_detail_activities(self) -> None:
|
||||
if self.activities_df is None or self.activities_df.empty:
|
||||
logger.info("No activities to export.")
|
||||
await self.mount(Static("No activities to export.", classes="notification"))
|
||||
return
|
||||
if not self.selected_row_ids:
|
||||
logger.info("No rows selected for export.")
|
||||
await self.mount(
|
||||
Static("No rows selected for export.", classes="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)
|
||||
selected_df = self.get_selected_data()
|
||||
selected_df.to_csv(file_path, index=False)
|
||||
logger.info("Exported selected activities to %s", file_path)
|
||||
await self.mount(
|
||||
Static(
|
||||
f"Exported selected activities to: {filename}",
|
||||
classes="notification",
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to export detail activities: %s", exc)
|
||||
await self.mount(
|
||||
Static(
|
||||
"¢ Failed to export activities; check logs.",
|
||||
classes="notification",
|
||||
)
|
||||
)
|
||||
|
||||
# Helper methods
|
||||
def get_selected_data(self) -> pd.DataFrame:
|
||||
"""Return a DataFrame of the selected rows."""
|
||||
if not self.selected_row_ids:
|
||||
return pd.DataFrame()
|
||||
# Filter by selected row IDs and drop the internal _row_id column
|
||||
selected_df = self.activities_df[
|
||||
self.activities_df["_row_id"].isin(self.selected_row_ids)
|
||||
].copy()
|
||||
if "_row_id" in selected_df.columns:
|
||||
selected_df = selected_df.drop(columns=["_row_id"])
|
||||
return selected_df
|
||||
|
||||
def get_selected_records(self) -> list[dict]:
|
||||
"""Return selected rows as a list of dicts."""
|
||||
if not self.selected_row_ids:
|
||||
return []
|
||||
# Filter by selected row IDs and drop the internal _row_id column
|
||||
selected_df = self.activities_df[
|
||||
self.activities_df["_row_id"].isin(self.selected_row_ids)
|
||||
].copy()
|
||||
if "_row_id" in selected_df.columns:
|
||||
selected_df = selected_df.drop(columns=["_row_id"])
|
||||
return selected_df.to_dict(orient="records")
|
||||
|
||||
|
||||
class ActivityDetailScreen(Screen):
|
||||
"""
|
||||
Screen that wraps ActivityDetailWidget. Expects a DataFrame passed on init.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("e", "export", "Export"),
|
||||
Binding("a", "select_all", "Select All"),
|
||||
Binding("n", "select_none", "Select None"),
|
||||
]
|
||||
|
||||
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_go_back(self) -> None:
|
||||
try:
|
||||
await self.app.pop_screen()
|
||||
except Exception:
|
||||
logger.debug("ActivityDetailScreen.action_go_back pop_screen failed.")
|
||||
|
||||
async def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
await self.app.pop_screen()
|
||||
|
||||
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()
|
||||
|
||||
async def action_select_all(self) -> None:
|
||||
"""Handle 'a' key for select all."""
|
||||
if hasattr(self, "widget") and self.widget is not None:
|
||||
await self.widget._select_all()
|
||||
|
||||
async def action_select_none(self) -> None:
|
||||
"""Handle 'n' key for select none."""
|
||||
if hasattr(self, "widget") and self.widget is not None:
|
||||
await self.widget._select_none()
|
||||
|
||||
|
||||
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("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("r", "refresh_sessions", "Refresh Sessions"),
|
||||
Binding("e", "export_activities", "Export activities"),
|
||||
]
|
||||
|
||||
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)
|
||||
|
||||
async def action_go_back(self) -> None:
|
||||
"""Go back one 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()
|
||||
|
||||
# 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)
|
||||
|
||||
# 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)
|
||||
@@ -0,0 +1,392 @@
|
||||
# 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()
|
||||
@@ -0,0 +1,52 @@
|
||||
# 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 typing import List, Optional
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.screen import Screen
|
||||
|
||||
from models.agent import Agent
|
||||
from TUI.Widgets.OTP_generate import OTPGenerator
|
||||
|
||||
|
||||
class OTPWorkflowScreen(Screen):
|
||||
"""Screen that handles the OTP generation workflow without agent selection."""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
def __init__(self, selected_agents: Optional[List[Agent]]):
|
||||
super().__init__()
|
||||
self.selected_agents = selected_agents
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Directly show the OTP generator for the selected agents."""
|
||||
yield OTPGenerator(self.selected_agents)
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle escape key to go back one screen."""
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Handle q key to go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
|
||||
"""Handle OTP generation request - pass it up to the app level if needed."""
|
||||
File diff suppressed because it is too large
Load Diff
@@ -23,9 +23,11 @@ the policy selection workflow.
|
||||
import logging
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Footer, Header
|
||||
|
||||
from widgets.policyselector import PolicySelector
|
||||
from TUI.Widgets.policyselector import PolicySelector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -42,6 +44,11 @@ class PolicySelectorScreen(Screen):
|
||||
agent_move_operations: Reference to the parent AgentMoveOperations widget.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
CSS = """
|
||||
Screen {
|
||||
layout: vertical;
|
||||
@@ -68,7 +75,18 @@ class PolicySelectorScreen(Screen):
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Create the PolicySelector widget."""
|
||||
yield Header(show_clock=True)
|
||||
yield PolicySelector(self.policies)
|
||||
yield Footer()
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle escape key to go back one screen."""
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Handle q key to go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_policy_selector_policy_selected(
|
||||
self, message: PolicySelector.PolicySelected
|
||||
@@ -0,0 +1,914 @@
|
||||
# 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/>.
|
||||
|
||||
"""
|
||||
Quiet Agent Workflow Screen Module
|
||||
|
||||
Provides a TUI workflow for identifying quiet agents and moving them to target policies.
|
||||
This screen replaces the legacy quietAgent.py with a comprehensive TUI interface that:
|
||||
1. Allows selection of an initial policy to analyze
|
||||
2. Categorizes devices into "Enforce Ready" and "Non-Enforce Ready" based on activity
|
||||
3. Allows users to select target policies for each category
|
||||
4. Uses the API to move devices to their target policies
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import logging
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
import pandas as pd
|
||||
from textual.app import ComposeResult
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.reactive import reactive
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Input, Static
|
||||
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from TUI.Widgets.policyselector import PolicySelector
|
||||
from utils.configmanager import load_env
|
||||
from utils.executionfetcher import getExecutions
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QuietAgentWorkflowScreen(Screen):
|
||||
"""
|
||||
A Textual screen for the Quiet Agent analysis and migration workflow.
|
||||
|
||||
This screen provides a multi-step workflow:
|
||||
1. Select initial policy to analyze
|
||||
2. Configure analysis parameters (history period and quiet time period)
|
||||
3. View categorized agents (enforce ready vs. non-enforce ready)
|
||||
4. Select target policies for each category
|
||||
5. Execute agent migrations
|
||||
|
||||
Attributes:
|
||||
api (AirlockAPIWrapper): API wrapper for Airlock operations
|
||||
policies (List[Policy]): List of all available policies
|
||||
selected_policy (Optional[Policy]): The initially selected policy to analyze
|
||||
history_days (int): Number of days of history to pull (default: 150, range: 1-365)
|
||||
quiet_days (int): Number of days without execution to be considered quiet (default: 45, range: 1-365)
|
||||
agents_df (Optional[pd.DataFrame]): DataFrame of all agents with analysis results
|
||||
enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents ready for enforcement
|
||||
non_enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents not ready for enforcement
|
||||
workflow_stage (str): Current stage of the workflow
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
("escape", "go_back", "Back"),
|
||||
("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
workflow_stage = reactive("select_policy") # Tracks current workflow stage
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper, policies: List[Policy]):
|
||||
"""
|
||||
Initialize the QuietAgentWorkflowScreen.
|
||||
|
||||
Args:
|
||||
api (AirlockAPIWrapper): API wrapper for Airlock operations
|
||||
policies (List[Policy]): List of all available policies
|
||||
"""
|
||||
super().__init__()
|
||||
self.api = api
|
||||
self.policies = policies
|
||||
self.selected_policy: Optional[Policy] = None
|
||||
self.history_days = 150 # Default value, user-selectable
|
||||
self.quiet_days = 45 # Default value
|
||||
self.agents_df: Optional[pd.DataFrame] = None
|
||||
self.enforce_ready_df: Optional[pd.DataFrame] = None
|
||||
self.non_enforce_ready_df: Optional[pd.DataFrame] = None
|
||||
self.enforce_ready_target_policy: Optional[Policy] = None
|
||||
self.non_enforce_ready_target_policy: Optional[Policy] = None
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Build the UI layout for the workflow screen."""
|
||||
# Include Header and Footer like other standalone screens
|
||||
yield Header(show_clock=True, icon="⚙️")
|
||||
|
||||
# Title area
|
||||
title = Static("Quiet Agent Workflow", id="workflow_title")
|
||||
title.styles.margin = (0, 0, 0, 1)
|
||||
yield title
|
||||
|
||||
# Status area
|
||||
status = Static("Step 1: Select Policy to Analyze", id="workflow_status")
|
||||
status.styles.margin = (0, 0, 1, 1)
|
||||
yield status
|
||||
|
||||
# Content area - dynamically populated based on workflow stage
|
||||
yield Vertical(id="content_area")
|
||||
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Initialize the screen when mounted."""
|
||||
# Show initial policy selection
|
||||
self._show_policy_selection()
|
||||
|
||||
def watch_workflow_stage(self, old_value: str, new_value: str) -> None:
|
||||
"""React to workflow stage changes."""
|
||||
logger.debug(f"Workflow stage changed from {old_value} to {new_value}")
|
||||
self._update_status_message()
|
||||
|
||||
def _update_status_message(self) -> None:
|
||||
"""Update the status message based on current workflow stage."""
|
||||
status_widget = self.query_one("#workflow_status", Static)
|
||||
|
||||
stage_messages = {
|
||||
"select_policy": "Step 1: Select Policy to Analyze",
|
||||
"select_history_days": "Step 2: Configure Analysis Parameters",
|
||||
"analyzing": "Analyzing agent activity...",
|
||||
"view_results": "Step 3: Review Categorized Agents",
|
||||
"select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents",
|
||||
"select_non_enforce_target": "Step 5: Select Target Policy for Non-Enforce Ready Agents",
|
||||
"confirm_migration": "Step 6: Confirm and Execute Migration",
|
||||
"executing": "Executing agent migrations...",
|
||||
"complete": "Migration Complete",
|
||||
}
|
||||
|
||||
status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage"))
|
||||
|
||||
def _show_policy_selection(self) -> None:
|
||||
"""Show the initial policy selection screen."""
|
||||
self.workflow_stage = "select_policy"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Create policy selector widget
|
||||
policy_selector = PolicySelector(self.policies)
|
||||
content.mount(policy_selector)
|
||||
|
||||
def on_policy_selector_policy_selected(
|
||||
self, message: PolicySelector.PolicySelected
|
||||
) -> None:
|
||||
"""Handle policy selection from PolicySelector widget."""
|
||||
# Handle based on current workflow stage
|
||||
if self.workflow_stage == "select_policy":
|
||||
# Initial policy selection for analysis
|
||||
self.selected_policy = message.policy
|
||||
logger.info(f"Selected policy for analysis: {self.selected_policy.name}")
|
||||
self._show_history_days_selection()
|
||||
elif self.workflow_stage == "select_enforce_target":
|
||||
# Target policy selection for enforce ready agents
|
||||
self.enforce_ready_target_policy = message.policy
|
||||
logger.info(
|
||||
f"Selected target policy for enforce ready: {message.policy.name}"
|
||||
)
|
||||
self._show_non_enforce_target_selection()
|
||||
elif self.workflow_stage == "select_non_enforce_target":
|
||||
# Target policy selection for non-enforce ready agents
|
||||
self.non_enforce_ready_target_policy = message.policy
|
||||
logger.info(
|
||||
f"Selected target policy for non-enforce ready: {message.policy.name}"
|
||||
)
|
||||
self._show_migration_confirmation()
|
||||
|
||||
def _show_history_days_selection(self) -> None:
|
||||
"""Show the history days and quiet days selection screen."""
|
||||
self.workflow_stage = "select_history_days"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Create info text
|
||||
info_widget = Static(
|
||||
f"Policy Selected: {self.selected_policy.name}\n\n"
|
||||
"Configure Analysis Parameters:",
|
||||
id="analysis_params_info",
|
||||
)
|
||||
info_widget.styles.margin = (0, 0, 2, 0)
|
||||
content.mount(info_widget)
|
||||
|
||||
# Create input container
|
||||
input_container = Vertical(id="analysis_params_input_container")
|
||||
input_container.styles.height = "auto"
|
||||
content.mount(input_container)
|
||||
|
||||
# History days label
|
||||
history_label = Static("History Period (days of execution history to pull):")
|
||||
history_label.styles.margin = (0, 0, 1, 0)
|
||||
input_container.mount(history_label)
|
||||
|
||||
# Add history days input field
|
||||
history_input = Input(
|
||||
placeholder="Enter days (1-365, default: 150)",
|
||||
value="150",
|
||||
id="history_days_input",
|
||||
)
|
||||
history_input.styles.width = "50"
|
||||
history_input.styles.margin = (0, 0, 2, 0)
|
||||
input_container.mount(history_input)
|
||||
|
||||
# Quiet days label
|
||||
quiet_label = Static(
|
||||
"Quiet Time Period (days without execution to be considered quiet):"
|
||||
)
|
||||
quiet_label.styles.margin = (0, 0, 1, 0)
|
||||
input_container.mount(quiet_label)
|
||||
|
||||
# Add quiet days input field
|
||||
quiet_input = Input(
|
||||
placeholder="Enter days (1-365, default: 45)",
|
||||
value="45",
|
||||
id="quiet_days_input",
|
||||
)
|
||||
quiet_input.styles.width = "50"
|
||||
quiet_input.styles.margin = (0, 0, 2, 0)
|
||||
input_container.mount(quiet_input)
|
||||
|
||||
# Add submit button
|
||||
submit_btn = Button(
|
||||
"Continue",
|
||||
id="analysis_params_submit",
|
||||
variant="primary",
|
||||
)
|
||||
submit_btn.styles.width = "50"
|
||||
submit_btn.styles.margin = (1, 0, 0, 0)
|
||||
input_container.mount(submit_btn)
|
||||
|
||||
# Focus the first input field
|
||||
history_input.focus()
|
||||
|
||||
def _validate_and_submit_history_days(self) -> None:
|
||||
"""Validate and submit the history days and quiet days inputs."""
|
||||
try:
|
||||
history_input = self.query_one("#history_days_input", Input)
|
||||
quiet_input = self.query_one("#quiet_days_input", Input)
|
||||
|
||||
history_value = history_input.value.strip()
|
||||
quiet_value = quiet_input.value.strip()
|
||||
|
||||
# Validate history days
|
||||
if not history_value:
|
||||
self.app.notify(
|
||||
"Please enter a history period value", severity="error", timeout=3
|
||||
)
|
||||
history_input.focus()
|
||||
return
|
||||
|
||||
try:
|
||||
history_days = int(history_value)
|
||||
except ValueError:
|
||||
self.app.notify(
|
||||
"Please enter a valid number for history period",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
history_input.focus()
|
||||
return
|
||||
|
||||
if history_days < 1 or history_days > 365:
|
||||
self.app.notify(
|
||||
"History period must be between 1 and 365 days",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
history_input.focus()
|
||||
return
|
||||
|
||||
# Validate quiet days
|
||||
if not quiet_value:
|
||||
self.app.notify(
|
||||
"Please enter a quiet time period value",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
quiet_input.focus()
|
||||
return
|
||||
|
||||
try:
|
||||
quiet_days = int(quiet_value)
|
||||
except ValueError:
|
||||
self.app.notify(
|
||||
"Please enter a valid number for quiet time period",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
quiet_input.focus()
|
||||
return
|
||||
|
||||
if quiet_days < 1 or quiet_days > 365:
|
||||
self.app.notify(
|
||||
"Quiet time period must be between 1 and 365 days",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
quiet_input.focus()
|
||||
return
|
||||
|
||||
# Check that quiet days doesn't exceed history days
|
||||
if quiet_days > history_days:
|
||||
self.app.notify(
|
||||
"Quiet time period cannot exceed history period",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
quiet_input.focus()
|
||||
return
|
||||
|
||||
# All validation passed
|
||||
self.history_days = history_days
|
||||
self.quiet_days = quiet_days
|
||||
logger.info(
|
||||
f"Selected history days: {history_days}, quiet days: {quiet_days}"
|
||||
)
|
||||
self._start_analysis()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error validating analysis parameters: {e}")
|
||||
self.app.notify(f"Error: {str(e)}", severity="error", timeout=3)
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button press events."""
|
||||
button_id = event.button.id
|
||||
|
||||
# Analysis parameters submit button
|
||||
if button_id == "analysis_params_submit":
|
||||
self._validate_and_submit_history_days()
|
||||
return
|
||||
|
||||
# Navigation buttons
|
||||
if button_id == "select_enforce_target_btn":
|
||||
self._show_enforce_target_selection()
|
||||
return
|
||||
|
||||
if button_id == "select_non_enforce_target_btn":
|
||||
self._show_non_enforce_target_selection()
|
||||
return
|
||||
|
||||
if button_id == "skip_enforce_target_btn":
|
||||
# Skip enforce ready target selection
|
||||
self.enforce_ready_target_policy = None
|
||||
self._show_non_enforce_target_selection()
|
||||
return
|
||||
|
||||
if button_id == "skip_non_enforce_target_btn":
|
||||
# Skip non-enforce ready target selection
|
||||
self.non_enforce_ready_target_policy = None
|
||||
self._show_migration_confirmation()
|
||||
return
|
||||
|
||||
if button_id == "confirm_migration_btn":
|
||||
self._execute_migration()
|
||||
return
|
||||
|
||||
if button_id == "cancel_migration_btn":
|
||||
self._show_results()
|
||||
return
|
||||
|
||||
if button_id == "export_results_btn":
|
||||
self._export_results()
|
||||
return
|
||||
|
||||
if button_id == "start_over_btn":
|
||||
self._show_policy_selection()
|
||||
return
|
||||
|
||||
def on_input_submitted(self, event: Input.Submitted) -> None:
|
||||
"""Handle input submission (Enter key pressed)."""
|
||||
if event.input.id in ["history_days_input", "quiet_days_input"]:
|
||||
self._validate_and_submit_history_days()
|
||||
|
||||
def _start_analysis(self) -> None:
|
||||
"""Start the agent activity analysis."""
|
||||
# Show notification that analysis is starting
|
||||
self.app.notify(
|
||||
"Starting analysis - this may take several minutes for large policies",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
# Clear the screen to provide a blank canvas for Rust progress output
|
||||
# (Rust output displays over the TUI, so we clear everything except header/footer)
|
||||
try:
|
||||
# Clear title
|
||||
title_widget = self.query_one("#workflow_title", Static)
|
||||
title_widget.update("")
|
||||
|
||||
# Clear status
|
||||
status_widget = self.query_one("#workflow_status", Static)
|
||||
status_widget.update("")
|
||||
|
||||
# Clear content area
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not clear screen for analysis: {e}")
|
||||
|
||||
# Delay the analysis start to ensure UI refresh completes first
|
||||
# This prevents Rust output from starting before the screen is cleared
|
||||
self.set_timer(0.5, self._perform_analysis_worker)
|
||||
|
||||
def _perform_analysis_worker(self) -> None:
|
||||
"""Perform the actual agent activity analysis."""
|
||||
try:
|
||||
# Get agents in the selected policy
|
||||
agents = self.api.agents_find_by_group(self.selected_policy.groupid)
|
||||
|
||||
if agents.empty:
|
||||
self.app.notify(
|
||||
f"No agents found in policy: {self.selected_policy.name}",
|
||||
severity="warning",
|
||||
timeout=5,
|
||||
)
|
||||
self._show_policy_selection()
|
||||
return
|
||||
|
||||
# Get execution history (this shows progress bars in terminal via airlock_libs)
|
||||
policy_exec_history = getExecutions(
|
||||
self.api, self.selected_policy, [1, 2, 6, 7], self.history_days
|
||||
)
|
||||
|
||||
if policy_exec_history.empty:
|
||||
logger.info(
|
||||
"No execution history found for the selected policy and time range."
|
||||
)
|
||||
# All agents are quiet (no executions)
|
||||
agents["execution_count"] = 0
|
||||
agents["days_since"] = None
|
||||
agents["required_quiet"] = self.quiet_days
|
||||
agents["enforce_ready"] = True
|
||||
else:
|
||||
# Convert datetime column
|
||||
policy_exec_history["datetime"] = pd.to_datetime(
|
||||
policy_exec_history["datetime"],
|
||||
format="%Y-%m-%dT%H:%M:%SZ",
|
||||
utc=True,
|
||||
)
|
||||
|
||||
# Calculate days ago
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
|
||||
lambda dt: (now - dt).days
|
||||
)
|
||||
|
||||
# Count total executions per hostname
|
||||
hostname_counts = policy_exec_history["hostname"].value_counts()
|
||||
agents["execution_count"] = (
|
||||
agents["hostname"].map(hostname_counts).fillna(0).astype(int)
|
||||
)
|
||||
|
||||
# Find most recent execution per hostname
|
||||
most_recent_exec = policy_exec_history.sort_values(
|
||||
by="days_ago"
|
||||
).drop_duplicates(subset="hostname", keep="first")
|
||||
|
||||
# Map most recent execution age to agents
|
||||
agents["days_since"] = agents["hostname"].map(
|
||||
most_recent_exec.set_index("hostname")["days_ago"]
|
||||
)
|
||||
|
||||
# Check for enforcement readiness
|
||||
agents["required_quiet"] = self.quiet_days
|
||||
agents["enforce_ready"] = agents["days_since"].apply(
|
||||
lambda x: True if pd.isna(x) or x > self.quiet_days else False
|
||||
)
|
||||
|
||||
# Sort agents
|
||||
agents = agents.sort_values(
|
||||
by=["execution_count", "hostname"], ascending=[True, True]
|
||||
)
|
||||
|
||||
# Store the results
|
||||
self.agents_df = agents
|
||||
|
||||
# Categorize agents into DataFrames
|
||||
self.enforce_ready_df = agents[agents["enforce_ready"]].copy()
|
||||
self.non_enforce_ready_df = agents[~agents["enforce_ready"]].copy()
|
||||
|
||||
logger.info(
|
||||
f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, "
|
||||
f"{len(self.non_enforce_ready_df)} non-enforce ready"
|
||||
)
|
||||
|
||||
self.app.notify(
|
||||
f"Analysis complete! Found {len(self.enforce_ready_df)} enforce ready, "
|
||||
f"{len(self.non_enforce_ready_df)} not ready",
|
||||
severity="success",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
# Show results
|
||||
self._show_results()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during analysis: {e}", exc_info=True)
|
||||
self.app.notify(f"Analysis failed: {str(e)}", severity="error", timeout=5)
|
||||
self._show_policy_selection()
|
||||
|
||||
def _show_results(self) -> None:
|
||||
"""Show the categorized results."""
|
||||
self.workflow_stage = "view_results"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Create results display container and mount it first
|
||||
results_container = Vertical(id="results_container")
|
||||
results_container.styles.height = "auto"
|
||||
results_container.styles.margin = (1, 1)
|
||||
content.mount(results_container)
|
||||
|
||||
# Summary statistics
|
||||
total_agents = len(self.enforce_ready_df) + len(self.non_enforce_ready_df)
|
||||
ready_count = len(self.enforce_ready_df)
|
||||
not_ready_count = len(self.non_enforce_ready_df)
|
||||
ready_percentage = (ready_count / total_agents * 100) if total_agents > 0 else 0
|
||||
|
||||
summary = Static(
|
||||
f"Analysis Results for: {self.selected_policy.name}\n\n"
|
||||
f"Total Agents: {total_agents}\n"
|
||||
f"Enforce Ready: {ready_count} ({ready_percentage:.1f}%)\n"
|
||||
f"Not Ready: {not_ready_count} ({100 - ready_percentage:.1f}%)\n\n"
|
||||
f"Quiet Threshold: {self.quiet_days} days\n"
|
||||
f"History Period: {self.history_days} days",
|
||||
id="results_summary",
|
||||
)
|
||||
summary.styles.margin = (0, 0, 2, 0)
|
||||
results_container.mount(summary)
|
||||
|
||||
# Action buttons
|
||||
button_container = Horizontal(id="results_buttons")
|
||||
button_container.styles.height = "auto"
|
||||
results_container.mount(button_container)
|
||||
|
||||
if ready_count > 0:
|
||||
enforce_btn = Button(
|
||||
f"Select Target for Enforce Ready ({ready_count})",
|
||||
id="select_enforce_target_btn",
|
||||
)
|
||||
enforce_btn.styles.margin = (0, 1, 1, 0)
|
||||
button_container.mount(enforce_btn)
|
||||
|
||||
if not_ready_count > 0:
|
||||
non_enforce_btn = Button(
|
||||
f"Select Target for Non-Enforce Ready ({not_ready_count})",
|
||||
id="select_non_enforce_target_btn",
|
||||
)
|
||||
non_enforce_btn.styles.margin = (0, 1, 1, 0)
|
||||
button_container.mount(non_enforce_btn)
|
||||
|
||||
export_btn = Button("Export Results", id="export_results_btn")
|
||||
export_btn.styles.margin = (0, 1, 1, 0)
|
||||
button_container.mount(export_btn)
|
||||
|
||||
start_over_btn = Button("Start Over", id="start_over_btn")
|
||||
start_over_btn.styles.margin = (0, 0, 1, 0)
|
||||
button_container.mount(start_over_btn)
|
||||
|
||||
# Tables showing agents
|
||||
tables_container = Horizontal()
|
||||
tables_container.styles.height = "1fr"
|
||||
results_container.mount(tables_container)
|
||||
|
||||
# Enforce Ready table
|
||||
if ready_count > 0:
|
||||
enforce_col = Vertical()
|
||||
enforce_col.styles.width = "1fr"
|
||||
enforce_col.styles.margin = (1, 1, 0, 0)
|
||||
tables_container.mount(enforce_col)
|
||||
|
||||
enforce_label = Static("Enforce Ready Agents")
|
||||
enforce_label.styles.margin = (0, 0, 1, 0)
|
||||
enforce_col.mount(enforce_label)
|
||||
|
||||
enforce_table = DataTable(id="enforce_ready_table")
|
||||
enforce_table.styles.height = "1fr"
|
||||
enforce_table.add_columns("Hostname", "Last Exec (days)")
|
||||
|
||||
# Display first 50 agents
|
||||
for idx, row in self.enforce_ready_df.head(50).iterrows():
|
||||
days_since = row["days_since"]
|
||||
days_str = f"{int(days_since)}" if not pd.isna(days_since) else "Never"
|
||||
enforce_table.add_row(row["hostname"], days_str)
|
||||
|
||||
if len(self.enforce_ready_df) > 50:
|
||||
enforce_table.add_row(
|
||||
f"... and {len(self.enforce_ready_df) - 50} more", ""
|
||||
)
|
||||
|
||||
enforce_col.mount(enforce_table)
|
||||
|
||||
# Non-Enforce Ready table
|
||||
if not_ready_count > 0:
|
||||
non_enforce_col = Vertical()
|
||||
non_enforce_col.styles.width = "1fr"
|
||||
non_enforce_col.styles.margin = (1, 0, 0, 1)
|
||||
tables_container.mount(non_enforce_col)
|
||||
|
||||
non_enforce_label = Static("Non-Enforce Ready Agents")
|
||||
non_enforce_label.styles.margin = (0, 0, 1, 0)
|
||||
non_enforce_col.mount(non_enforce_label)
|
||||
|
||||
non_enforce_table = DataTable(id="non_enforce_ready_table")
|
||||
non_enforce_table.styles.height = "1fr"
|
||||
non_enforce_table.add_columns("Hostname", "Last Exec (days)")
|
||||
|
||||
# Display first 50 agents
|
||||
for idx, row in self.non_enforce_ready_df.head(50).iterrows():
|
||||
days_since = row["days_since"]
|
||||
days_str = f"{int(days_since)}" if not pd.isna(days_since) else "N/A"
|
||||
non_enforce_table.add_row(row["hostname"], days_str)
|
||||
|
||||
if len(self.non_enforce_ready_df) > 50:
|
||||
non_enforce_table.add_row(
|
||||
f"... and {len(self.non_enforce_ready_df) - 50} more", ""
|
||||
)
|
||||
|
||||
non_enforce_col.mount(non_enforce_table)
|
||||
|
||||
def _show_enforce_target_selection(self) -> None:
|
||||
"""Show policy selection for enforce ready agents."""
|
||||
self.workflow_stage = "select_enforce_target"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Info message
|
||||
info = Static(
|
||||
f"Select target policy for {len(self.enforce_ready_df)} Enforce Ready agents\n"
|
||||
f"Source Policy: {self.selected_policy.name}",
|
||||
id="enforce_target_info",
|
||||
)
|
||||
info.styles.margin = (0, 0, 2, 0)
|
||||
content.mount(info)
|
||||
|
||||
# Policy selector
|
||||
policy_selector = PolicySelector(self.policies)
|
||||
content.mount(policy_selector)
|
||||
|
||||
# Skip button
|
||||
skip_btn = Button("Skip - No Migration", id="skip_enforce_target_btn")
|
||||
skip_btn.styles.width = "50%"
|
||||
skip_btn.styles.margin = (2, 0, 0, 0)
|
||||
content.mount(skip_btn)
|
||||
|
||||
def _show_non_enforce_target_selection(self) -> None:
|
||||
"""Show policy selection for non-enforce ready agents."""
|
||||
self.workflow_stage = "select_non_enforce_target"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Info message
|
||||
info = Static(
|
||||
f"Select target policy for {len(self.non_enforce_ready_df)} Non-Enforce Ready agents\n"
|
||||
f"Source Policy: {self.selected_policy.name}",
|
||||
id="non_enforce_target_info",
|
||||
)
|
||||
info.styles.margin = (0, 0, 2, 0)
|
||||
content.mount(info)
|
||||
|
||||
# Policy selector
|
||||
policy_selector = PolicySelector(self.policies)
|
||||
content.mount(policy_selector)
|
||||
|
||||
# Skip button
|
||||
skip_btn = Button("Skip - No Migration", id="skip_non_enforce_target_btn")
|
||||
skip_btn.styles.width = "50%"
|
||||
skip_btn.styles.margin = (2, 0, 0, 0)
|
||||
content.mount(skip_btn)
|
||||
|
||||
def _show_migration_confirmation(self) -> None:
|
||||
"""Show migration confirmation screen."""
|
||||
self.workflow_stage = "confirm_migration"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Build confirmation message
|
||||
confirmation_lines = [
|
||||
"Migration Summary\n",
|
||||
f"Source Policy: {self.selected_policy.name}\n",
|
||||
]
|
||||
|
||||
if self.enforce_ready_target_policy:
|
||||
confirmation_lines.append(
|
||||
f"\nEnforce Ready Migration:\n"
|
||||
f"Agents: {len(self.enforce_ready_df)}\n"
|
||||
f"Target: {self.enforce_ready_target_policy.name}\n"
|
||||
)
|
||||
|
||||
if self.non_enforce_ready_target_policy:
|
||||
confirmation_lines.append(
|
||||
f"\nNon-Enforce Ready Migration:\n"
|
||||
f"Agents: {len(self.non_enforce_ready_df)}\n"
|
||||
f"Target: {self.non_enforce_ready_target_policy.name}\n"
|
||||
)
|
||||
|
||||
if (
|
||||
not self.enforce_ready_target_policy
|
||||
and not self.non_enforce_ready_target_policy
|
||||
):
|
||||
confirmation_lines.append("\nNo migrations will be performed.")
|
||||
|
||||
confirmation = Static("".join(confirmation_lines), id="migration_confirmation")
|
||||
confirmation.styles.margin = (1, 1, 2, 1)
|
||||
content.mount(confirmation)
|
||||
|
||||
# Action buttons - mount container first, then add buttons
|
||||
button_container = Horizontal(id="confirmation_buttons")
|
||||
button_container.styles.height = "auto"
|
||||
button_container.styles.margin = (1, 1)
|
||||
content.mount(button_container)
|
||||
|
||||
if self.enforce_ready_target_policy or self.non_enforce_ready_target_policy:
|
||||
confirm_btn = Button("Confirm Migration", id="confirm_migration_btn")
|
||||
confirm_btn.styles.margin = (0, 1, 0, 0)
|
||||
button_container.mount(confirm_btn)
|
||||
|
||||
cancel_btn = Button("Cancel", id="cancel_migration_btn")
|
||||
button_container.mount(cancel_btn)
|
||||
|
||||
def _execute_migration(self) -> None:
|
||||
"""Execute the agent migrations."""
|
||||
self.workflow_stage = "executing"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Show executing message
|
||||
executing_msg = Static(
|
||||
"Executing agent migrations...\nPlease wait...",
|
||||
id="executing_message",
|
||||
)
|
||||
executing_msg.styles.margin = (2, 1)
|
||||
content.mount(executing_msg)
|
||||
|
||||
# Perform migrations asynchronously
|
||||
self.call_later(self._perform_migrations)
|
||||
|
||||
def _perform_migrations(self) -> None:
|
||||
"""Perform the actual agent migrations."""
|
||||
successful_migrations = []
|
||||
failed_migrations = []
|
||||
|
||||
try:
|
||||
# Migrate enforce ready agents
|
||||
if self.enforce_ready_target_policy:
|
||||
for idx, row in self.enforce_ready_df.iterrows():
|
||||
try:
|
||||
self.api.agent_move(
|
||||
row["agentid"], self.enforce_ready_target_policy.groupid
|
||||
)
|
||||
successful_migrations.append(
|
||||
(row["hostname"], self.enforce_ready_target_policy.name)
|
||||
)
|
||||
logger.debug(
|
||||
f"Moved {row['hostname']} to {self.enforce_ready_target_policy.name}"
|
||||
)
|
||||
except Exception as e:
|
||||
failed_migrations.append((row["hostname"], str(e)))
|
||||
logger.error(f"Failed to move {row['hostname']}: {e}")
|
||||
|
||||
# Migrate non-enforce ready agents
|
||||
if self.non_enforce_ready_target_policy:
|
||||
for idx, row in self.non_enforce_ready_df.iterrows():
|
||||
try:
|
||||
self.api.agent_move(
|
||||
row["agentid"], self.non_enforce_ready_target_policy.groupid
|
||||
)
|
||||
successful_migrations.append(
|
||||
(row["hostname"], self.non_enforce_ready_target_policy.name)
|
||||
)
|
||||
logger.debug(
|
||||
f"Moved {row['hostname']} to {self.non_enforce_ready_target_policy.name}"
|
||||
)
|
||||
except Exception as e:
|
||||
failed_migrations.append((row["hostname"], str(e)))
|
||||
logger.error(f"Failed to move {row['hostname']}: {e}")
|
||||
|
||||
# Show completion results
|
||||
self._show_completion_results(successful_migrations, failed_migrations)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during migration execution: {e}", exc_info=True)
|
||||
self.app.notify(f"Migration failed: {str(e)}", severity="error", timeout=5)
|
||||
self._show_results()
|
||||
|
||||
def _show_completion_results(
|
||||
self, successful: List[tuple], failed: List[tuple]
|
||||
) -> None:
|
||||
"""Show migration completion results."""
|
||||
self.workflow_stage = "complete"
|
||||
content = self.query_one("#content_area", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
# Results summary
|
||||
total_attempted = len(successful) + len(failed)
|
||||
success_rate = (
|
||||
(len(successful) / total_attempted * 100) if total_attempted > 0 else 0
|
||||
)
|
||||
|
||||
results = Static(
|
||||
f"Migration Complete\n\n"
|
||||
f"Total Agents Migrated: {len(successful)}\n"
|
||||
f"Failed Migrations: {len(failed)}\n"
|
||||
f"Success Rate: {success_rate:.1f}%",
|
||||
id="completion_summary",
|
||||
)
|
||||
results.styles.margin = (1, 1, 2, 1)
|
||||
content.mount(results)
|
||||
|
||||
# Details tables
|
||||
if successful:
|
||||
success_container = Vertical()
|
||||
success_container.styles.margin = (0, 1)
|
||||
content.mount(success_container)
|
||||
|
||||
success_label = Static("Successful Migrations")
|
||||
success_label.styles.margin = (0, 0, 1, 0)
|
||||
success_container.mount(success_label)
|
||||
|
||||
success_table = DataTable(id="success_table")
|
||||
success_table.styles.height = "auto"
|
||||
success_table.add_columns("Hostname", "Target Policy")
|
||||
|
||||
for hostname, target_policy in successful[:25]: # Show first 25
|
||||
success_table.add_row(hostname, target_policy)
|
||||
|
||||
if len(successful) > 25:
|
||||
success_table.add_row(f"... and {len(successful) - 25} more", "")
|
||||
|
||||
success_container.mount(success_table)
|
||||
|
||||
if failed:
|
||||
failed_container = Vertical()
|
||||
failed_container.styles.margin = (2, 1, 0, 1)
|
||||
content.mount(failed_container)
|
||||
|
||||
failed_label = Static("Failed Migrations")
|
||||
failed_label.styles.margin = (0, 0, 1, 0)
|
||||
failed_container.mount(failed_label)
|
||||
|
||||
failed_table = DataTable(id="failed_table")
|
||||
failed_table.styles.height = "auto"
|
||||
failed_table.add_columns("Hostname", "Error")
|
||||
|
||||
for hostname, error in failed[:25]: # Show first 25
|
||||
failed_table.add_row(hostname, error[:50]) # Truncate error
|
||||
|
||||
if len(failed) > 25:
|
||||
failed_table.add_row(f"... and {len(failed) - 25} more", "")
|
||||
|
||||
failed_container.mount(failed_table)
|
||||
|
||||
# Action button
|
||||
done_btn = Button("Done", id="start_over_btn")
|
||||
done_btn.styles.width = "50%"
|
||||
done_btn.styles.margin = (2, 0, 0, 0)
|
||||
content.mount(done_btn)
|
||||
|
||||
def _export_results(self) -> None:
|
||||
"""Export analysis results to CSV."""
|
||||
try:
|
||||
working_dir = load_env("WORKING_DIR") or os.getcwd()
|
||||
filename = os.path.join(
|
||||
working_dir,
|
||||
f"{self.selected_policy.name}_quiet_analysis_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
|
||||
)
|
||||
|
||||
self.agents_df.to_csv(filename, index=False)
|
||||
logger.info(f"Exported results to {filename}")
|
||||
self.app.notify(
|
||||
f"Results exported to:\n{filename}",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to export results: {e}")
|
||||
self.app.notify(f"Export failed: {str(e)}", severity="error", timeout=5)
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle back/escape action."""
|
||||
# Depending on stage, go back to previous stage or exit
|
||||
if self.workflow_stage in ["select_policy", "view_results", "complete"]:
|
||||
self.app.pop_screen()
|
||||
elif self.workflow_stage == "select_history_days":
|
||||
self._show_policy_selection()
|
||||
elif self.workflow_stage == "select_enforce_target":
|
||||
self._show_results()
|
||||
elif self.workflow_stage == "select_non_enforce_target":
|
||||
if self.enforce_ready_target_policy:
|
||||
self._show_enforce_target_selection()
|
||||
else:
|
||||
self._show_results()
|
||||
elif self.workflow_stage == "confirm_migration":
|
||||
self._show_non_enforce_target_selection()
|
||||
else:
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
@@ -1,5 +1,20 @@
|
||||
# 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/>.
|
||||
|
||||
import logging
|
||||
from typing import List
|
||||
from typing import List, Optional
|
||||
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.css.query import NoMatches
|
||||
@@ -23,15 +38,20 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OTPGenerator(Widget):
|
||||
"""Widget for generating OTPs for selected devices."""
|
||||
|
||||
# Reactive properties to track form completion
|
||||
requestor_filled = reactive(False)
|
||||
reasoning_filled = reactive(False)
|
||||
duration_selected = reactive(True) # Default is selected
|
||||
otp_generated = reactive(False)
|
||||
|
||||
class OTPInfo(Message):
|
||||
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__()
|
||||
self.devices = devices
|
||||
@@ -135,14 +155,10 @@ class OTPGenerator(Widget):
|
||||
button_row.styles.height = "auto"
|
||||
button_row.styles.margin = (1, 0, 0, 0)
|
||||
|
||||
back_button = Button("← Back", id="back_button")
|
||||
back_button.styles.width = "1fr"
|
||||
yield back_button
|
||||
|
||||
generate_button = Button(
|
||||
"Generate OTP", id="generate_button", variant="primary"
|
||||
)
|
||||
generate_button.styles.width = "2fr"
|
||||
generate_button.styles.width = "100%"
|
||||
yield generate_button
|
||||
|
||||
# Right side - Show device list initially, then output after generation
|
||||
@@ -165,7 +181,7 @@ class OTPGenerator(Widget):
|
||||
|
||||
# Show device list initially
|
||||
device_list_text = "\n".join(
|
||||
f"• {device.hostname}" for device in self.devices
|
||||
f"{device.hostname}" for device in self.devices
|
||||
)
|
||||
device_display = Static(device_list_text, id="device_display")
|
||||
yield device_display
|
||||
@@ -193,14 +209,7 @@ class OTPGenerator(Widget):
|
||||
def on_button_pressed(self, event: Button.Pressed):
|
||||
btn_id = event.button.id
|
||||
|
||||
if btn_id == "back_button":
|
||||
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "copy_clipboard_button":
|
||||
if btn_id == "copy_clipboard_button":
|
||||
try:
|
||||
output_area = self.query_one("#otp_output", TextArea)
|
||||
text_to_copy = output_area.text
|
||||
@@ -209,7 +218,7 @@ class OTPGenerator(Widget):
|
||||
|
||||
pyperclip.copy(text_to_copy)
|
||||
self.app.notify(
|
||||
"✅ Copied to clipboard!", severity="information", timeout=2
|
||||
"✓ Copied to clipboard!", severity="information", timeout=2
|
||||
)
|
||||
except ImportError:
|
||||
self.app.notify(
|
||||
@@ -243,7 +252,7 @@ class OTPGenerator(Widget):
|
||||
self.otp_generated = True
|
||||
|
||||
# Access API from the app - this is the key change!
|
||||
api = self.app.api
|
||||
api = self.app.api # type: ignore
|
||||
|
||||
output_lines = [
|
||||
"Requested OTP Codes:",
|
||||
@@ -1,22 +1,18 @@
|
||||
"""
|
||||
Agent Move Operations Widget Module
|
||||
# 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/>.
|
||||
|
||||
This module provides a Textual-based UI widget for performing bulk operations on
|
||||
agent devices in the Airlock system. It allows users to:
|
||||
- View selected agents and their current policy assignments
|
||||
- Move agents to local approval mode with OTP enforcement
|
||||
- Toggle agents between audit and enforcement policy modes
|
||||
- Select and move agents to alternate policies (future implementation)
|
||||
|
||||
The widget tracks operation state, manages button availability, and displays
|
||||
results with success/failure summaries that can be copied to clipboard.
|
||||
|
||||
Dependencies:
|
||||
- textual: TUI framework for building the widget and UI components
|
||||
- models.agent: Agent model class
|
||||
- services.agenthandler: Core agent operation functions
|
||||
- flows.localApproval: Local approval workflow handling
|
||||
"""
|
||||
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime
|
||||
@@ -30,16 +26,74 @@ from textual.css.query import NoMatches
|
||||
from textual.message import Message
|
||||
from textual.reactive import reactive
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Button, DataTable, Header, Static, TextArea
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea
|
||||
|
||||
from models.agent import Agent
|
||||
from screens.otpworkflowscreen import OTPWorkflowScreen
|
||||
from screens.policyselectorscreen import PolicySelectorScreen
|
||||
from widgets.OTP_generate import OTPGenerator
|
||||
from services.API import AirlockAPIWrapper
|
||||
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
|
||||
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
|
||||
from TUI.Screens.policyselectorscreen import PolicySelectorScreen
|
||||
from TUI.Widgets.OTP_generate import OTPGenerator
|
||||
from utils.configmanager import get_system_json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def moveAgentToRelatedPolicy(
|
||||
api: AirlockAPIWrapper,
|
||||
agent: Agent,
|
||||
mode: str = "audit",
|
||||
):
|
||||
"""
|
||||
Moves an agent between audit and enforcement policies based on the mode.
|
||||
|
||||
Args:
|
||||
api: AirlockAPIWrapper instance.
|
||||
agent: Agent object.
|
||||
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
|
||||
"""
|
||||
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
|
||||
if mode == "audit":
|
||||
if agent.groupid in policy_relationship_map:
|
||||
target_policy = policy_relationship_map[agent.groupid]
|
||||
elif agent.groupid in policy_relationship_map.values():
|
||||
logger.debug(
|
||||
f"Agent {agent.hostname} is already in an audit group. No action needed."
|
||||
)
|
||||
print(
|
||||
f"Agent {agent.hostname} is already in an audit group. No action needed."
|
||||
)
|
||||
return
|
||||
else:
|
||||
logger.warning(
|
||||
f"Error: No corresponding audit policy found for groupid: {agent.groupid}."
|
||||
)
|
||||
return
|
||||
|
||||
elif mode == "enforcement":
|
||||
inverse_map = {v: k for k, v in policy_relationship_map.items()}
|
||||
if agent.groupid in inverse_map:
|
||||
target_policy = inverse_map[agent.groupid]
|
||||
elif agent.groupid in inverse_map.values():
|
||||
logger.info(
|
||||
f"Agent {agent.hostname} is already in an enforcement group. No action needed."
|
||||
)
|
||||
return
|
||||
else:
|
||||
logger.warning(
|
||||
f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}."
|
||||
)
|
||||
return
|
||||
|
||||
else:
|
||||
logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.")
|
||||
return
|
||||
|
||||
result = api.agent_move(agent.agentid, target_policy)
|
||||
return result
|
||||
|
||||
|
||||
class AgentMoveOperations(Widget):
|
||||
"""
|
||||
A Textual widget for managing bulk agent operations and policy migrations.
|
||||
@@ -144,6 +198,7 @@ class AgentMoveOperations(Widget):
|
||||
toggle_enforcement_btn = self.query_one("#toggle_enforcement_btn", Button)
|
||||
other_policy_btn = self.query_one("#other_policy_btn", Button)
|
||||
otp_gen_btn = self.query_one("#otp_gen_btn", Button)
|
||||
exec_history_btn = self.query_one("#exec_history_btn", Button)
|
||||
|
||||
# If operation in progress, disable all
|
||||
if self.operation_in_progress:
|
||||
@@ -152,6 +207,7 @@ class AgentMoveOperations(Widget):
|
||||
local_approval_btn.disabled = True
|
||||
toggle_enforcement_btn.disabled = True
|
||||
other_policy_btn.disabled = True
|
||||
exec_history_btn.disabled = True
|
||||
else:
|
||||
# If an operation was selected, disable
|
||||
if self.selected_operation:
|
||||
@@ -166,6 +222,9 @@ class AgentMoveOperations(Widget):
|
||||
other_policy_btn.disabled = (
|
||||
self.selected_operation == "other_policy"
|
||||
)
|
||||
exec_history_btn.disabled = (
|
||||
self.selected_operation == "exec_history"
|
||||
)
|
||||
else:
|
||||
# Enable all buttons
|
||||
otp_gen_btn = False
|
||||
@@ -173,6 +232,7 @@ class AgentMoveOperations(Widget):
|
||||
local_approval_btn.disabled = False
|
||||
toggle_enforcement_btn.disabled = False
|
||||
other_policy_btn.disabled = False
|
||||
exec_history_btn.disabled = False
|
||||
|
||||
except NoMatches:
|
||||
pass
|
||||
@@ -252,7 +312,7 @@ class AgentMoveOperations(Widget):
|
||||
- Operations panel: 1/3 width
|
||||
- Results area: Initially hidden, shown after operation completion
|
||||
"""
|
||||
yield Header(show_clock=True, icon="⚙")
|
||||
yield Header(show_clock=True, icon="⚙️")
|
||||
title_text = Static(
|
||||
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
|
||||
id="move_ops_title",
|
||||
@@ -289,7 +349,7 @@ class AgentMoveOperations(Widget):
|
||||
yield operations_label
|
||||
|
||||
# Operation buttons
|
||||
export_csv_btn = Button("📈 Export CSV", id="export_csv_btn")
|
||||
export_csv_btn = Button("📄 Export CSV", id="export_csv_btn")
|
||||
export_csv_btn.styles.width = "100%"
|
||||
export_csv_btn.styles.margin = (0, 0, 1, 0)
|
||||
yield export_csv_btn
|
||||
@@ -320,15 +380,19 @@ class AgentMoveOperations(Widget):
|
||||
other_policy_btn.styles.margin = (0, 0, 1, 0)
|
||||
yield other_policy_btn
|
||||
|
||||
exec_history_btn = Button(
|
||||
"📊 View Execution History", id="exec_history_btn"
|
||||
)
|
||||
exec_history_btn.styles.width = "100%"
|
||||
exec_history_btn.styles.margin = (0, 0, 1, 0)
|
||||
yield exec_history_btn
|
||||
|
||||
# Status label
|
||||
status_label = Static("", id="status_label")
|
||||
status_label.styles.margin = (2, 0, 0, 0)
|
||||
yield status_label
|
||||
|
||||
back_button = Button("← Back", id="back_button")
|
||||
back_button.styles.width = "50%"
|
||||
back_button.styles.margin = (0, 1, 1, 0)
|
||||
yield back_button
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""
|
||||
@@ -363,7 +427,7 @@ class AgentMoveOperations(Widget):
|
||||
Handle button press events from the widget.
|
||||
|
||||
This Textual event handler routes button presses to appropriate actions:
|
||||
- back_button: Pop this screen (return to parent)
|
||||
|
||||
- copy_results_btn: Copy results text to clipboard (requires pyperclip)
|
||||
- local_approval_btn: Start local approval operation
|
||||
- toggle_enforcement_btn: Start toggle audit/enforcement operation
|
||||
@@ -377,12 +441,7 @@ class AgentMoveOperations(Widget):
|
||||
|
||||
btn_id = event.button.id
|
||||
|
||||
if btn_id == "back_button":
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "copy_results_btn":
|
||||
if btn_id == "copy_results_btn":
|
||||
try:
|
||||
results_text = self.query_one("#results_text", TextArea)
|
||||
import pyperclip
|
||||
@@ -399,7 +458,7 @@ class AgentMoveOperations(Widget):
|
||||
severity="warning",
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(f"⌠Failed to copy: {str(e)}", severity="error")
|
||||
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
|
||||
event.stop()
|
||||
elif btn_id == "export_csv_btn":
|
||||
self._start_export_csv_operation()
|
||||
@@ -419,6 +478,9 @@ class AgentMoveOperations(Widget):
|
||||
elif btn_id == "otp_gen_btn":
|
||||
self._start_OTP_gen_operation()
|
||||
event.stop()
|
||||
elif btn_id == "exec_history_btn":
|
||||
self._start_execution_history_operation()
|
||||
event.stop()
|
||||
|
||||
def _start_local_approval_operation(self) -> None:
|
||||
"""
|
||||
@@ -458,8 +520,6 @@ class AgentMoveOperations(Widget):
|
||||
try:
|
||||
import time
|
||||
|
||||
from services.agenthandler import moveAgentToRelatedPolicy
|
||||
|
||||
# Generate batch ID
|
||||
batch = int(time.time())
|
||||
duration = 360 # Default 6 hours, could make this configurable
|
||||
@@ -503,7 +563,6 @@ class AgentMoveOperations(Widget):
|
||||
self.selected_operation = "export_csv"
|
||||
self.operation_in_progress = True
|
||||
successful = []
|
||||
unsuccessful = []
|
||||
status_label = self.query_one("#status_label", Static)
|
||||
status_label.update("Exporting CSV...")
|
||||
self.app.refresh_data()
|
||||
@@ -577,7 +636,7 @@ class AgentMoveOperations(Widget):
|
||||
self.operation_in_progress = True
|
||||
|
||||
status_label = self.query_one("#status_label", Static)
|
||||
status_label.update("â³ Toggling enforcement mode...")
|
||||
status_label.update("🔄 Toggling enforcement mode...")
|
||||
|
||||
# Get API from app
|
||||
api = self.app.api
|
||||
@@ -586,10 +645,7 @@ class AgentMoveOperations(Widget):
|
||||
unsuccessful = []
|
||||
|
||||
try:
|
||||
from services.agenthandler import moveAgentToRelatedPolicy
|
||||
from utils.configmanager import get_protected_json
|
||||
|
||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
|
||||
for agent in self.agents:
|
||||
try:
|
||||
@@ -618,7 +674,7 @@ class AgentMoveOperations(Widget):
|
||||
return
|
||||
|
||||
self.operation_in_progress = False
|
||||
status_label.update("✅ Operation complete!")
|
||||
status_label.update("✅ Operation complete!")
|
||||
|
||||
# Display results in the widget
|
||||
self._display_results("Toggle Audit/Enforcement", successful, unsuccessful)
|
||||
@@ -694,6 +750,35 @@ class AgentMoveOperations(Widget):
|
||||
|
||||
self.app.push_screen(OTPWorkflowScreen(self.agents))
|
||||
|
||||
def _start_execution_history_operation(self) -> None:
|
||||
"""
|
||||
Launch the execution history viewer for selected agents.
|
||||
|
||||
This operation opens a new screen that allows the user to:
|
||||
1. Select a date range for execution history
|
||||
2. Fetch execution logs for all selected agents
|
||||
3. View the results in a table
|
||||
4. Export the results to CSV
|
||||
|
||||
The screen is pushed onto the screen stack, allowing the user to return
|
||||
to this screen when done.
|
||||
"""
|
||||
status_label = self.query_one("#status_label", Static)
|
||||
status_label.update("Opening execution history viewer...")
|
||||
|
||||
try:
|
||||
# Push the execution history screen
|
||||
self.app.push_screen(ExecutionHistoryScreen(self.agents))
|
||||
logger.info(
|
||||
f"Opened execution history viewer for {len(self.agents)} agents"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to open execution history viewer: {e}")
|
||||
status_label.update(f"❌ Error: {str(e)}")
|
||||
self.app.notify(
|
||||
f"Failed to open execution history: {str(e)}", severity="error"
|
||||
)
|
||||
|
||||
def _execute_move_to_policy(self, target_policy) -> None:
|
||||
"""
|
||||
Execute the actual move of agents to the selected policy.
|
||||
@@ -715,7 +800,7 @@ class AgentMoveOperations(Widget):
|
||||
for agent in self.agents:
|
||||
try:
|
||||
# Move agent to target policy
|
||||
result = api.agent_move(agent.agentid, target_policy.groupid)
|
||||
api.agent_move(agent.agentid, target_policy.groupid)
|
||||
successful.append((agent, f"Moved to {target_policy.name}"))
|
||||
logger.info(
|
||||
f"Successfully moved {agent.hostname} to policy {target_policy.name}"
|
||||
@@ -1,6 +1,22 @@
|
||||
# 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/>.
|
||||
|
||||
import difflib
|
||||
from pathlib import Path
|
||||
import re
|
||||
from typing import List
|
||||
from typing import List, Optional
|
||||
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.css.query import NoMatches
|
||||
@@ -20,12 +36,14 @@ from models.agent import Agent
|
||||
|
||||
|
||||
class MultiAgentSelector(Widget):
|
||||
"""Widget for selecting multiple agents from a list."""
|
||||
|
||||
class AgentsSelected(Message):
|
||||
def __init__(self, selected_agents: List[Agent]):
|
||||
super().__init__()
|
||||
self.selected_agents = selected_agents
|
||||
|
||||
def __init__(self, all_agents: List[Agent]):
|
||||
def __init__(self, all_agents: Optional[List[Agent]]):
|
||||
super().__init__()
|
||||
self.all_agents = all_agents
|
||||
self._match_type = "exact"
|
||||
@@ -40,7 +58,7 @@ class MultiAgentSelector(Widget):
|
||||
|
||||
def compose(self):
|
||||
yield Header(show_clock=True, icon="⚙")
|
||||
title_text = Static("🖧 Agent Selector", id="selector_title")
|
||||
title_text = Static("🖥️ Agent Selector", id="selector_title")
|
||||
title_text.styles.margin = (0, 0, 0, 1)
|
||||
yield title_text
|
||||
|
||||
@@ -60,7 +78,7 @@ class MultiAgentSelector(Widget):
|
||||
text_area.styles.overflow_y = "auto"
|
||||
yield text_area
|
||||
|
||||
with Horizontal(id="switch_search_container") as switch_search:
|
||||
with Horizontal(id="switch_container"):
|
||||
switch = Switch(value=False, id="match_switch")
|
||||
switch.styles.width = "auto"
|
||||
switch.styles.margin = (1, 0, 0, 0)
|
||||
@@ -72,8 +90,13 @@ class MultiAgentSelector(Widget):
|
||||
switch_label.styles.margin = (2, 1, 0, 0)
|
||||
yield switch_label
|
||||
|
||||
with Horizontal(id="action_buttons_container"):
|
||||
load_file = Button("📂 Load File", id="load_file_button")
|
||||
load_file.styles.margin = (1, 1, 0, 1)
|
||||
yield load_file
|
||||
|
||||
search = Button("🔍 Search", id="search_button")
|
||||
search.styles.margin = (1, 0, 0, 0)
|
||||
search.styles.margin = (1, 0, 0, 1)
|
||||
yield search
|
||||
|
||||
with Horizontal() as select_buttons:
|
||||
@@ -91,11 +114,6 @@ class MultiAgentSelector(Widget):
|
||||
button_row.styles.height = "auto"
|
||||
button_row.styles.margin = (1, 0, 0, 0)
|
||||
|
||||
back_button = Button("← Back", id="back_button")
|
||||
back_button.styles.width = "1fr"
|
||||
back_button.styles.margin = (0, 0, 0, 1)
|
||||
yield back_button
|
||||
|
||||
submit_button = Button(
|
||||
"▶ Select & Continue", id="submit_selection", variant="primary"
|
||||
)
|
||||
@@ -123,10 +141,7 @@ class MultiAgentSelector(Widget):
|
||||
match_list = self.query_one("#match_results", SelectionList)
|
||||
except NoMatches:
|
||||
return
|
||||
if btn_id == "back_button":
|
||||
self.app.pop_screen()
|
||||
event.stop()
|
||||
elif btn_id == "select_all":
|
||||
if btn_id == "select_all":
|
||||
match_list.select_all()
|
||||
event.stop()
|
||||
elif btn_id == "select_none":
|
||||
@@ -143,6 +158,9 @@ class MultiAgentSelector(Widget):
|
||||
]
|
||||
self.post_message(self.AgentsSelected(selected_agents))
|
||||
event.stop()
|
||||
elif btn_id == "load_file_button":
|
||||
self._load_from_file()
|
||||
event.stop()
|
||||
elif btn_id == "search_button":
|
||||
self.update_matches()
|
||||
event.stop()
|
||||
@@ -157,7 +175,7 @@ class MultiAgentSelector(Widget):
|
||||
match_list.add_option((name, name))
|
||||
unmatched_label = self.query_one("#unmatched_label", Static)
|
||||
if unmatched:
|
||||
unmatched_label.update(f"⚠️ No matches for: {', '.join(unmatched)}")
|
||||
unmatched_label.update(f"âš ï¸ No matches for: {', '.join(unmatched)}")
|
||||
else:
|
||||
unmatched_label.update("")
|
||||
|
||||
@@ -205,3 +223,118 @@ class MultiAgentSelector(Widget):
|
||||
else:
|
||||
unmatched.append(name)
|
||||
return sorted(matched), unmatched
|
||||
|
||||
def _load_from_file(self):
|
||||
"""Safely load device names from a text file."""
|
||||
try:
|
||||
# Import here to avoid issues if tkinter isn't available
|
||||
import tkinter as tk
|
||||
from tkinter import filedialog
|
||||
|
||||
# Create file dialog
|
||||
root = tk.Tk()
|
||||
root.withdraw()
|
||||
|
||||
file_path = filedialog.askopenfilename(
|
||||
title="Select device list file",
|
||||
filetypes=[
|
||||
("Text files", "*.txt"),
|
||||
("CSV files", "*.csv"),
|
||||
("All files", "*.*"),
|
||||
],
|
||||
)
|
||||
|
||||
if not file_path:
|
||||
# User cancelled
|
||||
return
|
||||
|
||||
# Validate file path
|
||||
path_obj = Path(file_path)
|
||||
if not path_obj.exists():
|
||||
self.app.notify("File does not exist", severity="error", timeout=3)
|
||||
return
|
||||
|
||||
if not path_obj.is_file():
|
||||
self.app.notify(
|
||||
"Selected path is not a file", severity="error", timeout=3
|
||||
)
|
||||
return
|
||||
|
||||
# Check file size (limit to 1 MB for safety)
|
||||
file_size = path_obj.stat().st_size
|
||||
if file_size > 1_000_000: # 1 MB
|
||||
self.app.notify(
|
||||
f"File too large ({file_size:,} bytes). Maximum 1 MB.",
|
||||
severity="error",
|
||||
timeout=5,
|
||||
)
|
||||
return
|
||||
|
||||
# Read file with proper encoding to preserve emojis
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
except UnicodeDecodeError:
|
||||
# Try with different encoding if UTF-8 fails
|
||||
try:
|
||||
with open(file_path, "r", encoding="latin-1") as f:
|
||||
content = f.read()
|
||||
self.app.notify(
|
||||
"File loaded with Latin-1 encoding (UTF-8 failed)",
|
||||
severity="warning",
|
||||
timeout=3,
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(
|
||||
f"Error reading file: {str(e)}", severity="error", timeout=5
|
||||
)
|
||||
return
|
||||
|
||||
# Validate and sanitize content
|
||||
lines = content.split("\n")
|
||||
valid_lines = []
|
||||
invalid_count = 0
|
||||
|
||||
# Pattern for valid hostnames/device names
|
||||
# Allows: letters, numbers, hyphens, underscores, periods, and Unicode chars
|
||||
hostname_pattern = re.compile(r"^[\w\-\.\u0080-\uFFFF]+$", re.UNICODE)
|
||||
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue # Skip empty lines
|
||||
|
||||
# Check if line looks like a valid hostname/device name
|
||||
if hostname_pattern.match(line):
|
||||
valid_lines.append(line)
|
||||
else:
|
||||
invalid_count += 1
|
||||
# Log but don't add invalid entries
|
||||
|
||||
if not valid_lines:
|
||||
self.app.notify(
|
||||
"No valid device names found in file", severity="warning", timeout=3
|
||||
)
|
||||
return
|
||||
|
||||
# Update text area with validated content
|
||||
text_area = self.query_one("#device_input", TextArea)
|
||||
text_area.text = "\n".join(valid_lines)
|
||||
|
||||
# Show notification
|
||||
msg = f"✅ Loaded {len(valid_lines)} devices from file"
|
||||
if invalid_count > 0:
|
||||
msg += f" ({invalid_count} invalid entries skipped)"
|
||||
|
||||
self.app.notify(msg, severity="information", timeout=5)
|
||||
|
||||
except ImportError:
|
||||
self.app.notify(
|
||||
"tkinter not available - cannot open file dialog",
|
||||
severity="error",
|
||||
timeout=3,
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(
|
||||
f"Error loading file: {str(e)}", severity="error", timeout=5
|
||||
)
|
||||
@@ -1,10 +1,17 @@
|
||||
"""
|
||||
Policy Selector Widget Module
|
||||
|
||||
Provides a Textual widget for selecting target policies for bulk agent operations.
|
||||
Allows users to browse available policies and select one as the destination for
|
||||
moving agents. Automatically excludes parent/logical policies.
|
||||
"""
|
||||
# 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/>
|
||||
|
||||
import logging
|
||||
import re
|
||||
@@ -14,7 +21,7 @@ import pandas as pd
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.message import Message
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea
|
||||
from textual.widgets import Button, DataTable, Static, TextArea
|
||||
|
||||
from models.policy import Policy
|
||||
|
||||
@@ -34,7 +41,7 @@ class PolicySelector(Widget):
|
||||
- Wildcard filtering (* and ?)
|
||||
- Interactive table for policy browsing
|
||||
- Explicit confirm button for selection
|
||||
- Cancel/back button to dismiss
|
||||
- Use escape key to go back
|
||||
|
||||
Attributes:
|
||||
policies (list[Policy]): List of available Policy objects to display.
|
||||
@@ -91,11 +98,10 @@ class PolicySelector(Widget):
|
||||
- Clear Filter button
|
||||
- Confirm Selection button
|
||||
- Policy table displaying available policies
|
||||
- Back buttons for navigation
|
||||
- Use escape key to go back
|
||||
"""
|
||||
yield Header(show_clock=True, icon="⚙")
|
||||
title_text = Static(
|
||||
"🎯 Select Target Policy",
|
||||
"Select Target Policy",
|
||||
id="policy_selector_title",
|
||||
)
|
||||
title_text.styles.margin = (0, 0, 1, 0)
|
||||
@@ -126,12 +132,12 @@ class PolicySelector(Widget):
|
||||
filter_help.styles.margin = (0, 0, 1, 0)
|
||||
yield filter_help
|
||||
|
||||
apply_button = Button("✓ Apply Filter", id="filter_button")
|
||||
apply_button = Button("🔍 Apply Filter", id="filter_button")
|
||||
apply_button.styles.width = "100%"
|
||||
apply_button.styles.margin = (0, 0, 1, 0)
|
||||
yield apply_button
|
||||
|
||||
clear_button = Button("Clear Filter", id="clear_filter_button")
|
||||
clear_button = Button("🧹 Clear Filter", id="clear_filter_button")
|
||||
clear_button.styles.width = "100%"
|
||||
clear_button.styles.margin = (0, 0, 1, 0)
|
||||
yield clear_button
|
||||
@@ -145,11 +151,6 @@ class PolicySelector(Widget):
|
||||
selected_label.styles.margin = (2, 0, 1, 0)
|
||||
yield selected_label
|
||||
|
||||
cancel_button = Button("← Back", id="back_button")
|
||||
cancel_button.styles.width = "100%"
|
||||
cancel_button.styles.margin = (1, 0, 1, 0)
|
||||
yield cancel_button
|
||||
|
||||
# Right side - Policy table
|
||||
with Vertical() as right_side:
|
||||
right_side.styles.width = "2fr"
|
||||
@@ -164,8 +165,6 @@ class PolicySelector(Widget):
|
||||
policy_table.styles.margin = (1, 0, 1, 0)
|
||||
yield policy_table
|
||||
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""
|
||||
Initialize the policy table when the widget is mounted.
|
||||
@@ -225,7 +224,6 @@ class PolicySelector(Widget):
|
||||
Handle button press events from the widget.
|
||||
|
||||
Routes to:
|
||||
- back_button (Cancel): Pop screen without selecting
|
||||
- filter_button (Apply Filter): Filter policies with wildcard support
|
||||
- clear_filter_button: Clear filter and show all policies
|
||||
- confirm_button: Confirm selection and post message
|
||||
@@ -235,12 +233,7 @@ class PolicySelector(Widget):
|
||||
"""
|
||||
btn_id = event.button.id
|
||||
|
||||
if btn_id == "back_button":
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "filter_button":
|
||||
if btn_id == "filter_button":
|
||||
self._apply_filter()
|
||||
event.stop()
|
||||
|
||||
@@ -286,7 +279,7 @@ class PolicySelector(Widget):
|
||||
if self.selected_policy:
|
||||
# Update selection display
|
||||
label = self.query_one("#selected_policy_label", Static)
|
||||
label.update(f"✓ Selected: {self.selected_policy.name}")
|
||||
label.update(f"Selected: {self.selected_policy.name}")
|
||||
|
||||
# Log for debugging
|
||||
logger.debug(
|
||||
@@ -328,7 +321,7 @@ class PolicySelector(Widget):
|
||||
|
||||
if highlighted_name:
|
||||
label = self.query_one("#selected_policy_label", Static)
|
||||
label.update(f"→ Highlighting: {highlighted_name}")
|
||||
label.update(f"Highlighting: {highlighted_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling row highlight: {e}")
|
||||
@@ -402,7 +395,9 @@ class PolicySelector(Widget):
|
||||
)
|
||||
|
||||
displayed_count = len(self._displayed_policies)
|
||||
status_text = f"📊 Showing {displayed_count} of {len(self._filtered_policies)} policies"
|
||||
status_text = (
|
||||
f"Showing {displayed_count} of {len(self._filtered_policies)} policies"
|
||||
)
|
||||
self.app.notify(status_text, severity="information", timeout=2)
|
||||
|
||||
# Clear selection when filter is applied
|
||||
@@ -477,7 +472,7 @@ class PolicySelector(Widget):
|
||||
"""
|
||||
if self.selected_policy is None:
|
||||
self.app.notify(
|
||||
"⚠️ Please select a policy first by clicking on a row in the table",
|
||||
"Please select a policy first by clicking on a row in the table",
|
||||
severity="warning",
|
||||
timeout=3,
|
||||
)
|
||||
@@ -486,6 +481,6 @@ class PolicySelector(Widget):
|
||||
# Log confirmation for debugging
|
||||
logger.info(f"Confirming selection of policy: {self.selected_policy.name}")
|
||||
self.app.notify(
|
||||
f"✅ Confirmed: {self.selected_policy.name}", severity="success", timeout=2
|
||||
f"Confirmed: {self.selected_policy.name}", severity="success", timeout=2
|
||||
)
|
||||
self.post_message(self.PolicySelected(self.selected_policy))
|
||||
@@ -1,10 +1,26 @@
|
||||
# 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 collections import defaultdict
|
||||
import logging
|
||||
|
||||
from rich.text import Text
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.message import Message
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Input, OptionList, Static, Tree
|
||||
from textual.widgets import Button, Input, OptionList, Static, Switch, Tree
|
||||
from textual.widgets.option_list import Option
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -13,35 +29,95 @@ logger = logging.getLogger(__name__)
|
||||
class PolicyTreeWidget(Widget):
|
||||
"""Widget for displaying and searching a hierarchical policy tree."""
|
||||
|
||||
class ViewExecutionHistory(Message):
|
||||
"""Message sent when user wants to view execution history for a device."""
|
||||
|
||||
def __init__(self, device):
|
||||
super().__init__()
|
||||
self.device = device
|
||||
|
||||
class GenerateOTP(Message):
|
||||
"""Message sent when user wants to generate OTP for a device."""
|
||||
|
||||
def __init__(self, device):
|
||||
super().__init__()
|
||||
self.device = device
|
||||
|
||||
class ToggleEnforcement(Message):
|
||||
"""Message sent when user wants to toggle audit/enforcement for a device."""
|
||||
|
||||
def __init__(self, device):
|
||||
super().__init__()
|
||||
self.device = device
|
||||
|
||||
def __init__(self, policies, devices):
|
||||
super().__init__()
|
||||
self.policies = policies
|
||||
self.devices = devices
|
||||
self.last_highlighted_node = None
|
||||
self.leaf_counts = defaultdict(int)
|
||||
self.match_type = "Count" # Default to sorting by count
|
||||
self.selected_device = None # Track currently selected device
|
||||
|
||||
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.styles.width = "2fr"
|
||||
policy_tree.styles.height = "100%"
|
||||
|
||||
# Create the search box and details pane
|
||||
label = Static("Device Search:")
|
||||
search_box = Input(
|
||||
placeholder="Search policies or devices...", id="tree_search"
|
||||
)
|
||||
exec_history_button = Button(
|
||||
"📊 Execution History", id="view_exec_history_button", disabled=True
|
||||
)
|
||||
exec_history_button.styles.margin = (0, 1, 0, 0) # Right margin
|
||||
|
||||
otp_button = Button("🎫 Generate OTP", id="generate_otp_button", disabled=True)
|
||||
otp_button.styles.margin = (0, 1, 0, 0) # Right margin
|
||||
|
||||
toggle_enforcement_button = Button(
|
||||
"🔄 Toggle Enforcement/Audit", id="toggle_enforcement_button", disabled=True
|
||||
)
|
||||
# No right margin on last button
|
||||
details_pane = Static("", id="details_pane")
|
||||
|
||||
# Layout the UI
|
||||
with Horizontal():
|
||||
yield policy_tree
|
||||
with Vertical() as right_pane:
|
||||
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 search_box
|
||||
# Action buttons in a horizontal row
|
||||
with Horizontal() as button_row:
|
||||
button_row.styles.height = "auto"
|
||||
yield exec_history_button
|
||||
yield otp_button
|
||||
yield toggle_enforcement_button
|
||||
yield details_pane
|
||||
|
||||
def on_mount(self) -> None:
|
||||
self._precompute_leaf_counts()
|
||||
|
||||
# Update root label with total leaf count
|
||||
total_leaves = sum(
|
||||
self.leaf_counts.get(policy.groupid, 0)
|
||||
@@ -50,8 +126,35 @@ class PolicyTreeWidget(Widget):
|
||||
)
|
||||
policy_tree = self.query_one("#policy_tree", Tree)
|
||||
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
|
||||
|
||||
self._build_tree()
|
||||
# Expand the root node
|
||||
policy_tree.root.expand()
|
||||
|
||||
def refresh_data(self, policies, devices):
|
||||
"""Refresh the widget with new data and rebuild the tree."""
|
||||
self.policies = policies
|
||||
self.devices = devices
|
||||
self.selected_device = None
|
||||
|
||||
# Disable all buttons since selection is lost
|
||||
try:
|
||||
self.query_one("#view_exec_history_button", Button).disabled = True
|
||||
self.query_one("#generate_otp_button", Button).disabled = True
|
||||
self.query_one("#toggle_enforcement_button", Button).disabled = True
|
||||
except:
|
||||
pass
|
||||
|
||||
# Rebuild tree with new data
|
||||
self._precompute_leaf_counts()
|
||||
total_leaves = sum(
|
||||
self.leaf_counts.get(policy.groupid, 0)
|
||||
for policy in self.policies
|
||||
if policy.parent == "global-policy-settings"
|
||||
)
|
||||
policy_tree = self.query_one("#policy_tree", Tree)
|
||||
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
|
||||
self._build_tree()
|
||||
policy_tree.root.expand()
|
||||
|
||||
def _precompute_leaf_counts(self):
|
||||
"""Precompute leaf counts for each policy group."""
|
||||
@@ -76,15 +179,21 @@ class PolicyTreeWidget(Widget):
|
||||
|
||||
def _build_tree(self):
|
||||
policy_tree = self.query_one("#policy_tree", Tree)
|
||||
policy_tree.clear() # Clear existing nodes
|
||||
node_map = {}
|
||||
|
||||
# Sort top-level policies
|
||||
top_policies = [
|
||||
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(
|
||||
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:
|
||||
label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})"
|
||||
@@ -98,9 +207,13 @@ class PolicyTreeWidget(Widget):
|
||||
children_by_parent[policy.parent].append(policy)
|
||||
|
||||
for parent_id, children in children_by_parent.items():
|
||||
if getattr(self, "match_type", "Count") == "Count":
|
||||
children.sort(
|
||||
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)
|
||||
if parent_node:
|
||||
for policy in children:
|
||||
@@ -108,11 +221,16 @@ class PolicyTreeWidget(Widget):
|
||||
node = parent_node.add(label=label, data=policy)
|
||||
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:
|
||||
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)
|
||||
if parent_node:
|
||||
for device in devices:
|
||||
parent_node.add(label=device.hostname, data=device)
|
||||
|
||||
def _collect_tree_nodes(self, node, all_nodes):
|
||||
@@ -133,6 +251,9 @@ class PolicyTreeWidget(Widget):
|
||||
node = message.node
|
||||
data = node.data
|
||||
details_pane = self.query_one("#details_pane", Static)
|
||||
exec_history_button = self.query_one("#view_exec_history_button", Button)
|
||||
otp_button = self.query_one("#generate_otp_button", Button)
|
||||
toggle_enforcement_button = self.query_one("#toggle_enforcement_button", Button)
|
||||
|
||||
if self.last_highlighted_node is not None:
|
||||
original_label = str(self.last_highlighted_node.label).strip()
|
||||
@@ -147,6 +268,20 @@ class PolicyTreeWidget(Widget):
|
||||
node.set_label(highlighted_label)
|
||||
self.last_highlighted_node = node
|
||||
|
||||
# Check if selected node is a device (has Agent data)
|
||||
from models.agent import Agent
|
||||
|
||||
if data and isinstance(data, Agent):
|
||||
self.selected_device = data
|
||||
exec_history_button.disabled = False
|
||||
otp_button.disabled = False
|
||||
toggle_enforcement_button.disabled = False
|
||||
else:
|
||||
self.selected_device = None
|
||||
exec_history_button.disabled = True
|
||||
otp_button.disabled = True
|
||||
toggle_enforcement_button.disabled = True
|
||||
|
||||
if data:
|
||||
details = "\n".join(
|
||||
f"{key}: {value}" for key, value in data.__dict__.items()
|
||||
@@ -157,6 +292,13 @@ class PolicyTreeWidget(Widget):
|
||||
|
||||
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:
|
||||
self._remove_match_selector()
|
||||
|
||||
@@ -229,3 +371,18 @@ class PolicyTreeWidget(Widget):
|
||||
option_list.remove()
|
||||
except:
|
||||
pass
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button presses."""
|
||||
if event.button.id == "view_exec_history_button":
|
||||
if self.selected_device:
|
||||
self.post_message(self.ViewExecutionHistory(self.selected_device))
|
||||
event.stop()
|
||||
elif event.button.id == "generate_otp_button":
|
||||
if self.selected_device:
|
||||
self.post_message(self.GenerateOTP(self.selected_device))
|
||||
event.stop()
|
||||
elif event.button.id == "toggle_enforcement_button":
|
||||
if self.selected_device:
|
||||
self.post_message(self.ToggleEnforcement(self.selected_device))
|
||||
event.stop()
|
||||
@@ -1,3 +1,18 @@
|
||||
# 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/>.
|
||||
|
||||
import logging
|
||||
|
||||
from textual.containers import Horizontal, Vertical
|
||||
@@ -59,15 +74,6 @@ class ResultsDisplay(Widget):
|
||||
margin-top: 1;
|
||||
width: 100%;
|
||||
}
|
||||
|
||||
#button_row {
|
||||
height: auto;
|
||||
margin: 1 0 0 0;
|
||||
}
|
||||
|
||||
#back_button {
|
||||
width: 1fr;
|
||||
}
|
||||
"""
|
||||
|
||||
class CopySuccess(Message):
|
||||
@@ -95,9 +101,9 @@ class ResultsDisplay(Widget):
|
||||
|
||||
def compose(self):
|
||||
with Vertical(id="results_screen"):
|
||||
yield Header(show_clock=True, icon="⚙")
|
||||
yield Header(show_clock=True, icon="⚙️")
|
||||
# Title
|
||||
title = Static(f"📊 {self.operation} - Results", id="results_title")
|
||||
title = Static(f"{self.operation} - Results", id="results_title")
|
||||
yield title
|
||||
|
||||
# Two-column layout
|
||||
@@ -107,7 +113,7 @@ class ResultsDisplay(Widget):
|
||||
yield Static("✅ Successful", id="success_label")
|
||||
yield Static(self.successful_results, id="success_results")
|
||||
yield Button(
|
||||
"📋✅ Copy Success List",
|
||||
"Copy Success List",
|
||||
id="copy_success",
|
||||
classes="copy_button",
|
||||
)
|
||||
@@ -117,15 +123,11 @@ class ResultsDisplay(Widget):
|
||||
yield Static("❌ Failed", id="failure_label")
|
||||
yield Static(self.unsuccessful_results, id="failure_results")
|
||||
yield Button(
|
||||
"📋❌ Copy Failure List",
|
||||
"Copy Failure List",
|
||||
id="copy_failure",
|
||||
classes="copy_button",
|
||||
)
|
||||
|
||||
# Back Button
|
||||
with Horizontal(id="button_row"):
|
||||
back_button = Button("← Back", id="back_button")
|
||||
yield back_button
|
||||
yield Footer()
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
@@ -138,18 +140,18 @@ class ResultsDisplay(Widget):
|
||||
|
||||
pyperclip.copy(str(success_widget.renderable))
|
||||
self.app.notify(
|
||||
"✅ Success list copied to clipboard!",
|
||||
"Success list copied to clipboard!",
|
||||
severity="information",
|
||||
timeout=2,
|
||||
)
|
||||
self.post_message(self.CopySuccess())
|
||||
except ImportError:
|
||||
self.app.notify(
|
||||
"âš ï¸ pyperclip not installed. Run: pip install pyperclip",
|
||||
"❌ pyperclip not installed. Run: pip install pyperclip",
|
||||
severity="warning",
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(f"⌠Failed to copy: {str(e)}", severity="error")
|
||||
self.app.notify(f"¢ Failed to copy: {str(e)}", severity="error")
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "copy_failure":
|
||||
@@ -159,20 +161,16 @@ class ResultsDisplay(Widget):
|
||||
|
||||
pyperclip.copy(str(failure_widget.renderable))
|
||||
self.app.notify(
|
||||
"✅ Failure list copied to clipboard!",
|
||||
"Failure list copied to clipboard!",
|
||||
severity="information",
|
||||
timeout=2,
|
||||
)
|
||||
self.post_message(self.CopyFailure())
|
||||
except ImportError:
|
||||
self.app.notify(
|
||||
"âš ï¸ pyperclip not installed. Run: pip install pyperclip",
|
||||
"❌ pyperclip not installed. Run: pip install pyperclip",
|
||||
severity="warning",
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(f"⌠Failed to copy: {str(e)}", severity="error")
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "back_button":
|
||||
self.app.pop_screen()
|
||||
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
|
||||
event.stop()
|
||||
@@ -0,0 +1,261 @@
|
||||
# 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/>.
|
||||
|
||||
import datetime
|
||||
import logging
|
||||
|
||||
from bson import ObjectId
|
||||
from textual.app import ComposeResult
|
||||
from textual.containers import Container, Vertical
|
||||
from textual.widgets import Button, DataTable, Input, Static
|
||||
|
||||
from services.API import AirlockAPIWrapper
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def skipback(days):
|
||||
"""
|
||||
Generate a MongoDB ObjectId for a given number of days ago from today.
|
||||
"""
|
||||
adjusted_days = days
|
||||
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
|
||||
days=adjusted_days
|
||||
)
|
||||
timestamp = int(date_days_ago.timestamp())
|
||||
hex_timestamp = format(timestamp, "08x")
|
||||
objectid_hex = hex_timestamp + "0000000000000000"
|
||||
return ObjectId(objectid_hex)
|
||||
|
||||
|
||||
class ServerLogWidget(Vertical):
|
||||
"""Widget for displaying server activity logs in a DataTable."""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
ServerLogWidget {
|
||||
width: 100%;
|
||||
height: 100%;
|
||||
}
|
||||
|
||||
ServerLogWidget #status_bar {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
background: $surface;
|
||||
padding: 1;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
ServerLogWidget DataTable {
|
||||
height: 1fr;
|
||||
border: solid $primary;
|
||||
}
|
||||
|
||||
ServerLogWidget #button_container {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
layout: horizontal;
|
||||
padding: 1;
|
||||
align: left middle;
|
||||
}
|
||||
|
||||
ServerLogWidget .filter_label {
|
||||
width: auto;
|
||||
height: 3;
|
||||
content-align: left middle;
|
||||
padding-right: 1;
|
||||
}
|
||||
|
||||
ServerLogWidget #filter_input {
|
||||
width: 40;
|
||||
margin-right: 1;
|
||||
}
|
||||
|
||||
ServerLogWidget Button {
|
||||
margin-right: 1;
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper):
|
||||
super().__init__()
|
||||
self.api = api
|
||||
self.all_logs = [] # Store all logs for filtering
|
||||
self.columns = [] # Store column names
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Static("Loading server logs (last 72 hours)...", id="status_bar")
|
||||
yield DataTable(id="server_log_table")
|
||||
with Container(id="button_container"):
|
||||
yield Static("Filter:", classes="filter_label")
|
||||
yield Input(placeholder="Filter (use * and ? wildcards)", id="filter_input")
|
||||
yield Button("Refresh", id="refresh_button", variant="primary")
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Initialize the DataTable and load server logs."""
|
||||
self.load_logs()
|
||||
|
||||
def load_logs(self) -> None:
|
||||
"""Load server logs from the API and populate the DataTable."""
|
||||
table = self.query_one("#server_log_table", DataTable)
|
||||
status = self.query_one("#status_bar", Static)
|
||||
|
||||
try:
|
||||
status.update("⏳ Loading server logs (last 72 hours)...")
|
||||
|
||||
# Create a fake checkpoint for 3 days ago (72 hours)
|
||||
checkpoint = str(skipback(3))
|
||||
|
||||
# Get server logs from API
|
||||
logs = self.api.server_logs(checkpoint=checkpoint)
|
||||
|
||||
if not logs:
|
||||
status.update("ℹï¸ No server logs found in the last 72 hours.")
|
||||
table.clear(columns=True)
|
||||
self.all_logs = []
|
||||
self.columns = []
|
||||
return
|
||||
|
||||
# Store all logs for filtering
|
||||
self.all_logs = logs
|
||||
|
||||
# Clear existing data
|
||||
table.clear(columns=True)
|
||||
|
||||
# Add columns based on the first log entry
|
||||
if logs:
|
||||
first_log = logs[0]
|
||||
self.columns = [col for col in first_log.keys() if col != "checkpoint"]
|
||||
|
||||
for col in self.columns:
|
||||
table.add_column(col, key=col)
|
||||
|
||||
# Add rows in reverse order so newest entries are at the top
|
||||
for log_entry in reversed(logs):
|
||||
row_data = []
|
||||
for col in self.columns:
|
||||
value = log_entry.get(col, "")
|
||||
# Format datetime column to be more readable
|
||||
if col == "datetime" and value:
|
||||
try:
|
||||
# Parse ISO format and convert to readable format
|
||||
dt = datetime.datetime.fromisoformat(
|
||||
str(value).replace("Z", "+00:00")
|
||||
)
|
||||
value = dt.strftime("%Y-%m-%d %H:%M:%S")
|
||||
except Exception:
|
||||
# If parsing fails, just use the original value
|
||||
pass
|
||||
row_data.append(str(value))
|
||||
table.add_row(*row_data)
|
||||
|
||||
status.update(
|
||||
f"✅ Loaded {len(logs)} log entries from the last 72 hours"
|
||||
)
|
||||
logger.info(f"Loaded {len(logs)} server log entries")
|
||||
else:
|
||||
status.update("ℹï¸ No log entries found.")
|
||||
self.all_logs = []
|
||||
self.columns = []
|
||||
|
||||
except Exception as exc:
|
||||
error_msg = f"❌ Error loading server logs: {exc}"
|
||||
status.update(error_msg)
|
||||
logger.error(f"Failed to load server logs: {exc}", exc_info=True)
|
||||
table.clear(columns=True)
|
||||
self.all_logs = []
|
||||
self.columns = []
|
||||
|
||||
def filter_logs(self, filter_text: str) -> None:
|
||||
"""Filter the logs based on the filter text with wildcard support."""
|
||||
import fnmatch
|
||||
|
||||
table = self.query_one("#server_log_table", DataTable)
|
||||
status = self.query_one("#status_bar", Static)
|
||||
|
||||
if not self.all_logs:
|
||||
return
|
||||
|
||||
# Clear existing data
|
||||
table.clear(columns=True)
|
||||
|
||||
# Re-add columns
|
||||
for col in self.columns:
|
||||
table.add_column(col, key=col)
|
||||
|
||||
# Filter logs
|
||||
filtered_logs = []
|
||||
if filter_text.strip():
|
||||
filter_pattern = filter_text.strip().lower()
|
||||
for log_entry in self.all_logs:
|
||||
# Check if any field matches the filter pattern
|
||||
match = False
|
||||
for col in self.columns:
|
||||
value = str(log_entry.get(col, "")).lower()
|
||||
if fnmatch.fnmatch(value, filter_pattern):
|
||||
match = True
|
||||
break
|
||||
if match:
|
||||
filtered_logs.append(log_entry)
|
||||
else:
|
||||
# No filter, show all logs
|
||||
filtered_logs = self.all_logs
|
||||
|
||||
# Add filtered rows in reverse order
|
||||
for log_entry in reversed(filtered_logs):
|
||||
row_data = []
|
||||
for col in self.columns:
|
||||
value = log_entry.get(col, "")
|
||||
# Format datetime column to be more readable
|
||||
if col == "datetime" and value:
|
||||
try:
|
||||
dt = datetime.datetime.fromisoformat(
|
||||
str(value).replace("Z", "+00:00")
|
||||
)
|
||||
value = dt.strftime("%Y-%m-%d %H:%M:%S")
|
||||
except Exception:
|
||||
pass
|
||||
row_data.append(str(value))
|
||||
table.add_row(*row_data)
|
||||
|
||||
if filter_text.strip():
|
||||
status.update(
|
||||
f"✅ Showing {len(filtered_logs)} of {len(self.all_logs)} log entries (filtered)"
|
||||
)
|
||||
else:
|
||||
status.update(
|
||||
f"✅ Loaded {len(self.all_logs)} log entries from the last 72 hours"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Filtered to {len(filtered_logs)} entries with pattern: {filter_text}"
|
||||
)
|
||||
|
||||
def on_input_changed(self, event: Input.Changed) -> None:
|
||||
"""Handle filter input changes."""
|
||||
if event.input.id == "filter_input":
|
||||
self.filter_logs(event.value)
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button presses."""
|
||||
button_id = event.button.id
|
||||
|
||||
if button_id == "refresh_button":
|
||||
self.load_logs()
|
||||
# Clear the filter input when refreshing
|
||||
try:
|
||||
filter_input = self.query_one("#filter_input", Input)
|
||||
filter_input.value = ""
|
||||
except Exception:
|
||||
pass
|
||||
event.stop()
|
||||
@@ -0,0 +1,567 @@
|
||||
# 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/>.
|
||||
|
||||
"""
|
||||
Settings widget combining theme selection and update checking.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import webbrowser
|
||||
|
||||
from textual.containers import Horizontal, Vertical, VerticalScroll
|
||||
from textual.message import Message
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Button, Rule, Static
|
||||
|
||||
from utils.versionchecker import (
|
||||
RELEASES_PAGE_URL,
|
||||
UpdateCheckResult,
|
||||
check_for_updates,
|
||||
get_current_version,
|
||||
get_version_checker,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SettingsWidget(Widget):
|
||||
"""Widget for application settings including themes and updates."""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
SettingsWidget {
|
||||
height: 1fr;
|
||||
}
|
||||
|
||||
/* Update section buttons - add margin between them */
|
||||
#update_buttons Button {
|
||||
margin-right: 1;
|
||||
}
|
||||
|
||||
/* Theme buttons - consistent width within columns, slightly smaller */
|
||||
.theme_btn {
|
||||
width: 100%;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
/* Column headers */
|
||||
.theme_column_header {
|
||||
text-align: center;
|
||||
text-style: bold;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
/* Section titles */
|
||||
.settings_section_title {
|
||||
text-style: bold;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
/* Theme columns - reduce overall width */
|
||||
#theme_columns {
|
||||
width: 80%;
|
||||
}
|
||||
|
||||
/* Theme columns spacing */
|
||||
#dark_themes_col1, #dark_themes_col2 {
|
||||
margin-right: 1;
|
||||
}
|
||||
|
||||
#light_themes_col {
|
||||
margin-left: 1;
|
||||
}
|
||||
"""
|
||||
|
||||
class ThemeSelected(Message):
|
||||
"""Message posted when a theme is selected."""
|
||||
|
||||
def __init__(self, theme_name: str):
|
||||
super().__init__()
|
||||
self.theme_name = theme_name
|
||||
|
||||
# Dark themes - Column 1
|
||||
DARK_THEMES_COL1 = [
|
||||
("Textual Dark", "textual-dark"),
|
||||
("Nord", "nord"),
|
||||
("Gruvbox", "gruvbox"),
|
||||
("Dracula", "dracula"),
|
||||
]
|
||||
|
||||
# Dark themes - Column 2
|
||||
DARK_THEMES_COL2 = [
|
||||
("Catppuccin Mocha", "catppuccin-mocha"),
|
||||
("Tokyo Night", "tokyo-night"),
|
||||
("Monokai", "monokai"),
|
||||
]
|
||||
|
||||
# Light themes (third column)
|
||||
LIGHT_THEMES = [
|
||||
("Textual Light", "textual-light"),
|
||||
("Flexoki", "flexoki"),
|
||||
("Catppuccin Latte", "catppuccin-latte"),
|
||||
("Solarized Light", "solarized-light"),
|
||||
]
|
||||
|
||||
# Combined for backward compatibility
|
||||
DARK_THEMES = DARK_THEMES_COL1 + DARK_THEMES_COL2
|
||||
AVAILABLE_THEMES = DARK_THEMES + LIGHT_THEMES
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._update_result: UpdateCheckResult | None = None
|
||||
self._checking = False
|
||||
|
||||
def compose(self):
|
||||
# Wrap everything in a scrollable container with auto height children
|
||||
with VerticalScroll(id="settings_scroll"):
|
||||
# Version & Updates Section
|
||||
with Vertical(id="updates_section") as updates:
|
||||
updates.styles.height = "auto"
|
||||
yield Static(
|
||||
"📦 Version & Updates",
|
||||
id="updates_title",
|
||||
classes="settings_section_title",
|
||||
)
|
||||
|
||||
version_text = f"Current Version: v{get_current_version()}"
|
||||
yield Static(version_text, id="current_version")
|
||||
|
||||
with Horizontal(id="update_buttons") as btn_row:
|
||||
btn_row.styles.height = "auto"
|
||||
yield Button("🔍 Check for Updates", id="check_updates_btn")
|
||||
yield Button("📥 View Releases", id="view_releases_btn")
|
||||
|
||||
yield Static("", id="update_status")
|
||||
|
||||
yield Rule()
|
||||
|
||||
# Theme Section - Three columns: Dark 1, Dark 2, Light
|
||||
with Vertical(id="themes_section") as themes:
|
||||
themes.styles.height = "auto"
|
||||
yield Static(
|
||||
"🎨 Theme Options",
|
||||
id="theme_title",
|
||||
classes="settings_section_title",
|
||||
)
|
||||
|
||||
with Horizontal(id="theme_columns") as cols:
|
||||
cols.styles.height = "auto"
|
||||
|
||||
# Dark themes section (2 columns under one header)
|
||||
with Vertical(id="dark_themes_section") as dark_section:
|
||||
dark_section.styles.width = "2fr"
|
||||
dark_section.styles.height = "auto"
|
||||
yield Static(
|
||||
"🌙 Dark Themes",
|
||||
classes="theme_column_header",
|
||||
id="dark_header",
|
||||
)
|
||||
|
||||
with Horizontal(id="dark_columns") as dark_cols:
|
||||
dark_cols.styles.height = "auto"
|
||||
|
||||
# Dark themes column 1
|
||||
with Vertical(id="dark_themes_col1") as dark_col1:
|
||||
dark_col1.styles.width = "1fr"
|
||||
dark_col1.styles.height = "auto"
|
||||
for label, btn_id in self.DARK_THEMES_COL1:
|
||||
yield Button(
|
||||
label,
|
||||
id=f"set_theme_{btn_id}",
|
||||
classes="theme_btn",
|
||||
)
|
||||
|
||||
# Dark themes column 2
|
||||
with Vertical(id="dark_themes_col2") as dark_col2:
|
||||
dark_col2.styles.width = "1fr"
|
||||
dark_col2.styles.height = "auto"
|
||||
for label, btn_id in self.DARK_THEMES_COL2:
|
||||
yield Button(
|
||||
label,
|
||||
id=f"set_theme_{btn_id}",
|
||||
classes="theme_btn",
|
||||
)
|
||||
|
||||
# Light themes column
|
||||
with Vertical(id="light_themes_col") as light_col:
|
||||
light_col.styles.width = "1fr"
|
||||
light_col.styles.height = "auto"
|
||||
yield Static("☀️ Light Themes", classes="theme_column_header")
|
||||
for label, btn_id in self.LIGHT_THEMES:
|
||||
yield Button(
|
||||
label, id=f"set_theme_{btn_id}", classes="theme_btn"
|
||||
)
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Check for cached update result on mount."""
|
||||
checker = get_version_checker()
|
||||
cached_result = checker.get_last_result()
|
||||
if cached_result and cached_result.update_available:
|
||||
self._update_result = cached_result
|
||||
self._show_update_available(cached_result)
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
button_id = event.button.id
|
||||
|
||||
if button_id == "check_updates_btn":
|
||||
self._check_for_updates()
|
||||
event.stop()
|
||||
elif button_id == "view_releases_btn":
|
||||
self._open_releases_page()
|
||||
event.stop()
|
||||
elif button_id == "download_update_btn":
|
||||
self._download_update()
|
||||
event.stop()
|
||||
elif button_id == "dismiss_update_btn":
|
||||
self._dismiss_update()
|
||||
event.stop()
|
||||
elif button_id and button_id.startswith("set_theme_"):
|
||||
theme_name = button_id.replace("set_theme_", "")
|
||||
self.post_message(self.ThemeSelected(theme_name))
|
||||
event.stop()
|
||||
|
||||
def _check_for_updates(self) -> None:
|
||||
"""Check for updates and update UI."""
|
||||
if self._checking:
|
||||
return
|
||||
|
||||
self._checking = True
|
||||
status = self.query_one("#update_status", Static)
|
||||
check_btn = self.query_one("#check_updates_btn", Button)
|
||||
|
||||
# Show checking status
|
||||
check_btn.disabled = True
|
||||
check_btn.label = "⏳ Checking..."
|
||||
status.update("🔄 Checking for updates...")
|
||||
|
||||
# Run check in worker to avoid blocking UI
|
||||
self.run_worker(self._do_update_check, exclusive=True)
|
||||
|
||||
async def _do_update_check(self) -> None:
|
||||
"""Worker to perform update check."""
|
||||
try:
|
||||
result = check_for_updates()
|
||||
self._update_result = result
|
||||
|
||||
# Since we're in an async worker (not a thread), we can call directly
|
||||
self._update_check_complete(result)
|
||||
except Exception as e:
|
||||
logger.error(f"Update check failed: {e}")
|
||||
self._update_check_failed(str(e))
|
||||
finally:
|
||||
self._checking = False
|
||||
|
||||
def _update_check_complete(self, result: UpdateCheckResult) -> None:
|
||||
"""Handle completed update check."""
|
||||
check_btn = self.query_one("#check_updates_btn", Button)
|
||||
check_btn.disabled = False
|
||||
check_btn.label = "🔍 Check for Updates"
|
||||
|
||||
if result.error:
|
||||
self._update_check_failed(result.error)
|
||||
return
|
||||
|
||||
if result.update_available:
|
||||
self._show_update_available(result)
|
||||
self.app.notify(
|
||||
f"🆕 Update available: {result.latest_version}",
|
||||
title="Update Available",
|
||||
severity="information",
|
||||
timeout=8,
|
||||
)
|
||||
else:
|
||||
status = self.query_one("#update_status", Static)
|
||||
status.update(f"✅ Loxide is up to date (v{result.current_version})")
|
||||
self.app.notify(
|
||||
"✅ Loxide is up to date!",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
def _update_check_failed(self, error: str) -> None:
|
||||
"""Handle failed update check."""
|
||||
check_btn = self.query_one("#check_updates_btn", Button)
|
||||
check_btn.disabled = False
|
||||
check_btn.label = "🔍 Check for Updates"
|
||||
|
||||
status = self.query_one("#update_status", Static)
|
||||
status.update(f"⚠️ Could not check for updates: {error}")
|
||||
|
||||
def _show_update_available(self, result: UpdateCheckResult) -> None:
|
||||
"""Show update available UI with release notes."""
|
||||
status = self.query_one("#update_status", Static)
|
||||
|
||||
msg = f"🆕 New version available: {result.latest_version}\n"
|
||||
msg += f" Current: v{result.current_version}"
|
||||
|
||||
if result.release_info and result.release_info.body:
|
||||
# Show release notes (truncate if very long)
|
||||
notes = result.release_info.body.strip()
|
||||
# Limit to ~500 chars to avoid overwhelming the UI
|
||||
if len(notes) > 500:
|
||||
notes = notes[:500] + "\n..."
|
||||
msg += f"\n\n📋 Release Notes:\n{notes}"
|
||||
|
||||
status.update(msg)
|
||||
|
||||
# Add download/dismiss buttons if not already there
|
||||
try:
|
||||
self.query_one("#download_update_btn")
|
||||
except Exception:
|
||||
# Buttons don't exist, add them
|
||||
button_container = self.query_one("#update_buttons", Horizontal)
|
||||
download_btn = Button(
|
||||
"📥 Download Update", id="download_update_btn", variant="success"
|
||||
)
|
||||
dismiss_btn = Button(
|
||||
"✖ Dismiss", id="dismiss_update_btn", variant="default"
|
||||
)
|
||||
button_container.mount(download_btn)
|
||||
button_container.mount(dismiss_btn)
|
||||
|
||||
def _open_releases_page(self) -> None:
|
||||
"""Open the releases page in browser."""
|
||||
try:
|
||||
webbrowser.open(RELEASES_PAGE_URL)
|
||||
self.app.notify("📂 Opened releases page in browser", timeout=3)
|
||||
except Exception as e:
|
||||
logger.error(f"Could not open browser: {e}")
|
||||
self.app.notify(f"⚠️ Could not open browser: {e}", severity="warning")
|
||||
|
||||
def _download_update(self) -> None:
|
||||
"""Download the update exe file."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
if not self._update_result or not self._update_result.release_info:
|
||||
self.app.notify("⚠️ No update information available", severity="warning")
|
||||
return
|
||||
|
||||
download_url = self._update_result.release_info.download_url
|
||||
if not download_url:
|
||||
# Fall back to opening the release page
|
||||
url = self._update_result.release_info.html_url
|
||||
try:
|
||||
webbrowser.open(url)
|
||||
self.app.notify(
|
||||
"📥 Opened download page in browser (no direct download available)",
|
||||
timeout=5,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Could not open browser: {e}")
|
||||
self.app.notify(f"⚠️ Could not open browser: {e}", severity="warning")
|
||||
return
|
||||
|
||||
# Determine destination path
|
||||
if os.name == "nt": # Windows
|
||||
downloads_dir = Path.home() / "Downloads"
|
||||
else:
|
||||
downloads_dir = Path.home() / "Downloads"
|
||||
if not downloads_dir.exists():
|
||||
downloads_dir = Path.home()
|
||||
|
||||
# Extract filename from URL
|
||||
filename = download_url.split("/")[-1]
|
||||
if not filename.endswith(".exe"):
|
||||
filename = f"Loxide_{self._update_result.latest_version}.exe"
|
||||
|
||||
dest_path = downloads_dir / filename
|
||||
|
||||
# Disable the button while downloading
|
||||
try:
|
||||
btn = self.query_one("#download_update_btn", Button)
|
||||
btn.disabled = True
|
||||
btn.label = "⏳ Downloading..."
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.app.notify(
|
||||
f"📥 Downloading to:\n{dest_path}", title="Download Starting", timeout=5
|
||||
)
|
||||
|
||||
# Small delay so user sees the "downloading to" toast before download completes
|
||||
self.set_timer(
|
||||
0.5,
|
||||
lambda: self.run_worker(
|
||||
self._do_download(download_url, dest_path), exclusive=True
|
||||
),
|
||||
)
|
||||
|
||||
async def _do_download(self, download_url: str, dest_path) -> None:
|
||||
"""Worker to download the update file."""
|
||||
try:
|
||||
# Download the file
|
||||
import requests
|
||||
|
||||
response = requests.get(download_url, stream=True, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
with open(dest_path, "wb") as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
|
||||
# Success
|
||||
self.app.notify(
|
||||
f"✅ Downloaded to:\n{dest_path}",
|
||||
title="Download Complete",
|
||||
severity="information",
|
||||
timeout=10,
|
||||
)
|
||||
logger.info(f"Update downloaded to {dest_path}")
|
||||
|
||||
# Re-enable button
|
||||
try:
|
||||
btn = self.query_one("#download_update_btn", Button)
|
||||
btn.disabled = False
|
||||
btn.label = "📥 Download Again"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Download failed: {e}")
|
||||
self.app.notify(f"❌ Download failed: {e}", severity="error", timeout=10)
|
||||
|
||||
# Re-enable button
|
||||
try:
|
||||
btn = self.query_one("#download_update_btn", Button)
|
||||
btn.disabled = False
|
||||
btn.label = "📥 Download Update"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _dismiss_update(self) -> None:
|
||||
"""Dismiss the current update notification."""
|
||||
if self._update_result and self._update_result.latest_version:
|
||||
checker = get_version_checker()
|
||||
checker.dismiss_update(self._update_result.latest_version)
|
||||
|
||||
# Remove the extra buttons
|
||||
try:
|
||||
self.query_one("#download_update_btn").remove()
|
||||
self.query_one("#dismiss_update_btn").remove()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
status = self.query_one("#update_status", Static)
|
||||
status.update(f"✓ Dismissed update {self._update_result.latest_version}")
|
||||
self._update_result = None
|
||||
|
||||
|
||||
# Keep ThemeSelector as a standalone for backward compatibility
|
||||
class ThemeSelector(Widget):
|
||||
"""Widget for selecting and applying Textual themes.
|
||||
|
||||
DEPRECATED: Use SettingsWidget instead for combined settings UI.
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
ThemeSelector {
|
||||
height: 1fr;
|
||||
}
|
||||
|
||||
/* Theme buttons - consistent width within columns */
|
||||
.theme_btn {
|
||||
width: 100%;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
/* Column headers */
|
||||
.theme_column_header {
|
||||
text-align: center;
|
||||
text-style: bold;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
/* Theme columns - reduce overall width */
|
||||
#theme_columns {
|
||||
width: 80%;
|
||||
}
|
||||
|
||||
/* Theme columns spacing */
|
||||
#dark_themes_col1, #dark_themes_col2 {
|
||||
margin-right: 1;
|
||||
}
|
||||
|
||||
#light_themes_col {
|
||||
margin-left: 1;
|
||||
}
|
||||
"""
|
||||
|
||||
class ThemeSelected(Message):
|
||||
"""Message posted when a theme is selected."""
|
||||
|
||||
def __init__(self, theme_name: str):
|
||||
super().__init__()
|
||||
self.theme_name = theme_name
|
||||
|
||||
DARK_THEMES_COL1 = SettingsWidget.DARK_THEMES_COL1
|
||||
DARK_THEMES_COL2 = SettingsWidget.DARK_THEMES_COL2
|
||||
DARK_THEMES = SettingsWidget.DARK_THEMES
|
||||
LIGHT_THEMES = SettingsWidget.LIGHT_THEMES
|
||||
AVAILABLE_THEMES = SettingsWidget.AVAILABLE_THEMES
|
||||
|
||||
def compose(self):
|
||||
with VerticalScroll(id="theme_scroll"):
|
||||
yield Static("Theme Options", id="theme_title")
|
||||
|
||||
with Horizontal(id="theme_columns") as cols:
|
||||
cols.styles.height = "auto"
|
||||
|
||||
# Dark themes section (2 columns under one header)
|
||||
with Vertical(id="dark_themes_section") as dark_section:
|
||||
dark_section.styles.width = "2fr"
|
||||
dark_section.styles.height = "auto"
|
||||
yield Static(
|
||||
"🌙 Dark Themes",
|
||||
classes="theme_column_header",
|
||||
id="dark_header",
|
||||
)
|
||||
|
||||
with Horizontal(id="dark_columns") as dark_cols:
|
||||
dark_cols.styles.height = "auto"
|
||||
|
||||
# Dark themes column 1
|
||||
with Vertical(id="dark_themes_col1") as dark_col1:
|
||||
dark_col1.styles.width = "1fr"
|
||||
dark_col1.styles.height = "auto"
|
||||
for label, btn_id in self.DARK_THEMES_COL1:
|
||||
yield Button(
|
||||
label, id=f"set_theme_{btn_id}", classes="theme_btn"
|
||||
)
|
||||
|
||||
# Dark themes column 2
|
||||
with Vertical(id="dark_themes_col2") as dark_col2:
|
||||
dark_col2.styles.width = "1fr"
|
||||
dark_col2.styles.height = "auto"
|
||||
for label, btn_id in self.DARK_THEMES_COL2:
|
||||
yield Button(
|
||||
label, id=f"set_theme_{btn_id}", classes="theme_btn"
|
||||
)
|
||||
|
||||
# Light themes column
|
||||
with Vertical(id="light_themes_col") as light_col:
|
||||
light_col.styles.width = "1fr"
|
||||
light_col.styles.height = "auto"
|
||||
yield Static("☀️ Light Themes", classes="theme_column_header")
|
||||
for label, btn_id in self.LIGHT_THEMES:
|
||||
yield Button(
|
||||
label, id=f"set_theme_{btn_id}", classes="theme_btn"
|
||||
)
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
button_id = event.button.id
|
||||
if button_id and button_id.startswith("set_theme_"):
|
||||
theme_name = button_id.replace("set_theme_", "")
|
||||
self.post_message(self.ThemeSelected(theme_name))
|
||||
@@ -1,3 +1,18 @@
|
||||
# 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 textual.containers import Vertical
|
||||
from textual.message import Message
|
||||
from textual.widget import Widget
|
||||
@@ -1,3 +1,4 @@
|
||||
/target
|
||||
build.sh
|
||||
pythontest.py
|
||||
changelog.md
|
||||
Generated
+1239
-336
File diff suppressed because it is too large
Load Diff
+16
-6
@@ -1,21 +1,31 @@
|
||||
[package]
|
||||
name = "airlock_libs"
|
||||
version = "2.0.0"
|
||||
version = "7.4.1"
|
||||
edition = "2024"
|
||||
|
||||
[lib]
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
chrono = "0.4.42"
|
||||
indicatif = "0.18.2"
|
||||
mongodb = "3.3.0"
|
||||
opentelemetry = { version = "0.27.0", features = ["logs", "metrics", "trace"] }
|
||||
opentelemetry-otlp = { version = "0.27.0", features = ["trace", "metrics", "grpc-tonic", "http-proto", "tls", "reqwest-client", "reqwest-rustls"] }
|
||||
opentelemetry-semantic-conventions = { version = "0.27.0" }
|
||||
opentelemetry-proto = { version = "0.27.0"}
|
||||
pyo3 = { version = "0.27.0", features = ["extension-module", "generate-import-lib"] }
|
||||
reqwest = { version = "0.12.24", features = ["json", "native-tls"] }
|
||||
reqwest = { version = "0.12.24", features = ["json", "native-tls", "rustls-tls"] }
|
||||
serde = "1.0.228"
|
||||
serde-pyobject = "0.8.0"
|
||||
serde_json = "1.0.145"
|
||||
tokio = { version = "1.48.0", features = ["full"] }
|
||||
tonic = { version = "0.12.3", features = ["tls-roots"] }
|
||||
tracing = "0.1.41"
|
||||
tracing-subscriber = "0.3.20"
|
||||
tracing-opentelemetry = "0.32.0"
|
||||
crossbeam = "0.8.4"
|
||||
log = "0.4.29"
|
||||
flexi_logger = "0.31.7"
|
||||
opentelemetry-appender-log = "0.27.0"
|
||||
opentelemetry_sdk = { version = "0.27.0", features = ["rt-tokio", "testing", "trace"] }
|
||||
|
||||
[package.metadata.maturin]
|
||||
generate-abi-stubs = true
|
||||
@@ -28,4 +38,4 @@ codegen-units = 1
|
||||
panic = 'abort'
|
||||
strip = true
|
||||
debug-assertions = false
|
||||
overflow-checks = false
|
||||
overflow-checks = true
|
||||
|
||||
@@ -1,40 +1,47 @@
|
||||
from typing import Dict, List
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
def pull_policy_exec_histories(
|
||||
self, type: List[str], checkpoint: str, policy: List[str]
|
||||
api,
|
||||
type: str,
|
||||
days: int,
|
||||
policy_name: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Retrieve execution history logs."""
|
||||
|
||||
def api(AirlockAPIWrapper):
|
||||
"""
|
||||
An implementation of the python AirlockAPIWrapper class to pass Python data into Rust
|
||||
Pull execution history for policies.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
base_url : str
|
||||
(Required) Base URL of the Airlock API, this should be in your .env file.
|
||||
api_key : str
|
||||
(Required) API Key for your profile in airlock, this should be in your credential manager.
|
||||
headers : {"X-APIKey": self.api_key}
|
||||
api : AirlockAPIWrapper
|
||||
The API wrapper instance
|
||||
policy_name : Optional[str]
|
||||
Name of the policy to query. If None, returns history for all policies.
|
||||
type : str
|
||||
JSON-style string list of execution types, e.g., "[1,2,3]"
|
||||
days : int
|
||||
Number of days to look back
|
||||
|
||||
```def __init__(self, base_url: str, api_key: str):
|
||||
self.base_url = base_ur.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.headers = {"X-APIKey": self.api_key}
|
||||
```
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
JSON string containing execution history response with structure:
|
||||
{"response": {"exechistories": [...]}}
|
||||
"""
|
||||
...
|
||||
|
||||
def history_logging(
|
||||
api,
|
||||
exec_types: str,
|
||||
checkpoint_number: str,
|
||||
policy_names: str,
|
||||
policy_names: Optional[str] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Query execution history logs from the Airlock API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
api : AirlockAPIWrapper
|
||||
The API wrapper instance
|
||||
|
||||
exec_types : str
|
||||
A JSON-style string list of execution types to retrieve.
|
||||
Example: "[3,5,8]"
|
||||
@@ -42,17 +49,30 @@ def history_logging(
|
||||
- 1 = Blocked Execution
|
||||
- 2 = Untrusted Execution [Audit]
|
||||
- 3 = Untrusted Execution [OTP]
|
||||
- 4 = Trusted Path Execution
|
||||
- 5 = Trusted Publisher Execution
|
||||
- 6 = Blocklist Execution
|
||||
- 7 = Blocklist Execution [Audit]
|
||||
- 8 = Trusted Process Execution
|
||||
(etc.)
|
||||
- 9 = Constrained Execution
|
||||
- 10 = Trusted Metadata Execution
|
||||
- 11 = Trusted Browser Execution
|
||||
- 12 = Blocked Browser Execution
|
||||
- 13 = Untrusted Browser Execution [Audit]
|
||||
- 14 = Untrusted Browser Execution [OTP]
|
||||
- 15 = Blocklist Browser Execution [Audit]
|
||||
- 16 = Blocklist Browser Execution
|
||||
- 17 = Trusted Installer Execution
|
||||
- 18 = Trusted Browser Metadata Execution
|
||||
|
||||
checkpoint_number : str
|
||||
The checkpoint ID. Used to fetch results after a certain event.
|
||||
Example: "601d275487bacb01e3470713"
|
||||
|
||||
policy_names : str
|
||||
policy_names : Optional[str], default None
|
||||
A comma-separated or JSON-style list of policy group names.
|
||||
Example: "Apple Mac" or "["Apple Mac", "Servers London"]"
|
||||
If None or not provided, retrieves execution history for ALL policies.
|
||||
Example: "Apple Mac" or '["Apple Mac", "Servers London"]'
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -76,14 +96,29 @@ def history_logging(
|
||||
- datetime: str
|
||||
- ip: str
|
||||
- localip: str
|
||||
|
||||
Raises
|
||||
------
|
||||
RuntimeError
|
||||
If the request fails or the response cannot be parsed.
|
||||
|
||||
Example
|
||||
-------
|
||||
>>> histories = await airlock_libs.history_logging("[3,5,8]", "601d275487bacb01e3470713", "Apple Mac")
|
||||
Examples
|
||||
--------
|
||||
>>> # Get execution history for ALL policies
|
||||
>>> histories = airlock_libs.history_logging(
|
||||
... api,
|
||||
... "[1,2,3]",
|
||||
... "601d275487bacb01e3470713",
|
||||
... None # or omit this parameter
|
||||
... )
|
||||
>>>
|
||||
>>> # Get execution history for a specific policy
|
||||
>>> histories = airlock_libs.history_logging(
|
||||
... api,
|
||||
... "[3,5,8]",
|
||||
... "601d275487bacb01e3470713",
|
||||
... "Apple Mac"
|
||||
... )
|
||||
>>> print(histories[0]["filename"])
|
||||
'chrome.exe'
|
||||
"""
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "airlock_libs"
|
||||
version = "2.0.0"
|
||||
version = "7.4.1"
|
||||
description = "Airlock Digital API Wrapper"
|
||||
readme = "README.md"
|
||||
license = { text = "AGPL-3.0-only" }
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use pyo3::prelude::*;
|
||||
mod services;
|
||||
pub mod modules;
|
||||
pub mod prelude;
|
||||
pub mod services;
|
||||
#[pymodule]
|
||||
fn airlock_libs(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
|
||||
m.add_function(wrap_pyfunction!(services::pull_policy_exec_histories, py)?)?;
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
use crate::prelude::*;
|
||||
use crate::services::get_base_directory;
|
||||
#[allow(non_snake_case)]
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub struct TelemetryConfig {
|
||||
pub TELEMETRY: bool,
|
||||
pub TELEM_URL: Option<String>,
|
||||
}
|
||||
|
||||
impl TelemetryConfig {
|
||||
pub fn init_tracer() -> opentelemetry_sdk::trace::TracerProvider {
|
||||
let cfg: TelemetryConfig = TelemetryConfig::load();
|
||||
if !cfg.TELEMETRY {
|
||||
return TracerProvider::builder().build();
|
||||
}
|
||||
let endpoint = cfg.TELEM_URL.unwrap_or_default();
|
||||
let channel = Channel::from_shared(endpoint.clone())
|
||||
.unwrap()
|
||||
.tls_config(ClientTlsConfig::new().with_native_roots())
|
||||
.unwrap()
|
||||
.connect_lazy();
|
||||
let exporter = opentelemetry_otlp::SpanExporter::builder()
|
||||
.with_tonic()
|
||||
.with_endpoint(endpoint.clone())
|
||||
.with_channel(channel)
|
||||
.build()
|
||||
.expect("Failed to build exporter");
|
||||
opentelemetry_sdk::trace::TracerProvider::builder()
|
||||
.with_simple_exporter(exporter)
|
||||
.with_resource(Resource::new(vec![KeyValue::new(
|
||||
"service.name",
|
||||
"LoxideLibs",
|
||||
)]))
|
||||
.build()
|
||||
}
|
||||
fn load() -> Self {
|
||||
let cfg_path = get_base_directory().join("config\\user_config.json");
|
||||
if !cfg_path.exists() {
|
||||
return Self {
|
||||
TELEMETRY: false,
|
||||
TELEM_URL: None,
|
||||
};
|
||||
}
|
||||
match fs::read_to_string(&cfg_path) {
|
||||
Ok(contents) => serde_json::from_str::<Self>(&contents).unwrap_or(Self {
|
||||
TELEMETRY: false,
|
||||
TELEM_URL: None,
|
||||
}),
|
||||
Err(_) => Self {
|
||||
TELEMETRY: false,
|
||||
TELEM_URL: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ApiResponse {
|
||||
pub(crate) error: String,
|
||||
pub(crate) response: ExecHistories,
|
||||
}
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ExecHistories {
|
||||
pub(crate) exechistories: Vec<Group>,
|
||||
}
|
||||
#[derive(Debug, Deserialize, Serialize, Clone)]
|
||||
pub struct Group {
|
||||
pub(crate) checkpoint: String,
|
||||
#[serde(rename = "type")]
|
||||
pub(crate) exectype: u8,
|
||||
pub(crate) username: String,
|
||||
pub(crate) hostname: String,
|
||||
pub(crate) netdomain: String,
|
||||
pub(crate) filename: String,
|
||||
pub(crate) ppolicy: String,
|
||||
pub(crate) policyname: String,
|
||||
pub(crate) policyver: String,
|
||||
pub(crate) commandline: String,
|
||||
pub(crate) publisher: String,
|
||||
pub(crate) pprocess: String,
|
||||
pub(crate) gprocess: String,
|
||||
pub(crate) sha256: String,
|
||||
pub(crate) datetime: String,
|
||||
pub(crate) md5: String,
|
||||
pub(crate) sha128: String,
|
||||
pub(crate) sha384: String,
|
||||
pub(crate) sha512: String,
|
||||
pub(crate) ip: String,
|
||||
pub(crate) localip: String,
|
||||
}
|
||||
|
||||
pub struct PyData {
|
||||
pub headers: reqwest::header::HeaderMap,
|
||||
pub base_url: String,
|
||||
}
|
||||
|
||||
impl PyData {
|
||||
pub fn extract_data(py: Python<'_>, obj: &Py<PyAny>) -> Self {
|
||||
let headers_raw = obj.getattr(py, "headers").unwrap().to_string();
|
||||
let headers_json = headers_raw.replace('\'', "\"");
|
||||
let parsed: Value = serde_json::from_str(&headers_json).unwrap();
|
||||
let mut header_map = HeaderMap::new();
|
||||
if let Some(obj) = parsed.as_object() {
|
||||
for (key, val) in obj {
|
||||
if let Some(v) = val.as_str() {
|
||||
let header_name = HeaderName::from_str(key).unwrap();
|
||||
let header_value: HeaderValue = HeaderValue::from_str(v).unwrap();
|
||||
header_map.insert(header_name, header_value);
|
||||
}
|
||||
}
|
||||
}
|
||||
let base_url = obj.getattr(py, "base_url").unwrap().to_string();
|
||||
Self {
|
||||
headers: header_map,
|
||||
base_url,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SkipBack;
|
||||
|
||||
impl SkipBack {
|
||||
pub fn find_checkpoint(days: i64) -> ObjectId {
|
||||
let date_days_ago = Local::now() - Duration::days(days);
|
||||
let timestamp = date_days_ago.timestamp() as u32;
|
||||
let mut hex_timestamp = String::new();
|
||||
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap();
|
||||
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
|
||||
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod datatypes;
|
||||
@@ -0,0 +1,32 @@
|
||||
pub use chrono::{Duration, Local, NaiveDate};
|
||||
pub use crossbeam::channel::unbounded;
|
||||
pub use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
|
||||
pub use mongodb::bson::oid::ObjectId;
|
||||
pub use opentelemetry::global::GlobalTracerProvider;
|
||||
pub use opentelemetry::trace::noop::NoopTracerProvider;
|
||||
pub use opentelemetry::trace::{Status, TraceContextExt, Tracer};
|
||||
pub use opentelemetry::*;
|
||||
pub use opentelemetry_otlp::ExportConfig;
|
||||
pub use opentelemetry_otlp::WithExportConfig;
|
||||
pub use opentelemetry_otlp::WithTonicConfig;
|
||||
pub use opentelemetry_sdk::Resource;
|
||||
pub use opentelemetry_sdk::trace::{Config, TracerProvider};
|
||||
pub use pyo3::{prelude::*, types::PyString};
|
||||
pub use reqwest::{
|
||||
Client,
|
||||
header::{HeaderMap, HeaderName, HeaderValue},
|
||||
};
|
||||
pub use serde::{Deserialize, Serialize};
|
||||
pub use serde_json::Value;
|
||||
pub use std::sync::{Arc, Mutex};
|
||||
pub use std::thread;
|
||||
pub use std::{
|
||||
collections::HashMap,
|
||||
env,
|
||||
fmt::Write,
|
||||
fs::{self, File},
|
||||
io::{Read, Seek, SeekFrom},
|
||||
path::PathBuf,
|
||||
str::FromStr,
|
||||
};
|
||||
pub use tonic::transport::{Channel, ClientTlsConfig};
|
||||
+259
-152
@@ -1,131 +1,131 @@
|
||||
use chrono::{Duration, Local, NaiveDate};
|
||||
use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
|
||||
use mongodb::bson::oid::ObjectId;
|
||||
use pyo3::{prelude::*, types::PyString};
|
||||
use reqwest::{
|
||||
Client,
|
||||
header::{HeaderMap, HeaderName, HeaderValue},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
env,
|
||||
fmt::Write,
|
||||
fs::{self, File},
|
||||
io::{Read, Seek, SeekFrom},
|
||||
path::PathBuf,
|
||||
str::FromStr,
|
||||
};
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct ApiResponse {
|
||||
error: String,
|
||||
response: ExecHistories,
|
||||
}
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct ExecHistories {
|
||||
exechistories: Vec<Group>,
|
||||
}
|
||||
#[derive(Debug, Deserialize, Serialize, Clone)]
|
||||
struct Group {
|
||||
checkpoint: String,
|
||||
#[serde(rename = "type")]
|
||||
exectype: u8,
|
||||
username: String,
|
||||
hostname: String,
|
||||
netdomain: String,
|
||||
filename: String,
|
||||
ppolicy: String,
|
||||
policyname: String,
|
||||
policyver: String,
|
||||
commandline: String,
|
||||
publisher: String,
|
||||
pprocess: String,
|
||||
gprocess: String,
|
||||
sha256: String,
|
||||
datetime: String,
|
||||
md5: String,
|
||||
sha128: String,
|
||||
sha384: String,
|
||||
sha512: String,
|
||||
ip: String,
|
||||
localip: String,
|
||||
}
|
||||
use crate::modules::datatypes::*;
|
||||
use crate::prelude::*;
|
||||
use opentelemetry::trace::SpanContext;
|
||||
|
||||
#[pyfunction]
|
||||
pub fn pull_policy_exec_histories(
|
||||
py: Python<'_>,
|
||||
py_self: Py<PyAny>,
|
||||
policy_names: String,
|
||||
policy_names: Option<String>,
|
||||
exec_types: String,
|
||||
days: i64,
|
||||
) -> Py<PyString> {
|
||||
println!();
|
||||
let data: PyData = PyData::extract_data(py, &py_self);
|
||||
let headers: HeaderMap = data.headers;
|
||||
let base_url: String = data.base_url;
|
||||
let handle: thread::JoinHandle<String> = std::thread::spawn(move || {
|
||||
let rt: tokio::runtime::Runtime = match tokio::runtime::Runtime::new() {
|
||||
Ok(rt) => rt,
|
||||
Err(e) => {
|
||||
println!("Failed to build Tokio Runtime: {:?}", e);
|
||||
std::process::abort();
|
||||
}
|
||||
};
|
||||
let tracer_provider = rt.block_on(async { TelemetryConfig::init_tracer() });
|
||||
global::set_tracer_provider(tracer_provider.clone());
|
||||
let tracer: global::BoxedTracer = global::tracer("tracer");
|
||||
let _cx: Context = Context::new();
|
||||
let file_path: PathBuf = format!(
|
||||
"{}\\cache\\chunkinator.json",
|
||||
get_base_directory().display()
|
||||
)
|
||||
.into();
|
||||
let writeable_filepath = file_path.clone();
|
||||
if !file_path.exists() {
|
||||
if let Some(parent_dir) = file_path.parent()
|
||||
if !&file_path.exists() {
|
||||
if let Some(parent_dir) = &file_path.parent()
|
||||
&& !parent_dir.exists()
|
||||
{
|
||||
fs::create_dir_all(parent_dir).unwrap();
|
||||
match fs::create_dir_all(parent_dir) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("Failed to Create Directory {:?}: {}", parent_dir, e);
|
||||
std::process::abort();
|
||||
}
|
||||
fs::File::create(file_path).unwrap();
|
||||
}
|
||||
let data = ApiResponse {
|
||||
}
|
||||
match fs::File::create(&file_path) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("Failed to Create Directory {:?}: {}", &file_path, e);
|
||||
std::process::abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
let data: ApiResponse = ApiResponse {
|
||||
error: "Success".to_string(),
|
||||
response: ExecHistories {
|
||||
exechistories: vec![],
|
||||
},
|
||||
};
|
||||
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
|
||||
fs::write(writeable_filepath.clone(), data_write).unwrap();
|
||||
let mut checkpoint_number: String = skipback(days).to_string();
|
||||
let multi_progress = MultiProgress::new();
|
||||
multi_progress.set_draw_target(ProgressDrawTarget::stdout());
|
||||
let progress_bar = multi_progress.add(ProgressBar::new(100));
|
||||
progress_bar.set_style(
|
||||
ProgressStyle::default_bar()
|
||||
.template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
|
||||
.unwrap(),
|
||||
);
|
||||
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
|
||||
let client = build_client(py, &py_self);
|
||||
let api: Py<PyAny> = py_self;
|
||||
let cutoff = Local::now().naive_local() - Duration::days(days);
|
||||
let mut f = File::open(&writeable_filepath).unwrap();
|
||||
loop {
|
||||
f.seek(SeekFrom::Start(0)).unwrap();
|
||||
let execution_histories = history_logging(
|
||||
py,
|
||||
&api,
|
||||
&exec_types,
|
||||
&checkpoint_number,
|
||||
&policy_names,
|
||||
&client,
|
||||
);
|
||||
let parsed_responses = execution_histories.response.exechistories;
|
||||
if parsed_responses.is_empty() {
|
||||
break;
|
||||
let writeable_filepath: PathBuf = file_path.clone();
|
||||
let data_write: String = serde_json::to_string_pretty(&data).expect("Failed to serialize");
|
||||
match fs::write(writeable_filepath.clone(), data_write) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
|
||||
std::process::abort();
|
||||
}
|
||||
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() {
|
||||
let mut contents = String::new();
|
||||
f.read_to_string(&mut contents).unwrap();
|
||||
let existing_data: ApiResponse =
|
||||
}
|
||||
let mut checkpoint_number: String = SkipBack::find_checkpoint(days).to_string();
|
||||
let progress_bar = Arc::new(Mutex::new(ProgressBar::new(100)));
|
||||
progress_bar
|
||||
.lock()
|
||||
.unwrap()
|
||||
.set_draw_target(ProgressDrawTarget::stderr());
|
||||
progress_bar.lock().unwrap().set_style(
|
||||
ProgressStyle::default_bar()
|
||||
.template("Total - Policy Name: {msg}: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
|
||||
.unwrap().progress_chars("⣿⣦⣀")
|
||||
);
|
||||
let client: Client = tracer.in_span("Building HTTP Client", |cx| {
|
||||
let client_result: Result<Client, reqwest::Error> = build_client(headers);
|
||||
match client_result {
|
||||
Ok(client_result) => {
|
||||
cx.span().add_event(
|
||||
"info",
|
||||
vec![KeyValue::new(
|
||||
"Client Built Successfully",
|
||||
format!("{:?}", client_result),
|
||||
)],
|
||||
);
|
||||
client_result
|
||||
}
|
||||
Err(client_result) => {
|
||||
cx.span().add_event(
|
||||
"warn",
|
||||
vec![KeyValue::new(
|
||||
"Client Failed to Build",
|
||||
format!("{:?}", &client_result),
|
||||
)],
|
||||
);
|
||||
cx.span()
|
||||
.set_status(Status::error("Client Failed to Build"));
|
||||
println!("Failed to Build Client: {:?}", client_result);
|
||||
std::process::abort();
|
||||
}
|
||||
}
|
||||
});
|
||||
let cutoff: chrono::NaiveDateTime =
|
||||
Local::now().naive_local() - chrono::Duration::days(days);
|
||||
let (tx, rx) = unbounded::<(SpanContext, Vec<Group>)>();
|
||||
let pb_clone = progress_bar.clone();
|
||||
thread::spawn(move || {
|
||||
let tracer = global::tracer("loxide");
|
||||
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
|
||||
{
|
||||
let contents: String = fs::read_to_string(&writeable_filepath).unwrap_or_default();
|
||||
let existing: ApiResponse =
|
||||
serde_json::from_str(&contents).unwrap_or(ApiResponse {
|
||||
error: "Success".to_string(),
|
||||
response: ExecHistories {
|
||||
exechistories: vec![],
|
||||
},
|
||||
});
|
||||
existing_data
|
||||
existing
|
||||
.response
|
||||
.exechistories
|
||||
.into_iter()
|
||||
.map(|entry| {
|
||||
.map(|entry: Group| {
|
||||
(
|
||||
(
|
||||
entry.sha256.clone(),
|
||||
@@ -139,15 +139,27 @@ pub fn pull_policy_exec_histories(
|
||||
} else {
|
||||
HashMap::new()
|
||||
};
|
||||
for (index, executions) in parsed_responses.iter().enumerate() {
|
||||
while let Ok((parent_spancontext, parsed_responses)) = rx.recv() {
|
||||
let parent_ctx = Context::new().with_remote_span_context(parent_spancontext);
|
||||
let span = tracer.build_with_context(
|
||||
tracer
|
||||
.span_builder("Deduplicate and Write")
|
||||
.with_kind(trace::SpanKind::Consumer),
|
||||
&parent_ctx,
|
||||
);
|
||||
let cx = Context::current_with_span(span);
|
||||
cx.span().add_event(
|
||||
"Received Data from Producer",
|
||||
vec![KeyValue::new(
|
||||
"Items to Process",
|
||||
parsed_responses.len().to_string(),
|
||||
)],
|
||||
);
|
||||
for executions in parsed_responses {
|
||||
if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if index == parsed_responses.len() - 1 {
|
||||
checkpoint_number = executions.checkpoint.clone();
|
||||
break;
|
||||
}
|
||||
let history_date = match NaiveDate::parse_from_str(
|
||||
let history_date: NaiveDate = match NaiveDate::parse_from_str(
|
||||
&executions.datetime.replace(" +0000 UTC", ""),
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
) {
|
||||
@@ -155,7 +167,7 @@ pub fn pull_policy_exec_histories(
|
||||
Err(_) => continue,
|
||||
};
|
||||
if history_date >= cutoff.into() {
|
||||
let key = (
|
||||
let key: (String, String, String) = (
|
||||
executions.sha256.clone(),
|
||||
executions.filename.clone(),
|
||||
executions.hostname.clone(),
|
||||
@@ -163,72 +175,172 @@ pub fn pull_policy_exec_histories(
|
||||
seen.entry(key).or_insert(executions.clone());
|
||||
}
|
||||
}
|
||||
let final_response = ApiResponse {
|
||||
let final_response: ApiResponse = ApiResponse {
|
||||
error: "Success".to_string(),
|
||||
response: ExecHistories {
|
||||
exechistories: seen.values().cloned().collect(),
|
||||
},
|
||||
};
|
||||
let data_write = serde_json::to_string_pretty(&final_response).unwrap();
|
||||
fs::write(&writeable_filepath, data_write).unwrap();
|
||||
if let Some(last_item) = &final_response.response.exechistories.last()
|
||||
let data_write: String = serde_json::to_string_pretty(&final_response).unwrap();
|
||||
match fs::write(&writeable_filepath, data_write) {
|
||||
Ok(_) => {
|
||||
cx.span().add_event(
|
||||
"Writing Data to File",
|
||||
vec![KeyValue::new("Success", "Ok".to_string())],
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
cx.span().add_event(
|
||||
"Writing Data to File",
|
||||
vec![KeyValue::new("Failed", e.to_string())],
|
||||
);
|
||||
cx.span()
|
||||
.set_status(Status::error("Failed to Write to File"));
|
||||
}
|
||||
}
|
||||
cx.span().add_event(
|
||||
"Finished Deduplicating Data",
|
||||
vec![KeyValue::new(
|
||||
"Items Successfully Processed",
|
||||
seen.len().to_string(),
|
||||
)],
|
||||
);
|
||||
}
|
||||
});
|
||||
let mut first_date: Option<NaiveDate> = None;
|
||||
tracer.in_span("Airlock Data Retreival", |cx| {
|
||||
pb_clone
|
||||
.lock()
|
||||
.unwrap()
|
||||
.enable_steady_tick(std::time::Duration::from_millis(100));
|
||||
pb_clone
|
||||
.lock()
|
||||
.unwrap()
|
||||
.set_message(policy_names.clone().unwrap_or("Statistics".to_string()));
|
||||
let span: opentelemetry::trace::SpanRef<'_> = cx.span();
|
||||
span.set_attribute(KeyValue::new("Days", days.to_string()));
|
||||
span.set_attribute(KeyValue::new(
|
||||
"Policy Name",
|
||||
policy_names
|
||||
.clone()
|
||||
.unwrap_or("Statistics Monitoring".to_string()),
|
||||
));
|
||||
loop {
|
||||
let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
|
||||
cx.span().add_event(
|
||||
"Retrieving Responses from API",
|
||||
vec![KeyValue::new(
|
||||
"Checkpoint Number",
|
||||
checkpoint_number.to_string(),
|
||||
)],
|
||||
);
|
||||
let results: ApiResponse = rt.block_on(history_logging(
|
||||
&base_url,
|
||||
&exec_types,
|
||||
&checkpoint_number,
|
||||
&policy_names,
|
||||
&client,
|
||||
));
|
||||
cx.span().add_event(
|
||||
"Got Responses from API",
|
||||
vec![KeyValue::new(
|
||||
"Items in Response",
|
||||
results.response.exechistories.len().to_string(),
|
||||
)],
|
||||
);
|
||||
cx.span().set_status(Status::Ok);
|
||||
cx.span().set_attribute(KeyValue::new(
|
||||
"items_in_response",
|
||||
results.response.exechistories.len().to_string(),
|
||||
));
|
||||
results
|
||||
});
|
||||
let parsed_responses: Vec<Group> = execution_histories.response.exechistories;
|
||||
if parsed_responses.is_empty() {
|
||||
break;
|
||||
}
|
||||
match tx.send((cx.span().span_context().clone(), parsed_responses.clone())) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
cx.span().add_event(
|
||||
"Failed to Send Items to Processor",
|
||||
vec![KeyValue::new("Response from Processor", e.to_string())],
|
||||
);
|
||||
cx.span().set_status(Status::error("Processor Failed"))
|
||||
}
|
||||
}
|
||||
checkpoint_number = parsed_responses.last().unwrap().checkpoint.clone();
|
||||
if let Some(last_item) = parsed_responses.last()
|
||||
&& let Ok(last_date) = NaiveDate::parse_from_str(
|
||||
&last_item.datetime.replace(" +0000 UTC", ""),
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
)
|
||||
{
|
||||
let date_diff = Local::now().naive_local().date() - last_date;
|
||||
let percentage_diff = (days - date_diff.num_days()) as f64 / days as f64 * 100.0;
|
||||
progress_bar.set_position(percentage_diff.round() as u64);
|
||||
progress_bar.set_message("Total Percent Complete");
|
||||
if first_date.is_none() {
|
||||
first_date = Some(last_date);
|
||||
}
|
||||
if let Some(base_date) = first_date {
|
||||
let date_diff: chrono::TimeDelta = last_date - base_date;
|
||||
let total_span: i64 =
|
||||
(Local::now().naive_local().date() - base_date).num_days();
|
||||
let percentage: u64 = ((date_diff.num_days() as f64 / total_span as f64)
|
||||
* 100.0)
|
||||
.clamp(0.0, 100.0)
|
||||
.round() as u64;
|
||||
pb_clone.lock().unwrap().set_position(percentage);
|
||||
}
|
||||
}
|
||||
progress_bar.finish_with_message("All Checkpoints Complete");
|
||||
let return_data = fs::read_to_string(&writeable_filepath).unwrap();
|
||||
PyString::new(py, &return_data).into()
|
||||
}
|
||||
});
|
||||
progress_bar
|
||||
.lock()
|
||||
.unwrap()
|
||||
.finish_with_message("All Checkpoints Complete");
|
||||
let return_data: String = match fs::read_to_string(file_path.clone()) {
|
||||
Ok(return_data) => return_data,
|
||||
Err(e) => {
|
||||
println!("Failed to read data from: {:?}: {}", &file_path, e);
|
||||
std::process::abort();
|
||||
}
|
||||
};
|
||||
tracer_provider
|
||||
.shutdown()
|
||||
.expect("Failed to Shutdown Tracer Provdier");
|
||||
drop(tx);
|
||||
return_data.to_string()
|
||||
});
|
||||
let gil_value: String = handle.join().unwrap();
|
||||
Python::attach(|py: Python<'_>| PyString::new(py, &gil_value).into())
|
||||
}
|
||||
|
||||
fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Client {
|
||||
let headers = py_self.getattr(py, "headers").unwrap().to_string();
|
||||
let headers_replace = headers.replace('\'', "\"");
|
||||
let parsed: Value = serde_json::from_str(headers_replace.as_str()).unwrap();
|
||||
let mut header_map = HeaderMap::new();
|
||||
if let Some(obj) = parsed.as_object() {
|
||||
for (_key, value) in obj {
|
||||
if let Some(v) = value.as_str() {
|
||||
let val = HeaderValue::from_str(v).unwrap();
|
||||
header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_client(headers: HeaderMap) -> Result<reqwest::Client, reqwest::Error> {
|
||||
Client::builder()
|
||||
.danger_accept_invalid_certs(true)
|
||||
.default_headers(header_map)
|
||||
.timeout(std::time::Duration::from_secs(300))
|
||||
.default_headers(headers)
|
||||
.build()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
#[tracing::instrument(name = "history_logging")]
|
||||
async fn history_logging(
|
||||
py: Python<'_>,
|
||||
py_self: &Py<PyAny>,
|
||||
base_url: &String,
|
||||
exec_types: &String,
|
||||
checkpoint_number: &String,
|
||||
policy_names: &String,
|
||||
policy_names: &Option<String>,
|
||||
client: &Client,
|
||||
) -> ApiResponse {
|
||||
let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
|
||||
let policy_json = match policy_names {
|
||||
Some(name) => format!(r#"[ "{}" ]"#, name), // JSON array with one element
|
||||
None => "[]".to_string(), // Empty JSON array
|
||||
};
|
||||
let payload = format!(
|
||||
r#"{{
|
||||
"type": {},
|
||||
"checkpoint": "{}",
|
||||
"policy": ["{}"]
|
||||
"policy": {}
|
||||
}}"#,
|
||||
exec_types, checkpoint_number, policy_names
|
||||
exec_types, checkpoint_number, policy_json
|
||||
);
|
||||
let res = client
|
||||
let res: Result<reqwest::Response, reqwest::Error> = client
|
||||
.post(format!("{}/v1/logging/exechistories", base_url))
|
||||
.body(payload)
|
||||
.send()
|
||||
@@ -237,7 +349,7 @@ async fn history_logging(
|
||||
Ok(res) => {
|
||||
let first_response: ApiResponse = serde_json::from_str(&res.text().await.unwrap())
|
||||
.expect("Failed to retrieve response from API");
|
||||
return first_response;
|
||||
first_response
|
||||
}
|
||||
Err(_res) => {
|
||||
let failed_response: ApiResponse = ApiResponse {
|
||||
@@ -246,12 +358,12 @@ async fn history_logging(
|
||||
exechistories: vec![],
|
||||
},
|
||||
};
|
||||
return failed_response;
|
||||
failed_response
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_base_directory() -> PathBuf {
|
||||
pub fn get_base_directory() -> PathBuf {
|
||||
let home = env::var_os("HOME")
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| env::var_os("USERPROFILE").map(PathBuf::from))
|
||||
@@ -264,15 +376,10 @@ fn get_base_directory() -> PathBuf {
|
||||
.unwrap_or_else(|| home.join("AppData").join("Roaming"));
|
||||
appdata.join("Loxide")
|
||||
}
|
||||
_ => home.join(".local").join("share").join("Loxide"),
|
||||
"linux" => home.join(".local").join("share").join("Loxide"),
|
||||
_ => {
|
||||
println!("{} is currently not compatible with LoxideLibs", os);
|
||||
std::process::abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn skipback(days: i64) -> ObjectId {
|
||||
let date_days_ago = Local::now() - Duration::days(days);
|
||||
let timestamp = date_days_ago.timestamp() as u32;
|
||||
let mut hex_timestamp = String::new();
|
||||
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap();
|
||||
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
|
||||
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
|
||||
}
|
||||
|
||||
@@ -1,14 +1,39 @@
|
||||
{
|
||||
"APPNAME": "AirlockTools",
|
||||
"APPNAME": "Loxide",
|
||||
"URL": "https://server:3129",
|
||||
"LOG_LEVEL": "INFO",
|
||||
"BAD_PATH_PARTS": ["users","wwwroot","windows\\temp","windows\\task","windows\\system32","startup", "windows\\fonts","Recycle.Bin","AppData","programdata", "Solarwinds","kaseya"],
|
||||
"BAD_PUBLISHERS": ["Brave", "Zoom", "GlavSoft", "VNC"],
|
||||
"PUPS":["logmein","invalid","nmap","LTSvc","VNC","Kaseya","Solarwinds","mRemoteNG"],
|
||||
"BAD_PATH_PARTS": [
|
||||
"users",
|
||||
"wwwroot",
|
||||
"windows\\temp",
|
||||
"windows\\task",
|
||||
"windows\\system32",
|
||||
"startup",
|
||||
"windows\\fonts",
|
||||
"Recycle.Bin",
|
||||
"AppData",
|
||||
"programdata",
|
||||
"Solarwinds",
|
||||
"kaseya"
|
||||
],
|
||||
"BAD_PUBLISHERS": [
|
||||
"Brave",
|
||||
"Zoom",
|
||||
"GlavSoft",
|
||||
"VNC"
|
||||
],
|
||||
"PUPS": [
|
||||
"logmein",
|
||||
"invalid",
|
||||
"nmap",
|
||||
"LTSvc",
|
||||
"VNC",
|
||||
"Kaseya",
|
||||
"Solarwinds",
|
||||
"mRemoteNG"
|
||||
],
|
||||
"PATH_EXCLUSION_CONST": 4,
|
||||
"MIN_FILES_FOR_PATH": 4,
|
||||
"VT_THREAT_TOLERANCE": 4,
|
||||
"POLICY_MAP_ENF_AUD": {
|
||||
|
||||
}
|
||||
"POLICY_MAP_ENF_AUD": {}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,476 @@
|
||||
# Loxide
|
||||
## API Reference Document
|
||||
|
||||
**Version 1.0 | December 2025**
|
||||
|
||||
---
|
||||
|
||||
## 1. Overview
|
||||
|
||||
This document provides a complete reference for the AirlockAPIWrapper class and all Airlock API endpoints used by Loxide.
|
||||
|
||||
### 1.1 Authentication
|
||||
|
||||
All API requests require an API key passed via the X-APIKey header:
|
||||
|
||||
```
|
||||
Headers: { "X-APIKey": "your-api-key-here" }
|
||||
```
|
||||
|
||||
### 1.2 Base URL
|
||||
|
||||
The base URL is configured in `system_config.json` and typically follows the pattern:
|
||||
|
||||
```
|
||||
https://airlock.example.com/api
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Agent Management
|
||||
|
||||
### 2.1 agent_find_all()
|
||||
|
||||
Retrieve all registered agents.
|
||||
|
||||
**Endpoint:** `POST /v1/agent/find`
|
||||
**Payload:** `{}`
|
||||
**Returns:** DataFrame with agent records
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `hostname` | str | Device hostname |
|
||||
| `agentid` | str | Unique identifier |
|
||||
| `clientversion` | str | Agent version |
|
||||
| `groupid` | str | Policy group ID |
|
||||
| `status` | int | 0=Offline, 1=Online, 2=Hidden, 3=Safemode |
|
||||
| `lastcheckin` | str | Last check-in timestamp |
|
||||
| `ip` | str | External IP address |
|
||||
| `localip` | str | Internal IP address |
|
||||
| `domain` | str | Network domain |
|
||||
| `os` | str | Operating system |
|
||||
| `username` | str | Logged-in user |
|
||||
| `freespace` | int | Available disk space |
|
||||
| `policyversion` | str | Active policy version |
|
||||
|
||||
### 2.2 agent_find_by_hostname(hostname: str)
|
||||
|
||||
Find agents matching a hostname pattern.
|
||||
|
||||
**Endpoint:** `POST /v1/agent/find`
|
||||
**Payload:** `{ "hostname": "<pattern>" }`
|
||||
**Returns:** DataFrame with matching agents
|
||||
|
||||
### 2.3 agent_find_by_id(agentid: str)
|
||||
|
||||
Find agent by unique ID.
|
||||
|
||||
**Endpoint:** `POST /v1/agent/find`
|
||||
**Payload:** `{ "agentid": "<id>" }`
|
||||
**Returns:** DataFrame with agent record
|
||||
|
||||
### 2.4 agent_find_by_status(status: int)
|
||||
|
||||
Find agents by status code.
|
||||
|
||||
**Endpoint:** `POST /v1/agent/find`
|
||||
**Payload:** `{ "status": <code> }`
|
||||
**Status Codes:**
|
||||
- `0` - Offline
|
||||
- `1` - Online
|
||||
- `2` - Hidden
|
||||
- `3` - Safemode
|
||||
|
||||
### 2.5 agent_move(agentid: str, groupid: str)
|
||||
|
||||
Move an agent to a different policy group.
|
||||
|
||||
**Endpoint:** `POST /v1/agent/move`
|
||||
**Payload:** `{ "agentid": "<id>", "groupid": "<target_group>" }`
|
||||
**Returns:** dict with operation result
|
||||
|
||||
### 2.6 agents_find_by_group(groupid: str)
|
||||
|
||||
Find all agents in a policy group.
|
||||
|
||||
**Endpoint:** `POST /v1/agent/find`
|
||||
**Payload:** `{ "groupid": "<id>" }`
|
||||
**Returns:** DataFrame with agents
|
||||
|
||||
---
|
||||
|
||||
## 3. Policy Management
|
||||
|
||||
### 3.1 policy_find_all()
|
||||
|
||||
Retrieve all policy groups.
|
||||
|
||||
**Endpoint:** `POST /v1/group`
|
||||
**Payload:** `{}`
|
||||
**Returns:** DataFrame with policy records
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `name` | str | Policy display name |
|
||||
| `groupid` | int | Unique group identifier |
|
||||
| `hidden` | bool | Visibility flag |
|
||||
| `parent` | str | Parent policy name (if child) |
|
||||
|
||||
### 3.2 policy_set_auditmode(groupid: str, auditmode: str)
|
||||
|
||||
Toggle policy between audit and enforcement modes.
|
||||
|
||||
**Endpoint:** `POST /v1/group/settings/auditmode`
|
||||
**Payload:** `{ "groupid": "<id>", "auditmode": "<mode>" }`
|
||||
**Mode Values:**
|
||||
- `"1"` - Audit mode (log only)
|
||||
- `"0"` - Enforcement mode (block)
|
||||
|
||||
### 3.3 policy_list_agents(groupid: str)
|
||||
|
||||
List all agents assigned to a policy group.
|
||||
|
||||
**Endpoint:** `POST /v1/group/agents`
|
||||
**Payload:** `{ "groupid": "<id>" }`
|
||||
**Returns:** DataFrame with agents
|
||||
|
||||
### 3.4 policy_list_allowlists(groupid: str)
|
||||
|
||||
List allowlists assigned to a policy group.
|
||||
|
||||
**Endpoint:** `POST /v1/group/policies`
|
||||
**Payload:** `{ "groupid": "<id>" }`
|
||||
**Returns:** DataFrame with applications
|
||||
|
||||
### 3.5 policy_clone(source_groupid: str, target_groupid: str)
|
||||
|
||||
Clone a policy from one group to another.
|
||||
|
||||
**Endpoint:** `POST /v1/group/assign`
|
||||
**Payload:** `{ "groupid": "<source>", "targetgroupid": "<target>" }`
|
||||
|
||||
### 3.6 policy_add_path_exclusions(groupid: str, paths: List[str])
|
||||
|
||||
Add path exclusions to a policy group.
|
||||
|
||||
**Endpoint:** `POST /v1/group/path/add`
|
||||
**Payload:** `{ "groupid": "<id>", "path": ["<path1>", "<path2>"] }`
|
||||
|
||||
### 3.7 policy_add_publishers(groupid: str, publishers: List[str])
|
||||
|
||||
Add trusted publishers to a policy group.
|
||||
|
||||
**Endpoint:** `POST /v1/group/publisher/add`
|
||||
**Payload:** `{ "groupid": "<id>", "publisher": ["<pub1>", "<pub2>"] }`
|
||||
|
||||
---
|
||||
|
||||
## 4. OTP Management
|
||||
|
||||
### 4.1 otp_generate(agentid: str, duration: int, purpose: str)
|
||||
|
||||
Generate a new One-Time Password for an agent.
|
||||
|
||||
**Endpoint:** `POST /v1/otp/retrieve`
|
||||
**Payload:**
|
||||
```json
|
||||
{
|
||||
"agentid": "<id>",
|
||||
"duration": "<minutes>",
|
||||
"purpose": "<description>"
|
||||
}
|
||||
```
|
||||
**Returns:** str - The generated OTP code
|
||||
|
||||
### 4.2 otp_find_active()
|
||||
|
||||
Retrieve all active OTP sessions.
|
||||
|
||||
**Endpoint:** `POST /v1/otp/usage`
|
||||
**Payload:** `{ "status": "1" }`
|
||||
**Returns:** DataFrame with OTP records
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `otpid` | str | OTP session identifier |
|
||||
| `agentid` | str | Associated agent ID |
|
||||
| `hostname` | str | Agent hostname |
|
||||
| `purpose` | str | OTP purpose description |
|
||||
| `granted` | str | Grant timestamp |
|
||||
| `expires` | str | Expiration timestamp |
|
||||
|
||||
### 4.3 otp_find_awaiting()
|
||||
|
||||
Retrieve OTPs awaiting activation.
|
||||
|
||||
**Endpoint:** `POST /v1/otp/usage`
|
||||
**Payload:** `{ "status": "0" }`
|
||||
|
||||
### 4.4 otp_find_enforced()
|
||||
|
||||
Retrieve enforced OTPs.
|
||||
|
||||
**Endpoint:** `POST /v1/otp/usage`
|
||||
**Payload:** `{ "status": "2" }`
|
||||
|
||||
### 4.5 otp_find_revoked()
|
||||
|
||||
Retrieve revoked OTPs.
|
||||
|
||||
**Endpoint:** `POST /v1/otp/usage`
|
||||
**Payload:** `{ "status": "3" }`
|
||||
|
||||
### 4.6 otp_find_by_agent(agentid: str)
|
||||
|
||||
Retrieve OTPs for a specific agent.
|
||||
|
||||
**Endpoint:** `POST /v1/otp/usage`
|
||||
**Payload:** `{ "agentid": "<id>" }`
|
||||
|
||||
### 4.7 otp_revoke(otpid: str)
|
||||
|
||||
Revoke an active OTP session.
|
||||
|
||||
**Endpoint:** `POST /v1/otp/revoke`
|
||||
**Payload:** `{ "otpid": "<id>" }`
|
||||
**Returns:** dict with operation result
|
||||
|
||||
### 4.8 otp_validate(otpcode: str)
|
||||
|
||||
Validate an OTP code.
|
||||
|
||||
**Endpoint:** `POST /v1/otp/validate`
|
||||
**Payload:** `{ "otpcode": "<code>" }`
|
||||
**Returns:** dict indicating validity
|
||||
|
||||
### 4.9 otp_get_activities(otpid: str)
|
||||
|
||||
Retrieve activity log for a specific OTP.
|
||||
|
||||
**Endpoint:** `POST /v1/otp/activities`
|
||||
**Payload:** `{ "otpid": "<id>" }`
|
||||
**Returns:** DataFrame with OTP activities
|
||||
|
||||
---
|
||||
|
||||
## 5. Execution History
|
||||
|
||||
### 5.1 history_execution(today: str, date_selected: str, agent_name: str)
|
||||
|
||||
Retrieve execution history for a specific agent.
|
||||
|
||||
**Endpoint:** `POST /v1/getexechistory`
|
||||
**Payload:**
|
||||
```json
|
||||
{
|
||||
"datefrom": "<YYYY-MM-DD>",
|
||||
"dateto": "<YYYY-MM-DD>",
|
||||
"hostname": "<name>"
|
||||
}
|
||||
```
|
||||
**Returns:** List[Dict] with execution records
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `type` | int | Execution type code |
|
||||
| `hostname` | str | Device hostname |
|
||||
| `username` | str | User who executed |
|
||||
| `filename` | str | Executed filename |
|
||||
| `sha256` | str | File hash |
|
||||
| `publisher` | str | Code signer |
|
||||
| `datetime` | str | Execution timestamp |
|
||||
| `policyname` | str | Active policy |
|
||||
| `policyver` | str | Policy version |
|
||||
| `commandline` | str | Full command line |
|
||||
| `pprocess` | str | Parent process |
|
||||
|
||||
### 5.2 Execution Type Codes
|
||||
|
||||
| Code | Description |
|
||||
|------|-------------|
|
||||
| 0 | Trusted Execution |
|
||||
| 1 | Blocked Execution |
|
||||
| 2 | Untrusted Execution [Audit] |
|
||||
| 3 | Untrusted Execution [OTP] |
|
||||
| 4 | Trusted Path Execution |
|
||||
| 5 | Trusted Publisher Execution |
|
||||
| 6 | Blocklist Execution |
|
||||
| 7 | Blocklist Execution [Audit] |
|
||||
| 8 | Trusted Process Execution |
|
||||
| 9 | Constrained Execution |
|
||||
| 10 | Trusted Metadata Execution |
|
||||
| 11 | Trusted Browser Execution |
|
||||
| 12 | Blocked Browser Execution |
|
||||
| 13 | Untrusted Browser Execution [Audit] |
|
||||
| 14 | Untrusted Browser Execution [OTP] |
|
||||
| 15 | Blocklist Browser Execution [Audit] |
|
||||
| 16 | Blocklist Browser Execution |
|
||||
| 17 | Trusted Installer Execution |
|
||||
| 18 | Trusted Browser Metadata Execution |
|
||||
|
||||
### 5.3 history_logging(type: List[str], checkpoint: str, policy: Optional[List[str]])
|
||||
|
||||
Retrieve execution history logs with pagination.
|
||||
|
||||
**Endpoint:** `POST /v1/logging/exechistories`
|
||||
**Payload:**
|
||||
```json
|
||||
{
|
||||
"type": ["1", "2", "3"],
|
||||
"checkpoint": "<checkpoint_id>",
|
||||
"policy": ["<policy_name>"]
|
||||
}
|
||||
```
|
||||
**Returns:** str with execution histories
|
||||
|
||||
---
|
||||
|
||||
## 6. Server Logs
|
||||
|
||||
### 6.1 server_logs(checkpoint: Optional[str])
|
||||
|
||||
Retrieve server activity logs.
|
||||
|
||||
**Endpoint:** `POST /v1/logging/svractivities`
|
||||
**Payload:** `{}` or `{ "checkpoint": "<last_checkpoint>" }`
|
||||
**Returns:** str with server activities
|
||||
|
||||
The checkpoint parameter enables pagination for large result sets. Pass the last checkpoint from a previous call to get subsequent records.
|
||||
|
||||
---
|
||||
|
||||
## 7. Allowlist and Blocklist Management
|
||||
|
||||
### 7.1 allowlist_find_all()
|
||||
|
||||
Retrieve all allowlist applications.
|
||||
|
||||
**Endpoint:** `POST /v1/application`
|
||||
**Payload:** `{}`
|
||||
**Returns:** DataFrame with application records
|
||||
|
||||
### 7.2 allowlist_export(applicationid: str)
|
||||
|
||||
Export allowlist as XML.
|
||||
|
||||
**Endpoint:** `POST /v1/application/export`
|
||||
**Payload:** `{ "applicationid": "<id>" }`
|
||||
**Returns:** bytes - XML content
|
||||
|
||||
### 7.3 baseline_find_all()
|
||||
|
||||
Retrieve all baselines.
|
||||
|
||||
**Endpoint:** `POST /v1/baseline`
|
||||
**Payload:** `{}`
|
||||
**Returns:** DataFrame with baseline records
|
||||
|
||||
### 7.4 baseline_export(baselineid: str)
|
||||
|
||||
Export baseline as XML.
|
||||
|
||||
**Endpoint:** `POST /v1/baseline/export`
|
||||
**Payload:** `{ "baselineid": "<id>" }`
|
||||
**Returns:** bytes - XML content
|
||||
|
||||
### 7.5 blocklist_find_all()
|
||||
|
||||
Retrieve all blocklists.
|
||||
|
||||
**Endpoint:** `POST /v1/blocklist`
|
||||
**Payload:** `{}`
|
||||
**Returns:** DataFrame with blocklist records
|
||||
|
||||
### 7.6 blocklist_export(blocklistid: str)
|
||||
|
||||
Export blocklist as XML.
|
||||
|
||||
**Endpoint:** `POST /v1/blocklist/export`
|
||||
**Payload:** `{ "blocklistid": "<id>" }`
|
||||
**Returns:** bytes - XML content
|
||||
|
||||
### 7.7 hash_add_to_allowlist(applicationid: str, hashes: List[str])
|
||||
|
||||
Add hashes to an allowlist.
|
||||
|
||||
**Endpoint:** `POST /v1/hash/application/add`
|
||||
**Payload:** `{ "applicationid": "<id>", "hashes": ["<sha256>", ...] }`
|
||||
**Returns:** dict with operation result
|
||||
|
||||
### 7.8 hash_query(hashes: List[str])
|
||||
|
||||
Query information about specific hashes.
|
||||
|
||||
**Endpoint:** `POST /v1/hash/query`
|
||||
**Payload:** `{ "hashes": ["<sha256>", ...] }`
|
||||
**Returns:** DataFrame with hash records
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `sha256` | str | Hash value |
|
||||
| `filename` | str | Associated filename |
|
||||
| `publisher` | str | Code signer (or "Not Signed") |
|
||||
| `reputation` | dict | VirusTotal scan results |
|
||||
| `applications` | str | Associated allowlists |
|
||||
| `baselines` | str | Associated baselines |
|
||||
| `blocklists` | str | Associated blocklists |
|
||||
|
||||
---
|
||||
|
||||
## 8. Rust Backend (airlock_libs)
|
||||
|
||||
The `airlock_libs` package provides Rust-accelerated functions for performance-critical operations.
|
||||
|
||||
### 8.1 pull_policy_exec_histories(api, type, days, policy_name)
|
||||
|
||||
Pull execution history for policies with optimized performance.
|
||||
|
||||
| Parameter | Type | Description |
|
||||
|-----------|------|-------------|
|
||||
| `api` | AirlockAPIWrapper | API wrapper instance |
|
||||
| `type` | str | JSON list of exec types, e.g., "[1,2,3]" |
|
||||
| `days` | int | Days to look back |
|
||||
| `policy_name` | Optional[str] | Specific policy or None for all |
|
||||
|
||||
**Returns:** JSON string with execution history
|
||||
|
||||
### 8.2 history_logging(api, exec_types, checkpoint_number, policy_names)
|
||||
|
||||
Query execution logs with pagination support.
|
||||
|
||||
| Parameter | Type | Description |
|
||||
|-----------|------|-------------|
|
||||
| `api` | AirlockAPIWrapper | API wrapper instance |
|
||||
| `exec_types` | str | JSON list of exec types, e.g., "[3,5,8]" |
|
||||
| `checkpoint_number` | str | Checkpoint ID for pagination |
|
||||
| `policy_names` | Optional[str] | Comma-separated policy names or None |
|
||||
|
||||
**Returns:** List[Dict] - Execution history records
|
||||
|
||||
---
|
||||
|
||||
## 9. Error Handling
|
||||
|
||||
All API methods may raise:
|
||||
|
||||
- `requests.exceptions.RequestException` - Network or HTTP errors
|
||||
- `ValueError` - Invalid response format
|
||||
- `KeyError` - Missing expected fields in response
|
||||
|
||||
Recommended pattern:
|
||||
|
||||
```python
|
||||
try:
|
||||
result = api.agent_find_all()
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"API request failed: {e}")
|
||||
# Handle error appropriately
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 10. License
|
||||
|
||||
Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
|
||||
GNU Affero General Public License v3.0
|
||||
@@ -0,0 +1,510 @@
|
||||
# Loxide
|
||||
## User Stories and Use Cases
|
||||
|
||||
**Version 1.0 | December 2025**
|
||||
|
||||
---
|
||||
|
||||
## Epic 1: Multi-Agent Operations
|
||||
|
||||
**Epic Statement:** As a administrator, I need to perform bulk operations on multiple endpoints efficiently.
|
||||
|
||||
---
|
||||
|
||||
### US-1.1: Agent Selection
|
||||
|
||||
**User Story:** As a administrator, I want to select multiple agents using various methods so that I can perform bulk operations efficiently.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Can paste a list of device names (700+ lines)
|
||||
- ✅ Can use wildcards (* and ?) for pattern matching
|
||||
- ✅ Can import device names from a file
|
||||
- ✅ Can toggle between exact and fuzzy matching
|
||||
- ✅ Unmatched entries are clearly displayed
|
||||
- ✅ Can select/deselect all matched agents
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 8
|
||||
|
||||
---
|
||||
|
||||
### US-1.2: Agent Policy Move
|
||||
|
||||
**User Story:** As a administrator, I want to move selected agents to a different policy group so that I can organize endpoints by security requirements.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Can select destination policy from list
|
||||
- ✅ Move operation provides progress feedback
|
||||
- ✅ Success/failure results are color-coded
|
||||
- ✅ Can export results to CSV
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
### US-1.3: Toggle Enforcement Mode
|
||||
|
||||
**User Story:** As a administrator, I want to toggle agents between audit and enforcement mode so that I can gradually roll out policy enforcement.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Clear indication of current mode
|
||||
- ✅ Confirmation before mode change
|
||||
- ✅ Results displayed after operation
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 3
|
||||
|
||||
---
|
||||
|
||||
### US-1.4: View Agent Execution History
|
||||
|
||||
**User Story:** As a administrator, I want to view execution history for selected agents so that I can understand what applications are running.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Date range selector (1-365 days)
|
||||
- ✅ Results displayed in DataTable
|
||||
- ✅ Export to CSV functionality
|
||||
- ✅ Sortable columns
|
||||
|
||||
**Priority:** Medium
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
## Epic 2: Policy Preparation
|
||||
|
||||
**Epic Statement:** As a administrator, I need to prepare policies for enforcement by analyzing execution history.
|
||||
|
||||
---
|
||||
|
||||
### US-2.1: Policy Selection
|
||||
|
||||
**User Story:** As a administrator, I want to select source policies for analysis so that I can review their execution history.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Multi-select with checkboxes
|
||||
- ✅ Filter/search capability
|
||||
- ✅ Policy hierarchy visible
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
### US-2.2: Configure Analysis Parameters
|
||||
|
||||
**User Story:** As a administrator, I want to configure history days and allowlists so that I can customize the analysis scope.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ History days configurable from 1 to 365
|
||||
- ✅ Default value clearly indicated
|
||||
- ✅ Allowlist selection available
|
||||
- ✅ Input validation with helpful error messages
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 3
|
||||
|
||||
---
|
||||
|
||||
### US-2.3: Review Analysis Results
|
||||
|
||||
**User Story:** As a administrator, I want to review categorized execution history so that I can make informed enforcement decisions.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Results categorized: Approved, Unapproved, Needs Review
|
||||
- ✅ Hash reputation data displayed
|
||||
- ✅ Publisher information shown
|
||||
- ✅ Export to CSV available
|
||||
- ✅ Color-coded categories
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 8
|
||||
|
||||
---
|
||||
|
||||
### US-2.4: Add Hashes to Allowlist
|
||||
|
||||
**User Story:** As a administrator, I want to add approved hashes to an allowlist so that they won't be blocked after enforcement.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Select hashes from analysis results
|
||||
- ✅ Choose target allowlist
|
||||
- ✅ Confirmation before adding
|
||||
- ✅ Success/failure feedback
|
||||
|
||||
**Priority:** Medium
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
## Epic 3: Quiet Agent Detection
|
||||
|
||||
**Epic Statement:** As a administrator, I need to identify inactive endpoints ready for enforcement.
|
||||
|
||||
---
|
||||
|
||||
### US-3.1: Configure Quiet Threshold
|
||||
|
||||
**User Story:** As a administrator, I want to set quiet day thresholds so that I can define what constitutes an inactive agent.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Configurable quiet days parameter
|
||||
- ✅ History days parameter
|
||||
- ✅ Clear explanation of thresholds
|
||||
- ✅ Input validation
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 3
|
||||
|
||||
---
|
||||
|
||||
### US-3.2: Run Quiet Agent Analysis
|
||||
|
||||
**User Story:** As a administrator, I want the system to analyze agent activity using high-performance Rust code so that I can quickly identify quiet agents.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Progress bar during analysis
|
||||
- ✅ Console output for Rust progress
|
||||
- ✅ Non-blocking UI during analysis
|
||||
- ✅ Clear completion notification
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 8
|
||||
|
||||
---
|
||||
|
||||
### US-3.3: View Quiet Agents
|
||||
|
||||
**User Story:** As a administrator, I want to see a list of quiet agents so that I can move them to enforcement.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Results displayed in sortable table
|
||||
- ✅ Can select agents for bulk move
|
||||
- ✅ Last activity date shown
|
||||
- ✅ Export to CSV
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
## Epic 4: OTP Management
|
||||
|
||||
**Epic Statement:** As an support technician, I need to manage temporary policy bypasses for end users.
|
||||
|
||||
---
|
||||
|
||||
### US-4.1: Generate OTP
|
||||
|
||||
**User Story:** As an support technician, I want to generate an OTP for an agent so that a user can temporarily bypass policy restrictions.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Select agent by hostname
|
||||
- ✅ Specify duration in minutes
|
||||
- ✅ Enter purpose/ticket number
|
||||
- ✅ OTP code displayed clearly
|
||||
- ✅ Copy to clipboard functionality
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
### US-4.2: View Active OTPs
|
||||
|
||||
**User Story:** As a administrator, I want to view all active OTP sessions so that I can monitor temporary policy bypasses.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ List all active OTPs
|
||||
- ✅ Show hostname, purpose, expiration
|
||||
- ✅ Filter by status
|
||||
- ✅ Refresh capability (r key)
|
||||
- ✅ Sortable columns
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
### US-4.3: Revoke OTP
|
||||
|
||||
**User Story:** As a administrator, I want to revoke an active OTP so that I can end a temporary bypass immediately.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Select OTPs for revocation (checkbox)
|
||||
- ✅ Bulk revoke capability
|
||||
- ✅ Confirmation before revocation
|
||||
- ✅ Results displayed after operation
|
||||
- ✅ Auto-refresh list after revocation
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
### US-4.4: View OTP Activities
|
||||
|
||||
**User Story:** As a administrator, I want to see what applications were executed during an OTP session so that I can audit temporary bypasses.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Select OTP to view
|
||||
- ✅ Display execution list
|
||||
- ✅ Show file, hash, timestamp
|
||||
- ✅ Export capability
|
||||
|
||||
**Priority:** Medium
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
## Epic 5: Execution History
|
||||
|
||||
**Epic Statement:** As a analyst, I need to investigate execution events on endpoints.
|
||||
|
||||
---
|
||||
|
||||
### US-5.1: Query Execution History
|
||||
|
||||
**User Story:** As a analyst, I want to query execution history for specific agents so that I can investigate security events.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Select agent by hostname
|
||||
- ✅ Configure date range (start/end)
|
||||
- ✅ Filter by execution type
|
||||
- ✅ Results in sortable DataTable
|
||||
- ✅ Pagination for large results
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 8
|
||||
|
||||
---
|
||||
|
||||
### US-5.2: Export History
|
||||
|
||||
**User Story:** As a analyst, I want to export execution history to CSV so that I can perform offline analysis.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Export button available (e key)
|
||||
- ✅ All visible columns included
|
||||
- ✅ Proper CSV formatting
|
||||
- ✅ Timestamp in filename
|
||||
- ✅ Notification on success
|
||||
|
||||
**Priority:** Medium
|
||||
**Story Points:** 3
|
||||
|
||||
---
|
||||
|
||||
### US-5.3: Hash Reputation Lookup
|
||||
|
||||
**User Story:** As a analyst, I want to see reputation data for executed files so that I can assess risk.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ VirusTotal score displayed
|
||||
- ✅ Publisher information shown
|
||||
- ✅ Known allowlist membership indicated
|
||||
- ✅ Risk categorization (approved/unapproved/needs_review)
|
||||
|
||||
**Priority:** Medium
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
## Epic 6: Server Monitoring
|
||||
|
||||
**Epic Statement:** As a system administrator, I need to monitor Airlock server activity.
|
||||
|
||||
---
|
||||
|
||||
### US-6.1: View Server Logs
|
||||
|
||||
**User Story:** As a system administrator, I want to view recent server activity so that I can monitor system health.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Default view of last 72 hours
|
||||
- ✅ Datetime properly formatted (YYYY-MM-DD HH:MM:SS)
|
||||
- ✅ Refresh via keyboard shortcut (r)
|
||||
- ✅ Escape to return to main menu
|
||||
- ✅ Auto-scroll to latest entries
|
||||
|
||||
**Priority:** Medium
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
## Epic 7: Application Configuration
|
||||
|
||||
**Epic Statement:** As a power user, I need to customize the application to my preferences.
|
||||
|
||||
---
|
||||
|
||||
### US-7.1: Theme Selection
|
||||
|
||||
**User Story:** As a power user, I want to change the UI theme so that I can work comfortably in different lighting conditions.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Multiple themes available (textual-dark, gruvbox, retro-terminal, amber-terminal)
|
||||
- ✅ Theme persists across sessions
|
||||
- ✅ Preview before applying
|
||||
- ✅ Saved to user config
|
||||
|
||||
**Priority:** Low
|
||||
**Story Points:** 3
|
||||
|
||||
---
|
||||
|
||||
### US-7.2: Working Directory Access
|
||||
|
||||
**User Story:** As a power user, I want to access my working directory from within the application so that I can manage exported files.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Directory tree visible in main menu
|
||||
- ✅ Keyboard shortcut to open in file manager (f)
|
||||
- ✅ Configurable working directory
|
||||
- ✅ Auto-create directory structure
|
||||
|
||||
**Priority:** Low
|
||||
**Story Points:** 3
|
||||
|
||||
---
|
||||
|
||||
### US-7.3: Credential Management
|
||||
|
||||
**User Story:** As a user, I want my API credentials stored securely so that I don't have to enter them every time.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ First-time setup prompts for API key
|
||||
- ✅ Master password protects credentials
|
||||
- ✅ Password complexity requirements enforced
|
||||
- ✅ 3 retry attempts on wrong password
|
||||
- ✅ Platform-native keyring used
|
||||
|
||||
**Priority:** High
|
||||
**Story Points:** 8
|
||||
|
||||
---
|
||||
|
||||
## Epic 8: Statistics and Reporting
|
||||
|
||||
**Epic Statement:** As a administrator, I need visibility into my environment's security posture.
|
||||
|
||||
---
|
||||
|
||||
### US-8.1: View System Statistics
|
||||
|
||||
**User Story:** As a administrator, I want to see an overview of agents and policies so that I can understand my environment.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Total agent count
|
||||
- ✅ Total policy count
|
||||
- ✅ Agent status breakdown (Online/Offline/Hidden/Safemode)
|
||||
- ✅ Visual charts using plotext
|
||||
- ✅ Configurable time range (1/7/30 days)
|
||||
|
||||
**Priority:** Medium
|
||||
**Story Points:** 5
|
||||
|
||||
---
|
||||
|
||||
### US-8.2: View Execution Statistics
|
||||
|
||||
**User Story:** As a administrator, I want to see execution statistics so that I can identify trends.
|
||||
|
||||
**Acceptance Criteria:**
|
||||
- ✅ Execution counts by type
|
||||
- ✅ Top executed files
|
||||
- ✅ Top blocked files
|
||||
- ✅ Visual bar charts
|
||||
- ✅ Refresh capability
|
||||
|
||||
**Priority:** Medium
|
||||
**Story Points:** 5
|
||||
|
||||
|
||||
---
|
||||
|
||||
## Phase 2: LEMON Integration
|
||||
|
||||
**Project:** LEMON (Loxide Execution MONitoring)
|
||||
**Previous Codename:** Overlock
|
||||
**Status:** Phase 2 - Separate Project
|
||||
|
||||
LEMON user stories are documented separately in `LEMON_03_User_Stories.md`.
|
||||
|
||||
**Summary of Phase 2 Loxide Stories:**
|
||||
- LEMON Sessions Screen - View/create/cancel sessions
|
||||
- LEMON Hash Review Screen - Approve/reject pending hashes
|
||||
- Certificate Setup - Configure mTLS authentication
|
||||
- Audit Chain Verification - Verify log integrity
|
||||
|
||||
See LEMON documentation for complete user stories.
|
||||
|
||||
---
|
||||
|
||||
## Acceptance Test Scenarios
|
||||
|
||||
### Scenario: Bulk Agent Move
|
||||
|
||||
```gherkin
|
||||
Given I am on the Multi-Agent Operations screen
|
||||
When I paste 100 device names into the selector
|
||||
And I click Search
|
||||
Then I should see matched agents in the selection list
|
||||
When I select 50 agents using checkboxes
|
||||
And I choose "Production Policy" as destination
|
||||
And I click "Move Selected"
|
||||
Then I should see a progress indicator
|
||||
And I should see 50 success results with green indicators
|
||||
And I should be able to export results to CSV
|
||||
```
|
||||
|
||||
### Scenario: OTP Generation and Revocation
|
||||
|
||||
```gherkin
|
||||
Given I am on the OTP Management screen
|
||||
When I search for agent "DESKTOP-001"
|
||||
And I set duration to 60 minutes
|
||||
And I enter purpose "Ticket #12345 - Software installation"
|
||||
And I click Generate OTP
|
||||
Then I should see an 8-character OTP code
|
||||
And I should be able to copy it to clipboard
|
||||
|
||||
Given the OTP is active
|
||||
When I navigate to OTP Revoke screen
|
||||
And I select the OTP for "DESKTOP-001"
|
||||
And I click "Revoke Selected"
|
||||
Then I should see confirmation dialog
|
||||
When I confirm revocation
|
||||
Then I should see success message
|
||||
And the OTP should no longer appear in active list
|
||||
```
|
||||
|
||||
### Scenario: Policy Preparation Workflow
|
||||
|
||||
```gherkin
|
||||
Given I am on the Policy Prep screen
|
||||
When I select "Audit Policy A" and "Audit Policy B"
|
||||
And I set history days to 30
|
||||
And I click Next
|
||||
Then I should see execution history being fetched
|
||||
|
||||
When the fetch completes
|
||||
Then I should see categorized results:
|
||||
| Category | Count |
|
||||
| Approved | 150 |
|
||||
| Unapproved | 25 |
|
||||
| Needs Review | 10 |
|
||||
|
||||
When I click "Export to CSV"
|
||||
Then I should see a file saved notification
|
||||
And the CSV should contain all execution records
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## License
|
||||
|
||||
Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
|
||||
GNU Affero General Public License v3.0
|
||||
+356
-35
@@ -1,59 +1,380 @@
|
||||
# Loxide
|
||||
|
||||
# 🛡️ Loxide
|
||||
A Python/Rust toolkit for secure, auditable, and automated Airlock agent and policy management. Designed for enterprise environments, Loxide provides advanced policy workflows, device tracking, and a modern terminal-based interface for security administrators who need bulk operations and streamlined workflows not available in the standard Airlock console.
|
||||
|
||||
Python/Rust/Oxide toolkit for secure, auditable, and automated airlock agent and policy management. Designed for enterprise environments, it supports advanced policy workflows, device tracking, and terminal-based interaction.
|
||||

|
||||

|
||||

|
||||

|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
## 🚀 Features
|
||||
- 🔍 **Fuzzy Device Search**
|
||||
Quickly locate devices using partial or approximate matches.
|
||||
|
||||
- 📦 **Batch Move Devices**
|
||||
Move multiple devices between groups or policies easily.
|
||||
|
||||
- 🔄 **Toggle Enforcement/Audit Policies**
|
||||
Seamlessly switch devices between enforcement and audit modes.
|
||||
|
||||
- 🕵️♂️ **Device History Search**
|
||||
Track agent executions.
|
||||
|
||||
- 🧰 **Prepare Policies for Enforcement**
|
||||
Validate and stage policies before pushing them to enforcement.
|
||||
|
||||
- 💤 **Find Quiet Hosts**
|
||||
Identify devices ready for enforcement.
|
||||
|
||||
- 🎛️ **TUI**
|
||||
Navigate with arrow keys and F-key shortcuts using a custom ANSI-colored terminal UI.
|
||||
- [Features](#features)
|
||||
- [Installation](#installation)
|
||||
- [Configuration](#configuration)
|
||||
- [Usage](#usage)
|
||||
- [Workflows](#workflows)
|
||||
- [Architecture](#architecture)
|
||||
- [Keyboard Reference](#keyboard-reference)
|
||||
- [Roadmap](#roadmap)
|
||||
- [Authors](#authors)
|
||||
- [License](#license)
|
||||
|
||||
---
|
||||
|
||||
## 🧭 Roadmap
|
||||
## Features
|
||||
|
||||
- ⚙️ **Rust-based Async API Calls**
|
||||
Improve performance and concurrency with a Rust-powered backend.
|
||||
### Multi-Agent Operations
|
||||
|
||||
- ✅ **Carbon Black-style Local Approval**
|
||||
Enable local user approvals for policy exceptions and enforcement actions.
|
||||
Select and operate on multiple endpoints simultaneously using hostname, wildcard patterns (`*`, `?`), or file import. Supports lists of 700+ devices with fuzzy matching.
|
||||
|
||||
- 📊 **Audit Logging & Export**
|
||||
Add detailed logging and export capabilities for compliance and analysis.
|
||||
- Move agents between policy groups in bulk
|
||||
- Toggle between audit and enforcement modes
|
||||
- Generate OTPs for multiple agents at once
|
||||
- View execution history across selected devices
|
||||
- Export operation results to CSV
|
||||
|
||||
### Policy Preparation Workflow
|
||||
|
||||
A guided multi-step wizard for safely transitioning policies from audit to enforcement:
|
||||
|
||||
1. Select source policies for analysis (multi-select with checkboxes)
|
||||
2. Choose destination policy and allowlist
|
||||
3. Configure analysis parameters (1-365 day history window)
|
||||
4. Fetch and categorize execution history data
|
||||
5. Review unapproved applications by category
|
||||
6. Add approved hashes to allowlists
|
||||
7. Export analysis results for documentation
|
||||
|
||||
### Quiet Agent Detection
|
||||
|
||||
Identify inactive endpoints that are ready for enforcement using high-performance Rust-powered analysis:
|
||||
|
||||
- Configure quiet day threshold and history parameters
|
||||
- Analyze agent activity patterns efficiently
|
||||
- Review quiet agent list with policy information
|
||||
- Bulk move quiet agents to enforcement policies
|
||||
|
||||
### OTP Management
|
||||
|
||||
Complete One-Time Password lifecycle management for handling application blocking issues:
|
||||
|
||||
- Generate OTPs with configurable duration and purpose
|
||||
- Monitor active OTP sessions in real-time
|
||||
- View applications executed during OTP windows
|
||||
- Revoke active OTPs when needed
|
||||
- Analyze historical OTP activity
|
||||
|
||||
### Execution History Analysis
|
||||
|
||||
Detailed analysis of application execution events across your environment:
|
||||
|
||||
- Date range selection (configurable 1-365 days)
|
||||
- Filter by agent, policy, filename, publisher, or hash
|
||||
- Categorize executions by approval status (18 execution types)
|
||||
- VirusTotal reputation data via Airlock's integration
|
||||
- Export results to CSV for external analysis
|
||||
|
||||
### Server Activity Monitoring
|
||||
|
||||
Real-time visibility into Airlock server operations:
|
||||
|
||||
- DataTable display of server activity logs
|
||||
- Configurable time window (default 72 hours)
|
||||
- Auto-refresh capability
|
||||
- Search and filter functionality
|
||||
|
||||
### Modern Terminal Interface
|
||||
|
||||
Built with Textual 6.5.0 for a responsive, keyboard-driven experience:
|
||||
|
||||
- Full keyboard navigation with discoverable shortcuts
|
||||
- Tab-based workflow organization
|
||||
- Toast notifications for operations
|
||||
- Color-coded status indicators
|
||||
- Multiple theme options (Gruvbox, Amber Terminal, Retro Terminal)
|
||||
|
||||
---
|
||||
|
||||
## 🧑💻 Requirements
|
||||
## Installation
|
||||
|
||||
[airlock_libs](https://git.racooncity.org/brotoskyj/-/packages/pypi/airlock-libs/)
|
||||
### Prerequisites
|
||||
|
||||
- Python 3.10 or higher
|
||||
- Access to Airlock API server
|
||||
- Network access to private PyPI server (for airlock_libs)
|
||||
- Modern terminal with Unicode support (80x24 minimum, 120x40 recommended)
|
||||
|
||||
### Quick Start
|
||||
|
||||
```bash
|
||||
# Clone the repository
|
||||
git clone https://git.racooncity.org/brotoskyj/loxide.git
|
||||
cd loxide
|
||||
|
||||
# Install dependencies
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Run the application
|
||||
python Loxide.py
|
||||
```
|
||||
|
||||
### Dependencies
|
||||
|
||||
```
|
||||
# Core TUI Framework
|
||||
textual==6.5.0
|
||||
|
||||
# API and Data Handling
|
||||
requests==2.32.5
|
||||
pandas==2.3.3
|
||||
numpy==2.3.4
|
||||
|
||||
# Database
|
||||
pymongo==4.15.3
|
||||
|
||||
# Security
|
||||
cryptography==46.0.3
|
||||
keyring==25.6.0
|
||||
|
||||
# Utilities
|
||||
python-dotenv==1.2.1
|
||||
tqdm==4.67.1
|
||||
urllib3==2.5.0
|
||||
pyperclip==1.11.0
|
||||
|
||||
# Private Package (Rust integration)
|
||||
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
|
||||
airlock_libs==5.2.1
|
||||
```
|
||||
|
||||
See also: [airlock_libs documentation](https://git.racooncity.org/brotoskyj/-/packages/pypi/airlock-libs/)
|
||||
|
||||
---
|
||||
|
||||
## 📜 License
|
||||
## Configuration
|
||||
|
||||
### First Run
|
||||
|
||||
On first launch, Loxide will prompt for your Airlock API credentials. These are securely stored using your platform's native keyring (Windows Credential Manager, macOS Keychain, or Linux Secret Service) with an additional encryption layer.
|
||||
|
||||
### Configuration Files
|
||||
|
||||
Loxide uses two configuration files:
|
||||
|
||||
| File | Location | Purpose |
|
||||
|------|----------|---------|
|
||||
| `system_config.json` | Application directory | System defaults, bundled with application |
|
||||
| `user_config.json` | User home directory | User preferences, created on first run |
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Description |
|
||||
|----------|-------------|
|
||||
| `LOXIDE_LOG_LEVEL` | Logging verbosity (DEBUG, INFO, WARNING, ERROR) |
|
||||
| `HTTPS_PROXY` | Proxy server for Airlock API connections |
|
||||
|
||||
---
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
python Loxide.py
|
||||
```
|
||||
|
||||
### Main Interface
|
||||
|
||||
The application opens to a tabbed interface with the following sections:
|
||||
|
||||
| Tab | Description |
|
||||
|-----|-------------|
|
||||
| Dashboard | Summary statistics and environment overview |
|
||||
| Agent Operations | Multi-agent selection and bulk operations |
|
||||
| Policy Prep | Policy enforcement preparation wizard |
|
||||
| Quiet Agents | Inactive agent detection workflow |
|
||||
| OTP Management | OTP generation, monitoring, and revocation |
|
||||
|
||||
### User Classes
|
||||
|
||||
Loxide supports three primary user roles:
|
||||
|
||||
- **Administrator** — Full access to all functionality, policy management, bulk operations
|
||||
- **Analyst** — Monitoring, investigation, and reporting (execution history, OTP activity, logs)
|
||||
- **Support** — Targeted operations like OTP generation for end-user issues
|
||||
|
||||
---
|
||||
|
||||
## Workflows
|
||||
|
||||
### Moving Agents to a New Policy
|
||||
|
||||
1. Navigate to **Agent Operations** tab
|
||||
2. Enter device names (paste list, use wildcards, or load from file)
|
||||
3. Click **Search** to find matching agents
|
||||
4. Select agents using checkboxes (or Select All)
|
||||
5. Click **Move to Policy**
|
||||
6. Select destination policy from the tree
|
||||
7. Confirm the operation
|
||||
|
||||
### Preparing a Policy for Enforcement
|
||||
|
||||
1. Navigate to **Policy Prep** tab
|
||||
2. Select source policies to analyze
|
||||
3. Choose destination enforcement policy
|
||||
4. Set history days (how far back to analyze)
|
||||
5. Select target allowlist for approved hashes
|
||||
6. Click **Fetch Data** to retrieve execution history
|
||||
7. Review unapproved applications by category
|
||||
8. Add approved hashes to the allowlist
|
||||
9. Export results for documentation
|
||||
|
||||
### Generating an OTP
|
||||
|
||||
1. Navigate to **Agent Operations** tab
|
||||
2. Search for and select the target agent
|
||||
3. Click **Generate OTP**
|
||||
4. Set duration (in minutes) and purpose
|
||||
5. Copy the generated OTP code
|
||||
|
||||
---
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
Loxide/
|
||||
├── Loxide.py # Main application entry point
|
||||
├── API.py # AirlockAPIWrapper - all API calls
|
||||
├── setup.py # Initialization, logging, config loading
|
||||
├── configmanager.py # System/user configuration management
|
||||
│
|
||||
├── models/ # Data models
|
||||
│ ├── agent.py # Agent data model
|
||||
│ ├── policy.py # Policy data model
|
||||
│ └── execution.py # Execution event model
|
||||
│
|
||||
├── services/ # Business logic
|
||||
│ └── security.py # Credential management
|
||||
│
|
||||
├── utils/ # Helper functions
|
||||
│ └── utils.py # Utility functions
|
||||
│
|
||||
└── TUI/
|
||||
├── Screens/ # Full-page workflow screens
|
||||
│ ├── policyprepworkflowscreen.py
|
||||
│ ├── quietagentworkflowscreen.py
|
||||
│ ├── moveagentworkflowscreen.py
|
||||
│ ├── executionhistoryscreen.py
|
||||
│ ├── otpworkflowscreen.py
|
||||
│ ├── otpactivityscreen.py
|
||||
│ ├── otprevokescreen.py
|
||||
│ └── allowlistselectionscreen.py
|
||||
│
|
||||
├── Widgets/ # Reusable UI components
|
||||
│ ├── policyselector.py
|
||||
│ ├── policytreewidget.py
|
||||
│ ├── multiagentselector.py
|
||||
│ ├── agentmoveoperations.py
|
||||
│ ├── serverlogwidget.py
|
||||
│ ├── resultsdisplay.py
|
||||
│ └── themeselector.py
|
||||
│
|
||||
└── Themes/ # Color themes
|
||||
├── theme_amber_terminal.py
|
||||
└── theme_retro_terminal.py
|
||||
```
|
||||
|
||||
### Design Principles
|
||||
|
||||
- **Layered Architecture** — Clear separation between presentation, business logic, data access, and infrastructure
|
||||
- **Screen vs Widget Pattern** — Screens have headers/footers and are pushed; widgets are embedded components
|
||||
- **Secure by Default** — API keys encrypted at rest, secure credential storage
|
||||
- **Cross-Platform** — Works on Windows 10/11 and Linux (Ubuntu 24+)
|
||||
|
||||
---
|
||||
|
||||
## Keyboard Reference
|
||||
|
||||
### Global
|
||||
|
||||
| Key | Action |
|
||||
|-----|--------|
|
||||
| `q` | Quit application |
|
||||
| `Tab` | Next widget |
|
||||
| `Shift+Tab` | Previous widget |
|
||||
| `←` `→` | Navigate tabs |
|
||||
| `?` | Show help |
|
||||
|
||||
### Screens
|
||||
|
||||
| Key | Action |
|
||||
|-----|--------|
|
||||
| `Escape` | Go back / Cancel |
|
||||
| `r` | Refresh data |
|
||||
| `e` | Export to CSV |
|
||||
|
||||
### Data Tables
|
||||
|
||||
| Key | Action |
|
||||
|-----|--------|
|
||||
| `↑` `↓` | Move cursor |
|
||||
| `Page Up/Down` | Scroll page |
|
||||
| `Home` / `End` | First / last row |
|
||||
| `Space` | Toggle selection |
|
||||
| `Enter` | Select / activate |
|
||||
| `a` | Select all |
|
||||
| `n` | Deselect all |
|
||||
|
||||
---
|
||||
|
||||
## Roadmap
|
||||
|
||||
### Planned Features
|
||||
|
||||
- **Carbon Black-style Local Approval** — Enable local user approvals for policy exceptions and enforcement actions
|
||||
|
||||
- **Audit Logging & Export** — Detailed logging and export capabilities for compliance and analysis
|
||||
|
||||
- **LEMON Integration** — Phase 2 integration with LEMON for enhanced analytics
|
||||
|
||||
- **Offline Mode** — Cached data review when Airlock server is unreachable
|
||||
|
||||
---
|
||||
|
||||
## Performance
|
||||
|
||||
| Metric | Target |
|
||||
|--------|--------|
|
||||
| Application startup | < 5 seconds |
|
||||
| Agent list load (1000 agents) | < 10 seconds |
|
||||
| Execution history (30 days, single agent) | < 15 seconds |
|
||||
| UI response time | < 100 milliseconds |
|
||||
| CSV export (10,000 rows) | < 30 seconds |
|
||||
| Memory usage (idle) | < 200 MB |
|
||||
|
||||
---
|
||||
|
||||
## Authors
|
||||
|
||||
- **Brandon Wickline** — Lead Python Developer
|
||||
- **James Brotosky** — Lead Rust Developer
|
||||
|
||||
---
|
||||
|
||||
## License
|
||||
|
||||
**Loxide** is licensed under the **GNU Affero General Public License v3.0**.
|
||||
|
||||
You may copy, distribute, and modify the software under the terms of the AGPL-3.0 license.
|
||||
|
||||
See the [LICENSE](LICENSE.md) file for full details, or visit
|
||||
[https://www.gnu.org/license/agpl-3.0.html](https://www.gnu.org/license/agpl-3.0.html)
|
||||
```
|
||||
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.
|
||||
```
|
||||
|
||||
See the [LICENSE](LICENSE.md) file for full details, or visit [gnu.org/licenses/agpl-3.0.html](https://www.gnu.org/licenses/agpl-3.0.html)
|
||||
@@ -1,242 +0,0 @@
|
||||
"""
|
||||
This module handles the creation of local approval requests.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
from models.agent import Agent
|
||||
from services.agenthandler import moveAgentToRelatedPolicy, selectAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_protected_json
|
||||
from utils.utils import colorText, get_sanitized_input
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LocalApprovalRequestor:
|
||||
"""Handles creation of local approval requests in Loxide."""
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper, username: str = None):
|
||||
"""
|
||||
Initialize the local approval requestor.
|
||||
|
||||
Args:
|
||||
api: AirlockAPIWrapper instance
|
||||
username: Username creating the approvals (for tracking)
|
||||
"""
|
||||
self.api = api
|
||||
self.policy_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
self.username = (
|
||||
username or os.getenv("USERNAME") or os.getenv("USER") or "unknown"
|
||||
)
|
||||
|
||||
def create_local_approval(
|
||||
self, agent_id: str, duration_minutes: int, batch_id: Optional[int] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Create a single local approval request.
|
||||
|
||||
Args:
|
||||
agent_id: Agent ID to create approval for
|
||||
duration_minutes: Duration of approval in minutes
|
||||
batch_id: Optional batch identifier (defaults to timestamp)
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if batch_id is None:
|
||||
batch_id = int(time.time())
|
||||
|
||||
purpose = (
|
||||
f"🎫 Local Approval 🎫 - {duration_minutes} mins - "
|
||||
f"batch:{batch_id} Client:{agent_id} User:{self.username}"
|
||||
)
|
||||
|
||||
try:
|
||||
self.api.otp_generate(agent_id, duration_minutes, purpose)
|
||||
logger.info(
|
||||
f"Generated local approval for {agent_id}, batch {batch_id}, by {self.username}"
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to generate local approval for {agent_id}: {e}")
|
||||
return False
|
||||
|
||||
def move_agent_to_audit(self, agent: Agent) -> bool:
|
||||
"""
|
||||
Move an agent to its corresponding audit policy.
|
||||
|
||||
Args:
|
||||
agent: Agent object to move
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
moveAgentToRelatedPolicy(self.api, agent, "audit")
|
||||
logger.info(f"Moved {agent.hostname} to audit policy")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to move {agent.hostname} to audit: {e}")
|
||||
return False
|
||||
|
||||
def create_local_approval_batch(
|
||||
self,
|
||||
agents: List[Agent],
|
||||
duration_minutes: int,
|
||||
) -> tuple[int, int, int]:
|
||||
"""
|
||||
Create local approvals for multiple agents and move them to audit.
|
||||
|
||||
Args:
|
||||
agents: List of Agent objects
|
||||
duration_minutes: Duration of approval in minutes
|
||||
db_path: Optional path to database for history tracking
|
||||
|
||||
Returns:
|
||||
Tuple of (batch_id, success_count, failure_count)
|
||||
"""
|
||||
batch_id = int(time.time())
|
||||
success_count = 0
|
||||
failure_count = 0
|
||||
|
||||
print(colorText(f"\n📦 Processing batch {batch_id}...", "cyan"))
|
||||
print(colorText(f"👤 Requested by: {self.username}", "cyan"))
|
||||
print(
|
||||
colorText(f"📊 Moving {len(agents)} agent(s) to local approval\n", "cyan")
|
||||
)
|
||||
|
||||
for agent in agents:
|
||||
try:
|
||||
# Create local approval
|
||||
approval_success = self.create_local_approval(
|
||||
agent.agentid, duration_minutes, batch_id
|
||||
)
|
||||
|
||||
if not approval_success:
|
||||
raise Exception("Failed to create local approval")
|
||||
|
||||
# Move to audit policy
|
||||
move_success = self.move_agent_to_audit(agent)
|
||||
|
||||
if not move_success:
|
||||
raise Exception("Failed to move to audit policy")
|
||||
|
||||
print(colorText(f"✓ {agent.hostname}", "green"))
|
||||
success_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(colorText(f"✗ {agent.hostname}: {e}", "red"))
|
||||
logger.error(f"Error processing agent {agent.hostname}: {e}")
|
||||
failure_count += 1
|
||||
|
||||
return batch_id, success_count, failure_count
|
||||
|
||||
def interactive_local_approval(self):
|
||||
"""
|
||||
Interactive workflow to create local approvals for selected agents.
|
||||
|
||||
This prompts the user to select a duration and agents, then creates
|
||||
the local approvals and moves agents to audit policies.
|
||||
"""
|
||||
# Duration options in minutes
|
||||
duration_options = [
|
||||
(15, "15 minutes"),
|
||||
(60, "1 hour"),
|
||||
(360, "6 hours"),
|
||||
(1440, "1 day"),
|
||||
(10080, "1 week"),
|
||||
]
|
||||
|
||||
# Display duration options
|
||||
print(colorText("\n⏱️ Select Local Approval Duration:", "white"))
|
||||
print(colorText("=" * 50, "white"))
|
||||
|
||||
for i, (minutes, label) in enumerate(duration_options, start=1):
|
||||
print(f" {i}. {label} ({minutes} minutes)")
|
||||
|
||||
print(colorText("=" * 50, "white"))
|
||||
|
||||
# Get user selection
|
||||
try:
|
||||
choice = int(get_sanitized_input("\nEnter the number of your choice: "))
|
||||
|
||||
if 1 <= choice <= len(duration_options):
|
||||
duration_minutes, duration_label = duration_options[choice - 1]
|
||||
print(colorText(f"✓ Selected: {duration_label}", "green"))
|
||||
logger.info(f"User selected duration: {duration_minutes} minutes")
|
||||
else:
|
||||
print(colorText("❌ Invalid choice.", "red"))
|
||||
logger.warning("Invalid duration choice")
|
||||
return
|
||||
|
||||
except ValueError:
|
||||
print(colorText("❌ Invalid input. Please enter a number.", "red"))
|
||||
logger.warning("Invalid input for duration selection")
|
||||
return
|
||||
|
||||
# Select agents
|
||||
print(colorText("\n🎯 Select Agents for Local Approval:", "white"))
|
||||
agents = selectAgents(self.api)
|
||||
|
||||
if not agents:
|
||||
print(colorText("❌ No agents found or error retrieving agents.", "red"))
|
||||
logger.warning("No agents selected or error retrieving agents")
|
||||
return
|
||||
|
||||
# Confirm with user
|
||||
print(colorText("\n📋 Summary:", "cyan"))
|
||||
print(colorText(f" Duration: {duration_label}", "white"))
|
||||
print(colorText(f" Agents: {len(agents)}", "white"))
|
||||
|
||||
confirm = get_sanitized_input("\nProceed? (y/n): ").lower()
|
||||
|
||||
if confirm != "y":
|
||||
print(colorText("❌ Operation cancelled.", "yellow"))
|
||||
return
|
||||
|
||||
# Process the batch
|
||||
batch_id, success_count, failure_count = self.create_local_approval_batch(
|
||||
agents, duration_minutes
|
||||
)
|
||||
|
||||
# Display summary
|
||||
self._display_summary(batch_id, duration_label, success_count, failure_count)
|
||||
|
||||
def _display_summary(
|
||||
self, batch_id: int, duration_label: str, success_count: int, failure_count: int
|
||||
):
|
||||
"""
|
||||
Display operation summary.
|
||||
|
||||
Args:
|
||||
batch_id: Batch identifier
|
||||
duration_label: Human-readable duration
|
||||
success_count: Number of successful operations
|
||||
failure_count: Number of failed operations
|
||||
"""
|
||||
print(colorText(f"\n{'=' * 60}", "white"))
|
||||
print(colorText("📊 Local Approval Summary", "cyan"))
|
||||
print(colorText("=" * 60, "white"))
|
||||
|
||||
print(colorText(f"✓ Successfully processed: {success_count}", "green"))
|
||||
|
||||
if failure_count > 0:
|
||||
print(colorText(f"✗ Failed: {failure_count}", "red"))
|
||||
|
||||
print(colorText(f"\n📦 Batch ID: {batch_id}", "cyan"))
|
||||
print(colorText(f"⏱️ Duration: {duration_label}", "cyan"))
|
||||
|
||||
print(colorText("=" * 60, "white"))
|
||||
print(colorText("\n💡 Next Steps:", "yellow"))
|
||||
print(colorText(" • Agents have been moved to audit policies", "white"))
|
||||
print(colorText(" • Local approvals are active", "white"))
|
||||
print(
|
||||
colorText(
|
||||
f" • Agents will return to enforcement after {duration_label}", "white"
|
||||
)
|
||||
)
|
||||
print(colorText("=" * 60 + "\n", "white"))
|
||||
-187
@@ -1,187 +0,0 @@
|
||||
# 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 datetime import datetime
|
||||
import logging
|
||||
import os
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from services.agenthandler import selectAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import load_env
|
||||
from utils.selector import Selector
|
||||
from utils.utils import colorText, get_sanitized_input
|
||||
|
||||
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):
|
||||
activeagents = api.otp_find_active()
|
||||
awaitingagents = api.otp_find_awaiting()
|
||||
enforcedagents = api.otp_find_enforced()
|
||||
revokedagents = api.otp_find_revoked()
|
||||
|
||||
# Add a 'status' column to each DataFrame
|
||||
activeagents["status"] = "active"
|
||||
awaitingagents["status"] = "awaiting"
|
||||
enforcedagents["status"] = "enforced"
|
||||
revokedagents["status"] = "revoked"
|
||||
|
||||
# Combine all into one DataFrame
|
||||
combined_agents = pd.concat(
|
||||
[activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True
|
||||
)
|
||||
combined_agents = combined_agents.sort_values(by="otpid", ascending=False)
|
||||
|
||||
# Optionally, select specific hosts
|
||||
user_input = (
|
||||
get_sanitized_input("\nWould you like to search for a specific device? (y/n): ")
|
||||
.strip()
|
||||
.lower()
|
||||
)
|
||||
if user_input == "y":
|
||||
agentnames = []
|
||||
agents = selectAgents(api)
|
||||
for agent in agents:
|
||||
agentnames.append(agent.hostname)
|
||||
|
||||
combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)]
|
||||
|
||||
# Present and select rows
|
||||
selected_rows = Selector.select_dataframe_with_mode(
|
||||
combined_agents,
|
||||
columns=["otpid", "hostname", "status", "purpose", "granted"],
|
||||
header="OTP Sessions",
|
||||
)
|
||||
combined_df = pd.DataFrame()
|
||||
|
||||
for row in selected_rows:
|
||||
otpid = row["otpid"]
|
||||
hostname = row["hostname"]
|
||||
result = api.otp_get_activities(otpid)
|
||||
result["hostname"] = hostname
|
||||
if not result.empty:
|
||||
logger.info(f"Activities for {hostname} (otpid: {otpid}):\n{result}")
|
||||
combined_df = pd.concat([combined_df, result], ignore_index=True)
|
||||
else:
|
||||
logger.info(f"No activities found for {hostname} (otpid: {otpid})")
|
||||
|
||||
user_input = (
|
||||
get_sanitized_input(
|
||||
"\nWould you like to export the results to a CSV file? (y/n): "
|
||||
)
|
||||
.strip()
|
||||
.lower()
|
||||
)
|
||||
if user_input == "y":
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
filename = f"otp_activities_{timestamp}.csv"
|
||||
file_path = os.path.join(str(working_dir), filename)
|
||||
|
||||
combined_df.to_csv(file_path, index=False)
|
||||
logging.info(f"Exported Data to {file_path}")
|
||||
|
||||
print(
|
||||
colorText(
|
||||
f"\n✅ OTP Activity exported to: {working_dir}\\{filename}",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
logging.debug("User declined to export the DataFrame.")
|
||||
|
||||
|
||||
def otp_revoke(api: AirlockAPIWrapper):
|
||||
|
||||
activeagents = api.otp_find_active()
|
||||
awaitingagents = api.otp_find_awaiting()
|
||||
|
||||
activeagents["status"] = "active"
|
||||
awaitingagents["status"] = "awaiting"
|
||||
|
||||
combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True)
|
||||
combined_agents = combined_agents.sort_values(by="otpid", ascending=False)
|
||||
|
||||
# Combine all into one DataFrame
|
||||
combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True)
|
||||
combined_agents = combined_agents.sort_values(by="otpid", ascending=False)
|
||||
|
||||
# Optionally, select specific hosts
|
||||
user_input = (
|
||||
get_sanitized_input("\nWould you like to search for a specific device? (y/n): ")
|
||||
.strip()
|
||||
.lower()
|
||||
)
|
||||
if user_input == "y":
|
||||
agentnames = []
|
||||
agents = selectAgents(api)
|
||||
for agent in agents:
|
||||
agentnames.append(agent.hostname)
|
||||
|
||||
combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)]
|
||||
|
||||
# Present and select rows
|
||||
selected_rows = Selector.select_dataframe_with_mode(
|
||||
combined_agents,
|
||||
columns=["otpid", "hostname", "status", "purpose", "granted"],
|
||||
header="OTP Sessions",
|
||||
)
|
||||
|
||||
for row in selected_rows:
|
||||
otpid = row["otpid"]
|
||||
hostname = row["hostname"]
|
||||
result = api.otp_revoke(otpid)
|
||||
logger.info(f"{hostname} (otpid: {otpid}):\n{result}")
|
||||
@@ -1,870 +0,0 @@
|
||||
# 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/>.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import os.path
|
||||
import re
|
||||
from typing import List
|
||||
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
from models.execution import ExecutionHistoryRecord
|
||||
from models.policy import Allowlist, Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_protected_value, load_env, load_env_json
|
||||
from utils.selector import Selector
|
||||
from utils.utils import (
|
||||
areYouSure,
|
||||
clear_screen,
|
||||
colorText,
|
||||
formatHTML,
|
||||
get_sanitized_input,
|
||||
locked,
|
||||
open_directory,
|
||||
print_x_wide,
|
||||
regulator,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
|
||||
|
||||
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||||
logger.debug("Prompting for Policies")
|
||||
print(colorText("Please select policy/policies", "white"))
|
||||
selected = Selector.select_objects(policies, allow_multiple, prompt_each=True)
|
||||
|
||||
if selected is None:
|
||||
return []
|
||||
|
||||
# Normalize to always return a list
|
||||
logger.debug("Returning {selected.dict}")
|
||||
return selected if isinstance(selected, list) else [selected]
|
||||
|
||||
|
||||
def selectAllowlists(
|
||||
api: AirlockAPIWrapper, policy=all, allow_multiple=True
|
||||
) -> List[Allowlist]:
|
||||
if policy == "all":
|
||||
allowlists = [
|
||||
Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()
|
||||
]
|
||||
else:
|
||||
allowlists = [
|
||||
Allowlist(**row.to_dict())
|
||||
for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()
|
||||
]
|
||||
logger.debug("Prompting for Allowlist(s)")
|
||||
print(colorText("Please select allowlist(s)", "white"))
|
||||
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
|
||||
|
||||
if selected is None:
|
||||
return []
|
||||
|
||||
# Normalize to always return a list
|
||||
logger.debug(f"Returning {selected}")
|
||||
return selected if isinstance(selected, list) else [selected]
|
||||
|
||||
|
||||
def sortHashes(
|
||||
api: AirlockAPIWrapper, selected_policies: List[Policy], type=[1, 2, 6, 7]
|
||||
):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
|
||||
logger.debug(f"{history_days} day selected for history")
|
||||
|
||||
if history_days is None:
|
||||
logging.warning("No history range selected. Aborting.")
|
||||
return
|
||||
|
||||
policy_executions = ExecutionHistoryRecord.from_policies(
|
||||
api, selected_policies, type_=type, history_days=history_days
|
||||
)
|
||||
|
||||
logger.debug(f"Executions contains {policy_executions}")
|
||||
|
||||
enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(
|
||||
api, policy_executions
|
||||
)
|
||||
categorized_executions = (
|
||||
ExecutionHistoryRecord.categorize_executions_by_hash_decision(
|
||||
enriched_executions
|
||||
)
|
||||
)
|
||||
approved, unapproved, needs_review, unknown = (
|
||||
ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions)
|
||||
)
|
||||
|
||||
categories = {
|
||||
"needs_review": needs_review,
|
||||
"approved": approved,
|
||||
"unapproved": unapproved,
|
||||
"leftover": unknown,
|
||||
}
|
||||
|
||||
for label, records in categories.items():
|
||||
if not records:
|
||||
continue # Skip empty or falsy categories
|
||||
|
||||
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
|
||||
html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html"
|
||||
|
||||
# Convert ExecutionHistoryRecord objects to dictionaries
|
||||
df = pd.DataFrame([r.__dict__ for r in records])
|
||||
|
||||
# Optional: flatten hash_obj if needed
|
||||
if not df.empty and "hash_obj" in df.columns:
|
||||
hash_df = df["hash_obj"].apply(lambda h: h.to_dict() if h else {})
|
||||
df = pd.concat([df.drop(columns=["hash_obj"]), hash_df], axis=1)
|
||||
|
||||
# Save to CSV
|
||||
df.to_csv(csv_path, index=False)
|
||||
logger.info(f"Saved {label} executions to {csv_path}")
|
||||
|
||||
# Generate HTML
|
||||
formatHTML(df, html_path)
|
||||
logger.info(f"Generated HTML report at {html_path}")
|
||||
|
||||
|
||||
def buildPathsandPublishers(selected_policies: List[Policy], split):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
df1 = pd.DataFrame()
|
||||
df2 = pd.DataFrame()
|
||||
all_approved_hashes = pd.DataFrame()
|
||||
path1 = (
|
||||
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
|
||||
)
|
||||
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
|
||||
path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type=int)
|
||||
|
||||
if os.path.exists(path1):
|
||||
df1 = pd.read_csv(path1)
|
||||
else:
|
||||
logger.warning(f"File not found: {path1}")
|
||||
|
||||
if os.path.exists(path2):
|
||||
df2 = pd.read_csv(path2)
|
||||
else:
|
||||
logger.warning(f"File not found: {path2}")
|
||||
|
||||
if df1.empty and df2.empty:
|
||||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||
all_approved_hashes = pd.DataFrame()
|
||||
logger.debug(all_approved_hashes.head)
|
||||
else:
|
||||
all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
|
||||
if "filename" in all_approved_hashes.columns:
|
||||
all_approved_hashes = all_approved_hashes.sort_values(by="filename")
|
||||
else:
|
||||
logger.warning(
|
||||
"Warning: 'filename' column not found in concatenated DataFrame."
|
||||
)
|
||||
|
||||
if not all_approved_hashes.empty and path_exclusion_constant:
|
||||
|
||||
primary_path_exclusions = calculatePath(
|
||||
all_approved_hashes,
|
||||
path_exclusion_constant,
|
||||
split,
|
||||
)
|
||||
remaining_hashes = all_approved_hashes[
|
||||
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
|
||||
]
|
||||
secondary_path_exclusions = calculatePath(
|
||||
remaining_hashes, (path_exclusion_constant - 1), split
|
||||
)
|
||||
remaining_hashes = remaining_hashes[
|
||||
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
|
||||
]
|
||||
dataframes = {
|
||||
"all_approved_hashes": all_approved_hashes,
|
||||
"primary_Paths": primary_path_exclusions,
|
||||
"secondary_Paths": secondary_path_exclusions,
|
||||
"hashes_not_approvable_by_path": remaining_hashes,
|
||||
}
|
||||
logger.debug("Preparing to sort dataframes")
|
||||
for name, df in dataframes.items():
|
||||
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||||
if "hashes" in name:
|
||||
df.sort_values(by="filename", inplace=True)
|
||||
else:
|
||||
df.sort_values(by="longestcfp", inplace=True)
|
||||
|
||||
df.to_csv(
|
||||
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv",
|
||||
index=False,
|
||||
)
|
||||
formatHTML(
|
||||
df,
|
||||
f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html",
|
||||
)
|
||||
|
||||
if not all_approved_hashes.empty:
|
||||
# Drop all not signed, only keep unique values
|
||||
publist = all_approved_hashes[
|
||||
all_approved_hashes["publisher"] != "Not Signed"
|
||||
].drop_duplicates(subset=["publisher"])
|
||||
# Remove Bad publisher if somehow they made it this far
|
||||
pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
|
||||
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
|
||||
publist = publist[["publisher"]]
|
||||
publist.sort_values(by="publisher", inplace=True)
|
||||
publist.to_csv(
|
||||
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv",
|
||||
index=False,
|
||||
)
|
||||
else:
|
||||
logger.debug("Approved Hashes list appears empty")
|
||||
|
||||
|
||||
def buildPreflights(selected_policies: List[Policy]):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
df1 = pd.DataFrame()
|
||||
df2 = pd.DataFrame()
|
||||
approved_hashes = pd.DataFrame()
|
||||
approved_publishers = pd.DataFrame()
|
||||
|
||||
hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv"
|
||||
path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
|
||||
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv"
|
||||
publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv"
|
||||
|
||||
# Read in and combine the two path generations
|
||||
if os.path.exists(path1):
|
||||
df1 = pd.read_csv(path1)
|
||||
else:
|
||||
logger.warning(f"File not found: {path1}")
|
||||
|
||||
if os.path.exists(path2):
|
||||
df2 = pd.read_csv(path2)
|
||||
else:
|
||||
logger.warning(f"File not found: {path2}")
|
||||
|
||||
if df1.empty and df2.empty:
|
||||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||
approved_paths = pd.DataFrame()
|
||||
else:
|
||||
approved_paths = pd.concat([df1, df2], ignore_index=True)
|
||||
|
||||
approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first")
|
||||
|
||||
# We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions.
|
||||
if os.path.exists(hash):
|
||||
hashes = pd.read_csv(hash)
|
||||
approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])]
|
||||
|
||||
approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first")
|
||||
|
||||
else:
|
||||
logger.warning(f"File not found: {hash}")
|
||||
|
||||
if os.path.exists(publishers):
|
||||
approved_publishers = pd.read_csv(publishers)
|
||||
|
||||
else:
|
||||
logger.warning(f"File not found: {publishers}")
|
||||
|
||||
dataframes = {
|
||||
"approved_paths": approved_paths,
|
||||
"approved_hashes": approved_hashes,
|
||||
"approved_publishers": approved_publishers,
|
||||
}
|
||||
|
||||
for name, df in dataframes.items():
|
||||
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||||
if name == "approved_paths":
|
||||
df.sort_values(by="longestcfp", inplace=True)
|
||||
elif name == "approved_hashes":
|
||||
df.sort_values(by="filename", inplace=True)
|
||||
elif name == "approved_publishers":
|
||||
df.sort_values(by="publisher", inplace=True)
|
||||
|
||||
df.to_csv(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv",
|
||||
index=False,
|
||||
)
|
||||
formatHTML(
|
||||
df,
|
||||
f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html",
|
||||
)
|
||||
|
||||
|
||||
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
|
||||
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type=int)
|
||||
|
||||
def clean_split(path):
|
||||
if not isinstance(path, (str, bytes, os.PathLike)):
|
||||
return []
|
||||
parts = str(os.path.normpath(path)).split(os.sep)
|
||||
parts = [p for p in parts if p] # Remove empty strings
|
||||
return parts
|
||||
|
||||
# Diagnostic: log any non-string entries
|
||||
non_string_entries = df[
|
||||
~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))
|
||||
]
|
||||
if not non_string_entries.empty:
|
||||
print(f"[WARNING] Non-string entries found in column '{col}':")
|
||||
print(non_string_entries)
|
||||
|
||||
df = df.copy()
|
||||
split_paths = df[col].apply(clean_split)
|
||||
|
||||
if min_files_for_path is not None:
|
||||
df = df[
|
||||
split_paths.apply(lambda parts: len(parts) >= min_files_for_path)
|
||||
].copy()
|
||||
split_paths = split_paths[df.index]
|
||||
|
||||
df["group_key"] = split_paths.apply(
|
||||
lambda parts: os.sep.join(parts[:path_exclusion_constant])
|
||||
)
|
||||
grouped = df.groupby("group_key")
|
||||
new_rows = []
|
||||
|
||||
for _, group_df in grouped:
|
||||
paths = group_df[col].tolist()
|
||||
split_parts = [clean_split(p) for p in paths]
|
||||
|
||||
def longest_common_prefix(paths):
|
||||
if not paths:
|
||||
return []
|
||||
prefix = paths[0]
|
||||
for path in paths[1:]:
|
||||
prefix = [a for a, b in zip(prefix, path) if a == b]
|
||||
if not prefix:
|
||||
break
|
||||
return prefix
|
||||
|
||||
common_prefix = longest_common_prefix(split_parts)
|
||||
prefix_str = os.sep.join(common_prefix)
|
||||
|
||||
for i, parts in enumerate(split_parts):
|
||||
filename = parts[-1]
|
||||
middle = (
|
||||
os.sep.join(parts[len(common_prefix) : -1])
|
||||
if len(parts) > len(common_prefix) + 1
|
||||
else ""
|
||||
)
|
||||
row = group_df.iloc[i].copy()
|
||||
row["longestcfp"] = prefix_str
|
||||
row["middle"] = middle
|
||||
row["filename_only"] = filename
|
||||
row["file_extension"] = os.path.splitext(filename)[1].lower()
|
||||
new_rows.append(row)
|
||||
|
||||
return pd.DataFrame(new_rows).drop(columns=["group_key"])
|
||||
|
||||
|
||||
def calculatePath(approved_hashes, path_exclusion_constant, split):
|
||||
if split:
|
||||
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
|
||||
else:
|
||||
dfs_by_policy = [approved_hashes]
|
||||
|
||||
badpathparts = load_env_json("BAD_PATH_PARTS", "[]")
|
||||
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type=int)
|
||||
|
||||
processed_dfs = []
|
||||
|
||||
for df in dfs_by_policy:
|
||||
haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename")
|
||||
haslcp = haslcp.drop_duplicates()
|
||||
|
||||
forbidden = regulator(badpathparts, True)
|
||||
forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
|
||||
|
||||
logger.debug("Removing forbidden filepaths for path exceptions")
|
||||
print(colorText("Removing forbidden filepaths for path exceptions", "green"))
|
||||
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
|
||||
|
||||
lcp_not_forbidden_review = lcp_not_forbidden[
|
||||
[
|
||||
"policyname",
|
||||
"longestcfp",
|
||||
"middle",
|
||||
"filename_only",
|
||||
"file_extension",
|
||||
"sha256",
|
||||
]
|
||||
]
|
||||
|
||||
unique_sha_counts = (
|
||||
lcp_not_forbidden_review.groupby("longestcfp")["sha256"]
|
||||
.nunique()
|
||||
.reset_index()
|
||||
)
|
||||
unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]
|
||||
|
||||
lcp_not_forbidden_review = lcp_not_forbidden_review.merge(
|
||||
unique_sha_counts, on="longestcfp", how="left"
|
||||
)
|
||||
lcp_not_forbidden_review = lcp_not_forbidden_review[
|
||||
lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path
|
||||
]
|
||||
processed_dfs.append(lcp_not_forbidden_review)
|
||||
|
||||
pathExclusions = pd.concat(processed_dfs, ignore_index=True)
|
||||
|
||||
return pathExclusions
|
||||
|
||||
|
||||
def testChange(selected_policies, destination_policy, destination_allowlist):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
logger.info("These path exclusions would be added to:")
|
||||
logger.info(destination_policy)
|
||||
|
||||
pathexclusions = pd.read_csv(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
|
||||
)
|
||||
hashes = pd.read_csv(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
|
||||
)
|
||||
|
||||
unique_combinations = pathexclusions[
|
||||
["longestcfp", "file_extension"]
|
||||
].drop_duplicates()
|
||||
|
||||
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
|
||||
processed_paths = [
|
||||
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
|
||||
for path, ext in unique_combinations.itertuples(index=False, name=None)
|
||||
]
|
||||
|
||||
for path in processed_paths:
|
||||
logger.info(path)
|
||||
|
||||
print(colorText("These publishers would added", "yellow"))
|
||||
processed_publishers = []
|
||||
if os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
|
||||
):
|
||||
publishers = pd.read_csv(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
|
||||
)
|
||||
if publishers.empty:
|
||||
print(colorText("The publishers list is empty.", "red"))
|
||||
else:
|
||||
processed_publishers = (
|
||||
publishers[publishers["publisher"] != "Not Signed"]["publisher"]
|
||||
.drop_duplicates()
|
||||
.tolist()
|
||||
)
|
||||
for publisher in processed_publishers:
|
||||
print(publisher)
|
||||
|
||||
print(colorText("These hashes would be added to:", "yellow"))
|
||||
print(destination_allowlist)
|
||||
|
||||
processed_hashes = hashes["sha256"].unique().tolist()
|
||||
print_x_wide(processed_hashes, 3)
|
||||
|
||||
return processed_paths, processed_hashes, processed_publishers
|
||||
|
||||
|
||||
def menu_policy_enforce(
|
||||
api: AirlockAPIWrapper,
|
||||
): # TODO Need to clean up 6 and 7 into functions
|
||||
selected_policies = []
|
||||
destination_policy = []
|
||||
destination_allowlist = []
|
||||
processed_paths = []
|
||||
processed_hashes = []
|
||||
processed_publishers = []
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
while True:
|
||||
printEnforceChecklist(
|
||||
selected_policies, destination_policy, destination_allowlist
|
||||
)
|
||||
choice = get_sanitized_input("\nEnter your choice: ")
|
||||
|
||||
if choice == "1":
|
||||
clear_screen()
|
||||
selected_policies = selectPolicies(api, True)
|
||||
|
||||
elif choice == "2":
|
||||
clear_screen()
|
||||
print(
|
||||
colorText(
|
||||
"Please choose destination_name Policy for Path Exclusions", "white"
|
||||
)
|
||||
)
|
||||
|
||||
destination_policy = selectPolicies(api, False)
|
||||
|
||||
print(colorText("Please choose Allowlist for Hashes", "white"))
|
||||
|
||||
destination_allowlist = selectAllowlists(api, destination_policy, False)
|
||||
|
||||
elif choice == "3":
|
||||
clear_screen()
|
||||
sortHashes(
|
||||
api,
|
||||
selected_policies,
|
||||
type=[1, 2, 6, 7],
|
||||
)
|
||||
|
||||
elif choice == "4":
|
||||
clear_screen()
|
||||
if os.path.exists(
|
||||
f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"
|
||||
):
|
||||
buildPathsandPublishers(selected_policies, False)
|
||||
else:
|
||||
print(
|
||||
"File not found. Please make sure it's saved correctly and try again."
|
||||
)
|
||||
|
||||
elif choice == "5":
|
||||
clear_screen()
|
||||
if os.path.exists(
|
||||
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
|
||||
) and os.path.exists(
|
||||
f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
|
||||
):
|
||||
buildPreflights(selected_policies)
|
||||
else:
|
||||
print(
|
||||
"File not found. Please make sure it's saved correctly and try again."
|
||||
)
|
||||
|
||||
elif choice == "6":
|
||||
clear_screen()
|
||||
if (
|
||||
os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
|
||||
)
|
||||
and os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
|
||||
)
|
||||
and destination_policy
|
||||
and destination_allowlist
|
||||
):
|
||||
processed_paths, processed_hashes, processed_publishers = testChange(
|
||||
selected_policies, destination_policy, destination_allowlist
|
||||
)
|
||||
else:
|
||||
# Log which condition(s) failed
|
||||
missing_items = []
|
||||
if not os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
|
||||
):
|
||||
missing_items.append("approved_paths.csv not found")
|
||||
if not os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
|
||||
):
|
||||
missing_items.append("approved_hashes.csv not found")
|
||||
if not destination_policy:
|
||||
missing_items.append("destination_policy is empty or None")
|
||||
if not destination_allowlist:
|
||||
missing_items.append("destination_allowlist is empty or None")
|
||||
|
||||
logger.error("Preflight check failed due to the following:")
|
||||
for item in missing_items:
|
||||
logger.error(f" - {item}")
|
||||
|
||||
elif choice == "7":
|
||||
clear_screen()
|
||||
areYouSure()
|
||||
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
if (
|
||||
processed_paths
|
||||
and processed_hashes
|
||||
and processed_publishers
|
||||
and destination_policy
|
||||
and destination_allowlist
|
||||
and confirmation.strip() == "I AGREE"
|
||||
):
|
||||
print(colorText("Proceeding with the code...", "yellow"))
|
||||
api.hash_add_to_allowlist(
|
||||
destination_allowlist[0].applicationid, processed_hashes
|
||||
)
|
||||
api.policy_add_path_exclusions(
|
||||
destination_policy[0].groupid, processed_paths
|
||||
)
|
||||
if processed_publishers:
|
||||
api.policy_add_publishers(
|
||||
destination_policy[0].groupid, processed_publishers
|
||||
)
|
||||
|
||||
locked()
|
||||
|
||||
else:
|
||||
logger.error("Confirmation block failed. Reasons:")
|
||||
if not processed_publishers or processed_hashes or processed_paths:
|
||||
logger.error(" - Test not performed.")
|
||||
if not destination_policy:
|
||||
logger.error(" - `destination_policy` is missing or invalid.")
|
||||
if not destination_allowlist:
|
||||
logger.error(" - `destination_allowlist` is missing or invalid.")
|
||||
if confirmation.strip() != "I AGREE":
|
||||
logger.error(
|
||||
" - User did not confirm with 'I AGREE'. Received: '%s'",
|
||||
confirmation.strip(),
|
||||
)
|
||||
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
elif choice.upper() == "B":
|
||||
break
|
||||
|
||||
else:
|
||||
print(colorText("Invalid choice. Please try again.", "red"))
|
||||
|
||||
|
||||
def section_header(title):
|
||||
print(
|
||||
colorText(
|
||||
"\n --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(f" ------------- {title} -------------", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
section_header("🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒")
|
||||
print(
|
||||
colorText(
|
||||
"\nSequentially follow these steps to prepare a policy for enforcement:",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
|
||||
# Step 1: Originating Policies
|
||||
print(
|
||||
colorText(
|
||||
"\n1. Choose which policy or policies to gather execution info from", "cyan"
|
||||
)
|
||||
)
|
||||
if not selected_policies:
|
||||
print(colorText(" [✗] No policies have been chosen", "red"))
|
||||
else:
|
||||
print(colorText("The following policies have been chosen:", "green"))
|
||||
for policy in selected_policies:
|
||||
print(colorText(f" [✓] {policy.name}", "green"))
|
||||
|
||||
# Step 2: Destination Policy and Allowlist
|
||||
print(
|
||||
colorText("2. Choose the destination policy and associated allowlist", "cyan")
|
||||
)
|
||||
if destination_policy:
|
||||
print(
|
||||
colorText(
|
||||
f" [✓] {destination_policy[0].name} has been selected as the destination policy",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(colorText(" [✗] No destination policy has been chosen", "red"))
|
||||
|
||||
if destination_allowlist:
|
||||
print(
|
||||
colorText(
|
||||
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(colorText(" [✗] No allowlist has been chosen", "red"))
|
||||
|
||||
# Step 3: Data Preparation
|
||||
print(
|
||||
colorText(
|
||||
f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
if selected_policies:
|
||||
policy_id = selected_policies[0].name
|
||||
review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv"
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Data has been fetched"
|
||||
if os.path.exists(review_path)
|
||||
else " [✗] Data has not been fetched"
|
||||
),
|
||||
"green" if os.path.exists(review_path) else "red",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check data fetch status", "red"
|
||||
)
|
||||
)
|
||||
|
||||
# Step 4: Manual Review
|
||||
print(colorText("4. Manually review the files:", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing hashes you do not approve of", "cyan"
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" This will start the process to generate possible filepath approvals",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
if selected_policies:
|
||||
policy_id = selected_policies[0].name
|
||||
approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv"
|
||||
second_review_path = (
|
||||
f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv"
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Reviewed hashes have been loaded"
|
||||
if os.path.exists(approved_path)
|
||||
else " [✗] Reviewed hashes have not been loaded"
|
||||
),
|
||||
"green" if os.path.exists(approved_path) else "red",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Path review list created"
|
||||
if os.path.exists(second_review_path)
|
||||
else " [✗] Path review list has not been created"
|
||||
),
|
||||
"green" if os.path.exists(second_review_path) else "red",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check reviewed hashes or path list",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
|
||||
# Step 5: Path Review
|
||||
print(
|
||||
colorText(
|
||||
f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing path exclusions or publishers you do not approve of.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
f" When complete, save the files to {working_dir}\\data\\Approved",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(" Choose this option when done to build your preflights", "cyan")
|
||||
)
|
||||
|
||||
if selected_policies:
|
||||
policy_id = selected_policies[0].name
|
||||
reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv"
|
||||
preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv"
|
||||
preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv"
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Reviewed path list detected"
|
||||
if os.path.exists(reviewed_path)
|
||||
else " [✗] Path review list has not been detected"
|
||||
),
|
||||
"green" if os.path.exists(reviewed_path) else "red",
|
||||
)
|
||||
)
|
||||
preflight_ready = os.path.exists(preflight_paths) and os.path.exists(
|
||||
preflight_hashes
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Preflight Path Exclusion List has been generated"
|
||||
if preflight_ready
|
||||
else " [✗] Preflight Path Exclusion List has not been generated"
|
||||
),
|
||||
"green" if preflight_ready else "red",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check preflight status", "red"
|
||||
)
|
||||
)
|
||||
|
||||
# Final Steps
|
||||
print(
|
||||
colorText(
|
||||
"6. Test ------------------------------------------------------", "cyan"
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Prints to console the changes that would be made, must be done to proceed. ",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"7. Liftoff ------------------------------------------------------", "cyan"
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Apply path exclusions and approved publishers to selected policy",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(" Apply approved hashes to allowlist", "cyan"))
|
||||
|
||||
# Utility Options
|
||||
print(colorText("F. 📂 - Open Working Directory", "cyan"))
|
||||
print(colorText("B. 🔚 - Back", "cyan"))
|
||||
@@ -1,134 +0,0 @@
|
||||
# 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/>.
|
||||
|
||||
import datetime
|
||||
import logging
|
||||
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
from flows.prepPolicy import selectPolicies
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.policyhandler import getPolicyInfo
|
||||
from utils.configmanager import load_env
|
||||
from utils.selector import Selector
|
||||
from utils.utils import colorText, get_sanitized_input
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
def findQuietAgents(api: AirlockAPIWrapper):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
# Get policy selection and agent list
|
||||
selected_policy = selectPolicies(api, False)
|
||||
if selected_policy:
|
||||
agents = api.agents_find_by_group(selected_policy[0].groupid)
|
||||
|
||||
# Prompt user for history range
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
required_quiet = Selector.select_value(
|
||||
prompt="Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–365): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
|
||||
confirm = Selector.confirm(
|
||||
f"Do you wish to proceed to pull history for {selected_policy[0].name}? Y/N : "
|
||||
)
|
||||
# Get execution history as a DataFrame
|
||||
if confirm:
|
||||
policy_exec_history = getPolicyInfo(
|
||||
api, selected_policy[0], [1, 2, 6, 7], history_days
|
||||
)
|
||||
|
||||
if policy_exec_history.empty:
|
||||
logging.info(
|
||||
"No execution history found for the selected policy and time range."
|
||||
)
|
||||
get_sanitized_input("Press enter to continue")
|
||||
return
|
||||
|
||||
# Convert 'datetime' column to timezone-aware datetime objects
|
||||
policy_exec_history["datetime"] = pd.to_datetime(
|
||||
policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True
|
||||
)
|
||||
|
||||
# Get current UTC time
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
|
||||
# Calculate days ago
|
||||
policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
|
||||
lambda dt: (now - dt).days
|
||||
)
|
||||
|
||||
# Count total executions per hostname
|
||||
hostname_counts = policy_exec_history["hostname"].value_counts()
|
||||
|
||||
# Map execution counts to agents
|
||||
agents["execution_count"] = (
|
||||
agents["hostname"].map(hostname_counts).fillna(0).astype(int)
|
||||
)
|
||||
|
||||
# Find most recent execution per hostname
|
||||
most_recent_exec = policy_exec_history.sort_values(
|
||||
by="days_ago"
|
||||
).drop_duplicates(subset="hostname", keep="first")
|
||||
|
||||
# Map most recent execution age to agents
|
||||
agents["days_since"] = agents["hostname"].map(
|
||||
most_recent_exec.set_index("hostname")["days_ago"]
|
||||
)
|
||||
|
||||
# Check for enforcement readiness
|
||||
agents["required_quiet"] = required_quiet
|
||||
agents["enforce_ready"] = agents["days_since"].apply(
|
||||
lambda x: True if pd.isna(x) or x > required_quiet else False
|
||||
)
|
||||
|
||||
# Sort agents by execution count and hostname
|
||||
agents = agents.sort_values(
|
||||
by=["execution_count", "hostname"], ascending=[True, True]
|
||||
)
|
||||
|
||||
# Save to CSV
|
||||
filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv"
|
||||
logging.debug(f"Saving CSV to {filename}")
|
||||
print(colorText(f"Saving CSV to {filename}", "green"))
|
||||
agents.to_csv(filename, index=False)
|
||||
|
||||
# Summary statistics
|
||||
total_agents = len(agents)
|
||||
ready_agents = agents["enforce_ready"].sum()
|
||||
not_ready_agents = total_agents - ready_agents
|
||||
ready_percentage = (ready_agents / total_agents) * 100
|
||||
|
||||
# Print results
|
||||
|
||||
message = (
|
||||
f"Total agents: {total_agents}\n"
|
||||
f"Agents marked as 'enforce_ready': {ready_agents}\n"
|
||||
f"Agents not ready: {not_ready_agents}\n"
|
||||
f"Percentage ready for enforcement: {ready_percentage:.2f}%"
|
||||
)
|
||||
logger.debug(message)
|
||||
colorText(message, "green")
|
||||
get_sanitized_input("Press enter to continue")
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 118 KiB |
+8
-14
@@ -28,8 +28,8 @@ import pandas as pd
|
||||
|
||||
import airlock_libs
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_protected_value, load_env_json
|
||||
from utils.utils import colorText, regulator
|
||||
from utils.configmanager import get_system_list, get_system_value
|
||||
from utils.utils import regulator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -90,9 +90,9 @@ class Hash:
|
||||
|
||||
@classmethod
|
||||
def categorize_hashes(cls, hashes):
|
||||
threat_tolerance = get_protected_value("VT_THREAT_TOLERANCE", cast_type=int)
|
||||
bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
|
||||
pups_pattern = regulator(load_env_json("PUPS", "[]"))
|
||||
threat_tolerance = get_system_value("VT_THREAT_TOLERANCE", cast_type=int)
|
||||
bad_publishers_pattern = regulator(get_system_list("BAD_PUBLISHERS"))
|
||||
pups_pattern = regulator(get_system_list("PUPS"))
|
||||
|
||||
approved_count = 0
|
||||
unapproved_count = 0
|
||||
@@ -293,12 +293,6 @@ class ExecutionHistoryRecord:
|
||||
logger.debug(
|
||||
f"Staging of Execution history for policy: {policy.name} is complete"
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
f"Staging of Execution history for policy: {policy.name} is complete",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
|
||||
return executions
|
||||
|
||||
@@ -384,9 +378,9 @@ class ExecutionHistoryRecord:
|
||||
Returns:
|
||||
List[ExecutionHistoryRecord]: The same list, with hash_obj.at_decision updated.
|
||||
"""
|
||||
threat_tolerance = get_protected_value("VT_THREAT_TOLERANCE", cast_type=int)
|
||||
bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
|
||||
pups_pattern = regulator(load_env_json("PUPS", "[]"))
|
||||
threat_tolerance = get_system_value("VT_THREAT_TOLERANCE", cast_type=int)
|
||||
bad_publishers_pattern = regulator(get_system_list("BAD_PUBLISHERS"))
|
||||
pups_pattern = regulator(get_system_list("PUPS"))
|
||||
|
||||
approved_count = 0
|
||||
unapproved_count = 0
|
||||
|
||||
+19
-7
@@ -1,14 +1,26 @@
|
||||
# Core TUI dependencies
|
||||
textual==6.5.0
|
||||
|
||||
# API and data handling
|
||||
Requests==2.32.5
|
||||
pandas==2.3.3
|
||||
numpy==2.3.4
|
||||
|
||||
# Database
|
||||
pymongo==4.15.3
|
||||
|
||||
# Security and encryption
|
||||
cryptography==46.0.3
|
||||
keyring==25.6.0
|
||||
numpy==2.3.4
|
||||
pandas==2.3.3
|
||||
pymongo==4.15.3
|
||||
|
||||
# Environment management
|
||||
python-dotenv==1.2.1
|
||||
Requests==2.32.5
|
||||
textual==6.5.0
|
||||
tqdm==4.67.1
|
||||
|
||||
# Utilities
|
||||
urllib3==2.5.0
|
||||
plotext==5.3.2
|
||||
pyperclip==1.11.0
|
||||
|
||||
# Custom/Private packages
|
||||
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
|
||||
airlock_libs==2.0.0
|
||||
airlock_libs==7.4.1
|
||||
@@ -1,61 +0,0 @@
|
||||
from typing import List
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.screen import Screen
|
||||
|
||||
from models.agent import Agent
|
||||
from widgets.agentmoveoperations import AgentMoveOperations
|
||||
from widgets.multiagentselector import MultiAgentSelector
|
||||
from widgets.resultsdisplay import ResultsDisplay
|
||||
|
||||
|
||||
class MoveAgentWorkflowScreen(Screen):
|
||||
"""Screen that handles the agent movement workflow."""
|
||||
|
||||
def __init__(self, all_agents: List[Agent]):
|
||||
super().__init__()
|
||||
self.all_agents = all_agents
|
||||
self.selected_agents = None
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Start with the multi-agent selector."""
|
||||
yield MultiAgentSelector(self.all_agents)
|
||||
|
||||
def on_multi_agent_selector_agents_selected(
|
||||
self, message: MultiAgentSelector.AgentsSelected
|
||||
) -> None:
|
||||
"""Handle selected agents - switch to operations screen."""
|
||||
self.selected_agents = message.selected_agents
|
||||
|
||||
# Remove the MultiAgentSelector
|
||||
selector = self.query_one(MultiAgentSelector)
|
||||
selector.remove()
|
||||
|
||||
# Mount the AgentMoveOperations with the selected Agent objects
|
||||
self.mount(AgentMoveOperations(self.selected_agents))
|
||||
|
||||
def on_agent_move_operations_operation_complete(
|
||||
self, message: AgentMoveOperations.OperationComplete
|
||||
) -> None:
|
||||
"""Handle completion of move operation - transition to results screen."""
|
||||
# Format successful results
|
||||
success_lines = []
|
||||
for agent, result in message.successful:
|
||||
success_lines.append(f"✓ {agent.hostname}")
|
||||
|
||||
# Format unsuccessful results
|
||||
failure_lines = []
|
||||
for agent, error in message.unsuccessful:
|
||||
failure_lines.append(f"✗ {agent.hostname}: {error}")
|
||||
|
||||
successful_text = "\n".join(success_lines) if success_lines else "(none)"
|
||||
unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)"
|
||||
|
||||
# Remove the operations widget
|
||||
ops_widget = self.query_one(AgentMoveOperations)
|
||||
ops_widget.remove()
|
||||
|
||||
# Mount the results display
|
||||
self.mount(
|
||||
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
|
||||
)
|
||||
@@ -1,24 +0,0 @@
|
||||
# otp_workflow_screen.py
|
||||
|
||||
from typing import List
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.screen import Screen
|
||||
|
||||
from models.agent import Agent
|
||||
from widgets.OTP_generate import OTPGenerator
|
||||
|
||||
|
||||
class OTPWorkflowScreen(Screen):
|
||||
"""Screen that handles the OTP generation workflow without agent selection."""
|
||||
|
||||
def __init__(self, selected_agents: List[Agent]):
|
||||
super().__init__()
|
||||
self.selected_agents = selected_agents
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Directly show the OTP generator for the selected agents."""
|
||||
yield OTPGenerator(self.selected_agents)
|
||||
|
||||
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
|
||||
"""Handle OTP generation request - pass it up to the app level if needed."""
|
||||
+72
-2
@@ -64,6 +64,23 @@ class AirlockAPIWrapper:
|
||||
logger.error(f"API request failed: {e}")
|
||||
raise
|
||||
|
||||
def _post_raw(self, endpoint: str, payload: Optional[dict] = None) -> bytes:
|
||||
"""
|
||||
Send POST request and return raw response content (bytes).
|
||||
Useful for XML endpoints.
|
||||
"""
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
data = json.dumps(payload or {})
|
||||
try:
|
||||
logger.debug(f"POST Request to {url} with payload: {payload}")
|
||||
response = requests.post(url, headers=self.headers, data=data, verify=False)
|
||||
response.raise_for_status()
|
||||
logger.debug(f"Raw response received from {url}")
|
||||
return response.content # bytes
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"API request failed: {e}")
|
||||
raise
|
||||
|
||||
# Allowlist Management
|
||||
def allowlist_find_all(self) -> pd.DataFrame:
|
||||
"""
|
||||
@@ -75,6 +92,12 @@ class AirlockAPIWrapper:
|
||||
result = self._post("/v1/application", {})
|
||||
return pd.DataFrame(result["response"]["applications"])
|
||||
|
||||
def allowlist_export(self, applicationid) -> bytes:
|
||||
"""Return Allowlist XML as bytes to save to file"""
|
||||
payload = {"applicationid": applicationid}
|
||||
result = self._post_raw("/v1/application/export", payload)
|
||||
return result # should be bytes
|
||||
|
||||
# Agent Management
|
||||
def agent_find_all(self) -> pd.DataFrame:
|
||||
"""Retrieve all agents."""
|
||||
@@ -116,6 +139,30 @@ class AirlockAPIWrapper:
|
||||
result = self._post("/v1/agent/find", payload)
|
||||
return pd.DataFrame(result["response"]["agents"])
|
||||
|
||||
# Baseline Managment
|
||||
def baseline_find_all(self) -> pd.DataFrame:
|
||||
"""Retrieve all Baselines."""
|
||||
result = self._post("/v1/baseline", {})
|
||||
return pd.DataFrame(result["response"]["baselines"])
|
||||
|
||||
def baseline_export(self, baselineid) -> bytes:
|
||||
"""Return Baseline XML as bytes to save to file"""
|
||||
payload = {"baselineid": baselineid}
|
||||
result = self._post_raw("/v1/baseline/export", payload)
|
||||
return result
|
||||
|
||||
# Blocklist Managment
|
||||
def blocklist_find_all(self) -> pd.DataFrame:
|
||||
"""Retrieve all Baselines."""
|
||||
result = self._post("/v1/blocklist", {})
|
||||
return pd.DataFrame(result["response"]["blocklists"])
|
||||
|
||||
def blocklist_export(self, blocklistid) -> bytes:
|
||||
"""Return Blocklist XML as bytes to save to file"""
|
||||
payload = {"blocklistid": blocklistid}
|
||||
result = self._post_raw("/v1/blocklist/export", payload)
|
||||
return result
|
||||
|
||||
# Hash Management
|
||||
def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict:
|
||||
"""Add hashes to the allowlist for a specific application."""
|
||||
@@ -253,12 +300,27 @@ class AirlockAPIWrapper:
|
||||
}
|
||||
return self._post("/v1/group/settings/script_custom", payload)
|
||||
|
||||
def policy_set_upgradetarget(
|
||||
self,
|
||||
groupid: str,
|
||||
windows: str,
|
||||
macos: str,
|
||||
) -> dict:
|
||||
payload = {
|
||||
"groupid": groupid,
|
||||
"windows": windows,
|
||||
"macos": macos,
|
||||
}
|
||||
return self._post("/v1/group/settings/selfupgrade/target", payload)
|
||||
|
||||
# Execution History
|
||||
def history_logging(
|
||||
self, type: List[str], checkpoint: str, policy: List[str]
|
||||
self, type: List[str], checkpoint: str, policy: Optional[List[str]] = None
|
||||
) -> str:
|
||||
"""Retrieve execution history logs."""
|
||||
payload = {"type": type, "checkpoint": checkpoint, "policy": policy}
|
||||
payload = {"type": type, "checkpoint": checkpoint}
|
||||
if policy is not None:
|
||||
payload["policy"] = policy
|
||||
result = self._post("/v1/logging/exechistories", payload)
|
||||
return result["response"]["exechistories"]
|
||||
|
||||
@@ -291,6 +353,14 @@ class AirlockAPIWrapper:
|
||||
result = self._post("/v1/getexechistory", payload)
|
||||
return result["response"]["exechistory"]
|
||||
|
||||
def server_logs(self, checkpoint: str | None = None) -> str:
|
||||
"""Retrieves Server Activity History Logs."""
|
||||
payload = {}
|
||||
if checkpoint is not None:
|
||||
payload["checkpoint"] = checkpoint
|
||||
result = self._post("/v1/logging/svractivities?checkpoint", payload)
|
||||
return result["response"]["svractivities"]
|
||||
|
||||
|
||||
"""
|
||||
from services.API import AirlockAPIWrapper
|
||||
|
||||
@@ -1,388 +0,0 @@
|
||||
# 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 dataclasses import asdict
|
||||
from datetime import datetime, timedelta
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import List
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from flows.prepPolicy import selectPolicies
|
||||
from models.agent import Agent
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_protected_json, load_env
|
||||
from utils.selector import Selector
|
||||
from utils.utils import colorText, get_sanitized_input
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||||
agents = selectAgents(api)
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
|
||||
if not agents or not history_days:
|
||||
print(colorText("No agents selected or invalid history range.", "red"))
|
||||
return
|
||||
|
||||
historical_date = (datetime.now() - timedelta(days=history_days)).strftime(
|
||||
"%Y-%m-%d"
|
||||
)
|
||||
today = datetime.now().strftime("%Y-%m-%d")
|
||||
|
||||
all_history = []
|
||||
|
||||
for agent in agents:
|
||||
try:
|
||||
exechistory = api.history_execution(today, historical_date, agent.hostname)
|
||||
except Exception as e:
|
||||
print(
|
||||
colorText(
|
||||
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
if isinstance(exechistory, list):
|
||||
for block in exechistory:
|
||||
record = {
|
||||
"Command": block.get("commandline", "N/A"),
|
||||
"Date": block.get("datetime", "N/A"),
|
||||
"Filename": block.get("filename", "N/A"),
|
||||
"Policy Name": block.get("policyname", "N/A"),
|
||||
"Hostname": block.get("hostname", "N/A"),
|
||||
"Hash": block.get("sha256", "N/A"),
|
||||
}
|
||||
all_history.append(record)
|
||||
|
||||
if not outputjson:
|
||||
for key, value in record.items():
|
||||
print(colorText(f"{key}: {value}", "green"))
|
||||
print("\n")
|
||||
else:
|
||||
print(
|
||||
colorText(f"No execution history found for {agent.hostname}.", "yellow")
|
||||
)
|
||||
|
||||
if outputjson:
|
||||
print(json.dumps(all_history, indent=2))
|
||||
|
||||
|
||||
def findAllAgents(api):
|
||||
# Step 1: Load data from API
|
||||
policies = [Policy(**row["data"]) for _, row in api.policy_find_all().iterrows()]
|
||||
agents = [Agent(**row["data"]) for _, row in api.agent_find_all().iterrows()]
|
||||
|
||||
for agent in agents:
|
||||
agent.enrich_with_policies(policies)
|
||||
|
||||
return agents
|
||||
|
||||
|
||||
def findAgents(api, return_dataframe):
|
||||
agents = selectAgents(api)
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
if not agents:
|
||||
logging.warning("No agents or policies found.")
|
||||
print("No agents matched the criteria.")
|
||||
return
|
||||
|
||||
# Convert enriched agents to DataFrame
|
||||
agent_dicts = [asdict(agent) for agent in agents]
|
||||
agent_df = pd.DataFrame(agent_dicts)
|
||||
|
||||
if return_dataframe:
|
||||
logging.debug("Returning DataFrame to caller.")
|
||||
return agent_df
|
||||
|
||||
# Otherwise, print and optionally export
|
||||
print(agent_df)
|
||||
logging.debug("Displayed DataFrame to console.")
|
||||
|
||||
user_input = (
|
||||
get_sanitized_input(
|
||||
"\nWould you like to export the results to a CSV file? (y/n): "
|
||||
)
|
||||
.strip()
|
||||
.lower()
|
||||
)
|
||||
if user_input == "y":
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
filename = f"agentsearch_{timestamp}.csv"
|
||||
file_path = os.path.join(str(working_dir), filename)
|
||||
|
||||
agent_df.to_csv(file_path, index=False)
|
||||
logging.info(f"Exported DataFrame to {file_path}")
|
||||
|
||||
print(
|
||||
colorText(
|
||||
f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
logging.debug("User declined to export the DataFrame.")
|
||||
|
||||
|
||||
def collect_device_names() -> List[str]:
|
||||
print(colorText("🔍 Device Search", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
"Enter the device hostnames you'd like to search for, one per line.", "cyan"
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"When you're done, press Enter twice (Three times if you have a single device).\n",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText("Example:", "cyan"))
|
||||
print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan"))
|
||||
print(colorText("Paste or type your device names below:", "white"))
|
||||
|
||||
device_input_lines = []
|
||||
empty_line_count = 0
|
||||
valid_line_pattern = re.compile(r"^[a-zA-Z0-9_\- ]+$")
|
||||
|
||||
while True:
|
||||
line = get_sanitized_input("")
|
||||
stripped_line = line.strip()
|
||||
|
||||
if stripped_line == "":
|
||||
empty_line_count += 1
|
||||
if empty_line_count == 2:
|
||||
break
|
||||
continue
|
||||
else:
|
||||
empty_line_count = 0
|
||||
|
||||
if valid_line_pattern.match(stripped_line):
|
||||
device_input_lines.append(stripped_line)
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
|
||||
return [name for name in device_input_lines if name]
|
||||
|
||||
|
||||
def choose_match_type() -> bool:
|
||||
print(colorText("Use exact match? (Y for exact, N for fuzzy):", "white"))
|
||||
return get_sanitized_input("").strip().lower() in ["y", "yes"]
|
||||
|
||||
|
||||
def match_agents(
|
||||
device_names: List[str], agents: List["Agent"], use_exact: bool
|
||||
) -> List["Agent"]:
|
||||
if use_exact:
|
||||
return [
|
||||
agent
|
||||
for agent in agents
|
||||
if agent.hostname.lower() in [name.lower() for name in device_names]
|
||||
]
|
||||
else:
|
||||
pattern = "|".join(map(re.escape, device_names))
|
||||
regex = re.compile(pattern, re.IGNORECASE)
|
||||
return [agent for agent in agents if regex.search(agent.hostname)]
|
||||
|
||||
|
||||
def show_unmatched(
|
||||
device_names: List[str], matched_agents: List["Agent"], use_exact: bool
|
||||
):
|
||||
if use_exact:
|
||||
unmatched = [
|
||||
name
|
||||
for name in device_names
|
||||
if not any(
|
||||
agent.hostname.lower() == name.lower() for agent in matched_agents
|
||||
)
|
||||
]
|
||||
else:
|
||||
unmatched = [
|
||||
name
|
||||
for name in device_names
|
||||
if not any(
|
||||
re.search(re.escape(name), agent.hostname, re.IGNORECASE)
|
||||
for agent in matched_agents
|
||||
)
|
||||
]
|
||||
|
||||
if unmatched:
|
||||
logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
|
||||
print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
|
||||
|
||||
|
||||
def enrich_agents(agents: List["Agent"], policies: List["Policy"]):
|
||||
for agent in agents:
|
||||
agent.enrich_with_policies(policies)
|
||||
|
||||
|
||||
def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
||||
device_names = collect_device_names()
|
||||
if not device_names:
|
||||
logger.debug("No device names entered")
|
||||
print(colorText("⚠️ No device names entered.", "red"))
|
||||
return []
|
||||
|
||||
use_exact = choose_match_type()
|
||||
|
||||
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||||
agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()]
|
||||
matched_agents = match_agents(device_names, agents, use_exact)
|
||||
matched_agents.sort(key=lambda agent: agent.hostname.lower())
|
||||
|
||||
show_unmatched(device_names, matched_agents, use_exact)
|
||||
|
||||
if not matched_agents:
|
||||
logger.debug("❌ No matching devices found.")
|
||||
print(colorText("❌ No matching devices found.", "red"))
|
||||
return []
|
||||
|
||||
print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green"))
|
||||
logger.info("Matched agent hostnames:")
|
||||
rows = (len(matched_agents) + 2) // 3 # 3 columns
|
||||
for row in range(rows):
|
||||
line = ""
|
||||
for col in range(3):
|
||||
idx = row + col * rows
|
||||
if idx < len(matched_agents):
|
||||
line += f"{matched_agents[idx].hostname:<30}"
|
||||
logger.info(line)
|
||||
|
||||
matched_agents = Selector.select_with_mode(
|
||||
matched_agents,
|
||||
label_func=lambda agent: agent.hostname,
|
||||
header="Matched Devices:",
|
||||
)
|
||||
|
||||
if not matched_agents:
|
||||
logger.debug("❌ No matching devices remain after refinement.")
|
||||
print(colorText("❌ No matching devices remain after refinement.", "red"))
|
||||
return []
|
||||
|
||||
enrich_agents(matched_agents, policies)
|
||||
return matched_agents
|
||||
|
||||
|
||||
def moveAgentToRelatedPolicy(
|
||||
api: AirlockAPIWrapper,
|
||||
agent: Agent,
|
||||
mode: str = "audit",
|
||||
):
|
||||
"""
|
||||
Moves an agent between audit and enforcement policies based on the mode.
|
||||
|
||||
Args:
|
||||
api: AirlockAPIWrapper instance.
|
||||
agent: Agent object.
|
||||
policy_relationship_map: Dict mapping enforcement → audit.
|
||||
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
|
||||
"""
|
||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
|
||||
if mode == "audit":
|
||||
if agent.groupid in policy_relationship_map:
|
||||
target_policy = policy_relationship_map[agent.groupid]
|
||||
elif agent.groupid in policy_relationship_map.values():
|
||||
logger.debug(
|
||||
f"Agent {agent.hostname} is already in an audit group. No action needed."
|
||||
)
|
||||
print(
|
||||
f"Agent {agent.hostname} is already in an audit group. No action needed."
|
||||
)
|
||||
return
|
||||
else:
|
||||
logger.warning(
|
||||
f"Error: No corresponding audit policy found for groupid: {agent.groupid}."
|
||||
)
|
||||
return
|
||||
|
||||
elif mode == "enforcement":
|
||||
inverse_map = {v: k for k, v in policy_relationship_map.items()}
|
||||
if agent.groupid in inverse_map:
|
||||
target_policy = inverse_map[agent.groupid]
|
||||
elif agent.groupid in inverse_map.values():
|
||||
logger.info(
|
||||
f"Agent {agent.hostname} is already in an enforcement group. No action needed."
|
||||
)
|
||||
return
|
||||
else:
|
||||
logger.warning(
|
||||
f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}."
|
||||
)
|
||||
return
|
||||
|
||||
else:
|
||||
logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.")
|
||||
return
|
||||
|
||||
result = api.agent_move(agent.agentid, target_policy)
|
||||
return result
|
||||
|
||||
|
||||
def toggleEnforcement(api: AirlockAPIWrapper):
|
||||
choices = ["Audit", "Enforcement", "Exit"]
|
||||
print(colorText("Move devices to which state?:", "yellow"))
|
||||
direction = Selector.select_string(choices, False, False)
|
||||
if direction == "Exit":
|
||||
pass
|
||||
else:
|
||||
devices = selectAgents(api)
|
||||
for device in devices:
|
||||
print(device.hostname)
|
||||
confirm = Selector.confirm(
|
||||
"Would you like to continue with these devices? Y/N: "
|
||||
)
|
||||
if direction and devices and confirm:
|
||||
for device in devices:
|
||||
result = moveAgentToRelatedPolicy(api, device, str(direction).lower())
|
||||
logger.info(f"{device.hostname}: result: {result}")
|
||||
get_sanitized_input("Press enter to continue")
|
||||
|
||||
|
||||
def moveAgents(api: AirlockAPIWrapper):
|
||||
devices = selectAgents(api)
|
||||
for device in devices:
|
||||
print(device.hostname)
|
||||
confirm_devices = Selector.confirm(
|
||||
"Would you like to continue with these devices? Y/N: "
|
||||
)
|
||||
if devices and confirm_devices:
|
||||
policies = selectPolicies(api, False)
|
||||
confirm_move = Selector.confirm(
|
||||
f"Would you like to move these devices to {policies[0].name}?"
|
||||
)
|
||||
if confirm_move:
|
||||
for device in devices:
|
||||
result = api.agent_move(device.agentid, policies[0].groupid)
|
||||
logger.info(f"{device.hostname}: result: {result}")
|
||||
else:
|
||||
logger.info("Exiting without change")
|
||||
get_sanitized_input("Press enter to continue")
|
||||
@@ -1,249 +0,0 @@
|
||||
# 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/>.
|
||||
|
||||
|
||||
import datetime
|
||||
import gc
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
from bson import ObjectId
|
||||
import pandas as pd
|
||||
import tqdm
|
||||
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_protected_json
|
||||
from utils.setup import get_base_directory
|
||||
from utils.utils import areYouSure, colorText, get_sanitized_input
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def pullPolicyExechistories(
|
||||
api: AirlockAPIWrapper,
|
||||
policy: Policy,
|
||||
type: list,
|
||||
days,
|
||||
outputjson: bool,
|
||||
):
|
||||
|
||||
file_path = f"{get_base_directory()}\\cache\\chunkinator.json"
|
||||
|
||||
# Ensure the file exists
|
||||
if not os.path.exists(file_path):
|
||||
with open(file_path, "w") as file:
|
||||
json.dump({"error": "Success", "response": {"exechistories": []}}, file)
|
||||
logger.debug(f"File '{file_path}' has been created.")
|
||||
else:
|
||||
logger.debug(f"File '{file_path}' already exists.")
|
||||
|
||||
checkpoint = str(skipback(days))
|
||||
json_output = {"error": "Success", "response": {"exechistories": []}}
|
||||
|
||||
with tqdm.tqdm(
|
||||
file=sys.stdout,
|
||||
leave=True,
|
||||
total=10000,
|
||||
desc=f"Checkpoint Progress: {checkpoint}",
|
||||
colour="blue",
|
||||
initial=1,
|
||||
) as filebar:
|
||||
with tqdm.tqdm(
|
||||
file=sys.stdout,
|
||||
leave=True,
|
||||
total=100,
|
||||
desc=f"Total of {policy} Complete: ",
|
||||
) as pbar:
|
||||
while True:
|
||||
histories = api.history_logging(
|
||||
type=type, checkpoint=checkpoint, policy=[policy.name]
|
||||
)
|
||||
|
||||
# Ensure histories is a list of dictionaries
|
||||
if not isinstance(histories, list) or not all(
|
||||
isinstance(h, dict) for h in histories
|
||||
):
|
||||
logger.error(
|
||||
"Unexpected response format from API. Expected list of dictionaries."
|
||||
)
|
||||
break
|
||||
|
||||
filebar.total = len(histories)
|
||||
|
||||
if not histories:
|
||||
break
|
||||
|
||||
for index, history_item in enumerate(histories):
|
||||
if (
|
||||
"checkpoint" not in history_item
|
||||
or "datetime" not in history_item
|
||||
):
|
||||
continue # Skip malformed entries
|
||||
|
||||
# Update checkpoint on last item
|
||||
if index == len(histories) - 1:
|
||||
checkpoint = history_item[
|
||||
"checkpoint"
|
||||
] # pyright: ignore[reportArgumentType]
|
||||
filebar.desc = f"Checkpoint Progress: {checkpoint}"
|
||||
break
|
||||
|
||||
try:
|
||||
history_date = datetime.datetime.strptime(
|
||||
history_item["datetime"].replace(
|
||||
" +0000 UTC", ""
|
||||
), # pyright: ignore[reportArgumentType]
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
).date()
|
||||
except ValueError:
|
||||
continue # Skip if date format is invalid
|
||||
|
||||
if (
|
||||
datetime.date.today() - datetime.timedelta(days=days)
|
||||
) <= history_date:
|
||||
json_output["response"]["exechistories"].append(history_item)
|
||||
|
||||
filebar.update(1)
|
||||
filebar.refresh()
|
||||
|
||||
# Deduplicate entries
|
||||
seen = {}
|
||||
if os.path.exists(file_path):
|
||||
with open(file_path, "r") as file:
|
||||
existing_data = json.load(file)
|
||||
combined = (
|
||||
existing_data["response"]["exechistories"]
|
||||
+ json_output["response"]["exechistories"]
|
||||
)
|
||||
else:
|
||||
combined = json_output["response"]["exechistories"]
|
||||
|
||||
for entry in combined:
|
||||
key = (
|
||||
entry.get("sha256"),
|
||||
entry.get("filename"),
|
||||
entry.get("hostname"),
|
||||
)
|
||||
seen[key] = entry
|
||||
|
||||
deduplicated = list(seen.values())
|
||||
with open(file_path, "w") as file:
|
||||
json.dump(
|
||||
{
|
||||
"error": "Success",
|
||||
"response": {"exechistories": deduplicated},
|
||||
},
|
||||
file,
|
||||
)
|
||||
|
||||
json_output["response"]["exechistories"].clear()
|
||||
|
||||
# Update progress bar based on last valid item
|
||||
try:
|
||||
last_date = datetime.datetime.strptime(
|
||||
history_item["datetime"].replace(" +0000 UTC", ""), # type: ignore
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
).date()
|
||||
date_diff = datetime.date.today() - last_date
|
||||
percentage_diff = (
|
||||
((days + 10) - date_diff.days) / (days + 10)
|
||||
) * 100
|
||||
pbar.n = round(percentage_diff)
|
||||
pbar.set_description_str(f"Total of {policy} Complete: ")
|
||||
pbar.refresh()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
filebar.n = 1
|
||||
|
||||
# Final output
|
||||
with open(file_path, "r") as file:
|
||||
final_output = json.load(file)
|
||||
os.remove(file_path)
|
||||
|
||||
return json.dumps(final_output) if outputjson else None
|
||||
|
||||
|
||||
def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
|
||||
import airlock_libs
|
||||
|
||||
executionhist_policy = pd.DataFrame()
|
||||
exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days)
|
||||
if exehist is not None:
|
||||
data = json.loads(exehist)
|
||||
executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
|
||||
if not executionhist_policy.empty:
|
||||
executionhist_policy = executionhist_policy[
|
||||
[
|
||||
"datetime",
|
||||
"sha256",
|
||||
"publisher",
|
||||
"filename",
|
||||
"hostname",
|
||||
"username",
|
||||
"pprocess",
|
||||
"gprocess",
|
||||
"commandline",
|
||||
]
|
||||
]
|
||||
executionhist_policy["policy"] = policy # Add policy column here
|
||||
executionhist_policy = executionhist_policy.drop_duplicates(
|
||||
subset=["sha256", "filename", "hostname"]
|
||||
)
|
||||
executionhist_policy = executionhist_policy.sort_values(
|
||||
by=["sha256", "filename"]
|
||||
)
|
||||
logger.debug(f"Staging of Execution history for policy: {policy} is complete")
|
||||
print(
|
||||
colorText(
|
||||
f"Staging of Execution history for policy: {policy} is complete",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
del data
|
||||
del exehist
|
||||
gc.collect()
|
||||
return executionhist_policy
|
||||
|
||||
|
||||
def skipback(days):
|
||||
"""
|
||||
Generate a MongoDB ObjectId for a given number of days ago from today.
|
||||
"""
|
||||
adjusted_days = days
|
||||
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
|
||||
days=adjusted_days
|
||||
)
|
||||
timestamp = int(date_days_ago.timestamp())
|
||||
hex_timestamp = format(timestamp, "08x")
|
||||
objectid_hex = hex_timestamp + "0000000000000000"
|
||||
return ObjectId(objectid_hex)
|
||||
|
||||
|
||||
def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
|
||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
for enforcement_policy, audit_policy in policy_relationship_map.items():
|
||||
api.policy_clone(enforcement_policy, audit_policy)
|
||||
api.policy_set_auditmode(audit_policy, "1")
|
||||
|
||||
|
||||
def confirmUpdateAfromE(api: AirlockAPIWrapper):
|
||||
areYouSure()
|
||||
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
if confirmation.strip() == "I AGREE":
|
||||
updateAuditPoliciesFromEnforcementPolices(api)
|
||||
@@ -169,17 +169,3 @@ def getAPI(USERNAME, SERVICE_NAME):
|
||||
logging.warning(
|
||||
"Password does not meet complexity requirements. Try again."
|
||||
)
|
||||
|
||||
|
||||
class APIKeyManager:
|
||||
_api_key = None
|
||||
|
||||
@classmethod
|
||||
def load(cls, service: str, username: str, password: str):
|
||||
cls._api_key = retrieve_api_key(service, username, password)
|
||||
|
||||
@classmethod
|
||||
def get(cls) -> str:
|
||||
if cls._api_key is None:
|
||||
raise ValueError("API key not loaded. Call APIKeyManager.load() first.")
|
||||
return cls._api_key
|
||||
|
||||
+241
-42
@@ -18,117 +18,280 @@ import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from typing import Callable, Optional, TypeVar
|
||||
from typing import Any, Callable, Optional, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PROTECTED_KEYS = [
|
||||
# System config keys - these are immutable and come from system_config.json (bundled in exe)
|
||||
SYSTEM_CONFIG_KEYS = [
|
||||
"URL",
|
||||
"APPNAME",
|
||||
"LOG_LEVEL",
|
||||
"BAD_PATH_PARTS",
|
||||
"BAD_PUBLISHERS",
|
||||
"PUPS",
|
||||
"PATH_EXCLUSION_CONST",
|
||||
"MIN_FILES_FOR_PATH",
|
||||
"VT_THREAT_TOLERANCE",
|
||||
"POLICY_MAP_ENF_AUD",
|
||||
]
|
||||
|
||||
_protected_config = {}
|
||||
# User config keys - these can be changed by the end user
|
||||
USER_CONFIG_KEYS = [
|
||||
"TELEMETRY", # User opt-in/out for telemetry
|
||||
"TELEM_URL",
|
||||
"TEXTUAL_THEME", # UI theme preference
|
||||
"EXTRAS", # Feature flags
|
||||
]
|
||||
|
||||
# In-memory config storage
|
||||
_system_config = {}
|
||||
_user_config = {}
|
||||
|
||||
|
||||
def get_system_config_path() -> Path:
|
||||
"""
|
||||
Get path to system_config.json.
|
||||
Priority:
|
||||
1. Bundled in exe (_MEIPASS)
|
||||
2. Next to this file (development)
|
||||
"""
|
||||
# Check inside bundled EXE directory first
|
||||
bundled_dir = Path(getattr(sys, "_MEIPASS", ""))
|
||||
bundled_path = bundled_dir / "system_config.json"
|
||||
if bundled_path.exists():
|
||||
return bundled_path
|
||||
|
||||
# Fallback to external location
|
||||
# Fallback to development location (next to this file)
|
||||
return Path(__file__).parent.parent / "system_config.json"
|
||||
|
||||
|
||||
def load_protected_config() -> dict:
|
||||
global _protected_config
|
||||
def load_system_config() -> dict:
|
||||
"""
|
||||
Load system configuration from system_config.json.
|
||||
This should only be called once at startup.
|
||||
Returns the full system config dict.
|
||||
"""
|
||||
global _system_config
|
||||
|
||||
try:
|
||||
with open(get_system_config_path(), "r") as f:
|
||||
system_config = json.load(f)
|
||||
config_path = get_system_config_path()
|
||||
with open(config_path, "r") as f:
|
||||
_system_config = json.load(f)
|
||||
logger.debug(f"✅ Loaded system config from {config_path}")
|
||||
except FileNotFoundError:
|
||||
logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
|
||||
system_config = {
|
||||
logger.warning("⚠️ system_config.json not found. Using minimal defaults.")
|
||||
# Minimal defaults for development without system_config.json
|
||||
_system_config = {
|
||||
"APPNAME": "Loxide",
|
||||
"LOG_LEVEL": "INFO",
|
||||
"PATH_EXCLUSION_CONST": 4,
|
||||
"MIN_FILES_FOR_PATH": 4,
|
||||
"VT_THREAT_TOLERANCE": 4,
|
||||
"POLICY_MAP_ENF_AUD": {"enforced_id": "audit_id"},
|
||||
"POLICY_MAP_ENF_AUD": {},
|
||||
}
|
||||
|
||||
_protected_config = {key: system_config[key] for key in PROTECTED_KEYS}
|
||||
return _protected_config
|
||||
return _system_config
|
||||
|
||||
|
||||
def get_protected_value(
|
||||
def get_system_value(
|
||||
key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None
|
||||
) -> Optional[T]:
|
||||
value = _protected_config.get(key)
|
||||
"""
|
||||
Get a value from system config (immutable).
|
||||
|
||||
Parameters:
|
||||
key: The config key to retrieve
|
||||
cast_type: Function to cast the value to desired type
|
||||
default: Default value if key not found
|
||||
|
||||
Returns:
|
||||
The config value cast to the desired type, or default
|
||||
"""
|
||||
value = _system_config.get(key)
|
||||
if value is None:
|
||||
logging.warning(f"Protected config key '{key}' not found.")
|
||||
logger.warning(f"System config key '{key}' not found.")
|
||||
return default
|
||||
|
||||
try:
|
||||
if isinstance(value, str):
|
||||
value = value.strip("'\"")
|
||||
return cast_type(value)
|
||||
except (ValueError, TypeError):
|
||||
logging.warning(
|
||||
f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}."
|
||||
logger.warning(
|
||||
f"Invalid value for system key '{key}': {value}. Expected type {cast_type.__name__}."
|
||||
)
|
||||
return default
|
||||
|
||||
|
||||
def get_protected_json(key: str, default: str = "{}") -> dict:
|
||||
raw = _protected_config.get(key, default)
|
||||
def get_system_json(key: str, default: Optional[dict] = None) -> dict:
|
||||
"""
|
||||
Get a JSON/dict value from system config.
|
||||
Handles both dict values and JSON strings.
|
||||
"""
|
||||
if default is None:
|
||||
default = {}
|
||||
|
||||
raw = _system_config.get(key, default)
|
||||
if isinstance(raw, dict):
|
||||
return raw
|
||||
|
||||
try:
|
||||
return json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
try:
|
||||
escaped = raw.encode("unicode_escape").decode("utf-8")
|
||||
return json.loads(escaped)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to parse protected JSON key '{key}': {e}")
|
||||
return json.loads(default)
|
||||
except (json.JSONDecodeError, TypeError) as e:
|
||||
logger.error(f"Failed to parse system JSON key '{key}': {e}")
|
||||
return default
|
||||
|
||||
|
||||
def load_env_json(key: str, default: str):
|
||||
raw = os.getenv(key, default)
|
||||
def get_system_list(key: str, default: Optional[list] = None) -> list:
|
||||
"""
|
||||
Get a list value from system config.
|
||||
Handles both list values and JSON strings.
|
||||
|
||||
Parameters:
|
||||
key: The config key to retrieve
|
||||
default: Default value if key not found or parsing fails
|
||||
|
||||
Returns:
|
||||
The list value or default
|
||||
"""
|
||||
if default is None:
|
||||
default = []
|
||||
|
||||
raw = _system_config.get(key, default)
|
||||
if isinstance(raw, list):
|
||||
return raw
|
||||
|
||||
try:
|
||||
return json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
result = json.loads(raw) if isinstance(raw, str) else raw
|
||||
if isinstance(result, list):
|
||||
return result
|
||||
logger.warning(f"System config key '{key}' is not a list: {type(result)}")
|
||||
return default
|
||||
except (json.JSONDecodeError, TypeError) as e:
|
||||
logger.error(f"Failed to parse system list key '{key}': {e}")
|
||||
return default
|
||||
|
||||
|
||||
def load_user_config(config_dir: Path) -> dict:
|
||||
"""
|
||||
Load user configuration from user_config.json.
|
||||
Creates the file with defaults if it doesn't exist.
|
||||
|
||||
Parameters:
|
||||
config_dir: Directory containing user_config.json
|
||||
|
||||
Returns:
|
||||
The user config dict
|
||||
"""
|
||||
global _user_config
|
||||
|
||||
user_config_path = config_dir / "user_config.json"
|
||||
|
||||
if not user_config_path.exists():
|
||||
# Create default user config
|
||||
default_user_config = {
|
||||
"TELEMETRY": False,
|
||||
"TELEM_URL": "",
|
||||
"TEXTUAL_THEME": "gruvbox",
|
||||
"EXTRAS": "NOTTODAY",
|
||||
}
|
||||
user_config_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(user_config_path, "w") as f:
|
||||
json.dump(default_user_config, f, indent=4)
|
||||
logger.debug(f"Created default user config at {user_config_path}")
|
||||
_user_config = default_user_config
|
||||
else:
|
||||
with open(user_config_path, "r") as f:
|
||||
_user_config = json.load(f)
|
||||
logger.debug(f"✅ Loaded user config from {user_config_path}")
|
||||
|
||||
return _user_config
|
||||
|
||||
|
||||
def save_user_config(config_dir: Path, updates: dict) -> None:
|
||||
"""
|
||||
Save updates to user configuration.
|
||||
Only keys in USER_CONFIG_KEYS are allowed.
|
||||
|
||||
Parameters:
|
||||
config_dir: Directory containing user_config.json
|
||||
updates: Dict of key-value pairs to update
|
||||
"""
|
||||
global _user_config
|
||||
|
||||
# Validate that only user-configurable keys are being updated
|
||||
invalid_keys = [k for k in updates.keys() if k not in USER_CONFIG_KEYS]
|
||||
if invalid_keys:
|
||||
logger.error(f"Attempted to save invalid user config keys: {invalid_keys}")
|
||||
raise ValueError(f"Cannot modify system config keys: {invalid_keys}")
|
||||
|
||||
# Update in-memory config
|
||||
_user_config.update(updates)
|
||||
|
||||
# Write to file
|
||||
user_config_path = config_dir / "user_config.json"
|
||||
user_config_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(user_config_path, "w") as f:
|
||||
json.dump(_user_config, f, indent=4)
|
||||
|
||||
logger.debug(f"✅ Saved user config to {user_config_path}: {updates}")
|
||||
|
||||
|
||||
def get_user_value(
|
||||
key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None
|
||||
) -> Optional[T]:
|
||||
"""
|
||||
Get a value from user config (mutable).
|
||||
|
||||
Parameters:
|
||||
key: The config key to retrieve
|
||||
cast_type: Function to cast the value to desired type
|
||||
default: Default value if key not found
|
||||
|
||||
Returns:
|
||||
The config value cast to the desired type, or default
|
||||
"""
|
||||
value = _user_config.get(key)
|
||||
if value is None:
|
||||
logger.warning(f"User config key '{key}' not found.")
|
||||
return default
|
||||
|
||||
try:
|
||||
escaped = raw.encode("unicode_escape").decode("utf-8")
|
||||
return json.loads(escaped)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to parse {key}: {e}")
|
||||
return json.loads(default)
|
||||
if isinstance(value, str):
|
||||
value = value.strip("'\"")
|
||||
return cast_type(value)
|
||||
except (ValueError, TypeError):
|
||||
logger.warning(
|
||||
f"Invalid value for user key '{key}': {value}. Expected type {cast_type.__name__}."
|
||||
)
|
||||
return default
|
||||
|
||||
|
||||
def load_env(
|
||||
key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None
|
||||
) -> Optional[T]:
|
||||
"""
|
||||
Safely retrieves an environment variable and casts it to the desired type.
|
||||
Safely retrieves an environment variable from .env and casts it to the desired type.
|
||||
This should ONLY be used for runtime/dynamic values like WORKING_DIR.
|
||||
|
||||
For system config, use get_system_value().
|
||||
For user config, use get_user_value().
|
||||
|
||||
Parameters:
|
||||
key (str): The name of the environment variable.
|
||||
cast_type (Callable[[str], T], optional): Function to cast the value. Defaults to str.
|
||||
default (Optional[T], optional): Default value if the variable is not set or invalid.
|
||||
key: The name of the environment variable
|
||||
cast_type: Function to cast the value. Defaults to str
|
||||
default: Default value if the variable is not set or invalid
|
||||
|
||||
Returns:
|
||||
Optional[T]: The casted value or the default.
|
||||
The casted value or the default
|
||||
"""
|
||||
value = os.getenv(key)
|
||||
if value is None:
|
||||
logger.warning(f"Environment variable '{key}' not set.")
|
||||
logger.debug(f"Environment variable '{key}' not set, using default.")
|
||||
return default
|
||||
|
||||
try:
|
||||
value = value.strip("'\"") # Strip surrounding quotes
|
||||
return cast_type(value)
|
||||
@@ -137,3 +300,39 @@ def load_env(
|
||||
f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}."
|
||||
)
|
||||
return default
|
||||
|
||||
|
||||
def load_env_json(key: str, default: str = "[]") -> Any:
|
||||
"""
|
||||
Load a JSON value from environment or system config.
|
||||
|
||||
DEPRECATED: This function is kept for backward compatibility.
|
||||
- For system config lists (BAD_PUBLISHERS, PUPS, BAD_PATH_PARTS), use get_system_list()
|
||||
- For system config dicts, use get_system_json()
|
||||
- For actual .env JSON values, parse manually
|
||||
|
||||
This function automatically redirects known system config keys to system config.
|
||||
"""
|
||||
# Known system config list keys - redirect to system config
|
||||
system_list_keys = ["BAD_PUBLISHERS", "PUPS", "BAD_PATH_PARTS"]
|
||||
if key in system_list_keys:
|
||||
logger.debug(f"Redirecting load_env_json('{key}') to get_system_list()")
|
||||
return get_system_list(key, json.loads(default) if default else [])
|
||||
|
||||
# Known system config dict keys - redirect to system config
|
||||
system_dict_keys = ["POLICY_MAP_ENF_AUD"]
|
||||
if key in system_dict_keys:
|
||||
logger.debug(f"Redirecting load_env_json('{key}') to get_system_json()")
|
||||
return get_system_json(key, json.loads(default) if default else {})
|
||||
|
||||
# Fall back to reading from .env (backward compatibility for unknown keys)
|
||||
raw = os.getenv(key, default)
|
||||
try:
|
||||
return json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
try:
|
||||
escaped = raw.encode("unicode_escape").decode("utf-8")
|
||||
return json.loads(escaped)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse {key}: {e}")
|
||||
return json.loads(default)
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
# 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/>.
|
||||
|
||||
|
||||
import datetime
|
||||
import gc
|
||||
import json
|
||||
import logging
|
||||
|
||||
from bson import ObjectId
|
||||
import pandas as pd
|
||||
|
||||
from services.API import AirlockAPIWrapper
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def getExecutions(api: AirlockAPIWrapper, policy, type, days):
|
||||
import airlock_libs
|
||||
|
||||
executionhist_policy = pd.DataFrame()
|
||||
exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days)
|
||||
if exehist is not None:
|
||||
data = json.loads(exehist)
|
||||
executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
|
||||
if not executionhist_policy.empty:
|
||||
executionhist_policy = executionhist_policy[
|
||||
[
|
||||
"datetime",
|
||||
"sha256",
|
||||
"publisher",
|
||||
"filename",
|
||||
"hostname",
|
||||
"username",
|
||||
"pprocess",
|
||||
"gprocess",
|
||||
"commandline",
|
||||
]
|
||||
]
|
||||
executionhist_policy["policy"] = policy # Add policy column here
|
||||
executionhist_policy = executionhist_policy.drop_duplicates(
|
||||
subset=["sha256", "filename", "hostname"]
|
||||
)
|
||||
executionhist_policy = executionhist_policy.sort_values(
|
||||
by=["sha256", "filename"]
|
||||
)
|
||||
logger.debug(f"Staging of Execution history for policy: {policy} is complete")
|
||||
|
||||
del data
|
||||
del exehist
|
||||
gc.collect()
|
||||
return executionhist_policy
|
||||
|
||||
|
||||
def skipback(days):
|
||||
"""
|
||||
Generate a MongoDB ObjectId for a given number of days ago from today.
|
||||
"""
|
||||
adjusted_days = days
|
||||
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
|
||||
days=adjusted_days
|
||||
)
|
||||
timestamp = int(date_days_ago.timestamp())
|
||||
hex_timestamp = format(timestamp, "08x")
|
||||
objectid_hex = hex_timestamp + "0000000000000000"
|
||||
return ObjectId(objectid_hex)
|
||||
@@ -1,353 +0,0 @@
|
||||
# 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/>.
|
||||
import logging
|
||||
from typing import Any, Callable, List, Optional, Union
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from utils.utils import colorText, get_sanitized_input
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Selector:
|
||||
@staticmethod
|
||||
def _get_sorted_items(
|
||||
items: List[Any], label_func: Callable[[Any], str]
|
||||
) -> List[Any]:
|
||||
return sorted(items, key=lambda item: label_func(item).lower())
|
||||
|
||||
@staticmethod
|
||||
def _display_choices(
|
||||
items: List[Any],
|
||||
label_func: Callable[[Any], str],
|
||||
num_columns: int = 3,
|
||||
header: str = "Available Choices:",
|
||||
) -> None:
|
||||
# Force single column if items are DataFrame rows
|
||||
|
||||
if items and isinstance(items[0], (pd.Series, dict)):
|
||||
num_columns = 1
|
||||
|
||||
rows = (len(items) + num_columns - 1) // num_columns
|
||||
print(f"\n{header}")
|
||||
for row in range(rows):
|
||||
line = ""
|
||||
for col in range(num_columns):
|
||||
idx = row + col * rows
|
||||
if idx < len(items):
|
||||
label = label_func(items[idx])
|
||||
line += f"{idx + 1}: {label:<30}"
|
||||
print(line)
|
||||
|
||||
@staticmethod
|
||||
def _display_selected_items(
|
||||
selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 3
|
||||
) -> None:
|
||||
print(colorText("\nCurrent selections:", "cyan"))
|
||||
if not selected:
|
||||
print(" (none)")
|
||||
return
|
||||
sorted_selected = sorted(selected, key=lambda item: label_func(item).lower())
|
||||
rows = (len(sorted_selected) + num_columns - 1) // num_columns
|
||||
for row in range(rows):
|
||||
line = ""
|
||||
for col in range(num_columns):
|
||||
idx = row + col * rows
|
||||
if idx < len(sorted_selected):
|
||||
label = label_func(sorted_selected[idx])
|
||||
line += f"{label:<30}"
|
||||
print(line)
|
||||
|
||||
@staticmethod
|
||||
def _parse_selection_input(input_str: str, max_index: int) -> List[int]:
|
||||
selections = []
|
||||
for part in input_str.split(","):
|
||||
part = part.strip()
|
||||
if "-" in part:
|
||||
try:
|
||||
start, end = map(int, part.split("-"))
|
||||
selections.extend(range(start, end + 1))
|
||||
except ValueError:
|
||||
continue
|
||||
elif part.isdigit():
|
||||
selections.append(int(part))
|
||||
return [i for i in selections if 1 <= i <= max_index]
|
||||
|
||||
@staticmethod
|
||||
def _select_from_list(
|
||||
items: List[Any],
|
||||
label_func: Callable[[Any], str],
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False,
|
||||
header: str = "Available Choices:",
|
||||
num_columns: int = 3,
|
||||
) -> Union[Optional[Any], List[Any]]:
|
||||
if not items:
|
||||
logger.warning("No items available for selection.")
|
||||
return None
|
||||
|
||||
full_sorted_items = Selector._get_sorted_items(items, label_func)
|
||||
remaining_items = full_sorted_items.copy()
|
||||
selected = []
|
||||
|
||||
if allow_multiple:
|
||||
while True:
|
||||
Selector._display_choices(
|
||||
remaining_items, label_func, num_columns=num_columns, header=header
|
||||
)
|
||||
Selector._display_selected_items(
|
||||
selected, label_func, num_columns=num_columns
|
||||
)
|
||||
choice = (
|
||||
get_sanitized_input(
|
||||
"Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: "
|
||||
)
|
||||
.strip()
|
||||
.lower()
|
||||
)
|
||||
if choice == "q":
|
||||
break
|
||||
elif choice == "r":
|
||||
selected.clear()
|
||||
remaining_items = full_sorted_items.copy()
|
||||
print(colorText("🔄 Selections reset.", "yellow"))
|
||||
continue
|
||||
indices = Selector._parse_selection_input(choice, len(remaining_items))
|
||||
newly_selected = []
|
||||
for index in indices:
|
||||
item = remaining_items[index - 1]
|
||||
if item not in selected:
|
||||
selected.append(item)
|
||||
newly_selected.append(item)
|
||||
if prompt_each:
|
||||
logger.info(f"Selected: {label_func(item)}")
|
||||
else:
|
||||
logger.warning("Item already selected.")
|
||||
remaining_items = [
|
||||
item for item in remaining_items if item not in newly_selected
|
||||
]
|
||||
return selected if selected else None
|
||||
else:
|
||||
Selector._display_choices(
|
||||
full_sorted_items, label_func, num_columns=num_columns, header=header
|
||||
)
|
||||
try:
|
||||
choice = int(get_sanitized_input("Select one item by number: "))
|
||||
if 1 <= choice <= len(full_sorted_items):
|
||||
selected_item = full_sorted_items[choice - 1]
|
||||
logger.info(f"Selected: {label_func(selected_item)}")
|
||||
return selected_item
|
||||
else:
|
||||
logger.warning("Selection out of range.")
|
||||
except ValueError:
|
||||
logger.warning("Invalid input.")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def select_with_mode(
|
||||
items: List[Any],
|
||||
label_func: Callable[[Any], str],
|
||||
header: str = "Available Choices:",
|
||||
) -> List[Any]:
|
||||
print(
|
||||
colorText(
|
||||
"Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
mode = get_sanitized_input("").strip().lower()
|
||||
if mode == "a":
|
||||
return items
|
||||
selected = Selector._select_from_list(
|
||||
items,
|
||||
label_func=label_func,
|
||||
allow_multiple=True,
|
||||
prompt_each=False,
|
||||
header=header,
|
||||
)
|
||||
if not selected:
|
||||
return items
|
||||
if mode == "i":
|
||||
print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
|
||||
return selected
|
||||
elif mode == "e":
|
||||
print(colorText(f"🚫 Excluded {len(selected)} item(s).", "yellow"))
|
||||
return [item for item in items if item not in selected]
|
||||
else:
|
||||
print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
|
||||
return items
|
||||
|
||||
@staticmethod
|
||||
def select_objects(
|
||||
objects: List[Any], allow_multiple: bool = False, prompt_each: bool = False
|
||||
) -> Union[Optional[Any], List[Any]]:
|
||||
return Selector._select_from_list(
|
||||
objects,
|
||||
label_func=lambda obj: getattr(obj, "name", str(obj)),
|
||||
allow_multiple=allow_multiple,
|
||||
prompt_each=prompt_each,
|
||||
header="Available Objects:",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def select_string(
|
||||
options: List[str], allow_multiple: bool = False, prompt_each: bool = False
|
||||
) -> Union[Optional[str], List[str]]:
|
||||
return Selector._select_from_list(
|
||||
options,
|
||||
label_func=str,
|
||||
allow_multiple=allow_multiple,
|
||||
prompt_each=prompt_each,
|
||||
header="Available Options:",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def select_int(
|
||||
options: List[int], allow_multiple: bool = False, prompt_each: bool = False
|
||||
) -> Union[Optional[int], List[int]]:
|
||||
return Selector._select_from_list(
|
||||
options,
|
||||
label_func=lambda x: str(x),
|
||||
allow_multiple=allow_multiple,
|
||||
prompt_each=prompt_each,
|
||||
header="Available Integers:",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def select_value(
|
||||
prompt: str,
|
||||
value_type: type = int,
|
||||
valid_range: Optional[tuple] = None,
|
||||
allow_quit: bool = False,
|
||||
) -> Optional[Any]:
|
||||
while True:
|
||||
user_input = get_sanitized_input(prompt).strip().lower()
|
||||
if allow_quit and user_input == "q":
|
||||
logger.info("User opted to quit value selection.")
|
||||
return None
|
||||
try:
|
||||
value = value_type(user_input)
|
||||
if valid_range:
|
||||
min_val, max_val = valid_range
|
||||
if not (min_val <= value <= max_val):
|
||||
logger.warning(f"Value out of range ({min_val}–{max_val}).")
|
||||
continue
|
||||
logger.info(f"User selected value: {value}")
|
||||
return value
|
||||
except ValueError:
|
||||
logger.warning(f"Invalid input. Expected a {value_type.__name__}.")
|
||||
|
||||
@staticmethod
|
||||
def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool:
|
||||
while True:
|
||||
response = get_sanitized_input(prompt).strip().lower()
|
||||
if response in ["y", "yes"]:
|
||||
logger.info("User confirmed action.")
|
||||
return True
|
||||
elif response in ["n", "no"]:
|
||||
logger.info("User declined action.")
|
||||
return False
|
||||
else:
|
||||
logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")
|
||||
|
||||
@staticmethod
|
||||
def select_dataframe_rows(
|
||||
df: pd.DataFrame,
|
||||
columns: Optional[List[str]] = None,
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False,
|
||||
header: str = "Available Rows:",
|
||||
) -> List[pd.Series]:
|
||||
if df.empty:
|
||||
print("DataFrame is empty.")
|
||||
return []
|
||||
|
||||
if columns:
|
||||
df = df[columns]
|
||||
|
||||
items = [row for _, row in df.iterrows()]
|
||||
label_func = lambda row: str(row.to_dict())
|
||||
|
||||
result = Selector._select_from_list(
|
||||
items,
|
||||
label_func=label_func,
|
||||
allow_multiple=allow_multiple,
|
||||
prompt_each=prompt_each,
|
||||
header=header,
|
||||
)
|
||||
|
||||
if isinstance(result, pd.Series):
|
||||
return [result]
|
||||
elif isinstance(result, list):
|
||||
return result
|
||||
else:
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def select_dataframe_with_mode(
|
||||
df: pd.DataFrame,
|
||||
columns: Optional[List[str]] = None,
|
||||
header: str = "Available Rows:",
|
||||
) -> List[pd.Series]:
|
||||
if df.empty:
|
||||
print("⚠️ DataFrame is empty.")
|
||||
return []
|
||||
|
||||
# Filter columns if specified
|
||||
if columns:
|
||||
df = df[columns]
|
||||
|
||||
items = df.to_dict("records")
|
||||
label_func = lambda row: " | ".join(str(row[col]) for col in df.columns)
|
||||
|
||||
# Show rows first
|
||||
print(colorText(header, "cyan"))
|
||||
for i, row in enumerate(items):
|
||||
print(f"{i}: {label_func(row)}")
|
||||
|
||||
# Prompt for mode once
|
||||
print(
|
||||
colorText(
|
||||
"\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
mode = get_sanitized_input("").strip().lower()
|
||||
|
||||
if mode == "a":
|
||||
return [pd.Series(row) for row in items]
|
||||
|
||||
# Prompt for selection only once
|
||||
selected = Selector._select_from_list(
|
||||
items,
|
||||
label_func=label_func,
|
||||
allow_multiple=True,
|
||||
prompt_each=False,
|
||||
header=header,
|
||||
)
|
||||
|
||||
if not selected:
|
||||
return [pd.Series(row) for row in items]
|
||||
|
||||
if mode == "i":
|
||||
print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
|
||||
return [pd.Series(row) for row in selected]
|
||||
elif mode == "e":
|
||||
print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow"))
|
||||
return [pd.Series(row) for row in items if row not in selected]
|
||||
else:
|
||||
print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
|
||||
return []
|
||||
+120
-91
@@ -13,18 +13,63 @@
|
||||
# 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/>.
|
||||
|
||||
import json
|
||||
import logging
|
||||
import logging.config
|
||||
import logging.handlers
|
||||
import os
|
||||
from pathlib import Path
|
||||
import platform
|
||||
import sys
|
||||
|
||||
from dotenv import load_dotenv, set_key
|
||||
|
||||
from utils.configmanager import PROTECTED_KEYS, load_protected_config
|
||||
from utils.configmanager import (
|
||||
get_system_value,
|
||||
load_system_config,
|
||||
load_user_config,
|
||||
)
|
||||
|
||||
|
||||
class TextualNotificationHandler(logging.Handler):
|
||||
"""
|
||||
Custom logging handler that sends ERROR, WARNING, and CRITICAL logs
|
||||
to Textual toast notifications.
|
||||
"""
|
||||
|
||||
def __init__(self, app):
|
||||
super().__init__()
|
||||
self.app = app
|
||||
|
||||
def emit(self, record):
|
||||
try:
|
||||
# Only handle ERROR, WARNING, and CRITICAL
|
||||
if record.levelno >= logging.WARNING:
|
||||
# Format the message
|
||||
msg = self.format(record)
|
||||
|
||||
# Map log levels to Textual severity
|
||||
severity_map = {
|
||||
logging.WARNING: "warning",
|
||||
logging.ERROR: "error",
|
||||
logging.CRITICAL: "error",
|
||||
}
|
||||
severity = severity_map.get(record.levelno, "information")
|
||||
|
||||
# Send to Textual notification
|
||||
# Use call_from_thread if logging from non-main thread
|
||||
try:
|
||||
self.app.notify(msg, severity=severity, timeout=5)
|
||||
except Exception:
|
||||
# If we're not on the main thread, schedule it
|
||||
try:
|
||||
self.app.call_from_thread(
|
||||
self.app.notify, msg, severity=severity, timeout=5
|
||||
)
|
||||
except Exception:
|
||||
# Silently fail to avoid breaking the logging system
|
||||
pass
|
||||
except Exception:
|
||||
# Silently fail to avoid breaking the logging system
|
||||
pass
|
||||
|
||||
|
||||
def get_base_directory() -> Path:
|
||||
@@ -38,42 +83,46 @@ def get_base_directory() -> Path:
|
||||
return home / ".local" / "share" / "Loxide"
|
||||
|
||||
|
||||
def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
|
||||
def configure_logging(log_dir: Path, cache_dir: Path, log_level: str = "INFO"):
|
||||
"""
|
||||
Configure logging and return function to attach notification handler.
|
||||
|
||||
Args:
|
||||
log_dir: Directory for log files
|
||||
cache_dir: Directory for cache files (used by version checker)
|
||||
log_level: Logging level string
|
||||
|
||||
Returns:
|
||||
Function to attach notification handler to Textual app
|
||||
"""
|
||||
log_file = log_dir / "Loxide.log"
|
||||
|
||||
config = {
|
||||
"version": 1, # Required key for dictConfig format version
|
||||
"disable_existing_loggers": False, # Keeps existing loggers active
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
"formatters": {
|
||||
"detailed": {
|
||||
"format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
# Includes timestamp, logger name, level, and message
|
||||
},
|
||||
"simple": {
|
||||
"format": "%(levelname)s - %(message)s"
|
||||
# Minimal format for console output
|
||||
},
|
||||
"simple": {"format": "%(levelname)s - %(message)s"},
|
||||
"toast": {"format": "%(name)s: %(message)s"}, # Simpler format for toasts
|
||||
},
|
||||
"handlers": {
|
||||
"file": {
|
||||
"class": "logging.handlers.TimedRotatingFileHandler",
|
||||
"filename": str(log_file),
|
||||
"when": "midnight", # Rotate logs at midnight
|
||||
"interval": 1, # Every 1 day
|
||||
"backupCount": 7, # Keep 7 days of logs
|
||||
"encoding": "utf-8", # Ensure UTF-8 encoding
|
||||
"level": "DEBUG", # Always log DEBUG and above
|
||||
"formatter": "detailed", # Use detailed format
|
||||
},
|
||||
"console": {
|
||||
"class": "logging.StreamHandler",
|
||||
"level": log_level.upper(), # Configurable log level
|
||||
"formatter": "simple", # Use simple format
|
||||
"when": "midnight",
|
||||
"interval": 1,
|
||||
"backupCount": 7,
|
||||
"encoding": "utf-8",
|
||||
"level": "DEBUG",
|
||||
"formatter": "detailed",
|
||||
},
|
||||
# REMOVED console handler - it interferes with Textual TUI
|
||||
},
|
||||
"root": {
|
||||
"level": "DEBUG", # Root logger level
|
||||
"handlers": ["file", "console"], # Attach both handlers
|
||||
"level": "DEBUG",
|
||||
"handlers": ["file"], # Only use file handler, not console
|
||||
},
|
||||
}
|
||||
|
||||
@@ -94,56 +143,43 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
|
||||
logging.config.dictConfig(config)
|
||||
logging.getLogger().debug("✅ Logging configured.")
|
||||
|
||||
# Return a function to attach the notification handler once the app is created
|
||||
def attach_notification_handler(app):
|
||||
"""Attach the Textual notification handler and version checker to the app."""
|
||||
# Attach logging handler
|
||||
handler = TextualNotificationHandler(app)
|
||||
handler.setLevel(logging.WARNING) # Only WARNING and above
|
||||
formatter = logging.Formatter("%(name)s: %(message)s")
|
||||
handler.setFormatter(formatter)
|
||||
logging.getLogger().addHandler(handler)
|
||||
logging.getLogger().debug("✅ Textual notification handler attached.")
|
||||
|
||||
def get_system_config_path() -> Path:
|
||||
base_path = Path(
|
||||
getattr(sys, "_MEIPASS", os.path.dirname(os.path.abspath(__file__)))
|
||||
)
|
||||
return base_path.parent / "system_config.json"
|
||||
|
||||
|
||||
def load_system_config() -> dict:
|
||||
# Attach version checker (checks in background, notifies if update available)
|
||||
try:
|
||||
config_path = get_system_config_path()
|
||||
with open(config_path, "r") as f:
|
||||
return json.load(f)
|
||||
except FileNotFoundError:
|
||||
logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
|
||||
return {
|
||||
"APPNAME": "Loxide",
|
||||
"LOG_LEVEL": "DEBUG",
|
||||
"PATH_EXCLUSION_CONST": 4,
|
||||
"MIN_FILES_FOR_PATH": 4,
|
||||
"VT_THREAT_TOLERANCE": 4,
|
||||
"POLICY_MAP_ENF_AUD": {"enforced_id": "audit_id"},
|
||||
}
|
||||
from utils.versionchecker import create_update_notifier
|
||||
|
||||
|
||||
def load_user_config(config_dir: Path) -> dict:
|
||||
user_config_path = config_dir / "user_config.json"
|
||||
if not user_config_path.exists():
|
||||
default_user_config = {"URL": "", "LOG_LEVEL": "INFO"}
|
||||
with open(user_config_path, "w") as f:
|
||||
json.dump(default_user_config, f, indent=4)
|
||||
logging.debug(f"Created user config at {user_config_path}")
|
||||
with open(user_config_path, "r") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def write_config_to_env(config: dict, env_path: Path):
|
||||
for key, value in config.items():
|
||||
if key in PROTECTED_KEYS:
|
||||
continue # Skip protected keys
|
||||
try:
|
||||
serialized = (
|
||||
json.dumps(value) if isinstance(value, (list, dict)) else str(value)
|
||||
)
|
||||
set_key(env_path, key, serialized)
|
||||
create_update_notifier(app, cache_dir=cache_dir)
|
||||
logging.getLogger().debug("✅ Version checker attached.")
|
||||
except ImportError as e:
|
||||
logging.getLogger().debug(f"Version checker not available: {e}")
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to write {key} to .env: {e}")
|
||||
logging.getLogger().warning(f"Could not initialize version checker: {e}")
|
||||
|
||||
return attach_notification_handler
|
||||
|
||||
|
||||
def setup():
|
||||
"""
|
||||
Initialize the application environment:
|
||||
1. Create directory structure
|
||||
2. Load system config (immutable, from system_config.json)
|
||||
3. Load user config (mutable, from user_config.json)
|
||||
4. Configure logging
|
||||
5. Set up .env with WORKING_DIR only
|
||||
|
||||
Returns:
|
||||
attach_notification_handler: Function to attach notification handler to TUI app
|
||||
"""
|
||||
base_dir = get_base_directory()
|
||||
dirs = {
|
||||
"config": base_dir / "config",
|
||||
@@ -155,25 +191,34 @@ def setup():
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
logging.debug(f"{name.capitalize()} directory ensured at: {path}")
|
||||
|
||||
system_config = load_system_config()
|
||||
configure_logging(dirs["logs"], system_config.get("LOG_LEVEL", "DEBUG"))
|
||||
# Load system config (immutable)
|
||||
load_system_config()
|
||||
|
||||
# Configure logging with system-defined log level
|
||||
log_level = get_system_value("LOG_LEVEL", str, "INFO")
|
||||
attach_handler = configure_logging(dirs["logs"], dirs["cache"], log_level)
|
||||
|
||||
# Load user config (mutable)
|
||||
load_user_config(dirs["config"])
|
||||
|
||||
# Set up .env file - ONLY for WORKING_DIR (runtime-configurable value)
|
||||
env_path = base_dir / ".env"
|
||||
if not env_path.exists():
|
||||
env_path.touch()
|
||||
load_dotenv(dotenv_path=env_path, override=True)
|
||||
|
||||
# Set up working directory (only dynamic value in .env)
|
||||
working_dir = Path(os.getenv("WORKING_DIR") or (base_dir / "data"))
|
||||
working_dir.mkdir(parents=True, exist_ok=True)
|
||||
set_key(env_path, "WORKING_DIR", str(working_dir))
|
||||
set_key(str(env_path), "WORKING_DIR", str(working_dir))
|
||||
os.environ["WORKING_DIR"] = str(working_dir)
|
||||
logging.debug(f"Working directory set to: {working_dir}")
|
||||
|
||||
# Create folder structure in working directory
|
||||
folders_structure = {
|
||||
"Approved": [],
|
||||
"Needs_Review": ["Review_First", "Review_Second", "HTML"],
|
||||
"Preflight": ["HTML"],
|
||||
"Archived": [],
|
||||
}
|
||||
|
||||
for folder_name, subfolders in folders_structure.items():
|
||||
@@ -183,25 +228,9 @@ def setup():
|
||||
for subfolder in subfolders:
|
||||
subfolder_path = folder_path / subfolder
|
||||
subfolder_path.mkdir(parents=True, exist_ok=True)
|
||||
logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
|
||||
logging.debug(f"'{subfolder}' subfolder created at: {subfolder_path}")
|
||||
|
||||
user_config = load_user_config(dirs["config"])
|
||||
merged_config = {**system_config, **user_config}
|
||||
logging.info("✅ Setup complete")
|
||||
|
||||
protected_config = load_protected_config()
|
||||
merged_config.update(protected_config)
|
||||
|
||||
# ✅ URL resolution order: system_config → .env → user prompt
|
||||
url = system_config.get("URL")
|
||||
if not url:
|
||||
url = os.getenv("URL")
|
||||
if not url:
|
||||
url = input(
|
||||
"🌐 Enter the service URL (e.g., https://example.com/api): "
|
||||
).strip()
|
||||
merged_config["URL"] = url
|
||||
set_key(env_path, "URL", url)
|
||||
os.environ["URL"] = url
|
||||
logging.debug(f"Service URL set to: {url}")
|
||||
|
||||
write_config_to_env(merged_config, env_path)
|
||||
# Return the attach handler function
|
||||
return attach_handler
|
||||
|
||||
-516
@@ -1,516 +0,0 @@
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import dotenv
|
||||
from dotenv import set_key
|
||||
from textual.app import App, ComposeResult
|
||||
from textual.containers import Vertical
|
||||
from textual.reactive import reactive
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import (
|
||||
Button,
|
||||
DirectoryTree,
|
||||
Footer,
|
||||
Header,
|
||||
Static,
|
||||
Tab,
|
||||
Tabs,
|
||||
)
|
||||
|
||||
from flows.otp import otp_activities_by_agent, otp_revoke
|
||||
from flows.prepPolicy import menu_policy_enforce
|
||||
from flows.quietAgent import findQuietAgents
|
||||
from models.agent import Agent
|
||||
from models.policy import Policy
|
||||
from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
||||
from screens.otpworkflowscreen import OTPWorkflowScreen
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.policyhandler import confirmUpdateAfromE
|
||||
from utils.configmanager import load_env
|
||||
from utils.setup import get_base_directory, load_user_config
|
||||
from utils.utils import open_directory
|
||||
from widgets.agentmoveoperations import AgentMoveOperations
|
||||
from widgets.amber_terminal_theme import get_amber_terminal_theme
|
||||
from widgets.multiagentselector import MultiAgentSelector
|
||||
from widgets.OTP_generate import OTPGenerator
|
||||
from widgets.policytreewidget import PolicyTreeWidget
|
||||
from widgets.resultsdisplay import ResultsDisplay
|
||||
from widgets.retro_terminal_theme import get_retro_terminal_theme
|
||||
from widgets.themeselector import ThemeSelector
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GLOBAL STASH
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PENDING_JOB = None
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
|
||||
# ---------------------------------------------------------------------------
|
||||
def _persist_user_theme(theme_name: str) -> None:
|
||||
"""
|
||||
Store the chosen Textual theme in the user's config:
|
||||
<base>/config/user_config.json
|
||||
and also mirror to <base>/.env so load_env(...) sees it.
|
||||
"""
|
||||
base_dir = get_base_directory()
|
||||
config_dir = base_dir / "config"
|
||||
user_config_path = config_dir / "user_config.json"
|
||||
env_path = base_dir / ".env"
|
||||
|
||||
# ensure dirs / files exist similarly to setup()
|
||||
config_dir.mkdir(parents=True, exist_ok=True)
|
||||
if not user_config_path.exists():
|
||||
# minimal default like your load_user_config does
|
||||
user_config_path.write_text(
|
||||
'{"URL": "", "LOG_LEVEL": "INFO"}\n', encoding="utf-8"
|
||||
)
|
||||
|
||||
# load existing user config
|
||||
user_conf = load_user_config(config_dir)
|
||||
user_conf["TEXTUAL_THEME"] = theme_name
|
||||
|
||||
# write it back
|
||||
user_config_path.write_text(
|
||||
# pretty print so it stays human-readable
|
||||
__import__("json").dumps(user_conf, indent=4),
|
||||
encoding="utf-8",
|
||||
)
|
||||
logger.debug("Updated user_config.json with TEXTUAL_THEME=%s", theme_name)
|
||||
|
||||
# mirror to .env (like setup.write_config_to_env does)
|
||||
env_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not env_path.exists():
|
||||
env_path.touch()
|
||||
try:
|
||||
set_key(str(env_path), "TEXTUAL_THEME", theme_name)
|
||||
except Exception as exc: # keep going even if .env write fails
|
||||
logger.warning("Failed to mirror TEXTUAL_THEME to .env: %s", exc)
|
||||
|
||||
# reload so load_env(...) sees the new value right now
|
||||
dotenv.load_dotenv(dotenv_path=env_path, override=True)
|
||||
logger.debug("Reloaded .env from %s", env_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1) SCREEN
|
||||
# ---------------------------------------------------------------------------
|
||||
class MainMenuScreen(Screen):
|
||||
current_tab = reactive("")
|
||||
|
||||
BUTTON_DEFS = {
|
||||
"agent_actions": [
|
||||
(
|
||||
"🖥️ - Find, Move, or Generate OTP for Agents",
|
||||
"move_agent_workflow_button",
|
||||
),
|
||||
("🔇 - Find Quiet Hosts", "find_quiet_button"),
|
||||
],
|
||||
"policy": [
|
||||
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
|
||||
("🔄 - Update Audit Policies", "policy_audit_update_button"),
|
||||
("📊 - OTP Activities By Agent", "otp_activities_button"),
|
||||
("❌ - Revoke OTPs", "otp_revoke_button"),
|
||||
],
|
||||
}
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper) -> None:
|
||||
super().__init__()
|
||||
self.api = api
|
||||
self.extras = load_env("EXTRAS")
|
||||
wd = load_env("WORKING_DIR") or os.getcwd()
|
||||
if not os.path.isdir(wd):
|
||||
wd = os.getcwd()
|
||||
self.working_dir = wd
|
||||
|
||||
def _make_buttons_for(self, tab_id: str) -> Vertical:
|
||||
defs = self.BUTTON_DEFS.get(tab_id, [])
|
||||
buttons = []
|
||||
for label, btn_id in defs:
|
||||
btn = Button(label, id=btn_id)
|
||||
btn.styles.width = "100%"
|
||||
buttons.append(btn)
|
||||
return Vertical(*buttons)
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Header(show_clock=True, icon="⚙")
|
||||
|
||||
tabs = [
|
||||
Tab("Tree View", id="p_tree"),
|
||||
Tab("Agents", id="agent_actions"),
|
||||
Tab("Directory", id="dir"),
|
||||
Tab("Settings", id="settings"),
|
||||
]
|
||||
|
||||
if self.extras == "POLICYPREP":
|
||||
tabs.insert(2, Tab("Policy Prep", id="policy"))
|
||||
|
||||
yield Tabs(*tabs, id="tabs")
|
||||
yield Vertical(id="content")
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
self.switch_tab("agent_actions")
|
||||
|
||||
# focus helpers
|
||||
def _get_content_buttons(self) -> list[Button]:
|
||||
content = self.query_one("#content", Vertical)
|
||||
return list(content.query(Button))
|
||||
|
||||
def _focus_first_button(self) -> None:
|
||||
buttons = self._get_content_buttons()
|
||||
if buttons:
|
||||
buttons[0].focus()
|
||||
|
||||
def _focus_tabs(self) -> None:
|
||||
tabs = self.query_one("#tabs", Tabs)
|
||||
tabs.focus()
|
||||
|
||||
def _focus_nearby_button(self, direction: int) -> None:
|
||||
buttons = self._get_content_buttons()
|
||||
if not buttons:
|
||||
return
|
||||
|
||||
try:
|
||||
current = next(i for i, b in enumerate(buttons) if b.has_focus)
|
||||
except StopIteration:
|
||||
if direction > 0:
|
||||
buttons[0].focus()
|
||||
else:
|
||||
buttons[-1].focus()
|
||||
return
|
||||
|
||||
if direction < 0 and current == 0:
|
||||
self._focus_tabs()
|
||||
return
|
||||
|
||||
new_index = current + direction
|
||||
if 0 <= new_index < len(buttons):
|
||||
buttons[new_index].focus()
|
||||
|
||||
def switch_tab(self, tab_id: str) -> None:
|
||||
self.current_tab = tab_id
|
||||
content = self.query_one("#content", Vertical)
|
||||
content.remove_children()
|
||||
|
||||
if tab_id in self.BUTTON_DEFS:
|
||||
content.mount(self._make_buttons_for(tab_id))
|
||||
self.call_later(self._focus_first_button)
|
||||
elif tab_id == "dir":
|
||||
content.mount(DirectoryTree(self.working_dir, id="dir_tree"))
|
||||
elif tab_id == "p_tree":
|
||||
content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
|
||||
elif tab_id == "settings":
|
||||
content.mount(ThemeSelector())
|
||||
else:
|
||||
content.mount(Static(f"Unknown tab: {tab_id}"))
|
||||
|
||||
def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
|
||||
self.switch_tab(event.tab.id)
|
||||
|
||||
def on_multi_agent_selector_agents_selected(
|
||||
self, message: MultiAgentSelector.AgentsSelected
|
||||
) -> None:
|
||||
"""Handle selected agents from AgentSelector."""
|
||||
global _PENDING_JOB
|
||||
selected_agents = message.selected_agents
|
||||
logger.info("Selected agents: %s", selected_agents)
|
||||
# TODO: Implement actual handling of selected agents
|
||||
_PENDING_JOB = ("multi_agent_action", selected_agents)
|
||||
self.app.exit()
|
||||
|
||||
def on_theme_selector_theme_selected(
|
||||
self, message: ThemeSelector.ThemeSelected
|
||||
) -> None:
|
||||
"""Handle theme selection from ThemeSelector."""
|
||||
global _PENDING_JOB
|
||||
_persist_user_theme(message.theme_name)
|
||||
_PENDING_JOB = ("restart",)
|
||||
self.app.exit()
|
||||
|
||||
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
|
||||
"""Handle OTP generation request from the workflow."""
|
||||
global _PENDING_JOB
|
||||
|
||||
# Log what we received
|
||||
logger.info(
|
||||
"OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d",
|
||||
len(message.devices),
|
||||
message.requestor,
|
||||
message.reasoning,
|
||||
message.duration,
|
||||
)
|
||||
|
||||
# Set up the job to run the OTP generation
|
||||
_PENDING_JOB = (
|
||||
"otp_workflow",
|
||||
message.devices,
|
||||
message.requestor,
|
||||
message.reasoning,
|
||||
message.duration,
|
||||
)
|
||||
|
||||
self.app.exit()
|
||||
|
||||
def on_agent_move_operations_operation_complete(
|
||||
self, message: AgentMoveOperations.OperationComplete
|
||||
) -> None:
|
||||
"""Handle completion of agent move operation - show results."""
|
||||
logger.info(
|
||||
"Agent move operation completed: %s, %d successful, %d unsuccessful",
|
||||
message.operation,
|
||||
len(message.successful),
|
||||
len(message.unsuccessful),
|
||||
)
|
||||
|
||||
# Format results for display
|
||||
successful_text = "\n".join(
|
||||
[f"{agent.hostname}" for agent, _ in message.successful]
|
||||
)
|
||||
unsuccessful_text = "\n".join(
|
||||
[f"{agent.hostname}: {error}" for agent, error in message.unsuccessful]
|
||||
)
|
||||
|
||||
# Remove the operations widget
|
||||
try:
|
||||
ops_widget = self.query_one(AgentMoveOperations)
|
||||
ops_widget.remove()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Show results
|
||||
self.query_one("#content", Vertical).mount(
|
||||
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
|
||||
)
|
||||
|
||||
def on_results_display_go_back(self, message: ResultsDisplay.GoBack) -> None:
|
||||
"""Handle back button from results display."""
|
||||
try:
|
||||
results_widget = self.query_one(ResultsDisplay)
|
||||
results_widget.remove()
|
||||
except Exception:
|
||||
pass
|
||||
# Return to main menu
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_directory_tree_file_selected(
|
||||
self, event: DirectoryTree.FileSelected
|
||||
) -> None:
|
||||
path = event.path
|
||||
logger.debug("Directory file selected: %s", path)
|
||||
try:
|
||||
open_directory(str(path))
|
||||
except Exception as exc:
|
||||
logger.error("Failed to open %s: %s", path, exc)
|
||||
self.app.bell()
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
global _PENDING_JOB
|
||||
button_id = event.button.id
|
||||
logger.debug("Button pressed: %s", button_id)
|
||||
|
||||
match button_id:
|
||||
case "move_agent_workflow_button":
|
||||
# Push Move Agent workflow screen
|
||||
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
|
||||
event.stop()
|
||||
return # Don't exit the app
|
||||
case "otp_generate_button":
|
||||
# NEW: Push OTP workflow screen instead of legacy function
|
||||
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
|
||||
event.stop()
|
||||
return # Don't exit the app
|
||||
case "find_quiet_button":
|
||||
_PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {})
|
||||
case "otp_activities_button":
|
||||
_PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {})
|
||||
case "otp_revoke_button":
|
||||
_PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {})
|
||||
case "policy_prep_button":
|
||||
_PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {})
|
||||
case "policy_audit_update_button":
|
||||
_PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {})
|
||||
case _:
|
||||
self.app.bell()
|
||||
logger.warning("Unknown button pressed: %s", button_id)
|
||||
return
|
||||
|
||||
logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB)
|
||||
self.app.exit()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2) APP
|
||||
# ---------------------------------------------------------------------------
|
||||
class Loxide(App):
|
||||
CSS = """
|
||||
#logo {
|
||||
width: 100%;
|
||||
content-align: center middle;
|
||||
text-align: center;
|
||||
}
|
||||
"""
|
||||
BINDINGS = [
|
||||
("q", "quit", "Quit"),
|
||||
("d", "open_dir", "Open Directory"),
|
||||
]
|
||||
|
||||
def __init__(self, api: AirlockAPIWrapper):
|
||||
self._textual_theme = load_env("TEXTUAL_THEME") or "nord"
|
||||
super().__init__()
|
||||
self.api = api
|
||||
wd = load_env("WORKING_DIR") or os.getcwd()
|
||||
if not os.path.isdir(wd):
|
||||
wd = os.getcwd()
|
||||
self.working_dir = wd
|
||||
# Initial data load
|
||||
self.refresh_data()
|
||||
|
||||
def refresh_data(self) -> None:
|
||||
"""Public method to refresh policies and devices from the API."""
|
||||
try:
|
||||
self.policies = [
|
||||
Policy(**row.to_dict())
|
||||
for _, row in self.api.policy_find_all().iterrows()
|
||||
]
|
||||
self.devices = [
|
||||
Agent(**row.to_dict())
|
||||
for _, row in self.api.agent_find_all().iterrows()
|
||||
]
|
||||
if self.policies and self.devices:
|
||||
for agent in self.devices:
|
||||
agent.enrich_with_policies(self.policies)
|
||||
logger.debug(
|
||||
f"Enriched {len(self.devices)} agents with policy information"
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to load policies/devices: %s", exc)
|
||||
self.policies = None
|
||||
self.devices = None
|
||||
|
||||
def on_mount(self, api: AirlockAPIWrapper) -> None:
|
||||
self.register_theme(get_retro_terminal_theme())
|
||||
self.register_theme(get_amber_terminal_theme())
|
||||
self.theme = self._textual_theme
|
||||
self.push_screen(MainMenuScreen(api))
|
||||
|
||||
def action_quit(self) -> None:
|
||||
global _PENDING_JOB
|
||||
_PENDING_JOB = None
|
||||
self.exit()
|
||||
|
||||
def action_open_dir(self) -> None:
|
||||
# Refresh data before proceeding
|
||||
self.refresh_data()
|
||||
screen = self.screen_stack[-1]
|
||||
if isinstance(screen, MainMenuScreen):
|
||||
if screen.current_tab != "dir":
|
||||
screen.switch_tab("dir")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3) TERMINAL + LEGACY
|
||||
# ---------------------------------------------------------------------------
|
||||
def _restore_terminal_for_legacy() -> None:
|
||||
sys.stdout.write("\033[?1049l")
|
||||
sys.stdout.write("\033[?25h")
|
||||
sys.stdout.write("\033[0m")
|
||||
sys.stdout.write("\033[?1000l\033[?1002l\033[?1003l\033[?1006l")
|
||||
sys.stdout.write("\033[2J\033[H")
|
||||
sys.stdout.flush()
|
||||
if os.name == "nt":
|
||||
try:
|
||||
import ctypes
|
||||
|
||||
kernel32 = ctypes.windll.kernel32
|
||||
handle = kernel32.GetStdHandle(-11)
|
||||
mode = ctypes.c_ulong()
|
||||
if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
|
||||
kernel32.SetConsoleMode(handle, mode.value | 0x0004)
|
||||
except Exception as exc:
|
||||
logger.debug("VT enable on Windows failed: %s", exc)
|
||||
|
||||
|
||||
def _run_legacy_job(func, args, kwargs) -> None:
|
||||
logger.debug("Running legacy job: %s", getattr(func, "__name__", func))
|
||||
_restore_terminal_for_legacy()
|
||||
try:
|
||||
func(*args, **kwargs)
|
||||
finally:
|
||||
try:
|
||||
input("\nPress Enter to return to the UI...")
|
||||
except EOFError:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4) PUBLIC ENTRYPOINT
|
||||
# ---------------------------------------------------------------------------
|
||||
def run_Loxide(api: AirlockAPIWrapper) -> None:
|
||||
global _PENDING_JOB
|
||||
base_dir = get_base_directory()
|
||||
env_path = base_dir / ".env"
|
||||
dotenv.load_dotenv(dotenv_path=env_path, override=True)
|
||||
|
||||
max_attempts = 5
|
||||
attempts = 0
|
||||
|
||||
while attempts < max_attempts:
|
||||
attempts += 1
|
||||
logger.debug("Starting job loop iteration (attempt %d)", attempts)
|
||||
_PENDING_JOB = None
|
||||
app = Loxide(api)
|
||||
|
||||
try:
|
||||
app.run()
|
||||
except SystemExit as exc:
|
||||
if exc.code != 0:
|
||||
logger.debug("Caught SystemExit from Textual: %s", exc)
|
||||
raise
|
||||
|
||||
job = _PENDING_JOB
|
||||
logger.debug("After app.run(), _PENDING_JOB = %r", job)
|
||||
|
||||
if not job:
|
||||
logger.debug("No job pending, exiting loop")
|
||||
break
|
||||
|
||||
if job[0] == "legacy":
|
||||
_, func, args, kwargs = job
|
||||
_run_legacy_job(func, args, kwargs)
|
||||
continue
|
||||
|
||||
if job[0] == "restart":
|
||||
logger.debug("Restarting job loop")
|
||||
continue
|
||||
|
||||
if job[0] == "multi_agent_action":
|
||||
logger.info("Multi-agent action with selected agents: %s", job[1])
|
||||
continue
|
||||
|
||||
if job[0] == "otp_workflow":
|
||||
_, devices, requestor, reasoning, duration = job
|
||||
|
||||
def otp_generate_with_params():
|
||||
# Your OTP logic here
|
||||
pass
|
||||
|
||||
_run_legacy_job(otp_generate_with_params, (), {})
|
||||
continue
|
||||
|
||||
logger.error("Unknown job type: %r", job)
|
||||
break
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5) DEV
|
||||
# ---------------------------------------------------------------------------
|
||||
if __name__ == "__main__":
|
||||
api = AirlockAPIWrapper()
|
||||
run_Loxide(api)
|
||||
-437
@@ -19,81 +19,10 @@ import os
|
||||
import platform
|
||||
import re
|
||||
import subprocess
|
||||
import tempfile
|
||||
import tkinter as tk
|
||||
from tkinter import filedialog
|
||||
|
||||
import pandas as pd
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def import_to_dataframe(file_path: str) -> pd.DataFrame:
|
||||
df = pd.DataFrame()
|
||||
|
||||
try:
|
||||
if not os.path.exists(file_path):
|
||||
print(colorText(f"Error: File '{file_path}' does not exist.", "red"))
|
||||
return df
|
||||
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
|
||||
if ext == ".csv":
|
||||
df = pd.read_csv(file_path)
|
||||
elif ext == ".parquet":
|
||||
df = pd.read_parquet(file_path)
|
||||
else:
|
||||
print(colorText(f"Error: Unsupported file extension '{ext}'.", "red"))
|
||||
return df
|
||||
|
||||
if df.empty:
|
||||
print(colorText("Error: File has headers but no data rows.", "red"))
|
||||
else:
|
||||
print(colorText(f"Data loaded successfully from {file_path}", "green"))
|
||||
|
||||
return df
|
||||
|
||||
except pd.errors.EmptyDataError:
|
||||
print(
|
||||
colorText(
|
||||
"Notice: CSV file is completely empty, falling back to empty frame",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
return pd.DataFrame()
|
||||
|
||||
except Exception as e:
|
||||
print(colorText(f"Error reading file: {e}", "red"))
|
||||
return pd.DataFrame()
|
||||
|
||||
|
||||
def choose_directory():
|
||||
root = tk.Tk()
|
||||
root.withdraw() # Hide the main window
|
||||
directory = filedialog.askdirectory(title="Select a Directory")
|
||||
print("Selected directory:", directory)
|
||||
return directory
|
||||
|
||||
|
||||
def choose_file(initial_directory=None, required_substring=None):
|
||||
"""Open a file dialog and ensure the selected file contains a required substring."""
|
||||
while True:
|
||||
root = tk.Tk()
|
||||
root.withdraw() # Hide the main window
|
||||
file_path = filedialog.askopenfilename(initialdir=initial_directory)
|
||||
|
||||
if not file_path:
|
||||
print("No file selected.")
|
||||
return None
|
||||
|
||||
if required_substring and required_substring not in file_path:
|
||||
print(
|
||||
f"The selected file must contain '{required_substring}' in its path or name. Please try again."
|
||||
)
|
||||
else:
|
||||
return file_path
|
||||
|
||||
|
||||
def get_sanitized_input(prompt: str) -> str:
|
||||
while True:
|
||||
user_input = input(prompt)
|
||||
@@ -151,235 +80,6 @@ def irtang():
|
||||
)
|
||||
|
||||
|
||||
def section_header(title):
|
||||
print(
|
||||
colorText(
|
||||
"\n --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(f" ------------- {title} -------------", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def areYouSure():
|
||||
print(
|
||||
colorText(
|
||||
"🛑****************************************************************************************************************************************🛑",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"⚠️=========================================================================================================================================⚠️",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"🛑========================================================================================================================================🛑",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"⚠️-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------⚠️",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"🛑========================================================================================================================================🛑",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"⚠️=========================================================================================================================================⚠️",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"🛑****************************************************************************************************************************************🛑",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def locked():
|
||||
print(
|
||||
colorText(
|
||||
r"""
|
||||
████████████████████████████████████████████████████████████████
|
||||
███ ██
|
||||
██ ██████ ███
|
||||
██ ████████████ ███
|
||||
██ ████ ███ ███
|
||||
██ ███ ███ ███
|
||||
██ ███ ███ ███
|
||||
██ ▒████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ██████████████████████ ███
|
||||
██ ███
|
||||
███ ███
|
||||
████████████████████████████████████████████████████████████████████
|
||||
▒██████████████████████████████████████████████████████████████████▒
|
||||
▒████
|
||||
▒████
|
||||
▓██████████████████████████████████████████
|
||||
█████████████████████████████████████████████░
|
||||
""",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def printDeviceEnforceChecklist():
|
||||
print(
|
||||
colorText(
|
||||
"\n --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"\nSequentually follow these steps to prepare a policy for enforcement:",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"\n1. Choose which originating policy or policies to move to enforcement",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText("3. Manually review the files:", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" 'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" If metarules need to be created, please make note of them, and remove the row from the csv.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" When complete, save both csv files to the directory 'approved' and choose this option.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
"4. Manually review the file 'needs_approved\\paths_needing_review.csv'",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing path exclusions you do not approve of",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" When complete, save the csv file to the directory 'approved'", "cyan"
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Do the same process with the list of publishers forthe same directories",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(" Preflight Lists will be generated", "cyan"))
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"5. Choose the destination policy and parent and child allow list", "cyan"
|
||||
)
|
||||
)
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"6. Test ------------------------------------------------------", "cyan"
|
||||
)
|
||||
)
|
||||
print(colorText(" Print rather than apply selected data.", "cyan"))
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"7. Liftoff ------------------------------------------------------", "cyan"
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Apply path exclusions according to allowed and approved paths",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(" Apply signed or attested hashes to Parent Allow List", "cyan")
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Apply approved, but unsigned hashes to the Child Allow List", "cyan"
|
||||
)
|
||||
)
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"R. Remove/Reset Generated data - will prompt to allow keeping execution history",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
print(colorText("B. Back", "cyan"))
|
||||
|
||||
|
||||
def colorText(text, color):
|
||||
colors = {
|
||||
"red": "\033[91m",
|
||||
@@ -394,133 +94,6 @@ def colorText(text, color):
|
||||
return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"
|
||||
|
||||
|
||||
def formatHTML(df, output_html_path=None, overwrite=True):
|
||||
from datetime import datetime
|
||||
|
||||
# Get current date and filename for subtitle
|
||||
today = datetime.now().strftime("%d %B %Y") # Changed to "Day Month Year"
|
||||
filename = output_html_path.replace(".html", "") if output_html_path else "Report"
|
||||
|
||||
dark_css = """
|
||||
<style>
|
||||
body {
|
||||
background-color: #000000;
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
|
||||
color: #f8f8f2;
|
||||
}
|
||||
.header {
|
||||
text-align: center;
|
||||
margin: 20px auto;
|
||||
padding: 10px;
|
||||
border-bottom: 2px solid #ffd700;
|
||||
max-width: 95%;
|
||||
}
|
||||
.header h1 {
|
||||
color: #ffd700;
|
||||
margin: 0;
|
||||
font-size: 32px;
|
||||
}
|
||||
.header p {
|
||||
color: #00bfff;
|
||||
margin: 5px 0 0 0;
|
||||
font-size: 18px;
|
||||
}
|
||||
.table-container {
|
||||
overflow-y: scroll;
|
||||
margin: 0 auto;
|
||||
width: 95%;
|
||||
max-height: calc(80vh - 100px);
|
||||
display: block;
|
||||
border: 1px solid #3a3a4d;
|
||||
margin-bottom: 0;
|
||||
}
|
||||
table {
|
||||
border-collapse: collapse;
|
||||
font-size: 14px;
|
||||
background-color: #1e1e2f;
|
||||
color: #f8f8f2;
|
||||
width: max-content;
|
||||
}
|
||||
th, td {
|
||||
border: 1px solid #3a3a4d;
|
||||
text-align: left;
|
||||
padding: 10px;
|
||||
max-width: 300px;
|
||||
word-wrap: break-word;
|
||||
overflow-wrap: break-word;
|
||||
}
|
||||
/* First column: no wrap */
|
||||
td:nth-child(1), th:nth-child(1) {
|
||||
white-space: nowrap;
|
||||
max-width: none !important;
|
||||
word-wrap: normal !important;
|
||||
}
|
||||
th {
|
||||
background-color: #2e2e40;
|
||||
color: #ffd700;
|
||||
position: sticky;
|
||||
top: 0;
|
||||
z-index: 10;
|
||||
}
|
||||
tr:nth-child(even) {
|
||||
background-color: #262638;
|
||||
}
|
||||
tr:hover {
|
||||
background-color: #33334d;
|
||||
color: #00bfff;
|
||||
}
|
||||
/* Custom scrollbar styling */
|
||||
.table-container::-webkit-scrollbar {
|
||||
width: 12px;
|
||||
}
|
||||
.table-container::-webkit-scrollbar-track {
|
||||
background: #1e1e2f;
|
||||
}
|
||||
.table-container::-webkit-scrollbar-thumb {
|
||||
background-color: #3a3a4d;
|
||||
border-radius: 6px;
|
||||
}
|
||||
</style>
|
||||
"""
|
||||
|
||||
header = f"""
|
||||
<div class="header">
|
||||
<h1>Airlock Tools</h1>
|
||||
<p>{filename} - {today}</p>
|
||||
</div>
|
||||
"""
|
||||
|
||||
html_table = df.to_html(index=False, escape=False)
|
||||
styled_html = (
|
||||
f"<html>\n"
|
||||
f"<head><title>Airlock Tools Report</title></head>\n"
|
||||
f"<body>\n"
|
||||
f"{dark_css}\n"
|
||||
f"{header}\n"
|
||||
f"<div class='table-container'>\n"
|
||||
f" {html_table}\n"
|
||||
f"</div>\n"
|
||||
f"</body>\n"
|
||||
f"</html>"
|
||||
)
|
||||
if output_html_path:
|
||||
with open(output_html_path, "w", encoding="utf-8") as f:
|
||||
f.write(styled_html)
|
||||
print(f"✅ Styled table saved to '{output_html_path}'")
|
||||
elif overwrite:
|
||||
with tempfile.NamedTemporaryFile(
|
||||
suffix=".html", delete=False, mode="w", encoding="utf-8"
|
||||
) as f:
|
||||
f.write(styled_html)
|
||||
temp_path = f.name
|
||||
|
||||
print(f"✅ Styled table saved to temporary file: {temp_path}")
|
||||
else:
|
||||
return styled_html
|
||||
|
||||
|
||||
def open_directory(path):
|
||||
system = platform.system()
|
||||
|
||||
@@ -530,13 +103,3 @@ def open_directory(path):
|
||||
subprocess.run(["xdg-open", path])
|
||||
else:
|
||||
raise OSError(f"Unsupported operating system: {system}")
|
||||
|
||||
|
||||
def print_x_wide(items: list, width: int):
|
||||
for i in range(0, len(items), width):
|
||||
row = items[i : i + width]
|
||||
print(" | ".join(row))
|
||||
|
||||
|
||||
def clear_screen():
|
||||
os.system("cls" if os.name == "nt" else "clear")
|
||||
|
||||
@@ -0,0 +1,560 @@
|
||||
# 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/>.
|
||||
|
||||
"""
|
||||
Version checking and update notification system for Loxide.
|
||||
|
||||
Checks against Gitea releases at:
|
||||
https://git.racooncity.org/brotoskyj/AirlockTools/releases
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import re
|
||||
import threading
|
||||
from typing import Callable, Optional
|
||||
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Current application version - UPDATE THIS ON EACH RELEASE
|
||||
__version__ = "1.0.0"
|
||||
|
||||
# Gitea release API configuration
|
||||
GITEA_API_BASE = "https://git.racooncity.org/api/v1"
|
||||
REPO_OWNER = "brotoskyj"
|
||||
REPO_NAME = "AirlockTools"
|
||||
RELEASES_URL = f"{GITEA_API_BASE}/repos/{REPO_OWNER}/{REPO_NAME}/releases"
|
||||
RELEASES_PAGE_URL = f"https://git.racooncity.org/{REPO_OWNER}/{REPO_NAME}/releases"
|
||||
|
||||
# How often to check for updates (in hours)
|
||||
CHECK_INTERVAL_HOURS = 24
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReleaseInfo:
|
||||
"""Information about a release."""
|
||||
|
||||
tag_name: str
|
||||
version: tuple # Parsed semantic version (major, minor, patch)
|
||||
name: str
|
||||
body: str # Release notes
|
||||
published_at: datetime
|
||||
html_url: str
|
||||
download_url: Optional[str] = None # URL to download the release asset
|
||||
is_prerelease: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class UpdateCheckResult:
|
||||
"""Result of an update check."""
|
||||
|
||||
current_version: str
|
||||
latest_version: Optional[str]
|
||||
update_available: bool
|
||||
release_info: Optional[ReleaseInfo]
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
def parse_version(version_str: str) -> Optional[tuple]:
|
||||
"""
|
||||
Parse a version string into a comparable tuple.
|
||||
Supports formats: v1.2.3, 1.2.3, v1.2, 1.2
|
||||
|
||||
Returns (major, minor, patch) tuple or None if parsing fails.
|
||||
"""
|
||||
if not version_str:
|
||||
return None
|
||||
|
||||
# Strip 'v' prefix if present
|
||||
clean = version_str.lstrip("vV").strip()
|
||||
|
||||
# Match semantic version pattern
|
||||
match = re.match(r"^(\d+)(?:\.(\d+))?(?:\.(\d+))?", clean)
|
||||
if not match:
|
||||
return None
|
||||
|
||||
major = int(match.group(1))
|
||||
minor = int(match.group(2)) if match.group(2) else 0
|
||||
patch = int(match.group(3)) if match.group(3) else 0
|
||||
|
||||
return (major, minor, patch)
|
||||
|
||||
|
||||
def compare_versions(v1: tuple, v2: tuple) -> int:
|
||||
"""
|
||||
Compare two version tuples.
|
||||
|
||||
Returns:
|
||||
-1 if v1 < v2
|
||||
0 if v1 == v2
|
||||
1 if v1 > v2
|
||||
"""
|
||||
for a, b in zip(v1, v2):
|
||||
if a < b:
|
||||
return -1
|
||||
if a > b:
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
def get_current_version() -> str:
|
||||
"""Get the current application version."""
|
||||
return __version__
|
||||
|
||||
|
||||
def _parse_release_response(release_data: dict) -> Optional[ReleaseInfo]:
|
||||
"""Parse a release from Gitea API response."""
|
||||
try:
|
||||
tag_name = release_data.get("tag_name", "")
|
||||
version = parse_version(tag_name)
|
||||
if not version:
|
||||
logger.debug(f"Could not parse version from tag: {tag_name}")
|
||||
return None
|
||||
|
||||
# Parse published date
|
||||
published_str = release_data.get("published_at", "")
|
||||
try:
|
||||
published_at = datetime.fromisoformat(published_str.replace("Z", "+00:00"))
|
||||
except (ValueError, AttributeError):
|
||||
published_at = datetime.now(UTC)
|
||||
|
||||
# Get download URL from assets if available
|
||||
download_url = None
|
||||
assets = release_data.get("assets", [])
|
||||
for asset in assets:
|
||||
# Prefer .exe or .zip files
|
||||
name = asset.get("name", "").lower()
|
||||
if name.endswith((".exe", ".zip", ".msi")):
|
||||
download_url = asset.get("browser_download_url")
|
||||
break
|
||||
|
||||
return ReleaseInfo(
|
||||
tag_name=tag_name,
|
||||
version=version,
|
||||
name=release_data.get("name", tag_name),
|
||||
body=release_data.get("body", ""),
|
||||
published_at=published_at,
|
||||
html_url=release_data.get("html_url", RELEASES_PAGE_URL),
|
||||
download_url=download_url,
|
||||
is_prerelease=release_data.get("prerelease", False),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to parse release data: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def fetch_latest_release(
|
||||
include_prerelease: bool = False, timeout: int = 10
|
||||
) -> Optional[ReleaseInfo]:
|
||||
"""
|
||||
Fetch the latest release from Gitea.
|
||||
|
||||
Args:
|
||||
include_prerelease: Whether to include pre-release versions
|
||||
timeout: Request timeout in seconds
|
||||
|
||||
Returns:
|
||||
ReleaseInfo for the latest release, or None if fetch fails
|
||||
"""
|
||||
try:
|
||||
response = requests.get(
|
||||
RELEASES_URL,
|
||||
params={"limit": 10}, # Get last 10 releases to find latest stable
|
||||
timeout=timeout,
|
||||
headers={"Accept": "application/json"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
releases = response.json()
|
||||
if not releases:
|
||||
logger.debug("No releases found")
|
||||
return None
|
||||
|
||||
# Find the latest release (first non-prerelease if we're excluding them)
|
||||
for release_data in releases:
|
||||
release_info = _parse_release_response(release_data)
|
||||
if release_info is None:
|
||||
continue
|
||||
|
||||
if include_prerelease or not release_info.is_prerelease:
|
||||
return release_info
|
||||
|
||||
# If all are prereleases and we're excluding them, return the first one anyway
|
||||
# but log a warning
|
||||
if releases:
|
||||
logger.debug("All releases are pre-releases")
|
||||
return _parse_release_response(releases[0])
|
||||
|
||||
return None
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
logger.warning("Timeout fetching releases from Gitea")
|
||||
return None
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.warning(f"Failed to fetch releases: {e}")
|
||||
return None
|
||||
except (json.JSONDecodeError, KeyError) as e:
|
||||
logger.warning(f"Failed to parse release response: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def check_for_updates(include_prerelease: bool = False) -> UpdateCheckResult:
|
||||
"""
|
||||
Check if a newer version is available.
|
||||
|
||||
Args:
|
||||
include_prerelease: Whether to consider pre-release versions
|
||||
|
||||
Returns:
|
||||
UpdateCheckResult with the check results
|
||||
"""
|
||||
current = get_current_version()
|
||||
current_parsed = parse_version(current)
|
||||
|
||||
if not current_parsed:
|
||||
return UpdateCheckResult(
|
||||
current_version=current,
|
||||
latest_version=None,
|
||||
update_available=False,
|
||||
release_info=None,
|
||||
error="Could not parse current version",
|
||||
)
|
||||
|
||||
release_info = fetch_latest_release(include_prerelease=include_prerelease)
|
||||
|
||||
if release_info is None:
|
||||
return UpdateCheckResult(
|
||||
current_version=current,
|
||||
latest_version=None,
|
||||
update_available=False,
|
||||
release_info=None,
|
||||
error="Could not fetch release information",
|
||||
)
|
||||
|
||||
is_newer = compare_versions(release_info.version, current_parsed) > 0
|
||||
|
||||
return UpdateCheckResult(
|
||||
current_version=current,
|
||||
latest_version=release_info.tag_name,
|
||||
update_available=is_newer,
|
||||
release_info=release_info,
|
||||
)
|
||||
|
||||
|
||||
class VersionChecker:
|
||||
"""
|
||||
Background version checker that periodically checks for updates
|
||||
and can notify the application when updates are available.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cache_dir: Optional[Path] = None,
|
||||
check_interval_hours: int = CHECK_INTERVAL_HOURS,
|
||||
on_update_available: Optional[Callable[[UpdateCheckResult], None]] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the version checker.
|
||||
|
||||
Args:
|
||||
cache_dir: Directory to store last check timestamp
|
||||
check_interval_hours: Hours between automatic checks
|
||||
on_update_available: Callback when update is available
|
||||
"""
|
||||
self.cache_dir = cache_dir
|
||||
self.check_interval = timedelta(hours=check_interval_hours)
|
||||
self.on_update_available = on_update_available
|
||||
self._last_check: Optional[datetime] = None
|
||||
self._last_result: Optional[UpdateCheckResult] = None
|
||||
self._check_thread: Optional[threading.Thread] = None
|
||||
self._dismissed_version: Optional[str] = None
|
||||
|
||||
# Load cached state
|
||||
self._load_cache()
|
||||
|
||||
@property
|
||||
def cache_file(self) -> Optional[Path]:
|
||||
if self.cache_dir:
|
||||
return self.cache_dir / "version_check_cache.json"
|
||||
return None
|
||||
|
||||
def _load_cache(self) -> None:
|
||||
"""Load cached check state."""
|
||||
if not self.cache_file or not self.cache_file.exists():
|
||||
return
|
||||
|
||||
try:
|
||||
with open(self.cache_file, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
if "last_check" in data:
|
||||
self._last_check = datetime.fromisoformat(data["last_check"])
|
||||
# Don't load dismissed_version - dismiss is session-only
|
||||
|
||||
except (json.JSONDecodeError, ValueError, OSError) as e:
|
||||
logger.debug(f"Could not load version check cache: {e}")
|
||||
|
||||
def _save_cache(self) -> None:
|
||||
"""Save check state to cache."""
|
||||
if not self.cache_file:
|
||||
return
|
||||
|
||||
try:
|
||||
self.cache_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
data = {}
|
||||
if self._last_check:
|
||||
data["last_check"] = self._last_check.isoformat()
|
||||
# Don't save dismissed_version - dismiss is session-only
|
||||
|
||||
with open(self.cache_file, "w") as f:
|
||||
json.dump(data, f)
|
||||
|
||||
except OSError as e:
|
||||
logger.debug(f"Could not save version check cache: {e}")
|
||||
|
||||
def should_check(self) -> bool:
|
||||
"""Determine if enough time has passed to check again."""
|
||||
if self._last_check is None:
|
||||
return True
|
||||
|
||||
elapsed = datetime.now(UTC) - self._last_check
|
||||
return elapsed >= self.check_interval
|
||||
|
||||
def check_now(
|
||||
self, force: bool = False, include_prerelease: bool = False
|
||||
) -> UpdateCheckResult:
|
||||
"""
|
||||
Check for updates immediately.
|
||||
|
||||
Args:
|
||||
force: Check even if recently checked
|
||||
include_prerelease: Include pre-release versions
|
||||
|
||||
Returns:
|
||||
UpdateCheckResult
|
||||
"""
|
||||
if not force and not self.should_check() and self._last_result:
|
||||
return self._last_result
|
||||
|
||||
result = check_for_updates(include_prerelease=include_prerelease)
|
||||
self._last_check = datetime.now(UTC)
|
||||
self._last_result = result
|
||||
self._save_cache()
|
||||
|
||||
# Notify if update available and not dismissed
|
||||
if (
|
||||
result.update_available
|
||||
and self.on_update_available
|
||||
and result.latest_version != self._dismissed_version
|
||||
):
|
||||
self.on_update_available(result)
|
||||
|
||||
return result
|
||||
|
||||
def check_async(
|
||||
self, force: bool = False, include_prerelease: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
Check for updates in background thread.
|
||||
|
||||
Args:
|
||||
force: Check even if recently checked
|
||||
include_prerelease: Include pre-release versions
|
||||
"""
|
||||
if self._check_thread and self._check_thread.is_alive():
|
||||
return # Already checking
|
||||
|
||||
if not force and not self.should_check():
|
||||
return # Too soon to check again
|
||||
|
||||
def _check():
|
||||
try:
|
||||
self.check_now(force=True, include_prerelease=include_prerelease)
|
||||
except Exception as e:
|
||||
logger.debug(f"Background version check failed: {e}")
|
||||
|
||||
self._check_thread = threading.Thread(target=_check, daemon=True)
|
||||
self._check_thread.start()
|
||||
|
||||
def dismiss_update(self, version: str) -> None:
|
||||
"""
|
||||
Dismiss update notification for a specific version.
|
||||
Only lasts for the current session - will nag again on next startup.
|
||||
|
||||
Args:
|
||||
version: Version to dismiss (e.g., "v1.2.3")
|
||||
"""
|
||||
# Session-only dismiss - don't save to cache
|
||||
self._dismissed_version = version
|
||||
|
||||
def clear_dismissed(self) -> None:
|
||||
"""Clear the dismissed version so user will be nagged again."""
|
||||
self._dismissed_version = None
|
||||
|
||||
def get_last_result(self) -> Optional[UpdateCheckResult]:
|
||||
"""Get the result of the last check."""
|
||||
return self._last_result
|
||||
|
||||
|
||||
# Global instance for easy access
|
||||
_global_checker: Optional[VersionChecker] = None
|
||||
|
||||
|
||||
def get_version_checker(
|
||||
cache_dir: Optional[Path] = None,
|
||||
on_update_available: Optional[Callable[[UpdateCheckResult], None]] = None,
|
||||
) -> VersionChecker:
|
||||
"""
|
||||
Get or create the global version checker instance.
|
||||
|
||||
Args:
|
||||
cache_dir: Directory for caching (only used on first call)
|
||||
on_update_available: Callback for updates (only used on first call)
|
||||
|
||||
Returns:
|
||||
The global VersionChecker instance
|
||||
"""
|
||||
global _global_checker
|
||||
|
||||
if _global_checker is None:
|
||||
_global_checker = VersionChecker(
|
||||
cache_dir=cache_dir,
|
||||
on_update_available=on_update_available,
|
||||
)
|
||||
|
||||
return _global_checker
|
||||
|
||||
|
||||
def format_update_message(result: UpdateCheckResult, short: bool = False) -> str:
|
||||
"""
|
||||
Format a human-readable update message.
|
||||
|
||||
Args:
|
||||
result: The update check result
|
||||
short: Whether to use a short format
|
||||
|
||||
Returns:
|
||||
Formatted message string
|
||||
"""
|
||||
if not result.update_available:
|
||||
return f"✅ Loxide is up to date (v{result.current_version})"
|
||||
|
||||
if short:
|
||||
return f"🆕 Update available: {result.latest_version}"
|
||||
|
||||
msg = f"🆕 Loxide {result.latest_version} is available! (current: v{result.current_version})"
|
||||
|
||||
if result.release_info:
|
||||
msg += f"\n📥 Download: {result.release_info.html_url}"
|
||||
|
||||
# Include release notes preview if available
|
||||
if result.release_info.body:
|
||||
notes = result.release_info.body.strip()
|
||||
# Truncate if too long
|
||||
if len(notes) > 200:
|
||||
notes = notes[:200] + "..."
|
||||
msg += f"\n\n📋 Release Notes:\n{notes}"
|
||||
|
||||
return msg
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Textual TUI Integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def create_update_notifier(
|
||||
app, cache_dir: Optional[Path] = None, nag_on_startup: bool = True
|
||||
):
|
||||
"""
|
||||
Create a version checker that notifies via Textual toast notifications.
|
||||
|
||||
This should be called after the Textual app is created.
|
||||
|
||||
Args:
|
||||
app: The Textual App instance
|
||||
cache_dir: Directory for caching check state
|
||||
nag_on_startup: Always show notification on startup if update available
|
||||
|
||||
Returns:
|
||||
The VersionChecker instance
|
||||
"""
|
||||
|
||||
def on_update_available(result: UpdateCheckResult):
|
||||
"""Callback when update is available - show toast notification."""
|
||||
try:
|
||||
msg = f"🆕 Update available: {result.latest_version}\nGo to Settings to download"
|
||||
try:
|
||||
app.notify(
|
||||
msg, title="Loxide Update Available", severity="warning", timeout=15
|
||||
)
|
||||
except RuntimeError:
|
||||
app.call_from_thread(
|
||||
app.notify,
|
||||
msg,
|
||||
title="Loxide Update Available",
|
||||
severity="warning",
|
||||
timeout=15,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not show update notification: {e}")
|
||||
|
||||
checker = get_version_checker(
|
||||
cache_dir=cache_dir,
|
||||
on_update_available=on_update_available,
|
||||
)
|
||||
|
||||
# Store checker on app so Loxide.on_mount can use it
|
||||
if nag_on_startup:
|
||||
app._version_checker = checker
|
||||
app._version_nag_shown = False
|
||||
|
||||
return checker
|
||||
|
||||
|
||||
def check_for_updates_startup(
|
||||
cache_dir: Optional[Path] = None,
|
||||
) -> Optional[UpdateCheckResult]:
|
||||
"""
|
||||
Check for updates during application startup.
|
||||
|
||||
This performs a synchronous check but respects the cache interval,
|
||||
so it will only actually query the network once per CHECK_INTERVAL_HOURS.
|
||||
|
||||
Returns the result if an update is available, None otherwise.
|
||||
|
||||
Example usage:
|
||||
result = check_for_updates_startup(cache_dir)
|
||||
if result and result.update_available:
|
||||
print(format_update_message(result))
|
||||
"""
|
||||
checker = get_version_checker(cache_dir=cache_dir)
|
||||
|
||||
# Only check if enough time has passed (uses cache)
|
||||
if not checker.should_check():
|
||||
result = checker.get_last_result()
|
||||
if result and result.update_available:
|
||||
return result
|
||||
return None
|
||||
|
||||
result = checker.check_now(force=False)
|
||||
if result.update_available:
|
||||
return result
|
||||
return None
|
||||
@@ -1,35 +0,0 @@
|
||||
from textual.color import Color
|
||||
from textual.theme import Theme
|
||||
|
||||
|
||||
def get_amber_terminal_theme():
|
||||
"""Amber CRT theme with compensated brightness for blending."""
|
||||
return Theme(
|
||||
name="amber-terminal",
|
||||
background=Color.parse("#000000"), # pure black
|
||||
primary=Color.parse("#ffb733"), # bright amber
|
||||
secondary=Color.parse("#e69500"), # strong amber
|
||||
success=Color.parse("#ffb733"),
|
||||
warning=Color.parse("#ffff66"),
|
||||
error=Color.parse("#ff3300"),
|
||||
surface=Color.parse("#3a1f00"), # brighter brown for blending
|
||||
)
|
||||
|
||||
|
||||
AMBER_TERMINAL_CSS = """
|
||||
Screen {
|
||||
align: center middle;
|
||||
background: #000000; /* force black */
|
||||
color: #ffb733; /* force amber text */
|
||||
}
|
||||
|
||||
.widget {
|
||||
border: tall #ffb733; /* force amber border */
|
||||
background: #3a1f00; /* compensated surface */
|
||||
width: 80%;
|
||||
}
|
||||
|
||||
* {
|
||||
font-family: "Courier New", monospace;
|
||||
}
|
||||
"""
|
||||
@@ -1,38 +0,0 @@
|
||||
from textual.color import Color
|
||||
|
||||
|
||||
def get_retro_terminal_theme():
|
||||
from textual.theme import Theme
|
||||
|
||||
return Theme(
|
||||
name="retro-terminal",
|
||||
background=Color.parse("#000000"),
|
||||
primary=Color.parse("#00ff00"),
|
||||
secondary=Color.parse("#00aa00"),
|
||||
success=Color.parse("#00ff00"),
|
||||
warning=Color.parse("#ffff00"),
|
||||
error=Color.parse("#ff0000"),
|
||||
surface=Color.parse("#071802"),
|
||||
)
|
||||
|
||||
|
||||
RETRO_TERMINAL_CSS = """
|
||||
/* Retro terminal CRT effect */
|
||||
Screen {
|
||||
align: center middle;
|
||||
background: $background;
|
||||
color: $text;
|
||||
}
|
||||
|
||||
/* Blocky, pixelated widgets */
|
||||
.widget {
|
||||
border: tall $primary;
|
||||
background: $surface;
|
||||
width: 80%;
|
||||
}
|
||||
|
||||
/* Monospaced font */
|
||||
* {
|
||||
font-family: "Courier New", monospace;
|
||||
}
|
||||
"""
|
||||
Reference in New Issue
Block a user