Compare commits

..

27 Commits

Author SHA1 Message Date
brotoskyj 729b45f52a Features - Optional Policy Name in Execution Histories
Build Library / Build Library (push) Successful in 4m51s
closes #48
2025-12-18 10:14:58 -05:00
brotoskyj 66bb21ed88 Styling - Progress Bar
Build Library / Build Library (push) Successful in 5m21s
Changed progress bar indicators and added policy name to the progress bar
closes #50
2025-12-17 16:38:53 -05:00
Zarithas a7b659c951 policyprepworkflow: enhancements and fixes
- Added row copy functionality (Ctrl+C)
- Improved row selection visual contrast
- Fixed data state issues when navigating between stages

Server Log tab
- Added new Server Log tab to the main application
- Implemented live filtering with wildcard support
- Enabled auto-refresh capability

Multiagent selector
- Added file loading support for device lists
2025-12-17 12:38:31 -05:00
brotoskyj 53f0b548b0 Bug Fixes
Build Library / Build Library (push) Successful in 4m51s
Changed days to days.to_string() to fix type confusion
closes #47
2025-12-17 11:02:13 -05:00
Zarithas 66eb101c5d Fixed Bitbake 2025-12-16 17:07:35 -05:00
Zarithas 22101c1eba Merge branch 'Zar-Branch' of https://git.racooncity.org/brotoskyj/Airlocktools into Zar-Branch
git commit -m "feat: Multiple UI improvements and new server log functionality

- Add Server Log tab with DataTable display of server activity logs

- Fix keyboard navigation bug in agents tab

- Add execution history viewer for selected agents

- Improve policy tree widget functionality by adding single device operations

- Integrate logging notifications into TUI
  - Add TextualNotificationHandler to setup.py
  - Display ERROR/WARNING/CRITICAL logs as toast notifications
  - Remove terminal output to prevent interference with TUI
  - Logs still written to Loxide.log file

closes #45"
2025-12-16 16:17:25 -05:00
Zarithas fc17c869fc feat: Multiple UI improvements and new server log functionality
- Add Server Log tab with DataTable display of server activity logs

- Fix keyboard navigation bug in agents tab

- Add execution history viewer for selected agents

- Improve policy tree widget functionality by adding single device operations

- Integrate logging notifications into TUI
  - Add TextualNotificationHandler to setup.py
  - Display ERROR/WARNING/CRITICAL logs as toast notifications
  - Remove terminal output to prevent interference with TUI
  - Logs still written to Loxide.log file

closes #45
2025-12-16 16:16:49 -05:00
brotoskyj 1dbbcff5d5 Merge remote-tracking branch 'origin/RustImplementation' into RustImplementation
Build Library / Build Library (push) Successful in 5m18s
2025-12-15 17:04:32 -05:00
brotoskyj f080b0034f Changes:
Features
Moved the init_tracer() function to an implementation in TelemetryConfig for cleaner main file
closes #43

Bug Fixes
Telemetry is now opt-in again. Due to changes in opentelemetry, this required a major overhaul of the telemetryconfig function.
closes #44
2025-12-15 17:04:17 -05:00
Zarithas 57d0f12000 feat: Complete Policy Prep Workflow with UX upgrades, Liftoff API, and TUI merge
- Added intro screen with workflow overview, time estimate, and onboarding controls
- Improved visuals: cleaner checkboxes (/), better loading screen layout
- Enforced mandatory tab reviews for critical steps with warnings and blocked navigation
- Optimized logging: INFO for milestones, DEBUG for internals; cleaner production logs
- Implemented Liftoff API integration: paths, publishers, hashes with granular error handling
- Color-coded completion feedback ( success,  failure,  partial) and detailed summaries
- Consolidated architecture: merged TUI.py into Loxide.py (single entry point, no circular imports)
- Fixed race condition in table creation with concurrency locks
2025-12-15 17:01:58 -05:00
brotoskyj 0dbc744471 Merge remote-tracking branch 'origin/RustImplementation' into RustImplementation 2025-12-15 14:14:16 -05:00
brotoskyj 7a912bddab Bug Fixes
Security Vulnerabilities Patched
RUSTSEC-2025-0009 - Some AES functions may panic when overflow checking is enabled
RUSTSEC-2024-0336 - rustls::Connection::Common::complete_io could fall into an infinite loop based on network input
closes #47
2025-12-15 14:13:57 -05:00
brotoskyj 24211c318b Merge remote-tracking branch 'origin/RustImplementation' into RustImplementation
Build Library / Build Library (push) Successful in 5m2s
2025-12-15 14:12:30 -05:00
brotoskyj 630e0a3cdf Bug Fixes
Security Vulnerabilities Patched
RUSTSEC-2025-0009 - Some AES functions may panic when overflow checking is enabled
RUSTSEC-2024-0336 - rustls::Connection::Common::complete_io could fall into an infinite loop based on network input
2025-12-15 14:12:03 -05:00
Zarithas 797d0f4462 fix(policy-prep): implement table editors and workflow improvements
- Add table editors for Policy Prep workflow
- 'Add to policy' remains a placeholder
- Apply planned tweaks:
  - Replace ballot checkbox with ✓ for selection
  - Relocate loading screen text to bottom:
    'Building Path exclusions and publisher lists...
     This may take a moment for large datasets.'
  - Ensure interaction with all tables before allowing review steps
  - Move excessive logging to debug level
  - Add Step 0 to explain process before user begins

