Compare commits

...

25 Commits

Author SHA1 Message Date
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
brotoskyj f3c1d97d28 Implemented Threaded Rust Process
Build Library / Build Library (push) Successful in 6m18s
It still locks the GIL, but this is the ground work for the non blocking GUI. Python objects are converted to Rust objects in the beginning of the function call so they can be sent to threads safely.
2025-12-02 10:53:25 -05:00
brotoskyj 99d7b5f74e Merge remote-tracking branch 'refs/remotes/origin/RustImplementation' into RustImplementation 2025-11-25 15:30:07 -05:00
brotoskyj 6327adeabf Refactor of Loxide Libs
Removed spaces, imports, and converters from code
2025-11-25 15:29:55 -05:00
brotoskyj b289e9324c Slight Changes
Build Library / Build Library (push) Successful in 5m15s
Refactor of Loxide Libs
Removed spaces, imports, and converters from code
2025-11-25 15:27:16 -05:00
brotoskyj a80c2ca1e1 Features
Build Library / Build Library (push) Successful in 4m46s
closes #34
Removed a lot of unwraps, most remaining unwraps will likely stay in the code, as it will be expected behavior to panic and crash rather than a system level abort/exit event
2025-11-24 11:36:50 -05:00
brotoskyj dd206eb272 Feature
closes #37
2025-11-24 10:35:38 -05:00
brotoskyj 87ab1e3b28 Telemetry Config Change
Build Library / Build Library (push) Successful in 4m54s
Added implementation instead of separate function to load Telemetry configuration
2025-11-21 14:00:09 -05:00
brotoskyj 303ecd8368 Removed a lot of unwraps
Build Library / Build Library (push) Successful in 5m2s
Removing the unwraps now has better error handling and will cause the program to crash with crash dumps instead of panic!
2025-11-21 12:20:11 -05:00
39 changed files with 5878 additions and 1748 deletions
+1 -2
View File
@@ -30,7 +30,7 @@ from services.API import AirlockAPIWrapper
from services.security import getAPI from services.security import getAPI
from TUI.TUI import run_Loxide from TUI.TUI import run_Loxide
from utils.configmanager import get_system_value from utils.configmanager import get_system_value
from utils.setup import get_base_directory, setup from utils.setup import setup
from utils.utils import irtang from utils.utils import irtang
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
@@ -40,7 +40,6 @@ def main():
irtang() irtang()
# Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
setup() setup()
base_dir = get_base_directory()
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
try: try:
@@ -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 __future__ import annotations
import logging import logging
@@ -150,14 +164,11 @@ class AllowlistSelectionWidget(Static):
# Action buttons at bottom # Action buttons at bottom
with Horizontal(id="action_buttons"): 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.add_btn = Button(" Add to Allowlist", id="add_to_allowlist_btn")
self.back_btn.styles.width = "50%" self.add_btn.styles.width = "100%"
self.add_btn.styles.width = "50%"
self.add_btn.disabled = True # Disabled until allowlist selected self.add_btn.disabled = True # Disabled until allowlist selected
yield self.back_btn
yield self.add_btn yield self.add_btn
async def on_mount(self) -> None: async def on_mount(self) -> None:
@@ -380,9 +391,9 @@ class AllowlistSelectionWidget(Static):
if found_col: if found_col:
self.hash_column = 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: else:
preview_lines.append("⚠️ **No hash column found**\n") preview_lines.append(" **No hash column found**\n")
preview_lines.append("Available columns:\n") preview_lines.append("Available columns:\n")
for col in self.selected_data.columns: for col in self.selected_data.columns:
if col != "_row_id": if col != "_row_id":
@@ -464,7 +475,7 @@ class AllowlistSelectionWidget(Static):
self.selected_allowlist = self.allowlists[actual_allowlist_index] self.selected_allowlist = self.allowlists[actual_allowlist_index]
self.add_btn.disabled = False self.add_btn.disabled = False
self.add_btn.label = ( 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 # Update preview with selection
@@ -523,11 +534,6 @@ class AllowlistSelectionWidget(Static):
btn = getattr(event, "button", None) or getattr(event, "sender", None) btn = getattr(event, "button", None) or getattr(event, "sender", None)
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", 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": if btn is self.refresh_btn or btn_id == "refresh_allowlists_btn":
await self.load_allowlists() await self.load_allowlists()
event.stop() event.stop()
@@ -553,7 +559,7 @@ class AllowlistSelectionWidget(Static):
try: try:
# Disable button during operation # Disable button during operation
self.add_btn.disabled = True self.add_btn.disabled = True
self.add_btn.label = "Adding hashes..." self.add_btn.label = "Adding hashes..."
# Call API to add hashes # Call API to add hashes
app_id = self.selected_allowlist.get("applicationid") 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) result = self.api.hash_add_to_allowlist(app_id, self.hashes_to_add)
logger.debug(f"Hash adding api call: {result}")
# Success notification # Success notification
self.app.notify( 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", title="Success",
severity="information", severity="information",
timeout=5, timeout=5,
@@ -575,7 +581,7 @@ class AllowlistSelectionWidget(Static):
# Update preview to show success # Update preview to show success
self.preview_area.text = ( self.preview_area.text = (
f"## SUCCESS\n\n" f"## SUCCESS\n\n"
f"Added **{len(self.hashes_to_add)} hashes** to allowlist:\n" f"Added **{len(self.hashes_to_add)} hashes** to allowlist:\n"
f"**{allowlist_name}** (ID: {app_id})\n\n" f"**{allowlist_name}** (ID: {app_id})\n\n"
f"### Operation Details:\n" f"### Operation Details:\n"
@@ -586,13 +592,13 @@ class AllowlistSelectionWidget(Static):
) )
# Change button to "Done" # 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 self.add_btn.disabled = True
except Exception as exc: except Exception as exc:
logger.exception(f"Failed to add hashes to allowlist: {exc}") logger.exception(f"Failed to add hashes to allowlist: {exc}")
self.app.notify( self.app.notify(
f"Failed to add hashes: {str(exc)}", f"Failed to add hashes: {str(exc)}",
title="Error", title="Error",
severity="error", severity="error",
timeout=10, timeout=10,
@@ -600,7 +606,7 @@ class AllowlistSelectionWidget(Static):
# Re-enable button # Re-enable button
self.add_btn.disabled = False self.add_btn.disabled = False
self.add_btn.label = "Retry Add to Allowlist" self.add_btn.label = "Retry Add to Allowlist"
class AllowlistSelectionScreen(Screen): class AllowlistSelectionScreen(Screen):
@@ -609,9 +615,9 @@ class AllowlistSelectionScreen(Screen):
""" """
BINDINGS = [ BINDINGS = [
Binding("b", "back", "Back"), Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
Binding("r", "refresh", "Refresh Allowlists"), Binding("r", "refresh", "Refresh Allowlists"),
Binding("enter", "confirm", "Add to Allowlist"),
] ]
def __init__( def __init__(
@@ -641,10 +647,15 @@ class AllowlistSelectionScreen(Screen):
yield self.widget yield self.widget
yield Footer() yield Footer()
async def action_back(self) -> None: async def action_go_back(self) -> None:
"""Go back to previous screen.""" """Go back to previous screen."""
await self.app.pop_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: async def action_refresh(self) -> None:
"""Refresh the allowlists.""" """Refresh the allowlists."""
if hasattr(self, "widget") and self.widget: if hasattr(self, "widget") and self.widget:
+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 __future__ import annotations
from datetime import datetime from datetime import datetime
@@ -11,7 +25,7 @@ from textual.containers import Horizontal, Vertical
from textual.screen import Screen from textual.screen import Screen
from textual.widgets import Button, DataTable, Footer, Header, Static 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 from utils.configmanager import load_env
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -31,8 +45,7 @@ class OTPActivitiesWidget(Static):
""" """
Reusable widget that contains the sessions table (left) and an Activity Preview (right). 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 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 with Continue button.
currently-loaded activities.
""" """
DEFAULT_CSS = """ DEFAULT_CSS = """
@@ -85,15 +98,10 @@ class OTPActivitiesWidget(Static):
with Vertical(id="activity_preview_container"): with Vertical(id="activity_preview_container"):
self.activities_table = DataTable(id="activity_preview_table") self.activities_table = DataTable(id="activity_preview_table")
yield self.activities_table yield self.activities_table
# Buttons area at the bottom (Back, Continue) # Button area at the bottom (Continue)
with Horizontal(id="activity_buttons"): 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") self.continue_btn = Button("Continue", id="activity_continue_btn")
# Stretch buttons nicely self.continue_btn.styles.width = "100%"
self.back_btn.styles.width = "50%"
self.continue_btn.styles.width = "50%"
yield self.back_btn
yield self.continue_btn yield self.continue_btn
async def on_mount(self) -> None: async def on_mount(self) -> None:
@@ -122,7 +130,7 @@ class OTPActivitiesWidget(Static):
async def on_button_pressed(self, event) -> None: # type: ignore[override] 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 # Try to resolve the button object from the event
btn = ( btn = (
@@ -136,18 +144,12 @@ class OTPActivitiesWidget(Static):
or getattr(event, "button_id", None) or getattr(event, "button_id", None)
or getattr(event, "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 ---- # ---- Continue ----
if btn is self.continue_btn or btn_id == getattr(self.continue_btn, "id", None): 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: if self._activities_df is None or self._activities_df.empty:
logger.info("Continue pressed but no activities loaded.") logger.info("Continue pressed but no activities loaded.")
await self.post_message( self.app.notify(
Static("No activities loaded to continue with.") "No activities loaded to continue with.", severity="warning"
) )
return return
# Copy activities DataFrame to pass to new screen # Copy activities DataFrame to pass to new screen
@@ -463,7 +465,7 @@ class OTPActivitiesWidget(Static):
try: try:
self._activities_df.to_csv(file_path, index=False) self._activities_df.to_csv(file_path, index=False)
logger.info("Exported activities to %s", file_path) 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: except Exception as exc:
logger.exception("Failed to export activities to %s: %s", file_path, exc) logger.exception("Failed to export activities to %s: %s", file_path, exc)
await self.post_message(Static("Failed to export activities; check logs.")) await self.post_message(Static("Failed to export activities; check logs."))
@@ -472,7 +474,7 @@ class OTPActivitiesWidget(Static):
class ActivityDetailWidget(Static): class ActivityDetailWidget(Static):
""" """
Interactive widget for Activity Detail screen. 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. Now includes Select All/None and Add to Allowlist functionality.
""" """
@@ -529,12 +531,10 @@ class ActivityDetailWidget(Static):
# Original buttons at bottom # Original buttons at bottom
with Horizontal(id="detail_buttons"): with Horizontal(id="detail_buttons"):
self.detail_back_btn = Button("Back", id="detail_back_btn")
self.add_allowlist_btn = Button( 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.add_allowlist_btn
yield self.detail_back_btn
async def on_mount(self) -> None: async def on_mount(self) -> None:
await self._build_table(rebuild=True) await self._build_table(rebuild=True)
@@ -547,12 +547,12 @@ class ActivityDetailWidget(Static):
# Update button labels with count # Update button labels with count
count = len(self.selected_row_ids) count = len(self.selected_row_ids)
total = len(self.activities_df) len(self.activities_df)
if has_selection: 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: 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: async def _build_table(self, rebuild: bool = True) -> None:
"""Rebuild the DataTable. If rebuild=False, only refresh rows.""" """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)) vals.append("" if pd.isna(v) else str(v))
# Check if this row is selected # 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 # Add row to table
row_key = self.detail_table.add_row(checkbox, *vals) 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 self.detail_table.update_cell(row_key, "select", "") # Unchecked
else: else:
self.selected_row_ids.add(row_id) 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() self._update_button_states()
@@ -652,19 +652,13 @@ class ActivityDetailWidget(Static):
logger.exception("Failed to sort by column %s: %s", column_key, exc) logger.exception("Failed to sort by column %s: %s", column_key, exc)
return return
# ✅ Only refresh rows, not columns # Only refresh rows, not columns
await self._build_table(rebuild=False) await self._build_table(rebuild=False)
async def on_button_pressed(self, event) -> None: async def on_button_pressed(self, event) -> None:
btn = getattr(event, "button", None) or getattr(event, "sender", None) btn = getattr(event, "button", None) or getattr(event, "sender", None)
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", 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": if btn is self.add_allowlist_btn or btn_id == "add_allowlist_btn":
await self._open_allowlist_screen() await self._open_allowlist_screen()
return return
@@ -676,7 +670,7 @@ class ActivityDetailWidget(Static):
# Update all checkboxes in the table # Update all checkboxes in the table
for row_key, row_id in self.row_key_to_id.items(): 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() self._update_button_states()
logger.info(f"Selected all {len(self.selected_row_ids)} rows") logger.info(f"Selected all {len(self.selected_row_ids)} rows")
@@ -688,7 +682,7 @@ class ActivityDetailWidget(Static):
# Update all checkboxes in the table # Update all checkboxes in the table
for row_key, row_id in self.row_key_to_id.items(): 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() self._update_button_states()
logger.info("Cleared all selections") logger.info("Cleared all selections")
@@ -730,14 +724,12 @@ class ActivityDetailWidget(Static):
async def _export_detail_activities(self) -> None: async def _export_detail_activities(self) -> None:
if self.activities_df is None or self.activities_df.empty: if self.activities_df is None or self.activities_df.empty:
logger.info("No activities to export.") logger.info("No activities to export.")
await self.mount( await self.mount(Static("No activities to export.", classes="notification"))
Static("❌ No activities to export.", classes="notification")
)
return return
if not self.selected_row_ids: if not self.selected_row_ids:
logger.info("No rows selected for export.") logger.info("No rows selected for export.")
await self.mount( await self.mount(
Static("❌ No rows selected for export.", classes="notification") Static("No rows selected for export.", classes="notification")
) )
return return
try: try:
@@ -750,7 +742,7 @@ class ActivityDetailWidget(Static):
logger.info("Exported selected activities to %s", file_path) logger.info("Exported selected activities to %s", file_path)
await self.mount( await self.mount(
Static( Static(
f"Exported selected activities to: {filename}", f"Exported selected activities to: {filename}",
classes="notification", classes="notification",
) )
) )
@@ -758,12 +750,12 @@ class ActivityDetailWidget(Static):
logger.exception("Failed to export detail activities: %s", exc) logger.exception("Failed to export detail activities: %s", exc)
await self.mount( await self.mount(
Static( Static(
"❌ Failed to export activities; check logs.", "¢ Failed to export activities; check logs.",
classes="notification", classes="notification",
) )
) )
# ✅ Helper methods # Helper methods
def get_selected_data(self) -> pd.DataFrame: def get_selected_data(self) -> pd.DataFrame:
"""Return a DataFrame of the selected rows.""" """Return a DataFrame of the selected rows."""
if not self.selected_row_ids: if not self.selected_row_ids:
@@ -795,7 +787,8 @@ class ActivityDetailScreen(Screen):
""" """
BINDINGS = [ BINDINGS = [
Binding("b", "back", "Back"), Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
Binding("e", "export", "Export"), Binding("e", "export", "Export"),
Binding("a", "select_all", "Select All"), Binding("a", "select_all", "Select All"),
Binding("n", "select_none", "Select None"), Binding("n", "select_none", "Select None"),
@@ -819,11 +812,16 @@ class ActivityDetailScreen(Screen):
yield self.widget yield self.widget
yield Footer() yield Footer()
async def action_back(self) -> None: async def action_go_back(self) -> None:
try: try:
await self.app.pop_screen() await self.app.pop_screen()
except Exception: 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: async def action_export(self) -> None:
# Delegate to widget export helper # Delegate to widget export helper
@@ -851,9 +849,10 @@ class OTPActivitiesScreen(Screen):
""" """
BINDINGS = [ BINDINGS = [
Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
Binding("r", "refresh_sessions", "Refresh Sessions"), Binding("r", "refresh_sessions", "Refresh Sessions"),
Binding("e", "export_activities", "Export activities"), Binding("e", "export_activities", "Export activities"),
Binding("q", "quit", "Quit"),
] ]
def compose(self) -> ComposeResult: def compose(self) -> ComposeResult:
@@ -875,6 +874,15 @@ class OTPActivitiesScreen(Screen):
else: else:
await self.widget.load_sessions_from_api(api) 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 # Simple actions bound to keys
async def action_refresh_sessions(self) -> None: async def action_refresh_sessions(self) -> None:
api = getattr(self.app, "api", None) api = getattr(self.app, "api", None)
@@ -884,10 +892,6 @@ class OTPActivitiesScreen(Screen):
logger.info("Refreshing OTP sessions via API.") logger.info("Refreshing OTP sessions via API.")
await self.widget.load_sessions_from_api(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: # 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: async def fetch_activities_for_otpid(self, otpid, hostname=None) -> None:
api = getattr(self.app, "api", 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 import logging
from textual.app import ComposeResult from textual.app import ComposeResult
from textual.binding import Binding
from textual.screen import Screen 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__) logger = logging.getLogger(__name__)
@@ -42,6 +44,11 @@ class PolicySelectorScreen(Screen):
agent_move_operations: Reference to the parent AgentMoveOperations widget. agent_move_operations: Reference to the parent AgentMoveOperations widget.
""" """
BINDINGS = [
Binding("escape", "go_back", "Back"),
Binding("q", "main_menu", "Main Menu"),
]
CSS = """ CSS = """
Screen { Screen {
layout: vertical; layout: vertical;
@@ -68,7 +75,18 @@ class PolicySelectorScreen(Screen):
def compose(self) -> ComposeResult: def compose(self) -> ComposeResult:
"""Create the PolicySelector widget.""" """Create the PolicySelector widget."""
yield Header(show_clock=True)
yield PolicySelector(self.policies) 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( def on_policy_selector_policy_selected(
self, message: PolicySelector.PolicySelected self, message: PolicySelector.PolicySelected
@@ -34,12 +34,12 @@ from textual.app import ComposeResult
from textual.containers import Horizontal, Vertical from textual.containers import Horizontal, Vertical
from textual.reactive import reactive from textual.reactive import reactive
from textual.screen import Screen 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 models.policy import Policy
from services.API import AirlockAPIWrapper from services.API import AirlockAPIWrapper
from services.policyhandler import getPolicyInfo from services.policyhandler import getPolicyInfo
from TUI.policyselector import PolicySelector from TUI.Widgets.policyselector import PolicySelector
from utils.configmanager import load_env from utils.configmanager import load_env
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -51,16 +51,17 @@ class QuietAgentWorkflowScreen(Screen):
This screen provides a multi-step workflow: This screen provides a multi-step workflow:
1. Select initial policy to analyze 1. Select initial policy to analyze
2. View categorized agents (enforce ready vs. non-enforce ready) 2. Configure analysis parameters (history period and quiet time period)
3. Select target policies for each category 3. View categorized agents (enforce ready vs. non-enforce ready)
4. Execute agent migrations 4. Select target policies for each category
5. Execute agent migrations
Attributes: Attributes:
api (AirlockAPIWrapper): API wrapper for Airlock operations api (AirlockAPIWrapper): API wrapper for Airlock operations
policies (List[Policy]): List of all available policies policies (List[Policy]): List of all available policies
selected_policy (Optional[Policy]): The initially selected policy to analyze selected_policy (Optional[Policy]): The initially selected policy to analyze
history_days (int): Number of days of history to pull (default: 150) 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) 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 agents_df (Optional[pd.DataFrame]): DataFrame of all agents with analysis results
enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents ready for enforcement 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 non_enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents not ready for enforcement
@@ -69,6 +70,7 @@ class QuietAgentWorkflowScreen(Screen):
BINDINGS = [ BINDINGS = [
("escape", "go_back", "Back"), ("escape", "go_back", "Back"),
("q", "main_menu", "Main Menu"),
] ]
workflow_stage = reactive("select_policy") # Tracks current workflow stage workflow_stage = reactive("select_policy") # Tracks current workflow stage
@@ -85,7 +87,7 @@ class QuietAgentWorkflowScreen(Screen):
self.api = api self.api = api
self.policies = policies self.policies = policies
self.selected_policy: Optional[Policy] = None 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.quiet_days = 45 # Default value
self.agents_df: Optional[pd.DataFrame] = None self.agents_df: Optional[pd.DataFrame] = None
self.enforce_ready_df: Optional[pd.DataFrame] = None self.enforce_ready_df: Optional[pd.DataFrame] = None
@@ -96,10 +98,10 @@ class QuietAgentWorkflowScreen(Screen):
def compose(self) -> ComposeResult: def compose(self) -> ComposeResult:
"""Build the UI layout for the workflow screen.""" """Build the UI layout for the workflow screen."""
# Include Header and Footer like other standalone screens # Include Header and Footer like other standalone screens
yield Header(show_clock=True, icon="") yield Header(show_clock=True, icon="")
# Title area # 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) title.styles.margin = (0, 0, 0, 1)
yield title yield title
@@ -129,14 +131,14 @@ class QuietAgentWorkflowScreen(Screen):
stage_messages = { stage_messages = {
"select_policy": "Step 1: Select Policy to Analyze", "select_policy": "Step 1: Select Policy to Analyze",
"select_quiet_days": "Step 2: Select Quiet Time Period", "select_history_days": "Step 2: Configure Analysis Parameters",
"analyzing": "📊 Analyzing agent activity...", "analyzing": "Analyzing agent activity...",
"view_results": "Step 3: Review Categorized Agents", "view_results": "Step 3: Review Categorized Agents",
"select_enforce_target": "Step 4: Select Target Policy for Enforce Ready 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", "select_non_enforce_target": "Step 5: Select Target Policy for Non-Enforce Ready Agents",
"confirm_migration": "Step 6: Confirm and Execute Migration", "confirm_migration": "Step 6: Confirm and Execute Migration",
"executing": "Executing agent migrations...", "executing": "Executing agent migrations...",
"complete": "Migration Complete", "complete": "Migration Complete",
} }
status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage")) status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage"))
@@ -160,7 +162,7 @@ class QuietAgentWorkflowScreen(Screen):
# Initial policy selection for analysis # Initial policy selection for analysis
self.selected_policy = message.policy self.selected_policy = message.policy
logger.info(f"Selected policy for analysis: {self.selected_policy.name}") 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": elif self.workflow_stage == "select_enforce_target":
# Target policy selection for enforce ready agents # Target policy selection for enforce ready agents
self.enforce_ready_target_policy = message.policy self.enforce_ready_target_policy = message.policy
@@ -176,64 +178,170 @@ class QuietAgentWorkflowScreen(Screen):
) )
self._show_migration_confirmation() self._show_migration_confirmation()
def _show_quiet_days_selection(self) -> None: def _show_history_days_selection(self) -> None:
"""Show the quiet days selection screen.""" """Show the history days and quiet days selection screen."""
self.workflow_stage = "select_quiet_days" self.workflow_stage = "select_history_days"
content = self.query_one("#content_area", Vertical) content = self.query_one("#content_area", Vertical)
content.remove_children() content.remove_children()
# Create info text # Create info text
info_widget = Static( info_widget = Static(
f"Policy Selected: {self.selected_policy.name}\n\n" f"Policy Selected: {self.selected_policy.name}\n\n"
f"History Period: {self.history_days} days\n\n" "Configure Analysis Parameters:",
"Select quiet time period (days without untrusted execution):", id="analysis_params_info",
id="quiet_days_info",
) )
info_widget.styles.margin = (0, 0, 2, 0) info_widget.styles.margin = (0, 0, 2, 0)
content.mount(info_widget) content.mount(info_widget)
# Create button container and mount it first # Create input container
button_container = Vertical(id="quiet_days_buttons") input_container = Vertical(id="analysis_params_input_container")
button_container.styles.height = "auto" input_container.styles.height = "auto"
content.mount(button_container) content.mount(input_container)
# Now add buttons to the mounted container # History days label
for days in [15, 30, 45, 60]: history_label = Static("History Period (days of execution history to pull):")
btn = Button( history_label.styles.margin = (0, 0, 1, 0)
f"{days} days {'(Default)' if days == 45 else ''}", input_container.mount(history_label)
id=f"quiet_days_{days}",
classes="quiet_day_btn", # 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%" history_input.styles.width = "50"
btn.styles.margin = (0, 0, 1, 0) history_input.styles.margin = (0, 0, 2, 0)
button_container.mount(btn) input_container.mount(history_input)
back_btn = Button("← Back", id="back_to_policy_selection") # Quiet days label
back_btn.styles.width = "100%" quiet_label = Static(
back_btn.styles.margin = (2, 0, 0, 0) "Quiet Time Period (days without execution to be considered quiet):"
button_container.mount(back_btn) )
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: def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button press events.""" """Handle button press events."""
button_id = event.button.id button_id = event.button.id
# Quiet days selection buttons # Analysis parameters submit button
if button_id and button_id.startswith("quiet_days_"): if button_id == "analysis_params_submit":
days = int(button_id.split("_")[-1]) self._validate_and_submit_history_days()
self.quiet_days = days
logger.info(f"Selected quiet days: {days}")
self._start_analysis()
return return
# Navigation buttons # 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": if button_id == "select_enforce_target_btn":
self._show_enforce_target_selection() self._show_enforce_target_selection()
return return
@@ -270,46 +378,44 @@ class QuietAgentWorkflowScreen(Screen):
self._show_policy_selection() self._show_policy_selection()
return 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: def _start_analysis(self) -> None:
"""Start the agent activity analysis.""" """Start the agent activity analysis."""
self.workflow_stage = "analyzing" # Show notification that analysis is starting
content = self.query_one("#content_area", Vertical)
content.remove_children()
# Show analyzing message with detailed steps
analyzing_msg = Static(
f"📊 Analyzing Agent Activity\n"
f"{'=' * 50}\n\n"
f"Policy: {self.selected_policy.name}\n"
f"History Period: {self.history_days} days\n"
f"Quiet Threshold: {self.quiet_days} days\n\n"
f"Progress:\n"
f"⏳ Step 1/4: Fetching agents from policy...\n"
f"⏱️ Step 2/4: Pulling execution history (this may take a moment)...\n"
f"⏱️ Step 3/4: Analyzing activity patterns...\n"
f"⏱️ Step 4/4: Categorizing agents...\n\n"
f"Please wait - this operation cannot be cancelled.",
id="analyzing_message",
)
analyzing_msg.styles.margin = (2, 1)
content.mount(analyzing_msg)
# Show notification
self.app.notify( self.app.notify(
"Starting analysis - this may take several minutes for large policies", "Starting analysis - this may take several minutes for large policies",
severity="information", severity="information",
timeout=5, timeout=5,
) )
# Perform the analysis asynchronously # Clear the screen to provide a blank canvas for Rust progress output
self.call_later(self._perform_analysis) # (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.""" """Perform the actual agent activity analysis."""
try: try:
# Update status: Fetching agents
self._update_analysis_status("Step 1/4: Fetching agents from policy...")
# Get agents in the selected policy # Get agents in the selected policy
agents = self.api.agents_find_by_group(self.selected_policy.groupid) agents = self.api.agents_find_by_group(self.selected_policy.groupid)
@@ -322,32 +428,11 @@ class QuietAgentWorkflowScreen(Screen):
self._show_policy_selection() self._show_policy_selection()
return 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) # Get execution history (this shows progress bars in terminal via airlock_libs)
policy_exec_history = getPolicyInfo( policy_exec_history = getPolicyInfo(
self.api, self.selected_policy, [1, 2, 6, 7], self.history_days 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: if policy_exec_history.empty:
logger.info( logger.info(
"No execution history found for the selected policy and time range." "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 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 # Sort agents
agents = agents.sort_values( agents = agents.sort_values(
by=["execution_count", "hostname"], ascending=[True, True] by=["execution_count", "hostname"], ascending=[True, True]
@@ -405,8 +487,8 @@ class QuietAgentWorkflowScreen(Screen):
self.agents_df = agents self.agents_df = agents
# Categorize agents into DataFrames # Categorize agents into DataFrames
self.enforce_ready_df = agents[agents["enforce_ready"] == True].copy() self.enforce_ready_df = agents[agents["enforce_ready"]].copy()
self.non_enforce_ready_df = agents[agents["enforce_ready"] == False].copy() self.non_enforce_ready_df = agents[~agents["enforce_ready"]].copy()
logger.info( logger.info(
f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, " 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.app.notify(f"Analysis failed: {str(e)}", severity="error", timeout=5)
self._show_policy_selection() 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: def _show_results(self) -> None:
"""Show the categorized results.""" """Show the categorized results."""
self.workflow_stage = "view_results" self.workflow_stage = "view_results"
@@ -469,9 +530,9 @@ class QuietAgentWorkflowScreen(Screen):
summary = Static( summary = Static(
f"Analysis Results for: {self.selected_policy.name}\n\n" f"Analysis Results for: {self.selected_policy.name}\n\n"
f"📊 Total Agents: {total_agents}\n" f"Total Agents: {total_agents}\n"
f"Enforce Ready: {ready_count} ({ready_percentage:.1f}%)\n" f"Enforce Ready: {ready_count} ({ready_percentage:.1f}%)\n"
f"Not Ready: {not_ready_count} ({100 - ready_percentage:.1f}%)\n\n" f"Not Ready: {not_ready_count} ({100 - ready_percentage:.1f}%)\n\n"
f"Quiet Threshold: {self.quiet_days} days\n" f"Quiet Threshold: {self.quiet_days} days\n"
f"History Period: {self.history_days} days", f"History Period: {self.history_days} days",
id="results_summary", id="results_summary",
@@ -500,11 +561,11 @@ class QuietAgentWorkflowScreen(Screen):
non_enforce_btn.styles.margin = (0, 1, 1, 0) non_enforce_btn.styles.margin = (0, 1, 1, 0)
button_container.mount(non_enforce_btn) 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) export_btn.styles.margin = (0, 1, 1, 0)
button_container.mount(export_btn) 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) start_over_btn.styles.margin = (0, 0, 1, 0)
button_container.mount(start_over_btn) button_container.mount(start_over_btn)
@@ -520,7 +581,7 @@ class QuietAgentWorkflowScreen(Screen):
enforce_col.styles.margin = (1, 1, 0, 0) enforce_col.styles.margin = (1, 1, 0, 0)
tables_container.mount(enforce_col) 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_label.styles.margin = (0, 0, 1, 0)
enforce_col.mount(enforce_label) enforce_col.mount(enforce_label)
@@ -548,7 +609,7 @@ class QuietAgentWorkflowScreen(Screen):
non_enforce_col.styles.margin = (1, 0, 0, 1) non_enforce_col.styles.margin = (1, 0, 0, 1)
tables_container.mount(non_enforce_col) 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_label.styles.margin = (0, 0, 1, 0)
non_enforce_col.mount(non_enforce_label) non_enforce_col.mount(non_enforce_label)
@@ -589,7 +650,7 @@ class QuietAgentWorkflowScreen(Screen):
content.mount(policy_selector) content.mount(policy_selector)
# Skip button # 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.width = "50%"
skip_btn.styles.margin = (2, 0, 0, 0) skip_btn.styles.margin = (2, 0, 0, 0)
content.mount(skip_btn) content.mount(skip_btn)
@@ -614,7 +675,7 @@ class QuietAgentWorkflowScreen(Screen):
content.mount(policy_selector) content.mount(policy_selector)
# Skip button # 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.width = "50%"
skip_btn.styles.margin = (2, 0, 0, 0) skip_btn.styles.margin = (2, 0, 0, 0)
content.mount(skip_btn) content.mount(skip_btn)
@@ -627,29 +688,29 @@ class QuietAgentWorkflowScreen(Screen):
# Build confirmation message # Build confirmation message
confirmation_lines = [ confirmation_lines = [
"🔐 Migration Summary\n", "Migration Summary\n",
f"Source Policy: {self.selected_policy.name}\n", f"Source Policy: {self.selected_policy.name}\n",
] ]
if self.enforce_ready_target_policy: if self.enforce_ready_target_policy:
confirmation_lines.append( confirmation_lines.append(
f"\nEnforce Ready Migration:\n" f"\nEnforce Ready Migration:\n"
f"Agents: {len(self.enforce_ready_df)}\n" f"Agents: {len(self.enforce_ready_df)}\n"
f"Target: {self.enforce_ready_target_policy.name}\n" f"Target: {self.enforce_ready_target_policy.name}\n"
) )
if self.non_enforce_ready_target_policy: if self.non_enforce_ready_target_policy:
confirmation_lines.append( confirmation_lines.append(
f"\nNon-Enforce Ready Migration:\n" f"\nNon-Enforce Ready Migration:\n"
f"Agents: {len(self.non_enforce_ready_df)}\n" f"Agents: {len(self.non_enforce_ready_df)}\n"
f"Target: {self.non_enforce_ready_target_policy.name}\n" f"Target: {self.non_enforce_ready_target_policy.name}\n"
) )
if ( if (
not self.enforce_ready_target_policy not self.enforce_ready_target_policy
and not self.non_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 = Static("".join(confirmation_lines), id="migration_confirmation")
confirmation.styles.margin = (1, 1, 2, 1) confirmation.styles.margin = (1, 1, 2, 1)
@@ -662,11 +723,11 @@ class QuietAgentWorkflowScreen(Screen):
content.mount(button_container) content.mount(button_container)
if self.enforce_ready_target_policy or self.non_enforce_ready_target_policy: 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) confirm_btn.styles.margin = (0, 1, 0, 0)
button_container.mount(confirm_btn) 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) button_container.mount(cancel_btn)
def _execute_migration(self) -> None: def _execute_migration(self) -> None:
@@ -677,7 +738,7 @@ class QuietAgentWorkflowScreen(Screen):
# Show executing message # Show executing message
executing_msg = Static( executing_msg = Static(
"Executing agent migrations...\nPlease wait...", "Executing agent migrations...\nPlease wait...",
id="executing_message", id="executing_message",
) )
executing_msg.styles.margin = (2, 1) executing_msg.styles.margin = (2, 1)
@@ -696,7 +757,7 @@ class QuietAgentWorkflowScreen(Screen):
if self.enforce_ready_target_policy: if self.enforce_ready_target_policy:
for idx, row in self.enforce_ready_df.iterrows(): for idx, row in self.enforce_ready_df.iterrows():
try: try:
result = self.api.agent_move( self.api.agent_move(
row["agentid"], self.enforce_ready_target_policy.groupid row["agentid"], self.enforce_ready_target_policy.groupid
) )
successful_migrations.append( successful_migrations.append(
@@ -713,7 +774,7 @@ class QuietAgentWorkflowScreen(Screen):
if self.non_enforce_ready_target_policy: if self.non_enforce_ready_target_policy:
for idx, row in self.non_enforce_ready_df.iterrows(): for idx, row in self.non_enforce_ready_df.iterrows():
try: try:
result = self.api.agent_move( self.api.agent_move(
row["agentid"], self.non_enforce_ready_target_policy.groupid row["agentid"], self.non_enforce_ready_target_policy.groupid
) )
successful_migrations.append( successful_migrations.append(
@@ -749,7 +810,7 @@ class QuietAgentWorkflowScreen(Screen):
) )
results = Static( results = Static(
f"Migration Complete\n\n" f"Migration Complete\n\n"
f"Total Agents Migrated: {len(successful)}\n" f"Total Agents Migrated: {len(successful)}\n"
f"Failed Migrations: {len(failed)}\n" f"Failed Migrations: {len(failed)}\n"
f"Success Rate: {success_rate:.1f}%", f"Success Rate: {success_rate:.1f}%",
@@ -764,7 +825,7 @@ class QuietAgentWorkflowScreen(Screen):
success_container.styles.margin = (0, 1) success_container.styles.margin = (0, 1)
content.mount(success_container) content.mount(success_container)
success_label = Static("Successful Migrations") success_label = Static("Successful Migrations")
success_label.styles.margin = (0, 0, 1, 0) success_label.styles.margin = (0, 0, 1, 0)
success_container.mount(success_label) success_container.mount(success_label)
@@ -785,7 +846,7 @@ class QuietAgentWorkflowScreen(Screen):
failed_container.styles.margin = (2, 1, 0, 1) failed_container.styles.margin = (2, 1, 0, 1)
content.mount(failed_container) content.mount(failed_container)
failed_label = Static("Failed Migrations") failed_label = Static("Failed Migrations")
failed_label.styles.margin = (0, 0, 1, 0) failed_label.styles.margin = (0, 0, 1, 0)
failed_container.mount(failed_label) failed_container.mount(failed_label)
@@ -802,7 +863,7 @@ class QuietAgentWorkflowScreen(Screen):
failed_container.mount(failed_table) failed_container.mount(failed_table)
# Action button # 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.width = "50%"
done_btn.styles.margin = (2, 0, 0, 0) done_btn.styles.margin = (2, 0, 0, 0)
content.mount(done_btn) content.mount(done_btn)
@@ -833,7 +894,7 @@ class QuietAgentWorkflowScreen(Screen):
# Depending on stage, go back to previous stage or exit # Depending on stage, go back to previous stage or exit
if self.workflow_stage in ["select_policy", "view_results", "complete"]: if self.workflow_stage in ["select_policy", "view_results", "complete"]:
self.app.pop_screen() self.app.pop_screen()
elif self.workflow_stage == "select_quiet_days": elif self.workflow_stage == "select_history_days":
self._show_policy_selection() self._show_policy_selection()
elif self.workflow_stage == "select_enforce_target": elif self.workflow_stage == "select_enforce_target":
self._show_results() self._show_results()
@@ -846,3 +907,8 @@ class QuietAgentWorkflowScreen(Screen):
self._show_non_enforce_target_selection() self._show_non_enforce_target_selection()
else: else:
self.app.pop_screen() 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()
+64 -131
View File
@@ -1,6 +1,20 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import logging import logging
import os import os
import sys
from typing import Optional from typing import Optional
import dotenv import dotenv
@@ -19,24 +33,22 @@ from textual.widgets import (
Tabs, Tabs,
) )
from flows.otp import otp_revoke
from flows.prepPolicy import menu_policy_enforce
from models.agent import Agent from models.agent import Agent
from models.policy import Policy from models.policy import Policy
from services.API import AirlockAPIWrapper from services.API import AirlockAPIWrapper
from services.policyhandler import confirmUpdateAfromE from TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
from TUI.agentmoveoperations import AgentMoveOperations from TUI.Screens.otpactivityscreen import OTPActivitiesScreen
from TUI.moveagentworkflowscreen import MoveAgentWorkflowScreen from TUI.Screens.otprevokescreen import OTPRevokeScreen
from TUI.multiagentselector import MultiAgentSelector from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
from TUI.OTP_generate import OTPGenerator from TUI.Screens.policyprepworkflowscreen import PolicyPrepWorkflowScreen
from TUI.otpactivityscreen import OTPActivitiesScreen from TUI.Screens.quietagentworkflowscreen import QuietAgentWorkflowScreen
from TUI.otpworkflowscreen import OTPWorkflowScreen from TUI.Themes.theme_amber_terminal import get_amber_terminal_theme
from TUI.policytreewidget import PolicyTreeWidget from TUI.Themes.theme_retro_terminal import get_retro_terminal_theme
from TUI.quietagentworkflowscreen import QuietAgentWorkflowScreen from TUI.Themes.themeselector import ThemeSelector
from TUI.resultsdisplay import ResultsDisplay from TUI.Widgets.agentmoveoperations import AgentMoveOperations
from TUI.theme_amber_terminal import get_amber_terminal_theme from TUI.Widgets.multiagentselector import MultiAgentSelector
from TUI.theme_retro_terminal import get_retro_terminal_theme from TUI.Widgets.policytreewidget import PolicyTreeWidget
from TUI.themeselector import ThemeSelector from TUI.Widgets.resultsdisplay import ResultsDisplay
from utils.configmanager import get_user_value, load_env, save_user_config from utils.configmanager import get_user_value, load_env, save_user_config
from utils.setup import get_base_directory from utils.setup import get_base_directory
from utils.utils import open_directory from utils.utils import open_directory
@@ -47,7 +59,7 @@ dotenv.load_dotenv()
# GLOBAL STASH # GLOBAL STASH
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
_PENDING_JOB = None _APP_RESTART_REASON = None
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -80,16 +92,15 @@ class MainMenuScreen(Screen):
BUTTON_DEFS = { BUTTON_DEFS = {
"agent_actions": [ "agent_actions": [
( (
"🖥️ - Find, Move, or Generate OTP for Agents", "🖥️ - Find agent, Move agent, or Generate One Time Pass",
"move_agent_workflow_button", "move_agent_workflow_button",
), ),
("📊 - Review and appove OTP Activities", "otp_activities_button"), ("🎫 - Review and approve OTP Activities", "otp_activities_button"),
("📇 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"), ("🛑 - Revoke Active OTP Session", "otp_revoke_button"),
], ],
"policy": [ "policy": [
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"), ("⚖️ - Prepare Policy For Enforcement", "policy_prep_button"),
("🔄 - Update Audit Policies", "policy_audit_update_button"), ("🔕 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"),
("❌ - Revoke OTPs", "otp_revoke_button"),
], ],
} }
@@ -128,7 +139,6 @@ class MainMenuScreen(Screen):
yield Footer() yield Footer()
def on_mount(self) -> None: def on_mount(self) -> None:
api = self.app.api
self.switch_tab("agent_actions") self.switch_tab("agent_actions")
# focus helpers # focus helpers
@@ -191,42 +201,20 @@ class MainMenuScreen(Screen):
self, message: MultiAgentSelector.AgentsSelected self, message: MultiAgentSelector.AgentsSelected
) -> None: ) -> None:
"""Handle selected agents from AgentSelector.""" """Handle selected agents from AgentSelector."""
global _PENDING_JOB global _APP_RESTART_REASON
selected_agents = message.selected_agents selected_agents = message.selected_agents
logger.info("Selected agents: %s", selected_agents) logger.info("Selected agents: %s", selected_agents)
# TODO: Implement actual handling of selected agents # TODO: Implement actual handling of selected agents
_PENDING_JOB = ("multi_agent_action", selected_agents) _APP_RESTART_REASON = ("multi_agent_action", selected_agents)
self.app.exit() self.app.exit()
def on_theme_selector_theme_selected( def on_theme_selector_theme_selected(
self, message: ThemeSelector.ThemeSelected self, message: ThemeSelector.ThemeSelected
) -> None: ) -> None:
"""Handle theme selection from ThemeSelector.""" """Handle theme selection from ThemeSelector."""
global _PENDING_JOB global _APP_RESTART_REASON
_persist_user_theme(message.theme_name) _persist_user_theme(message.theme_name)
_PENDING_JOB = ("restart",) _APP_RESTART_REASON = ("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() self.app.exit()
def on_agent_move_operations_operation_complete( def on_agent_move_operations_operation_complete(
@@ -282,7 +270,6 @@ class MainMenuScreen(Screen):
self.app.bell() self.app.bell()
def on_button_pressed(self, event: Button.Pressed) -> None: def on_button_pressed(self, event: Button.Pressed) -> None:
global _PENDING_JOB
button_id = event.button.id button_id = event.button.id
logger.debug("Button pressed: %s", button_id) logger.debug("Button pressed: %s", button_id)
@@ -308,27 +295,23 @@ class MainMenuScreen(Screen):
return return
case "otp_revoke_button": case "otp_revoke_button":
_PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {}) self.app.push_screen(OTPRevokeScreen())
event.stop()
return
case "policy_prep_button": case "policy_prep_button":
_PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {}) # Use the new TUI workflow screen instead of legacy
self.app.push_screen(
case "policy_audit_update_button": PolicyPrepWorkflowScreen(self.app.api, self.app.policies)
_PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {}) )
event.stop()
return
case _: case _:
self.app.bell() self.app.bell()
logger.warning("Unknown button pressed: %s", button_id) logger.warning("Unknown button pressed: %s", button_id)
return 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 # 2) APP
@@ -353,7 +336,7 @@ class Loxide(App[Message]):
] ]
def __init__(self, api: AirlockAPIWrapper): def __init__(self, api: AirlockAPIWrapper):
self._textual_theme = get_user_value("TEXTUAL_THEME", str, "nord") self._textual_theme = get_user_value("TEXTUAL_THEME", str, "textual-dark")
super().__init__() super().__init__()
self.api = api self.api = api
wd = load_env("WORKING_DIR") or os.getcwd() wd = load_env("WORKING_DIR") or os.getcwd()
@@ -395,8 +378,8 @@ class Loxide(App[Message]):
self.refresh_data() self.refresh_data()
def action_quit(self) -> None: def action_quit(self) -> None:
global _PENDING_JOB global _APP_RESTART_REASON
_PENDING_JOB = None _APP_RESTART_REASON = None
self.exit() self.exit()
def action_open_fe(self) -> None: def action_open_fe(self) -> None:
@@ -410,45 +393,10 @@ class Loxide(App[Message]):
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# 3) TERMINAL + LEGACY # 3) PUBLIC ENTRYPOINT
# ---------------------------------------------------------------------------
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: def run_Loxide(api: AirlockAPIWrapper) -> None:
global _PENDING_JOB global _APP_RESTART_REASON
base_dir = get_base_directory() base_dir = get_base_directory()
env_path = base_dir / ".env" env_path = base_dir / ".env"
dotenv.load_dotenv(dotenv_path=env_path, override=True) dotenv.load_dotenv(dotenv_path=env_path, override=True)
@@ -458,8 +406,8 @@ def run_Loxide(api: AirlockAPIWrapper) -> None:
while attempts < max_attempts: while attempts < max_attempts:
attempts += 1 attempts += 1
logger.debug("Starting job loop iteration (attempt %d)", attempts) logger.debug("Starting app loop iteration (attempt %d)", attempts)
_PENDING_JOB = None _APP_RESTART_REASON = None
app = Loxide(api) app = Loxide(api)
try: try:
@@ -469,42 +417,27 @@ def run_Loxide(api: AirlockAPIWrapper) -> None:
logger.debug("Caught SystemExit from Textual: %s", exc) logger.debug("Caught SystemExit from Textual: %s", exc)
raise raise
job = _PENDING_JOB reason = _APP_RESTART_REASON
logger.debug("After app.run(), _PENDING_JOB = %r", job) logger.debug("After app.run(), _APP_RESTART_REASON = %r", reason)
if not job: if not reason:
logger.debug("No job pending, exiting loop") logger.debug("No restart reason, exiting loop")
break break
if job[0] == "legacy": if reason[0] == "restart":
_, func, args, kwargs = job logger.debug("Restarting app loop")
_run_legacy_job(func, args, kwargs)
continue continue
if job[0] == "restart": if reason[0] == "multi_agent_action":
logger.debug("Restarting job loop") logger.info("Multi-agent action with selected agents: %s", reason[1])
continue continue
if job[0] == "multi_agent_action": logger.error("Unknown restart reason: %r", reason)
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 break
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# 5) DEV # 4) DEV
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
if __name__ == "__main__": if __name__ == "__main__":
api = AirlockAPIWrapper() api = AirlockAPIWrapper()
@@ -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.color import Color
from textual.theme import Theme 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 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.containers import Vertical
from textual.message import Message from textual.message import Message
from textual.widget import Widget 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 import logging
from typing import List, Optional from typing import List, Optional
@@ -23,6 +38,8 @@ logger = logging.getLogger(__name__)
class OTPGenerator(Widget): class OTPGenerator(Widget):
"""Widget for generating OTPs for selected devices."""
# Reactive properties to track form completion # Reactive properties to track form completion
requestor_filled = reactive(False) requestor_filled = reactive(False)
reasoning_filled = reactive(False) reasoning_filled = reactive(False)
@@ -139,14 +156,10 @@ class OTPGenerator(Widget):
button_row.styles.height = "auto" button_row.styles.height = "auto"
button_row.styles.margin = (1, 0, 0, 0) 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_button = Button(
"Generate OTP", id="generate_button", variant="primary" "Generate OTP", id="generate_button", variant="primary"
) )
generate_button.styles.width = "2fr" generate_button.styles.width = "100%"
yield generate_button yield generate_button
# Right side - Show device list initially, then output after generation # Right side - Show device list initially, then output after generation
@@ -169,7 +182,7 @@ class OTPGenerator(Widget):
# Show device list initially # Show device list initially
device_list_text = "\n".join( 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") device_display = Static(device_list_text, id="device_display")
yield device_display yield device_display
@@ -197,14 +210,7 @@ class OTPGenerator(Widget):
def on_button_pressed(self, event: Button.Pressed): def on_button_pressed(self, event: Button.Pressed):
btn_id = event.button.id btn_id = event.button.id
if btn_id == "back_button": if btn_id == "copy_clipboard_button":
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "copy_clipboard_button":
try: try:
output_area = self.query_one("#otp_output", TextArea) output_area = self.query_one("#otp_output", TextArea)
text_to_copy = output_area.text text_to_copy = output_area.text
@@ -213,7 +219,7 @@ class OTPGenerator(Widget):
pyperclip.copy(text_to_copy) pyperclip.copy(text_to_copy)
self.app.notify( self.app.notify(
" Copied to clipboard!", severity="information", timeout=2 " Copied to clipboard!", severity="information", timeout=2
) )
except ImportError: except ImportError:
self.app.notify( 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 dataclasses import asdict
from datetime import datetime from datetime import datetime
import logging import logging
@@ -10,12 +26,12 @@ from textual.css.query import NoMatches
from textual.message import Message from textual.message import Message
from textual.reactive import reactive from textual.reactive import reactive
from textual.widget import Widget 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 models.agent import Agent
from TUI.OTP_generate import OTPGenerator from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
from TUI.otpworkflowscreen import OTPWorkflowScreen from TUI.Screens.policyselectorscreen import PolicySelectorScreen
from TUI.policyselectorscreen import PolicySelectorScreen from TUI.Widgets.OTP_generate import OTPGenerator
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -183,21 +199,21 @@ class AgentMoveOperations(Widget):
f"Operation: {operation_name}", f"Operation: {operation_name}",
f"{'=' * 50}", f"{'=' * 50}",
"", "",
f"✅ Successful ({len(successful)}):", f" Successful ({len(successful)}):",
] ]
if successful: if successful:
for agent, result in successful: for agent, result in successful:
results_lines.append(f" ✅ {agent.hostname}") results_lines.append(f" {agent.hostname}")
else: else:
results_lines.append(" (none)") results_lines.append(" (none)")
results_lines.append("") results_lines.append("")
results_lines.append(f"❌ Failed ({len(unsuccessful)}):") results_lines.append(f" Failed ({len(unsuccessful)}):")
if unsuccessful: if unsuccessful:
for agent, error in unsuccessful: for agent, error in unsuccessful:
results_lines.append(f" ❌ {agent.hostname}: {error}") results_lines.append(f" {agent.hostname}: {error}")
else: else:
results_lines.append(" (none)") results_lines.append(" (none)")
@@ -232,9 +248,9 @@ class AgentMoveOperations(Widget):
- Operations panel: 1/3 width - Operations panel: 1/3 width
- Results area: Initially hidden, shown after operation completion - Results area: Initially hidden, shown after operation completion
""" """
yield Header(show_clock=True, icon="âš™") yield Header(show_clock=True, icon="⚙️")
title_text = Static( 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", id="move_ops_title",
) )
title_text.styles.margin = (0, 0, 1, 0) title_text.styles.margin = (0, 0, 1, 0)
@@ -269,32 +285,32 @@ class AgentMoveOperations(Widget):
yield operations_label yield operations_label
# Operation buttons # 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.width = "100%"
export_csv_btn.styles.margin = (0, 0, 1, 0) export_csv_btn.styles.margin = (0, 0, 1, 0)
yield export_csv_btn yield export_csv_btn
local_approval_btn = Button( 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.width = "100%"
local_approval_btn.styles.margin = (0, 0, 1, 0) local_approval_btn.styles.margin = (0, 0, 1, 0)
yield local_approval_btn 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.width = "100%"
otp_gen_btn.styles.margin = (0, 0, 1, 0) otp_gen_btn.styles.margin = (0, 0, 1, 0)
yield otp_gen_btn yield otp_gen_btn
toggle_enforcement_btn = Button( 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.width = "100%"
toggle_enforcement_btn.styles.margin = (0, 0, 1, 0) toggle_enforcement_btn.styles.margin = (0, 0, 1, 0)
yield toggle_enforcement_btn yield toggle_enforcement_btn
other_policy_btn = Button( 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.width = "100%"
other_policy_btn.styles.margin = (0, 0, 1, 0) other_policy_btn.styles.margin = (0, 0, 1, 0)
@@ -305,10 +321,7 @@ class AgentMoveOperations(Widget):
status_label.styles.margin = (2, 0, 0, 0) status_label.styles.margin = (2, 0, 0, 0)
yield status_label yield status_label
back_button = Button("← Back", id="back_button") yield Footer()
back_button.styles.width = "50%"
back_button.styles.margin = (0, 1, 1, 0)
yield back_button
def on_mount(self) -> None: def on_mount(self) -> None:
""" """
@@ -343,7 +356,7 @@ class AgentMoveOperations(Widget):
Handle button press events from the widget. Handle button press events from the widget.
This Textual event handler routes button presses to appropriate actions: 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) - copy_results_btn: Copy results text to clipboard (requires pyperclip)
- local_approval_btn: Start local approval operation - local_approval_btn: Start local approval operation
- toggle_enforcement_btn: Start toggle audit/enforcement operation - toggle_enforcement_btn: Start toggle audit/enforcement operation
@@ -357,29 +370,24 @@ class AgentMoveOperations(Widget):
btn_id = event.button.id btn_id = event.button.id
if btn_id == "back_button": if btn_id == "copy_results_btn":
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "copy_results_btn":
try: try:
results_text = self.query_one("#results_text", TextArea) results_text = self.query_one("#results_text", TextArea)
import pyperclip import pyperclip
pyperclip.copy(results_text.text) pyperclip.copy(results_text.text)
self.app.notify( self.app.notify(
"📋✅ Results copied to clipboard!", "📋✅ Results copied to clipboard!",
severity="information", severity="information",
timeout=2, timeout=2,
) )
except ImportError: except ImportError:
self.app.notify( self.app.notify(
"❌ pyperclip not installed. Run: pip install pyperclip", " pyperclip not installed. Run: pip install pyperclip",
severity="warning", severity="warning",
) )
except Exception as e: 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() event.stop()
elif btn_id == "export_csv_btn": elif btn_id == "export_csv_btn":
self._start_export_csv_operation() self._start_export_csv_operation()
@@ -427,7 +435,7 @@ class AgentMoveOperations(Widget):
self.operation_in_progress = True self.operation_in_progress = True
status_label = self.query_one("#status_label", Static) 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 # Get API from app
api = self.app.api api = self.app.api
@@ -462,12 +470,12 @@ class AgentMoveOperations(Widget):
except Exception as e: except Exception as e:
logger.error(f"Error during local approval operation: {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 self.operation_in_progress = False
return return
self.operation_in_progress = False self.operation_in_progress = False
status_label.update("✅ Operation complete!") status_label.update(" Operation complete!")
# Display results in the widget # Display results in the widget
self._display_results("Local Approval Mode", successful, unsuccessful) self._display_results("Local Approval Mode", successful, unsuccessful)
@@ -483,7 +491,6 @@ class AgentMoveOperations(Widget):
self.selected_operation = "export_csv" self.selected_operation = "export_csv"
self.operation_in_progress = True self.operation_in_progress = True
successful = [] successful = []
unsuccessful = []
status_label = self.query_one("#status_label", Static) status_label = self.query_one("#status_label", Static)
status_label.update("Exporting CSV...") status_label.update("Exporting CSV...")
self.app.refresh_data() self.app.refresh_data()
@@ -511,9 +518,9 @@ class AgentMoveOperations(Widget):
file_path = os.path.join(str(path), filename) file_path = os.path.join(str(path), filename)
df.to_csv(file_path, index=False) df.to_csv(file_path, index=False)
successful.append(file_path) successful.append(file_path)
status_label.update(f"✅ Exported to {file_path}") status_label.update(f" Exported to {file_path}")
except Exception: except Exception:
status_label.update("❌ Failed") status_label.update(" Failed")
self.operation_in_progress = False self.operation_in_progress = False
@@ -557,7 +564,7 @@ class AgentMoveOperations(Widget):
self.operation_in_progress = True self.operation_in_progress = True
status_label = self.query_one("#status_label", Static) status_label = self.query_one("#status_label", Static)
status_label.update("⏳ Toggling enforcement mode...") status_label.update("🔄 Toggling enforcement mode...")
# Get API from app # Get API from app
api = self.app.api api = self.app.api
@@ -593,12 +600,12 @@ class AgentMoveOperations(Widget):
except Exception as e: except Exception as e:
logger.error(f"Error during toggle enforcement operation: {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 self.operation_in_progress = False
return return
self.operation_in_progress = False self.operation_in_progress = False
status_label.update("✅ Operation complete!") status_label.update(" Operation complete!")
# Display results in the widget # Display results in the widget
self._display_results("Toggle Audit/Enforcement", successful, unsuccessful) self._display_results("Toggle Audit/Enforcement", successful, unsuccessful)
@@ -663,7 +670,7 @@ class AgentMoveOperations(Widget):
except Exception as e: except Exception as e:
logger.error(f"Error loading policies: {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.operation_in_progress = False
self.selected_operation = "" self.selected_operation = ""
self.app.notify(f"Failed to load policies: {str(e)}", severity="error") self.app.notify(f"Failed to load policies: {str(e)}", severity="error")
@@ -695,7 +702,7 @@ class AgentMoveOperations(Widget):
for agent in self.agents: for agent in self.agents:
try: try:
# Move agent to target policy # 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}")) successful.append((agent, f"Moved to {target_policy.name}"))
logger.info( logger.info(
f"Successfully moved {agent.hostname} to policy {target_policy.name}" f"Successfully moved {agent.hostname} to policy {target_policy.name}"
@@ -1,3 +1,18 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import difflib import difflib
import re import re
from typing import List, Optional from typing import List, Optional
@@ -20,6 +35,8 @@ from models.agent import Agent
class MultiAgentSelector(Widget): class MultiAgentSelector(Widget):
"""Widget for selecting multiple agents from a list."""
class AgentsSelected(Message): class AgentsSelected(Message):
def __init__(self, selected_agents: List[Agent]): def __init__(self, selected_agents: List[Agent]):
super().__init__() super().__init__()
@@ -40,7 +57,7 @@ class MultiAgentSelector(Widget):
def compose(self): def compose(self):
yield Header(show_clock=True, icon="") 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) title_text.styles.margin = (0, 0, 0, 1)
yield title_text yield title_text
@@ -60,7 +77,7 @@ class MultiAgentSelector(Widget):
text_area.styles.overflow_y = "auto" text_area.styles.overflow_y = "auto"
yield text_area yield text_area
with Horizontal(id="switch_search_container") as switch_search: with Horizontal(id="switch_search_container"):
switch = Switch(value=False, id="match_switch") switch = Switch(value=False, id="match_switch")
switch.styles.width = "auto" switch.styles.width = "auto"
switch.styles.margin = (1, 0, 0, 0) switch.styles.margin = (1, 0, 0, 0)
@@ -91,11 +108,6 @@ class MultiAgentSelector(Widget):
button_row.styles.height = "auto" button_row.styles.height = "auto"
button_row.styles.margin = (1, 0, 0, 0) 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( submit_button = Button(
"▶ Select & Continue", id="submit_selection", variant="primary" "▶ Select & Continue", id="submit_selection", variant="primary"
) )
@@ -123,10 +135,7 @@ class MultiAgentSelector(Widget):
match_list = self.query_one("#match_results", SelectionList) match_list = self.query_one("#match_results", SelectionList)
except NoMatches: except NoMatches:
return return
if btn_id == "back_button": if btn_id == "select_all":
self.app.pop_screen()
event.stop()
elif btn_id == "select_all":
match_list.select_all() match_list.select_all()
event.stop() event.stop()
elif btn_id == "select_none": elif btn_id == "select_none":
@@ -1,10 +1,17 @@
""" # Copyright (C) 2025 James Brotosky, Brandon Wickline
Policy Selector Widget Module #
# This program is free software: you can redistribute it and/or modify
Provides a Textual widget for selecting target policies for bulk agent operations. # it under the terms of the GNU Affero General Public License as published
Allows users to browse available policies and select one as the destination for # by the Free Software Foundation, either version 3 of the License, or
moving agents. Automatically excludes parent/logical policies. # (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 logging
import re import re
@@ -34,7 +41,7 @@ class PolicySelector(Widget):
- Wildcard filtering (* and ?) - Wildcard filtering (* and ?)
- Interactive table for policy browsing - Interactive table for policy browsing
- Explicit confirm button for selection - Explicit confirm button for selection
- Cancel/back button to dismiss - Use escape key to go back
Attributes: Attributes:
policies (list[Policy]): List of available Policy objects to display. policies (list[Policy]): List of available Policy objects to display.
@@ -91,10 +98,10 @@ class PolicySelector(Widget):
- Clear Filter button - Clear Filter button
- Confirm Selection button - Confirm Selection button
- Policy table displaying available policies - Policy table displaying available policies
- Back buttons for navigation - Use escape key to go back
""" """
title_text = Static( title_text = Static(
"🎯 Select Target Policy", "Select Target Policy",
id="policy_selector_title", id="policy_selector_title",
) )
title_text.styles.margin = (0, 0, 1, 0) title_text.styles.margin = (0, 0, 1, 0)
@@ -125,12 +132,12 @@ class PolicySelector(Widget):
filter_help.styles.margin = (0, 0, 1, 0) filter_help.styles.margin = (0, 0, 1, 0)
yield filter_help 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.width = "100%"
apply_button.styles.margin = (0, 0, 1, 0) apply_button.styles.margin = (0, 0, 1, 0)
yield apply_button 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.width = "100%"
clear_button.styles.margin = (0, 0, 1, 0) clear_button.styles.margin = (0, 0, 1, 0)
yield clear_button yield clear_button
@@ -144,11 +151,6 @@ class PolicySelector(Widget):
selected_label.styles.margin = (2, 0, 1, 0) selected_label.styles.margin = (2, 0, 1, 0)
yield selected_label 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 # Right side - Policy table
with Vertical() as right_side: with Vertical() as right_side:
right_side.styles.width = "2fr" right_side.styles.width = "2fr"
@@ -222,7 +224,6 @@ class PolicySelector(Widget):
Handle button press events from the widget. Handle button press events from the widget.
Routes to: Routes to:
- back_button (Cancel): Pop screen without selecting
- filter_button (Apply Filter): Filter policies with wildcard support - filter_button (Apply Filter): Filter policies with wildcard support
- clear_filter_button: Clear filter and show all policies - clear_filter_button: Clear filter and show all policies
- confirm_button: Confirm selection and post message - confirm_button: Confirm selection and post message
@@ -232,12 +233,7 @@ class PolicySelector(Widget):
""" """
btn_id = event.button.id btn_id = event.button.id
if btn_id == "back_button": if btn_id == "filter_button":
while len(self.app.screen_stack) > 2:
self.app.pop_screen()
event.stop()
elif btn_id == "filter_button":
self._apply_filter() self._apply_filter()
event.stop() event.stop()
@@ -283,7 +279,7 @@ class PolicySelector(Widget):
if self.selected_policy: if self.selected_policy:
# Update selection display # Update selection display
label = self.query_one("#selected_policy_label", Static) 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 # Log for debugging
logger.debug( logger.debug(
@@ -325,7 +321,7 @@ class PolicySelector(Widget):
if highlighted_name: if highlighted_name:
label = self.query_one("#selected_policy_label", Static) label = self.query_one("#selected_policy_label", Static)
label.update(f"Highlighting: {highlighted_name}") label.update(f"Highlighting: {highlighted_name}")
except Exception as e: except Exception as e:
logger.error(f"Error handling row highlight: {e}") logger.error(f"Error handling row highlight: {e}")
@@ -399,7 +395,9 @@ class PolicySelector(Widget):
) )
displayed_count = len(self._displayed_policies) 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) self.app.notify(status_text, severity="information", timeout=2)
# Clear selection when filter is applied # Clear selection when filter is applied
@@ -474,7 +472,7 @@ class PolicySelector(Widget):
""" """
if self.selected_policy is None: if self.selected_policy is None:
self.app.notify( 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", severity="warning",
timeout=3, timeout=3,
) )
@@ -483,6 +481,6 @@ class PolicySelector(Widget):
# Log confirmation for debugging # Log confirmation for debugging
logger.info(f"Confirming selection of policy: {self.selected_policy.name}") logger.info(f"Confirming selection of policy: {self.selected_policy.name}")
self.app.notify( 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)) self.post_message(self.PolicySelected(self.selected_policy))
@@ -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 collections import defaultdict from collections import defaultdict
import logging import logging
+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 import logging
from textual.containers import Horizontal, Vertical from textual.containers import Horizontal, Vertical
@@ -59,15 +74,6 @@ class ResultsDisplay(Widget):
margin-top: 1; margin-top: 1;
width: 100%; width: 100%;
} }
#button_row {
height: auto;
margin: 1 0 0 0;
}
#back_button {
width: 1fr;
}
""" """
class CopySuccess(Message): class CopySuccess(Message):
@@ -95,9 +101,9 @@ class ResultsDisplay(Widget):
def compose(self): def compose(self):
with Vertical(id="results_screen"): with Vertical(id="results_screen"):
yield Header(show_clock=True, icon="") yield Header(show_clock=True, icon="")
# Title # Title
title = Static(f"📊 {self.operation} - Results", id="results_title") title = Static(f"{self.operation} - Results", id="results_title")
yield title yield title
# Two-column layout # Two-column layout
@@ -107,7 +113,7 @@ class ResultsDisplay(Widget):
yield Static("✅ Successful", id="success_label") yield Static("✅ Successful", id="success_label")
yield Static(self.successful_results, id="success_results") yield Static(self.successful_results, id="success_results")
yield Button( yield Button(
"📋✅ Copy Success List", "Copy Success List",
id="copy_success", id="copy_success",
classes="copy_button", classes="copy_button",
) )
@@ -117,15 +123,11 @@ class ResultsDisplay(Widget):
yield Static("❌ Failed", id="failure_label") yield Static("❌ Failed", id="failure_label")
yield Static(self.unsuccessful_results, id="failure_results") yield Static(self.unsuccessful_results, id="failure_results")
yield Button( yield Button(
"📋❌ Copy Failure List", "Copy Failure List",
id="copy_failure", id="copy_failure",
classes="copy_button", classes="copy_button",
) )
# Back Button
with Horizontal(id="button_row"):
back_button = Button("← Back", id="back_button")
yield back_button
yield Footer() yield Footer()
def on_button_pressed(self, event: Button.Pressed) -> None: def on_button_pressed(self, event: Button.Pressed) -> None:
@@ -138,18 +140,18 @@ class ResultsDisplay(Widget):
pyperclip.copy(str(success_widget.renderable)) pyperclip.copy(str(success_widget.renderable))
self.app.notify( self.app.notify(
"✅ Success list copied to clipboard!", "Success list copied to clipboard!",
severity="information", severity="information",
timeout=2, timeout=2,
) )
self.post_message(self.CopySuccess()) self.post_message(self.CopySuccess())
except ImportError: except ImportError:
self.app.notify( self.app.notify(
"⚠️ pyperclip not installed. Run: pip install pyperclip", " pyperclip not installed. Run: pip install pyperclip",
severity="warning", severity="warning",
) )
except Exception as e: 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() event.stop()
elif btn_id == "copy_failure": elif btn_id == "copy_failure":
@@ -159,20 +161,16 @@ class ResultsDisplay(Widget):
pyperclip.copy(str(failure_widget.renderable)) pyperclip.copy(str(failure_widget.renderable))
self.app.notify( self.app.notify(
"✅ Failure list copied to clipboard!", "Failure list copied to clipboard!",
severity="information", severity="information",
timeout=2, timeout=2,
) )
self.post_message(self.CopyFailure()) self.post_message(self.CopyFailure())
except ImportError: except ImportError:
self.app.notify( self.app.notify(
"⚠️ pyperclip not installed. Run: pip install pyperclip", " pyperclip not installed. Run: pip install pyperclip",
severity="warning", severity="warning",
) )
except Exception as e: 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 == "back_button":
self.app.pop_screen()
event.stop() 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."""
+1
View File
@@ -1,3 +1,4 @@
/target /target
build.sh build.sh
pythontest.py pythontest.py
changelog.md
+559 -801
View File
File diff suppressed because it is too large Load Diff
+14 -12
View File
@@ -1,29 +1,31 @@
[package] [package]
name = "airlock_libs" name = "signoz_test"
version = "3.1.2" version = "6.0.0"
edition = "2024" edition = "2024"
[lib]
crate-type = ["cdylib"]
[dependencies] [dependencies]
chrono = "0.4.42" chrono = "0.4.42"
indicatif = "0.18.2" indicatif = "0.18.2"
mongodb = "3.3.0" mongodb = "3.3.0"
opentelemetry = { version = "0.18.0", features = ["rt-tokio", "metrics", "trace"] } opentelemetry = { version = "0.27.0", features = ["logs", "metrics", "trace"] }
opentelemetry-otlp = { version = "0.11.0", features = ["trace", "metrics"] } opentelemetry-otlp = { version = "0.27.0", features = ["trace", "metrics", "grpc-tonic", "http-proto", "tls", "reqwest-client", "reqwest-rustls"] }
opentelemetry-semantic-conventions = { version = "0.10.0" } opentelemetry-semantic-conventions = { version = "0.27.0" }
opentelemetry-proto = { version = "0.1.0"} opentelemetry-proto = { version = "0.27.0"}
pyo3 = { version = "0.27.0", features = ["extension-module", "generate-import-lib"] } 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 = "1.0.228"
serde-pyobject = "0.8.0" serde-pyobject = "0.8.0"
serde_json = "1.0.145" serde_json = "1.0.145"
tokio = { version = "1.48.0", features = ["full"] } 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 = "0.1.41"
tracing-subscriber = "0.3.20" tracing-subscriber = "0.3.20"
tracing-opentelemetry = "0.32.0" tracing-opentelemetry = "0.32.0"
crossbeam = "0.8.4"
log = "0.4.29"
flexi_logger = "0.31.7"
opentelemetry-appender-log = "0.27.0"
opentelemetry_sdk = { version = "0.27.0", features = ["rt-tokio", "trace"] }
[package.metadata.maturin] [package.metadata.maturin]
generate-abi-stubs = true generate-abi-stubs = true
@@ -36,4 +38,4 @@ codegen-units = 1
panic = 'abort' panic = 'abort'
strip = true strip = true
debug-assertions = false debug-assertions = false
overflow-checks = false overflow-checks = true
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "maturin"
[project] [project]
name = "airlock_libs" name = "airlock_libs"
version = "3.1.2" version = "6.0.0"
description = "Airlock Digital API Wrapper" description = "Airlock Digital API Wrapper"
readme = "README.md" readme = "README.md"
license = { text = "AGPL-3.0-only" } license = { text = "AGPL-3.0-only" }
+3 -1
View File
@@ -1,5 +1,7 @@
use pyo3::prelude::*; use pyo3::prelude::*;
mod services; pub mod modules;
pub mod prelude;
pub mod services;
#[pymodule] #[pymodule]
fn airlock_libs(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> { fn airlock_libs(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(services::pull_policy_exec_histories, py)?)?; m.add_function(wrap_pyfunction!(services::pull_policy_exec_histories, py)?)?;
+106
View File
@@ -0,0 +1,106 @@
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 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;
+27
View File
@@ -0,0 +1,27 @@
pub use chrono::{Duration, Local, NaiveDate};
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_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::{
collections::HashMap,
env,
fmt::Write,
fs::{self, File},
io::{Read, Seek, SeekFrom},
path::PathBuf,
str::FromStr,
};
+182 -216
View File
@@ -1,71 +1,10 @@
use chrono::{Duration, Local, NaiveDate}; use crate::modules::datatypes::*;
use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle}; use crate::prelude::*;
use mongodb::bson::oid::ObjectId; use crossbeam::channel::unbounded;
use opentelemetry::global::shutdown_tracer_provider; use opentelemetry_otlp::WithTonicConfig;
use opentelemetry::sdk::Resource; use std::sync::{Arc, Mutex};
use opentelemetry::trace::noop::NoopTracerProvider; use std::thread;
use opentelemetry::trace::{Status, TraceContextExt, TraceError}; use tonic::transport::{Channel, ClientTlsConfig};
use opentelemetry::{Context, KeyValue, sdk::trace as sdktrace, trace::Tracer};
use opentelemetry::{Key, global};
use opentelemetry_otlp::WithExportConfig;
use pyo3::{prelude::*, types::PyString};
use reqwest::{
Client,
header::{HeaderMap, HeaderName, HeaderValue},
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{
collections::HashMap,
env,
fmt::Write,
fs::{self, File},
io::{Read, Seek, SeekFrom},
path::PathBuf,
str::FromStr,
};
#[derive(Deserialize, Debug)]
struct TelemetryConfig {
TELEMETRY: bool,
TELEM_URL: Option<String>,
}
#[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,
}
#[pyfunction] #[pyfunction]
pub fn pull_policy_exec_histories( pub fn pull_policy_exec_histories(
py: Python<'_>, py: Python<'_>,
@@ -74,46 +13,75 @@ pub fn pull_policy_exec_histories(
exec_types: String, exec_types: String,
days: i64, days: i64,
) -> Py<PyString> { ) -> Py<PyString> {
let rt = tokio::runtime::Runtime::new().unwrap(); println!();
rt.block_on(async { let data: PyData = PyData::extract_data(py, &py_self);
let _ = init_tracer(); let headers: HeaderMap = data.headers;
}); let base_url: String = data.base_url;
let tracer = global::tracer("global_tracer"); let handle: thread::JoinHandle<String> = std::thread::spawn(move || {
let _cx = Context::new(); let rt: tokio::runtime::Runtime = match tokio::runtime::Runtime::new() {
Ok(rt) => rt,
Err(e) => {
println!("Failed to build Tokio Runtime: {:?}", e);
std::process::abort();
}
};
let tracer_provider = rt.block_on(async { 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!( let file_path: PathBuf = format!(
"{}\\cache\\chunkinator.json", "{}\\cache\\chunkinator.json",
get_base_directory().display() get_base_directory().display()
) )
.into(); .into();
let writeable_filepath = file_path.clone(); if !&file_path.exists() {
if !file_path.exists() { if let Some(parent_dir) = &file_path.parent()
if let Some(parent_dir) = file_path.parent()
&& !parent_dir.exists() && !parent_dir.exists()
{ {
fs::create_dir_all(parent_dir).unwrap(); match fs::create_dir_all(parent_dir) {
Ok(_) => {}
Err(e) => {
println!("Failed to Create Directory {:?}: {}", parent_dir, e);
std::process::abort();
} }
fs::File::create(file_path).unwrap();
} }
let data = ApiResponse { }
match fs::File::create(&file_path) {
Ok(_) => {}
Err(e) => {
println!("Failed to Create Directory {:?}: {}", &file_path, e);
std::process::abort();
}
}
}
let data: ApiResponse = ApiResponse {
error: "Success".to_string(), error: "Success".to_string(),
response: ExecHistories { response: ExecHistories {
exechistories: vec![], exechistories: vec![],
}, },
}; };
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize"); let writeable_filepath: PathBuf = file_path.clone();
fs::write(writeable_filepath.clone(), data_write).unwrap(); let data_write: String = serde_json::to_string_pretty(&data).expect("Failed to serialize");
let mut checkpoint_number: String = skipback(days).to_string(); match fs::write(writeable_filepath.clone(), data_write) {
let multi_progress = MultiProgress::new(); Ok(_) => {}
multi_progress.set_draw_target(ProgressDrawTarget::stdout()); Err(e) => {
let progress_bar = multi_progress.add(ProgressBar::new(100)); println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
progress_bar.set_style( std::process::abort();
}
}
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() ProgressStyle::default_bar()
.template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}") .template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len} {message}")
.unwrap(), .unwrap(),
); );
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100)); let client: Client = tracer.in_span("Building HTTP Client", |cx| {
let client = tracer.in_span("Building HTTP Client", |cx| { let client_result: Result<Client, reqwest::Error> = build_client(headers);
let client_result = build_client(py, &py_self);
match client_result { match client_result {
Ok(client_result) => { Ok(client_result) => {
cx.span().add_event( cx.span().add_event(
@@ -135,53 +103,31 @@ pub fn pull_policy_exec_histories(
); );
cx.span() cx.span()
.set_status(Status::error("Client Failed to Build")); .set_status(Status::error("Client Failed to Build"));
panic!("Failed to Build Client: {:?}", client_result); println!("Failed to Build Client: {:?}", client_result);
std::process::abort();
} }
} }
}); });
let api: Py<PyAny> = py_self; let cutoff: chrono::NaiveDateTime =
let cutoff = Local::now().naive_local() - Duration::days(days); Local::now().naive_local() - chrono::Duration::days(days);
let mut f = File::open(&writeable_filepath).unwrap(); let (tx, rx) = unbounded::<Vec<Group>>();
tracer.in_span("Airlock Data Retreival", |cx| { let pb_clone = progress_bar.clone();
let span = cx.span(); thread::spawn(move || {
span.set_attribute(Key::new("Days").string(days.to_string().to_string()));
loop {
f.seek(SeekFrom::Start(0)).unwrap();
let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
let results: ApiResponse = history_logging(
py,
&api,
&exec_types,
&checkpoint_number,
&policy_names,
&client,
);
cx.span().set_attribute(KeyValue::new(
"items_in_response",
results.response.exechistories.len().to_string(),
));
results
});
let parsed_responses = execution_histories.response.exechistories;
if parsed_responses.is_empty() {
break;
}
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
{ {
let mut contents = String::new(); let contents: String = fs::read_to_string(&writeable_filepath).unwrap_or_default();
f.read_to_string(&mut contents).unwrap(); let existing: ApiResponse =
let existing_data: ApiResponse =
serde_json::from_str(&contents).unwrap_or(ApiResponse { serde_json::from_str(&contents).unwrap_or(ApiResponse {
error: "Success".to_string(), error: "Success".to_string(),
response: ExecHistories { response: ExecHistories {
exechistories: vec![], exechistories: vec![],
}, },
}); });
existing_data existing
.response .response
.exechistories .exechistories
.into_iter() .into_iter()
.map(|entry| { .map(|entry: Group| {
( (
( (
entry.sha256.clone(), entry.sha256.clone(),
@@ -195,15 +141,12 @@ pub fn pull_policy_exec_histories(
} else { } else {
HashMap::new() HashMap::new()
}; };
for (index, executions) in parsed_responses.iter().enumerate() { while let Ok(parsed_responses) = rx.recv() {
for executions in parsed_responses {
if executions.checkpoint.is_empty() || executions.datetime.is_empty() { if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
continue; continue;
} }
if index == parsed_responses.len() - 1 { let history_date: NaiveDate = match NaiveDate::parse_from_str(
checkpoint_number = executions.checkpoint.clone();
break;
}
let history_date = match NaiveDate::parse_from_str(
&executions.datetime.replace(" +0000 UTC", ""), &executions.datetime.replace(" +0000 UTC", ""),
"%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%SZ",
) { ) {
@@ -211,7 +154,7 @@ pub fn pull_policy_exec_histories(
Err(_) => continue, Err(_) => continue,
}; };
if history_date >= cutoff.into() { if history_date >= cutoff.into() {
let key = ( let key: (String, String, String) = (
executions.sha256.clone(), executions.sha256.clone(),
executions.filename.clone(), executions.filename.clone(),
executions.hostname.clone(), executions.hostname.clone(),
@@ -219,64 +162,113 @@ pub fn pull_policy_exec_histories(
seen.entry(key).or_insert(executions.clone()); seen.entry(key).or_insert(executions.clone());
} }
} }
let final_response = ApiResponse { let final_response: ApiResponse = ApiResponse {
error: "Success".to_string(), error: "Success".to_string(),
response: ExecHistories { response: ExecHistories {
exechistories: seen.values().cloned().collect(), exechistories: seen.values().cloned().collect(),
}, },
}; };
let data_write = serde_json::to_string_pretty(&final_response).unwrap(); let data_write: String = serde_json::to_string_pretty(&final_response).unwrap();
fs::write(&writeable_filepath, data_write).unwrap(); match fs::write(&writeable_filepath, data_write) {
if let Some(last_item) = &final_response.response.exechistories.last() Ok(_) => {}
Err(e) => {
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));
let span: opentelemetry::trace::SpanRef<'_> = cx.span();
//span.set_attribute(Key::new("Days").string(days.to_string()));
//span.set_attribute(Key::new("Days"));
//span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
span.set_attribute(KeyValue::new("Days", days));
span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
loop {
let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
let results: ApiResponse = history_logging(
&base_url,
&exec_types,
&checkpoint_number,
&policy_names,
&client,
);
cx.span().set_attribute(KeyValue::new(
"items_in_response",
results.response.exechistories.len().to_string(),
));
results
});
let parsed_responses: Vec<Group> = execution_histories.response.exechistories;
if parsed_responses.is_empty() {
break;
}
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( && let Ok(last_date) = NaiveDate::parse_from_str(
&last_item.datetime.replace(" +0000 UTC", ""), &last_item.datetime.replace(" +0000 UTC", ""),
"%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%SZ",
) )
{ {
let date_diff = Local::now().naive_local().date() - last_date; if first_date.is_none() {
let percentage_diff = (days - date_diff.num_days()) as f64 / days as f64 * 100.0; first_date = Some(last_date);
progress_bar.set_position(percentage_diff.round() as u64); }
progress_bar.set_message("Total Percent Complete"); 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"); progress_bar
let return_data = fs::read_to_string(&writeable_filepath).unwrap(); .lock()
shutdown_tracer_provider(); .unwrap()
PyString::new(py, &return_data).into() .finish_with_message("All Checkpoints Complete");
} let return_data: String = match fs::read_to_string(file_path.clone()) {
Ok(return_data) => return_data,
fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Result<reqwest::Client, reqwest::Error> { Err(e) => {
let headers = py_self.getattr(py, "headers").unwrap().to_string(); println!("Failed to read data from: {:?}: {}", &file_path, e);
let headers_replace = headers.replace('\'', "\""); std::process::abort();
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() { tracer_provider
for (_key, value) in obj { .shutdown()
if let Some(v) = value.as_str() { .expect("Failed to Shutdown Tracer Provdier");
let val = HeaderValue::from_str(v).unwrap(); drop(tx);
header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val); return_data.to_string()
} });
} 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> {
Client::builder() Client::builder()
.danger_accept_invalid_certs(true) .danger_accept_invalid_certs(true)
.default_headers(header_map) .default_headers(headers)
.timeout(std::time::Duration::from_secs(300)) .timeout(std::time::Duration::from_secs(300))
.build() .build()
} }
#[tokio::main] #[tokio::main]
async fn history_logging( async fn history_logging(
py: Python<'_>, base_url: &String,
py_self: &Py<PyAny>,
exec_types: &String, exec_types: &String,
checkpoint_number: &String, checkpoint_number: &String,
policy_names: &String, policy_names: &String,
client: &Client, client: &Client,
) -> ApiResponse { ) -> ApiResponse {
let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
let payload = format!( let payload = format!(
r#"{{ r#"{{
"type": {}, "type": {},
@@ -285,7 +277,7 @@ async fn history_logging(
}}"#, }}"#,
exec_types, checkpoint_number, policy_names exec_types, checkpoint_number, policy_names
); );
let res = client let res: Result<reqwest::Response, reqwest::Error> = client
.post(format!("{}/v1/logging/exechistories", base_url)) .post(format!("{}/v1/logging/exechistories", base_url))
.body(payload) .body(payload)
.send() .send()
@@ -308,7 +300,7 @@ async fn history_logging(
} }
} }
fn get_base_directory() -> PathBuf { pub fn get_base_directory() -> PathBuf {
let home = env::var_os("HOME") let home = env::var_os("HOME")
.map(PathBuf::from) .map(PathBuf::from)
.or_else(|| env::var_os("USERPROFILE").map(PathBuf::from)) .or_else(|| env::var_os("USERPROFILE").map(PathBuf::from))
@@ -321,60 +313,34 @@ fn get_base_directory() -> PathBuf {
.unwrap_or_else(|| home.join("AppData").join("Roaming")); .unwrap_or_else(|| home.join("AppData").join("Roaming"));
appdata.join("Loxide") 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 { fn init_tracer() -> opentelemetry_sdk::trace::TracerProvider {
let date_days_ago = Local::now() - Duration::days(days); let cfg: TelemetryConfig = TelemetryConfig::load();
let timestamp = date_days_ago.timestamp() as u32; let endpoint = cfg.TELEM_URL.unwrap_or_default().clone();
let mut hex_timestamp = String::new(); let channel_endpoint = endpoint.clone();
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap(); let channel = Channel::from_shared(channel_endpoint.clone())
let objectid_hex = format!("{}0000000000000000", hex_timestamp); .unwrap()
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex") .tls_config(ClientTlsConfig::new().with_native_roots())
} .unwrap()
.connect_lazy();
fn load_telemetry_config() -> TelemetryConfig { let exporter = opentelemetry_otlp::SpanExporter::builder()
let cfg_path = get_base_directory().join("config\\user_config.json"); .with_tonic()
if !cfg_path.exists() { .with_endpoint(endpoint.clone())
return TelemetryConfig { .with_channel(channel)
TELEMETRY: false, .build()
TELEM_URL: None, .expect("Failed to build exporter");
}; opentelemetry_sdk::trace::TracerProvider::builder()
} .with_simple_exporter(exporter)
match fs::read_to_string(&cfg_path) { .with_resource(Resource::new(vec![KeyValue::new(
Ok(contents) => { "service.name",
serde_json::from_str::<TelemetryConfig>(&contents).unwrap_or(TelemetryConfig { "LoxideLibs",
TELEMETRY: false, )]))
TELEM_URL: None, .build()
})
}
Err(_) => TelemetryConfig {
TELEMETRY: false,
TELEM_URL: None,
},
}
}
fn init_tracer() -> Result<Option<sdktrace::Tracer>, TraceError> {
let cfg = load_telemetry_config();
if !cfg.TELEMETRY {
global::set_tracer_provider(NoopTracerProvider::new());
return Ok(None);
}
let endpoint = cfg.TELEM_URL.unwrap_or_default();
let tracer =
opentelemetry_otlp::new_pipeline()
.tracing()
.with_exporter(
opentelemetry_otlp::new_exporter()
.tonic()
.with_endpoint(endpoint),
)
.with_trace_config(sdktrace::config().with_resource(Resource::new(vec![
KeyValue::new("service.name", "LoxideLibs"),
])))
.install_simple()
.unwrap();
Ok(Some(tracer))
} }
+37 -28
View File
@@ -1,6 +1,17 @@
""" # Copyright (C) 2025 James Brotosky, Brandon Wickline
This module handles the creation of local approval requests. #
""" # 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 logging
import os import os
@@ -51,7 +62,7 @@ class LocalApprovalRequestor:
batch_id = int(time.time()) batch_id = int(time.time())
purpose = ( purpose = (
f"🎫 Local Approval 🎫 - {duration_minutes} mins - " f" Local Approval - {duration_minutes} mins - "
f"batch:{batch_id} Client:{agent_id} User:{self.username}" f"batch:{batch_id} Client:{agent_id} User:{self.username}"
) )
@@ -103,11 +114,9 @@ class LocalApprovalRequestor:
success_count = 0 success_count = 0
failure_count = 0 failure_count = 0
print(colorText(f"\n📦 Processing batch {batch_id}...", "cyan")) print(colorText(f"\n Processing batch {batch_id}...", "cyan"))
print(colorText(f"👤 Requested by: {self.username}", "cyan")) print(colorText(f" Requested by: {self.username}", "cyan"))
print( print(colorText(f" Moving {len(agents)} agent(s) to local approval\n", "cyan"))
colorText(f"📊 Moving {len(agents)} agent(s) to local approval\n", "cyan")
)
for agent in agents: for agent in agents:
try: try:
@@ -125,11 +134,11 @@ class LocalApprovalRequestor:
if not move_success: if not move_success:
raise Exception("Failed to move to audit policy") raise Exception("Failed to move to audit policy")
print(colorText(f"✓ {agent.hostname}", "green")) print(colorText(f" {agent.hostname}", "green"))
success_count += 1 success_count += 1
except Exception as e: 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}") logger.error(f"Error processing agent {agent.hostname}: {e}")
failure_count += 1 failure_count += 1
@@ -152,7 +161,7 @@ class LocalApprovalRequestor:
] ]
# Display duration options # Display duration options
print(colorText("\n⏱️ Select Local Approval Duration:", "white")) print(colorText("\n Select Local Approval Duration:", "white"))
print(colorText("=" * 50, "white")) print(colorText("=" * 50, "white"))
for i, (minutes, label) in enumerate(duration_options, start=1): for i, (minutes, label) in enumerate(duration_options, start=1):
@@ -166,36 +175,36 @@ class LocalApprovalRequestor:
if 1 <= choice <= len(duration_options): if 1 <= choice <= len(duration_options):
duration_minutes, duration_label = duration_options[choice - 1] 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") logger.info(f"User selected duration: {duration_minutes} minutes")
else: else:
print(colorText("❌ Invalid choice.", "red")) print(colorText(" Invalid choice.", "red"))
logger.warning("Invalid duration choice") logger.warning("Invalid duration choice")
return return
except ValueError: 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") logger.warning("Invalid input for duration selection")
return return
# Select agents # Select agents
print(colorText("\n🎯 Select Agents for Local Approval:", "white")) print(colorText("\nSelect Agents for Local Approval:", "white"))
agents = selectAgents(self.api) agents = selectAgents(self.api)
if not agents: 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") logger.warning("No agents selected or error retrieving agents")
return return
# Confirm with user # Confirm with user
print(colorText("\n📋 Summary:", "cyan")) print(colorText("\nSummary:", "cyan"))
print(colorText(f" Duration: {duration_label}", "white")) print(colorText(f" Duration: {duration_label}", "white"))
print(colorText(f" Agents: {len(agents)}", "white")) print(colorText(f" Agents: {len(agents)}", "white"))
confirm = get_sanitized_input("\nProceed? (y/n): ").lower() confirm = get_sanitized_input("\nProceed? (y/n): ").lower()
if confirm != "y": if confirm != "y":
print(colorText("❌ Operation cancelled.", "yellow")) print(colorText(" Operation cancelled.", "yellow"))
return return
# Process the batch # Process the batch
@@ -219,24 +228,24 @@ class LocalApprovalRequestor:
failure_count: Number of failed operations failure_count: Number of failed operations
""" """
print(colorText(f"\n{'=' * 60}", "white")) print(colorText(f"\n{'=' * 60}", "white"))
print(colorText("📊 Local Approval Summary", "cyan")) print(colorText(" Local Approval Summary", "cyan"))
print(colorText("=" * 60, "white")) print(colorText("=" * 60, "white"))
print(colorText(f"✓ Successfully processed: {success_count}", "green")) print(colorText(f" Successfully processed: {success_count}", "green"))
if failure_count > 0: 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"\n Batch ID: {batch_id}", "cyan"))
print(colorText(f"⏱️ Duration: {duration_label}", "cyan")) print(colorText(f" Duration: {duration_label}", "cyan"))
print(colorText("=" * 60, "white")) print(colorText("=" * 60, "white"))
print(colorText("\n💡 Next Steps:", "yellow")) print(colorText("\n Next Steps:", "yellow"))
print(colorText(" • Agents have been moved to audit policies", "white")) print(colorText(" Agents have been moved to audit policies", "white"))
print(colorText(" • Local approvals are active", "white")) print(colorText(" Local approvals are active", "white"))
print( print(
colorText( colorText(
f" • Agents will return to enforcement after {duration_label}", f" Agents will return to enforcement after {duration_label}",
"white", "white",
) )
) )
+23 -23
View File
@@ -88,7 +88,7 @@ def sortHashes(
): ):
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
history_days = Selector.select_value( 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, value_type=int,
valid_range=(1, 150), valid_range=(1, 150),
) )
@@ -655,7 +655,7 @@ def section_header(title):
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist): def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR") working_dir = load_env("WORKING_DIR")
section_header("🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒") section_header("Prepare to Enforce Policy")
print( print(
colorText( colorText(
"\nSequentially follow these steps to prepare a policy for enforcement:", "\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: if not selected_policies:
print(colorText(" [✗] No policies have been chosen", "red")) print(colorText(" [] No policies have been chosen", "red"))
else: else:
print(colorText("The following policies have been chosen:", "green")) print(colorText("The following policies have been chosen:", "green"))
for policy in selected_policies: for policy in selected_policies:
print(colorText(f" [✓] {policy.name}", "green")) print(colorText(f" [] {policy.name}", "green"))
# Step 2: Destination Policy and Allowlist # Step 2: Destination Policy and Allowlist
print( print(
@@ -683,22 +683,22 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
if destination_policy: if destination_policy:
print( print(
colorText( 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", "green",
) )
) )
else: else:
print(colorText(" [✗] No destination policy has been chosen", "red")) print(colorText(" [] No destination policy has been chosen", "red"))
if destination_allowlist: if destination_allowlist:
print( print(
colorText( colorText(
f" [✓] {destination_allowlist[0].name} has been selected as allowlist", f" [] {destination_allowlist[0].name} has been selected as allowlist",
"green", "green",
) )
) )
else: else:
print(colorText(" [✗] No allowlist has been chosen", "red")) print(colorText(" [] No allowlist has been chosen", "red"))
# Step 3: Data Preparation # Step 3: Data Preparation
print( print(
@@ -713,9 +713,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print( print(
colorText( colorText(
( (
" [✓] Data has been fetched" " [] Data has been fetched"
if os.path.exists(review_path) 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", "green" if os.path.exists(review_path) else "red",
) )
@@ -723,7 +723,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
else: else:
print( print(
colorText( 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( print(
colorText( colorText(
( (
" [✓] Reviewed hashes have been loaded" " [] Reviewed hashes have been loaded"
if os.path.exists(approved_path) 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", "green" if os.path.exists(approved_path) else "red",
) )
@@ -766,9 +766,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print( print(
colorText( colorText(
( (
" [✓] Path review list created" " [] Path review list created"
if os.path.exists(second_review_path) 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", "green" if os.path.exists(second_review_path) else "red",
) )
@@ -776,7 +776,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
else: else:
print( print(
colorText( colorText(
" [✗] No policies selected, cannot check reviewed hashes or path list", " [] No policies selected, cannot check reviewed hashes or path list",
"red", "red",
) )
) )
@@ -812,9 +812,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print( print(
colorText( colorText(
( (
" [✓] Reviewed path list detected" " [] Reviewed path list detected"
if os.path.exists(reviewed_path) 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", "green" if os.path.exists(reviewed_path) else "red",
) )
@@ -825,9 +825,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
print( print(
colorText( colorText(
( (
" [✓] Preflight Path Exclusion List has been generated" " [] Preflight Path Exclusion List has been generated"
if preflight_ready 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", "green" if preflight_ready else "red",
) )
@@ -835,7 +835,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
else: else:
print( print(
colorText( 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")) print(colorText(" Apply approved hashes to allowlist", "cyan"))
# Utility Options # Utility Options
print(colorText("F. 📂 - Open Working Directory", "cyan")) print(colorText("F. Open Working Directory", "cyan"))
print(colorText("B. 🔚 - Back", "cyan")) print(colorText("B. Back", "cyan"))
+4 -4
View File
@@ -145,13 +145,13 @@ class Hash:
approved_count += 1 approved_count += 1
except (ValueError, TypeError): except (ValueError, TypeError):
logger.debug( 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" hash_obj.at_decision = "needs_review"
needs_review_count += 1 needs_review_count += 1
logger.debug( 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 return hashes
@@ -443,13 +443,13 @@ class ExecutionHistoryRecord:
approved_count += 1 approved_count += 1
except (ValueError, TypeError) as e: except (ValueError, TypeError) as e:
logger.debug( 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" hash_obj.at_decision = "needs_review"
needs_review_count += 1 needs_review_count += 1
logger.debug( 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}" f"Approved: {approved_count}, Unapproved: {unapproved_count}"
) )
+1 -1
View File
@@ -11,4 +11,4 @@ urllib3==2.5.0
pyperclip==1.11.0 pyperclip==1.11.0
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/ --extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
airlock_libs==3.1.2 airlock_libs==6.0.0
+7 -7
View File
@@ -38,9 +38,9 @@ logger = logging.getLogger(__name__)
def devicehistory(api: AirlockAPIWrapper, outputjson: bool): def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
agents = selectAgents(api) agents = selectAgents(api)
history_days = Selector.select_value( 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, value_type=int,
valid_range=(1, 150), valid_range=(1, 365),
) )
if not agents or not history_days: if not agents or not history_days:
@@ -139,7 +139,7 @@ def findAgents(api, return_dataframe):
print( print(
colorText( colorText(
f"\n✠Matched devices exported to: {working_dir}\\{filename}", f"\n✠Matched devices exported to: {working_dir}\\{filename}",
"green", "green",
) )
) )
@@ -148,7 +148,7 @@ def findAgents(api, return_dataframe):
def collect_device_names() -> List[str]: def collect_device_names() -> List[str]:
print(colorText("🍠Device Search", "cyan")) print(colorText("🖥��Â Device Search", "cyan"))
print( print(
colorText( colorText(
"Enter the device hostnames you'd like to search for, one per line.", "cyan" "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: else:
print( print(
colorText( 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", "yellow",
) )
) )
@@ -265,7 +265,7 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
print(colorText("❌ No matching devices found.", "red")) print(colorText("❌ No matching devices found.", "red"))
return [] 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:") logger.info("Matched agent hostnames:")
rows = (len(matched_agents) + 2) // 3 # 3 columns rows = (len(matched_agents) + 2) // 3 # 3 columns
for row in range(rows): for row in range(rows):
@@ -302,7 +302,7 @@ def moveAgentToRelatedPolicy(
Args: Args:
api: AirlockAPIWrapper instance. api: AirlockAPIWrapper instance.
agent: Agent object. 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. mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
""" """
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}") 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 models.policy import Policy
from services.API import AirlockAPIWrapper from services.API import AirlockAPIWrapper
from utils.configmanager import get_system_json
from utils.setup import get_base_directory 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__) logger = logging.getLogger(__name__)
@@ -233,17 +232,3 @@ def skipback(days):
hex_timestamp = format(timestamp, "08x") hex_timestamp = format(timestamp, "08x")
objectid_hex = hex_timestamp + "0000000000000000" objectid_hex = hex_timestamp + "0000000000000000"
return ObjectId(objectid_hex) 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")) print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
return selected return selected
elif mode == "e": 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] return [item for item in items if item not in selected]
else: else:
print(colorText("⚠️ Invalid mode. Returning all items.", "yellow")) print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
@@ -279,7 +279,9 @@ class Selector:
df = df[columns] df = df[columns]
items = [row for _, row in df.iterrows()] 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( result = Selector._select_from_list(
items, items,
@@ -311,7 +313,9 @@ class Selector:
df = df[columns] df = df[columns]
items = df.to_dict("records") 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 # Show rows first
print(colorText(header, "cyan")) print(colorText(header, "cyan"))
@@ -346,7 +350,7 @@ class Selector:
print(colorText(f"✅ Included {len(selected)} row(s).", "green")) print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
return [pd.Series(row) for row in selected] return [pd.Series(row) for row in selected]
elif mode == "e": 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] return [pd.Series(row) for row in items if row not in selected]
else: else:
print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow")) print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
+3 -3
View File
@@ -118,14 +118,14 @@ def setup():
logging.debug(f"{name.capitalize()} directory ensured at: {path}") logging.debug(f"{name.capitalize()} directory ensured at: {path}")
# Load system config (immutable) # Load system config (immutable)
system_config = load_system_config() load_system_config()
# Configure logging with system-defined log level # Configure logging with system-defined log level
log_level = get_system_value("LOG_LEVEL", str, "INFO") log_level = get_system_value("LOG_LEVEL", str, "INFO")
configure_logging(dirs["logs"], log_level) configure_logging(dirs["logs"], log_level)
# Load user config (mutable) # 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) # Set up .env file - ONLY for WORKING_DIR (runtime-configurable value)
env_path = base_dir / ".env" env_path = base_dir / ".env"
@@ -155,6 +155,6 @@ def setup():
for subfolder in subfolders: for subfolder in subfolders:
subfolder_path = folder_path / subfolder subfolder_path = folder_path / subfolder
subfolder_path.mkdir(parents=True, exist_ok=True) 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") logging.info("✅ Setup complete")