mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

RustImplementation #28

Merged
mysticmomba merged 37 commits from RustImplementation into master 2025-11-18 14:51:42 -05:00
Conversation 0 Commits 37 Files Changed 44
+2015 -748
21 changed files with 2015 additions and 748 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 5eff148900 - Show all commits
AirlockTools_Server.py
-93
View File
@@ -1,93 +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 Add 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 dotenv
import urllib3
import flows.localApproval as la
from Server.scheduler_async import (
recurring_job,
register_function,
reload_jobs,
start_scheduler,
)
from services.API import AirlockAPIWrapper
from services.policyhandler import updateAuditPoliciesFromEnforcementPolices
from services.security import getAPI
from utils.setup import setup
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
def main():
# Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
working_dir = setup()
logger = logging.getLogger(__name__)
dotenv.load_dotenv(dotenv_path=working_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 = AirlockAPIWrapper(
base_url=str(os.getenv("URL")),
api_key=getAPI(username, "AirlockTools"),
)
logger.info("Running non-interactively to start monitoring Airlock Changes")
register_function("monitorLA", la.scheduleAddingLAHashes)
register_function("updateAuditPolicies", updateAuditPoliciesFromEnforcementPolices)
if not os.path.exists("scheduling\\jobs.json"):
recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[api])
recurring_job(
"updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api]
)
else:
reload_jobs()
start_scheduler()
if __name__ == "__main__":
main()
IRT_icon_32-512.ico
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 21 KiB