Notes:
Further discussion needed on enforcing table interaction before review.
2025-12-11 16:55:12 -05:00
brotoskyj 59bb97ec4e Refactored LoxideLibs
Build Library / Build Library (push) Successful in 5m30s
1. Added compatibility check, LoxideLibs will now abort the entire program if OS is not linux or windows.
2. Changed the python data extraction compatibility layer, LoxideLibs was calling the extract data function twice, causing very slight overhead. I have now changed this so that the function returns a Struct that is now easily extractable via dot method notation.
2025-12-11 11:47:50 -05:00
brotoskyj 0ac3b54d89 Refactored Progress Bar
Build Library / Build Library (push) Successful in 5m54s
Refactored Progress Bar to remove multiprogress bar and only draw one instance. #36 is still open and not fixed with this push, but I believe this is the way to fix the issue.
Also implemented an Arc Mutex on the progress bar so it can be controlled via different threads.
2025-12-05 17:19:40 -05:00
Zarithas 154a7efcc8 Bug fix for Revoke OTP resolved, no longer crashes when select all is chosen when there are no active sessions, Swapped Revoke OTP and Quiet Hosts locations on menus 2025-12-05 15:30:24 -05:00
Zarithas ab5f00d8e7 Merge branch 'RustImplementation' of https://git.racooncity.org/brotoskyj/Airlocktools into RustImplementation 2025-12-05 15:06:06 -05:00
Zarithas 3ab803c12e Quiet Agent UI improvements 2025-12-05 15:05:46 -05:00
Zarithas 98cb23e5ea Bugfix for Issue 39.
Fixes:
brotoskyj/AirlockTools#39
2025-12-05 14:48:23 -05:00
brotoskyj 0aabbfd36e Style Change
Build Library / Build Library (push) Successful in 6m25s
Cleaned up services.rs file and removed whitespaces
Specified Data Types were needed instead of allowing the compiler to select types
2025-12-05 11:12:29 -05:00
brotoskyj b19eeb6c96 Fixed Progress Bar
Build Library / Build Library (push) Successful in 6m22s
Changed date math so that the last date from the first response from the API is the minimum value for the progress bar. This gives true progress percentages from the first date to the last date on the last request. #36
2025-12-04 11:57:49 -05:00
brotoskyj 76bd3a6087 Feature added - MPSC channel
Build Library / Build Library (push) Successful in 5m44s
closes #38
2025-12-03 17:54:57 -05:00
brotoskyj 6ab9de413f Refactored loxide libs for easier readability
Build Library / Build Library (push) Successful in 5m52s
2025-12-03 11:45:22 -05:00
Zarithas e5b4b9d959 Merge branch 'RustImplementation' of https://git.racooncity.org/brotoskyj/Airlocktools into RustImplementation 2025-12-02 16:23:28 -05:00
Zarithas 1d9caadaf3 Major step towards unification of the UI Implementation of the Back Feature, splitting of TUI files back into subfolders 2025-12-02 16:22:43 -05:00
42 changed files with 8272 additions and 1960 deletions
+538 -6
View File
@@ -23,24 +23,556 @@
import logging
import os
from typing import Optional
import dotenv
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 (
Button,
DirectoryTree,
Footer,
Header,
Static,
Tab,
Tabs,
)
import urllib3
from models.agent import Agent
from models.policy import Policy
from services.API import AirlockAPIWrapper
from services.security import getAPI
from TUI.TUI import run_Loxide
from utils.configmanager import get_system_value
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
from TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
from TUI.Screens.otpactivityscreen import OTPActivitiesScreen
from TUI.Screens.otprevokescreen import OTPRevokeScreen
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
from TUI.Screens.policyprepworkflowscreen import PolicyPrepWorkflowScreen
from TUI.Screens.quietagentworkflowscreen import QuietAgentWorkflowScreen
from TUI.Themes.theme_amber_terminal import get_amber_terminal_theme
from TUI.Themes.theme_retro_terminal import get_retro_terminal_theme
from TUI.Themes.themeselector import ThemeSelector
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
from TUI.Widgets.multiagentselector import MultiAgentSelector
from TUI.Widgets.policytreewidget import PolicyTreeWidget
from TUI.Widgets.resultsdisplay import ResultsDisplay
from TUI.Widgets.serverlogwidget import ServerLogWidget
from utils.configmanager import (
get_system_value,
get_user_value,
load_env,
save_user_config,
)
from utils.setup import get_base_directory, setup
from utils.utils import irtang
from utils.utils import irtang, open_directory
dotenv.load_dotenv()
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
# ---------------------------------------------------------------------------
# GLOBAL STASH
# ---------------------------------------------------------------------------
_APP_RESTART_REASON = None
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
# ---------------------------------------------------------------------------
def _persist_user_theme(theme_name: str) -> None:
"""
Store the chosen Textual theme in the user's config using the config manager.
No need to touch .env - config manager handles everything.
"""
base_dir = get_base_directory()
config_dir = base_dir / "config"
try:
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 = {
"agent_actions": [
{
"label": "🖥️ - Multi-Agent Operations",
"id": "move_agent_workflow_button",
"description": "Select agents to: Move policies, Generate OTPs, Toggle audit/enforcement, View history, Export data",
},
{
"label": "🎫 - Review and approve OTP Activities",
"id": "otp_activities_button",
},
{
"label": "🛑 - Revoke Active OTP Session",
"id": "otp_revoke_button",
},
],
"policy": [
{
"label": "⚖️ - Prepare Policy For Enforcement",
"id": "policy_prep_button",
},
{
"label": "🔕 - Find and Move Quiet Hosts to Enforcement",
"id": "find_quiet_button",
},
],
}
def __init__(self) -> None:
super().__init__()
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()
self.working_dir = wd
def _make_buttons_for(self, tab_id: str) -> Vertical:
defs = self.BUTTON_DEFS.get(tab_id, [])
widgets = []
for item in defs:
# Support both old tuple format and new dict format
if isinstance(item, dict):
label = item["label"]
btn_id = item["id"]
description = item.get("description")
else:
# Old tuple format: (label, id)
label, btn_id = item
description = None
btn = Button(label, id=btn_id)
btn.styles.width = "100%"
widgets.append(btn)
# Add description text if provided
if description:
desc_text = Static(description, classes="button_description")
desc_text.styles.width = "100%"
desc_text.styles.color = "ansi_bright_black"
desc_text.styles.text_align = "center"
desc_text.styles.margin = (0, 0, 1, 0)
widgets.append(desc_text)
return Vertical(*widgets)
def compose(self) -> ComposeResult:
yield Header(show_clock=True, icon="")
tabs = [
Tab("Agents", id="agent_actions"),
Tab("Tree View", id="p_tree"),
Tab("Server Log", id="server_log"),
Tab("Directory", id="dir"),
Tab("Settings", id="settings"),
]
if self.extras == "POLICYPREP":
tabs.insert(2, Tab("Policy Prep", id="policy"))
yield Tabs(*tabs, id="tabs")
yield Vertical(id="content")
yield Footer()
def on_mount(self) -> None:
self.switch_tab("agent_actions")
def on_key(self, event) -> None:
"""Handle up/down arrow keys for button navigation."""
if event.key == "down":
self._focus_nearby_button(1)
event.prevent_default()
event.stop()
elif event.key == "up":
self._focus_nearby_button(-1)
event.prevent_default()
event.stop()
# left/right are handled by Textual's default tab navigation
# focus helpers
def _get_content_buttons(self) -> list[Button]:
content = self.query_one("#content", Vertical)
return list(content.query(Button))
def _focus_first_button(self) -> None:
buttons = self._get_content_buttons()
if buttons:
buttons[0].focus()
def _focus_tabs(self) -> None:
tabs = self.query_one("#tabs", Tabs)
tabs.focus()
def _focus_nearby_button(self, direction: int) -> None:
buttons = self._get_content_buttons()
if not buttons:
return
try:
current = next(i for i, b in enumerate(buttons) if b.has_focus)
except StopIteration:
if direction > 0:
buttons[0].focus()
else:
buttons[-1].focus()
return
if direction < 0 and current == 0:
self._focus_tabs()
return
new_index = current + direction
if 0 <= new_index < len(buttons):
buttons[new_index].focus()
def switch_tab(self, tab_id: str) -> None:
self.current_tab = tab_id
content = self.query_one("#content", Vertical)
content.remove_children()
if tab_id in self.BUTTON_DEFS:
content.mount(self._make_buttons_for(tab_id))
elif tab_id == "server_log":
content.mount(ServerLogWidget(self.app.api))
elif tab_id == "dir":
content.mount(DirectoryTree(self.working_dir, id="dir_tree"))
elif tab_id == "p_tree":
content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
elif tab_id == "settings":
content.mount(ThemeSelector())
else:
content.mount(Static(f"Unknown tab: {tab_id}"))
def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
self.switch_tab(event.tab.id)
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
"""Handle selected agents from AgentSelector."""
global _APP_RESTART_REASON
selected_agents = message.selected_agents
logger.info("Selected agents: %s", selected_agents)
# TODO: Implement actual handling of selected agents
_APP_RESTART_REASON = ("multi_agent_action", selected_agents)
self.app.exit()
def on_theme_selector_theme_selected(
self, message: ThemeSelector.ThemeSelected
) -> None:
"""Handle theme selection from ThemeSelector."""
global _APP_RESTART_REASON
_persist_user_theme(message.theme_name)
_APP_RESTART_REASON = ("restart",)
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_policy_tree_widget_view_execution_history(
self, message: PolicyTreeWidget.ViewExecutionHistory
) -> None:
"""Handle request to view execution history for a device from tree view."""
logger.info("Viewing execution history for device: %s", message.device.hostname)
self.app.push_screen(ExecutionHistoryScreen([message.device]))
message.stop()
def on_policy_tree_widget_generate_otp(
self, message: PolicyTreeWidget.GenerateOTP
) -> None:
"""Handle request to generate OTP for a device from tree view."""
logger.info("Generating OTP for device: %s", message.device.hostname)
self.app.push_screen(OTPWorkflowScreen([message.device]))
message.stop()
def on_policy_tree_widget_toggle_enforcement(
self, message: PolicyTreeWidget.ToggleEnforcement
) -> None:
"""Handle request to toggle enforcement for a device from tree view."""
logger.info("Toggling enforcement for device: %s", message.device.hostname)
try:
from services.agenthandler import moveAgentToRelatedPolicy
from utils.configmanager import get_system_json
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
# Determine current mode and toggle
if message.device.groupid in policy_relationship_map:
# Currently in enforcement, move to audit
result = moveAgentToRelatedPolicy(self.app.api, message.device, "audit")
mode = "audit"
else:
# Currently in audit, move to enforcement
result = moveAgentToRelatedPolicy(
self.app.api, message.device, "enforcement"
)
mode = "enforcement"
logger.info(f"Successfully toggled {message.device.hostname} to {mode}")
# Refresh data at the app level
self.app.refresh_data()
# Refresh the tree widget with new data
try:
tree_widget = self.query_one(PolicyTreeWidget)
tree_widget.refresh_data(self.app.policies, self.app.devices)
except:
pass
except Exception as e:
logger.error(
f"Failed to toggle enforcement for {message.device.hostname}: {e}"
)
self.app.bell()
message.stop()
def on_directory_tree_file_selected(
self, event: DirectoryTree.FileSelected
) -> None:
path = event.path
logger.debug("Directory file selected: %s", path)
try:
open_directory(str(path))
except Exception as exc:
logger.error("Failed to open %s: %s", path, exc)
self.app.bell()
def on_button_pressed(self, event: Button.Pressed) -> None:
button_id = event.button.id
logger.debug("Button pressed: %s", button_id)
match button_id:
case "move_agent_workflow_button":
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
event.stop()
case "otp_generate_button":
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
event.stop()
case "find_quiet_button":
self.app.push_screen(
QuietAgentWorkflowScreen(self.app.api, self.app.policies)
)
event.stop()
return
case "otp_activities_button":
self.app.push_screen(OTPActivitiesScreen())
event.stop()
return
case "otp_revoke_button":
self.app.push_screen(OTPRevokeScreen())
event.stop()
return
case "policy_prep_button":
# Use the new TUI workflow screen instead of legacy
self.app.push_screen(
PolicyPrepWorkflowScreen(self.app.api, self.app.policies)
)
event.stop()
return
case _:
self.app.bell()
logger.warning("Unknown button pressed: %s", button_id)
return
# ---------------------------------------------------------------------------
# 2) APP
# ---------------------------------------------------------------------------
class Loxide(App[Message]):
api: AirlockAPIWrapper
working_dir: str
policies: Optional[list[Policy]]
devices: Optional[list[Agent]]
CSS = """
#logo {
width: 100%;
content-align: center middle;
text-align: center;
}
"""
BINDINGS = [
("q", "quit", "Quit"),
("f", "open_fe", "Launch Explorer"),
("r", "refresh", "Refresh"),
]
def __init__(self, api: AirlockAPIWrapper):
self._textual_theme = get_user_value("TEXTUAL_THEME", str, "textual-dark")
super().__init__()
self.api = api
wd = load_env("WORKING_DIR") or os.getcwd()
if not os.path.isdir(wd):
wd = os.getcwd()
self.working_dir = wd
# Initial data load
self.refresh_data()
def refresh_data(self) -> None:
"""Public method to refresh policies and devices from the API."""
try:
self.policies = [
Policy(**row.to_dict())
for _, row in self.api.policy_find_all().iterrows()
]
self.devices = [
Agent(**row.to_dict())
for _, row in self.api.agent_find_all().iterrows()
]
if self.policies and self.devices:
for agent in self.devices:
agent.enrich_with_policies(self.policies)
logger.debug(
f"Enriched {len(self.devices)} agents with policy information"
)
except Exception as exc:
logger.error("Failed to load policies/devices: %s", exc)
self.policies = None
self.devices = None
def on_mount(self, api: AirlockAPIWrapper) -> None:
self.register_theme(get_retro_terminal_theme())
self.register_theme(get_amber_terminal_theme())
self.theme = self._textual_theme
self.push_screen(MainMenuScreen())
def action_refresh(self) -> None:
self.refresh_data()
def action_quit(self) -> None:
global _APP_RESTART_REASON
_APP_RESTART_REASON = None
self.exit()
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
# ---------------------------------------------------------------------------
# 3) PUBLIC ENTRYPOINT - Updated to accept attach_notification_handler
# ---------------------------------------------------------------------------
def run_Loxide(api: AirlockAPIWrapper, attach_notification_handler=None) -> None:
global _APP_RESTART_REASON
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 app loop iteration (attempt %d)", attempts)
_APP_RESTART_REASON = None
app = Loxide(api)
# Attach the notification handler if provided
if attach_notification_handler:
attach_notification_handler(app)
try:
app.run()
except SystemExit as exc:
if exc.code != 0:
logger.debug("Caught SystemExit from Textual: %s", exc)
raise
reason = _APP_RESTART_REASON
logger.debug("After app.run(), _APP_RESTART_REASON = %r", reason)
if not reason:
logger.debug("No restart reason, exiting loop")
break
if reason[0] == "restart":
logger.debug("Restarting app loop")
continue
if reason[0] == "multi_agent_action":
logger.info("Multi-agent action with selected agents: %s", reason[1])
continue
logger.error("Unknown restart reason: %r", reason)
break
# ---------------------------------------------------------------------------
# 4) MAIN FUNCTION - Updated to get and pass attach_notification_handler
# ---------------------------------------------------------------------------
def main():
irtang()
# Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
setup()
base_dir = get_base_directory()
# setup() now returns a function to attach the notification handler
attach_notification_handler = setup()
logger = logging.getLogger(__name__)
try:
@@ -67,7 +599,7 @@ def main():
base_url=str(url),
api_key=api_key,
)
run_Loxide(api)
run_Loxide(api, attach_notification_handler)
if __name__ == "__main__":
@@ -1,3 +1,17 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from __future__ import annotations
import logging
@@ -150,14 +164,11 @@ class AllowlistSelectionWidget(Static):
# 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.styles.width = "100%"
self.add_btn.disabled = True # Disabled until allowlist selected
yield self.back_btn
yield self.add_btn
async def on_mount(self) -> None:
@@ -380,9 +391,9 @@ class AllowlistSelectionWidget(Static):
if found_col:
self.hash_column = found_col
preview_lines.append(f"✓ Found hash column: **{found_col}**\n")
preview_lines.append(f" Found hash column: **{found_col}**\n")
else:
preview_lines.append("⚠️ **No hash column found**\n")
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":
@@ -464,7 +475,7 @@ class AllowlistSelectionWidget(Static):
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')}'"
f" Add to '{self.selected_allowlist.get('name', 'Unknown')}'"
)
# Update preview with selection
@@ -523,11 +534,6 @@ class AllowlistSelectionWidget(Static):
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()
@@ -553,7 +559,7 @@ class AllowlistSelectionWidget(Static):
try:
# Disable button during operation
self.add_btn.disabled = True
self.add_btn.label = "Adding hashes..."
self.add_btn.label = "Adding hashes..."
# Call API to add hashes
app_id = self.selected_allowlist.get("applicationid")
@@ -564,10 +570,10 @@ class AllowlistSelectionWidget(Static):
)
result = self.api.hash_add_to_allowlist(app_id, self.hashes_to_add)
logger.debug(f"Hash adding api call: {result}")
# Success notification
self.app.notify(
f"Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'",
f"Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'",
title="Success",
severity="information",
timeout=5,
@@ -575,7 +581,7 @@ class AllowlistSelectionWidget(Static):
# Update preview to show success
self.preview_area.text = (
f"## SUCCESS\n\n"
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"
@@ -586,13 +592,13 @@ class AllowlistSelectionWidget(Static):
)
# Change button to "Done"
self.add_btn.label = "Done"
self.add_btn.label = "Done - Press q to return to main menu"
self.add_btn.disabled = True
except Exception as exc:
logger.exception(f"Failed to add hashes to allowlist: {exc}")
self.app.notify(
f"Failed to add hashes: {str(exc)}",
f"Failed to add hashes: {str(exc)}",
title="Error",
severity="error",
timeout=10,
@@ -600,7 +606,7 @@ class AllowlistSelectionWidget(Static):
# Re-enable button
self.add_btn.disabled = False
self.add_btn.label = "Retry Add to Allowlist"
self.add_btn.label = "Retry Add to Allowlist"
class AllowlistSelectionScreen(Screen):
@@ -609,9 +615,9 @@ class AllowlistSelectionScreen(Screen):
"""
BINDINGS = [
Binding("b", "back", "Back"),
Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
Binding("r", "refresh", "Refresh Allowlists"),
Binding("enter", "confirm", "Add to Allowlist"),
]
def __init__(
@@ -641,10 +647,15 @@ class AllowlistSelectionScreen(Screen):
yield self.widget
yield Footer()
async def action_back(self) -> None:
async def action_go_back(self) -> None:
"""Go back to previous screen."""
await self.app.pop_screen()
async def action_main_menu(self) -> None:
"""Go back to main menu."""
while len(self.app.screen_stack) > 2:
await self.app.pop_screen()
async def action_refresh(self) -> None:
"""Refresh the allowlists."""
if hasattr(self, "widget") and self.widget:
+639
View File
@@ -0,0 +1,639 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from datetime import datetime, timedelta
import logging
import os
from typing import List
import pandas as pd
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical
from textual.screen import Screen
from textual.widgets import (
Button,
DataTable,
Footer,
Header,
Label,
Select,
Static,
)
from models.agent import Agent
from models.execution import ExecutionHistoryRecord
from utils.configmanager import load_env
logger = logging.getLogger(__name__)
class ExecutionHistoryScreen(Screen):
"""
A screen for viewing and exporting execution history for selected agents.
This screen allows users to:
1. Select a start date and end date using dropdown selects
2. Fetch execution history for all selected agents
3. View the results in a DataTable
4. Export the results to CSV using a keybinding
Attributes:
agents (List[Agent]): List of agents to fetch execution history for
execution_data (pd.DataFrame): Combined execution history data
working_dir (str): Directory for CSV exports
"""
DEFAULT_CSS = """
ExecutionHistoryScreen {
align: center top;
}
#main_container {
width: 95%;
height: 1fr;
border: solid $primary;
padding: 1;
}
#title {
text-style: bold;
color: $text;
text-align: center;
margin-bottom: 1;
}
#date_container {
height: auto;
margin-bottom: 1;
}
#start_date_row, #end_date_row {
height: auto;
align-horizontal: left;
margin-bottom: 1;
}
.date_label {
width: 8;
margin-right: 1;
}
.date_selector {
width: 18;
margin: 0 1;
}
#quick_buttons_row {
height: auto;
align-horizontal: center;
margin-bottom: 1;
}
.quick_select_btn {
margin: 0 1;
}
#button_row {
height: auto;
align-horizontal: center;
margin-top: 1;
margin-bottom: 1;
}
Button {
margin: 0 1;
}
#status_label {
text-align: center;
color: $accent;
margin-bottom: 1;
}
#results_container {
height: 1fr;
display: none;
}
#results_button_row {
height: auto;
align-horizontal: center;
margin-bottom: 1;
}
#history_table {
height: 1fr;
border: solid $primary;
}
DataTable > .datatable--header {
text-style: bold;
background: $primary 20%;
}
"""
BINDINGS = [
Binding("escape", "close_screen", "Close"),
Binding("e", "export_csv", "Export CSV"),
Binding("q", "close_screen", "Quit"),
]
def __init__(self, agents: List[Agent]):
"""
Initialize the ExecutionHistoryScreen.
Args:
agents (List[Agent]): List of agents to fetch execution history for
"""
super().__init__()
self.agents = agents
self.execution_data = pd.DataFrame()
self.working_dir = load_env("WORKING_DIR") or os.getcwd()
# Generate dropdown options
today = datetime.now().date()
# Month options - format is (display_text, value)
self.month_options = [
("January", "01"),
("February", "02"),
("March", "03"),
("April", "04"),
("May", "05"),
("June", "06"),
("July", "07"),
("August", "08"),
("September", "09"),
("October", "10"),
("November", "11"),
("December", "12"),
]
# Day options (1-31) - format is (display_text, value)
self.day_options = [(f"{i}", f"{i:02d}") for i in range(1, 32)]
# Year options (current year back 5 years) - format is (display_text, value)
current_year = today.year
self.year_options = [
(str(year), str(year)) for year in range(current_year, current_year - 6, -1)
]
# Default dates: last 30 days
start_date = today - timedelta(days=30)
self.start_month = f"{start_date.month:02d}"
self.start_day = f"{start_date.day:02d}"
self.start_year = str(start_date.year)
self.end_month = f"{today.month:02d}"
self.end_day = f"{today.day:02d}"
self.end_year = str(today.year)
def compose(self) -> ComposeResult:
"""Build the UI layout."""
yield Header(show_clock=True, icon="📊")
with Vertical(id="main_container"):
title_text = f"Execution History - {len(self.agents)} Agent(s)"
yield Static(title_text, id="title")
# Date selection area
with Vertical(id="date_container"):
yield Label("Select Date Range:")
# Start date row
with Horizontal(id="start_date_row"):
yield Label("From:", classes="date_label")
yield Select(
options=self.month_options,
value=self.start_month,
id="start_month_select",
classes="date_selector",
)
yield Select(
options=self.day_options,
value=self.start_day,
id="start_day_select",
classes="date_selector",
)
yield Select(
options=self.year_options,
value=self.start_year,
id="start_year_select",
classes="date_selector",
)
# End date row
with Horizontal(id="end_date_row"):
yield Label("To:", classes="date_label")
yield Select(
options=self.month_options,
value=self.end_month,
id="end_month_select",
classes="date_selector",
)
yield Select(
options=self.day_options,
value=self.end_day,
id="end_day_select",
classes="date_selector",
)
yield Select(
options=self.year_options,
value=self.end_year,
id="end_year_select",
classes="date_selector",
)
# Quick select buttons
with Horizontal(id="quick_buttons_row"):
yield Button(
"1 Day",
id="quick_1day",
classes="quick_select_btn",
variant="default",
)
yield Button(
"1 Week",
id="quick_1week",
classes="quick_select_btn",
variant="default",
)
yield Button(
"30 Days",
id="quick_30days",
classes="quick_select_btn",
variant="default",
)
# Buttons
with Horizontal(id="button_row"):
yield Button("Fetch History", id="fetch_btn", variant="primary")
yield Button("Close", id="close_btn", variant="error")
# Status
yield Static(
"Select date range and click 'Fetch History'", id="status_label"
)
# Results container (hidden initially, shown after fetch)
with Vertical(id="results_container"):
with Horizontal(id="results_button_row"):
yield Button("Export CSV", id="export_btn", variant="success")
yield Button("Back", id="back_btn", variant="default")
yield DataTable(id="history_table")
yield Footer()
def on_mount(self) -> None:
"""Initialize the table when screen is mounted."""
table = self.query_one("#history_table", DataTable)
table.cursor_type = "row"
table.zebra_stripes = True
# Initially empty - will populate after fetch
logger.info(f"ExecutionHistoryScreen mounted with {len(self.agents)} agents")
def on_select_changed(self, event: Select.Changed) -> None:
"""Handle date selection changes."""
select_id = event.select.id
if select_id == "start_month_select":
self.start_month = event.value
logger.debug(f"Start month changed to: {self.start_month}")
elif select_id == "start_day_select":
self.start_day = event.value
logger.debug(f"Start day changed to: {self.start_day}")
elif select_id == "start_year_select":
self.start_year = event.value
logger.debug(f"Start year changed to: {self.start_year}")
elif select_id == "end_month_select":
self.end_month = event.value
logger.debug(f"End month changed to: {self.end_month}")
elif select_id == "end_day_select":
self.end_day = event.value
logger.debug(f"End day changed to: {self.end_day}")
elif select_id == "end_year_select":
self.end_year = event.value
logger.debug(f"End year changed to: {self.end_year}")
def _set_quick_date_range(self, days: int) -> None:
"""Set the date range based on quick select button."""
today = datetime.now().date()
start_date = today - timedelta(days=days)
# Update internal values
self.start_month = f"{start_date.month:02d}"
self.start_day = f"{start_date.day:02d}"
self.start_year = str(start_date.year)
self.end_month = f"{today.month:02d}"
self.end_day = f"{today.day:02d}"
self.end_year = str(today.year)
# Update the Select widgets
try:
self.query_one("#start_month_select", Select).value = self.start_month
self.query_one("#start_day_select", Select).value = self.start_day
self.query_one("#start_year_select", Select).value = self.start_year
self.query_one("#end_month_select", Select).value = self.end_month
self.query_one("#end_day_select", Select).value = self.end_day
self.query_one("#end_year_select", Select).value = self.end_year
self.app.notify(
f"Date range set to last {days} day(s)",
severity="information",
timeout=2,
)
logger.info(f"Quick select: Set date range to last {days} days")
except Exception as e:
logger.error(f"Failed to update date selects: {e}")
def _show_date_selection(self) -> None:
"""Show the date selection view and hide results."""
try:
self.query_one("#date_container").styles.display = "block"
self.query_one("#button_row").styles.display = "block"
self.query_one("#status_label").styles.display = "block"
self.query_one("#results_container").styles.display = "none"
except Exception as e:
logger.error(f"Failed to show date selection: {e}")
def _show_results(self) -> None:
"""Hide date selection view and show results."""
try:
self.query_one("#date_container").styles.display = "none"
self.query_one("#button_row").styles.display = "none"
self.query_one("#status_label").styles.display = "none"
self.query_one("#results_container").styles.display = "block"
except Exception as e:
logger.error(f"Failed to show results: {e}")
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button clicks."""
if event.button.id == "fetch_btn":
self._fetch_execution_history()
elif event.button.id == "export_btn":
self._export_to_csv()
elif event.button.id == "close_btn":
self.app.pop_screen()
elif event.button.id == "back_btn":
self._show_date_selection()
elif event.button.id == "quick_1day":
self._set_quick_date_range(days=1)
elif event.button.id == "quick_1week":
self._set_quick_date_range(days=7)
elif event.button.id == "quick_30days":
self._set_quick_date_range(days=30)
def _fetch_execution_history(self) -> None:
"""Fetch execution history for all selected agents."""
status_label = self.query_one("#status_label", Static)
status_label.update("⏳ Fetching execution history...")
# Disable buttons during fetch
fetch_btn = self.query_one("#fetch_btn", Button)
export_btn = self.query_one("#export_btn", Button)
fetch_btn.disabled = True
export_btn.disabled = True
api = self.app.api
all_history = []
try:
# Construct dates from dropdowns
start_date_str = f"{self.start_year}-{self.start_month}-{self.start_day}"
end_date_str = f"{self.end_year}-{self.end_month}-{self.end_day}"
# Validate dates
try:
start_dt = datetime.strptime(start_date_str, "%Y-%m-%d")
end_dt = datetime.strptime(end_date_str, "%Y-%m-%d")
except ValueError as e:
status_label.update(f"❌ Invalid date: {str(e)}")
fetch_btn.disabled = False
export_btn.disabled = False
self.app.notify(f"Invalid date selected: {str(e)}", severity="error")
return
if start_dt > end_dt:
status_label.update("❌ Error: Start date must be before end date")
fetch_btn.disabled = False
export_btn.disabled = False
return
# Fetch history for each agent
for i, agent in enumerate(self.agents):
try:
status_label.update(
f"⏳ Fetching history for {agent.hostname} ({i+1}/{len(self.agents)})..."
)
# Call API - note the API expects 'dateto' first, then 'datefrom'
history = api.history_execution(
today=end_date_str,
date_selected=start_date_str,
agent_name=agent.hostname,
)
if history:
# Add agent hostname to each record for identification
for record in history:
record["agent_hostname"] = agent.hostname
all_history.extend(history)
logger.info(
f"Fetched {len(history)} records for {agent.hostname}"
)
else:
logger.info(f"No history found for {agent.hostname}")
except Exception as e:
logger.error(f"Failed to fetch history for {agent.hostname}: {e}")
self.app.notify(
f"Warning: Failed to fetch history for {agent.hostname}",
severity="warning",
)
# Convert to DataFrame
if all_history:
status_label.update(
"⏳ Enriching execution data with hash information..."
)
# Normalize field names (handle API typos)
for record in all_history:
if "policver" in record and "policyver" not in record:
record["policyver"] = record.pop("policver")
# Convert dict records to ExecutionHistoryRecord objects
execution_records = []
for record in all_history:
try:
execution_records.append(ExecutionHistoryRecord(**record))
except TypeError as e:
logger.warning(f"Failed to create ExecutionHistoryRecord: {e}")
# If it fails, just keep the dict
continue
# Enrich with hash data if we have ExecutionHistoryRecord objects
if execution_records:
try:
enriched_records = ExecutionHistoryRecord.enrich_with_hashes(
api, execution_records
)
logger.info(
f"Enriched {len(enriched_records)} records with hash data"
)
# Convert back to DataFrame
self.execution_data = pd.DataFrame(
[r.__dict__ for r in enriched_records]
)
# Flatten hash_obj if present
if (
not self.execution_data.empty
and "hash_obj" in self.execution_data.columns
):
hash_df = self.execution_data["hash_obj"].apply(
lambda h: (
h.to_dict() if h and hasattr(h, "to_dict") else {}
)
)
self.execution_data = pd.concat(
[
self.execution_data.drop(columns=["hash_obj"]),
hash_df,
],
axis=1,
)
except Exception as e:
logger.warning(f"Failed to enrich with hashes: {e}")
# Fall back to plain DataFrame
self.execution_data = pd.DataFrame(all_history)
else:
# If we couldn't create any ExecutionHistoryRecord objects, just use raw data
self.execution_data = pd.DataFrame(all_history)
self._populate_table()
self._show_results() # Switch to results view
self.app.notify(
f"Successfully loaded {len(self.execution_data)} records",
severity="information",
)
else:
status_label.update(
"ℹ️ No execution history found for selected agents/dates"
)
self.app.notify("No execution history found", severity="information")
self.execution_data = pd.DataFrame()
except Exception as e:
logger.error(f"Error fetching execution history: {e}")
status_label.update(f"❌ Error: {str(e)}")
self.app.notify(f"Failed to fetch history: {str(e)}", severity="error")
finally:
# Re-enable buttons
fetch_btn.disabled = False
export_btn.disabled = False
def _populate_table(self) -> None:
"""Populate the DataTable with execution history data."""
table = self.query_one("#history_table", DataTable)
table.clear(columns=True)
if self.execution_data.empty:
return
# Define preferred column order (your specified order)
preferred_order = [
"policyname",
"policyver",
"hostname",
"username",
"publisher",
"filename",
"pprocess",
"gprocess",
"sha256",
"commandline",
"agent_hostname", # Our custom field
]
# Get available columns in preferred order, then add any remaining columns
available_cols = []
for col in preferred_order:
if col in self.execution_data.columns:
available_cols.append(col)
# Add any remaining columns not in preferred order
for col in self.execution_data.columns:
if col not in available_cols:
available_cols.append(col)
# Add columns to table
for col in available_cols:
table.add_column(col, key=col)
# Add rows
for idx, row in self.execution_data.iterrows():
row_data = []
for col in available_cols:
value = row[col]
# Convert to string, handle None/NaN
if pd.isna(value):
row_data.append("")
else:
row_data.append(str(value))
table.add_row(*row_data, key=str(idx))
logger.info(f"Populated table with {len(self.execution_data)} rows")
def _export_to_csv(self) -> None:
"""Export the current execution data to CSV."""
if self.execution_data.empty:
self.app.notify("No data to export", severity="warning")
return
try:
# Create filename with timestamp
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"execution_history_{timestamp}.csv"
filepath = os.path.join(self.working_dir, filename)
# Export to CSV
self.execution_data.to_csv(filepath, index=False, encoding="utf-8-sig")
self.app.notify(
f"✅ Exported {len(self.execution_data)} records to: {filepath}",
severity="information",
timeout=5,
)
logger.info(f"Exported execution history to: {filepath}")
except Exception as e:
logger.error(f"Failed to export CSV: {e}")
self.app.notify(f"Failed to export CSV: {str(e)}", severity="error")
def action_export_csv(self) -> None:
"""Keybinding action to export CSV."""
self._export_to_csv()
def action_close_screen(self) -> None:
"""Close this screen and return to previous."""
self.app.pop_screen()
+114
View File
@@ -0,0 +1,114 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from typing import List, Optional
from textual.app import ComposeResult
from textual.binding import Binding
from textual.css.query import NoMatches
from textual.screen import Screen
from models.agent import Agent
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
from TUI.Widgets.multiagentselector import MultiAgentSelector
from TUI.Widgets.resultsdisplay import ResultsDisplay
class MoveAgentWorkflowScreen(Screen):
"""Screen that handles the agent movement workflow."""
BINDINGS = [
Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
]
def __init__(self, all_agents: Optional[List[Agent]]):
super().__init__()
self.all_agents = all_agents
self.selected_agents = None
self.workflow_stage = "select_agents" # Track current stage
def compose(self) -> ComposeResult:
"""Start with the multi-agent selector."""
yield MultiAgentSelector(self.all_agents)
def action_go_back(self) -> None:
"""Handle escape key to go back one step within the workflow."""
if self.workflow_stage == "select_agents":
# At first stage, go back to main menu
self.app.pop_screen()
elif self.workflow_stage == "operations":
# Go back to agent selection
try:
ops_widget = self.query_one(AgentMoveOperations)
ops_widget.remove()
except NoMatches:
pass
self.mount(MultiAgentSelector(self.all_agents))
self.workflow_stage = "select_agents"
elif self.workflow_stage == "results":
# Go back to operations
try:
results_widget = self.query_one(ResultsDisplay)
results_widget.remove()
except NoMatches:
pass
self.mount(AgentMoveOperations(self.selected_agents))
self.workflow_stage = "operations"
def action_main_menu(self) -> None:
"""Handle q key to go back to main menu."""
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
"""Handle selected agents - switch to operations screen."""
self.selected_agents = message.selected_agents
# Remove the MultiAgentSelector
selector = self.query_one(MultiAgentSelector)
selector.remove()
# Mount the AgentMoveOperations with the selected Agent objects
self.mount(AgentMoveOperations(self.selected_agents))
self.workflow_stage = "operations"
def on_agent_move_operations_operation_complete(
self, message: AgentMoveOperations.OperationComplete
) -> None:
"""Handle completion of move operation - transition to results screen."""
# Format successful results
success_lines = []
for agent, result in message.successful:
success_lines.append(f"{agent.hostname}")
# Format unsuccessful results
failure_lines = []
for agent, error in message.unsuccessful:
failure_lines.append(f"❌ — {agent.hostname}: {error}")
successful_text = "\n".join(success_lines) if success_lines else "(none)"
unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)"
# Remove the operations widget
ops_widget = self.query_one(AgentMoveOperations)
ops_widget.remove()
# Mount the results display
self.mount(
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
)
self.workflow_stage = "results"
@@ -1,3 +1,17 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from __future__ import annotations
from datetime import datetime
@@ -11,7 +25,7 @@ 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 TUI.Screens.allowlistselectionscreen import AllowlistSelectionScreen
from utils.configmanager import load_env
logger = logging.getLogger(__name__)
@@ -31,8 +45,7 @@ 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.
with Continue button.
"""
DEFAULT_CSS = """
@@ -85,15 +98,10 @@ class OTPActivitiesWidget(Static):
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)
# Button area at the bottom (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
self.continue_btn.styles.width = "100%"
yield self.continue_btn
async def on_mount(self) -> None:
@@ -122,7 +130,7 @@ class OTPActivitiesWidget(Static):
async def on_button_pressed(self, event) -> None: # type: ignore[override]
"""
Handle Back / Continue buttons for the Activity Preview area.
Handle Continue button for the Activity Preview area.
"""
# Try to resolve the button object from the event
btn = (
@@ -136,18 +144,12 @@ class OTPActivitiesWidget(Static):
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.")
self.app.notify(
"No activities loaded to continue with.", severity="warning"
)
return
# Copy activities DataFrame to pass to new screen
@@ -463,7 +465,7 @@ class OTPActivitiesWidget(Static):
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}"))
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."))
@@ -472,7 +474,7 @@ class OTPActivitiesWidget(Static):
class ActivityDetailWidget(Static):
"""
Interactive widget for Activity Detail screen.
Shows the provided DataFrame in a DataTable and offers Export + Back buttons.
Shows the provided DataFrame in a DataTable and offers Export button.
Now includes Select All/None and Add to Allowlist functionality.
"""
@@ -529,12 +531,10 @@ class ActivityDetailWidget(Static):
# 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"
"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)
@@ -547,12 +547,12 @@ class ActivityDetailWidget(Static):
# Update button labels with count
count = len(self.selected_row_ids)
total = len(self.activities_df)
len(self.activities_df)
if has_selection:
self.add_allowlist_btn.label = f"📋 Add {count} Selected to Allowlist"
self.add_allowlist_btn.label = f"Add {count} Selected to Allowlist"
else:
self.add_allowlist_btn.label = "📋 Add Selected to Allowlist"
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."""
@@ -591,7 +591,7 @@ class ActivityDetailWidget(Static):
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 ""
checkbox = "" if row_id in self.selected_row_ids else ""
# Add row to table
row_key = self.detail_table.add_row(checkbox, *vals)
@@ -616,7 +616,7 @@ class ActivityDetailWidget(Static):
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.detail_table.update_cell(row_key, "select", "") # Checked
self._update_button_states()
@@ -652,19 +652,13 @@ class ActivityDetailWidget(Static):
logger.exception("Failed to sort by column %s: %s", column_key, exc)
return
# ✅ Only refresh rows, not columns
# 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
@@ -676,7 +670,7 @@ class ActivityDetailWidget(Static):
# 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.detail_table.update_cell(row_key, "select", "")
self._update_button_states()
logger.info(f"Selected all {len(self.selected_row_ids)} rows")
@@ -688,7 +682,7 @@ class ActivityDetailWidget(Static):
# 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.detail_table.update_cell(row_key, "select", "☑️")
self._update_button_states()
logger.info("Cleared all selections")
@@ -730,14 +724,12 @@ class ActivityDetailWidget(Static):
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")
)
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")
Static("No rows selected for export.", classes="notification")
)
return
try:
@@ -750,7 +742,7 @@ class ActivityDetailWidget(Static):
logger.info("Exported selected activities to %s", file_path)
await self.mount(
Static(
f"Exported selected activities to: {filename}",
f"Exported selected activities to: {filename}",
classes="notification",
)
)
@@ -758,12 +750,12 @@ class ActivityDetailWidget(Static):
logger.exception("Failed to export detail activities: %s", exc)
await self.mount(
Static(
"❌ Failed to export activities; check logs.",
"¢ Failed to export activities; check logs.",
classes="notification",
)
)
# ✅ Helper methods
# Helper methods
def get_selected_data(self) -> pd.DataFrame:
"""Return a DataFrame of the selected rows."""
if not self.selected_row_ids:
@@ -795,7 +787,8 @@ class ActivityDetailScreen(Screen):
"""
BINDINGS = [
Binding("b", "back", "Back"),
Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
Binding("e", "export", "Export"),
Binding("a", "select_all", "Select All"),
Binding("n", "select_none", "Select None"),
@@ -819,11 +812,16 @@ class ActivityDetailScreen(Screen):
yield self.widget
yield Footer()
async def action_back(self) -> None:
async def action_go_back(self) -> None:
try:
await self.app.pop_screen()
except Exception:
logger.debug("ActivityDetailScreen.action_back pop_screen failed.")
logger.debug("ActivityDetailScreen.action_go_back pop_screen failed.")
async def action_main_menu(self) -> None:
"""Go back to main menu."""
while len(self.app.screen_stack) > 2:
await self.app.pop_screen()
async def action_export(self) -> None:
# Delegate to widget export helper
@@ -851,9 +849,10 @@ class OTPActivitiesScreen(Screen):
"""
BINDINGS = [
Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
Binding("r", "refresh_sessions", "Refresh Sessions"),
Binding("e", "export_activities", "Export activities"),
Binding("q", "quit", "Quit"),
]
def compose(self) -> ComposeResult:
@@ -875,6 +874,15 @@ class OTPActivitiesScreen(Screen):
else:
await self.widget.load_sessions_from_api(api)
async def action_go_back(self) -> None:
"""Go back one screen."""
await self.app.pop_screen()
async def action_main_menu(self) -> None:
"""Go back to main menu."""
while len(self.app.screen_stack) > 2:
await self.app.pop_screen()
# Simple actions bound to keys
async def action_refresh_sessions(self) -> None:
api = getattr(self.app, "api", None)
@@ -884,10 +892,6 @@ class OTPActivitiesScreen(Screen):
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)
+392
View File
@@ -0,0 +1,392 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from __future__ import annotations
import logging
from typing import List, Optional
import pandas as pd
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical
from textual.message import Message
from textual.screen import Screen
from textual.widgets import Button, DataTable, Footer, Header, Static
logger = logging.getLogger(__name__)
class OTPRevokeWidget(Static):
"""
Widget for managing OTP session revocation.
Displays active OTP sessions and allows selection for revocation.
"""
class SessionsRevoked(Message):
"""Message sent when sessions are revoked."""
def __init__(self, revoked_sessions: List[dict]):
super().__init__()
self.revoked_sessions = revoked_sessions
DEFAULT_CSS = """
OTPRevokeWidget {
height: 1fr;
}
#main_container {
width: 100%;
height: 100%;
layout: vertical;
}
#sessions_container {
height: 1fr;
border: none;
padding: 1;
}
#button_container {
height: auto;
width: 100%;
padding: 1;
align: center middle;
}
#button_container Button {
min-width: 16;
margin: 0 1;
}
#result_container {
height: auto;
max-height: 10;
border: solid #444444;
padding: 1;
margin: 1;
overflow-y: auto;
}
.panel-title {
text-style: bold;
margin: 0 0 1 0;
}
"""
def compose(self) -> ComposeResult:
with Vertical(id="main_container"):
# Sessions table
yield Static("OTP Sessions", classes="panel-title")
with Vertical(id="sessions_container"):
self.sessions_table = DataTable(id="sessions_table")
self.sessions_table.styles.width = "100%"
self.sessions_table.styles.height = "1fr"
yield self.sessions_table
# Action buttons
with Horizontal(id="button_container"):
yield Button("Refresh", id="refresh_btn")
yield Button("Select All", id="select_all_btn")
yield Button("Clear Selection", id="select_none_btn")
yield Button("Revoke Selected", id="revoke_btn", variant="error")
# Results display
with Vertical(id="result_container"):
yield Static("Revocation Results", classes="panel-title")
self.results_display = Static("No actions performed yet.")
yield self.results_display
async def on_mount(self) -> None:
"""Initialize the widget when mounted."""
# Configure sessions table
self.sessions_table.clear()
self.sessions_table.add_columns(
"", "OTP ID", "Hostname", "Status", "Purpose", "Granted"
)
# Enable row selection with checkbox column
self.sessions_table.cursor_type = "row"
try:
self.sessions_table.zebra_stripes = True
except Exception:
pass
# Initialize state
self._sessions_df: Optional[pd.DataFrame] = None
self._filtered_df: Optional[pd.DataFrame] = None
self._selected_otpids: set = set()
async def load_sessions_from_api(self, api) -> None:
"""Load active OTP sessions from the API."""
try:
# Fetch only active sessions
active_df = api.otp_find_active()
# Ensure we have a DataFrame
if not isinstance(active_df, pd.DataFrame):
active_df = pd.DataFrame(active_df)
# Add status column
active_df["status"] = "active"
# Sort by otpid if column exists
if "otpid" in active_df.columns and not active_df.empty:
active_df = active_df.sort_values(by="otpid", ascending=False)
# Store the full dataframe
self._sessions_df = active_df
self._filtered_df = active_df.copy()
# Display in table
await self._refresh_table()
# Update status
active_count = len(active_df)
status_msg = f"Loaded {active_count} active sessions"
logger.info(status_msg)
self.results_display.update(status_msg)
except Exception as e:
logger.exception(f"Failed to load OTP sessions: {e}")
self.results_display.update(f"Error loading sessions: {str(e)}")
async def _refresh_table(self) -> None:
"""Refresh the table display with current filtered data."""
if self._filtered_df is None or self._filtered_df.empty:
self.sessions_table.clear()
return
# Ensure expected columns exist
expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]
for col in expected_cols:
if col not in self._filtered_df.columns:
self._filtered_df[col] = ""
# Clear and repopulate table
self.sessions_table.clear(columns=False)
for _, row in self._filtered_df.iterrows():
otpid = str(row.get("otpid", ""))
# Check if this row is selected
checkbox = "☑️" if otpid in self._selected_otpids else ""
self.sessions_table.add_row(
checkbox,
str(otpid),
str(row.get("hostname", "")),
str(row.get("status", "")),
str(row.get("purpose", "")),
str(row.get("granted", "")),
)
async def on_button_pressed(self, event) -> None:
"""Handle button presses."""
btn = event.button
if btn.id == "refresh_btn":
# Refresh sessions
api = getattr(self.app, "api", None)
if api:
await self.load_sessions_from_api(api)
elif btn.id == "select_all_btn":
# Select all visible rows
if (
self._filtered_df is not None
and not self._filtered_df.empty
and "otpid" in self._filtered_df.columns
):
self._selected_otpids = set(str(x) for x in self._filtered_df["otpid"])
await self._refresh_table()
elif btn.id == "select_none_btn":
# Clear selection
self._selected_otpids.clear()
await self._refresh_table()
elif btn.id == "revoke_btn":
# Revoke selected sessions
await self._revoke_selected()
async def on_data_table_row_selected(self, event) -> None:
"""Handle row selection in the table."""
if event.data_table != self.sessions_table:
return
try:
# Get the row index from the cursor row
row_index = self.sessions_table.cursor_row
if (
self._filtered_df is not None
and not self._filtered_df.empty
and "otpid" in self._filtered_df.columns
and row_index < len(self._filtered_df)
):
# Get the OTP ID for this row
otpid = str(self._filtered_df.iloc[row_index]["otpid"])
# Toggle selection
if otpid in self._selected_otpids:
self._selected_otpids.remove(otpid)
else:
self._selected_otpids.add(otpid)
# Refresh table to update checkbox
await self._refresh_table()
# Restore cursor position
self.sessions_table.move_cursor(row=row_index)
except Exception as e:
logger.exception(f"Error handling row selection: {e}")
async def _revoke_selected(self) -> None:
"""Revoke the selected OTP sessions."""
if not self._selected_otpids:
self.results_display.update("No sessions selected for revocation")
return
api = getattr(self.app, "api", None)
if not api:
self.results_display.update("API not available")
return
# Collect results
results = []
success_count = 0
failure_count = 0
for otpid in self._selected_otpids:
try:
# Get hostname for this session
hostname = "Unknown"
if self._sessions_df is not None:
# Convert otpid to same type as in DataFrame for comparison
otpid_compare = otpid
if len(self._sessions_df) > 0:
first_otpid = self._sessions_df["otpid"].iloc[0]
if isinstance(first_otpid, int):
otpid_compare = int(otpid)
match = self._sessions_df[
self._sessions_df["otpid"] == otpid_compare
]
if not match.empty:
hostname = match.iloc[0].get("hostname", "Unknown")
# Revoke the session
result = api.otp_revoke(otpid)
if result and result.get("status") != "error":
success_count += 1
results.append(f"Revoked OTP {otpid} for {hostname}")
logger.info(f"Revoked OTP {otpid} for {hostname}: {result}")
else:
failure_count += 1
error_msg = (
result.get("message", "Unknown error")
if result
else "No response"
)
results.append(
f"Failed to revoke OTP {otpid} for {hostname}: {error_msg}"
)
logger.error(f"Failed to revoke OTP {otpid}: {error_msg}")
except Exception as e:
failure_count += 1
results.append(f"Error revoking OTP {otpid}: {str(e)}")
logger.exception(f"Exception revoking OTP {otpid}: {e}")
# Update results display
summary = (
f"Revocation complete: {success_count} succeeded, {failure_count} failed\n"
)
details = "\n".join(results[-5:]) # Show last 5 results
if len(results) > 5:
details = f"... (showing last 5 of {len(results)} results)\n" + details
self.results_display.update(summary + details)
# Clear selection and refresh
self._selected_otpids.clear()
await self.load_sessions_from_api(api)
# Post message about revoked sessions
if success_count > 0:
self.post_message(self.SessionsRevoked(results))
class OTPRevokeScreen(Screen):
"""
Main screen for OTP session revocation workflow.
This replaces the otp_revoke function from otp.py.
"""
BINDINGS = [
Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
Binding("r", "refresh", "Refresh"),
Binding("a", "select_all", "Select All"),
Binding("n", "select_none", "Clear Selection"),
Binding("d", "revoke", "Revoke Selected"),
]
def compose(self) -> ComposeResult:
yield Header(show_clock=True)
self.widget = OTPRevokeWidget()
yield self.widget
yield Footer()
async def on_mount(self) -> None:
"""Load sessions when screen mounts."""
api = getattr(self.app, "api", None)
if api:
await self.widget.load_sessions_from_api(api)
else:
logger.warning("OTPRevokeScreen mounted but no self.app.api found.")
async def action_refresh(self) -> None:
"""Refresh the sessions list."""
api = getattr(self.app, "api", None)
if api:
await self.widget.load_sessions_from_api(api)
async def action_select_all(self) -> None:
"""Select all visible sessions."""
if (
self.widget._filtered_df is not None
and not self.widget._filtered_df.empty
and "otpid" in self.widget._filtered_df.columns
):
self.widget._selected_otpids = set(
str(x) for x in self.widget._filtered_df["otpid"]
)
await self.widget._refresh_table()
async def action_select_none(self) -> None:
"""Clear all selections."""
self.widget._selected_otpids.clear()
await self.widget._refresh_table()
async def action_revoke(self) -> None:
"""Revoke selected sessions."""
await self.widget._revoke_selected()
async def action_go_back(self) -> None:
"""Go back to previous screen."""
await self.app.pop_screen()
async def action_main_menu(self) -> None:
"""Go back to main menu."""
while len(self.app.screen_stack) > 2:
await self.app.pop_screen()
+52
View File
@@ -0,0 +1,52 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from typing import List, Optional
from textual.app import ComposeResult
from textual.binding import Binding
from textual.screen import Screen
from models.agent import Agent
from TUI.Widgets.OTP_generate import OTPGenerator
class OTPWorkflowScreen(Screen):
"""Screen that handles the OTP generation workflow without agent selection."""
BINDINGS = [
Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
]
def __init__(self, selected_agents: Optional[List[Agent]]):
super().__init__()
self.selected_agents = selected_agents
def compose(self) -> ComposeResult:
"""Directly show the OTP generator for the selected agents."""
yield OTPGenerator(self.selected_agents)
def action_go_back(self) -> None:
"""Handle escape key to go back one screen."""
self.app.pop_screen()
def action_main_menu(self) -> None:
"""Handle q key to go back to main menu."""
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
"""Handle OTP generation request - pass it up to the app level if needed."""
File diff suppressed because it is too large Load Diff
@@ -23,9 +23,11 @@ the policy selection workflow.
import logging
from textual.app import ComposeResult
from textual.binding import Binding
from textual.screen import Screen
from textual.widgets import Footer, Header
from TUI.policyselector import PolicySelector
from TUI.Widgets.policyselector import PolicySelector
logger = logging.getLogger(__name__)
@@ -42,6 +44,11 @@ class PolicySelectorScreen(Screen):
agent_move_operations: Reference to the parent AgentMoveOperations widget.
"""
BINDINGS = [
Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
]
CSS = """
Screen {
layout: vertical;
@@ -68,7 +75,18 @@ class PolicySelectorScreen(Screen):
def compose(self) -> ComposeResult:
"""Create the PolicySelector widget."""
yield Header(show_clock=True)
yield PolicySelector(self.policies)
yield Footer()
def action_go_back(self) -> None:
"""Handle escape key to go back one screen."""
self.app.pop_screen()
def action_main_menu(self) -> None:
"""Handle q key to go back to main menu."""
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
def on_policy_selector_policy_selected(
self, message: PolicySelector.PolicySelected
@@ -34,12 +34,12 @@ 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 textual.widgets import Button, DataTable, Footer, Header, Input, Static
from models.policy import Policy
from services.API import AirlockAPIWrapper
from services.policyhandler import getPolicyInfo
from TUI.policyselector import PolicySelector
from TUI.Widgets.policyselector import PolicySelector
from utils.configmanager import load_env
logger = logging.getLogger(__name__)
@@ -51,16 +51,17 @@ class QuietAgentWorkflowScreen(Screen):
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
2. Configure analysis parameters (history period and quiet time period)
3. View categorized agents (enforce ready vs. non-enforce ready)
4. Select target policies for each category
5. Execute agent migrations
Attributes:
api (AirlockAPIWrapper): API wrapper for Airlock operations
policies (List[Policy]): List of all available policies
selected_policy (Optional[Policy]): The initially selected policy to analyze
history_days (int): Number of days of history to pull (default: 150)
quiet_days (int): Number of days without execution to be considered quiet (default: 45)
history_days (int): Number of days of history to pull (default: 150, range: 1-365)
quiet_days (int): Number of days without execution to be considered quiet (default: 45, range: 1-365)
agents_df (Optional[pd.DataFrame]): DataFrame of all agents with analysis results
enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents ready for enforcement
non_enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents not ready for enforcement
@@ -69,6 +70,7 @@ class QuietAgentWorkflowScreen(Screen):
BINDINGS = [
("escape", "go_back", "Back"),
("q", "main_menu", "Main Menu"),
]
workflow_stage = reactive("select_policy") # Tracks current workflow stage
@@ -85,7 +87,7 @@ class QuietAgentWorkflowScreen(Screen):
self.api = api
self.policies = policies
self.selected_policy: Optional[Policy] = None
self.history_days = 150 # Fixed as per requirements
self.history_days = 150 # Default value, user-selectable
self.quiet_days = 45 # Default value
self.agents_df: Optional[pd.DataFrame] = None
self.enforce_ready_df: Optional[pd.DataFrame] = None
@@ -96,10 +98,10 @@ class QuietAgentWorkflowScreen(Screen):
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="")
yield Header(show_clock=True, icon="")
# Title area
title = Static("🔒 Quiet Agent Workflow", id="workflow_title")
title = Static("Quiet Agent Workflow", id="workflow_title")
title.styles.margin = (0, 0, 0, 1)
yield title
@@ -129,14 +131,14 @@ class QuietAgentWorkflowScreen(Screen):
stage_messages = {
"select_policy": "Step 1: Select Policy to Analyze",
"select_quiet_days": "Step 2: Select Quiet Time Period",
"analyzing": "📊 Analyzing agent activity...",
"select_history_days": "Step 2: Configure Analysis Parameters",
"analyzing": "Analyzing agent activity...",
"view_results": "Step 3: Review Categorized Agents",
"select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents",
"select_non_enforce_target": "Step 5: Select Target Policy for Non-Enforce Ready Agents",
"confirm_migration": "Step 6: Confirm and Execute Migration",
"executing": "Executing agent migrations...",
"complete": "Migration Complete",
"executing": "Executing agent migrations...",
"complete": "Migration Complete",
}
status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage"))
@@ -160,7 +162,7 @@ class QuietAgentWorkflowScreen(Screen):
# 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()
self._show_history_days_selection()
elif self.workflow_stage == "select_enforce_target":
# Target policy selection for enforce ready agents
self.enforce_ready_target_policy = message.policy
@@ -176,64 +178,170 @@ class QuietAgentWorkflowScreen(Screen):
)
self._show_migration_confirmation()
def _show_quiet_days_selection(self) -> None:
"""Show the quiet days selection screen."""
self.workflow_stage = "select_quiet_days"
def _show_history_days_selection(self) -> None:
"""Show the history days and quiet days selection screen."""
self.workflow_stage = "select_history_days"
content = self.query_one("#content_area", Vertical)
content.remove_children()
# Create info text
info_widget = Static(
f"Policy Selected: {self.selected_policy.name}\n\n"
f"History Period: {self.history_days} days\n\n"
"Select quiet time period (days without untrusted execution):",
id="quiet_days_info",
"Configure Analysis Parameters:",
id="analysis_params_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)
# Create input container
input_container = Vertical(id="analysis_params_input_container")
input_container.styles.height = "auto"
content.mount(input_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",
# History days label
history_label = Static("History Period (days of execution history to pull):")
history_label.styles.margin = (0, 0, 1, 0)
input_container.mount(history_label)
# Add history days input field
history_input = Input(
placeholder="Enter days (1-365, default: 150)",
value="150",
id="history_days_input",
)
btn.styles.width = "100%"
btn.styles.margin = (0, 0, 1, 0)
button_container.mount(btn)
history_input.styles.width = "50"
history_input.styles.margin = (0, 0, 2, 0)
input_container.mount(history_input)
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)
# Quiet days label
quiet_label = Static(
"Quiet Time Period (days without execution to be considered quiet):"
)
quiet_label.styles.margin = (0, 0, 1, 0)
input_container.mount(quiet_label)
# Add quiet days input field
quiet_input = Input(
placeholder="Enter days (1-365, default: 45)",
value="45",
id="quiet_days_input",
)
quiet_input.styles.width = "50"
quiet_input.styles.margin = (0, 0, 2, 0)
input_container.mount(quiet_input)
# Add submit button
submit_btn = Button(
"Continue",
id="analysis_params_submit",
variant="primary",
)
submit_btn.styles.width = "50"
submit_btn.styles.margin = (1, 0, 0, 0)
input_container.mount(submit_btn)
# Focus the first input field
history_input.focus()
def _validate_and_submit_history_days(self) -> None:
"""Validate and submit the history days and quiet days inputs."""
try:
history_input = self.query_one("#history_days_input", Input)
quiet_input = self.query_one("#quiet_days_input", Input)
history_value = history_input.value.strip()
quiet_value = quiet_input.value.strip()
# Validate history days
if not history_value:
self.app.notify(
"Please enter a history period value", severity="error", timeout=3
)
history_input.focus()
return
try:
history_days = int(history_value)
except ValueError:
self.app.notify(
"Please enter a valid number for history period",
severity="error",
timeout=3,
)
history_input.focus()
return
if history_days < 1 or history_days > 365:
self.app.notify(
"History period must be between 1 and 365 days",
severity="error",
timeout=3,
)
history_input.focus()
return
# Validate quiet days
if not quiet_value:
self.app.notify(
"Please enter a quiet time period value",
severity="error",
timeout=3,
)
quiet_input.focus()
return
try:
quiet_days = int(quiet_value)
except ValueError:
self.app.notify(
"Please enter a valid number for quiet time period",
severity="error",
timeout=3,
)
quiet_input.focus()
return
if quiet_days < 1 or quiet_days > 365:
self.app.notify(
"Quiet time period must be between 1 and 365 days",
severity="error",
timeout=3,
)
quiet_input.focus()
return
# Check that quiet days doesn't exceed history days
if quiet_days > history_days:
self.app.notify(
"Quiet time period cannot exceed history period",
severity="error",
timeout=3,
)
quiet_input.focus()
return
# All validation passed
self.history_days = history_days
self.quiet_days = quiet_days
logger.info(
f"Selected history days: {history_days}, quiet days: {quiet_days}"
)
self._start_analysis()
except Exception as e:
logger.error(f"Error validating analysis parameters: {e}")
self.app.notify(f"Error: {str(e)}", severity="error", timeout=3)
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button press events."""
button_id = event.button.id
# 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()
# Analysis parameters submit button
if button_id == "analysis_params_submit":
self._validate_and_submit_history_days()
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
@@ -270,46 +378,44 @@ class QuietAgentWorkflowScreen(Screen):
self._show_policy_selection()
return
def on_input_submitted(self, event: Input.Submitted) -> None:
"""Handle input submission (Enter key pressed)."""
if event.input.id in ["history_days_input", "quiet_days_input"]:
self._validate_and_submit_history_days()
def _start_analysis(self) -> None:
"""Start the agent activity analysis."""
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
# Show notification that analysis is starting
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)
# Clear the screen to provide a blank canvas for Rust progress output
# (Rust output displays over the TUI, so we clear everything except header/footer)
try:
# Clear title
title_widget = self.query_one("#workflow_title", Static)
title_widget.update("")
def _perform_analysis(self) -> None:
# Clear status
status_widget = self.query_one("#workflow_status", Static)
status_widget.update("")
# Clear content area
content = self.query_one("#content_area", Vertical)
content.remove_children()
except Exception as e:
logger.debug(f"Could not clear screen for analysis: {e}")
# Delay the analysis start to ensure UI refresh completes first
# This prevents Rust output from starting before the screen is cleared
self.set_timer(0.5, self._perform_analysis_worker)
def _perform_analysis_worker(self) -> None:
"""Perform the actual agent activity analysis."""
try:
# 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)
@@ -322,32 +428,11 @@ class QuietAgentWorkflowScreen(Screen):
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."
@@ -393,9 +478,6 @@ class QuietAgentWorkflowScreen(Screen):
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]
@@ -405,8 +487,8 @@ class QuietAgentWorkflowScreen(Screen):
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()
self.enforce_ready_df = agents[agents["enforce_ready"]].copy()
self.non_enforce_ready_df = agents[~agents["enforce_ready"]].copy()
logger.info(
f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, "
@@ -428,27 +510,6 @@ class QuietAgentWorkflowScreen(Screen):
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"
@@ -469,9 +530,9 @@ class QuietAgentWorkflowScreen(Screen):
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"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",
@@ -500,11 +561,11 @@ class QuietAgentWorkflowScreen(Screen):
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 = 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 = Button("Start Over", id="start_over_btn")
start_over_btn.styles.margin = (0, 0, 1, 0)
button_container.mount(start_over_btn)
@@ -520,7 +581,7 @@ class QuietAgentWorkflowScreen(Screen):
enforce_col.styles.margin = (1, 1, 0, 0)
tables_container.mount(enforce_col)
enforce_label = Static("Enforce Ready Agents")
enforce_label = Static("Enforce Ready Agents")
enforce_label.styles.margin = (0, 0, 1, 0)
enforce_col.mount(enforce_label)
@@ -548,7 +609,7 @@ class QuietAgentWorkflowScreen(Screen):
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 = Static("Non-Enforce Ready Agents")
non_enforce_label.styles.margin = (0, 0, 1, 0)
non_enforce_col.mount(non_enforce_label)
@@ -589,7 +650,7 @@ class QuietAgentWorkflowScreen(Screen):
content.mount(policy_selector)
# Skip button
skip_btn = Button("⭕️ Skip - No Migration", id="skip_enforce_target_btn")
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)
@@ -614,7 +675,7 @@ class QuietAgentWorkflowScreen(Screen):
content.mount(policy_selector)
# Skip button
skip_btn = Button("⭕️ Skip - No Migration", id="skip_non_enforce_target_btn")
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)
@@ -627,29 +688,29 @@ class QuietAgentWorkflowScreen(Screen):
# Build confirmation message
confirmation_lines = [
"🔐 Migration Summary\n",
"Migration Summary\n",
f"Source Policy: {self.selected_policy.name}\n",
]
if self.enforce_ready_target_policy:
confirmation_lines.append(
f"\nEnforce Ready Migration:\n"
f"Agents: {len(self.enforce_ready_df)}\n"
f"Target: {self.enforce_ready_target_policy.name}\n"
f"\nEnforce Ready Migration:\n"
f"Agents: {len(self.enforce_ready_df)}\n"
f"Target: {self.enforce_ready_target_policy.name}\n"
)
if self.non_enforce_ready_target_policy:
confirmation_lines.append(
f"\nNon-Enforce Ready Migration:\n"
f"Agents: {len(self.non_enforce_ready_df)}\n"
f"Target: {self.non_enforce_ready_target_policy.name}\n"
f"\nNon-Enforce Ready Migration:\n"
f"Agents: {len(self.non_enforce_ready_df)}\n"
f"Target: {self.non_enforce_ready_target_policy.name}\n"
)
if (
not self.enforce_ready_target_policy
and not self.non_enforce_ready_target_policy
):
confirmation_lines.append("\n⚠️ No migrations will be performed.")
confirmation_lines.append("\nNo migrations will be performed.")
confirmation = Static("".join(confirmation_lines), id="migration_confirmation")
confirmation.styles.margin = (1, 1, 2, 1)
@@ -662,11 +723,11 @@ class QuietAgentWorkflowScreen(Screen):
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 = 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")
cancel_btn = Button("Cancel", id="cancel_migration_btn")
button_container.mount(cancel_btn)
def _execute_migration(self) -> None:
@@ -677,7 +738,7 @@ class QuietAgentWorkflowScreen(Screen):
# Show executing message
executing_msg = Static(
"Executing agent migrations...\nPlease wait...",
"Executing agent migrations...\nPlease wait...",
id="executing_message",
)
executing_msg.styles.margin = (2, 1)
@@ -696,7 +757,7 @@ class QuietAgentWorkflowScreen(Screen):
if self.enforce_ready_target_policy:
for idx, row in self.enforce_ready_df.iterrows():
try:
result = self.api.agent_move(
self.api.agent_move(
row["agentid"], self.enforce_ready_target_policy.groupid
)
successful_migrations.append(
@@ -713,7 +774,7 @@ class QuietAgentWorkflowScreen(Screen):
if self.non_enforce_ready_target_policy:
for idx, row in self.non_enforce_ready_df.iterrows():
try:
result = self.api.agent_move(
self.api.agent_move(
row["agentid"], self.non_enforce_ready_target_policy.groupid
)
successful_migrations.append(
@@ -749,7 +810,7 @@ class QuietAgentWorkflowScreen(Screen):
)
results = Static(
f"Migration Complete\n\n"
f"Migration Complete\n\n"
f"Total Agents Migrated: {len(successful)}\n"
f"Failed Migrations: {len(failed)}\n"
f"Success Rate: {success_rate:.1f}%",
@@ -764,7 +825,7 @@ class QuietAgentWorkflowScreen(Screen):
success_container.styles.margin = (0, 1)
content.mount(success_container)
success_label = Static("Successful Migrations")
success_label = Static("Successful Migrations")
success_label.styles.margin = (0, 0, 1, 0)
success_container.mount(success_label)
@@ -785,7 +846,7 @@ class QuietAgentWorkflowScreen(Screen):
failed_container.styles.margin = (2, 1, 0, 1)
content.mount(failed_container)
failed_label = Static("Failed Migrations")
failed_label = Static("Failed Migrations")
failed_label.styles.margin = (0, 0, 1, 0)
failed_container.mount(failed_label)
@@ -802,7 +863,7 @@ class QuietAgentWorkflowScreen(Screen):
failed_container.mount(failed_table)
# Action button
done_btn = Button("✔ Done", id="start_over_btn")
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)
@@ -833,7 +894,7 @@ class QuietAgentWorkflowScreen(Screen):
# 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":
elif self.workflow_stage == "select_history_days":
self._show_policy_selection()
elif self.workflow_stage == "select_enforce_target":
self._show_results()
@@ -846,3 +907,8 @@ class QuietAgentWorkflowScreen(Screen):
self._show_non_enforce_target_selection()
else:
self.app.pop_screen()
def action_main_menu(self) -> None:
"""Go back to main menu."""
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
-511
View File
@@ -1,511 +0,0 @@
import logging
import os
import sys
from typing import Optional
import dotenv
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 (
Button,
DirectoryTree,
Footer,
Header,
Static,
Tab,
Tabs,
)
from flows.otp import otp_revoke
from flows.prepPolicy import menu_policy_enforce
from models.agent import Agent
from models.policy import Policy
from services.API import AirlockAPIWrapper
from services.policyhandler import confirmUpdateAfromE
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
dotenv.load_dotenv()
# ---------------------------------------------------------------------------
# GLOBAL STASH
# ---------------------------------------------------------------------------
_PENDING_JOB = None
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# helper to persist TEXTUAL_THEME to *user* config and mirror to .env
# ---------------------------------------------------------------------------
def _persist_user_theme(theme_name: str) -> None:
"""
Store the chosen Textual theme in the user's config using the config manager.
No need to touch .env - config manager handles everything.
"""
base_dir = get_base_directory()
config_dir = base_dir / "config"
try:
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 = {
"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) -> None:
super().__init__()
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()
self.working_dir = wd
def _make_buttons_for(self, tab_id: str) -> Vertical:
defs = self.BUTTON_DEFS.get(tab_id, [])
buttons = []
for label, btn_id in defs:
btn = Button(label, id=btn_id)
btn.styles.width = "100%"
buttons.append(btn)
return Vertical(*buttons)
def compose(self) -> ComposeResult:
yield Header(show_clock=True, icon="")
tabs = [
Tab("Tree View", id="p_tree"),
Tab("Agents", id="agent_actions"),
Tab("Directory", id="dir"),
Tab("Settings", id="settings"),
]
if self.extras == "POLICYPREP":
tabs.insert(2, Tab("Policy Prep", id="policy"))
yield Tabs(*tabs, id="tabs")
yield Vertical(id="content")
yield Footer()
def on_mount(self) -> None:
api = self.app.api
self.switch_tab("agent_actions")
# focus helpers
def _get_content_buttons(self) -> list[Button]:
content = self.query_one("#content", Vertical)
return list(content.query(Button))
def _focus_first_button(self) -> None:
buttons = self._get_content_buttons()
if buttons:
buttons[0].focus()
def _focus_tabs(self) -> None:
tabs = self.query_one("#tabs", Tabs)
tabs.focus()
def _focus_nearby_button(self, direction: int) -> None:
buttons = self._get_content_buttons()
if not buttons:
return
try:
current = next(i for i, b in enumerate(buttons) if b.has_focus)
except StopIteration:
if direction > 0:
buttons[0].focus()
else:
buttons[-1].focus()
return
if direction < 0 and current == 0:
self._focus_tabs()
return
new_index = current + direction
if 0 <= new_index < len(buttons):
buttons[new_index].focus()
def switch_tab(self, tab_id: str) -> None:
self.current_tab = tab_id
content = self.query_one("#content", Vertical)
content.remove_children()
if tab_id in self.BUTTON_DEFS:
content.mount(self._make_buttons_for(tab_id))
self.call_later(self._focus_first_button)
elif tab_id == "dir":
content.mount(DirectoryTree(self.working_dir, id="dir_tree"))
elif tab_id == "p_tree":
content.mount(PolicyTreeWidget(self.app.policies, self.app.devices))
elif tab_id == "settings":
content.mount(ThemeSelector())
else:
content.mount(Static(f"Unknown tab: {tab_id}"))
def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None:
self.switch_tab(event.tab.id)
def on_multi_agent_selector_agents_selected(
self, message: MultiAgentSelector.AgentsSelected
) -> None:
"""Handle selected agents from AgentSelector."""
global _PENDING_JOB
selected_agents = message.selected_agents
logger.info("Selected agents: %s", selected_agents)
# TODO: Implement actual handling of selected agents
_PENDING_JOB = ("multi_agent_action", selected_agents)
self.app.exit()
def on_theme_selector_theme_selected(
self, message: ThemeSelector.ThemeSelected
) -> None:
"""Handle theme selection from ThemeSelector."""
global _PENDING_JOB
_persist_user_theme(message.theme_name)
_PENDING_JOB = ("restart",)
self.app.exit()
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
"""Handle OTP generation request from the workflow."""
global _PENDING_JOB
logger.info(
"OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d",
len(message.devices),
message.requestor,
message.reasoning,
message.duration,
)
_PENDING_JOB = (
"otp_workflow",
message.devices,
message.requestor,
message.reasoning,
message.duration,
)
self.app.exit()
def on_agent_move_operations_operation_complete(
self, message: AgentMoveOperations.OperationComplete
) -> None:
"""Handle completion of agent move operation - show results."""
logger.info(
"Agent move operation completed: %s, %d successful, %d unsuccessful",
message.operation,
len(message.successful),
len(message.unsuccessful),
)
# Format results for display
successful_text = "\n".join(
[f"{agent.hostname}" for agent, _ in message.successful]
)
unsuccessful_text = "\n".join(
[f"{agent.hostname}: {error}" for agent, error in message.unsuccessful]
)
# Remove the operations widget
try:
ops_widget = self.query_one(AgentMoveOperations)
ops_widget.remove()
except Exception:
pass
# Show results
self.query_one("#content", Vertical).mount(
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
)
def on_results_display_go_back(self, message: ResultsDisplay.GoBack) -> None:
"""Handle back button from results display."""
try:
results_widget = self.query_one(ResultsDisplay)
results_widget.remove()
except Exception:
pass
# Return to main menu
self.app.pop_screen()
def on_directory_tree_file_selected(
self, event: DirectoryTree.FileSelected
) -> None:
path = event.path
logger.debug("Directory file selected: %s", path)
try:
open_directory(str(path))
except Exception as exc:
logger.error("Failed to open %s: %s", path, exc)
self.app.bell()
def on_button_pressed(self, event: Button.Pressed) -> None:
global _PENDING_JOB
button_id = event.button.id
logger.debug("Button pressed: %s", button_id)
match button_id:
case "move_agent_workflow_button":
self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices))
event.stop()
case "otp_generate_button":
self.app.push_screen(OTPWorkflowScreen(self.app.devices))
event.stop()
case "find_quiet_button":
self.app.push_screen(
QuietAgentWorkflowScreen(self.app.api, self.app.policies)
)
event.stop()
return
case "otp_activities_button":
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[Message]):
api: AirlockAPIWrapper
working_dir: str
policies: Optional[list[Policy]]
devices: Optional[list[Agent]]
CSS = """
#logo {
width: 100%;
content-align: center middle;
text-align: center;
}
"""
BINDINGS = [
("q", "quit", "Quit"),
("f", "open_fe", "Launch Explorer"),
("r", "refresh", "Refresh"),
]
def __init__(self, api: AirlockAPIWrapper):
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()
def refresh_data(self) -> None:
"""Public method to refresh policies and devices from the API."""
try:
self.policies = [
Policy(**row.to_dict())
for _, row in self.api.policy_find_all().iterrows()
]
self.devices = [
Agent(**row.to_dict())
for _, row in self.api.agent_find_all().iterrows()
]
if self.policies and self.devices:
for agent in self.devices:
agent.enrich_with_policies(self.policies)
logger.debug(
f"Enriched {len(self.devices)} agents with policy information"
)
except Exception as exc:
logger.error("Failed to load policies/devices: %s", exc)
self.policies = None
self.devices = None
def on_mount(self, api: AirlockAPIWrapper) -> None:
self.register_theme(get_retro_terminal_theme())
self.register_theme(get_amber_terminal_theme())
self.theme = self._textual_theme
self.push_screen(MainMenuScreen())
def action_refresh(self) -> None:
self.refresh_data()
def action_quit(self) -> None:
global _PENDING_JOB
_PENDING_JOB = None
self.exit()
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
# ---------------------------------------------------------------------------
# 3) TERMINAL + LEGACY
# ---------------------------------------------------------------------------
def _restore_terminal_for_legacy() -> None:
sys.stdout.write("\033[?1049l")
sys.stdout.write("\033[?25h")
sys.stdout.write("\033[0m")
sys.stdout.write("\033[?1000l\033[?1002l\033[?1003l\033[?1006l")
sys.stdout.write("\033[2J\033[H")
sys.stdout.flush()
if os.name == "nt":
try:
import ctypes
kernel32 = ctypes.windll.kernel32
handle = kernel32.GetStdHandle(-11)
mode = ctypes.c_ulong()
if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
kernel32.SetConsoleMode(handle, mode.value | 0x0004)
except Exception as exc:
logger.debug("VT enable on Windows failed: %s", exc)
def _run_legacy_job(func, args, kwargs) -> None:
logger.debug("Running legacy job: %s", getattr(func, "__name__", func))
_restore_terminal_for_legacy()
try:
func(*args, **kwargs)
finally:
try:
input("\nPress Enter to return to the UI...")
except EOFError:
pass
# ---------------------------------------------------------------------------
# 4) PUBLIC ENTRYPOINT
# ---------------------------------------------------------------------------
def run_Loxide(api: AirlockAPIWrapper) -> None:
global _PENDING_JOB
base_dir = get_base_directory()
env_path = base_dir / ".env"
dotenv.load_dotenv(dotenv_path=env_path, override=True)
max_attempts = 5
attempts = 0
while attempts < max_attempts:
attempts += 1
logger.debug("Starting job loop iteration (attempt %d)", attempts)
_PENDING_JOB = None
app = Loxide(api)
try:
app.run()
except SystemExit as exc:
if exc.code != 0:
logger.debug("Caught SystemExit from Textual: %s", exc)
raise
job = _PENDING_JOB
logger.debug("After app.run(), _PENDING_JOB = %r", job)
if not job:
logger.debug("No job pending, exiting loop")
break
if job[0] == "legacy":
_, func, args, kwargs = job
_run_legacy_job(func, args, kwargs)
continue
if job[0] == "restart":
logger.debug("Restarting job loop")
continue
if job[0] == "multi_agent_action":
logger.info("Multi-agent action with selected agents: %s", job[1])
continue
if job[0] == "otp_workflow":
_, devices, requestor, reasoning, duration = job
def otp_generate_with_params():
# Your OTP logic here
pass
_run_legacy_job(otp_generate_with_params, (), {})
continue
logger.error("Unknown job type: %r", job)
break
# ---------------------------------------------------------------------------
# 5) DEV
# ---------------------------------------------------------------------------
if __name__ == "__main__":
api = AirlockAPIWrapper()
run_Loxide(api)
@@ -1,3 +1,18 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from textual.color import Color
from textual.theme import Theme
@@ -1,3 +1,18 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from textual.color import Color
@@ -1,3 +1,18 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from textual.containers import Vertical
from textual.message import Message
from textual.widget import Widget
@@ -1,3 +1,18 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import logging
from typing import List, Optional
@@ -23,6 +38,8 @@ logger = logging.getLogger(__name__)
class OTPGenerator(Widget):
"""Widget for generating OTPs for selected devices."""
# Reactive properties to track form completion
requestor_filled = reactive(False)
reasoning_filled = reactive(False)
@@ -139,14 +156,10 @@ class OTPGenerator(Widget):
button_row.styles.height = "auto"
button_row.styles.margin = (1, 0, 0, 0)
back_button = Button("← Back", id="back_button")
back_button.styles.width = "1fr"
yield back_button
generate_button = Button(
"Generate OTP", id="generate_button", variant="primary"
)
generate_button.styles.width = "2fr"
generate_button.styles.width = "100%"
yield generate_button
# Right side - Show device list initially, then output after generation
@@ -169,7 +182,7 @@ class OTPGenerator(Widget):
# Show device list initially
device_list_text = "\n".join(
f"{device.hostname}" for device in self.devices
f"{device.hostname}" for device in self.devices
)
device_display = Static(device_list_text, id="device_display")
yield device_display
@@ -197,14 +210,7 @@ class OTPGenerator(Widget):
def on_button_pressed(self, event: Button.Pressed):
btn_id = event.button.id
if btn_id == "back_button":
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "copy_clipboard_button":
if btn_id == "copy_clipboard_button":
try:
output_area = self.query_one("#otp_output", TextArea)
text_to_copy = output_area.text
@@ -213,7 +219,7 @@ class OTPGenerator(Widget):
pyperclip.copy(text_to_copy)
self.app.notify(
" Copied to clipboard!", severity="information", timeout=2
" Copied to clipboard!", severity="information", timeout=2
)
except ImportError:
self.app.notify(
@@ -1,3 +1,19 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from dataclasses import asdict
from datetime import datetime
import logging
@@ -10,12 +26,13 @@ from textual.css.query import NoMatches
from textual.message import Message
from textual.reactive import reactive
from textual.widget import Widget
from textual.widgets import Button, DataTable, Header, Static, TextArea
from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea
from models.agent import Agent
from TUI.OTP_generate import OTPGenerator
from TUI.otpworkflowscreen import OTPWorkflowScreen
from TUI.policyselectorscreen import PolicySelectorScreen
from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
from TUI.Screens.policyselectorscreen import PolicySelectorScreen
from TUI.Widgets.OTP_generate import OTPGenerator
logger = logging.getLogger(__name__)
@@ -124,6 +141,7 @@ class AgentMoveOperations(Widget):
toggle_enforcement_btn = self.query_one("#toggle_enforcement_btn", Button)
other_policy_btn = self.query_one("#other_policy_btn", Button)
otp_gen_btn = self.query_one("#otp_gen_btn", Button)
exec_history_btn = self.query_one("#exec_history_btn", Button)
# If operation in progress, disable all
if self.operation_in_progress:
@@ -132,6 +150,7 @@ class AgentMoveOperations(Widget):
local_approval_btn.disabled = True
toggle_enforcement_btn.disabled = True
other_policy_btn.disabled = True
exec_history_btn.disabled = True
else:
# If an operation was selected, disable
if self.selected_operation:
@@ -146,6 +165,9 @@ class AgentMoveOperations(Widget):
other_policy_btn.disabled = (
self.selected_operation == "other_policy"
)
exec_history_btn.disabled = (
self.selected_operation == "exec_history"
)
else:
# Enable all buttons
otp_gen_btn = False
@@ -153,6 +175,7 @@ class AgentMoveOperations(Widget):
local_approval_btn.disabled = False
toggle_enforcement_btn.disabled = False
other_policy_btn.disabled = False
exec_history_btn.disabled = False
except NoMatches:
pass
@@ -183,21 +206,21 @@ class AgentMoveOperations(Widget):
f"Operation: {operation_name}",
f"{'=' * 50}",
"",
f"✅ Successful ({len(successful)}):",
f" Successful ({len(successful)}):",
]
if successful:
for agent, result in successful:
results_lines.append(f" ✅ {agent.hostname}")
results_lines.append(f" {agent.hostname}")
else:
results_lines.append(" (none)")
results_lines.append("")
results_lines.append(f"❌ Failed ({len(unsuccessful)}):")
results_lines.append(f" Failed ({len(unsuccessful)}):")
if unsuccessful:
for agent, error in unsuccessful:
results_lines.append(f" ❌ {agent.hostname}: {error}")
results_lines.append(f" {agent.hostname}: {error}")
else:
results_lines.append(" (none)")
@@ -232,9 +255,9 @@ class AgentMoveOperations(Widget):
- Operations panel: 1/3 width
- Results area: Initially hidden, shown after operation completion
"""
yield Header(show_clock=True, icon="âš™")
yield Header(show_clock=True, icon="⚙️")
title_text = Static(
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
id="move_ops_title",
)
title_text.styles.margin = (0, 0, 1, 0)
@@ -269,46 +292,50 @@ class AgentMoveOperations(Widget):
yield operations_label
# Operation buttons
export_csv_btn = Button("📈 Export CSV", id="export_csv_btn")
export_csv_btn = Button("📄 Export CSV", id="export_csv_btn")
export_csv_btn.styles.width = "100%"
export_csv_btn.styles.margin = (0, 0, 1, 0)
yield export_csv_btn
local_approval_btn = Button(
"✔️ Local Approval Mode", id="local_approval_btn"
"✔️ 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 = 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 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"
"🔀 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
exec_history_btn = Button(
"📊 View Execution History", id="exec_history_btn"
)
exec_history_btn.styles.width = "100%"
exec_history_btn.styles.margin = (0, 0, 1, 0)
yield exec_history_btn
# Status label
status_label = Static("", id="status_label")
status_label.styles.margin = (2, 0, 0, 0)
yield status_label
back_button = Button("← Back", id="back_button")
back_button.styles.width = "50%"
back_button.styles.margin = (0, 1, 1, 0)
yield back_button
yield Footer()
def on_mount(self) -> None:
"""
@@ -343,7 +370,7 @@ class AgentMoveOperations(Widget):
Handle button press events from the widget.
This Textual event handler routes button presses to appropriate actions:
- back_button: Pop this screen (return to parent)
- copy_results_btn: Copy results text to clipboard (requires pyperclip)
- local_approval_btn: Start local approval operation
- toggle_enforcement_btn: Start toggle audit/enforcement operation
@@ -357,29 +384,24 @@ class AgentMoveOperations(Widget):
btn_id = event.button.id
if btn_id == "back_button":
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "copy_results_btn":
if btn_id == "copy_results_btn":
try:
results_text = self.query_one("#results_text", TextArea)
import pyperclip
pyperclip.copy(results_text.text)
self.app.notify(
"📋✅ Results copied to clipboard!",
"📋✅ Results copied to clipboard!",
severity="information",
timeout=2,
)
except ImportError:
self.app.notify(
"❌ pyperclip not installed. Run: pip install pyperclip",
" pyperclip not installed. Run: pip install pyperclip",
severity="warning",
)
except Exception as e:
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
self.app.notify(f" Failed to copy: {str(e)}", severity="error")
event.stop()
elif btn_id == "export_csv_btn":
self._start_export_csv_operation()
@@ -399,6 +421,9 @@ class AgentMoveOperations(Widget):
elif btn_id == "otp_gen_btn":
self._start_OTP_gen_operation()
event.stop()
elif btn_id == "exec_history_btn":
self._start_execution_history_operation()
event.stop()
def _start_local_approval_operation(self) -> None:
"""
@@ -427,7 +452,7 @@ class AgentMoveOperations(Widget):
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("✔️ Moving agents to local approval...")
status_label.update("✔️ Moving agents to local approval...")
# Get API from app
api = self.app.api
@@ -462,12 +487,12 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error during local approval operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f" Error: {str(e)}")
self.operation_in_progress = False
return
self.operation_in_progress = False
status_label.update("✅ Operation complete!")
status_label.update(" Operation complete!")
# Display results in the widget
self._display_results("Local Approval Mode", successful, unsuccessful)
@@ -483,7 +508,6 @@ class AgentMoveOperations(Widget):
self.selected_operation = "export_csv"
self.operation_in_progress = True
successful = []
unsuccessful = []
status_label = self.query_one("#status_label", Static)
status_label.update("Exporting CSV...")
self.app.refresh_data()
@@ -511,9 +535,9 @@ class AgentMoveOperations(Widget):
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}")
status_label.update(f" Exported to {file_path}")
except Exception:
status_label.update("❌ Failed")
status_label.update(" Failed")
self.operation_in_progress = False
@@ -557,7 +581,7 @@ class AgentMoveOperations(Widget):
self.operation_in_progress = True
status_label = self.query_one("#status_label", Static)
status_label.update("⏳ Toggling enforcement mode...")
status_label.update("🔄 Toggling enforcement mode...")
# Get API from app
api = self.app.api
@@ -593,12 +617,12 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error during toggle enforcement operation: {e}")
status_label.update(f"❌ Error: {str(e)}")
status_label.update(f" Error: {str(e)}")
self.operation_in_progress = False
return
self.operation_in_progress = False
status_label.update("✅ Operation complete!")
status_label.update(" Operation complete!")
# Display results in the widget
self._display_results("Toggle Audit/Enforcement", successful, unsuccessful)
@@ -663,7 +687,7 @@ class AgentMoveOperations(Widget):
except Exception as e:
logger.error(f"Error loading policies: {e}")
status_label.update(f"❌ Error: {str(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")
@@ -674,6 +698,35 @@ class AgentMoveOperations(Widget):
self.app.push_screen(OTPWorkflowScreen(self.agents))
def _start_execution_history_operation(self) -> None:
"""
Launch the execution history viewer for selected agents.
This operation opens a new screen that allows the user to:
1. Select a date range for execution history
2. Fetch execution logs for all selected agents
3. View the results in a table
4. Export the results to CSV
The screen is pushed onto the screen stack, allowing the user to return
to this screen when done.
"""
status_label = self.query_one("#status_label", Static)
status_label.update("Opening execution history viewer...")
try:
# Push the execution history screen
self.app.push_screen(ExecutionHistoryScreen(self.agents))
logger.info(
f"Opened execution history viewer for {len(self.agents)} agents"
)
except Exception as e:
logger.error(f"Failed to open execution history viewer: {e}")
status_label.update(f"❌ Error: {str(e)}")
self.app.notify(
f"Failed to open execution history: {str(e)}", severity="error"
)
def _execute_move_to_policy(self, target_policy) -> None:
"""
Execute the actual move of agents to the selected policy.
@@ -695,7 +748,7 @@ class AgentMoveOperations(Widget):
for agent in self.agents:
try:
# Move agent to target policy
result = api.agent_move(agent.agentid, target_policy.groupid)
api.agent_move(agent.agentid, target_policy.groupid)
successful.append((agent, f"Moved to {target_policy.name}"))
logger.info(
f"Successfully moved {agent.hostname} to policy {target_policy.name}"
@@ -1,4 +1,20 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import difflib
from pathlib import Path
import re
from typing import List, Optional
@@ -20,6 +36,8 @@ from models.agent import Agent
class MultiAgentSelector(Widget):
"""Widget for selecting multiple agents from a list."""
class AgentsSelected(Message):
def __init__(self, selected_agents: List[Agent]):
super().__init__()
@@ -40,7 +58,7 @@ class MultiAgentSelector(Widget):
def compose(self):
yield Header(show_clock=True, icon="")
title_text = Static("🖧 Agent Selector", id="selector_title")
title_text = Static("🖥️ Agent Selector", id="selector_title")
title_text.styles.margin = (0, 0, 0, 1)
yield title_text
@@ -60,7 +78,7 @@ class MultiAgentSelector(Widget):
text_area.styles.overflow_y = "auto"
yield text_area
with Horizontal(id="switch_search_container") as switch_search:
with Horizontal(id="switch_container"):
switch = Switch(value=False, id="match_switch")
switch.styles.width = "auto"
switch.styles.margin = (1, 0, 0, 0)
@@ -72,8 +90,13 @@ class MultiAgentSelector(Widget):
switch_label.styles.margin = (2, 1, 0, 0)
yield switch_label
with Horizontal(id="action_buttons_container"):
load_file = Button("📂 Load File", id="load_file_button")
load_file.styles.margin = (1, 1, 0, 1)
yield load_file
search = Button("🔍 Search", id="search_button")
search.styles.margin = (1, 0, 0, 0)
search.styles.margin = (1, 0, 0, 1)
yield search
with Horizontal() as select_buttons:
@@ -91,11 +114,6 @@ class MultiAgentSelector(Widget):
button_row.styles.height = "auto"
button_row.styles.margin = (1, 0, 0, 0)
back_button = Button("← Back", id="back_button")
back_button.styles.width = "1fr"
back_button.styles.margin = (0, 0, 0, 1)
yield back_button
submit_button = Button(
"▶ Select & Continue", id="submit_selection", variant="primary"
)
@@ -123,10 +141,7 @@ class MultiAgentSelector(Widget):
match_list = self.query_one("#match_results", SelectionList)
except NoMatches:
return
if btn_id == "back_button":
self.app.pop_screen()
event.stop()
elif btn_id == "select_all":
if btn_id == "select_all":
match_list.select_all()
event.stop()
elif btn_id == "select_none":
@@ -143,6 +158,9 @@ class MultiAgentSelector(Widget):
]
self.post_message(self.AgentsSelected(selected_agents))
event.stop()
elif btn_id == "load_file_button":
self._load_from_file()
event.stop()
elif btn_id == "search_button":
self.update_matches()
event.stop()
@@ -157,7 +175,7 @@ class MultiAgentSelector(Widget):
match_list.add_option((name, name))
unmatched_label = self.query_one("#unmatched_label", Static)
if unmatched:
unmatched_label.update(f"⚠️ No matches for: {', '.join(unmatched)}")
unmatched_label.update(f"⚠️ No matches for: {', '.join(unmatched)}")
else:
unmatched_label.update("")
@@ -205,3 +223,118 @@ class MultiAgentSelector(Widget):
else:
unmatched.append(name)
return sorted(matched), unmatched
def _load_from_file(self):
"""Safely load device names from a text file."""
try:
# Import here to avoid issues if tkinter isn't available
import tkinter as tk
from tkinter import filedialog
# Create file dialog
root = tk.Tk()
root.withdraw()
file_path = filedialog.askopenfilename(
title="Select device list file",
filetypes=[
("Text files", "*.txt"),
("CSV files", "*.csv"),
("All files", "*.*"),
],
)
if not file_path:
# User cancelled
return
# Validate file path
path_obj = Path(file_path)
if not path_obj.exists():
self.app.notify("File does not exist", severity="error", timeout=3)
return
if not path_obj.is_file():
self.app.notify(
"Selected path is not a file", severity="error", timeout=3
)
return
# Check file size (limit to 1 MB for safety)
file_size = path_obj.stat().st_size
if file_size > 1_000_000: # 1 MB
self.app.notify(
f"File too large ({file_size:,} bytes). Maximum 1 MB.",
severity="error",
timeout=5,
)
return
# Read file with proper encoding to preserve emojis
try:
with open(file_path, "r", encoding="utf-8") as f:
content = f.read()
except UnicodeDecodeError:
# Try with different encoding if UTF-8 fails
try:
with open(file_path, "r", encoding="latin-1") as f:
content = f.read()
self.app.notify(
"File loaded with Latin-1 encoding (UTF-8 failed)",
severity="warning",
timeout=3,
)
except Exception as e:
self.app.notify(
f"Error reading file: {str(e)}", severity="error", timeout=5
)
return
# Validate and sanitize content
lines = content.split("\n")
valid_lines = []
invalid_count = 0
# Pattern for valid hostnames/device names
# Allows: letters, numbers, hyphens, underscores, periods, and Unicode chars
hostname_pattern = re.compile(r"^[\w\-\.\u0080-\uFFFF]+$", re.UNICODE)
for line in lines:
line = line.strip()
if not line:
continue # Skip empty lines
# Check if line looks like a valid hostname/device name
if hostname_pattern.match(line):
valid_lines.append(line)
else:
invalid_count += 1
# Log but don't add invalid entries
if not valid_lines:
self.app.notify(
"No valid device names found in file", severity="warning", timeout=3
)
return
# Update text area with validated content
text_area = self.query_one("#device_input", TextArea)
text_area.text = "\n".join(valid_lines)
# Show notification
msg = f"✅ Loaded {len(valid_lines)} devices from file"
if invalid_count > 0:
msg += f" ({invalid_count} invalid entries skipped)"
self.app.notify(msg, severity="information", timeout=5)
except ImportError:
self.app.notify(
"tkinter not available - cannot open file dialog",
severity="error",
timeout=3,
)
except Exception as e:
self.app.notify(
f"Error loading file: {str(e)}", severity="error", timeout=5
)
@@ -1,10 +1,17 @@
"""
Policy Selector Widget Module
Provides a Textual widget for selecting target policies for bulk agent operations.
Allows users to browse available policies and select one as the destination for
moving agents. Automatically excludes parent/logical policies.
"""
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>
import logging
import re
@@ -34,7 +41,7 @@ class PolicySelector(Widget):
- Wildcard filtering (* and ?)
- Interactive table for policy browsing
- Explicit confirm button for selection
- Cancel/back button to dismiss
- Use escape key to go back
Attributes:
policies (list[Policy]): List of available Policy objects to display.
@@ -91,10 +98,10 @@ class PolicySelector(Widget):
- Clear Filter button
- Confirm Selection button
- Policy table displaying available policies
- Back buttons for navigation
- Use escape key to go back
"""
title_text = Static(
"🎯 Select Target Policy",
"Select Target Policy",
id="policy_selector_title",
)
title_text.styles.margin = (0, 0, 1, 0)
@@ -125,12 +132,12 @@ class PolicySelector(Widget):
filter_help.styles.margin = (0, 0, 1, 0)
yield filter_help
apply_button = Button(" Apply Filter", id="filter_button")
apply_button = Button("🔍 Apply Filter", id="filter_button")
apply_button.styles.width = "100%"
apply_button.styles.margin = (0, 0, 1, 0)
yield apply_button
clear_button = Button("Clear Filter", id="clear_filter_button")
clear_button = Button("🧹 Clear Filter", id="clear_filter_button")
clear_button.styles.width = "100%"
clear_button.styles.margin = (0, 0, 1, 0)
yield clear_button
@@ -144,11 +151,6 @@ class PolicySelector(Widget):
selected_label.styles.margin = (2, 0, 1, 0)
yield selected_label
cancel_button = Button("← Back", id="back_button")
cancel_button.styles.width = "100%"
cancel_button.styles.margin = (1, 0, 1, 0)
yield cancel_button
# Right side - Policy table
with Vertical() as right_side:
right_side.styles.width = "2fr"
@@ -222,7 +224,6 @@ class PolicySelector(Widget):
Handle button press events from the widget.
Routes to:
- back_button (Cancel): Pop screen without selecting
- filter_button (Apply Filter): Filter policies with wildcard support
- clear_filter_button: Clear filter and show all policies
- confirm_button: Confirm selection and post message
@@ -232,12 +233,7 @@ class PolicySelector(Widget):
"""
btn_id = event.button.id
if btn_id == "back_button":
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "filter_button":
if btn_id == "filter_button":
self._apply_filter()
event.stop()
@@ -283,7 +279,7 @@ class PolicySelector(Widget):
if self.selected_policy:
# Update selection display
label = self.query_one("#selected_policy_label", Static)
label.update(f"Selected: {self.selected_policy.name}")
label.update(f"Selected: {self.selected_policy.name}")
# Log for debugging
logger.debug(
@@ -325,7 +321,7 @@ class PolicySelector(Widget):
if highlighted_name:
label = self.query_one("#selected_policy_label", Static)
label.update(f"Highlighting: {highlighted_name}")
label.update(f"Highlighting: {highlighted_name}")
except Exception as e:
logger.error(f"Error handling row highlight: {e}")
@@ -399,7 +395,9 @@ class PolicySelector(Widget):
)
displayed_count = len(self._displayed_policies)
status_text = f"📊 Showing {displayed_count} of {len(self._filtered_policies)} policies"
status_text = (
f"Showing {displayed_count} of {len(self._filtered_policies)} policies"
)
self.app.notify(status_text, severity="information", timeout=2)
# Clear selection when filter is applied
@@ -474,7 +472,7 @@ class PolicySelector(Widget):
"""
if self.selected_policy is None:
self.app.notify(
"⚠️ Please select a policy first by clicking on a row in the table",
"Please select a policy first by clicking on a row in the table",
severity="warning",
timeout=3,
)
@@ -483,6 +481,6 @@ class PolicySelector(Widget):
# Log confirmation for debugging
logger.info(f"Confirming selection of policy: {self.selected_policy.name}")
self.app.notify(
f"Confirmed: {self.selected_policy.name}", severity="success", timeout=2
f"Confirmed: {self.selected_policy.name}", severity="success", timeout=2
)
self.post_message(self.PolicySelected(self.selected_policy))
@@ -1,10 +1,26 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from collections import defaultdict
import logging
from rich.text import Text
from textual.containers import Horizontal, Vertical
from textual.message import Message
from textual.widget import Widget
from textual.widgets import Input, OptionList, Static, Switch, Tree
from textual.widgets import Button, Input, OptionList, Static, Switch, Tree
from textual.widgets.option_list import Option
logger = logging.getLogger(__name__)
@@ -13,6 +29,27 @@ logger = logging.getLogger(__name__)
class PolicyTreeWidget(Widget):
"""Widget for displaying and searching a hierarchical policy tree."""
class ViewExecutionHistory(Message):
"""Message sent when user wants to view execution history for a device."""
def __init__(self, device):
super().__init__()
self.device = device
class GenerateOTP(Message):
"""Message sent when user wants to generate OTP for a device."""
def __init__(self, device):
super().__init__()
self.device = device
class ToggleEnforcement(Message):
"""Message sent when user wants to toggle audit/enforcement for a device."""
def __init__(self, device):
super().__init__()
self.device = device
def __init__(self, policies, devices):
super().__init__()
self.policies = policies
@@ -20,6 +57,7 @@ class PolicyTreeWidget(Widget):
self.last_highlighted_node = None
self.leaf_counts = defaultdict(int)
self.match_type = "Count" # Default to sorting by count
self.selected_device = None # Track currently selected device
def compose(self):
# Create the switch and its label
@@ -41,6 +79,18 @@ class PolicyTreeWidget(Widget):
search_box = Input(
placeholder="Search policies or devices...", id="tree_search"
)
exec_history_button = Button(
"📊 Execution History", id="view_exec_history_button", disabled=True
)
exec_history_button.styles.margin = (0, 1, 0, 0) # Right margin
otp_button = Button("🎫 Generate OTP", id="generate_otp_button", disabled=True)
otp_button.styles.margin = (0, 1, 0, 0) # Right margin
toggle_enforcement_button = Button(
"🔄 Toggle Enforcement/Audit", id="toggle_enforcement_button", disabled=True
)
# No right margin on last button
details_pane = Static("", id="details_pane")
# Layout the UI
@@ -58,6 +108,12 @@ class PolicyTreeWidget(Widget):
# Add the search box and details pane
yield label
yield search_box
# Action buttons in a horizontal row
with Horizontal() as button_row:
button_row.styles.height = "auto"
yield exec_history_button
yield otp_button
yield toggle_enforcement_button
yield details_pane
def on_mount(self) -> None:
@@ -74,6 +130,32 @@ class PolicyTreeWidget(Widget):
# Expand the root node
policy_tree.root.expand()
def refresh_data(self, policies, devices):
"""Refresh the widget with new data and rebuild the tree."""
self.policies = policies
self.devices = devices
self.selected_device = None
# Disable all buttons since selection is lost
try:
self.query_one("#view_exec_history_button", Button).disabled = True
self.query_one("#generate_otp_button", Button).disabled = True
self.query_one("#toggle_enforcement_button", Button).disabled = True
except:
pass
# Rebuild tree with new data
self._precompute_leaf_counts()
total_leaves = sum(
self.leaf_counts.get(policy.groupid, 0)
for policy in self.policies
if policy.parent == "global-policy-settings"
)
policy_tree = self.query_one("#policy_tree", Tree)
policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})")
self._build_tree()
policy_tree.root.expand()
def _precompute_leaf_counts(self):
"""Precompute leaf counts for each policy group."""
device_counts = defaultdict(int)
@@ -169,6 +251,9 @@ class PolicyTreeWidget(Widget):
node = message.node
data = node.data
details_pane = self.query_one("#details_pane", Static)
exec_history_button = self.query_one("#view_exec_history_button", Button)
otp_button = self.query_one("#generate_otp_button", Button)
toggle_enforcement_button = self.query_one("#toggle_enforcement_button", Button)
if self.last_highlighted_node is not None:
original_label = str(self.last_highlighted_node.label).strip()
@@ -183,6 +268,20 @@ class PolicyTreeWidget(Widget):
node.set_label(highlighted_label)
self.last_highlighted_node = node
# Check if selected node is a device (has Agent data)
from models.agent import Agent
if data and isinstance(data, Agent):
self.selected_device = data
exec_history_button.disabled = False
otp_button.disabled = False
toggle_enforcement_button.disabled = False
else:
self.selected_device = None
exec_history_button.disabled = True
otp_button.disabled = True
toggle_enforcement_button.disabled = True
if data:
details = "\n".join(
f"{key}: {value}" for key, value in data.__dict__.items()
@@ -272,3 +371,18 @@ class PolicyTreeWidget(Widget):
option_list.remove()
except:
pass
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button presses."""
if event.button.id == "view_exec_history_button":
if self.selected_device:
self.post_message(self.ViewExecutionHistory(self.selected_device))
event.stop()
elif event.button.id == "generate_otp_button":
if self.selected_device:
self.post_message(self.GenerateOTP(self.selected_device))
event.stop()
elif event.button.id == "toggle_enforcement_button":
if self.selected_device:
self.post_message(self.ToggleEnforcement(self.selected_device))
event.stop()
+870
View File
@@ -0,0 +1,870 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import logging
import os
import os.path
import re
from typing import List
import dotenv
import pandas as pd
from models.execution import ExecutionHistoryRecord
from models.policy import Allowlist, Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import get_system_list, get_system_value, load_env
from utils.selector import Selector
from utils.utils import (
areYouSure,
clear_screen,
colorText,
formatHTML,
get_sanitized_input,
locked,
open_directory,
print_x_wide,
regulator,
)
logger = logging.getLogger(__name__)
dotenv.load_dotenv()
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
logger.debug("Prompting for Policies")
print(colorText("Please select policy/policies", "white"))
selected = Selector.select_objects(policies, allow_multiple, prompt_each=True)
if selected is None:
return []
# Normalize to always return a list
logger.debug("Returning {selected.dict}")
return selected if isinstance(selected, list) else [selected]
def selectAllowlists(
api: AirlockAPIWrapper, policy=all, allow_multiple=True
) -> List[Allowlist]:
if policy == "all":
allowlists = [
Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()
]
else:
allowlists = [
Allowlist(**row.to_dict())
for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()
]
logger.debug("Prompting for Allowlist(s)")
print(colorText("Please select allowlist(s)", "white"))
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
if selected is None:
return []
# Normalize to always return a list
logger.debug(f"Returning {selected}")
return selected if isinstance(selected, list) else [selected]
def sortHashes(
api: AirlockAPIWrapper, selected_policies: List[Policy], type=[1, 2, 6, 7]
):
working_dir = load_env("WORKING_DIR")
history_days = Selector.select_value(
prompt="Enter how many days of history to pull (1-365): ",
value_type=int,
valid_range=(1, 365),
)
logger.debug(f"{history_days} day selected for history")
if history_days is None:
logging.warning("No history range selected. Aborting.")
return
policy_executions = ExecutionHistoryRecord.from_policies(
api, selected_policies, type_=type, history_days=history_days
)
logger.debug(f"Executions contains {policy_executions}")
enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(
api, policy_executions
)
categorized_executions = (
ExecutionHistoryRecord.categorize_executions_by_hash_decision(
enriched_executions
)
)
approved, unapproved, needs_review, unknown = (
ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions)
)
categories = {
"needs_review": needs_review,
"approved": approved,
"unapproved": unapproved,
"leftover": unknown,
}
for label, records in categories.items():
if not records:
continue # Skip empty or falsy categories
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html"
# Convert ExecutionHistoryRecord objects to dictionaries
df = pd.DataFrame([r.__dict__ for r in records])
# Optional: flatten hash_obj if needed
if not df.empty and "hash_obj" in df.columns:
hash_df = df["hash_obj"].apply(lambda h: h.to_dict() if h else {})
df = pd.concat([df.drop(columns=["hash_obj"]), hash_df], axis=1)
# Save to CSV
df.to_csv(csv_path, index=False)
logger.info(f"Saved {label} executions to {csv_path}")
# Generate HTML
formatHTML(df, html_path)
logger.info(f"Generated HTML report at {html_path}")
def buildPathsandPublishers(selected_policies: List[Policy], split):
working_dir = load_env("WORKING_DIR")
df1 = pd.DataFrame()
df2 = pd.DataFrame()
all_approved_hashes = pd.DataFrame()
path1 = (
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
)
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
path_exclusion_constant = get_system_value("PATH_EXCLUSION_CONST", cast_type=int)
if os.path.exists(path1):
df1 = pd.read_csv(path1)
else:
logger.warning(f"File not found: {path1}")
if os.path.exists(path2):
df2 = pd.read_csv(path2)
else:
logger.warning(f"File not found: {path2}")
if df1.empty and df2.empty:
logger.warning("Both DataFrames are empty. Skipping sort.")
all_approved_hashes = pd.DataFrame()
logger.debug(all_approved_hashes.head)
else:
all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
if "filename" in all_approved_hashes.columns:
all_approved_hashes = all_approved_hashes.sort_values(by="filename")
else:
logger.warning(
"Warning: 'filename' column not found in concatenated DataFrame."
)
if not all_approved_hashes.empty and path_exclusion_constant:
primary_path_exclusions = calculatePath(
all_approved_hashes,
path_exclusion_constant,
split,
)
remaining_hashes = all_approved_hashes[
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
]
secondary_path_exclusions = calculatePath(
remaining_hashes, (path_exclusion_constant - 1), split
)
remaining_hashes = remaining_hashes[
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
]
dataframes = {
"all_approved_hashes": all_approved_hashes,
"primary_Paths": primary_path_exclusions,
"secondary_Paths": secondary_path_exclusions,
"hashes_not_approvable_by_path": remaining_hashes,
}
logger.debug("Preparing to sort dataframes")
for name, df in dataframes.items():
logger.debug(f" DataFrame headers: {list(df.columns)}")
if "hashes" in name:
df.sort_values(by="filename", inplace=True)
else:
df.sort_values(by="longestcfp", inplace=True)
df.to_csv(
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv",
index=False,
)
formatHTML(
df,
f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html",
)
if not all_approved_hashes.empty:
# Drop all not signed, only keep unique values
publist = all_approved_hashes[
all_approved_hashes["publisher"] != "Not Signed"
].drop_duplicates(subset=["publisher"])
# Remove Bad publisher if somehow they made it this far
pattern = regulator(get_system_list("BAD_PUBLISHERS"))
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
publist = publist[["publisher"]]
publist.sort_values(by="publisher", inplace=True)
publist.to_csv(
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv",
index=False,
)
else:
logger.debug("Approved Hashes list appears empty")
def buildPreflights(selected_policies: List[Policy]):
working_dir = load_env("WORKING_DIR")
df1 = pd.DataFrame()
df2 = pd.DataFrame()
approved_hashes = pd.DataFrame()
approved_publishers = pd.DataFrame()
hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv"
path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv"
publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv"
# Read in and combine the two path generations
if os.path.exists(path1):
df1 = pd.read_csv(path1)
else:
logger.warning(f"File not found: {path1}")
if os.path.exists(path2):
df2 = pd.read_csv(path2)
else:
logger.warning(f"File not found: {path2}")
if df1.empty and df2.empty:
logger.warning("Both DataFrames are empty. Skipping sort.")
approved_paths = pd.DataFrame()
else:
approved_paths = pd.concat([df1, df2], ignore_index=True)
approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first")
# We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions.
if os.path.exists(hash):
hashes = pd.read_csv(hash)
approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])]
approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first")
else:
logger.warning(f"File not found: {hash}")
if os.path.exists(publishers):
approved_publishers = pd.read_csv(publishers)
else:
logger.warning(f"File not found: {publishers}")
dataframes = {
"approved_paths": approved_paths,
"approved_hashes": approved_hashes,
"approved_publishers": approved_publishers,
}
for name, df in dataframes.items():
logger.debug(f" DataFrame headers: {list(df.columns)}")
if name == "approved_paths":
df.sort_values(by="longestcfp", inplace=True)
elif name == "approved_hashes":
df.sort_values(by="filename", inplace=True)
elif name == "approved_publishers":
df.sort_values(by="publisher", inplace=True)
df.to_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv",
index=False,
)
formatHTML(
df,
f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html",
)
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
def clean_split(path):
if not isinstance(path, (str, bytes, os.PathLike)):
return []
parts = str(os.path.normpath(path)).split(os.sep)
parts = [p for p in parts if p] # Remove empty strings
return parts
# Diagnostic: log any non-string entries
non_string_entries = df[
~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))
]
if not non_string_entries.empty:
print(f"[WARNING] Non-string entries found in column '{col}':")
print(non_string_entries)
df = df.copy()
split_paths = df[col].apply(clean_split)
if min_files_for_path is not None:
df = df[
split_paths.apply(lambda parts: len(parts) >= min_files_for_path)
].copy()
split_paths = split_paths[df.index]
df["group_key"] = split_paths.apply(
lambda parts: os.sep.join(parts[:path_exclusion_constant])
)
grouped = df.groupby("group_key")
new_rows = []
for _, group_df in grouped:
paths = group_df[col].tolist()
split_parts = [clean_split(p) for p in paths]
def longest_common_prefix(paths):
if not paths:
return []
prefix = paths[0]
for path in paths[1:]:
prefix = [a for a, b in zip(prefix, path) if a == b]
if not prefix:
break
return prefix
common_prefix = longest_common_prefix(split_parts)
prefix_str = os.sep.join(common_prefix)
for i, parts in enumerate(split_parts):
filename = parts[-1]
middle = (
os.sep.join(parts[len(common_prefix) : -1])
if len(parts) > len(common_prefix) + 1
else ""
)
row = group_df.iloc[i].copy()
row["longestcfp"] = prefix_str
row["middle"] = middle
row["filename_only"] = filename
row["file_extension"] = os.path.splitext(filename)[1].lower()
new_rows.append(row)
return pd.DataFrame(new_rows).drop(columns=["group_key"])
def calculatePath(approved_hashes, path_exclusion_constant, split):
if split:
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
else:
dfs_by_policy = [approved_hashes]
badpathparts = get_system_list("BAD_PATH_PARTS")
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
processed_dfs = []
for df in dfs_by_policy:
haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename")
haslcp = haslcp.drop_duplicates()
forbidden = regulator(badpathparts, True)
forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
logger.debug("Removing forbidden filepaths for path exceptions")
print(colorText("Removing forbidden filepaths for path exceptions", "green"))
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
lcp_not_forbidden_review = lcp_not_forbidden[
[
"policyname",
"longestcfp",
"middle",
"filename_only",
"file_extension",
"sha256",
]
]
unique_sha_counts = (
lcp_not_forbidden_review.groupby("longestcfp")["sha256"]
.nunique()
.reset_index()
)
unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]
lcp_not_forbidden_review = lcp_not_forbidden_review.merge(
unique_sha_counts, on="longestcfp", how="left"
)
lcp_not_forbidden_review = lcp_not_forbidden_review[
lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path
]
processed_dfs.append(lcp_not_forbidden_review)
pathExclusions = pd.concat(processed_dfs, ignore_index=True)
return pathExclusions
def testChange(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR")
logger.info("These path exclusions would be added to:")
logger.info(destination_policy)
pathexclusions = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
)
hashes = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
)
unique_combinations = pathexclusions[
["longestcfp", "file_extension"]
].drop_duplicates()
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
processed_paths = [
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
for path, ext in unique_combinations.itertuples(index=False, name=None)
]
for path in processed_paths:
logger.info(path)
print(colorText("These publishers would added", "yellow"))
processed_publishers = []
if os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
):
publishers = pd.read_csv(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
)
if publishers.empty:
print(colorText("The publishers list is empty.", "red"))
else:
processed_publishers = (
publishers[publishers["publisher"] != "Not Signed"]["publisher"]
.drop_duplicates()
.tolist()
)
for publisher in processed_publishers:
print(publisher)
print(colorText("These hashes would be added to:", "yellow"))
print(destination_allowlist)
processed_hashes = hashes["sha256"].unique().tolist()
print_x_wide(processed_hashes, 3)
return processed_paths, processed_hashes, processed_publishers
def menu_policy_enforce(
api: AirlockAPIWrapper,
): # TODO Need to clean up 6 and 7 into functions
selected_policies = []
destination_policy = []
destination_allowlist = []
processed_paths = []
processed_hashes = []
processed_publishers = []
working_dir = load_env("WORKING_DIR")
while True:
printEnforceChecklist(
selected_policies, destination_policy, destination_allowlist
)
choice = get_sanitized_input("\nEnter your choice: ")
if choice == "1":
clear_screen()
selected_policies = selectPolicies(api, True)
elif choice == "2":
clear_screen()
print(
colorText(
"Please choose destination_name Policy for Path Exclusions", "white"
)
)
destination_policy = selectPolicies(api, False)
print(colorText("Please choose Allowlist for Hashes", "white"))
destination_allowlist = selectAllowlists(api, destination_policy, False)
elif choice == "3":
clear_screen()
sortHashes(
api,
selected_policies,
type=[1, 2, 6, 7],
)
elif choice == "4":
clear_screen()
if os.path.exists(
f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"
):
buildPathsandPublishers(selected_policies, False)
else:
print(
"File not found. Please make sure it's saved correctly and try again."
)
elif choice == "5":
clear_screen()
if os.path.exists(
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
) and os.path.exists(
f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
):
buildPreflights(selected_policies)
else:
print(
"File not found. Please make sure it's saved correctly and try again."
)
elif choice == "6":
clear_screen()
if (
os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
)
and os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
)
and destination_policy
and destination_allowlist
):
processed_paths, processed_hashes, processed_publishers = testChange(
selected_policies, destination_policy, destination_allowlist
)
else:
# Log which condition(s) failed
missing_items = []
if not os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
):
missing_items.append("approved_paths.csv not found")
if not os.path.exists(
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
):
missing_items.append("approved_hashes.csv not found")
if not destination_policy:
missing_items.append("destination_policy is empty or None")
if not destination_allowlist:
missing_items.append("destination_allowlist is empty or None")
logger.error("Preflight check failed due to the following:")
for item in missing_items:
logger.error(f" - {item}")
elif choice == "7":
clear_screen()
areYouSure()
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
if (
processed_paths
and processed_hashes
and processed_publishers
and destination_policy
and destination_allowlist
and confirmation.strip() == "I AGREE"
):
print(colorText("Proceeding with the code...", "yellow"))
api.hash_add_to_allowlist(
destination_allowlist[0].applicationid, processed_hashes
)
api.policy_add_path_exclusions(
destination_policy[0].groupid, processed_paths
)
if processed_publishers:
api.policy_add_publishers(
destination_policy[0].groupid, processed_publishers
)
locked()
else:
logger.error("Confirmation block failed. Reasons:")
if not processed_publishers or processed_hashes or processed_paths:
logger.error(" - Test not performed.")
if not destination_policy:
logger.error(" - `destination_policy` is missing or invalid.")
if not destination_allowlist:
logger.error(" - `destination_allowlist` is missing or invalid.")
if confirmation.strip() != "I AGREE":
logger.error(
" - User did not confirm with 'I AGREE'. Received: '%s'",
confirmation.strip(),
)
elif choice.upper() == "F":
open_directory(working_dir)
elif choice.upper() == "B":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
def section_header(title):
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(colorText(f" ------------- {title} -------------", "cyan"))
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR")
section_header("Prepare to Enforce Policy")
print(
colorText(
"\nSequentially follow these steps to prepare a policy for enforcement:",
"white",
)
)
# Step 1: Originating Policies
print(
colorText(
"\n1. Choose which policy or policies to gather execution info from", "cyan"
)
)
if not selected_policies:
print(colorText(" [❌] No policies have been chosen", "red"))
else:
print(colorText("The following policies have been chosen:", "green"))
for policy in selected_policies:
print(colorText(f" [✅] {policy.name}", "green"))
# Step 2: Destination Policy and Allowlist
print(
colorText("2. Choose the destination policy and associated allowlist", "cyan")
)
if destination_policy:
print(
colorText(
f" [✅] {destination_policy[0].name} has been selected as the destination policy",
"green",
)
)
else:
print(colorText(" [❌] No destination policy has been chosen", "red"))
if destination_allowlist:
print(
colorText(
f" [✅] {destination_allowlist[0].name} has been selected as allowlist",
"green",
)
)
else:
print(colorText(" [❌] No allowlist has been chosen", "red"))
# Step 3: Data Preparation
print(
colorText(
f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review",
"cyan",
)
)
if selected_policies:
policy_id = selected_policies[0].name
review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv"
print(
colorText(
(
" [✅] Data has been fetched"
if os.path.exists(review_path)
else " [❌] Data has not been fetched"
),
"green" if os.path.exists(review_path) else "red",
)
)
else:
print(
colorText(
" [❌] No policies selected, cannot check data fetch status", "red"
)
)
# Step 4: Manual Review
print(colorText("4. Manually review the files:", "cyan"))
print(
colorText(
" Remove the rows containing hashes you do not approve of", "cyan"
)
)
print(
colorText(
f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.",
"cyan",
)
)
print(
colorText(
" This will start the process to generate possible filepath approvals",
"cyan",
)
)
if selected_policies:
policy_id = selected_policies[0].name
approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv"
second_review_path = (
f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv"
)
print(
colorText(
(
" [✅] Reviewed hashes have been loaded"
if os.path.exists(approved_path)
else " [❌] Reviewed hashes have not been loaded"
),
"green" if os.path.exists(approved_path) else "red",
)
)
print(
colorText(
(
" [✅] Path review list created"
if os.path.exists(second_review_path)
else " [❌] Path review list has not been created"
),
"green" if os.path.exists(second_review_path) else "red",
)
)
else:
print(
colorText(
" [❌] No policies selected, cannot check reviewed hashes or path list",
"red",
)
)
# Step 5: Path Review
print(
colorText(
f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\",
"cyan",
)
)
print(
colorText(
" Remove the rows containing path exclusions or publishers you do not approve of.",
"cyan",
)
)
print(
colorText(
f" When complete, save the files to {working_dir}\\data\\Approved",
"cyan",
)
)
print(
colorText(" Choose this option when done to build your preflights", "cyan")
)
if selected_policies:
policy_id = selected_policies[0].name
reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv"
preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv"
preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv"
print(
colorText(
(
" [✅] Reviewed path list detected"
if os.path.exists(reviewed_path)
else " [❌] Path review list has not been detected"
),
"green" if os.path.exists(reviewed_path) else "red",
)
)
preflight_ready = os.path.exists(preflight_paths) and os.path.exists(
preflight_hashes
)
print(
colorText(
(
" [✅] Preflight Path Exclusion List has been generated"
if preflight_ready
else " [❌] Preflight Path Exclusion List has not been generated"
),
"green" if preflight_ready else "red",
)
)
else:
print(
colorText(
" [❌] No policies selected, cannot check preflight status", "red"
)
)
# Final Steps
print(
colorText(
"6. Test ------------------------------------------------------", "cyan"
)
)
print(
colorText(
" Prints to console the changes that would be made, must be done to proceed. ",
"cyan",
)
)
print(
colorText(
"7. Liftoff ------------------------------------------------------", "cyan"
)
)
print(
colorText(
" Apply path exclusions and approved publishers to selected policy",
"cyan",
)
)
print(colorText(" Apply approved hashes to allowlist", "cyan"))
# Utility Options
print(colorText("F. Open Working Directory", "cyan"))
print(colorText("B. Back", "cyan"))
@@ -1,3 +1,18 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import logging
from textual.containers import Horizontal, Vertical
@@ -59,15 +74,6 @@ class ResultsDisplay(Widget):
margin-top: 1;
width: 100%;
}
#button_row {
height: auto;
margin: 1 0 0 0;
}
#back_button {
width: 1fr;
}
"""
class CopySuccess(Message):
@@ -95,9 +101,9 @@ class ResultsDisplay(Widget):
def compose(self):
with Vertical(id="results_screen"):
yield Header(show_clock=True, icon="")
yield Header(show_clock=True, icon="")
# Title
title = Static(f"📊 {self.operation} - Results", id="results_title")
title = Static(f"{self.operation} - Results", id="results_title")
yield title
# Two-column layout
@@ -107,7 +113,7 @@ class ResultsDisplay(Widget):
yield Static("✅ Successful", id="success_label")
yield Static(self.successful_results, id="success_results")
yield Button(
"📋✅ Copy Success List",
"Copy Success List",
id="copy_success",
classes="copy_button",
)
@@ -117,15 +123,11 @@ class ResultsDisplay(Widget):
yield Static("❌ Failed", id="failure_label")
yield Static(self.unsuccessful_results, id="failure_results")
yield Button(
"📋❌ Copy Failure List",
"Copy Failure List",
id="copy_failure",
classes="copy_button",
)
# Back Button
with Horizontal(id="button_row"):
back_button = Button("← Back", id="back_button")
yield back_button
yield Footer()
def on_button_pressed(self, event: Button.Pressed) -> None:
@@ -138,18 +140,18 @@ class ResultsDisplay(Widget):
pyperclip.copy(str(success_widget.renderable))
self.app.notify(
"✅ Success list copied to clipboard!",
"Success list copied to clipboard!",
severity="information",
timeout=2,
)
self.post_message(self.CopySuccess())
except ImportError:
self.app.notify(
"⚠️ pyperclip not installed. Run: pip install pyperclip",
" pyperclip not installed. Run: pip install pyperclip",
severity="warning",
)
except Exception as e:
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
self.app.notify(f"¢ Failed to copy: {str(e)}", severity="error")
event.stop()
elif btn_id == "copy_failure":
@@ -159,20 +161,16 @@ class ResultsDisplay(Widget):
pyperclip.copy(str(failure_widget.renderable))
self.app.notify(
"✅ Failure list copied to clipboard!",
"Failure list copied to clipboard!",
severity="information",
timeout=2,
)
self.post_message(self.CopyFailure())
except ImportError:
self.app.notify(
"⚠️ pyperclip not installed. Run: pip install pyperclip",
" pyperclip not installed. Run: pip install pyperclip",
severity="warning",
)
except Exception as e:
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
event.stop()
elif btn_id == "back_button":
self.app.pop_screen()
self.app.notify(f" Failed to copy: {str(e)}", severity="error")
event.stop()
+261
View File
@@ -0,0 +1,261 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import datetime
import logging
from bson import ObjectId
from textual.app import ComposeResult
from textual.containers import Container, Vertical
from textual.widgets import Button, DataTable, Input, Static
from services.API import AirlockAPIWrapper
logger = logging.getLogger(__name__)
def skipback(days):
"""
Generate a MongoDB ObjectId for a given number of days ago from today.
"""
adjusted_days = days
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
days=adjusted_days
)
timestamp = int(date_days_ago.timestamp())
hex_timestamp = format(timestamp, "08x")
objectid_hex = hex_timestamp + "0000000000000000"
return ObjectId(objectid_hex)
class ServerLogWidget(Vertical):
"""Widget for displaying server activity logs in a DataTable."""
DEFAULT_CSS = """
ServerLogWidget {
width: 100%;
height: 100%;
}
ServerLogWidget #status_bar {
width: 100%;
height: auto;
background: $surface;
padding: 1;
margin-bottom: 1;
}
ServerLogWidget DataTable {
height: 1fr;
border: solid $primary;
}
ServerLogWidget #button_container {
width: 100%;
height: auto;
layout: horizontal;
padding: 1;
align: left middle;
}
ServerLogWidget .filter_label {
width: auto;
height: 3;
content-align: left middle;
padding-right: 1;
}
ServerLogWidget #filter_input {
width: 40;
margin-right: 1;
}
ServerLogWidget Button {
margin-right: 1;
}
"""
def __init__(self, api: AirlockAPIWrapper):
super().__init__()
self.api = api
self.all_logs = [] # Store all logs for filtering
self.columns = [] # Store column names
def compose(self) -> ComposeResult:
yield Static("Loading server logs (last 72 hours)...", id="status_bar")
yield DataTable(id="server_log_table")
with Container(id="button_container"):
yield Static("Filter:", classes="filter_label")
yield Input(placeholder="Filter (use * and ? wildcards)", id="filter_input")
yield Button("Refresh", id="refresh_button", variant="primary")
def on_mount(self) -> None:
"""Initialize the DataTable and load server logs."""
self.load_logs()
def load_logs(self) -> None:
"""Load server logs from the API and populate the DataTable."""
table = self.query_one("#server_log_table", DataTable)
status = self.query_one("#status_bar", Static)
try:
status.update("⏳ Loading server logs (last 72 hours)...")
# Create a fake checkpoint for 3 days ago (72 hours)
checkpoint = str(skipback(3))
# Get server logs from API
logs = self.api.server_logs(checkpoint=checkpoint)
if not logs:
status.update("ℹï¸ No server logs found in the last 72 hours.")
table.clear(columns=True)
self.all_logs = []
self.columns = []
return
# Store all logs for filtering
self.all_logs = logs
# Clear existing data
table.clear(columns=True)
# Add columns based on the first log entry
if logs:
first_log = logs[0]
self.columns = [col for col in first_log.keys() if col != "checkpoint"]
for col in self.columns:
table.add_column(col, key=col)
# Add rows in reverse order so newest entries are at the top
for log_entry in reversed(logs):
row_data = []
for col in self.columns:
value = log_entry.get(col, "")
# Format datetime column to be more readable
if col == "datetime" and value:
try:
# Parse ISO format and convert to readable format
dt = datetime.datetime.fromisoformat(
str(value).replace("Z", "+00:00")
)
value = dt.strftime("%Y-%m-%d %H:%M:%S")
except Exception:
# If parsing fails, just use the original value
pass
row_data.append(str(value))
table.add_row(*row_data)
status.update(
f"✅ Loaded {len(logs)} log entries from the last 72 hours"
)
logger.info(f"Loaded {len(logs)} server log entries")
else:
status.update("ℹï¸ No log entries found.")
self.all_logs = []
self.columns = []
except Exception as exc:
error_msg = f"❌ Error loading server logs: {exc}"
status.update(error_msg)
logger.error(f"Failed to load server logs: {exc}", exc_info=True)
table.clear(columns=True)
self.all_logs = []
self.columns = []
def filter_logs(self, filter_text: str) -> None:
"""Filter the logs based on the filter text with wildcard support."""
import fnmatch
table = self.query_one("#server_log_table", DataTable)
status = self.query_one("#status_bar", Static)
if not self.all_logs:
return
# Clear existing data
table.clear(columns=True)
# Re-add columns
for col in self.columns:
table.add_column(col, key=col)
# Filter logs
filtered_logs = []
if filter_text.strip():
filter_pattern = filter_text.strip().lower()
for log_entry in self.all_logs:
# Check if any field matches the filter pattern
match = False
for col in self.columns:
value = str(log_entry.get(col, "")).lower()
if fnmatch.fnmatch(value, filter_pattern):
match = True
break
if match:
filtered_logs.append(log_entry)
else:
# No filter, show all logs
filtered_logs = self.all_logs
# Add filtered rows in reverse order
for log_entry in reversed(filtered_logs):
row_data = []
for col in self.columns:
value = log_entry.get(col, "")
# Format datetime column to be more readable
if col == "datetime" and value:
try:
dt = datetime.datetime.fromisoformat(
str(value).replace("Z", "+00:00")
)
value = dt.strftime("%Y-%m-%d %H:%M:%S")
except Exception:
pass
row_data.append(str(value))
table.add_row(*row_data)
if filter_text.strip():
status.update(
f"✅ Showing {len(filtered_logs)} of {len(self.all_logs)} log entries (filtered)"
)
else:
status.update(
f"✅ Loaded {len(self.all_logs)} log entries from the last 72 hours"
)
logger.info(
f"Filtered to {len(filtered_logs)} entries with pattern: {filter_text}"
)
def on_input_changed(self, event: Input.Changed) -> None:
"""Handle filter input changes."""
if event.input.id == "filter_input":
self.filter_logs(event.value)
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button presses."""
button_id = event.button.id
if button_id == "refresh_button":
self.load_logs()
# Clear the filter input when refreshing
try:
filter_input = self.query_one("#filter_input", Input)
filter_input.value = ""
except Exception:
pass
event.stop()
-61
View File
@@ -1,61 +0,0 @@
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)
)
-24
View File
@@ -1,24 +0,0 @@
# 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."""
+386 -587
View File
File diff suppressed because it is too large Load Diff
+13 -12
View File
@@ -1,30 +1,31 @@
[package]
name = "airlock_libs"
version = "5.0.0"
version = "7.0.0"
edition = "2024"
[lib]
crate-type = ["cdylib"]
[dependencies]
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"}
opentelemetry = { version = "0.27.0", features = ["logs", "metrics", "trace"] }
opentelemetry-otlp = { version = "0.27.0", features = ["trace", "metrics", "grpc-tonic", "http-proto", "tls", "reqwest-client", "reqwest-rustls"] }
opentelemetry-semantic-conventions = { version = "0.27.0" }
opentelemetry-proto = { version = "0.27.0"}
pyo3 = { version = "0.27.0", features = ["extension-module", "generate-import-lib"] }
reqwest = { version = "0.12.24", features = ["json", "native-tls"] }
reqwest = { version = "0.12.24", features = ["json", "native-tls", "rustls-tls"] }
serde = "1.0.228"
serde-pyobject = "0.8.0"
serde_json = "1.0.145"
tokio = { version = "1.48.0", features = ["full"] }
tonic = { version = "0.8.2", features = ["tls-roots"] }
tonic = { version = "0.12.3", features = ["tls-roots"] }
tracing = "0.1.41"
tracing-subscriber = "0.3.20"
tracing-opentelemetry = "0.32.0"
pyo3-async-runtimes = { version = "0.27.0", features = ["async-std", "tokio"] }
crossbeam = "0.8.4"
log = "0.4.29"
flexi_logger = "0.31.7"
opentelemetry-appender-log = "0.27.0"
opentelemetry_sdk = { version = "0.27.0", features = ["rt-tokio", "testing", "trace"] }
[package.metadata.maturin]
generate-abi-stubs = true
@@ -37,4 +38,4 @@ codegen-units = 1
panic = 'abort'
strip = true
debug-assertions = false
overflow-checks = false
overflow-checks = true
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "maturin"
[project]
name = "airlock_libs"
version = "5.0.0"
version = "7.0.0"
description = "Airlock Digital API Wrapper"
readme = "README.md"
license = { text = "AGPL-3.0-only" }
+3 -1
View File
@@ -1,5 +1,7 @@
use pyo3::prelude::*;
mod services;
pub mod modules;
pub mod prelude;
pub mod services;
#[pymodule]
fn airlock_libs(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(services::pull_policy_exec_histories, py)?)?;
+131
View File
@@ -0,0 +1,131 @@
use crate::prelude::*;
use crate::services::get_base_directory;
#[allow(non_snake_case)]
#[derive(Deserialize, Debug)]
pub struct TelemetryConfig {
pub TELEMETRY: bool,
pub TELEM_URL: Option<String>,
}
impl TelemetryConfig {
pub fn init_tracer() -> opentelemetry_sdk::trace::TracerProvider {
let cfg: TelemetryConfig = TelemetryConfig::load();
if !cfg.TELEMETRY {
return TracerProvider::builder().build();
}
let endpoint = cfg.TELEM_URL.unwrap_or_default();
let channel = Channel::from_shared(endpoint.clone())
.unwrap()
.tls_config(ClientTlsConfig::new().with_native_roots())
.unwrap()
.connect_lazy();
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_tonic()
.with_endpoint(endpoint.clone())
.with_channel(channel)
.build()
.expect("Failed to build exporter");
opentelemetry_sdk::trace::TracerProvider::builder()
.with_simple_exporter(exporter)
.with_resource(Resource::new(vec![KeyValue::new(
"service.name",
"LoxideLibs",
)]))
.build()
}
fn load() -> Self {
let cfg_path = get_base_directory().join("config\\user_config.json");
if !cfg_path.exists() {
return Self {
TELEMETRY: false,
TELEM_URL: None,
};
}
match fs::read_to_string(&cfg_path) {
Ok(contents) => serde_json::from_str::<Self>(&contents).unwrap_or(Self {
TELEMETRY: false,
TELEM_URL: None,
}),
Err(_) => Self {
TELEMETRY: false,
TELEM_URL: None,
},
}
}
}
#[derive(Debug, Deserialize, Serialize)]
pub struct ApiResponse {
pub(crate) error: String,
pub(crate) response: ExecHistories,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct ExecHistories {
pub(crate) exechistories: Vec<Group>,
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct Group {
pub(crate) checkpoint: String,
#[serde(rename = "type")]
pub(crate) exectype: u8,
pub(crate) username: String,
pub(crate) hostname: String,
pub(crate) netdomain: String,
pub(crate) filename: String,
pub(crate) ppolicy: String,
pub(crate) policyname: String,
pub(crate) policyver: String,
pub(crate) commandline: String,
pub(crate) publisher: String,
pub(crate) pprocess: String,
pub(crate) gprocess: String,
pub(crate) sha256: String,
pub(crate) datetime: String,
pub(crate) md5: String,
pub(crate) sha128: String,
pub(crate) sha384: String,
pub(crate) sha512: String,
pub(crate) ip: String,
pub(crate) localip: String,
}
pub struct PyData {
pub headers: reqwest::header::HeaderMap,
pub base_url: String,
}
impl PyData {
pub fn extract_data(py: Python<'_>, obj: &Py<PyAny>) -> Self {
let headers_raw = obj.getattr(py, "headers").unwrap().to_string();
let headers_json = headers_raw.replace('\'', "\"");
let parsed: Value = serde_json::from_str(&headers_json).unwrap();
let mut header_map = HeaderMap::new();
if let Some(obj) = parsed.as_object() {
for (key, val) in obj {
if let Some(v) = val.as_str() {
let header_name = HeaderName::from_str(key).unwrap();
let header_value: HeaderValue = HeaderValue::from_str(v).unwrap();
header_map.insert(header_name, header_value);
}
}
}
let base_url = obj.getattr(py, "base_url").unwrap().to_string();
Self {
headers: header_map,
base_url,
}
}
}
pub struct SkipBack;
impl SkipBack {
pub fn find_checkpoint(days: i64) -> ObjectId {
let date_days_ago = Local::now() - Duration::days(days);
let timestamp = date_days_ago.timestamp() as u32;
let mut hex_timestamp = String::new();
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap();
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
}
}
+1
View File
@@ -0,0 +1 @@
pub mod datatypes;
View File
+32
View File
@@ -0,0 +1,32 @@
pub use chrono::{Duration, Local, NaiveDate};
pub use crossbeam::channel::unbounded;
pub use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
pub use mongodb::bson::oid::ObjectId;
pub use opentelemetry::global::GlobalTracerProvider;
pub use opentelemetry::trace::noop::NoopTracerProvider;
pub use opentelemetry::trace::{Status, TraceContextExt, Tracer};
pub use opentelemetry::*;
pub use opentelemetry_otlp::ExportConfig;
pub use opentelemetry_otlp::WithExportConfig;
pub use opentelemetry_otlp::WithTonicConfig;
pub use opentelemetry_sdk::Resource;
pub use opentelemetry_sdk::trace::{Config, TracerProvider};
pub use pyo3::{prelude::*, types::PyString};
pub use reqwest::{
Client,
header::{HeaderMap, HeaderName, HeaderValue},
};
pub use serde::{Deserialize, Serialize};
pub use serde_json::Value;
pub use std::sync::{Arc, Mutex};
pub use std::thread;
pub use std::{
collections::HashMap,
env,
fmt::Write,
fs::{self, File},
io::{Read, Seek, SeekFrom},
path::PathBuf,
str::FromStr,
};
pub use tonic::transport::{Channel, ClientTlsConfig};
+152 -284
View File
@@ -1,163 +1,35 @@
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 pyo3_async_runtimes::async_std;
use reqwest::{
Client,
header::{HeaderMap, HeaderName, HeaderValue},
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{
collections::HashMap,
env,
fmt::Write,
fs::{self, File},
io::{Read, Seek, SeekFrom},
path::PathBuf,
str::FromStr,
};
#[allow(non_snake_case)]
#[derive(Deserialize, Debug)]
struct TelemetryConfig {
TELEMETRY: bool,
TELEM_URL: Option<String>,
}
impl TelemetryConfig {
pub fn load() -> Self {
let cfg_path = get_base_directory().join("config\\user_config.json");
if !cfg_path.exists() {
return Self {
TELEMETRY: false,
TELEM_URL: None,
};
}
match fs::read_to_string(&cfg_path) {
Ok(contents) => serde_json::from_str::<Self>(&contents).unwrap_or(Self {
TELEMETRY: false,
TELEM_URL: None,
}),
Err(_) => Self {
TELEMETRY: false,
TELEM_URL: None,
},
}
}
}
#[derive(Debug, Deserialize, Serialize)]
struct ApiResponse {
error: String,
response: ExecHistories,
}
#[derive(Debug, Deserialize, Serialize)]
struct ExecHistories {
exechistories: Vec<Group>,
}
#[derive(Debug, Deserialize, Serialize, Clone)]
struct Group {
checkpoint: String,
#[serde(rename = "type")]
exectype: u8,
username: String,
hostname: String,
netdomain: String,
filename: String,
ppolicy: String,
policyname: String,
policyver: String,
commandline: String,
publisher: String,
pprocess: String,
gprocess: String,
sha256: String,
datetime: String,
md5: String,
sha128: String,
sha384: String,
sha512: String,
ip: String,
localip: String,
}
enum ExtractedValues {
Headers(reqwest::header::HeaderMap),
BaseUrl(String),
}
trait Converter {
fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues;
}
struct PyData;
impl Converter for PyData {
fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues {
if extract_headers {
let headers = py_self.getattr(py, "headers").unwrap().to_string();
let headers_replace = headers.replace('\'', "\"");
let parsed: Value = serde_json::from_str(headers_replace.as_str()).unwrap();
let mut header_map = HeaderMap::new();
if let Some(obj) = parsed.as_object() {
for (_key, value) in obj {
if let Some(v) = value.as_str() {
let val = HeaderValue::from_str(v).unwrap();
header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val);
}
}
}
ExtractedValues::Headers(header_map)
} else {
let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
ExtractedValues::BaseUrl(base_url)
}
}
}
use crate::modules::datatypes::*;
use crate::prelude::*;
#[pyfunction]
pub fn pull_policy_exec_histories(
py: Python<'_>,
py_self: Py<PyAny>,
policy_names: String,
policy_names: Option<String>,
exec_types: String,
days: i64,
) -> Py<PyString> {
let headers: HeaderMap = match PyData::convert(py, &py_self, true) {
ExtractedValues::Headers(h) => h,
ExtractedValues::BaseUrl(_) => std::process::abort(),
};
let base_url = match PyData::convert(py, &py_self, false) {
ExtractedValues::Headers(_) => std::process::abort(),
ExtractedValues::BaseUrl(b) => b,
};
let handle = std::thread::spawn(move || {
let rt = match tokio::runtime::Runtime::new() {
println!();
let data: PyData = PyData::extract_data(py, &py_self);
let headers: HeaderMap = data.headers;
let base_url: String = data.base_url;
let handle: thread::JoinHandle<String> = std::thread::spawn(move || {
let rt: tokio::runtime::Runtime = match tokio::runtime::Runtime::new() {
Ok(rt) => rt,
Err(e) => {
println!("Failed to build Tokio Runtime: {:?}", e);
std::process::abort();
}
};
rt.block_on(async {
let _ = init_tracer();
});
let tracer = global::tracer("global_tracer");
let _cx = Context::new();
let tracer_provider = rt.block_on(async { TelemetryConfig::init_tracer() });
global::set_tracer_provider(tracer_provider.clone());
let tracer: global::BoxedTracer = global::tracer("tracer");
let _cx: Context = Context::new();
let file_path: PathBuf = format!(
"{}\\cache\\chunkinator.json",
get_base_directory().display()
)
.into();
let writeable_filepath = file_path.clone();
if !&file_path.exists() {
if let Some(parent_dir) = &file_path.parent()
&& !parent_dir.exists()
@@ -178,13 +50,14 @@ pub fn pull_policy_exec_histories(
}
}
}
let data = ApiResponse {
let data: ApiResponse = ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: vec![],
},
};
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
let writeable_filepath: PathBuf = file_path.clone();
let data_write: String = serde_json::to_string_pretty(&data).expect("Failed to serialize");
match fs::write(writeable_filepath.clone(), data_write) {
Ok(_) => {}
Err(e) => {
@@ -192,18 +65,19 @@ pub fn pull_policy_exec_histories(
std::process::abort();
}
}
let mut checkpoint_number: String = skipback(days).to_string();
let multi_progress = MultiProgress::new();
multi_progress.set_draw_target(ProgressDrawTarget::stdout());
let progress_bar = multi_progress.add(ProgressBar::new(100));
progress_bar.set_style(
let mut checkpoint_number: String = SkipBack::find_checkpoint(days).to_string();
let progress_bar = Arc::new(Mutex::new(ProgressBar::new(100)));
progress_bar
.lock()
.unwrap()
.set_draw_target(ProgressDrawTarget::stderr());
progress_bar.lock().unwrap().set_style(
ProgressStyle::default_bar()
.template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
.unwrap(),
.template("Total - Policy Name: {msg}: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
.unwrap().progress_chars("≫>")
);
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
let client = tracer.in_span("Building HTTP Client", |cx| {
let client_result = build_client(headers);
let client: Client = tracer.in_span("Building HTTP Client", |cx| {
let client_result: Result<Client, reqwest::Error> = build_client(headers);
match client_result {
Ok(client_result) => {
cx.span().add_event(
@@ -230,26 +104,94 @@ pub fn pull_policy_exec_histories(
}
}
});
let cutoff = Local::now().naive_local() - Duration::days(days);
let mut f = match File::open(&writeable_filepath) {
Ok(f) => f,
Err(e) => {
println!("Failed to Access {:?}: {}", &writeable_filepath, e);
std::process::abort();
}
let cutoff: chrono::NaiveDateTime =
Local::now().naive_local() - chrono::Duration::days(days);
let (tx, rx) = unbounded::<Vec<Group>>();
let pb_clone = progress_bar.clone();
thread::spawn(move || {
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
{
let contents: String = fs::read_to_string(&writeable_filepath).unwrap_or_default();
let existing: ApiResponse =
serde_json::from_str(&contents).unwrap_or(ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: vec![],
},
});
existing
.response
.exechistories
.into_iter()
.map(|entry: Group| {
(
(
entry.sha256.clone(),
entry.filename.clone(),
entry.hostname.clone(),
),
entry,
)
})
.collect()
} else {
HashMap::new()
};
tracer.in_span("Airlock Data Retreival", |cx| {
let span = cx.span();
span.set_attribute(Key::new("Days").string(days.to_string()));
span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
loop {
match f.seek(SeekFrom::Start(0)) {
while let Ok(parsed_responses) = rx.recv() {
for executions in parsed_responses {
if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
continue;
}
let history_date: NaiveDate = match NaiveDate::parse_from_str(
&executions.datetime.replace(" +0000 UTC", ""),
"%Y-%m-%dT%H:%M:%SZ",
) {
Ok(date) => date,
Err(_) => continue,
};
if history_date >= cutoff.into() {
let key: (String, String, String) = (
executions.sha256.clone(),
executions.filename.clone(),
executions.hostname.clone(),
);
seen.entry(key).or_insert(executions.clone());
}
}
let final_response: ApiResponse = ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: seen.values().cloned().collect(),
},
};
let data_write: String = serde_json::to_string_pretty(&final_response).unwrap();
match fs::write(&writeable_filepath, data_write) {
Ok(_) => {}
Err(e) => {
println!("Failed to seek start of {:?}: {}", f, e);
std::process::abort();
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
}
}
}
});
let mut first_date: Option<NaiveDate> = None;
tracer.in_span("Airlock Data Retreival", |cx| {
pb_clone
.lock()
.unwrap()
.enable_steady_tick(std::time::Duration::from_millis(100));
pb_clone
.lock()
.unwrap()
.set_message(policy_names.clone().unwrap_or("Statistics".to_string()));
let span: opentelemetry::trace::SpanRef<'_> = cx.span();
span.set_attribute(KeyValue::new("Days", days.to_string()));
span.set_attribute(KeyValue::new(
"Policy Name",
policy_names
.clone()
.unwrap_or("Statistics Monitoring".to_string()),
));
loop {
let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
let results: ApiResponse = history_logging(
&base_url,
@@ -264,103 +206,53 @@ pub fn pull_policy_exec_histories(
));
results
});
let parsed_responses = execution_histories.response.exechistories;
let parsed_responses: Vec<Group> = execution_histories.response.exechistories;
if parsed_responses.is_empty() {
break;
}
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 =
serde_json::from_str(&contents).unwrap_or(ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: vec![],
},
});
existing_data
.response
.exechistories
.into_iter()
.map(|entry| {
(
(
entry.sha256.clone(),
entry.filename.clone(),
entry.hostname.clone(),
),
entry,
)
})
.collect()
} else {
HashMap::new()
};
for (index, executions) in parsed_responses.iter().enumerate() {
if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
continue;
}
if index == parsed_responses.len() - 1 {
checkpoint_number = executions.checkpoint.clone();
break;
}
let history_date = match NaiveDate::parse_from_str(
&executions.datetime.replace(" +0000 UTC", ""),
"%Y-%m-%dT%H:%M:%SZ",
) {
Ok(date) => date,
Err(_) => continue,
};
if history_date >= cutoff.into() {
let key = (
executions.sha256.clone(),
executions.filename.clone(),
executions.hostname.clone(),
);
seen.entry(key).or_insert(executions.clone());
}
}
let final_response = ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: seen.values().cloned().collect(),
},
};
let data_write = serde_json::to_string_pretty(&final_response).unwrap();
match fs::write(&writeable_filepath, data_write) {
Ok(_) => {}
Err(e) => {
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
}
}
if let Some(last_item) = &final_response.response.exechistories.last()
tx.send(parsed_responses.clone()).unwrap();
checkpoint_number = parsed_responses.last().unwrap().checkpoint.clone();
if let Some(last_item) = parsed_responses.last()
&& let Ok(last_date) = NaiveDate::parse_from_str(
&last_item.datetime.replace(" +0000 UTC", ""),
"%Y-%m-%dT%H:%M:%SZ",
)
{
let date_diff = Local::now().naive_local().date() - last_date;
let percentage_diff =
(days - date_diff.num_days()) as f64 / days as f64 * 100.0;
progress_bar.set_position(percentage_diff.round() as u64);
progress_bar.set_message("Total Percent Complete");
if first_date.is_none() {
first_date = Some(last_date);
}
if let Some(base_date) = first_date {
let date_diff: chrono::TimeDelta = last_date - base_date;
let total_span: i64 =
(Local::now().naive_local().date() - base_date).num_days();
let percentage: u64 = ((date_diff.num_days() as f64 / total_span as f64)
* 100.0)
.clamp(0.0, 100.0)
.round() as u64;
pb_clone.lock().unwrap().set_position(percentage);
}
}
}
});
progress_bar.finish_with_message("All Checkpoints Complete");
let return_data = match fs::read_to_string(&writeable_filepath) {
progress_bar
.lock()
.unwrap()
.finish_with_message("All Checkpoints Complete");
let return_data: String = match fs::read_to_string(file_path.clone()) {
Ok(return_data) => return_data,
Err(e) => {
println!("Failed to read data from: {:?}: {}", &writeable_filepath, e);
println!("Failed to read data from: {:?}: {}", &file_path, e);
std::process::abort();
}
};
shutdown_tracer_provider();
tracer_provider
.shutdown()
.expect("Failed to Shutdown Tracer Provdier");
drop(tx);
return_data.to_string()
});
let gil_value = handle.join().unwrap();
Python::attach(|py| PyString::new(py, &gil_value).into())
let gil_value: String = handle.join().unwrap();
Python::attach(|py: Python<'_>| PyString::new(py, &gil_value).into())
}
fn build_client(headers: HeaderMap) -> Result<reqwest::Client, reqwest::Error> {
@@ -376,18 +268,22 @@ async fn history_logging(
base_url: &String,
exec_types: &String,
checkpoint_number: &String,
policy_names: &String,
policy_names: &Option<String>,
client: &Client,
) -> ApiResponse {
let policy_json = match policy_names {
Some(name) => format!(r#"[ "{}" ]"#, name), // JSON array with one element
None => "[]".to_string(), // Empty JSON array
};
let payload = format!(
r#"{{
"type": {},
"checkpoint": "{}",
"policy": ["{}"]
"policy": {}
}}"#,
exec_types, checkpoint_number, policy_names
exec_types, checkpoint_number, policy_json
);
let res = client
let res: Result<reqwest::Response, reqwest::Error> = client
.post(format!("{}/v1/logging/exechistories", base_url))
.body(payload)
.send()
@@ -410,7 +306,7 @@ async fn history_logging(
}
}
fn get_base_directory() -> PathBuf {
pub fn get_base_directory() -> PathBuf {
let home = env::var_os("HOME")
.map(PathBuf::from)
.or_else(|| env::var_os("USERPROFILE").map(PathBuf::from))
@@ -423,38 +319,10 @@ fn get_base_directory() -> PathBuf {
.unwrap_or_else(|| home.join("AppData").join("Roaming"));
appdata.join("Loxide")
}
_ => home.join(".local").join("share").join("Loxide"),
"linux" => home.join(".local").join("share").join("Loxide"),
_ => {
println!("{} is currently not compatible with LoxideLibs", os);
std::process::abort();
}
}
fn skipback(days: i64) -> ObjectId {
let date_days_ago = Local::now() - Duration::days(days);
let timestamp = date_days_ago.timestamp() as u32;
let mut hex_timestamp = String::new();
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap();
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
}
fn init_tracer() -> Result<Option<sdktrace::Tracer>, TraceError> {
let cfg = TelemetryConfig::load();
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))
}
+37 -28
View File
@@ -1,6 +1,17 @@
"""
This module handles the creation of local approval requests.
"""
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import logging
import os
@@ -51,7 +62,7 @@ class LocalApprovalRequestor:
batch_id = int(time.time())
purpose = (
f"🎫 Local Approval 🎫 - {duration_minutes} mins - "
f" Local Approval - {duration_minutes} mins - "
f"batch:{batch_id} Client:{agent_id} User:{self.username}"
)
@@ -103,11 +114,9 @@ class LocalApprovalRequestor:
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")
)
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:
@@ -125,11 +134,11 @@ class LocalApprovalRequestor:
if not move_success:
raise Exception("Failed to move to audit policy")
print(colorText(f"✓ {agent.hostname}", "green"))
print(colorText(f" {agent.hostname}", "green"))
success_count += 1
except Exception as e:
print(colorText(f"✗ {agent.hostname}: {e}", "red"))
print(colorText(f" {agent.hostname}: {e}", "red"))
logger.error(f"Error processing agent {agent.hostname}: {e}")
failure_count += 1
@@ -152,7 +161,7 @@ class LocalApprovalRequestor:
]
# Display duration options
print(colorText("\n⏱️ Select Local Approval Duration:", "white"))
print(colorText("\n Select Local Approval Duration:", "white"))
print(colorText("=" * 50, "white"))
for i, (minutes, label) in enumerate(duration_options, start=1):
@@ -166,36 +175,36 @@ class LocalApprovalRequestor:
if 1 <= choice <= len(duration_options):
duration_minutes, duration_label = duration_options[choice - 1]
print(colorText(f"✓ Selected: {duration_label}", "green"))
print(colorText(f" Selected: {duration_label}", "green"))
logger.info(f"User selected duration: {duration_minutes} minutes")
else:
print(colorText("❌ Invalid choice.", "red"))
print(colorText(" Invalid choice.", "red"))
logger.warning("Invalid duration choice")
return
except ValueError:
print(colorText("❌ Invalid input. Please enter a number.", "red"))
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"))
print(colorText("\nSelect Agents for Local Approval:", "white"))
agents = selectAgents(self.api)
if not agents:
print(colorText("❌ No agents found or error retrieving agents.", "red"))
print(colorText(" No agents found or error retrieving agents.", "red"))
logger.warning("No agents selected or error retrieving agents")
return
# Confirm with user
print(colorText("\n📋 Summary:", "cyan"))
print(colorText("\nSummary:", "cyan"))
print(colorText(f" Duration: {duration_label}", "white"))
print(colorText(f" Agents: {len(agents)}", "white"))
confirm = get_sanitized_input("\nProceed? (y/n): ").lower()
if confirm != "y":
print(colorText("❌ Operation cancelled.", "yellow"))
print(colorText(" Operation cancelled.", "yellow"))
return
# Process the batch
@@ -219,24 +228,24 @@ class LocalApprovalRequestor:
failure_count: Number of failed operations
"""
print(colorText(f"\n{'=' * 60}", "white"))
print(colorText("📊 Local Approval Summary", "cyan"))
print(colorText(" Local Approval Summary", "cyan"))
print(colorText("=" * 60, "white"))
print(colorText(f"✓ Successfully processed: {success_count}", "green"))
print(colorText(f" Successfully processed: {success_count}", "green"))
if failure_count > 0:
print(colorText(f"✗ Failed: {failure_count}", "red"))
print(colorText(f" Failed: {failure_count}", "red"))
print(colorText(f"\n📦 Batch ID: {batch_id}", "cyan"))
print(colorText(f"⏱️ Duration: {duration_label}", "cyan"))
print(colorText(f"\n Batch ID: {batch_id}", "cyan"))
print(colorText(f" Duration: {duration_label}", "cyan"))
print(colorText("=" * 60, "white"))
print(colorText("\n💡 Next Steps:", "yellow"))
print(colorText(" • Agents have been moved to audit policies", "white"))
print(colorText(" • Local approvals are active", "white"))
print(colorText("\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}",
f" Agents will return to enforcement after {duration_label}",
"white",
)
)
+23 -23
View File
@@ -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),
)
@@ -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"))
+4 -4
View File
@@ -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
@@ -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}"
)
+18 -6
View File
@@ -1,14 +1,26 @@
# Core TUI dependencies
textual==6.5.0
# API and data handling
Requests==2.32.5
pandas==2.3.3
numpy==2.3.4
# Database
pymongo==4.15.3
# Security and encryption
cryptography==46.0.3
keyring==25.6.0
numpy==2.3.4
pandas==2.3.3
pymongo==4.15.3
# Environment management
python-dotenv==1.2.1
Requests==2.32.5
textual==6.5.0
# Utilities
tqdm==4.67.1
urllib3==2.5.0
pyperclip==1.11.0
# Custom/Private packages
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
airlock_libs==5.0.0
airlock_libs==7.0.0
+8
View File
@@ -351,6 +351,14 @@ class AirlockAPIWrapper:
result = self._post("/v1/getexechistory", payload)
return result["response"]["exechistory"]
def server_logs(self, checkpoint: str | None = None) -> str:
"""Retrieves Server Activity History Logs."""
payload = {}
if checkpoint is not None:
payload["checkpoint"] = checkpoint
result = self._post("/v1/logging/svractivities?checkpoint", payload)
return result["response"]["svractivities"]
"""
from services.API import AirlockAPIWrapper
+7 -7
View File
@@ -38,9 +38,9 @@ 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–365): ",
value_type=int,
valid_range=(1, 150),
valid_range=(1, 365),
)
if not agents or not history_days:
@@ -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",
)
)
@@ -265,7 +265,7 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
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):
@@ -302,7 +302,7 @@ 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_system_json("POLICY_MAP_ENF_AUD", "{}")
+1 -16
View File
@@ -27,9 +27,8 @@ import tqdm
from models.policy import Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import get_system_json
from utils.setup import get_base_directory
from utils.utils import areYouSure, colorText, get_sanitized_input
from utils.utils import colorText
logger = logging.getLogger(__name__)
@@ -233,17 +232,3 @@ def skipback(days):
hex_timestamp = format(timestamp, "08x")
objectid_hex = hex_timestamp + "0000000000000000"
return ObjectId(objectid_hex)
def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
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")
def confirmUpdateAfromE(api: AirlockAPIWrapper):
areYouSure()
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
if confirmation.strip() == "I AGREE":
updateAuditPoliciesFromEnforcementPolices(api)
+8 -4
View File
@@ -184,7 +184,7 @@ class Selector:
print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
return selected
elif mode == "e":
print(colorText(f"🚫 Excluded {len(selected)} item(s).", "yellow"))
print(colorText(f"👫 Excluded {len(selected)} item(s).", "yellow"))
return [item for item in items if item not in selected]
else:
print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
@@ -279,7 +279,9 @@ class Selector:
df = df[columns]
items = [row for _, row in df.iterrows()]
label_func = lambda row: str(row.to_dict())
def label_func(row):
return str(row.to_dict())
result = Selector._select_from_list(
items,
@@ -311,7 +313,9 @@ class Selector:
df = df[columns]
items = df.to_dict("records")
label_func = lambda row: " | ".join(str(row[col]) for col in df.columns)
def label_func(row):
return " | ".join(str(row[col]) for col in df.columns)
# Show rows first
print(colorText(header, "cyan"))
@@ -346,7 +350,7 @@ class Selector:
print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
return [pd.Series(row) for row in selected]
elif mode == "e":
print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow"))
print(colorText(f"👫 Excluded {len(selected)} row(s).", "yellow"))
return [pd.Series(row) for row in items if row not in selected]
else:
print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
+78 -25
View File
@@ -29,6 +29,49 @@ from utils.configmanager import (
)
class TextualNotificationHandler(logging.Handler):
"""
Custom logging handler that sends ERROR, WARNING, and CRITICAL logs
to Textual toast notifications.
"""
def __init__(self, app):
super().__init__()
self.app = app
def emit(self, record):
try:
# Only handle ERROR, WARNING, and CRITICAL
if record.levelno >= logging.WARNING:
# Format the message
msg = self.format(record)
# Map log levels to Textual severity
severity_map = {
logging.WARNING: "warning",
logging.ERROR: "error",
logging.CRITICAL: "error",
}
severity = severity_map.get(record.levelno, "information")
# Send to Textual notification
# Use call_from_thread if logging from non-main thread
try:
self.app.notify(msg, severity=severity, timeout=5)
except Exception:
# If we're not on the main thread, schedule it
try:
self.app.call_from_thread(
self.app.notify, msg, severity=severity, timeout=5
)
except Exception:
# Silently fail to avoid breaking the logging system
pass
except Exception:
# Silently fail to avoid breaking the logging system
pass
def get_base_directory() -> Path:
system = platform.system()
home = Path.home()
@@ -44,38 +87,31 @@ def configure_logging(log_dir: Path, log_level: str = "INFO"):
log_file = log_dir / "Loxide.log"
config = {
"version": 1, # Required key for dictConfig format version
"disable_existing_loggers": False, # Keeps existing loggers active
"version": 1,
"disable_existing_loggers": False,
"formatters": {
"detailed": {
"format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
# Includes timestamp, logger name, level, and message
},
"simple": {
"format": "%(levelname)s - %(message)s"
# Minimal format for console output
},
"simple": {"format": "%(levelname)s - %(message)s"},
"toast": {"format": "%(name)s: %(message)s"}, # Simpler format for toasts
},
"handlers": {
"file": {
"class": "logging.handlers.TimedRotatingFileHandler",
"filename": str(log_file),
"when": "midnight", # Rotate logs at midnight
"interval": 1, # Every 1 day
"backupCount": 7, # Keep 7 days of logs
"encoding": "utf-8", # Ensure UTF-8 encoding
"level": "DEBUG", # Always log DEBUG and above to file
"formatter": "detailed", # Use detailed format
},
"console": {
"class": "logging.StreamHandler",
"level": log_level.upper(), # System-configured level for console
"formatter": "simple", # Use simple format
"when": "midnight",
"interval": 1,
"backupCount": 7,
"encoding": "utf-8",
"level": "DEBUG",
"formatter": "detailed",
},
# REMOVED console handler - it interferes with Textual TUI
},
"root": {
"level": "DEBUG", # Root logger level
"handlers": ["file", "console"], # Attach both handlers
"level": "DEBUG",
"handlers": ["file"], # Only use file handler, not console
},
}
@@ -96,6 +132,18 @@ def configure_logging(log_dir: Path, log_level: str = "INFO"):
logging.config.dictConfig(config)
logging.getLogger().debug("✅ Logging configured.")
# Return a function to attach the notification handler once the app is created
def attach_notification_handler(app):
"""Attach the Textual notification handler to the root logger."""
handler = TextualNotificationHandler(app)
handler.setLevel(logging.WARNING) # Only WARNING and above
formatter = logging.Formatter("%(name)s: %(message)s")
handler.setFormatter(formatter)
logging.getLogger().addHandler(handler)
logging.getLogger().debug("✅ Textual notification handler attached.")
return attach_notification_handler
def setup():
"""
@@ -105,6 +153,9 @@ def setup():
3. Load user config (mutable, from user_config.json)
4. Configure logging
5. Set up .env with WORKING_DIR only
Returns:
attach_notification_handler: Function to attach notification handler to TUI app
"""
base_dir = get_base_directory()
dirs = {
@@ -118,14 +169,14 @@ def setup():
logging.debug(f"{name.capitalize()} directory ensured at: {path}")
# Load system config (immutable)
system_config = load_system_config()
load_system_config()
# Configure logging with system-defined log level
log_level = get_system_value("LOG_LEVEL", str, "INFO")
configure_logging(dirs["logs"], log_level)
attach_handler = configure_logging(dirs["logs"], log_level)
# Load user config (mutable)
user_config = load_user_config(dirs["config"])
load_user_config(dirs["config"])
# Set up .env file - ONLY for WORKING_DIR (runtime-configurable value)
env_path = base_dir / ".env"
@@ -145,7 +196,6 @@ def setup():
"Approved": [],
"Needs_Review": ["Review_First", "Review_Second", "HTML"],
"Preflight": ["HTML"],
"Archived": [],
}
for folder_name, subfolders in folders_structure.items():
@@ -155,6 +205,9 @@ def setup():
for subfolder in subfolders:
subfolder_path = folder_path / subfolder
subfolder_path.mkdir(parents=True, exist_ok=True)
logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
logging.debug(f"'{subfolder}' subfolder created at: {subfolder_path}")
logging.info("✅ Setup complete")
# Return the attach handler function
return attach_handler