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
+6285 -1616
44 changed files with 6285 additions and 1616 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitea/workflows/loxide_lib.yml
+1 -1
View File
@@ -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
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

AirlockTools_Client.py → Loxide.py
+4 -15
View File
@@ -23,39 +23,28 @@
import logging
import os
import tempfile
import dotenv
import urllib3
from services.API import AirlockAPIWrapper
from services.security import getAPI
from TUI.TUI import run_Loxide
from utils.configmanager import get_system_value
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")
url = get_system_value("URL")
username = os.getenv("USERNAME")
if not url:
@@ -75,7 +64,7 @@ def main():
raise ValueError("API key for Loxide is missing.")
api = AirlockAPIWrapper(
base_url=str(os.getenv("URL")),
base_url=str(url),
api_key=api_key,
)
run_Loxide(api)
Loxide_Icon.ico
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 112 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.")
widgets/OTP_generate.py → TUI/OTP_generate.py
+26 -9
View File
@@ -1,12 +1,21 @@
import logging
from typing import List
from typing import List, Optional
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, Input, RadioButton, RadioSet, Static, TextArea
from textual.widgets import (
Button,
Footer,
Header,
Input,
RadioButton,
RadioSet,
Static,
TextArea,
)
from models.agent import Agent
@@ -22,7 +31,11 @@ class OTPGenerator(Widget):
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
@@ -70,6 +83,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 +179,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 +198,10 @@ class OTPGenerator(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_clipboard_button":
@@ -229,12 +247,11 @@ 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 = [
"=" * 60,
"OTP GENERATION RESULTS",
"=" * 60,
"Requested OTP Codes:",
"=" * 25,
]
otp_dict = {}
@@ -250,9 +267,9 @@ class OTPGenerator(Widget):
)
for hostname, otp_code in otp_dict.items():
output_lines.append(f"{hostname:30} | {otp_code}")
output_lines.append(f"{hostname} | {otp_code}")
output_lines.append("=" * 60)
output_lines.append("=" * 25)
result_text = "\n".join(output_lines)
self._show_result(result_text)
utils/tui.py → TUI/TUI.py
+155 -136
View File
@@ -1,11 +1,12 @@
import logging
import os
import sys
from typing import Optional
import dotenv
from dotenv import set_key
from textual.app import App, ComposeResult
from textual.containers import Vertical
from textual.message import Message
from textual.reactive import reactive
from textual.screen import Screen
from textual.widgets import (
@@ -18,22 +19,27 @@ from textual.widgets import (
Tabs,
)
from flows.otp import otp_activities_by_agent, otp_revoke
from flows.otp import 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.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 TUI.agentmoveoperations import AgentMoveOperations
from TUI.moveagentworkflowscreen import MoveAgentWorkflowScreen
from TUI.multiagentselector import MultiAgentSelector
from TUI.OTP_generate import OTPGenerator
from TUI.otpactivityscreen import OTPActivitiesScreen
from TUI.otpworkflowscreen import OTPWorkflowScreen
from TUI.policytreewidget import PolicyTreeWidget
from TUI.quietagentworkflowscreen import QuietAgentWorkflowScreen
from TUI.resultsdisplay import ResultsDisplay
from TUI.theme_amber_terminal import get_amber_terminal_theme
from TUI.theme_retro_terminal import get_retro_terminal_theme
from TUI.themeselector import ThemeSelector
from utils.configmanager import get_user_value, load_env, save_user_config
from utils.setup import get_base_directory
from utils.utils import open_directory
from widgets.multiagentselector import MultiAgentSelector
from widgets.OTP_generate import OTPGenerator
from widgets.policytreewidget import PolicyTreeWidget
from widgets.themeselector import ThemeSelector
dotenv.load_dotenv()
@@ -51,80 +57,45 @@ logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
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.
Store the chosen Textual theme in the user's config using the config manager.
No need to touch .env - config manager handles everything.
"""
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)
save_user_config(config_dir, {"TEXTUAL_THEME": theme_name})
logger.debug("Updated user config with TEXTUAL_THEME=%s", theme_name)
except Exception as exc:
logger.error("Failed to save TEXTUAL_THEME: %s", exc)
# ---------------------------------------------------------------------------
# 1) SCREEN
# ---------------------------------------------------------------------------
class MainMenuScreen(Screen):
api: AirlockAPIWrapper
current_tab = reactive("")
BUTTON_DEFS = {
"find": [
("🔍 - Device Search", "find_device_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"),
"agent_actions": [
(
"🖥️ - Find, Move, or Generate OTP for Agents",
"move_agent_workflow_button",
),
("📊 - Review and appove OTP Activities", "otp_activities_button"),
("📇 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"),
],
"policy": [
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
("🔄 - Update Audit Policies", "policy_audit_update_button"),
("❌ - Revoke OTPs", "otp_revoke_button"),
],
}
def __init__(self, api: AirlockAPIWrapper) -> None:
def __init__(self) -> None:
super().__init__()
self.api = api
self.extras = load_env("EXTRAS")
self.extras = get_user_value("EXTRAS", str, "NOTTODAY")
wd = load_env("WORKING_DIR") or os.getcwd()
if not os.path.isdir(wd):
wd = os.getcwd()
@@ -143,23 +114,22 @@ 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")
api = self.app.api
self.switch_tab("agent_actions")
# focus helpers
def _get_content_buttons(self) -> list[Button]:
@@ -220,7 +190,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)
@@ -241,7 +211,6 @@ class MainMenuScreen(Screen):
"""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),
@@ -250,7 +219,6 @@ class MainMenuScreen(Screen):
message.duration,
)
# Set up the job to run the OTP generation
_PENDING_JOB = (
"otp_workflow",
message.devices,
@@ -261,6 +229,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,47 +287,58 @@ 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":
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
event.stop()
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":
self.app.push_screen(
QuietAgentWorkflowScreen(self.app.api, self.app.policies)
)
event.stop()
return
case "otp_activities_button":
_PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {})
self.app.push_screen(OTPActivitiesScreen())
event.stop()
return
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
# Only exit the UI loop when we explicitly queued a legacy job.
# The original flow used `self.app.exit()` after setting _PENDING_JOB so
# the outer loop could run legacy code. Keep that behavior only for legacy jobs.
logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB)
if _PENDING_JOB and _PENDING_JOB[0] == "legacy":
# let the main loop pick up the legacy job
self.app.exit()
# ---------------------------------------------------------------------------
# 2) APP
# ---------------------------------------------------------------------------
class Loxide(App):
class Loxide(App[Message]):
api: AirlockAPIWrapper
working_dir: str
policies: Optional[list[Policy]]
devices: Optional[list[Agent]]
CSS = """
#logo {
width: 100%;
@@ -326,48 +346,67 @@ class Loxide(App):
text-align: center;
}
"""
BINDINGS = [
("q", "quit", "Quit"),
("d", "open_dir", "Open Directory"),
("f", "open_fe", "Launch Explorer"),
("r", "refresh", "Refresh"),
]
def __init__(self, api: AirlockAPIWrapper):
self._textual_theme = load_env("TEXTUAL_THEME") or "nord"
self._textual_theme = get_user_value("TEXTUAL_THEME", str, "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()
# 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))
self.push_screen(MainMenuScreen())
def action_refresh(self) -> None:
self.refresh_data()
def action_quit(self) -> None:
global _PENDING_JOB
_PENDING_JOB = None
self.exit()
def action_open_dir(self) -> None:
screen = self.screen_stack[-1]
if isinstance(screen, MainMenuScreen):
if screen.current_tab != "dir":
screen.switch_tab("dir")
def action_open_fe(self) -> None:
"""Open the working directory in the OS file manager (footer binding)."""
path_to_open = self.working_dir or os.getcwd()
try:
open_directory(path_to_open)
except Exception as exc:
logger.error("Failed to open directory %s: %s", path_to_open, exc)
self.bell() # optional feedback
# ---------------------------------------------------------------------------
@@ -380,7 +419,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 +435,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 +449,31 @@ def _run_legacy_job(func, args, kwargs) -> None:
# ---------------------------------------------------------------------------
def run_Loxide(api: AirlockAPIWrapper) -> None:
global _PENDING_JOB
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:
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 +482,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
TUI/agentmoveoperations.py
+733
View File
@@ -0,0 +1,733 @@
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 TUI.OTP_generate import OTPGenerator
from TUI.otpworkflowscreen import OTPWorkflowScreen
from TUI.policyselectorscreen import PolicySelectorScreen
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_system_json
policy_relationship_map = get_system_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,
)
)
TUI/allowlistselectionscreen.py
+657
View File
@@ -0,0 +1,657 @@
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.back_btn = Button("⬅ Back", id="back_btn")
self.add_btn = Button("➕ Add to Allowlist", id="add_to_allowlist_btn")
self.back_btn.styles.width = "50%"
self.add_btn.styles.width = "50%"
self.add_btn.disabled = True # Disabled until allowlist selected
yield self.back_btn
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.back_btn or btn_id == "back_btn":
await self.app.pop_screen()
event.stop()
return
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)
# 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"
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("b", "back", "Back"),
Binding("r", "refresh", "Refresh Allowlists"),
Binding("enter", "confirm", "Add to Allowlist"),
]
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_back(self) -> None:
"""Go back to previous screen."""
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()
TUI/moveagentworkflowscreen.py
+61
View File
@@ -0,0 +1,61 @@
from typing import List, Optional
from textual.app import ComposeResult
from textual.screen import Screen
from models.agent import Agent
from TUI.agentmoveoperations import AgentMoveOperations
from TUI.multiagentselector import MultiAgentSelector
from TUI.resultsdisplay import ResultsDisplay
class MoveAgentWorkflowScreen(Screen):
"""Screen that handles the agent movement workflow."""
def __init__(self, all_agents: Optional[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)
)
widgets/multiagentselector.py → TUI/multiagentselector.py
+20 -9
View File
@@ -1,12 +1,20 @@
import difflib
import re
from typing import List
from typing import List, Optional
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
@@ -17,7 +25,7 @@ class MultiAgentSelector(Widget):
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"
@@ -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"
TUI/otpactivityscreen.py
+897
View File
@@ -0,0 +1,897 @@
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.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 Back and Continue buttons. The Continue button pushes ActivityDetailScreen with the
currently-loaded activities.
"""
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
# Buttons area at the bottom (Back, Continue)
with Horizontal(id="activity_buttons"):
# Back takes left side, Continue right side
self.back_btn = Button("Back", id="activity_back_btn")
self.continue_btn = Button("Continue", id="activity_continue_btn")
# Stretch buttons nicely
self.back_btn.styles.width = "50%"
self.continue_btn.styles.width = "50%"
yield self.back_btn
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 Back / Continue buttons 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)
)
# ---- Back ----
if btn is self.back_btn or btn_id == getattr(self.back_btn, "id", None):
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
return
# ---- 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.")
await self.post_message(
Static("No activities loaded to continue with.")
)
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 + Back buttons.
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.detail_back_btn = Button("Back", id="detail_back_btn")
self.add_allowlist_btn = Button(
"📋 Add Selected to Allowlist", id="add_allowlist_btn"
)
yield self.add_allowlist_btn
yield self.detail_back_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)
total = 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.detail_back_btn or btn_id == "detail_back_btn":
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
return
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("b", "back", "Back"),
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_back(self) -> None:
try:
await self.app.pop_screen()
except Exception:
logger.debug("ActivityDetailScreen.action_back pop_screen failed.")
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("r", "refresh_sessions", "Refresh Sessions"),
Binding("e", "export_activities", "Export activities"),
Binding("q", "quit", "Quit"),
]
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)
# 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)
async def action_quit(self) -> None:
# Pop the screen or exit app
await self.app.pop_screen()
# 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)
TUI/otpworkflowscreen.py
+24
View File
@@ -0,0 +1,24 @@
# otp_workflow_screen.py
from typing import List, Optional
from textual.app import ComposeResult
from textual.screen import Screen
from models.agent import Agent
from TUI.OTP_generate import OTPGenerator
class OTPWorkflowScreen(Screen):
"""Screen that handles the OTP generation workflow without agent selection."""
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 on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
"""Handle OTP generation request - pass it up to the app level if needed."""
TUI/policyselector.py
+488
View File
@@ -0,0 +1,488 @@
"""
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, 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
"""
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
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))
TUI/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 TUI.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)
widgets/policytreewidget.py → TUI/policytreewidget.py
+112 -44
View File
@@ -1,9 +1,10 @@
from collections import defaultdict
import logging
from rich.text import Text
from textual.containers import Horizontal, Vertical
from textual.widget import Widget
from textual.widgets import Input, OptionList, Static, Tree
from textual.widgets import Input, OptionList, Static, Switch, Tree
from textual.widgets.option_list import Option
logger = logging.getLogger(__name__)
@@ -17,67 +18,145 @@ class PolicyTreeWidget(Widget):
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
def compose(self):
# Left: Policy Tree
policy_tree = Tree("Policies", id="policy_tree")
# 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%"
# Right: Search + Details
# Create the search box and details pane
label = Static("Device Search:")
search_box = Input(
placeholder="Search policies or devices...", id="tree_search"
)
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
yield details_pane
def on_mount(self) -> None:
"""Build the tree after mounting."""
self._build_tree()
def _build_tree(self) -> None:
"""Build the policy tree structure."""
self._precompute_leaf_counts()
# Update root label with total leaf count
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)
node_map = {}
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
self._build_tree()
# Expand the root node
policy_tree.root.expand()
def _precompute_leaf_counts(self):
"""Precompute leaf counts for each policy group."""
device_counts = defaultdict(int)
for device in self.devices:
device_counts[device.groupid] += 1
child_map = defaultdict(list)
for policy in self.policies:
child_map[policy.parent].append(policy.groupid)
def count_leaves(groupid):
count = device_counts[groupid]
for child_id in child_map.get(groupid, []):
count += count_leaves(child_id)
self.leaf_counts[groupid] = count
return count
# Top-level policies
for policy in self.policies:
if policy.parent == "global-policy-settings":
node = policy_tree.root.add(label=policy.name, data=policy)
count_leaves(policy.groupid)
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)})"
node = policy_tree.root.add(label=label, data=policy)
node_map[policy.groupid] = node
# Child policies
# Sort and add child policies
children_by_parent = defaultdict(list)
for policy in self.policies:
parent_id = policy.parent
if parent_id in node_map:
parent_node = node_map[parent_id]
node = parent_node.add(label=policy.name, data=policy)
if policy.parent != "global-policy-settings":
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:
label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})"
node = parent_node.add(label=label, data=policy)
node_map[policy.groupid] = node
# Devices under policies
# Add devices (leaf nodes) - always sort alphabetically
devices_by_group = defaultdict(list)
for device in self.devices:
group_id = device.groupid
if group_id in node_map:
parent_node = node_map[group_id]
label = device.hostname
parent_node.add(label=label, data=device)
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):
"""Helper to recursively collect all nodes from a tree."""
all_nodes.append(node)
for child in node.children:
self._collect_tree_nodes(child, all_nodes)
def _remove_match_selector(self):
"""Safely remove match selector widgets."""
try:
existing = self.query("#match_selector")
for widget in existing:
@@ -87,20 +166,16 @@ class PolicyTreeWidget(Widget):
logger.debug("Failed to remove match_selector: %s", exc)
def on_tree_node_selected(self, message: Tree.NodeSelected) -> None:
"""Handle tree node selection."""
node = message.node
data = node.data
details_pane = self.query_one("#details_pane", Static)
# Reset previous highlight
if self.last_highlighted_node is not None:
original_label = str(self.last_highlighted_node.label).strip()
# Remove any styling
if isinstance(self.last_highlighted_node.label, Text):
original_label = self.last_highlighted_node.label.plain
self.last_highlighted_node.set_label(original_label)
# Apply highlight to current node
label_text = str(node.label).strip()
if isinstance(node.label, Text):
label_text = node.label.plain
@@ -108,9 +183,7 @@ class PolicyTreeWidget(Widget):
node.set_label(highlighted_label)
self.last_highlighted_node = node
# Update details pane
if data:
# Work with dataclass objects using __dict__
details = "\n".join(
f"{key}: {value}" for key, value in data.__dict__.items()
)
@@ -118,12 +191,16 @@ class PolicyTreeWidget(Widget):
details = f"Selected: {node.label}"
details_pane.update(details)
# Stop event from bubbling
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:
"""Handle search input submission."""
# Remove existing match selector FIRST
self._remove_match_selector()
query = message.value.strip().lower()
@@ -138,7 +215,6 @@ class PolicyTreeWidget(Widget):
label_text = str(node.label).lower()
label_to_node[label_text] = node
if node.data:
# Use __dict__ for dataclass objects
data_dict = (
node.data.__dict__ if hasattr(node.data, "__dict__") else node.data
)
@@ -146,18 +222,15 @@ class PolicyTreeWidget(Widget):
if isinstance(value, str):
label_to_node[value.lower()] = node
# Wildcard-style substring match
matches = sorted([label for label in label_to_node if query in label])
if matches:
# Try to reuse existing match_selector or create new one
try:
option_list = self.query_one("#match_selector", OptionList)
option_list.clear_options()
option_list.display = True # Ensure it's visible
option_list.display = True
except:
option_list = OptionList(id="match_selector")
# Mount to the details pane's parent (the Vertical container)
details_pane.parent.mount(option_list)
for label in matches:
@@ -165,17 +238,14 @@ class PolicyTreeWidget(Widget):
details_pane.update(f"Found {len(matches)} matches. Select one below.")
else:
# Hide or remove the match_selector when no matches
self._remove_match_selector()
details_pane.update("No matches found.")
def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None:
"""Handle selection from search results."""
selected_id = event.option.id.replace("match_", "")
tree = self.query_one("#policy_tree", Tree)
details_pane = self.query_one("#details_pane", Static)
# Find the node
all_nodes = []
self._collect_tree_nodes(tree.root, all_nodes)
@@ -183,7 +253,6 @@ class PolicyTreeWidget(Widget):
match_node = label_to_node.get(selected_id.lower())
if match_node:
# Expand path (original working logic)
node = match_node
path = []
while node:
@@ -198,7 +267,6 @@ class PolicyTreeWidget(Widget):
match_node.set_label(Text(str(match_node.label), style="reverse bold"))
details_pane.update(f"Selected: {match_node.label}")
# Remove the match_selector after selection
try:
option_list = self.query_one("#match_selector", OptionList)
option_list.remove()
TUI/quietagentworkflowscreen.py
+848
View File
@@ -0,0 +1,848 @@
# 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, Static
from models.policy import Policy
from services.API import AirlockAPIWrapper
from services.policyhandler import getPolicyInfo
from TUI.policyselector import PolicySelector
from utils.configmanager import load_env
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. View categorized agents (enforce ready vs. non-enforce ready)
3. Select target policies for each category
4. 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)
quiet_days (int): Number of days without execution to be considered quiet (default: 45)
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"),
]
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 # Fixed as per requirements
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_quiet_days": "Step 2: Select Quiet Time Period",
"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_quiet_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_quiet_days_selection(self) -> None:
"""Show the quiet days selection screen."""
self.workflow_stage = "select_quiet_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"
f"History Period: {self.history_days} days\n\n"
"Select quiet time period (days without untrusted execution):",
id="quiet_days_info",
)
info_widget.styles.margin = (0, 0, 2, 0)
content.mount(info_widget)
# Create button container and mount it first
button_container = Vertical(id="quiet_days_buttons")
button_container.styles.height = "auto"
content.mount(button_container)
# Now add buttons to the mounted container
for days in [15, 30, 45, 60]:
btn = Button(
f"{days} days {'(Default)' if days == 45 else ''}",
id=f"quiet_days_{days}",
classes="quiet_day_btn",
)
btn.styles.width = "100%"
btn.styles.margin = (0, 0, 1, 0)
button_container.mount(btn)
back_btn = Button("← Back", id="back_to_policy_selection")
back_btn.styles.width = "100%"
back_btn.styles.margin = (2, 0, 0, 0)
button_container.mount(back_btn)
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button press events."""
button_id = event.button.id
# Quiet days selection buttons
if button_id and button_id.startswith("quiet_days_"):
days = int(button_id.split("_")[-1])
self.quiet_days = days
logger.info(f"Selected quiet days: {days}")
self._start_analysis()
return
# Navigation buttons
if button_id == "back_to_policy_selection":
self._show_policy_selection()
return
if button_id == "back_to_results":
self._show_results()
return
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 _start_analysis(self) -> None:
"""Start the agent activity analysis."""
self.workflow_stage = "analyzing"
content = self.query_one("#content_area", Vertical)
content.remove_children()
# Show analyzing message with detailed steps
analyzing_msg = Static(
f"📊 Analyzing Agent Activity\n"
f"{'=' * 50}\n\n"
f"Policy: {self.selected_policy.name}\n"
f"History Period: {self.history_days} days\n"
f"Quiet Threshold: {self.quiet_days} days\n\n"
f"Progress:\n"
f"⏳ Step 1/4: Fetching agents from policy...\n"
f"⏱️ Step 2/4: Pulling execution history (this may take a moment)...\n"
f"⏱️ Step 3/4: Analyzing activity patterns...\n"
f"⏱️ Step 4/4: Categorizing agents...\n\n"
f"Please wait - this operation cannot be cancelled.",
id="analyzing_message",
)
analyzing_msg.styles.margin = (2, 1)
content.mount(analyzing_msg)
# Show notification
self.app.notify(
"Starting analysis - this may take several minutes for large policies",
severity="information",
timeout=5,
)
# Perform the analysis asynchronously
self.call_later(self._perform_analysis)
def _perform_analysis(self) -> None:
"""Perform the actual agent activity analysis."""
try:
# Update status: Fetching agents
self._update_analysis_status("Step 1/4: Fetching agents from policy...")
# 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
agent_count = len(agents)
self.app.notify(
f"Found {agent_count} agents - fetching execution history...",
severity="information",
timeout=3,
)
# Update status: Pulling execution history
self._update_analysis_status(
f"Step 2/4: Pulling execution history for {agent_count} agents...\n"
f"(This may take several minutes - progress shown in terminal)"
)
# Get execution history (this shows progress bars in terminal via airlock_libs)
policy_exec_history = getPolicyInfo(
self.api, self.selected_policy, [1, 2, 6, 7], self.history_days
)
# Update status: Analyzing patterns
self._update_analysis_status("Step 3/4: Analyzing activity patterns...")
self.app.notify(
"History retrieved - analyzing patterns...",
severity="information",
timeout=2,
)
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
)
# Update status: Categorizing
self._update_analysis_status("Step 4/4: Categorizing agents...")
# 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"] == True].copy()
self.non_enforce_ready_df = agents[agents["enforce_ready"] == False].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 _update_analysis_status(self, status_text: str) -> None:
"""Update the analysis status message."""
try:
analyzing_msg = self.query_one("#analyzing_message", Static)
# Build updated message
updated_text = (
f"📊 Analyzing Agent Activity\n"
f"{'=' * 50}\n\n"
f"Policy: {self.selected_policy.name}\n"
f"History Period: {self.history_days} days\n"
f"Quiet Threshold: {self.quiet_days} days\n\n"
f"Progress:\n"
f"✅ {status_text}\n\n"
f"Please wait - this operation cannot be cancelled."
)
analyzing_msg.update(updated_text)
except Exception as e:
logger.debug(f"Could not update analysis status: {e}")
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"\n✅ Enforce 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"\n❌ Non-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("\n⚠️ No 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:
result = 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:
result = 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_quiet_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()
TUI/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()
TUI/theme_amber_terminal.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("#49331a"), # 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;
}
"""
TUI/theme_retro_terminal.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 → TUI/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)
airlock_libs/Cargo.lock
Generated
+1056 -191
View File
File diff suppressed because it is too large Load Diff
airlock_libs/Cargo.toml
+9 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "airlock_libs"
version = "2.0.0"
version = "3.1.2"
edition = "2024"
[lib]
@@ -10,12 +10,20 @@ crate-type = ["cdylib"]
chrono = "0.4.42"
indicatif = "0.18.2"
mongodb = "3.3.0"
opentelemetry = { version = "0.18.0", features = ["rt-tokio", "metrics", "trace"] }
opentelemetry-otlp = { version = "0.11.0", features = ["trace", "metrics"] }
opentelemetry-semantic-conventions = { version = "0.10.0" }
opentelemetry-proto = { version = "0.1.0"}
pyo3 = { version = "0.27.0", features = ["extension-module", "generate-import-lib"] }
reqwest = { version = "0.12.24", features = ["json", "native-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.8.2", features = ["tls-roots"] }
tracing = "0.1.41"
tracing-subscriber = "0.3.20"
tracing-opentelemetry = "0.32.0"
[package.metadata.maturin]
generate-abi-stubs = true
airlock_libs/pyproject.toml
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "maturin"
[project]
name = "airlock_libs"
version = "2.0.0"
version = "3.1.2"
description = "Airlock Digital API Wrapper"
readme = "README.md"
license = { text = "AGPL-3.0-only" }
airlock_libs/src/services.rs
+108 -6
View File
@@ -1,6 +1,13 @@
use chrono::{Duration, Local, NaiveDate};
use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
use mongodb::bson::oid::ObjectId;
use opentelemetry::global::shutdown_tracer_provider;
use opentelemetry::sdk::Resource;
use opentelemetry::trace::noop::NoopTracerProvider;
use opentelemetry::trace::{Status, TraceContextExt, TraceError};
use opentelemetry::{Context, KeyValue, sdk::trace as sdktrace, trace::Tracer};
use opentelemetry::{Key, global};
use opentelemetry_otlp::WithExportConfig;
use pyo3::{prelude::*, types::PyString};
use reqwest::{
Client,
@@ -18,6 +25,12 @@ use std::{
str::FromStr,
};
#[derive(Deserialize, Debug)]
struct TelemetryConfig {
TELEMETRY: bool,
TELEM_URL: Option<String>,
}
#[derive(Debug, Deserialize, Serialize)]
struct ApiResponse {
error: String,
@@ -61,6 +74,12 @@ pub fn pull_policy_exec_histories(
exec_types: String,
days: i64,
) -> Py<PyString> {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let _ = init_tracer();
});
let tracer = global::tracer("global_tracer");
let _cx = Context::new();
let file_path: PathBuf = format!(
"{}\\cache\\chunkinator.json",
get_base_directory().display()
@@ -93,13 +112,43 @@ pub fn pull_policy_exec_histories(
.unwrap(),
);
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
let client = build_client(py, &py_self);
let client = tracer.in_span("Building HTTP Client", |cx| {
let client_result = build_client(py, &py_self);
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"));
panic!("Failed to Build Client: {:?}", client_result);
}
}
});
let api: Py<PyAny> = py_self;
let cutoff = Local::now().naive_local() - Duration::days(days);
let mut f = File::open(&writeable_filepath).unwrap();
tracer.in_span("Airlock Data Retreival", |cx| {
let span = cx.span();
span.set_attribute(Key::new("Days").string(days.to_string().to_string()));
loop {
f.seek(SeekFrom::Start(0)).unwrap();
let execution_histories = history_logging(
let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
let results: ApiResponse = history_logging(
py,
&api,
&exec_types,
@@ -107,11 +156,18 @@ pub fn pull_policy_exec_histories(
&policy_names,
&client,
);
cx.span().set_attribute(KeyValue::new(
"items_in_response",
results.response.exechistories.len().to_string(),
));
results
});
let parsed_responses = execution_histories.response.exechistories;
if parsed_responses.is_empty() {
break;
}
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() {
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 =
@@ -183,12 +239,14 @@ pub fn pull_policy_exec_histories(
progress_bar.set_message("Total Percent Complete");
}
}
});
progress_bar.finish_with_message("All Checkpoints Complete");
let return_data = fs::read_to_string(&writeable_filepath).unwrap();
shutdown_tracer_provider();
PyString::new(py, &return_data).into()
}
fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Client {
fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Result<reqwest::Client, reqwest::Error> {
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();
@@ -205,9 +263,8 @@ fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Client {
Client::builder()
.danger_accept_invalid_certs(true)
.default_headers(header_map)
.timeout(std::time::Duration::from_secs(30))
.timeout(std::time::Duration::from_secs(300))
.build()
.unwrap()
}
#[tokio::main]
@@ -276,3 +333,48 @@ fn skipback(days: i64) -> ObjectId {
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
}
fn load_telemetry_config() -> TelemetryConfig {
let cfg_path = get_base_directory().join("config\\user_config.json");
if !cfg_path.exists() {
return TelemetryConfig {
TELEMETRY: false,
TELEM_URL: None,
};
}
match fs::read_to_string(&cfg_path) {
Ok(contents) => {
serde_json::from_str::<TelemetryConfig>(&contents).unwrap_or(TelemetryConfig {
TELEMETRY: false,
TELEM_URL: None,
})
}
Err(_) => TelemetryConfig {
TELEMETRY: false,
TELEM_URL: None,
},
}
}
fn init_tracer() -> Result<Option<sdktrace::Tracer>, TraceError> {
let cfg = load_telemetry_config();
if !cfg.TELEMETRY {
global::set_tracer_provider(NoopTracerProvider::new());
return Ok(None);
}
let endpoint = cfg.TELEM_URL.unwrap_or_default();
let tracer =
opentelemetry_otlp::new_pipeline()
.tracing()
.with_exporter(
opentelemetry_otlp::new_exporter()
.tonic()
.with_endpoint(endpoint),
)
.with_trace_config(sdktrace::config().with_resource(Resource::new(vec![
KeyValue::new("service.name", "LoxideLibs"),
])))
.install_simple()
.unwrap();
Ok(Some(tracer))
}
default_system_config.json
+32 -7
View File
@@ -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": {}
}
README.md → docs/README.md
View File
flows/localApproval.py
+201 -269
View File
@@ -1,311 +1,243 @@
# 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_system_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"
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_system_json("POLICY_MAP_ENF_AUD", "{}")
self.username = (
username or os.getenv("USERNAME") or os.getenv("USER") or "unknown"
)
previous_run.to_parquet(
f"{base_dir}\\cache\\last_local_approval.parquet", index=False
)
os.remove(f"{base_dir}\\cache\\newest_local_approval.parquet")
# Only keep rows presumably created by the generate local approval function
local_approval = local_approval[
local_approval["purpose"].str.startswith("🎫 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)
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"),
]
local_approval["batchid"] = local_approval["purpose"].apply(
lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1)
)
# Display duration options
print(colorText("\n⏱️ Select Local Approval Duration:", "white"))
print(colorText("=" * 50, "white"))
if not local_approval.empty:
local_approval.to_parquet(
f"{base_dir}\\cache\\newest_local_approval.parquet", index=False
)
for i, (minutes, label) in enumerate(duration_options, start=1):
print(f" {i}. {label} ({minutes} minutes)")
return local_approval
def scheduleAddingLAHashes(api: AirlockAPIWrapper):
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)
print(colorText("=" * 50, "white"))
# Get user selection
try:
register_function("add_hash", returnFromLocalApproval)
register_function("move_device", moveAgentToRelatedPolicy)
except Exception as e:
logger.warning(f"Failed to register functions: {e}")
return
choice = int(get_sanitized_input("\nEnter the number of your choice: "))
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)
)
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}"
)
except Exception as e:
logger.warning(f"Failed to process batch {batchid}: {e}")
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())
# 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}")
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.debug("Invalid Input")
return
except ValueError:
print(colorText("❌ Invalid input. Please enter a number.", "red"))
logger.debug("Invalid Input")
print(colorText("❌ Invalid choice.", "red"))
logger.warning("Invalid duration choice")
return
agents = selectAgents(api)
batch = int(time.time())
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.debug("No agents found or error retrieving agents")
print(colorText("❌ No agents found or error retrieving agents.", "red"))
logger.warning("No agents selected or error retrieving agents")
return
for agent in agents:
try:
addLocalApproval(api, batch, duration_selected, agent.agentid)
moveAgentToRelatedPolicy(api, agent, "audit")
except Exception as e:
print(colorText(f"❌ Error processing agent {agent.hostname}: {e}", "red"))
# 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()
def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid):
if confirm != "y":
print(colorText("❌ Operation cancelled.", "yellow"))
return
purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
api.otp_generate(agentid, duration_selected, purpose)
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", "{}")
# 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()
# Process the batch
batch_id, success_count, failure_count = self.create_local_approval_batch(
agents, duration_minutes
)
# Index last_agents by hostname for quick lookup
last_agent_map = {agent.hostname: agent for agent in last_agents}
# Display summary
self._display_summary(batch_id, duration_label, success_count, failure_count)
# Result buckets
newly_added = []
same_policy = []
moved_to_audit = []
moved_to_enforcement = []
unusual_move = []
def _display_summary(
self, batch_id: int, duration_label: str, success_count: int, failure_count: int
):
"""
Display operation summary.
for current in current_agents:
previous = last_agent_map.get(current.hostname)
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"))
if not previous:
newly_added.append(current)
continue
print(colorText(f"✓ Successfully processed: {success_count}", "green"))
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)
if failure_count > 0:
print(colorText(f"✗ Failed: {failure_count}", "red"))
# Return all five DataFrames
return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move
print(colorText(f"\n📦 Batch ID: {batch_id}", "cyan"))
print(colorText(f"⏱️ Duration: {duration_label}", "cyan"))
def getNewLocalApprovals(api: AirlockAPIWrapper):
working_dir = load_env("WORKING_DIR")
current_la = getLocalApprovals(api)
# 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)
# Create composite keys
current_la["key"] = (
current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str)
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",
)
old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str)
# Find new entries
new_entries = current_la[~current_la["key"].isin(old_la["key"])]
# 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]
# Save current approvals for next run
current_la.drop(columns=["key"], inplace=True)
current_la.to_parquet(old_la_path, index=False)
return recent_entries
print(colorText("=" * 60 + "\n", "white"))
flows/otp.py
+1 -119
View File
@@ -14,136 +14,18 @@
# 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
from utils.utils import 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()
flows/prepPolicy.py
+29 -29
View File
@@ -25,7 +25,7 @@ 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.configmanager import get_system_list, get_system_value, load_env
from utils.selector import Selector
from utils.utils import (
areYouSure,
@@ -88,7 +88,7 @@ def sortHashes(
):
working_dir = load_env("WORKING_DIR")
history_days = Selector.select_value(
prompt="Enter how many days of history to pull (1–150): ",
prompt="Enter how many days of history to pull (1–150): ",
value_type=int,
valid_range=(1, 150),
)
@@ -157,7 +157,7 @@ def buildPathsandPublishers(selected_policies: List[Policy], split):
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)
path_exclusion_constant = get_system_value("PATH_EXCLUSION_CONST", cast_type=int)
if os.path.exists(path1):
df1 = pd.read_csv(path1)
@@ -227,7 +227,7 @@ def buildPathsandPublishers(selected_policies: List[Policy], split):
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", "[]"))
pattern = regulator(get_system_list("BAD_PUBLISHERS"))
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
publist = publist[["publisher"]]
publist.sort_values(by="publisher", inplace=True)
@@ -313,7 +313,7 @@ def buildPreflights(selected_policies: List[Policy]):
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type=int)
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
def clean_split(path):
if not isinstance(path, (str, bytes, os.PathLike)):
@@ -385,8 +385,8 @@ def calculatePath(approved_hashes, path_exclusion_constant, split):
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)
badpathparts = get_system_list("BAD_PATH_PARTS")
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
processed_dfs = []
@@ -655,7 +655,7 @@ def section_header(title):
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR")
section_header("🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒")
section_header("🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒")
print(
colorText(
"\nSequentially follow these steps to prepare a policy for enforcement:",
@@ -670,11 +670,11 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
)
)
if not selected_policies:
print(colorText(" [✗] No policies have been chosen", "red"))
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"))
print(colorText(f" [✓] {policy.name}", "green"))
# Step 2: Destination Policy and Allowlist
print(
@@ -683,22 +683,22 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
if destination_policy:
print(
colorText(
f" [✓] {destination_policy[0].name} has been selected as the destination policy",
f" [✓] {destination_policy[0].name} has been selected as the destination policy",
"green",
)
)
else:
print(colorText(" [✗] No destination policy has been chosen", "red"))
print(colorText(" [✗] No destination policy has been chosen", "red"))
if destination_allowlist:
print(
colorText(
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
"green",
)
)
else:
print(colorText(" [✗] No allowlist has been chosen", "red"))
print(colorText(" [✗] No allowlist has been chosen", "red"))
# Step 3: Data Preparation
print(
@@ -713,9 +713,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print(
colorText(
(
" [✓] Data has been fetched"
" [✓] Data has been fetched"
if os.path.exists(review_path)
else " [✗] Data has not been fetched"
else " [✗] Data has not been fetched"
),
"green" if os.path.exists(review_path) else "red",
)
@@ -723,7 +723,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
else:
print(
colorText(
" [✗] No policies selected, cannot check data fetch status", "red"
" [✗] No policies selected, cannot check data fetch status", "red"
)
)
@@ -756,9 +756,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print(
colorText(
(
" [✓] Reviewed hashes have been loaded"
" [✓] Reviewed hashes have been loaded"
if os.path.exists(approved_path)
else " [✗] Reviewed hashes have not been loaded"
else " [✗] Reviewed hashes have not been loaded"
),
"green" if os.path.exists(approved_path) else "red",
)
@@ -766,9 +766,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print(
colorText(
(
" [✓] Path review list created"
" [✓] Path review list created"
if os.path.exists(second_review_path)
else " [✗] Path review list has not been created"
else " [✗] Path review list has not been created"
),
"green" if os.path.exists(second_review_path) else "red",
)
@@ -776,7 +776,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
else:
print(
colorText(
" [✗] No policies selected, cannot check reviewed hashes or path list",
" [✗] No policies selected, cannot check reviewed hashes or path list",
"red",
)
)
@@ -812,9 +812,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print(
colorText(
(
" [✓] Reviewed path list detected"
" [✓] Reviewed path list detected"
if os.path.exists(reviewed_path)
else " [✗] Path review list has not been detected"
else " [✗] Path review list has not been detected"
),
"green" if os.path.exists(reviewed_path) else "red",
)
@@ -825,9 +825,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print(
colorText(
(
" [✓] Preflight Path Exclusion List has been generated"
" [✓] Preflight Path Exclusion List has been generated"
if preflight_ready
else " [✗] Preflight Path Exclusion List has not been generated"
else " [✗] Preflight Path Exclusion List has not been generated"
),
"green" if preflight_ready else "red",
)
@@ -835,7 +835,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
else:
print(
colorText(
" [✗] No policies selected, cannot check preflight status", "red"
" [✗] No policies selected, cannot check preflight status", "red"
)
)
@@ -866,5 +866,5 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print(colorText(" Apply approved hashes to allowlist", "cyan"))
# Utility Options
print(colorText("F. 📂 - Open Working Directory", "cyan"))
print(colorText("B. 🔚 - Back", "cyan"))
print(colorText("F. 📂 - Open Working Directory", "cyan"))
print(colorText("B. 🔚 - Back", "cyan"))
flows/quietAgent.py
-134
View File
@@ -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")
models/execution.py
+11 -11
View File
@@ -28,7 +28,7 @@ import pandas as pd
import airlock_libs
from services.API import AirlockAPIWrapper
from utils.configmanager import get_protected_value, load_env_json
from utils.configmanager import get_system_list, get_system_value
from utils.utils import colorText, 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
@@ -145,13 +145,13 @@ class Hash:
approved_count += 1
except (ValueError, TypeError):
logger.debug(
"Needs Review: Scannermatch score is missing or invalid. — {e}"
"Needs Review: Scannermatch score is missing or invalid. — {e}"
)
hash_obj.at_decision = "needs_review"
needs_review_count += 1
logger.debug(
f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}"
f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}"
)
return hashes
@@ -384,9 +384,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
@@ -443,13 +443,13 @@ class ExecutionHistoryRecord:
approved_count += 1
except (ValueError, TypeError) as e:
logger.debug(
f"Needs Review: Scannermatch score is missing or invalid. — {e}"
f"Needs Review: Scannermatch score is missing or invalid. — {e}"
)
hash_obj.at_decision = "needs_review"
needs_review_count += 1
logger.debug(
f"Final counts — Needs Review: {needs_review_count}, "
f"Final counts — Needs Review: {needs_review_count}, "
f"Approved: {approved_count}, Unapproved: {unapproved_count}"
)
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)
requirements.txt
+1 -1
View File
@@ -11,4 +11,4 @@ urllib3==2.5.0
pyperclip==1.11.0
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
airlock_libs==2.0.0
airlock_libs==3.1.2
screens/otpworkflowscreen.py
-41
View File
@@ -1,41 +0,0 @@
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."""
def __init__(self, all_agents: List[Agent]):
super().__init__()
self.all_agents = all_agents
self.selected_devices = 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 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))
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
services/API.py
+60
View File
@@ -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,6 +300,19 @@ 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]
services/agenthandler.py
+16 -16
View File
@@ -28,7 +28,7 @@ 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.configmanager import get_system_json, load_env
from utils.selector import Selector
from utils.utils import colorText, get_sanitized_input
@@ -38,7 +38,7 @@ 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): ",
prompt="Enter how many days of history to pull (1–150): ",
value_type=int,
valid_range=(1, 150),
)
@@ -60,7 +60,7 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
except Exception as e:
print(
colorText(
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
)
)
continue
@@ -139,7 +139,7 @@ def findAgents(api, return_dataframe):
print(
colorText(
f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
"green",
)
)
@@ -148,7 +148,7 @@ def findAgents(api, return_dataframe):
def collect_device_names() -> List[str]:
print(colorText("🔍 Device Search", "cyan"))
print(colorText("🔍 Device Search", "cyan"))
print(
colorText(
"Enter the device hostnames you'd like to search for, one per line.", "cyan"
@@ -185,7 +185,7 @@ def collect_device_names() -> List[str]:
else:
print(
colorText(
f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
"yellow",
)
)
@@ -235,8 +235,8 @@ def show_unmatched(
]
if unmatched:
logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
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"]):
@@ -248,7 +248,7 @@ 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"))
print(colorText("⚠️ No device names entered.", "red"))
return []
use_exact = choose_match_type()
@@ -261,11 +261,11 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
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"))
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"))
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):
@@ -283,8 +283,8 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
)
if not matched_agents:
logger.debug("❌ No matching devices remain after refinement.")
print(colorText("❌ No matching devices remain after refinement.", "red"))
logger.debug("❌ No matching devices remain after refinement.")
print(colorText("❌ No matching devices remain after refinement.", "red"))
return []
enrich_agents(matched_agents, policies)
@@ -302,10 +302,10 @@ def moveAgentToRelatedPolicy(
Args:
api: AirlockAPIWrapper instance.
agent: Agent object.
policy_relationship_map: Dict mapping enforcement → audit.
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", "{}")
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
if mode == "audit":
if agent.groupid in policy_relationship_map:
services/policyhandler.py
+2 -2
View File
@@ -27,7 +27,7 @@ import tqdm
from models.policy import Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import get_protected_json
from utils.configmanager import get_system_json
from utils.setup import get_base_directory
from utils.utils import areYouSure, colorText, get_sanitized_input
@@ -236,7 +236,7 @@ def skipback(days):
def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
policy_relationship_map = get_system_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")
utils/configmanager.py
+248 -42
View File
@@ -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,46 @@ 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)
# Backwards compatibility aliases (deprecated - use get_system_value instead)
get_protected_value = get_system_value
get_protected_json = get_system_json
load_protected_config = load_system_config
PROTECTED_KEYS = SYSTEM_CONFIG_KEYS # For backwards compatibility
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/setup.py
+29 -76
View File
@@ -13,18 +13,20 @@
# 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,
)
def get_base_directory() -> Path:
@@ -38,7 +40,7 @@ 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, log_level: str = "INFO"):
log_file = log_dir / "Loxide.log"
config = {
@@ -62,12 +64,12 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
"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
"level": "DEBUG", # Always log DEBUG and above to file
"formatter": "detailed", # Use detailed format
},
"console": {
"class": "logging.StreamHandler",
"level": log_level.upper(), # Configurable log level
"level": log_level.upper(), # System-configured level for console
"formatter": "simple", # Use simple format
},
},
@@ -95,55 +97,15 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
logging.getLogger().debug("✅ Logging configured.")
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:
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"},
}
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)
except Exception as e:
logging.warning(f"Failed to write {key} to .env: {e}")
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
"""
base_dir = get_base_directory()
dirs = {
"config": base_dir / "config",
@@ -155,20 +117,30 @@ def setup():
path.mkdir(parents=True, exist_ok=True)
logging.debug(f"{name.capitalize()} directory ensured at: {path}")
# Load system config (immutable)
system_config = load_system_config()
configure_logging(dirs["logs"], system_config.get("LOG_LEVEL", "DEBUG"))
# Configure logging with system-defined log level
log_level = get_system_value("LOG_LEVEL", str, "INFO")
configure_logging(dirs["logs"], log_level)
# Load user config (mutable)
user_config = 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"],
@@ -185,23 +157,4 @@ def setup():
subfolder_path.mkdir(parents=True, exist_ok=True)
logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
user_config = load_user_config(dirs["config"])
merged_config = {**system_config, **user_config}
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)
logging.info("✅ Setup complete")
utils/test.py
View File
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(