Server/scheduler_async.py
-215
View File
@@ -1,215 +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 asyncio
import json
import logging
import os
from typing import Any, Callable, Dict, List
logger = logging.getLogger(__name__)
# Registry of functions that can be scheduled
FUNCTION_MAP: Dict[str, Callable] = {}
# Dictionary to manually track scheduled jobs by ID
scheduled_jobs: Dict[str, asyncio.TimerHandle] = {}
# Path to the JSON file for job persistence TODO - pin this to the correct place
JOBS_FILE = os.path.join(os.getcwd(), "jobs.json")
def register_function(name: str, func: Callable):
"""
Register a function so it can be called by name later.
Example:
register_function("say_hello", say_hello)
"""
FUNCTION_MAP[name] = func
def load_jobs() -> List[Dict[str, Any]]:
"""
Load jobs from the JSON file, or return [] if none exist.
"""
if not os.path.exists(JOBS_FILE):
return []
with open(JOBS_FILE, "r") as f:
return json.load(f)
def save_jobs(jobs: List[Dict[str, Any]]):
"""
Save jobs to the JSON file (overwrite).
"""
with open(JOBS_FILE, "w") as f:
json.dump(jobs, f, indent=4)
def cancel_job(job_id: str):
"""
Cancel a scheduled job by ID and remove it from the registry and persistence.
"""
handle = scheduled_jobs.pop(job_id, None)
if handle:
handle.cancel()
logger.info(f"Cancelled job '{job_id}'")
jobs = [j for j in load_jobs() if j.get("id") != job_id]
save_jobs(jobs)
def run_once_job(
job_id: str,
func_name: str,
delay_seconds: float,
args=None,
kwargs=None,
persist=True,
):
"""
Schedule a job to run once after a delay (in seconds).
"""
args = args or []
kwargs = kwargs or {}
def job_wrapper():
func = FUNCTION_MAP.get(func_name)
if func is None:
logger.error(f"Function '{func_name}' is not registered.")
return
func(*args, **kwargs)
cancel_job(job_id)
loop = asyncio.get_event_loop()
handle = loop.call_later(delay_seconds, job_wrapper)
scheduled_jobs[job_id] = handle
if persist:
jobs = [j for j in load_jobs() if j.get("id") != job_id]
jobs.append(
{
"id": job_id,
"type": "once",
"delay": delay_seconds,
"function": func_name,
"args": args,
"kwargs": kwargs,
}
)
save_jobs(jobs)
logger.info(
f"Scheduled one-time job '{job_id}' to run in {delay_seconds} seconds."
)
def recurring_job(
job_id: str, func_name: str, interval: float, args=None, kwargs=None, persist=True
):
"""
Schedule a recurring job.
"""
args = args or []
kwargs = kwargs or {}
def job_wrapper():
func = FUNCTION_MAP.get(func_name)
if func is None:
logger.error(f"Function '{func_name}' is not registered.")
return
func(*args, **kwargs)
# Reschedule the job
handle = asyncio.get_event_loop().call_later(interval, job_wrapper)
scheduled_jobs[job_id] = handle
cancel_job(job_id)
handle = asyncio.get_event_loop().call_later(interval, job_wrapper)
scheduled_jobs[job_id] = handle
if persist:
jobs = [j for j in load_jobs() if j.get("id") != job_id]
jobs.append(
{
"id": job_id,
"type": "recurring",
"interval": interval,
"function": func_name,
"args": args,
"kwargs": kwargs,
}
)
save_jobs(jobs)
logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.")
def reload_jobs():
"""
Reload jobs from JSON and reschedule them.
"""
jobs = load_jobs()
for job in jobs:
if job["type"] == "once":
run_once_job(
job["id"],
job["function"],
job["delay"],
job.get("args"),
job.get("kwargs"),
persist=False,
)
elif job["type"] == "recurring":
recurring_job(
job["id"],
job["function"],
job["interval"],
job.get("args"),
job.get("kwargs"),
persist=False,
)
async def start_scheduler():
"""
Start the asynchronous scheduler loop.
This function is a placeholder to keep the event loop alive.
Jobs are scheduled using asyncio.call_later and do not require polling.
"""
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
logger.critical("Scheduler stopped.")
"""
Start the asynchronous scheduler loop.
This function is a placeholder for compatibility. Since we use asyncio.call_later,
jobs are scheduled directly on the event loop and no polling is required.
Usage:
# In an async app (e.g., Textual)
asyncio.create_task(start_scheduler())
# Or in a standalone script
async def main():
await start_scheduler()
asyncio.run(main())
"""
try:
while True:
await asyncio.sleep(3600) # Sleep indefinitely; jobs run via call_later
except asyncio.CancelledError:
logger.critical("Scheduler stopped.")
airlock_libs/airlock_libs.pyi
+1 -1
View File
@@ -30,7 +30,7 @@ def history_logging(
checkpoint_number: str,
policy_names: str,
) -> List[Dict[str, Any]]:
"""
"""
Query execution history logs from the Airlock API.
Parameters
README.md → docs/README.md
View File
flows/localApproval.py
+193 -262
View File
@@ -1,311 +1,242 @@
# 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 handles the creation of local approval requests.
"""
import datetime
import logging
import os
import re
import time
import dotenv
import numpy as np
import pandas as pd
from typing import List, Optional
from models.agent import Agent
from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
from services.agenthandler import moveAgentToRelatedPolicy, selectAgents
from services.API import AirlockAPIWrapper
from utils.configmanager import get_protected_json, load_env, load_env_json
from utils.setup import get_base_directory
from utils.configmanager import get_protected_json
from utils.utils import colorText, get_sanitized_input
logger = logging.getLogger(__name__)
dotenv.load_dotenv()
class LocalApprovalRequestor:
"""Handles creation of local approval requests in Loxide."""
def getLocalApprovals(api: AirlockAPIWrapper):
base_dir = get_base_directory
result = api.otp_find_awaiting()
local_approval = pd.DataFrame(result["response"]["otpusage"])
if os.path.exists(f"{base_dir}\\cache\\newest_local_approval.parquet"):
previous_run = pd.read_parquet(
f"{base_dir}\\cache\\newest_local_approval.parquet"
)
previous_run.to_parquet(
f"{base_dir}\\cache\\last_local_approval.parquet", index=False
)
os.remove(f"{base_dir}\\cache\\newest_local_approval.parquet")
def __init__(self, api: AirlockAPIWrapper, username: str = None):
"""
Initialize the local approval requestor.
# Only keep rows presumably created by the generate local approval function
local_approval = local_approval[
local_approval["purpose"].str.startswith("🎫 Local Approval 🎫")
]
local_approval["batchid"] = local_approval["purpose"].apply(
lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1)
)
if not local_approval.empty:
local_approval.to_parquet(
f"{base_dir}\\cache\\newest_local_approval.parquet", index=False
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"
)
return local_approval
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)
def scheduleAddingLAHashes(api: AirlockAPIWrapper):
Returns:
True if successful, False otherwise
"""
if batch_id is None:
batch_id = int(time.time())
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
pups = load_env_json("PUPS", "[]")
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type=int)
purpose = (
f"🎫 Local Approval 🎫 - {duration_minutes} mins - "
f"batch:{batch_id} Client:{agent_id} User:{self.username}"
)
try:
register_function("add_hash", returnFromLocalApproval)
register_function("move_device", moveAgentToRelatedPolicy)
except Exception as e:
logger.warning(f"Failed to register functions: {e}")
return
try:
approvals_df = getNewLocalApprovals(api)
if approvals_df.empty:
logger.debug("No new local approvals found. Nothing to schedule.")
return
batches = approvals_df.groupby("batchid")
except Exception as e:
logger.warning(f"Failed to retrieve or group local approvals: {e}")
return
for batchid, batch_df in batches:
try:
duration_minutes = int(batch_df["duration"].iloc[0])
start_time = datetime.datetime.now()
run_time = start_time + datetime.timedelta(minutes=duration_minutes)
early_time = start_time + datetime.timedelta(
minutes=np.floor(duration_minutes * 0.95)
self.api.otp_generate(agent_id, duration_minutes, purpose)
logger.info(
f"Generated local approval for {agent_id}, batch {batch_id}, by {self.username}"
)
early_timestamp = early_time.timestamp()
run_timestamp = run_time.timestamp()
# Schedule add_hash job
try:
run_once_job(
f"add_hash_{batchid}",
"add_hash",
early_timestamp,
[
api,
batch_df,
policy_relationship_map,
bad_publisher_list,
pups,
threat_tolerance_constant,
],
None,
)
logger.debug(f"Scheduled add_hash for batch {batchid} at {early_time}")
except Exception:
logger.debug("Failed to schedule add_hash for batch {batchid}: {e}")
# Schedule move_device jobs
devices = batch_df["agentid"].drop_duplicates().tolist()
agents = []
for device in devices:
rows = api.agent_find_by_hostname(device).iterrows()
agents += [Agent(**row["data"]) for _, row in rows]
for agent in agents:
try:
run_once_job(
f"move_device_{agent.hostame}_{batchid}",
"move_device",
run_timestamp,
[api, agent, policy_relationship_map],
"enforcement",
)
print(
f"Scheduled move_device for device {agent.hostname} in batch {batchid} at {run_time}"
)
except Exception as e:
print(
f"Failed to schedule move_device for device {agent.hostname} in batch {batchid}: {e}"
)
return True
except Exception as e:
logger.warning(f"Failed to process batch {batchid}: {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.
def returnFromLocalApproval(
api,
device_df,
policy_relationship_map,
bad_publisher_list,
pups,
threat_tolerance_constant,
):
"""
# Get unique policy names from device list
policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist())
Args:
agent: Agent object to move
# Create inverse map to go from Audit to Enforcement
inverse_map = {v: k for k, v in policy_relationship_map.items()}
# Fetch all policies
all_policies = [Policy(row['groupid'], row['hidden'], row['name'], row['parent']) for _, row in api.policy_find_all().iterrows()]
# Define policy types
policy_types = [1, 2, 6, 7]
#TODO finish logic for adding hashes
"""
working_dir = load_env("WORKING_DIR")
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
pups = load_env_json("PUPS", "[]")
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE")
print(
f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}"
)
def moveToLocalApproval(api: AirlockAPIWrapper):
possible_durations = [15, 60, 360, 1440, 10080]
duration_selected = None
print(colorText("Please select a duration:", "white"))
for i, option in enumerate(possible_durations, start=1):
print(f"{i}. {option}")
try:
choice = int(get_sanitized_input("Enter the number of your choice:"))
if 1 <= choice <= len(possible_durations):
duration_selected = possible_durations[choice - 1]
print(colorText(f"You selected: {duration_selected}", "yellow"))
logger.debug(f"You selected: {duration_selected}")
else:
print(colorText("❌ Invalid choice.", "red"))
logger.debug("Invalid Input")
return
except ValueError:
print(colorText("❌ Invalid input. Please enter a number.", "red"))
logger.debug("Invalid Input")
return
agents = selectAgents(api)
batch = int(time.time())
if not agents:
print(colorText("❌ No agents found or error retrieving agents.", "red"))
logger.debug("No agents found or error retrieving agents")
return
for agent in agents:
Returns:
True if successful, False otherwise
"""
try:
addLocalApproval(api, batch, duration_selected, agent.agentid)
moveAgentToRelatedPolicy(api, agent, "audit")
moveAgentToRelatedPolicy(self.api, agent, "audit")
logger.info(f"Moved {agent.hostname} to audit policy")
return True
except Exception as e:
print(colorText(f"❌ Error processing agent {agent.hostname}: {e}", "red"))
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.
def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid):
Args:
agents: List of Agent objects
duration_minutes: Duration of approval in minutes
db_path: Optional path to database for history tracking
purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
api.otp_generate(agentid, duration_selected, purpose)
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")
)
def monitorAuditStatus(api: AirlockAPIWrapper):
current_agents = findAllAgents(api)
last_agents = []
if not last_agents:
last_agents = current_agents
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
for agent in agents:
try:
# Create local approval
approval_success = self.create_local_approval(
agent.agentid, duration_minutes, batch_id
)
# Reverse map for audit → enforcement
reverse_policy_map = {v: k for k, v in policy_relationship_map.items()}
known_transitions = set(policy_relationship_map.items()) | set(
reverse_policy_map.items()
)
if not approval_success:
raise Exception("Failed to create local approval")
# Index last_agents by hostname for quick lookup
last_agent_map = {agent.hostname: agent for agent in last_agents}
# Move to audit policy
move_success = self.move_agent_to_audit(agent)
# Result buckets
newly_added = []
same_policy = []
moved_to_audit = []
moved_to_enforcement = []
unusual_move = []
if not move_success:
raise Exception("Failed to move to audit policy")
for current in current_agents:
previous = last_agent_map.get(current.hostname)
print(colorText(f"✓ {agent.hostname}", "green"))
success_count += 1
if not previous:
newly_added.append(current)
continue
except Exception as e:
print(colorText(f"✗ {agent.hostname}: {e}", "red"))
logger.error(f"Error processing agent {agent.hostname}: {e}")
failure_count += 1
if current.groupid == previous.groupid:
same_policy.append(current)
elif (previous.groupid, current.groupid) in known_transitions:
moved_to_audit.append(current)
elif (current.groupid, previous.groupid) in known_transitions:
moved_to_enforcement.append(current)
else:
unusual_move.append(current)
return batch_id, success_count, failure_count
# Return all five DataFrames
return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move
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"),
]
def getNewLocalApprovals(api: AirlockAPIWrapper):
# Display duration options
print(colorText("\n⏱️ Select Local Approval Duration:", "white"))
print(colorText("=" * 50, "white"))
working_dir = load_env("WORKING_DIR")
current_la = getLocalApprovals(api)
for i, (minutes, label) in enumerate(duration_options, start=1):
print(f" {i}. {label} ({minutes} minutes)")
# Load old approval list
old_la_path = f"{working_dir}\\Scheduling\\last_local_approval.parquet"
if os.path.exists(old_la_path):
old_la = pd.read_parquet(old_la_path)
else:
old_la = pd.DataFrame(columns=current_la.columns)
print(colorText("=" * 50, "white"))
# Create composite keys
current_la["key"] = (
current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str)
)
old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str)
# Get user selection
try:
choice = int(get_sanitized_input("\nEnter the number of your choice: "))
# Find new entries
new_entries = current_la[~current_la["key"].isin(old_la["key"])]
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
# Convert 'granted' to datetime and filter by last 10 minutes
new_entries["granted"] = pd.to_datetime(
new_entries["granted"], utc=True, errors="coerce"
)
ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(
minutes=10
)
recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago]
except ValueError:
print(colorText("❌ Invalid input. Please enter a number.", "red"))
logger.warning("Invalid input for duration selection")
return
# Save current approvals for next run
current_la.drop(columns=["key"], inplace=True)
current_la.to_parquet(old_la_path, index=False)
# Select agents
print(colorText("\n🎯 Select Agents for Local Approval:", "white"))
agents = selectAgents(self.api)
return recent_entries
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"))
models/policy.py
+17 -14
View File
@@ -13,33 +13,36 @@
# 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
"""
Policy model representing policy data and relationships.
"""
class Policy:
def __init__(self, groupid, hidden, name, parent):
self.groupid = groupid
self.hidden = hidden
self.name = name
self.parent = parent
# policy.py
def __repr__(self):
# Show all current attributes, including dynamically added ones
from dataclasses import asdict, dataclass, field
import json
from typing import Optional
@dataclass(order=True)
class Policy:
name: str
groupid: int = field(compare=False)
hidden: bool = field(compare=False)
parent: Optional[str] = field(default=None, compare=False)
def __repr__(self) -> str:
attrs = ", ".join(
f"{key}={repr(value)}" for key, value in self.__dict__.items()
)
return f"<Execution({attrs})>"
def to_dict(self):
# Return all attributes as a dictionary
return self.__dict__
def to_dict(self) -> dict:
return asdict(self)
def to_json(self):
# Convert to JSON string, handling non-serializable types gracefully
def to_json(self) -> str:
return json.dumps(self.to_dict(), default=str)
screens/moveagentworkflowscreen.py
+61
View File
@@ -0,0 +1,61 @@
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)
)
screens/otpworkflowscreen.py
+8 -25
View File
@@ -1,41 +1,24 @@
# 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.multiagentselector import MultiAgentSelector
from widgets.OTP_generate import OTPGenerator
class OTPWorkflowScreen(Screen):
"""Screen that handles the OTP generation workflow."""
"""Screen that handles the OTP generation workflow without agent selection."""
def __init__(self, all_agents: List[Agent]):
def __init__(self, selected_agents: List[Agent]):
super().__init__()
self.all_agents = all_agents
self.selected_devices = None
self.selected_agents = selected_agents
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 OTP generator."""
self.selected_devices = message.selected_agents
# Remove the MultiAgentSelector
selector = self.query_one(MultiAgentSelector)
selector.remove()
# Mount the OTPGenerator with the selected Agent objects
# No need to pass API - it will access self.app.api directly
self.mount(OTPGenerator(self.selected_devices))
"""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 - call the actual OTP generation function."""
# This will be handled by the main app, but we can also do it here
# For now, just pass it up to the app level
pass
"""Handle OTP generation request - pass it up to the app level if needed."""
screens/policyselectorscreen.py
+91
View File
@@ -0,0 +1,91 @@
# 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/>.
"""
Policy Selector Screen Module
Provides a Textual Screen wrapper for the PolicySelector widget that manages
the policy selection workflow.
"""
import logging
from textual.app import ComposeResult
from textual.screen import Screen
from widgets.policyselector import PolicySelector
logger = logging.getLogger(__name__)
class PolicySelectorScreen(Screen):
"""
A Textual Screen for policy selection in agent move operations.
This screen wraps the PolicySelector widget and manages the workflow
of selecting a target policy for bulk agent movements.
Attributes:
policies: List of available policies (Policy objects or DataFrame).
agent_move_operations: Reference to the parent AgentMoveOperations widget.
"""
CSS = """
Screen {
layout: vertical;
background: $surface;
}
"""
def __init__(
self,
policies,
agent_move_operations=None,
):
"""
Initialize the PolicySelectorScreen.
Args:
policies: List of available policies to display.
agent_move_operations: Reference to parent AgentMoveOperations widget.
Used to call back when policy selection is confirmed.
"""
super().__init__()
self.policies = policies
self.agent_move_operations = agent_move_operations
def compose(self) -> ComposeResult:
"""Create the PolicySelector widget."""
yield PolicySelector(self.policies)
def on_policy_selector_policy_selected(
self, message: PolicySelector.PolicySelected
) -> None:
"""
Handle policy selection from the PolicySelector widget.
When a policy is selected, this handler:
1. Closes the selector screen
2. Calls the parent AgentMoveOperations to execute the move
Args:
message (PolicySelector.PolicySelected): Contains the selected policy.
"""
# Pop this screen to return to AgentMoveOperations
self.app.pop_screen()
# Call parent widget's method to execute the move
if self.agent_move_operations:
self.agent_move_operations._execute_move_to_policy(message.policy)
utils/selector.py
+3 -3
View File
@@ -33,7 +33,7 @@ class Selector:
def _display_choices(
items: List[Any],
label_func: Callable[[Any], str],
num_columns: int = 4,
num_columns: int = 3,
header: str = "Available Choices:",
) -> None:
# Force single column if items are DataFrame rows
@@ -54,7 +54,7 @@ class Selector:
@staticmethod
def _display_selected_items(
selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 4
selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 3
) -> None:
print(colorText("\nCurrent selections:", "cyan"))
if not selected:
@@ -93,7 +93,7 @@ class Selector:
allow_multiple: bool = False,
prompt_each: bool = False,
header: str = "Available Choices:",
num_columns: int = 4,
num_columns: int = 3,
) -> Union[Optional[Any], List[Any]]:
if not items:
logger.warning("No items available for selection.")
utils/tui.py
+99 -75
View File
@@ -23,16 +23,20 @@ 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.agenthandler import findAgents, moveAgents, toggleEnforcement
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()
@@ -101,23 +105,18 @@ class MainMenuScreen(Screen):
current_tab = reactive("")
BUTTON_DEFS = {
"find": [
("🔍 - Device Search", "find_device_button"),
"agent_actions": [
(
"🖥️ - Find, Move, or Generate OTP for Agents",
"move_agent_workflow_button",
),
("🔇 - Find Quiet Hosts", "find_quiet_button"),
],
"move": [
("✅ - Move to local approval", "move_local_button"),
("🔄 - Move to Audit/Enforcement", "move_audit_button"),
("🔀 - Move - Other", "move_other_button"),
],
"otp": [
("🎫 - Generate OTPs", "otp_generate_button"),
("📊 - OTP Activities By Agent", "otp_activities_button"),
("❌ - Revoke OTPs", "otp_revoke_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"),
],
}
@@ -143,23 +142,21 @@ class MainMenuScreen(Screen):
yield Header(show_clock=True, icon="⚙")
tabs = [
Tab("Policy Tree", id="p_tree"),
Tab("Device Search", id="find"),
Tab("Move Agent", id="move"),
Tab("OTP", id="otp"),
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(3, Tab("Policy Prep", id="policy"))
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("find")
self.switch_tab("agent_actions")
# focus helpers
def _get_content_buttons(self) -> list[Button]:
@@ -220,7 +217,7 @@ class MainMenuScreen(Screen):
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
"""Handle selected agents from MultiAgentSelector."""
"""Handle selected agents from AgentSelector."""
global _PENDING_JOB
selected_agents = message.selected_agents
logger.info("Selected agents: %s", selected_agents)
@@ -261,6 +258,47 @@ class MainMenuScreen(Screen):
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:
@@ -278,26 +316,18 @@ class MainMenuScreen(Screen):
logger.debug("Button pressed: %s", button_id)
match button_id:
case "find_device_button":
_PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {})
case "find_quiet_button":
_PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {})
case "move_local_button":
_PENDING_JOB = (
"legacy",
print,
("Move to local approval (placeholder)",),
{},
)
case "move_audit_button":
_PENDING_JOB = ("legacy", toggleEnforcement, (self.app.api,), {})
case "move_other_button":
_PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {})
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":
@@ -326,7 +356,6 @@ class Loxide(App):
text-align: center;
}
"""
BINDINGS = [
("q", "quit", "Quit"),
("d", "open_dir", "Open Directory"),
@@ -340,21 +369,34 @@ class Loxide(App):
if not os.path.isdir(wd):
wd = os.getcwd()
self.working_dir = wd
# Initial data load
self.refresh_data()
# Add error handling for API calls
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 api.policy_find_all().iterrows()
Policy(**row.to_dict())
for _, row in self.api.policy_find_all().iterrows()
]
self.devices = [
Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()
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))
@@ -364,6 +406,8 @@ class Loxide(App):
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":
@@ -380,7 +424,6 @@ def _restore_terminal_for_legacy() -> None:
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
@@ -397,7 +440,6 @@ def _restore_terminal_for_legacy() -> None:
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:
@@ -412,24 +454,31 @@ def _run_legacy_job(func, args, kwargs) -> None:
# ---------------------------------------------------------------------------
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)
while True:
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:
logger.debug("Caught SystemExit from Textual: %s", 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":
@@ -438,49 +487,24 @@ def run_Loxide(api: AirlockAPIWrapper) -> None:
continue
if job[0] == "restart":
# just loop again; fresh .env was already loaded at the top
logger.debug("Restarting job loop")
continue
if job[0] == "multi_agent_action":
# Handle multi-agent selection
logger.info("Multi-agent action with selected agents: %s", job[1])
continue
# NEW: Handle OTP workflow
if job[0] == "otp_workflow":
_, devices, requestor, reasoning, duration = job
# Call your OTP generation with the parameters
def otp_generate_with_params():
print(f"\n{'='*60}")
print("OTP GENERATION")
print(f"{'='*60}")
print(f"Requestor: {requestor}")
print(f"Reasoning: {reasoning}")
print(f"Duration: {duration} minutes")
print(f"\nGenerating OTPs for {len(devices)} devices:")
print(f"{'='*60}\n")
# Call your actual OTP generation function
# You'll need to adapt otp_generate to accept these parameters
# For now, this is a placeholder showing the structure
for device in devices:
print(f"Device: {device}")
print(f" Requestor: {requestor}")
print(f" Reason: {reasoning}")
print(f" Duration: {duration} minutes")
# TODO: Actually call your API to generate OTP
# result = api.generate_otp(device, requestor, reasoning, duration)
print()
print(f"{'='*60}")
print("OTP Generation Complete!")
print(f"{'='*60}")
# Your OTP logic here
pass
_run_legacy_job(otp_generate_with_params, (), {})
continue
logger.error("Unknown job type: %r", job)
break
utils/utils.py
-38
View File
@@ -151,44 +151,6 @@ def irtang():
)
def displayIntro():
print(
colorText(
r"""
_____ .__ .__ __ ___________ .__
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
\/ \/ \/ \/
""",
"cyan",
)
)
def welcome():
print(
colorText(
"=================================================================================",
"cyan",
)
)
print(
colorText(
"======================== Welcome to the Airlock API Tool ========================",
"cyan",
)
)
print(
colorText(
"=================================================================================",
"cyan",
)
)
def section_header(title):
print(
colorText(
widgets/OTP_generate.py
+16 -2
View File
@@ -6,7 +6,16 @@ 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, Input, RadioButton, RadioSet, Static, TextArea
from textual.widgets import (
Button,
Footer,
Header,
Input,
RadioButton,
RadioSet,
Static,
TextArea,
)
from models.agent import Agent
@@ -70,6 +79,7 @@ class OTPGenerator(Widget):
pass
def compose(self):
yield Header(show_clock=True, icon="⚙")
title_text = Static(
f"🎫 Generate One Time Passes for {len(self.devices)} device(s)",
id="otpgen_title",
@@ -165,6 +175,7 @@ class OTPGenerator(Widget):
copy_button.styles.margin = (1, 0, 0, 0)
copy_button.styles.display = "none"
yield copy_button
yield Footer()
def on_mount(self) -> None:
"""Set initial button state."""
@@ -183,7 +194,10 @@ class OTPGenerator(Widget):
btn_id = event.button.id
if btn_id == "back_button":
self.app.pop_screen()
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "copy_clipboard_button":
widgets/agentmoveoperations.py
+753
View File
@@ -0,0 +1,753 @@
"""
Agent Move Operations Widget Module
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
import logging
import os
from typing import List
import pandas as pd
from textual.containers import Horizontal, Vertical
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 models.agent import Agent
from screens.otpworkflowscreen import OTPWorkflowScreen
from screens.policyselectorscreen import PolicySelectorScreen
from widgets.OTP_generate import OTPGenerator
logger = logging.getLogger(__name__)
class AgentMoveOperations(Widget):
"""
A Textual widget for managing bulk agent operations and policy migrations.
This widget provides a comprehensive UI for performing operations on multiple
selected agents. It displays the list of target agents and provides buttons to
trigger various bulk operations like toggling policy modes or enabling local approval.
The widget manages its own state through reactive properties and provides real-time
feedback on operation progress and results. Operations are executed sequentially
per agent with error handling that tracks both successful and failed operations.
Attributes:
operation_in_progress (reactive[bool]): Tracks whether an operation is currently
executing. Used to disable buttons during execution.
selected_operation (reactive[str]): Tracks which operation type is currently
selected or in progress (e.g., "local_approval", "toggle_enforcement").
Example:
```python
agents = [agent1, agent2, agent3]
widget = AgentMoveOperations(agents)
```
"""
# Reactive property to track if an operation is in progress
operation_in_progress = reactive(False)
# Tracks the currently selected operation type
selected_operation = reactive("")
class OperationComplete(Message):
"""
Message posted when a bulk operation completes.
This message is broadcast to parent widgets/screens to notify them of
operation completion along with detailed results. It contains the list
of agents that were processed and the outcome for each.
Attributes:
operation (str): Name of the operation that completed (e.g., "Local Approval Mode").
agents (List[Agent]): List of all agents that were targeted by the operation.
successful (List[tuple]): List of (Agent, result_data) tuples for successfully
processed agents. Result data varies by operation type.
unsuccessful (List[tuple]): List of (Agent, error_message) tuples for agents
where the operation failed. Error message is a string explaining the failure.
"""
def __init__(
self,
operation: str,
agents: List[Agent],
successful: List[tuple],
unsuccessful: List[tuple],
):
super().__init__()
self.operation = operation
self.agents = agents
self.successful = successful # List of (agent, result) tuples
self.unsuccessful = unsuccessful # List of (agent, error) tuples
def __init__(self, agents: List[Agent]):
"""
Initialize the AgentMoveOperations widget.
Args:
agents (List[Agent]): List of Agent objects to perform operations on.
These agents will be displayed in the widget's agent table.
"""
super().__init__()
self.agents = agents
def watch_operation_in_progress(self, old_value: bool, new_value: bool) -> None:
"""
React to changes in the operation_in_progress reactive property.
This is called automatically by Textual when operation_in_progress changes.
It updates the button states to reflect whether an operation is running.
Args:
old_value (bool): Previous value of operation_in_progress.
new_value (bool): New value of operation_in_progress.
"""
self._update_button_states()
def _update_button_states(self) -> None:
"""
Update the enabled/disabled state of operation buttons based on current status.
This method implements the following logic:
- If an operation is in progress: disable all buttons
- If an operation is selected: disable only that operation's button
- If no operation is selected: enable all buttons
The state transitions prevent users from starting multiple operations
simultaneously and provide visual feedback on which operation is active.
Handles NoMatches exceptions gracefully in case buttons are not yet rendered.
"""
try:
export_csv_btn = self.query_one("#export_csv_btn", Button)
local_approval_btn = self.query_one("#local_approval_btn", Button)
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)
# If operation in progress, disable all
if self.operation_in_progress:
otp_gen_btn = True
export_csv_btn.disabled = True
local_approval_btn.disabled = True
toggle_enforcement_btn.disabled = True
other_policy_btn.disabled = True
else:
# If an operation was selected, disable
if self.selected_operation:
otp_gen_btn.disabled = self.selected_operation == "otp_gen"
export_csv_btn.disabled = self.selected_operation == "export_csv"
local_approval_btn.disabled = (
self.selected_operation == "local_approval"
)
toggle_enforcement_btn.disabled = (
self.selected_operation == "toggle_enforcement"
)
other_policy_btn.disabled = (
self.selected_operation == "other_policy"
)
else:
# Enable all buttons
otp_gen_btn = False
export_csv_btn = False
local_approval_btn.disabled = False
toggle_enforcement_btn.disabled = False
other_policy_btn.disabled = False
except NoMatches:
pass
def _display_results(
self, operation_name: str, successful: list, unsuccessful: list
) -> None:
"""
Display operation results in the results text area.
Formats the results into a human-readable summary including:
- Operation name and separator
- List of successful operations with agent hostnames
- List of failed operations with agent hostnames and error messages
- Summary statistics (total successful/failed count)
The results are displayed in the results_text TextArea widget and the
results container is made visible after being initially hidden.
Args:
operation_name (str): Human-readable name of the operation (e.g., "Local Approval Mode").
successful (list): List of (Agent, result_data) tuples for successful operations.
unsuccessful (list): List of (Agent, error_message) tuples for failed operations.
"""
try:
# Build results text
results_lines = [
f"Operation: {operation_name}",
f"{'=' * 50}",
"",
f"✅ Successful ({len(successful)}):",
]
if successful:
for agent, result in successful:
results_lines.append(f" ✅ {agent.hostname}")
else:
results_lines.append(" (none)")
results_lines.append("")
results_lines.append(f"❌ Failed ({len(unsuccessful)}):")
if unsuccessful:
for agent, error in unsuccessful:
results_lines.append(f" ❌ {agent.hostname}: {error}")
else:
results_lines.append(" (none)")
results_lines.append("")
results_lines.append(f"{'=' * 50}")
results_lines.append(
f"Total: {len(successful)} successful, {len(unsuccessful)} failed"
)
results_text_widget = self.query_one("#results_text", TextArea)
results_text_widget.text = "\n".join(results_lines)
# Show results container
results_container = self.query_one("#results_container", Vertical)
results_container.styles.display = "block"
except Exception as e:
logger.error(f"Error displaying results: {e}")
def compose(self):
"""
Build the UI layout for the AgentMoveOperations widget.
This method is called by Textual to create the widget's UI structure.
It builds a two-column layout with:
- Left side: Agent table showing selected agents and their current policies
- Right side: Operation buttons and results display area
- Bottom: Navigation buttons (Back)
The layout is responsive with:
- Agent table: 2/3 width
- Operations panel: 1/3 width
- Results area: Initially hidden, shown after operation completion
"""
yield Header(show_clock=True, icon="⚙")
title_text = Static(
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
id="move_ops_title",
)
title_text.styles.margin = (0, 0, 1, 0)
yield title_text
with Horizontal() as main_layout:
main_layout.styles.height = "auto"
# Left side - Agent list
with Vertical() as left_side:
left_side.styles.width = "3fr"
left_side.styles.height = "auto"
agents_label = Static("Selected Agents:")
agents_label.styles.margin = (0, 0, 0, 0)
yield agents_label
# Create a DataTable to show agents with their current policies
agent_table = DataTable(id="agent_table")
agent_table.styles.height = "1fr"
agent_table.styles.margin = (1, 0, 1, 0)
yield agent_table
# Right side - Operation buttons
with Vertical() as right_side:
right_side.styles.width = "2fr"
right_side.styles.margin = (0, 1, 0, 1)
right_side.styles.height = "auto"
operations_label = Static("Operations:")
operations_label.styles.margin = (0, 0, 1, 0)
yield operations_label
# Operation buttons
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
local_approval_btn = Button(
"✔️ Local Approval Mode", id="local_approval_btn"
)
local_approval_btn.styles.width = "100%"
local_approval_btn.styles.margin = (0, 0, 1, 0)
yield local_approval_btn
otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn")
otp_gen_btn.styles.width = "100%"
otp_gen_btn.styles.margin = (0, 0, 1, 0)
yield otp_gen_btn
toggle_enforcement_btn = Button(
"🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn"
)
toggle_enforcement_btn.styles.width = "100%"
toggle_enforcement_btn.styles.margin = (0, 0, 1, 0)
yield toggle_enforcement_btn
other_policy_btn = Button(
"🔀 Move to Other Policy", id="other_policy_btn"
)
other_policy_btn.styles.width = "100%"
other_policy_btn.styles.margin = (0, 0, 1, 0)
yield other_policy_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
def on_mount(self) -> None:
"""
Initialize widget after it has been mounted on the screen.
This Textual lifecycle method is called after the widget is added to the DOM.
It performs initialization tasks:
- Populates the agent table with columns for Hostname, Policy, and Status
- Adds rows to the table for each agent in self.agents
- Initializes button states based on current widget state
The agent table displays agent.hostname, agent.groupname (or "Unknown"),
and agent.status_text (or "Unknown") for each agent.
"""
table = self.query_one("#agent_table", DataTable)
table.add_columns("Hostname", "Current Policy", "Status")
for agent in self.agents:
table.add_row(
agent.hostname,
agent.groupname or "Unknown",
agent.status_text or "Unknown",
)
self._update_button_states()
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
"""Handle OTP generation request - call the actual OTP generation function."""
def on_button_pressed(self, event: Button.Pressed):
"""
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
- other_policy_btn: Start move to other policy operation
After handling, event.stop() is called to prevent event propagation.
Args:
event (Button.Pressed): The button press event containing the button reference.
"""
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":
try:
results_text = self.query_one("#results_text", TextArea)
import pyperclip
pyperclip.copy(results_text.text)
self.app.notify(
"📋✅ Results copied to clipboard!",
severity="information",
timeout=2,
)
except ImportError:
self.app.notify(
"❌ 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 == "export_csv_btn":
self._start_export_csv_operation()
event.stop()
elif btn_id == "local_approval_btn":
self._start_local_approval_operation()
event.stop()
elif btn_id == "toggle_enforcement_btn":
self._start_toggle_enforcement_operation()
event.stop()
elif btn_id == "other_policy_btn":
self._start_other_policy_operation()
event.stop()
elif btn_id == "otp_gen_btn":
self._start_OTP_gen_operation()
event.stop()
def _start_local_approval_operation(self) -> None:
"""
Execute the local approval mode operation on all selected agents.
This operation performs the following steps for each agent:
1. Generate a unique batch ID (current Unix timestamp)
2. Create a local approval OTP with default duration of 360 minutes (6 hours)
3. Move the agent to its related audit policy mode
The operation:
- Sets operation state flags (selected_operation, operation_in_progress)
- Updates the status label with progress indicator
- Iterates through all agents, tracking successful and unsuccessful operations
- Displays formatted results via _display_results()
- Posts an OperationComplete message for parent widget handling
Agents that fail are logged and added to the unsuccessful list with error details.
The operation completes and returns to a non-busy state regardless of individual
agent success/failure.
Note: The OTP duration (360 minutes) is currently hardcoded and could be
made configurable in future versions.
"""
self.selected_operation = "local_approval"
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("✔️ Moving agents to local approval...")
# Get API from app
api = self.app.api
successful = []
unsuccessful = []
try:
import time
from services.agenthandler import moveAgentToRelatedPolicy
# Generate batch ID
batch = int(time.time())
duration = 360 # Default 6 hours, could make this configurable
for agent in self.agents:
try:
# Add local approval OTP
addLocalApproval(api, batch, duration, agent.agentid)
# Move to audit mode
result = moveAgentToRelatedPolicy(api, agent, "audit")
successful.append((agent, result))
logger.info(
f"Successfully moved {agent.hostname} to local approval"
)
except Exception as e:
unsuccessful.append((agent, str(e)))
logger.error(
f"Failed to move {agent.hostname} to local approval: {e}"
)
except Exception as e:
logger.error(f"Error during local approval operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
return
self.operation_in_progress = False
status_label.update("✅ Operation complete!")
# Display results in the widget
self._display_results("Local Approval Mode", successful, unsuccessful)
# Also post message for potential parent handling
self.post_message(
self.OperationComplete(
"Local Approval Mode", self.agents, successful, unsuccessful
)
)
def _start_export_csv_operation(self) -> None:
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()
agents = self.agents
policies = self.app.policies
path = self.app.working_dir
try:
# Enrich each agent with policies and status text
for agent in agents:
agent.enrich_with_policies(policies)
# Convert each Agent to a dictionary, including all fields
data = []
for agent in agents:
row = asdict(agent)
# Remove the class-level status_map from the row
row.pop("status_map", None)
data.append(row)
# Create DataFrame
df = pd.DataFrame(data)
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
filename = f"agentsearch_{timestamp}.csv"
file_path = os.path.join(str(path), filename)
df.to_csv(file_path, index=False)
successful.append(file_path)
status_label.update(f"✅ Exported to {file_path}")
except Exception:
status_label.update("❌ Failed")
self.operation_in_progress = False
"""
# Display results in the widget
self._display_results("CSV Export", successful, unsuccessful)
# Also post message for potential parent handling
self.post_message(
self.OperationComplete(
"CSV Export", self.agents, successful, unsuccessful
)
)
"""
def _start_toggle_enforcement_operation(self) -> None:
"""
Toggle agents between enforcement and audit policy modes.
This operation intelligently switches each agent between enforcement and
audit modes based on its current state:
- If agent.groupid is in POLICY_MAP_ENF_AUD: currently enforcing , move to audit
- Otherwise: currently in audit, move to enforcement
The operation:
- Retrieves the enforcement/audit policy relationship map from protected config
- Sets operation state flags and updates status label
- Iterates through agents, determining current mode and toggling to opposite
- Tracks successful toggles with the new mode in the result message
- Logs both successes and failures
- Displays results and posts OperationComplete message
The policy relationship map (POLICY_MAP_ENF_AUD) must be present in protected
configuration and maps enforcement policy IDs to audit policy IDs. If the map
is empty or not found, all agents are assumed to be in audit mode and will
be moved to enforcement.
Returns to a non-busy state after completion regardless of individual results.
"""
self.selected_operation = "toggle_enforcement"
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("⏳ Toggling enforcement mode...")
# Get API from app
api = self.app.api
successful = []
unsuccessful = []
try:
from services.agenthandler import moveAgentToRelatedPolicy
from utils.configmanager import get_protected_json
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
for agent in self.agents:
try:
# Determine current mode and toggle
if agent.groupid in policy_relationship_map:
# Currently in enforcement, move to audit
result = moveAgentToRelatedPolicy(api, agent, "audit")
mode = "audit"
else:
# Currently in audit, move to enforcement
result = moveAgentToRelatedPolicy(api, agent, "enforcement")
mode = "enforcement"
successful.append((agent, f"Moved to {mode}: {result}"))
logger.info(f"Successfully toggled {agent.hostname} to {mode}")
self.app.refresh_data()
except Exception as e:
unsuccessful.append((agent, str(e)))
logger.error(f"Failed to toggle {agent.hostname}: {e}")
except Exception as e:
logger.error(f"Error during toggle enforcement operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
return
self.operation_in_progress = False
status_label.update("✅ Operation complete!")
# Display results in the widget
self._display_results("Toggle Audit/Enforcement", successful, unsuccessful)
# Also post message for potential parent handling
self.post_message(
self.OperationComplete(
"Toggle Audit/Enforcement", self.agents, successful, unsuccessful
)
)
def _start_other_policy_operation(self) -> None:
"""
Move agents to a user-selected policy (currently unimplemented).
This operation is intended to allow bulk movement of selected agents to any
alternative policy via a policy selection dialog. Currently, this feature
is not fully implemented.
Planned Implementation:
1. Push a new policy selector screen (TUI modal/overlay)
2. Allow user to choose target policy from available options
3. Move all selected agents to the chosen policy
4. Display results like other operations
Current Behavior:
- Sets selected_operation to "other_policy"
- Displays "Policy selection not yet implemented" status message
- Clears selected_operation without performing any action
TODO: Complete implementation by:
- Creating a policy selector screen component
- Implementing the policy selection logic
- Integrating with moveAgentToPolicy API call
- Adding proper result tracking and display
"""
self.selected_operation = "other_policy"
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("Loading available policies...")
try:
# Fetch all policies from API
api = self.app.api
# Fetch all available policies
all_policies_df = api.policy_find_all()
if all_policies_df.empty:
status_label.update("No policies available")
self.operation_in_progress = False
self.selected_operation = ""
return
# Create and push the policy selector screen
policy_selector_screen = PolicySelectorScreen(
policies=all_policies_df,
agent_move_operations=self,
)
self.app.push_screen(policy_selector_screen)
except Exception as e:
logger.error(f"Error loading policies: {e}")
status_label.update(f"❌ Error: {str(e)}")
self.operation_in_progress = False
self.selected_operation = ""
self.app.notify(f"Failed to load policies: {str(e)}", severity="error")
def _start_OTP_gen_operation(self) -> None:
status_label = self.query_one("#status_label", Static)
status_label.update("Generating OTP.")
self.app.push_screen(OTPWorkflowScreen(self.agents))
def _execute_move_to_policy(self, target_policy) -> None:
"""
Execute the actual move of agents to the selected policy.
Moves each agent sequentially to the target policy, tracking success/failure.
Updates the status label and displays results upon completion.
Args:
target_policy: The Policy object selected by the user.
"""
status_label = self.query_one("#status_label", Static)
status_label.update(f"Moving agents to {target_policy.name}...")
api = self.app.api
successful = []
unsuccessful = []
try:
for agent in self.agents:
try:
# Move agent to target policy
result = 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}"
)
except Exception as e:
unsuccessful.append((agent, str(e)))
logger.error(
f"Failed to move {agent.hostname} to policy {target_policy.name}: {e}"
)
except Exception as e:
logger.error(f"Error during move to policy operation: {e}")
status_label.update(f"Error: {str(e)}")
self.operation_in_progress = False
return
self.app.refresh_data()
self.operation_in_progress = False
status_label.update("Operation complete!")
# Display results in the widget
self._display_results(
f"Move to {target_policy.name}",
successful,
unsuccessful,
)
# Also post message for potential parent handling
self.post_message(
self.OperationComplete(
f"Move to {target_policy.name}",
self.agents,
successful,
unsuccessful,
)
)
widgets/amber_terminal_theme.py
+35
View File
@@ -0,0 +1,35 @@
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;
}
"""
widgets/multiagentselector.py
+18 -7
View File
@@ -6,7 +6,15 @@ from textual.containers import Horizontal, Vertical
from textual.css.query import NoMatches
from textual.message import Message
from textual.widget import Widget
from textual.widgets import Button, SelectionList, Static, Switch, TextArea
from textual.widgets import (
Button,
Footer,
Header,
SelectionList,
Static,
Switch,
TextArea,
)
from models.agent import Agent
@@ -31,7 +39,8 @@ class MultiAgentSelector(Widget):
self._match_type = value
def compose(self):
title_text = Static("🖧 Multi-Agent Selector", id="selector_title")
yield Header(show_clock=True, icon="⚙")
title_text = Static("🖧 Agent Selector", id="selector_title")
title_text.styles.margin = (0, 0, 0, 1)
yield title_text
@@ -70,20 +79,21 @@ class MultiAgentSelector(Widget):
with Horizontal() as select_buttons:
select_buttons.styles.margin = (0, 0, 0, 0)
select_all_button = Button("✅ Select All", id="select_all")
select_all_button.styles.margin = (1, 1, 0, 1)
yield select_all_button
select_none_button = Button("🚫 Select None", id="select_none")
select_none_button.styles.margin = (1, 0, 0, 1)
select_none_button.styles.margin = (1, 1, 0, 1)
yield select_none_button
select_all_button = Button("✅ Select All", id="select_all")
select_all_button.styles.margin = (1, 0, 0, 1)
yield select_all_button
with Horizontal() as button_row:
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(
@@ -98,6 +108,7 @@ class MultiAgentSelector(Widget):
right_pane.styles.width = "2fr"
yield SelectionList(id="match_results")
yield Static(id="unmatched_label")
yield Footer()
def on_switch_changed(self, event: Switch.Changed):
self.match_type = "fuzzy" if event.value else "exact"
widgets/policyselector.py
+491
View File
@@ -0,0 +1,491 @@
"""
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.
"""
import logging
import re
from typing import Optional
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 models.policy import Policy
logger = logging.getLogger(__name__)
class PolicySelector(Widget):
"""
A Textual widget for selecting a target policy for agent operations.
This widget displays available policies in a table and allows users to select
one policy as the destination for bulk agent movements. It automatically excludes:
- Parent/logical policies (where parent == "global-policy-settings")
- Specified policy IDs (e.g., the current policy)
Features:
- Wildcard filtering (* and ?)
- Interactive table for policy browsing
- Explicit confirm button for selection
- Cancel/back button to dismiss
Attributes:
policies (list[Policy]): List of available Policy objects to display.
excluded_policy_ids (set[str]): Set of policy IDs to exclude from selection.
selected_policy (Optional[Policy]): The currently selected policy (if any).
Automatically Filtered Out:
- Policies with parent == "global-policy-settings" (parent policies for organization)
- Any policies in excluded_policy_ids set
Example:
```python
policies = [policy1, policy2, policy3]
widget = PolicySelector(policies, excluded_policy_ids={current_policy.groupid})
```
"""
class PolicySelected(Message):
"""
Message posted when a policy is selected.
Attributes:
policy (Policy): The selected policy object.
"""
def __init__(self, policy: Policy):
super().__init__()
self.policy = policy
def __init__(self, policies: list):
"""
Initialize the PolicySelector widget.
Args:
policies (list): List of Policy objects or DataFrame rows to display.
Can be a list of Policy objects or a pandas DataFrame of policy data.
"""
super().__init__()
self.policies = policies
self.selected_policy: Optional[Policy] = None
self._filtered_policies = []
self._displayed_policies = [] # Track what's currently shown in the table
self._filter_text = ""
def compose(self):
"""
Build the UI layout for the PolicySelector widget.
The layout includes:
- Title indicating policy selection
- Search/filter text area with wildcard support
- Filter help text showing wildcard options
- Apply Filter button
- Clear Filter button
- Confirm Selection button
- Policy table displaying available policies
- Back buttons for navigation
"""
yield Header(show_clock=True, icon="⚙")
title_text = Static(
"🎯 Select Target Policy",
id="policy_selector_title",
)
title_text.styles.margin = (0, 0, 1, 0)
yield title_text
with Horizontal() as main_layout:
main_layout.styles.height = "auto"
# Left side - Filter and controls
with Vertical() as left_side:
left_side.styles.width = "1fr"
left_side.styles.height = "auto"
left_side.styles.margin = (0, 1, 0, 1)
filter_label = Static("Filter Policies:")
filter_label.styles.margin = (0, 0, 0, 0)
yield filter_label
filter_input = TextArea(
id="policy_filter",
text="",
)
filter_input.styles.height = 3
filter_input.styles.margin = (0, 0, 1, 0)
yield filter_input
filter_help = Static("(Use * and ? for wildcards)", id="filter_help")
filter_help.styles.margin = (0, 0, 1, 0)
yield filter_help
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.styles.width = "100%"
clear_button.styles.margin = (0, 0, 1, 0)
yield clear_button
confirm_button = Button("✅ Confirm Selection", id="confirm_button")
confirm_button.styles.width = "100%"
confirm_button.styles.margin = (1, 0, 1, 0)
yield confirm_button
selected_label = Static("", id="selected_policy_label")
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"
right_side.styles.height = "auto"
table_label = Static("Available Policies:")
table_label.styles.margin = (0, 0, 0, 0)
yield table_label
policy_table = DataTable(id="policy_table", cursor_type="row")
policy_table.styles.height = "1fr"
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.
Populates the table with column (Policy Name) and rows for each
available policy (excluding those in excluded_policy_ids and parent policies).
Sets up event handlers for table row selection.
Filters out:
- Parent policies (where parent == "global-policy-settings")
"""
table = self.query_one("#policy_table", DataTable)
# Configure table for row selection
table.cursor_type = "row"
table.zebra_stripes = True
# Only add Policy Name column
table.add_columns("Policy Name")
# Filter out excluded policies and convert to list if DataFrame
if isinstance(self.policies, pd.DataFrame):
policies_list = self.policies.to_dict("records")
else:
policies_list = self.policies
policies_list = sorted(policies_list)
self._filtered_policies = []
self._displayed_policies = [] # Initialize displayed list
for policy_data in policies_list:
# Handle both Policy objects and dict/DataFrame rows
if isinstance(policy_data, Policy):
policy_id = policy_data.groupid
policy_name = policy_data.name
parent = policy_data.parent
else:
policy_id = policy_data.get("groupid", "Unknown")
policy_name = policy_data.get("name", "Unknown")
parent = policy_data.get("parent", None)
# Skip parent policies (logical policies that shouldn't have devices)
if parent == "global-policy-settings":
logger.debug(f"Skipping parent policy: {policy_name}")
continue
self._filtered_policies.append(policy_data)
self._displayed_policies.append(policy_data) # Add to displayed list
table.add_row(
policy_name,
key=policy_id,
)
def on_button_pressed(self, event: Button.Pressed):
"""
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
Args:
event (Button.Pressed): The button press event.
"""
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":
self._apply_filter()
event.stop()
elif btn_id == "clear_filter_button":
self._clear_filter()
event.stop()
elif btn_id == "confirm_button":
self._confirm_selection()
event.stop()
def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
"""
Handle row selection in the policy table.
Updates the selected_policy and displays the selection in the UI.
Args:
event: DataTable.RowSelected event containing the selected row data.
"""
try:
# Get the row key from the event
row_key = event.row_key
if row_key is None:
return
# Find the policy with matching groupid
for policy_data in self._displayed_policies:
if isinstance(policy_data, Policy):
if policy_data.groupid == row_key.value:
self.selected_policy = policy_data
break
else:
if policy_data.get("groupid") == row_key.value:
self.selected_policy = Policy(
groupid=policy_data.get("groupid"),
hidden=policy_data.get("hidden", False),
name=policy_data.get("name"),
parent=policy_data.get("parent"),
)
break
if self.selected_policy:
# Update selection display
label = self.query_one("#selected_policy_label", Static)
label.update(f"✓ Selected: {self.selected_policy.name}")
# Log for debugging
logger.debug(
f"Selected policy: {self.selected_policy.name} (ID: {self.selected_policy.groupid})"
)
self.app.notify(
f"Selected: {self.selected_policy.name}",
severity="information",
timeout=1,
)
except Exception as e:
logger.error(f"Error handling row selection: {e}")
self.app.notify(f"Selection error: {str(e)}", severity="error")
def on_data_table_row_highlighted(self, event: DataTable.RowHighlighted) -> None:
"""
Handle row highlighting (cursor movement) in the table.
This provides immediate visual feedback when navigating rows.
"""
try:
# Get the row key from the event
row_key = event.row_key
if row_key is None:
return
# Find the highlighted policy
highlighted_name = None
for policy_data in self._displayed_policies:
if isinstance(policy_data, Policy):
if policy_data.groupid == row_key.value:
highlighted_name = policy_data.name
break
else:
if policy_data.get("groupid") == row_key.value:
highlighted_name = policy_data.get("name")
break
if highlighted_name:
label = self.query_one("#selected_policy_label", Static)
label.update(f"→ Highlighting: {highlighted_name}")
except Exception as e:
logger.error(f"Error handling row highlight: {e}")
def _apply_filter(self) -> None:
"""
Apply filter text to policy list with wildcard support.
Supports wildcards:
- * matches any sequence of characters
- ? matches a single character
Examples:
- "policy*" matches "policy_prod", "policy_dev", etc.
- "policy?" matches "policy1", "policy2", etc.
- "*audit*" matches anything containing "audit"
- "*test*" matches "AT Testing", "test_policy", etc.
Filters policies by name or ID (case-insensitive) and refreshes the table display
with only matching policies. Only filters from already-filtered list
(which excludes parent policies and excluded IDs).
"""
try:
filter_input = self.query_one("#policy_filter", TextArea)
filter_text = filter_input.text.strip()
table = self.query_one("#policy_table", DataTable)
table.clear()
# Clear the displayed policies list
self._displayed_policies = []
# Compile wildcard pattern if filter text is provided
pattern = None
if filter_text:
# Escape special regex chars but preserve wildcards
pattern_text = re.escape(filter_text.lower())
pattern_text = pattern_text.replace(r"\*", ".*").replace(r"\?", ".")
# Use search() for partial matching
pattern = re.compile(pattern_text, re.IGNORECASE)
# Filter policies based on search text
for policy_data in self._filtered_policies:
# Handle both Policy objects and dict/DataFrame rows
if isinstance(policy_data, Policy):
policy_name = policy_data.name.lower()
policy_id = policy_data.groupid.lower()
display_name = policy_data.name
key_id = policy_data.groupid
else:
policy_name = str(policy_data.get("name", "")).lower()
policy_id = str(policy_data.get("groupid", "Unknown")).lower()
display_name = policy_data.get("name")
key_id = policy_data.get("groupid")
# Match against filter text with wildcard support
if pattern:
# Use search() for partial matching
matches = pattern.search(policy_name) or pattern.search(policy_id)
else:
matches = True
if matches:
# Add to displayed policies list
self._displayed_policies.append(policy_data)
# Add row to table
table.add_row(
display_name,
key=key_id,
)
displayed_count = len(self._displayed_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
self.selected_policy = None
label = self.query_one("#selected_policy_label", Static)
label.update("")
except Exception as e:
logger.error(f"Error applying filter: {e}")
self.app.notify(f"❌ Filter error: {str(e)}", severity="error")
def _clear_filter(self) -> None:
"""
Clear the filter and display all available policies.
Resets the filter text and refreshes the table to show all policies
(already excluding parent policies and excluded IDs).
"""
try:
filter_input = self.query_one("#policy_filter", TextArea)
filter_input.text = ""
table = self.query_one("#policy_table", DataTable)
table.clear()
# Reset displayed policies to all filtered policies
self._displayed_policies = list(self._filtered_policies)
# Reload all policies
for policy_data in self._filtered_policies:
if isinstance(policy_data, Policy):
policy_id = policy_data.groupid
policy_name = policy_data.name
else:
policy_id = policy_data.get("groupid", "Unknown")
policy_name = policy_data.get("name", "Unknown")
# Add row with only policy name
table.add_row(
policy_name,
key=policy_id,
)
self.selected_policy = None
label = self.query_one("#selected_policy_label", Static)
label.update("")
except Exception as e:
logger.error(f"Error clearing filter: {e}")
def on_text_area_changed(self, event) -> None:
"""
Handle TextArea change events - specifically for Enter key in filter.
When the user types in the filter TextArea and the text ends with a newline,
treat it as pressing Enter and apply the filter.
"""
if event.text_area.id == "policy_filter":
# Check if the text ends with a newline (Enter was pressed)
if event.text_area.text.endswith("\n"):
# Remove the newline that was added
event.text_area.text = event.text_area.text.rstrip("\n")
# Apply the filter
self._apply_filter()
def _confirm_selection(self) -> None:
"""
Confirm the selected policy and post selection message.
Posts a PolicySelected message to the parent widget/screen with the
selected policy. If no policy is selected, displays an error notification.
"""
if self.selected_policy is None:
self.app.notify(
"⚠️ Please select a policy first by clicking on a row in the table",
severity="warning",
timeout=3,
)
return
# 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
)
self.post_message(self.PolicySelected(self.selected_policy))
widgets/resultsdisplay.py
+178
View File
@@ -0,0 +1,178 @@
import logging
from textual.containers import Horizontal, Vertical
from textual.message import Message
from textual.widget import Widget
from textual.widgets import Button, Footer, Header, Static
logger = logging.getLogger(__name__)
class ResultsDisplay(Widget):
"""Widget for displaying operation results in a two-column layout."""
CSS = """
ResultsDisplay {
height: 100%;
}
#results_screen {
height: 100%;
}
#results_title {
text-align: center;
margin: 1 0;
text-style: bold;
}
#results_layout {
height: 1fr;
margin: 1 0;
}
#left_column, #right_column {
width: 1fr;
height: 100%;
border: solid green;
padding: 1;
}
#right_column {
border: solid red;
}
#success_label, #failure_label {
text-style: bold;
margin-bottom: 1;
}
#success_results, #failure_results {
height: 1fr;
overflow-y: auto;
background: $surface;
border: round $primary;
padding: 1;
}
.copy_button {
margin-top: 1;
width: 100%;
}
#button_row {
height: auto;
margin: 1 0 0 0;
}
#back_button {
width: 1fr;
}
"""
class CopySuccess(Message):
"""Posted when success results are copied."""
pass
class CopyFailure(Message):
"""Posted when failure results are copied."""
pass
class GoBack(Message):
"""Posted when back button is pressed."""
pass
def __init__(
self, operation: str, successful_results: str, unsuccessful_results: str
) -> None:
super().__init__()
self.operation = operation
self.successful_results = successful_results
self.unsuccessful_results = unsuccessful_results
def compose(self):
with Vertical(id="results_screen"):
yield Header(show_clock=True, icon="⚙")
# Title
title = Static(f"📊 {self.operation} - Results", id="results_title")
yield title
# Two-column layout
with Horizontal(id="results_layout"):
# Left Column - Success
with Vertical(id="left_column"):
yield Static("✅ Successful", id="success_label")
yield Static(self.successful_results, id="success_results")
yield Button(
"📋✅ Copy Success List",
id="copy_success",
classes="copy_button",
)
# Right Column - Failure
with Vertical(id="right_column"):
yield Static("❌ Failed", id="failure_label")
yield Static(self.unsuccessful_results, id="failure_results")
yield Button(
"📋❌ 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:
btn_id = event.button.id
if btn_id == "copy_success":
success_widget = self.query_one("#success_results", Static)
try:
import pyperclip
pyperclip.copy(str(success_widget.renderable))
self.app.notify(
"✅ 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",
severity="warning",
)
except Exception as e:
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
event.stop()
elif btn_id == "copy_failure":
failure_widget = self.query_one("#failure_results", Static)
try:
import pyperclip
pyperclip.copy(str(failure_widget.renderable))
self.app.notify(
"✅ 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",
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()
event.stop()
widgets/retro_terminal_theme.py
+38
View File
@@ -0,0 +1,38 @@
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;
}
"""
widgets/themeselector.py
+13 -13
View File
@@ -15,26 +15,26 @@ class ThemeSelector(Widget):
self.theme_name = theme_name
AVAILABLE_THEMES = [
("textual-dark", "textual-dark"),
("textual-light", "textual-light"),
("nord", "nord"),
("gruvbox", "gruvbox"),
("catppuccin-mocha", "catppuccin-mocha"),
("dracula", "dracula"),
("tokyo-night", "tokyo-night"),
("monokai", "monokai"),
("flexoki", "flexoki"),
("catppuccin-latte", "catppuccin-latte"),
("solarized-light", "solarized-light"),
("Textual Dark", "textual-dark"),
("Textual Light", "textual-light"),
("Nord", "nord"),
("Gruvbox", "gruvbox"),
("Catppuccin Mocha", "catppuccin-mocha"),
("Dracula", "dracula"),
("Tokyo Night", "tokyo-night"),
("Monokai", "monokai"),
("Flexoki", "flexoki"),
("Catppuccin Latte", "catppuccin-latte"),
("Solarized Light", "solarized-light"),
("Retro Terminal", "retro-terminal"),
("Amber Terminal", "amber-terminal"), # your custom theme
]
def compose(self):
yield Static("Theme Options", id="theme_title")
with Vertical() as column:
column.styles.width = "1fr"
column.styles.height = "auto"
for label, btn_id in self.AVAILABLE_THEMES:
yield Button(label, id=f"set_theme_{btn_id}", compact=True)