Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 76bd3a6087 | |||
| 6ab9de413f | |||
| e5b4b9d959 | |||
| 1d9caadaf3 |
@@ -30,7 +30,7 @@ from services.API import AirlockAPIWrapper
|
||||
from services.security import getAPI
|
||||
from TUI.TUI import run_Loxide
|
||||
from utils.configmanager import get_system_value
|
||||
from utils.setup import get_base_directory, setup
|
||||
from utils.setup import setup
|
||||
from utils.utils import irtang
|
||||
|
||||
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
||||
@@ -40,7 +40,6 @@ def main():
|
||||
irtang()
|
||||
# Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
|
||||
setup()
|
||||
base_dir = get_base_directory()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
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
|
||||
|
||||
import logging
|
||||
@@ -150,14 +164,11 @@ class AllowlistSelectionWidget(Static):
|
||||
|
||||
# Action buttons at bottom
|
||||
with Horizontal(id="action_buttons"):
|
||||
self.back_btn = Button("⬅ Back", id="back_btn")
|
||||
self.add_btn = Button("➕ Add to Allowlist", id="add_to_allowlist_btn")
|
||||
|
||||
self.back_btn.styles.width = "50%"
|
||||
self.add_btn.styles.width = "50%"
|
||||
self.add_btn.styles.width = "100%"
|
||||
self.add_btn.disabled = True # Disabled until allowlist selected
|
||||
|
||||
yield self.back_btn
|
||||
yield self.add_btn
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
@@ -380,9 +391,9 @@ class AllowlistSelectionWidget(Static):
|
||||
|
||||
if found_col:
|
||||
self.hash_column = found_col
|
||||
preview_lines.append(f"âÅâ Found hash column: **{found_col}**\n")
|
||||
preview_lines.append(f"✅ Found hash column: **{found_col}**\n")
|
||||
else:
|
||||
preview_lines.append("⚠︠**No hash column found**\n")
|
||||
preview_lines.append("❌ **No hash column found**\n")
|
||||
preview_lines.append("Available columns:\n")
|
||||
for col in self.selected_data.columns:
|
||||
if col != "_row_id":
|
||||
@@ -464,7 +475,7 @@ class AllowlistSelectionWidget(Static):
|
||||
self.selected_allowlist = self.allowlists[actual_allowlist_index]
|
||||
self.add_btn.disabled = False
|
||||
self.add_btn.label = (
|
||||
f"â Add to '{self.selected_allowlist.get('name', 'Unknown')}'"
|
||||
f"➕ Add to '{self.selected_allowlist.get('name', 'Unknown')}'"
|
||||
)
|
||||
|
||||
# Update preview with selection
|
||||
@@ -523,11 +534,6 @@ class AllowlistSelectionWidget(Static):
|
||||
btn = getattr(event, "button", None) or getattr(event, "sender", None)
|
||||
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None)
|
||||
|
||||
if btn is self.back_btn or btn_id == "back_btn":
|
||||
await self.app.pop_screen()
|
||||
event.stop()
|
||||
return
|
||||
|
||||
if btn is self.refresh_btn or btn_id == "refresh_allowlists_btn":
|
||||
await self.load_allowlists()
|
||||
event.stop()
|
||||
@@ -553,7 +559,7 @@ class AllowlistSelectionWidget(Static):
|
||||
try:
|
||||
# Disable button during operation
|
||||
self.add_btn.disabled = True
|
||||
self.add_btn.label = "⏳ Adding hashes..."
|
||||
self.add_btn.label = "Adding hashes..."
|
||||
|
||||
# Call API to add hashes
|
||||
app_id = self.selected_allowlist.get("applicationid")
|
||||
@@ -564,10 +570,10 @@ class AllowlistSelectionWidget(Static):
|
||||
)
|
||||
|
||||
result = self.api.hash_add_to_allowlist(app_id, self.hashes_to_add)
|
||||
|
||||
logger.debug(f"Hash adding api call: {result}")
|
||||
# Success notification
|
||||
self.app.notify(
|
||||
f"✅ Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'",
|
||||
f"Successfully added {len(self.hashes_to_add)} hashes to '{allowlist_name}'",
|
||||
title="Success",
|
||||
severity="information",
|
||||
timeout=5,
|
||||
@@ -575,7 +581,7 @@ class AllowlistSelectionWidget(Static):
|
||||
|
||||
# Update preview to show success
|
||||
self.preview_area.text = (
|
||||
f"## ✅ SUCCESS\n\n"
|
||||
f"## SUCCESS\n\n"
|
||||
f"Added **{len(self.hashes_to_add)} hashes** to allowlist:\n"
|
||||
f"**{allowlist_name}** (ID: {app_id})\n\n"
|
||||
f"### Operation Details:\n"
|
||||
@@ -586,13 +592,13 @@ class AllowlistSelectionWidget(Static):
|
||||
)
|
||||
|
||||
# Change button to "Done"
|
||||
self.add_btn.label = "✅ Done"
|
||||
self.add_btn.label = "Done - Press q to return to main menu"
|
||||
self.add_btn.disabled = True
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception(f"Failed to add hashes to allowlist: {exc}")
|
||||
self.app.notify(
|
||||
f"❌ Failed to add hashes: {str(exc)}",
|
||||
f"Failed to add hashes: {str(exc)}",
|
||||
title="Error",
|
||||
severity="error",
|
||||
timeout=10,
|
||||
@@ -600,7 +606,7 @@ class AllowlistSelectionWidget(Static):
|
||||
|
||||
# Re-enable button
|
||||
self.add_btn.disabled = False
|
||||
self.add_btn.label = "⟳ Retry Add to Allowlist"
|
||||
self.add_btn.label = "Retry Add to Allowlist"
|
||||
|
||||
|
||||
class AllowlistSelectionScreen(Screen):
|
||||
@@ -609,9 +615,9 @@ class AllowlistSelectionScreen(Screen):
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("b", "back", "Back"),
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("r", "refresh", "Refresh Allowlists"),
|
||||
Binding("enter", "confirm", "Add to Allowlist"),
|
||||
]
|
||||
|
||||
def __init__(
|
||||
@@ -641,10 +647,15 @@ class AllowlistSelectionScreen(Screen):
|
||||
yield self.widget
|
||||
yield Footer()
|
||||
|
||||
async def action_back(self) -> None:
|
||||
async def action_go_back(self) -> None:
|
||||
"""Go back to previous screen."""
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_refresh(self) -> None:
|
||||
"""Refresh the allowlists."""
|
||||
if hasattr(self, "widget") and self.widget:
|
||||
@@ -0,0 +1,114 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
from typing import List, Optional
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.css.query import NoMatches
|
||||
from textual.screen import Screen
|
||||
|
||||
from models.agent import Agent
|
||||
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
|
||||
from TUI.Widgets.multiagentselector import MultiAgentSelector
|
||||
from TUI.Widgets.resultsdisplay import ResultsDisplay
|
||||
|
||||
|
||||
class MoveAgentWorkflowScreen(Screen):
|
||||
"""Screen that handles the agent movement workflow."""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
def __init__(self, all_agents: Optional[List[Agent]]):
|
||||
super().__init__()
|
||||
self.all_agents = all_agents
|
||||
self.selected_agents = None
|
||||
self.workflow_stage = "select_agents" # Track current stage
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Start with the multi-agent selector."""
|
||||
yield MultiAgentSelector(self.all_agents)
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle escape key to go back one step within the workflow."""
|
||||
if self.workflow_stage == "select_agents":
|
||||
# At first stage, go back to main menu
|
||||
self.app.pop_screen()
|
||||
elif self.workflow_stage == "operations":
|
||||
# Go back to agent selection
|
||||
try:
|
||||
ops_widget = self.query_one(AgentMoveOperations)
|
||||
ops_widget.remove()
|
||||
except NoMatches:
|
||||
pass
|
||||
self.mount(MultiAgentSelector(self.all_agents))
|
||||
self.workflow_stage = "select_agents"
|
||||
elif self.workflow_stage == "results":
|
||||
# Go back to operations
|
||||
try:
|
||||
results_widget = self.query_one(ResultsDisplay)
|
||||
results_widget.remove()
|
||||
except NoMatches:
|
||||
pass
|
||||
self.mount(AgentMoveOperations(self.selected_agents))
|
||||
self.workflow_stage = "operations"
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Handle q key to go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_multi_agent_selector_agents_selected(
|
||||
self, message: MultiAgentSelector.AgentsSelected
|
||||
) -> None:
|
||||
"""Handle selected agents - switch to operations screen."""
|
||||
self.selected_agents = message.selected_agents
|
||||
|
||||
# Remove the MultiAgentSelector
|
||||
selector = self.query_one(MultiAgentSelector)
|
||||
selector.remove()
|
||||
|
||||
# Mount the AgentMoveOperations with the selected Agent objects
|
||||
self.mount(AgentMoveOperations(self.selected_agents))
|
||||
self.workflow_stage = "operations"
|
||||
|
||||
def on_agent_move_operations_operation_complete(
|
||||
self, message: AgentMoveOperations.OperationComplete
|
||||
) -> None:
|
||||
"""Handle completion of move operation - transition to results screen."""
|
||||
# Format successful results
|
||||
success_lines = []
|
||||
for agent, result in message.successful:
|
||||
success_lines.append(f"✔ {agent.hostname}")
|
||||
|
||||
# Format unsuccessful results
|
||||
failure_lines = []
|
||||
for agent, error in message.unsuccessful:
|
||||
failure_lines.append(f"❌ — {agent.hostname}: {error}")
|
||||
|
||||
successful_text = "\n".join(success_lines) if success_lines else "(none)"
|
||||
unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)"
|
||||
|
||||
# Remove the operations widget
|
||||
ops_widget = self.query_one(AgentMoveOperations)
|
||||
ops_widget.remove()
|
||||
|
||||
# Mount the results display
|
||||
self.mount(
|
||||
ResultsDisplay(message.operation, successful_text, unsuccessful_text)
|
||||
)
|
||||
self.workflow_stage = "results"
|
||||
@@ -1,3 +1,17 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
@@ -11,7 +25,7 @@ from textual.containers import Horizontal, Vertical
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Static
|
||||
|
||||
from TUI.allowlistselectionscreen import AllowlistSelectionScreen
|
||||
from TUI.Screens.allowlistselectionscreen import AllowlistSelectionScreen
|
||||
from utils.configmanager import load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -31,8 +45,7 @@ class OTPActivitiesWidget(Static):
|
||||
"""
|
||||
Reusable widget that contains the sessions table (left) and an Activity Preview (right).
|
||||
The right side shows an Activity Preview that takes ~75% vertical space, and a lower area
|
||||
with Back and Continue buttons. The Continue button pushes ActivityDetailScreen with the
|
||||
currently-loaded activities.
|
||||
with Continue button.
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
@@ -85,15 +98,10 @@ class OTPActivitiesWidget(Static):
|
||||
with Vertical(id="activity_preview_container"):
|
||||
self.activities_table = DataTable(id="activity_preview_table")
|
||||
yield self.activities_table
|
||||
# Buttons area at the bottom (Back, Continue)
|
||||
# Button area at the bottom (Continue)
|
||||
with Horizontal(id="activity_buttons"):
|
||||
# Back takes left side, Continue right side
|
||||
self.back_btn = Button("Back", id="activity_back_btn")
|
||||
self.continue_btn = Button("Continue", id="activity_continue_btn")
|
||||
# Stretch buttons nicely
|
||||
self.back_btn.styles.width = "50%"
|
||||
self.continue_btn.styles.width = "50%"
|
||||
yield self.back_btn
|
||||
self.continue_btn.styles.width = "100%"
|
||||
yield self.continue_btn
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
@@ -122,7 +130,7 @@ class OTPActivitiesWidget(Static):
|
||||
|
||||
async def on_button_pressed(self, event) -> None: # type: ignore[override]
|
||||
"""
|
||||
Handle Back / Continue buttons for the Activity Preview area.
|
||||
Handle Continue button for the Activity Preview area.
|
||||
"""
|
||||
# Try to resolve the button object from the event
|
||||
btn = (
|
||||
@@ -136,18 +144,12 @@ class OTPActivitiesWidget(Static):
|
||||
or getattr(event, "button_id", None)
|
||||
or getattr(event, "id", None)
|
||||
)
|
||||
# ---- Back ----
|
||||
if btn is self.back_btn or btn_id == getattr(self.back_btn, "id", None):
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
event.stop()
|
||||
return
|
||||
# ---- Continue ----
|
||||
if btn is self.continue_btn or btn_id == getattr(self.continue_btn, "id", None):
|
||||
if self._activities_df is None or self._activities_df.empty:
|
||||
logger.info("Continue pressed but no activities loaded.")
|
||||
await self.post_message(
|
||||
Static("No activities loaded to continue with.")
|
||||
self.app.notify(
|
||||
"No activities loaded to continue with.", severity="warning"
|
||||
)
|
||||
return
|
||||
# Copy activities DataFrame to pass to new screen
|
||||
@@ -463,7 +465,7 @@ class OTPActivitiesWidget(Static):
|
||||
try:
|
||||
self._activities_df.to_csv(file_path, index=False)
|
||||
logger.info("Exported activities to %s", file_path)
|
||||
await self.post_message(Static(f"✅ Exported activities to: {file_path}"))
|
||||
await self.post_message(Static(f"Exported activities to: {file_path}"))
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to export activities to %s: %s", file_path, exc)
|
||||
await self.post_message(Static("Failed to export activities; check logs."))
|
||||
@@ -472,7 +474,7 @@ class OTPActivitiesWidget(Static):
|
||||
class ActivityDetailWidget(Static):
|
||||
"""
|
||||
Interactive widget for Activity Detail screen.
|
||||
Shows the provided DataFrame in a DataTable and offers Export + Back buttons.
|
||||
Shows the provided DataFrame in a DataTable and offers Export button.
|
||||
Now includes Select All/None and Add to Allowlist functionality.
|
||||
"""
|
||||
|
||||
@@ -529,12 +531,10 @@ class ActivityDetailWidget(Static):
|
||||
|
||||
# Original buttons at bottom
|
||||
with Horizontal(id="detail_buttons"):
|
||||
self.detail_back_btn = Button("Back", id="detail_back_btn")
|
||||
self.add_allowlist_btn = Button(
|
||||
"📋 Add Selected to Allowlist", id="add_allowlist_btn"
|
||||
"Add Selected to Allowlist", id="add_allowlist_btn"
|
||||
)
|
||||
yield self.add_allowlist_btn
|
||||
yield self.detail_back_btn
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
await self._build_table(rebuild=True)
|
||||
@@ -547,12 +547,12 @@ class ActivityDetailWidget(Static):
|
||||
|
||||
# Update button labels with count
|
||||
count = len(self.selected_row_ids)
|
||||
total = len(self.activities_df)
|
||||
len(self.activities_df)
|
||||
|
||||
if has_selection:
|
||||
self.add_allowlist_btn.label = f"📋 Add {count} Selected to Allowlist"
|
||||
self.add_allowlist_btn.label = f"Add {count} Selected to Allowlist"
|
||||
else:
|
||||
self.add_allowlist_btn.label = "📋 Add Selected to Allowlist"
|
||||
self.add_allowlist_btn.label = "Add Selected to Allowlist"
|
||||
|
||||
async def _build_table(self, rebuild: bool = True) -> None:
|
||||
"""Rebuild the DataTable. If rebuild=False, only refresh rows."""
|
||||
@@ -591,7 +591,7 @@ class ActivityDetailWidget(Static):
|
||||
vals.append("" if pd.isna(v) else str(v))
|
||||
|
||||
# Check if this row is selected
|
||||
checkbox = "☑" if row_id in self.selected_row_ids else "☐"
|
||||
checkbox = "☑️" if row_id in self.selected_row_ids else "☐"
|
||||
|
||||
# Add row to table
|
||||
row_key = self.detail_table.add_row(checkbox, *vals)
|
||||
@@ -616,7 +616,7 @@ class ActivityDetailWidget(Static):
|
||||
self.detail_table.update_cell(row_key, "select", "☐") # Unchecked
|
||||
else:
|
||||
self.selected_row_ids.add(row_id)
|
||||
self.detail_table.update_cell(row_key, "select", "☑") # Checked
|
||||
self.detail_table.update_cell(row_key, "select", "☑️") # Checked
|
||||
|
||||
self._update_button_states()
|
||||
|
||||
@@ -652,19 +652,13 @@ class ActivityDetailWidget(Static):
|
||||
logger.exception("Failed to sort by column %s: %s", column_key, exc)
|
||||
return
|
||||
|
||||
# ✅ Only refresh rows, not columns
|
||||
# Only refresh rows, not columns
|
||||
await self._build_table(rebuild=False)
|
||||
|
||||
async def on_button_pressed(self, event) -> None:
|
||||
btn = getattr(event, "button", None) or getattr(event, "sender", None)
|
||||
btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None)
|
||||
|
||||
if btn is self.detail_back_btn or btn_id == "detail_back_btn":
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
event.stop()
|
||||
return
|
||||
|
||||
if btn is self.add_allowlist_btn or btn_id == "add_allowlist_btn":
|
||||
await self._open_allowlist_screen()
|
||||
return
|
||||
@@ -676,7 +670,7 @@ class ActivityDetailWidget(Static):
|
||||
|
||||
# Update all checkboxes in the table
|
||||
for row_key, row_id in self.row_key_to_id.items():
|
||||
self.detail_table.update_cell(row_key, "select", "☑")
|
||||
self.detail_table.update_cell(row_key, "select", "☑️")
|
||||
|
||||
self._update_button_states()
|
||||
logger.info(f"Selected all {len(self.selected_row_ids)} rows")
|
||||
@@ -688,7 +682,7 @@ class ActivityDetailWidget(Static):
|
||||
|
||||
# Update all checkboxes in the table
|
||||
for row_key, row_id in self.row_key_to_id.items():
|
||||
self.detail_table.update_cell(row_key, "select", "☐")
|
||||
self.detail_table.update_cell(row_key, "select", "☑️")
|
||||
|
||||
self._update_button_states()
|
||||
logger.info("Cleared all selections")
|
||||
@@ -730,14 +724,12 @@ class ActivityDetailWidget(Static):
|
||||
async def _export_detail_activities(self) -> None:
|
||||
if self.activities_df is None or self.activities_df.empty:
|
||||
logger.info("No activities to export.")
|
||||
await self.mount(
|
||||
Static("⌠No activities to export.", classes="notification")
|
||||
)
|
||||
await self.mount(Static("No activities to export.", classes="notification"))
|
||||
return
|
||||
if not self.selected_row_ids:
|
||||
logger.info("No rows selected for export.")
|
||||
await self.mount(
|
||||
Static("⌠No rows selected for export.", classes="notification")
|
||||
Static("No rows selected for export.", classes="notification")
|
||||
)
|
||||
return
|
||||
try:
|
||||
@@ -750,7 +742,7 @@ class ActivityDetailWidget(Static):
|
||||
logger.info("Exported selected activities to %s", file_path)
|
||||
await self.mount(
|
||||
Static(
|
||||
f"✅ Exported selected activities to: {filename}",
|
||||
f"Exported selected activities to: {filename}",
|
||||
classes="notification",
|
||||
)
|
||||
)
|
||||
@@ -758,12 +750,12 @@ class ActivityDetailWidget(Static):
|
||||
logger.exception("Failed to export detail activities: %s", exc)
|
||||
await self.mount(
|
||||
Static(
|
||||
"⌠Failed to export activities; check logs.",
|
||||
"¢ Failed to export activities; check logs.",
|
||||
classes="notification",
|
||||
)
|
||||
)
|
||||
|
||||
# ✅ Helper methods
|
||||
# Helper methods
|
||||
def get_selected_data(self) -> pd.DataFrame:
|
||||
"""Return a DataFrame of the selected rows."""
|
||||
if not self.selected_row_ids:
|
||||
@@ -795,7 +787,8 @@ class ActivityDetailScreen(Screen):
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("b", "back", "Back"),
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("e", "export", "Export"),
|
||||
Binding("a", "select_all", "Select All"),
|
||||
Binding("n", "select_none", "Select None"),
|
||||
@@ -819,11 +812,16 @@ class ActivityDetailScreen(Screen):
|
||||
yield self.widget
|
||||
yield Footer()
|
||||
|
||||
async def action_back(self) -> None:
|
||||
async def action_go_back(self) -> None:
|
||||
try:
|
||||
await self.app.pop_screen()
|
||||
except Exception:
|
||||
logger.debug("ActivityDetailScreen.action_back pop_screen failed.")
|
||||
logger.debug("ActivityDetailScreen.action_go_back pop_screen failed.")
|
||||
|
||||
async def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_export(self) -> None:
|
||||
# Delegate to widget export helper
|
||||
@@ -851,9 +849,10 @@ class OTPActivitiesScreen(Screen):
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
Binding("r", "refresh_sessions", "Refresh Sessions"),
|
||||
Binding("e", "export_activities", "Export activities"),
|
||||
Binding("q", "quit", "Quit"),
|
||||
]
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
@@ -875,6 +874,15 @@ class OTPActivitiesScreen(Screen):
|
||||
else:
|
||||
await self.widget.load_sessions_from_api(api)
|
||||
|
||||
async def action_go_back(self) -> None:
|
||||
"""Go back one screen."""
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
await self.app.pop_screen()
|
||||
|
||||
# Simple actions bound to keys
|
||||
async def action_refresh_sessions(self) -> None:
|
||||
api = getattr(self.app, "api", None)
|
||||
@@ -884,10 +892,6 @@ class OTPActivitiesScreen(Screen):
|
||||
logger.info("Refreshing OTP sessions via API.")
|
||||
await self.widget.load_sessions_from_api(api)
|
||||
|
||||
async def action_quit(self) -> None:
|
||||
# Pop the screen or exit app
|
||||
await self.app.pop_screen()
|
||||
|
||||
# If you want an explicit method to fetch activities for a particular otpid from outside:
|
||||
async def fetch_activities_for_otpid(self, otpid, hostname=None) -> None:
|
||||
api = getattr(self.app, "api", None)
|
||||
@@ -0,0 +1,393 @@
|
||||
# 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;
|
||||
padding: 1;
|
||||
align: center middle;
|
||||
}
|
||||
#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"):
|
||||
self.refresh_button = Button("🔄 Refresh", id="refresh_btn")
|
||||
self.refresh_button.styles.width = "15%"
|
||||
self.refresh_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.refresh_button
|
||||
|
||||
self.select_all_button = Button("☑️ Select All", id="select_all_btn")
|
||||
self.select_all_button.styles.width = "15%"
|
||||
self.select_all_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.select_all_button
|
||||
|
||||
self.select_none_button = Button(
|
||||
"❌ Clear Selection", id="select_none_btn"
|
||||
)
|
||||
self.select_none_button.styles.width = "20%"
|
||||
self.select_none_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.select_none_button
|
||||
|
||||
self.revoke_button = Button(
|
||||
"🛑 Revoke Selected", id="revoke_btn", variant="error"
|
||||
)
|
||||
self.revoke_button.styles.width = "20%"
|
||||
self.revoke_button.styles.margin = (1, 1, 1, 1)
|
||||
yield self.revoke_button
|
||||
|
||||
# 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:
|
||||
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 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:
|
||||
self.widget._selected_otpids = set(
|
||||
str(x) for x in self.widget._filtered_df["otpid"]
|
||||
)
|
||||
await self.widget._refresh_table()
|
||||
|
||||
async def action_select_none(self) -> None:
|
||||
"""Clear all selections."""
|
||||
self.widget._selected_otpids.clear()
|
||||
await self.widget._refresh_table()
|
||||
|
||||
async def action_revoke(self) -> None:
|
||||
"""Revoke selected sessions."""
|
||||
await self.widget._revoke_selected()
|
||||
|
||||
async def action_go_back(self) -> None:
|
||||
"""Go back to previous screen."""
|
||||
await self.app.pop_screen()
|
||||
|
||||
async def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
await self.app.pop_screen()
|
||||
@@ -0,0 +1,52 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.screen import Screen
|
||||
|
||||
from models.agent import Agent
|
||||
from TUI.Widgets.OTP_generate import OTPGenerator
|
||||
|
||||
|
||||
class OTPWorkflowScreen(Screen):
|
||||
"""Screen that handles the OTP generation workflow without agent selection."""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
def __init__(self, selected_agents: Optional[List[Agent]]):
|
||||
super().__init__()
|
||||
self.selected_agents = selected_agents
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Directly show the OTP generator for the selected agents."""
|
||||
yield OTPGenerator(self.selected_agents)
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle escape key to go back one screen."""
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Handle q key to go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
|
||||
"""Handle OTP generation request - pass it up to the app level if needed."""
|
||||
File diff suppressed because it is too large
Load Diff
@@ -23,9 +23,11 @@ the policy selection workflow.
|
||||
import logging
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import Footer, Header
|
||||
|
||||
from TUI.policyselector import PolicySelector
|
||||
from TUI.Widgets.policyselector import PolicySelector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -42,6 +44,11 @@ class PolicySelectorScreen(Screen):
|
||||
agent_move_operations: Reference to the parent AgentMoveOperations widget.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "go_back", "Back"),
|
||||
Binding("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
CSS = """
|
||||
Screen {
|
||||
layout: vertical;
|
||||
@@ -68,7 +75,18 @@ class PolicySelectorScreen(Screen):
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Create the PolicySelector widget."""
|
||||
yield Header(show_clock=True)
|
||||
yield PolicySelector(self.policies)
|
||||
yield Footer()
|
||||
|
||||
def action_go_back(self) -> None:
|
||||
"""Handle escape key to go back one screen."""
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Handle q key to go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
def on_policy_selector_policy_selected(
|
||||
self, message: PolicySelector.PolicySelected
|
||||
@@ -39,7 +39,7 @@ from textual.widgets import Button, DataTable, Footer, Header, Static
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.policyhandler import getPolicyInfo
|
||||
from TUI.policyselector import PolicySelector
|
||||
from TUI.Widgets.policyselector import PolicySelector
|
||||
from utils.configmanager import load_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -69,6 +69,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
|
||||
BINDINGS = [
|
||||
("escape", "go_back", "Back"),
|
||||
("q", "main_menu", "Main Menu"),
|
||||
]
|
||||
|
||||
workflow_stage = reactive("select_policy") # Tracks current workflow stage
|
||||
@@ -96,10 +97,10 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Build the UI layout for the workflow screen."""
|
||||
# Include Header and Footer like other standalone screens
|
||||
yield Header(show_clock=True, icon="⚙")
|
||||
yield Header(show_clock=True, icon="⚙️")
|
||||
|
||||
# Title area
|
||||
title = Static("🔒 Quiet Agent Workflow", id="workflow_title")
|
||||
title = Static("Quiet Agent Workflow", id="workflow_title")
|
||||
title.styles.margin = (0, 0, 0, 1)
|
||||
yield title
|
||||
|
||||
@@ -130,13 +131,13 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
stage_messages = {
|
||||
"select_policy": "Step 1: Select Policy to Analyze",
|
||||
"select_quiet_days": "Step 2: Select Quiet Time Period",
|
||||
"analyzing": "📊 Analyzing agent activity...",
|
||||
"analyzing": "Analyzing agent activity...",
|
||||
"view_results": "Step 3: Review Categorized Agents",
|
||||
"select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents",
|
||||
"select_non_enforce_target": "Step 5: Select Target Policy for Non-Enforce Ready Agents",
|
||||
"confirm_migration": "Step 6: Confirm and Execute Migration",
|
||||
"executing": "⏳ Executing agent migrations...",
|
||||
"complete": "✅ Migration Complete",
|
||||
"executing": "Executing agent migrations...",
|
||||
"complete": "Migration Complete",
|
||||
}
|
||||
|
||||
status_widget.update(stage_messages.get(self.workflow_stage, "Unknown Stage"))
|
||||
@@ -208,11 +209,6 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
btn.styles.margin = (0, 0, 1, 0)
|
||||
button_container.mount(btn)
|
||||
|
||||
back_btn = Button("← Back", id="back_to_policy_selection")
|
||||
back_btn.styles.width = "100%"
|
||||
back_btn.styles.margin = (2, 0, 0, 0)
|
||||
button_container.mount(back_btn)
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button press events."""
|
||||
button_id = event.button.id
|
||||
@@ -226,14 +222,6 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
return
|
||||
|
||||
# Navigation buttons
|
||||
if button_id == "back_to_policy_selection":
|
||||
self._show_policy_selection()
|
||||
return
|
||||
|
||||
if button_id == "back_to_results":
|
||||
self._show_results()
|
||||
return
|
||||
|
||||
if button_id == "select_enforce_target_btn":
|
||||
self._show_enforce_target_selection()
|
||||
return
|
||||
@@ -278,16 +266,16 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
|
||||
# Show analyzing message with detailed steps
|
||||
analyzing_msg = Static(
|
||||
f"📊 Analyzing Agent Activity\n"
|
||||
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"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",
|
||||
)
|
||||
@@ -405,8 +393,8 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
self.agents_df = agents
|
||||
|
||||
# Categorize agents into DataFrames
|
||||
self.enforce_ready_df = agents[agents["enforce_ready"] == True].copy()
|
||||
self.non_enforce_ready_df = agents[agents["enforce_ready"] == False].copy()
|
||||
self.enforce_ready_df = agents[agents["enforce_ready"]].copy()
|
||||
self.non_enforce_ready_df = agents[not agents["enforce_ready"]].copy()
|
||||
|
||||
logger.info(
|
||||
f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, "
|
||||
@@ -435,13 +423,13 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
|
||||
# Build updated message
|
||||
updated_text = (
|
||||
f"📊 Analyzing Agent Activity\n"
|
||||
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"{status_text}\n\n"
|
||||
f"Please wait - this operation cannot be cancelled."
|
||||
)
|
||||
|
||||
@@ -469,9 +457,9 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
|
||||
summary = Static(
|
||||
f"Analysis Results for: {self.selected_policy.name}\n\n"
|
||||
f"📊 Total Agents: {total_agents}\n"
|
||||
f"✅ Enforce Ready: {ready_count} ({ready_percentage:.1f}%)\n"
|
||||
f"❌ Not Ready: {not_ready_count} ({100 - ready_percentage:.1f}%)\n\n"
|
||||
f"Total Agents: {total_agents}\n"
|
||||
f"Enforce Ready: {ready_count} ({ready_percentage:.1f}%)\n"
|
||||
f"Not Ready: {not_ready_count} ({100 - ready_percentage:.1f}%)\n\n"
|
||||
f"Quiet Threshold: {self.quiet_days} days\n"
|
||||
f"History Period: {self.history_days} days",
|
||||
id="results_summary",
|
||||
@@ -500,11 +488,11 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
non_enforce_btn.styles.margin = (0, 1, 1, 0)
|
||||
button_container.mount(non_enforce_btn)
|
||||
|
||||
export_btn = Button("💾 Export Results", id="export_results_btn")
|
||||
export_btn = Button("Export Results", id="export_results_btn")
|
||||
export_btn.styles.margin = (0, 1, 1, 0)
|
||||
button_container.mount(export_btn)
|
||||
|
||||
start_over_btn = Button("🔄 Start Over", id="start_over_btn")
|
||||
start_over_btn = Button("Start Over", id="start_over_btn")
|
||||
start_over_btn.styles.margin = (0, 0, 1, 0)
|
||||
button_container.mount(start_over_btn)
|
||||
|
||||
@@ -520,7 +508,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
enforce_col.styles.margin = (1, 1, 0, 0)
|
||||
tables_container.mount(enforce_col)
|
||||
|
||||
enforce_label = Static("✅ Enforce Ready Agents")
|
||||
enforce_label = Static("Enforce Ready Agents")
|
||||
enforce_label.styles.margin = (0, 0, 1, 0)
|
||||
enforce_col.mount(enforce_label)
|
||||
|
||||
@@ -548,7 +536,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
non_enforce_col.styles.margin = (1, 0, 0, 1)
|
||||
tables_container.mount(non_enforce_col)
|
||||
|
||||
non_enforce_label = Static("❌ Non-Enforce Ready Agents")
|
||||
non_enforce_label = Static("Non-Enforce Ready Agents")
|
||||
non_enforce_label.styles.margin = (0, 0, 1, 0)
|
||||
non_enforce_col.mount(non_enforce_label)
|
||||
|
||||
@@ -589,7 +577,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
content.mount(policy_selector)
|
||||
|
||||
# Skip button
|
||||
skip_btn = Button("⭕️ Skip - No Migration", id="skip_enforce_target_btn")
|
||||
skip_btn = Button("Skip - No Migration", id="skip_enforce_target_btn")
|
||||
skip_btn.styles.width = "50%"
|
||||
skip_btn.styles.margin = (2, 0, 0, 0)
|
||||
content.mount(skip_btn)
|
||||
@@ -614,7 +602,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
content.mount(policy_selector)
|
||||
|
||||
# Skip button
|
||||
skip_btn = Button("⭕️ Skip - No Migration", id="skip_non_enforce_target_btn")
|
||||
skip_btn = Button("Skip - No Migration", id="skip_non_enforce_target_btn")
|
||||
skip_btn.styles.width = "50%"
|
||||
skip_btn.styles.margin = (2, 0, 0, 0)
|
||||
content.mount(skip_btn)
|
||||
@@ -627,29 +615,29 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
|
||||
# Build confirmation message
|
||||
confirmation_lines = [
|
||||
"🔐 Migration Summary\n",
|
||||
"Migration Summary\n",
|
||||
f"Source Policy: {self.selected_policy.name}\n",
|
||||
]
|
||||
|
||||
if self.enforce_ready_target_policy:
|
||||
confirmation_lines.append(
|
||||
f"\n✅ Enforce Ready Migration:\n"
|
||||
f" • Agents: {len(self.enforce_ready_df)}\n"
|
||||
f" • Target: {self.enforce_ready_target_policy.name}\n"
|
||||
f"\nEnforce Ready Migration:\n"
|
||||
f"Agents: {len(self.enforce_ready_df)}\n"
|
||||
f"Target: {self.enforce_ready_target_policy.name}\n"
|
||||
)
|
||||
|
||||
if self.non_enforce_ready_target_policy:
|
||||
confirmation_lines.append(
|
||||
f"\n❌ Non-Enforce Ready Migration:\n"
|
||||
f" • Agents: {len(self.non_enforce_ready_df)}\n"
|
||||
f" • Target: {self.non_enforce_ready_target_policy.name}\n"
|
||||
f"\nNon-Enforce Ready Migration:\n"
|
||||
f"Agents: {len(self.non_enforce_ready_df)}\n"
|
||||
f"Target: {self.non_enforce_ready_target_policy.name}\n"
|
||||
)
|
||||
|
||||
if (
|
||||
not self.enforce_ready_target_policy
|
||||
and not self.non_enforce_ready_target_policy
|
||||
):
|
||||
confirmation_lines.append("\n⚠️ No migrations will be performed.")
|
||||
confirmation_lines.append("\nNo migrations will be performed.")
|
||||
|
||||
confirmation = Static("".join(confirmation_lines), id="migration_confirmation")
|
||||
confirmation.styles.margin = (1, 1, 2, 1)
|
||||
@@ -662,11 +650,11 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
content.mount(button_container)
|
||||
|
||||
if self.enforce_ready_target_policy or self.non_enforce_ready_target_policy:
|
||||
confirm_btn = Button("✅ Confirm Migration", id="confirm_migration_btn")
|
||||
confirm_btn = Button("Confirm Migration", id="confirm_migration_btn")
|
||||
confirm_btn.styles.margin = (0, 1, 0, 0)
|
||||
button_container.mount(confirm_btn)
|
||||
|
||||
cancel_btn = Button("❌ Cancel", id="cancel_migration_btn")
|
||||
cancel_btn = Button("Cancel", id="cancel_migration_btn")
|
||||
button_container.mount(cancel_btn)
|
||||
|
||||
def _execute_migration(self) -> None:
|
||||
@@ -677,7 +665,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
|
||||
# Show executing message
|
||||
executing_msg = Static(
|
||||
"⏳ Executing agent migrations...\nPlease wait...",
|
||||
"Executing agent migrations...\nPlease wait...",
|
||||
id="executing_message",
|
||||
)
|
||||
executing_msg.styles.margin = (2, 1)
|
||||
@@ -696,7 +684,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
if self.enforce_ready_target_policy:
|
||||
for idx, row in self.enforce_ready_df.iterrows():
|
||||
try:
|
||||
result = self.api.agent_move(
|
||||
self.api.agent_move(
|
||||
row["agentid"], self.enforce_ready_target_policy.groupid
|
||||
)
|
||||
successful_migrations.append(
|
||||
@@ -713,7 +701,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
if self.non_enforce_ready_target_policy:
|
||||
for idx, row in self.non_enforce_ready_df.iterrows():
|
||||
try:
|
||||
result = self.api.agent_move(
|
||||
self.api.agent_move(
|
||||
row["agentid"], self.non_enforce_ready_target_policy.groupid
|
||||
)
|
||||
successful_migrations.append(
|
||||
@@ -749,7 +737,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
)
|
||||
|
||||
results = Static(
|
||||
f"✅ Migration Complete\n\n"
|
||||
f"Migration Complete\n\n"
|
||||
f"Total Agents Migrated: {len(successful)}\n"
|
||||
f"Failed Migrations: {len(failed)}\n"
|
||||
f"Success Rate: {success_rate:.1f}%",
|
||||
@@ -764,7 +752,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
success_container.styles.margin = (0, 1)
|
||||
content.mount(success_container)
|
||||
|
||||
success_label = Static("✅ Successful Migrations")
|
||||
success_label = Static("Successful Migrations")
|
||||
success_label.styles.margin = (0, 0, 1, 0)
|
||||
success_container.mount(success_label)
|
||||
|
||||
@@ -785,7 +773,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
failed_container.styles.margin = (2, 1, 0, 1)
|
||||
content.mount(failed_container)
|
||||
|
||||
failed_label = Static("❌ Failed Migrations")
|
||||
failed_label = Static("Failed Migrations")
|
||||
failed_label.styles.margin = (0, 0, 1, 0)
|
||||
failed_container.mount(failed_label)
|
||||
|
||||
@@ -802,7 +790,7 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
failed_container.mount(failed_table)
|
||||
|
||||
# Action button
|
||||
done_btn = Button("✔ Done", id="start_over_btn")
|
||||
done_btn = Button("Done", id="start_over_btn")
|
||||
done_btn.styles.width = "50%"
|
||||
done_btn.styles.margin = (2, 0, 0, 0)
|
||||
content.mount(done_btn)
|
||||
@@ -846,3 +834,8 @@ class QuietAgentWorkflowScreen(Screen):
|
||||
self._show_non_enforce_target_selection()
|
||||
else:
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_main_menu(self) -> None:
|
||||
"""Go back to main menu."""
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
+63
-130
@@ -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 os
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
import dotenv
|
||||
@@ -19,24 +33,22 @@ from textual.widgets import (
|
||||
Tabs,
|
||||
)
|
||||
|
||||
from flows.otp import otp_revoke
|
||||
from flows.prepPolicy import menu_policy_enforce
|
||||
from models.agent import Agent
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.policyhandler import confirmUpdateAfromE
|
||||
from TUI.agentmoveoperations import AgentMoveOperations
|
||||
from TUI.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
||||
from TUI.multiagentselector import MultiAgentSelector
|
||||
from TUI.OTP_generate import OTPGenerator
|
||||
from TUI.otpactivityscreen import OTPActivitiesScreen
|
||||
from TUI.otpworkflowscreen import OTPWorkflowScreen
|
||||
from TUI.policytreewidget import PolicyTreeWidget
|
||||
from TUI.quietagentworkflowscreen import QuietAgentWorkflowScreen
|
||||
from TUI.resultsdisplay import ResultsDisplay
|
||||
from TUI.theme_amber_terminal import get_amber_terminal_theme
|
||||
from TUI.theme_retro_terminal import get_retro_terminal_theme
|
||||
from TUI.themeselector import ThemeSelector
|
||||
from TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen
|
||||
from TUI.Screens.otpactivityscreen import OTPActivitiesScreen
|
||||
from TUI.Screens.otprevokescreen import OTPRevokeScreen
|
||||
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
|
||||
from TUI.Screens.policyprepworkflowscreen import PolicyPrepWorkflowScreen
|
||||
from TUI.Screens.quietagentworkflowscreen import QuietAgentWorkflowScreen
|
||||
from TUI.Themes.theme_amber_terminal import get_amber_terminal_theme
|
||||
from TUI.Themes.theme_retro_terminal import get_retro_terminal_theme
|
||||
from TUI.Themes.themeselector import ThemeSelector
|
||||
from TUI.Widgets.agentmoveoperations import AgentMoveOperations
|
||||
from TUI.Widgets.multiagentselector import MultiAgentSelector
|
||||
from TUI.Widgets.policytreewidget import PolicyTreeWidget
|
||||
from TUI.Widgets.resultsdisplay import ResultsDisplay
|
||||
from utils.configmanager import get_user_value, load_env, save_user_config
|
||||
from utils.setup import get_base_directory
|
||||
from utils.utils import open_directory
|
||||
@@ -47,7 +59,7 @@ dotenv.load_dotenv()
|
||||
# GLOBAL STASH
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PENDING_JOB = None
|
||||
_APP_RESTART_REASON = None
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -83,13 +95,12 @@ class MainMenuScreen(Screen):
|
||||
"🖥️ - Find, Move, or Generate OTP for Agents",
|
||||
"move_agent_workflow_button",
|
||||
),
|
||||
("📊 - Review and appove OTP Activities", "otp_activities_button"),
|
||||
("📇 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"),
|
||||
("🎫 - Review and appove OTP Activities", "otp_activities_button"),
|
||||
("🔕 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"),
|
||||
],
|
||||
"policy": [
|
||||
("🔒 - Prepare Policy For Enforcement", "policy_prep_button"),
|
||||
("🔄 - Update Audit Policies", "policy_audit_update_button"),
|
||||
("❌ - Revoke OTPs", "otp_revoke_button"),
|
||||
("⚖️ - Prepare Policy For Enforcement", "policy_prep_button"),
|
||||
("🛑 - Revoke OTPs", "otp_revoke_button"),
|
||||
],
|
||||
}
|
||||
|
||||
@@ -128,7 +139,6 @@ class MainMenuScreen(Screen):
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
api = self.app.api
|
||||
self.switch_tab("agent_actions")
|
||||
|
||||
# focus helpers
|
||||
@@ -191,42 +201,20 @@ class MainMenuScreen(Screen):
|
||||
self, message: MultiAgentSelector.AgentsSelected
|
||||
) -> None:
|
||||
"""Handle selected agents from AgentSelector."""
|
||||
global _PENDING_JOB
|
||||
global _APP_RESTART_REASON
|
||||
selected_agents = message.selected_agents
|
||||
logger.info("Selected agents: %s", selected_agents)
|
||||
# TODO: Implement actual handling of selected agents
|
||||
_PENDING_JOB = ("multi_agent_action", selected_agents)
|
||||
_APP_RESTART_REASON = ("multi_agent_action", selected_agents)
|
||||
self.app.exit()
|
||||
|
||||
def on_theme_selector_theme_selected(
|
||||
self, message: ThemeSelector.ThemeSelected
|
||||
) -> None:
|
||||
"""Handle theme selection from ThemeSelector."""
|
||||
global _PENDING_JOB
|
||||
global _APP_RESTART_REASON
|
||||
_persist_user_theme(message.theme_name)
|
||||
_PENDING_JOB = ("restart",)
|
||||
self.app.exit()
|
||||
|
||||
def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None:
|
||||
"""Handle OTP generation request from the workflow."""
|
||||
global _PENDING_JOB
|
||||
|
||||
logger.info(
|
||||
"OTP Generation requested: %d devices, requestor=%s, reason=%s, duration=%d",
|
||||
len(message.devices),
|
||||
message.requestor,
|
||||
message.reasoning,
|
||||
message.duration,
|
||||
)
|
||||
|
||||
_PENDING_JOB = (
|
||||
"otp_workflow",
|
||||
message.devices,
|
||||
message.requestor,
|
||||
message.reasoning,
|
||||
message.duration,
|
||||
)
|
||||
|
||||
_APP_RESTART_REASON = ("restart",)
|
||||
self.app.exit()
|
||||
|
||||
def on_agent_move_operations_operation_complete(
|
||||
@@ -282,7 +270,6 @@ class MainMenuScreen(Screen):
|
||||
self.app.bell()
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
global _PENDING_JOB
|
||||
button_id = event.button.id
|
||||
logger.debug("Button pressed: %s", button_id)
|
||||
|
||||
@@ -308,27 +295,23 @@ class MainMenuScreen(Screen):
|
||||
return
|
||||
|
||||
case "otp_revoke_button":
|
||||
_PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {})
|
||||
self.app.push_screen(OTPRevokeScreen())
|
||||
event.stop()
|
||||
return
|
||||
|
||||
case "policy_prep_button":
|
||||
_PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {})
|
||||
|
||||
case "policy_audit_update_button":
|
||||
_PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {})
|
||||
# Use the new TUI workflow screen instead of legacy
|
||||
self.app.push_screen(
|
||||
PolicyPrepWorkflowScreen(self.app.api, self.app.policies)
|
||||
)
|
||||
event.stop()
|
||||
return
|
||||
|
||||
case _:
|
||||
self.app.bell()
|
||||
logger.warning("Unknown button pressed: %s", button_id)
|
||||
return
|
||||
|
||||
# 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
|
||||
@@ -353,7 +336,7 @@ class Loxide(App[Message]):
|
||||
]
|
||||
|
||||
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__()
|
||||
self.api = api
|
||||
wd = load_env("WORKING_DIR") or os.getcwd()
|
||||
@@ -395,8 +378,8 @@ class Loxide(App[Message]):
|
||||
self.refresh_data()
|
||||
|
||||
def action_quit(self) -> None:
|
||||
global _PENDING_JOB
|
||||
_PENDING_JOB = None
|
||||
global _APP_RESTART_REASON
|
||||
_APP_RESTART_REASON = None
|
||||
self.exit()
|
||||
|
||||
def action_open_fe(self) -> None:
|
||||
@@ -410,45 +393,10 @@ class Loxide(App[Message]):
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3) TERMINAL + LEGACY
|
||||
# ---------------------------------------------------------------------------
|
||||
def _restore_terminal_for_legacy() -> None:
|
||||
sys.stdout.write("\033[?1049l")
|
||||
sys.stdout.write("\033[?25h")
|
||||
sys.stdout.write("\033[0m")
|
||||
sys.stdout.write("\033[?1000l\033[?1002l\033[?1003l\033[?1006l")
|
||||
sys.stdout.write("\033[2J\033[H")
|
||||
sys.stdout.flush()
|
||||
if os.name == "nt":
|
||||
try:
|
||||
import ctypes
|
||||
|
||||
kernel32 = ctypes.windll.kernel32
|
||||
handle = kernel32.GetStdHandle(-11)
|
||||
mode = ctypes.c_ulong()
|
||||
if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
|
||||
kernel32.SetConsoleMode(handle, mode.value | 0x0004)
|
||||
except Exception as exc:
|
||||
logger.debug("VT enable on Windows failed: %s", exc)
|
||||
|
||||
|
||||
def _run_legacy_job(func, args, kwargs) -> None:
|
||||
logger.debug("Running legacy job: %s", getattr(func, "__name__", func))
|
||||
_restore_terminal_for_legacy()
|
||||
try:
|
||||
func(*args, **kwargs)
|
||||
finally:
|
||||
try:
|
||||
input("\nPress Enter to return to the UI...")
|
||||
except EOFError:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4) PUBLIC ENTRYPOINT
|
||||
# 3) PUBLIC ENTRYPOINT
|
||||
# ---------------------------------------------------------------------------
|
||||
def run_Loxide(api: AirlockAPIWrapper) -> None:
|
||||
global _PENDING_JOB
|
||||
global _APP_RESTART_REASON
|
||||
base_dir = get_base_directory()
|
||||
env_path = base_dir / ".env"
|
||||
dotenv.load_dotenv(dotenv_path=env_path, override=True)
|
||||
@@ -458,8 +406,8 @@ def run_Loxide(api: AirlockAPIWrapper) -> None:
|
||||
|
||||
while attempts < max_attempts:
|
||||
attempts += 1
|
||||
logger.debug("Starting job loop iteration (attempt %d)", attempts)
|
||||
_PENDING_JOB = None
|
||||
logger.debug("Starting app loop iteration (attempt %d)", attempts)
|
||||
_APP_RESTART_REASON = None
|
||||
app = Loxide(api)
|
||||
|
||||
try:
|
||||
@@ -469,42 +417,27 @@ def run_Loxide(api: AirlockAPIWrapper) -> None:
|
||||
logger.debug("Caught SystemExit from Textual: %s", exc)
|
||||
raise
|
||||
|
||||
job = _PENDING_JOB
|
||||
logger.debug("After app.run(), _PENDING_JOB = %r", job)
|
||||
reason = _APP_RESTART_REASON
|
||||
logger.debug("After app.run(), _APP_RESTART_REASON = %r", reason)
|
||||
|
||||
if not job:
|
||||
logger.debug("No job pending, exiting loop")
|
||||
if not reason:
|
||||
logger.debug("No restart reason, exiting loop")
|
||||
break
|
||||
|
||||
if job[0] == "legacy":
|
||||
_, func, args, kwargs = job
|
||||
_run_legacy_job(func, args, kwargs)
|
||||
if reason[0] == "restart":
|
||||
logger.debug("Restarting app loop")
|
||||
continue
|
||||
|
||||
if job[0] == "restart":
|
||||
logger.debug("Restarting job loop")
|
||||
if reason[0] == "multi_agent_action":
|
||||
logger.info("Multi-agent action with selected agents: %s", reason[1])
|
||||
continue
|
||||
|
||||
if job[0] == "multi_agent_action":
|
||||
logger.info("Multi-agent action with selected agents: %s", job[1])
|
||||
continue
|
||||
|
||||
if job[0] == "otp_workflow":
|
||||
_, devices, requestor, reasoning, duration = job
|
||||
|
||||
def otp_generate_with_params():
|
||||
# Your OTP logic here
|
||||
pass
|
||||
|
||||
_run_legacy_job(otp_generate_with_params, (), {})
|
||||
continue
|
||||
|
||||
logger.error("Unknown job type: %r", job)
|
||||
logger.error("Unknown restart reason: %r", reason)
|
||||
break
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5) DEV
|
||||
# 4) DEV
|
||||
# ---------------------------------------------------------------------------
|
||||
if __name__ == "__main__":
|
||||
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.theme import Theme
|
||||
|
||||
@@ -1,3 +1,18 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
from textual.color import Color
|
||||
|
||||
|
||||
@@ -1,3 +1,18 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
from textual.containers import Vertical
|
||||
from textual.message import Message
|
||||
from textual.widget import Widget
|
||||
@@ -1,3 +1,18 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
@@ -23,6 +38,8 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OTPGenerator(Widget):
|
||||
"""Widget for generating OTPs for selected devices."""
|
||||
|
||||
# Reactive properties to track form completion
|
||||
requestor_filled = reactive(False)
|
||||
reasoning_filled = reactive(False)
|
||||
@@ -139,14 +156,10 @@ class OTPGenerator(Widget):
|
||||
button_row.styles.height = "auto"
|
||||
button_row.styles.margin = (1, 0, 0, 0)
|
||||
|
||||
back_button = Button("← Back", id="back_button")
|
||||
back_button.styles.width = "1fr"
|
||||
yield back_button
|
||||
|
||||
generate_button = Button(
|
||||
"Generate OTP", id="generate_button", variant="primary"
|
||||
)
|
||||
generate_button.styles.width = "2fr"
|
||||
generate_button.styles.width = "100%"
|
||||
yield generate_button
|
||||
|
||||
# Right side - Show device list initially, then output after generation
|
||||
@@ -169,7 +182,7 @@ class OTPGenerator(Widget):
|
||||
|
||||
# Show device list initially
|
||||
device_list_text = "\n".join(
|
||||
f"• {device.hostname}" for device in self.devices
|
||||
f"{device.hostname}" for device in self.devices
|
||||
)
|
||||
device_display = Static(device_list_text, id="device_display")
|
||||
yield device_display
|
||||
@@ -197,14 +210,7 @@ class OTPGenerator(Widget):
|
||||
def on_button_pressed(self, event: Button.Pressed):
|
||||
btn_id = event.button.id
|
||||
|
||||
if btn_id == "back_button":
|
||||
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "copy_clipboard_button":
|
||||
if btn_id == "copy_clipboard_button":
|
||||
try:
|
||||
output_area = self.query_one("#otp_output", TextArea)
|
||||
text_to_copy = output_area.text
|
||||
@@ -213,7 +219,7 @@ class OTPGenerator(Widget):
|
||||
|
||||
pyperclip.copy(text_to_copy)
|
||||
self.app.notify(
|
||||
"✅ Copied to clipboard!", severity="information", timeout=2
|
||||
"✓ Copied to clipboard!", severity="information", timeout=2
|
||||
)
|
||||
except ImportError:
|
||||
self.app.notify(
|
||||
@@ -1,3 +1,19 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime
|
||||
import logging
|
||||
@@ -10,12 +26,12 @@ from textual.css.query import NoMatches
|
||||
from textual.message import Message
|
||||
from textual.reactive import reactive
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Button, DataTable, Header, Static, TextArea
|
||||
from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea
|
||||
|
||||
from models.agent import Agent
|
||||
from TUI.OTP_generate import OTPGenerator
|
||||
from TUI.otpworkflowscreen import OTPWorkflowScreen
|
||||
from TUI.policyselectorscreen import PolicySelectorScreen
|
||||
from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen
|
||||
from TUI.Screens.policyselectorscreen import PolicySelectorScreen
|
||||
from TUI.Widgets.OTP_generate import OTPGenerator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -183,21 +199,21 @@ class AgentMoveOperations(Widget):
|
||||
f"Operation: {operation_name}",
|
||||
f"{'=' * 50}",
|
||||
"",
|
||||
f"✅ Successful ({len(successful)}):",
|
||||
f"✅ Successful ({len(successful)}):",
|
||||
]
|
||||
|
||||
if successful:
|
||||
for agent, result in successful:
|
||||
results_lines.append(f" ✅ {agent.hostname}")
|
||||
results_lines.append(f" ✅ {agent.hostname}")
|
||||
else:
|
||||
results_lines.append(" (none)")
|
||||
|
||||
results_lines.append("")
|
||||
results_lines.append(f"⌠Failed ({len(unsuccessful)}):")
|
||||
results_lines.append(f"❌ Failed ({len(unsuccessful)}):")
|
||||
|
||||
if unsuccessful:
|
||||
for agent, error in unsuccessful:
|
||||
results_lines.append(f" ⌠{agent.hostname}: {error}")
|
||||
results_lines.append(f" ❌ {agent.hostname}: {error}")
|
||||
else:
|
||||
results_lines.append(" (none)")
|
||||
|
||||
@@ -232,9 +248,9 @@ class AgentMoveOperations(Widget):
|
||||
- Operations panel: 1/3 width
|
||||
- Results area: Initially hidden, shown after operation completion
|
||||
"""
|
||||
yield Header(show_clock=True, icon="âš™")
|
||||
yield Header(show_clock=True, icon="⚙️")
|
||||
title_text = Static(
|
||||
f"ðŸ–¥ï¸ Agent Operations - {len(self.agents)} device(s) selected",
|
||||
f"🖥️ Agent Operations - {len(self.agents)} device(s) selected",
|
||||
id="move_ops_title",
|
||||
)
|
||||
title_text.styles.margin = (0, 0, 1, 0)
|
||||
@@ -269,32 +285,32 @@ class AgentMoveOperations(Widget):
|
||||
yield operations_label
|
||||
|
||||
# Operation buttons
|
||||
export_csv_btn = Button("📈 Export CSV", id="export_csv_btn")
|
||||
export_csv_btn = Button("📄 Export CSV", id="export_csv_btn")
|
||||
export_csv_btn.styles.width = "100%"
|
||||
export_csv_btn.styles.margin = (0, 0, 1, 0)
|
||||
yield export_csv_btn
|
||||
|
||||
local_approval_btn = Button(
|
||||
"âœ”ï¸ Local Approval Mode", id="local_approval_btn"
|
||||
"✔️ Local Approval Mode", id="local_approval_btn"
|
||||
)
|
||||
local_approval_btn.styles.width = "100%"
|
||||
local_approval_btn.styles.margin = (0, 0, 1, 0)
|
||||
yield local_approval_btn
|
||||
|
||||
otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn")
|
||||
otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn")
|
||||
otp_gen_btn.styles.width = "100%"
|
||||
otp_gen_btn.styles.margin = (0, 0, 1, 0)
|
||||
yield otp_gen_btn
|
||||
|
||||
toggle_enforcement_btn = Button(
|
||||
"🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn"
|
||||
"🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn"
|
||||
)
|
||||
toggle_enforcement_btn.styles.width = "100%"
|
||||
toggle_enforcement_btn.styles.margin = (0, 0, 1, 0)
|
||||
yield toggle_enforcement_btn
|
||||
|
||||
other_policy_btn = Button(
|
||||
"🔀 Move to Other Policy", id="other_policy_btn"
|
||||
"🔀 Move to Other Policy", id="other_policy_btn"
|
||||
)
|
||||
other_policy_btn.styles.width = "100%"
|
||||
other_policy_btn.styles.margin = (0, 0, 1, 0)
|
||||
@@ -305,10 +321,7 @@ class AgentMoveOperations(Widget):
|
||||
status_label.styles.margin = (2, 0, 0, 0)
|
||||
yield status_label
|
||||
|
||||
back_button = Button("↠Back", id="back_button")
|
||||
back_button.styles.width = "50%"
|
||||
back_button.styles.margin = (0, 1, 1, 0)
|
||||
yield back_button
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""
|
||||
@@ -343,7 +356,7 @@ class AgentMoveOperations(Widget):
|
||||
Handle button press events from the widget.
|
||||
|
||||
This Textual event handler routes button presses to appropriate actions:
|
||||
- back_button: Pop this screen (return to parent)
|
||||
|
||||
- copy_results_btn: Copy results text to clipboard (requires pyperclip)
|
||||
- local_approval_btn: Start local approval operation
|
||||
- toggle_enforcement_btn: Start toggle audit/enforcement operation
|
||||
@@ -357,29 +370,24 @@ class AgentMoveOperations(Widget):
|
||||
|
||||
btn_id = event.button.id
|
||||
|
||||
if btn_id == "back_button":
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "copy_results_btn":
|
||||
if btn_id == "copy_results_btn":
|
||||
try:
|
||||
results_text = self.query_one("#results_text", TextArea)
|
||||
import pyperclip
|
||||
|
||||
pyperclip.copy(results_text.text)
|
||||
self.app.notify(
|
||||
"📋✅ Results copied to clipboard!",
|
||||
"📋✅ Results copied to clipboard!",
|
||||
severity="information",
|
||||
timeout=2,
|
||||
)
|
||||
except ImportError:
|
||||
self.app.notify(
|
||||
"⌠pyperclip not installed. Run: pip install pyperclip",
|
||||
"❌ pyperclip not installed. Run: pip install pyperclip",
|
||||
severity="warning",
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(f"âÂÅ’ Failed to copy: {str(e)}", severity="error")
|
||||
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
|
||||
event.stop()
|
||||
elif btn_id == "export_csv_btn":
|
||||
self._start_export_csv_operation()
|
||||
@@ -427,7 +435,7 @@ class AgentMoveOperations(Widget):
|
||||
self.operation_in_progress = True
|
||||
|
||||
status_label = self.query_one("#status_label", Static)
|
||||
status_label.update("âœ”ï¸ Moving agents to local approval...")
|
||||
status_label.update("✔️ Moving agents to local approval...")
|
||||
|
||||
# Get API from app
|
||||
api = self.app.api
|
||||
@@ -462,12 +470,12 @@ class AgentMoveOperations(Widget):
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during local approval operation: {e}")
|
||||
status_label.update(f"⌠Error: {str(e)}")
|
||||
status_label.update(f"❌ Error: {str(e)}")
|
||||
self.operation_in_progress = False
|
||||
return
|
||||
|
||||
self.operation_in_progress = False
|
||||
status_label.update("✅ Operation complete!")
|
||||
status_label.update("✅ Operation complete!")
|
||||
|
||||
# Display results in the widget
|
||||
self._display_results("Local Approval Mode", successful, unsuccessful)
|
||||
@@ -483,7 +491,6 @@ class AgentMoveOperations(Widget):
|
||||
self.selected_operation = "export_csv"
|
||||
self.operation_in_progress = True
|
||||
successful = []
|
||||
unsuccessful = []
|
||||
status_label = self.query_one("#status_label", Static)
|
||||
status_label.update("Exporting CSV...")
|
||||
self.app.refresh_data()
|
||||
@@ -511,9 +518,9 @@ class AgentMoveOperations(Widget):
|
||||
file_path = os.path.join(str(path), filename)
|
||||
df.to_csv(file_path, index=False)
|
||||
successful.append(file_path)
|
||||
status_label.update(f"✅ Exported to {file_path}")
|
||||
status_label.update(f"✅ Exported to {file_path}")
|
||||
except Exception:
|
||||
status_label.update("⌠Failed")
|
||||
status_label.update("❌ Failed")
|
||||
|
||||
self.operation_in_progress = False
|
||||
|
||||
@@ -557,7 +564,7 @@ class AgentMoveOperations(Widget):
|
||||
self.operation_in_progress = True
|
||||
|
||||
status_label = self.query_one("#status_label", Static)
|
||||
status_label.update("â³ Toggling enforcement mode...")
|
||||
status_label.update("🔄 Toggling enforcement mode...")
|
||||
|
||||
# Get API from app
|
||||
api = self.app.api
|
||||
@@ -593,12 +600,12 @@ class AgentMoveOperations(Widget):
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during toggle enforcement operation: {e}")
|
||||
status_label.update(f"⌠Error: {str(e)}")
|
||||
status_label.update(f"❌ Error: {str(e)}")
|
||||
self.operation_in_progress = False
|
||||
return
|
||||
|
||||
self.operation_in_progress = False
|
||||
status_label.update("✅ Operation complete!")
|
||||
status_label.update("✅ Operation complete!")
|
||||
|
||||
# Display results in the widget
|
||||
self._display_results("Toggle Audit/Enforcement", successful, unsuccessful)
|
||||
@@ -663,7 +670,7 @@ class AgentMoveOperations(Widget):
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading policies: {e}")
|
||||
status_label.update(f"⌠Error: {str(e)}")
|
||||
status_label.update(f"❌ Error: {str(e)}")
|
||||
self.operation_in_progress = False
|
||||
self.selected_operation = ""
|
||||
self.app.notify(f"Failed to load policies: {str(e)}", severity="error")
|
||||
@@ -695,7 +702,7 @@ class AgentMoveOperations(Widget):
|
||||
for agent in self.agents:
|
||||
try:
|
||||
# Move agent to target policy
|
||||
result = api.agent_move(agent.agentid, target_policy.groupid)
|
||||
api.agent_move(agent.agentid, target_policy.groupid)
|
||||
successful.append((agent, f"Moved to {target_policy.name}"))
|
||||
logger.info(
|
||||
f"Successfully moved {agent.hostname} to policy {target_policy.name}"
|
||||
@@ -1,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 re
|
||||
from typing import List, Optional
|
||||
@@ -20,6 +35,8 @@ from models.agent import Agent
|
||||
|
||||
|
||||
class MultiAgentSelector(Widget):
|
||||
"""Widget for selecting multiple agents from a list."""
|
||||
|
||||
class AgentsSelected(Message):
|
||||
def __init__(self, selected_agents: List[Agent]):
|
||||
super().__init__()
|
||||
@@ -40,7 +57,7 @@ class MultiAgentSelector(Widget):
|
||||
|
||||
def compose(self):
|
||||
yield Header(show_clock=True, icon="⚙")
|
||||
title_text = Static("🖧 Agent Selector", id="selector_title")
|
||||
title_text = Static("🖥️ Agent Selector", id="selector_title")
|
||||
title_text.styles.margin = (0, 0, 0, 1)
|
||||
yield title_text
|
||||
|
||||
@@ -60,7 +77,7 @@ class MultiAgentSelector(Widget):
|
||||
text_area.styles.overflow_y = "auto"
|
||||
yield text_area
|
||||
|
||||
with Horizontal(id="switch_search_container") as switch_search:
|
||||
with Horizontal(id="switch_search_container"):
|
||||
switch = Switch(value=False, id="match_switch")
|
||||
switch.styles.width = "auto"
|
||||
switch.styles.margin = (1, 0, 0, 0)
|
||||
@@ -91,11 +108,6 @@ class MultiAgentSelector(Widget):
|
||||
button_row.styles.height = "auto"
|
||||
button_row.styles.margin = (1, 0, 0, 0)
|
||||
|
||||
back_button = Button("← Back", id="back_button")
|
||||
back_button.styles.width = "1fr"
|
||||
back_button.styles.margin = (0, 0, 0, 1)
|
||||
yield back_button
|
||||
|
||||
submit_button = Button(
|
||||
"▶ Select & Continue", id="submit_selection", variant="primary"
|
||||
)
|
||||
@@ -123,10 +135,7 @@ class MultiAgentSelector(Widget):
|
||||
match_list = self.query_one("#match_results", SelectionList)
|
||||
except NoMatches:
|
||||
return
|
||||
if btn_id == "back_button":
|
||||
self.app.pop_screen()
|
||||
event.stop()
|
||||
elif btn_id == "select_all":
|
||||
if btn_id == "select_all":
|
||||
match_list.select_all()
|
||||
event.stop()
|
||||
elif btn_id == "select_none":
|
||||
@@ -1,10 +1,17 @@
|
||||
"""
|
||||
Policy Selector Widget Module
|
||||
|
||||
Provides a Textual widget for selecting target policies for bulk agent operations.
|
||||
Allows users to browse available policies and select one as the destination for
|
||||
moving agents. Automatically excludes parent/logical policies.
|
||||
"""
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>
|
||||
|
||||
import logging
|
||||
import re
|
||||
@@ -34,7 +41,7 @@ class PolicySelector(Widget):
|
||||
- Wildcard filtering (* and ?)
|
||||
- Interactive table for policy browsing
|
||||
- Explicit confirm button for selection
|
||||
- Cancel/back button to dismiss
|
||||
- Use escape key to go back
|
||||
|
||||
Attributes:
|
||||
policies (list[Policy]): List of available Policy objects to display.
|
||||
@@ -91,10 +98,10 @@ class PolicySelector(Widget):
|
||||
- Clear Filter button
|
||||
- Confirm Selection button
|
||||
- Policy table displaying available policies
|
||||
- Back buttons for navigation
|
||||
- Use escape key to go back
|
||||
"""
|
||||
title_text = Static(
|
||||
"🎯 Select Target Policy",
|
||||
"Select Target Policy",
|
||||
id="policy_selector_title",
|
||||
)
|
||||
title_text.styles.margin = (0, 0, 1, 0)
|
||||
@@ -125,12 +132,12 @@ class PolicySelector(Widget):
|
||||
filter_help.styles.margin = (0, 0, 1, 0)
|
||||
yield filter_help
|
||||
|
||||
apply_button = Button("✓ Apply Filter", id="filter_button")
|
||||
apply_button = Button("🔍 Apply Filter", id="filter_button")
|
||||
apply_button.styles.width = "100%"
|
||||
apply_button.styles.margin = (0, 0, 1, 0)
|
||||
yield apply_button
|
||||
|
||||
clear_button = Button("Clear Filter", id="clear_filter_button")
|
||||
clear_button = Button("🧹 Clear Filter", id="clear_filter_button")
|
||||
clear_button.styles.width = "100%"
|
||||
clear_button.styles.margin = (0, 0, 1, 0)
|
||||
yield clear_button
|
||||
@@ -144,11 +151,6 @@ class PolicySelector(Widget):
|
||||
selected_label.styles.margin = (2, 0, 1, 0)
|
||||
yield selected_label
|
||||
|
||||
cancel_button = Button("← Back", id="back_button")
|
||||
cancel_button.styles.width = "100%"
|
||||
cancel_button.styles.margin = (1, 0, 1, 0)
|
||||
yield cancel_button
|
||||
|
||||
# Right side - Policy table
|
||||
with Vertical() as right_side:
|
||||
right_side.styles.width = "2fr"
|
||||
@@ -222,7 +224,6 @@ class PolicySelector(Widget):
|
||||
Handle button press events from the widget.
|
||||
|
||||
Routes to:
|
||||
- back_button (Cancel): Pop screen without selecting
|
||||
- filter_button (Apply Filter): Filter policies with wildcard support
|
||||
- clear_filter_button: Clear filter and show all policies
|
||||
- confirm_button: Confirm selection and post message
|
||||
@@ -232,12 +233,7 @@ class PolicySelector(Widget):
|
||||
"""
|
||||
btn_id = event.button.id
|
||||
|
||||
if btn_id == "back_button":
|
||||
while len(self.app.screen_stack) > 2:
|
||||
self.app.pop_screen()
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "filter_button":
|
||||
if btn_id == "filter_button":
|
||||
self._apply_filter()
|
||||
event.stop()
|
||||
|
||||
@@ -283,7 +279,7 @@ class PolicySelector(Widget):
|
||||
if self.selected_policy:
|
||||
# Update selection display
|
||||
label = self.query_one("#selected_policy_label", Static)
|
||||
label.update(f"✓ Selected: {self.selected_policy.name}")
|
||||
label.update(f"Selected: {self.selected_policy.name}")
|
||||
|
||||
# Log for debugging
|
||||
logger.debug(
|
||||
@@ -325,7 +321,7 @@ class PolicySelector(Widget):
|
||||
|
||||
if highlighted_name:
|
||||
label = self.query_one("#selected_policy_label", Static)
|
||||
label.update(f"→ Highlighting: {highlighted_name}")
|
||||
label.update(f"Highlighting: {highlighted_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling row highlight: {e}")
|
||||
@@ -399,7 +395,9 @@ class PolicySelector(Widget):
|
||||
)
|
||||
|
||||
displayed_count = len(self._displayed_policies)
|
||||
status_text = f"📊 Showing {displayed_count} of {len(self._filtered_policies)} policies"
|
||||
status_text = (
|
||||
f"Showing {displayed_count} of {len(self._filtered_policies)} policies"
|
||||
)
|
||||
self.app.notify(status_text, severity="information", timeout=2)
|
||||
|
||||
# Clear selection when filter is applied
|
||||
@@ -474,7 +472,7 @@ class PolicySelector(Widget):
|
||||
"""
|
||||
if self.selected_policy is None:
|
||||
self.app.notify(
|
||||
"⚠️ Please select a policy first by clicking on a row in the table",
|
||||
"Please select a policy first by clicking on a row in the table",
|
||||
severity="warning",
|
||||
timeout=3,
|
||||
)
|
||||
@@ -483,6 +481,6 @@ class PolicySelector(Widget):
|
||||
# Log confirmation for debugging
|
||||
logger.info(f"Confirming selection of policy: {self.selected_policy.name}")
|
||||
self.app.notify(
|
||||
f"✅ Confirmed: {self.selected_policy.name}", severity="success", timeout=2
|
||||
f"Confirmed: {self.selected_policy.name}", severity="success", timeout=2
|
||||
)
|
||||
self.post_message(self.PolicySelected(self.selected_policy))
|
||||
@@ -1,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
|
||||
import logging
|
||||
|
||||
@@ -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–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
|
||||
logger.debug(f"{history_days} day selected for history")
|
||||
|
||||
if history_days is None:
|
||||
logging.warning("No history range selected. Aborting.")
|
||||
return
|
||||
|
||||
policy_executions = ExecutionHistoryRecord.from_policies(
|
||||
api, selected_policies, type_=type, history_days=history_days
|
||||
)
|
||||
|
||||
logger.debug(f"Executions contains {policy_executions}")
|
||||
|
||||
enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(
|
||||
api, policy_executions
|
||||
)
|
||||
categorized_executions = (
|
||||
ExecutionHistoryRecord.categorize_executions_by_hash_decision(
|
||||
enriched_executions
|
||||
)
|
||||
)
|
||||
approved, unapproved, needs_review, unknown = (
|
||||
ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions)
|
||||
)
|
||||
|
||||
categories = {
|
||||
"needs_review": needs_review,
|
||||
"approved": approved,
|
||||
"unapproved": unapproved,
|
||||
"leftover": unknown,
|
||||
}
|
||||
|
||||
for label, records in categories.items():
|
||||
if not records:
|
||||
continue # Skip empty or falsy categories
|
||||
|
||||
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
|
||||
html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html"
|
||||
|
||||
# Convert ExecutionHistoryRecord objects to dictionaries
|
||||
df = pd.DataFrame([r.__dict__ for r in records])
|
||||
|
||||
# Optional: flatten hash_obj if needed
|
||||
if not df.empty and "hash_obj" in df.columns:
|
||||
hash_df = df["hash_obj"].apply(lambda h: h.to_dict() if h else {})
|
||||
df = pd.concat([df.drop(columns=["hash_obj"]), hash_df], axis=1)
|
||||
|
||||
# Save to CSV
|
||||
df.to_csv(csv_path, index=False)
|
||||
logger.info(f"Saved {label} executions to {csv_path}")
|
||||
|
||||
# Generate HTML
|
||||
formatHTML(df, html_path)
|
||||
logger.info(f"Generated HTML report at {html_path}")
|
||||
|
||||
|
||||
def buildPathsandPublishers(selected_policies: List[Policy], split):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
df1 = pd.DataFrame()
|
||||
df2 = pd.DataFrame()
|
||||
all_approved_hashes = pd.DataFrame()
|
||||
path1 = (
|
||||
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
|
||||
)
|
||||
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
|
||||
path_exclusion_constant = get_system_value("PATH_EXCLUSION_CONST", cast_type=int)
|
||||
|
||||
if os.path.exists(path1):
|
||||
df1 = pd.read_csv(path1)
|
||||
else:
|
||||
logger.warning(f"File not found: {path1}")
|
||||
|
||||
if os.path.exists(path2):
|
||||
df2 = pd.read_csv(path2)
|
||||
else:
|
||||
logger.warning(f"File not found: {path2}")
|
||||
|
||||
if df1.empty and df2.empty:
|
||||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||
all_approved_hashes = pd.DataFrame()
|
||||
logger.debug(all_approved_hashes.head)
|
||||
else:
|
||||
all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
|
||||
if "filename" in all_approved_hashes.columns:
|
||||
all_approved_hashes = all_approved_hashes.sort_values(by="filename")
|
||||
else:
|
||||
logger.warning(
|
||||
"Warning: 'filename' column not found in concatenated DataFrame."
|
||||
)
|
||||
|
||||
if not all_approved_hashes.empty and path_exclusion_constant:
|
||||
|
||||
primary_path_exclusions = calculatePath(
|
||||
all_approved_hashes,
|
||||
path_exclusion_constant,
|
||||
split,
|
||||
)
|
||||
remaining_hashes = all_approved_hashes[
|
||||
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
|
||||
]
|
||||
secondary_path_exclusions = calculatePath(
|
||||
remaining_hashes, (path_exclusion_constant - 1), split
|
||||
)
|
||||
remaining_hashes = remaining_hashes[
|
||||
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
|
||||
]
|
||||
dataframes = {
|
||||
"all_approved_hashes": all_approved_hashes,
|
||||
"primary_Paths": primary_path_exclusions,
|
||||
"secondary_Paths": secondary_path_exclusions,
|
||||
"hashes_not_approvable_by_path": remaining_hashes,
|
||||
}
|
||||
logger.debug("Preparing to sort dataframes")
|
||||
for name, df in dataframes.items():
|
||||
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||||
if "hashes" in name:
|
||||
df.sort_values(by="filename", inplace=True)
|
||||
else:
|
||||
df.sort_values(by="longestcfp", inplace=True)
|
||||
|
||||
df.to_csv(
|
||||
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv",
|
||||
index=False,
|
||||
)
|
||||
formatHTML(
|
||||
df,
|
||||
f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html",
|
||||
)
|
||||
|
||||
if not all_approved_hashes.empty:
|
||||
# Drop all not signed, only keep unique values
|
||||
publist = all_approved_hashes[
|
||||
all_approved_hashes["publisher"] != "Not Signed"
|
||||
].drop_duplicates(subset=["publisher"])
|
||||
# Remove Bad publisher if somehow they made it this far
|
||||
pattern = regulator(get_system_list("BAD_PUBLISHERS"))
|
||||
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
|
||||
publist = publist[["publisher"]]
|
||||
publist.sort_values(by="publisher", inplace=True)
|
||||
publist.to_csv(
|
||||
f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv",
|
||||
index=False,
|
||||
)
|
||||
else:
|
||||
logger.debug("Approved Hashes list appears empty")
|
||||
|
||||
|
||||
def buildPreflights(selected_policies: List[Policy]):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
df1 = pd.DataFrame()
|
||||
df2 = pd.DataFrame()
|
||||
approved_hashes = pd.DataFrame()
|
||||
approved_publishers = pd.DataFrame()
|
||||
|
||||
hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv"
|
||||
path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
|
||||
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv"
|
||||
publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv"
|
||||
|
||||
# Read in and combine the two path generations
|
||||
if os.path.exists(path1):
|
||||
df1 = pd.read_csv(path1)
|
||||
else:
|
||||
logger.warning(f"File not found: {path1}")
|
||||
|
||||
if os.path.exists(path2):
|
||||
df2 = pd.read_csv(path2)
|
||||
else:
|
||||
logger.warning(f"File not found: {path2}")
|
||||
|
||||
if df1.empty and df2.empty:
|
||||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||||
approved_paths = pd.DataFrame()
|
||||
else:
|
||||
approved_paths = pd.concat([df1, df2], ignore_index=True)
|
||||
|
||||
approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first")
|
||||
|
||||
# We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions.
|
||||
if os.path.exists(hash):
|
||||
hashes = pd.read_csv(hash)
|
||||
approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])]
|
||||
|
||||
approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first")
|
||||
|
||||
else:
|
||||
logger.warning(f"File not found: {hash}")
|
||||
|
||||
if os.path.exists(publishers):
|
||||
approved_publishers = pd.read_csv(publishers)
|
||||
|
||||
else:
|
||||
logger.warning(f"File not found: {publishers}")
|
||||
|
||||
dataframes = {
|
||||
"approved_paths": approved_paths,
|
||||
"approved_hashes": approved_hashes,
|
||||
"approved_publishers": approved_publishers,
|
||||
}
|
||||
|
||||
for name, df in dataframes.items():
|
||||
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||||
if name == "approved_paths":
|
||||
df.sort_values(by="longestcfp", inplace=True)
|
||||
elif name == "approved_hashes":
|
||||
df.sort_values(by="filename", inplace=True)
|
||||
elif name == "approved_publishers":
|
||||
df.sort_values(by="publisher", inplace=True)
|
||||
|
||||
df.to_csv(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv",
|
||||
index=False,
|
||||
)
|
||||
formatHTML(
|
||||
df,
|
||||
f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html",
|
||||
)
|
||||
|
||||
|
||||
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
|
||||
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
|
||||
|
||||
def clean_split(path):
|
||||
if not isinstance(path, (str, bytes, os.PathLike)):
|
||||
return []
|
||||
parts = str(os.path.normpath(path)).split(os.sep)
|
||||
parts = [p for p in parts if p] # Remove empty strings
|
||||
return parts
|
||||
|
||||
# Diagnostic: log any non-string entries
|
||||
non_string_entries = df[
|
||||
~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))
|
||||
]
|
||||
if not non_string_entries.empty:
|
||||
print(f"[WARNING] Non-string entries found in column '{col}':")
|
||||
print(non_string_entries)
|
||||
|
||||
df = df.copy()
|
||||
split_paths = df[col].apply(clean_split)
|
||||
|
||||
if min_files_for_path is not None:
|
||||
df = df[
|
||||
split_paths.apply(lambda parts: len(parts) >= min_files_for_path)
|
||||
].copy()
|
||||
split_paths = split_paths[df.index]
|
||||
|
||||
df["group_key"] = split_paths.apply(
|
||||
lambda parts: os.sep.join(parts[:path_exclusion_constant])
|
||||
)
|
||||
grouped = df.groupby("group_key")
|
||||
new_rows = []
|
||||
|
||||
for _, group_df in grouped:
|
||||
paths = group_df[col].tolist()
|
||||
split_parts = [clean_split(p) for p in paths]
|
||||
|
||||
def longest_common_prefix(paths):
|
||||
if not paths:
|
||||
return []
|
||||
prefix = paths[0]
|
||||
for path in paths[1:]:
|
||||
prefix = [a for a, b in zip(prefix, path) if a == b]
|
||||
if not prefix:
|
||||
break
|
||||
return prefix
|
||||
|
||||
common_prefix = longest_common_prefix(split_parts)
|
||||
prefix_str = os.sep.join(common_prefix)
|
||||
|
||||
for i, parts in enumerate(split_parts):
|
||||
filename = parts[-1]
|
||||
middle = (
|
||||
os.sep.join(parts[len(common_prefix) : -1])
|
||||
if len(parts) > len(common_prefix) + 1
|
||||
else ""
|
||||
)
|
||||
row = group_df.iloc[i].copy()
|
||||
row["longestcfp"] = prefix_str
|
||||
row["middle"] = middle
|
||||
row["filename_only"] = filename
|
||||
row["file_extension"] = os.path.splitext(filename)[1].lower()
|
||||
new_rows.append(row)
|
||||
|
||||
return pd.DataFrame(new_rows).drop(columns=["group_key"])
|
||||
|
||||
|
||||
def calculatePath(approved_hashes, path_exclusion_constant, split):
|
||||
if split:
|
||||
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
|
||||
else:
|
||||
dfs_by_policy = [approved_hashes]
|
||||
|
||||
badpathparts = get_system_list("BAD_PATH_PARTS")
|
||||
min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int)
|
||||
|
||||
processed_dfs = []
|
||||
|
||||
for df in dfs_by_policy:
|
||||
haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename")
|
||||
haslcp = haslcp.drop_duplicates()
|
||||
|
||||
forbidden = regulator(badpathparts, True)
|
||||
forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
|
||||
|
||||
logger.debug("Removing forbidden filepaths for path exceptions")
|
||||
print(colorText("Removing forbidden filepaths for path exceptions", "green"))
|
||||
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
|
||||
|
||||
lcp_not_forbidden_review = lcp_not_forbidden[
|
||||
[
|
||||
"policyname",
|
||||
"longestcfp",
|
||||
"middle",
|
||||
"filename_only",
|
||||
"file_extension",
|
||||
"sha256",
|
||||
]
|
||||
]
|
||||
|
||||
unique_sha_counts = (
|
||||
lcp_not_forbidden_review.groupby("longestcfp")["sha256"]
|
||||
.nunique()
|
||||
.reset_index()
|
||||
)
|
||||
unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]
|
||||
|
||||
lcp_not_forbidden_review = lcp_not_forbidden_review.merge(
|
||||
unique_sha_counts, on="longestcfp", how="left"
|
||||
)
|
||||
lcp_not_forbidden_review = lcp_not_forbidden_review[
|
||||
lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path
|
||||
]
|
||||
processed_dfs.append(lcp_not_forbidden_review)
|
||||
|
||||
pathExclusions = pd.concat(processed_dfs, ignore_index=True)
|
||||
|
||||
return pathExclusions
|
||||
|
||||
|
||||
def testChange(selected_policies, destination_policy, destination_allowlist):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
logger.info("These path exclusions would be added to:")
|
||||
logger.info(destination_policy)
|
||||
|
||||
pathexclusions = pd.read_csv(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
|
||||
)
|
||||
hashes = pd.read_csv(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
|
||||
)
|
||||
|
||||
unique_combinations = pathexclusions[
|
||||
["longestcfp", "file_extension"]
|
||||
].drop_duplicates()
|
||||
|
||||
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
|
||||
processed_paths = [
|
||||
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
|
||||
for path, ext in unique_combinations.itertuples(index=False, name=None)
|
||||
]
|
||||
|
||||
for path in processed_paths:
|
||||
logger.info(path)
|
||||
|
||||
print(colorText("These publishers would added", "yellow"))
|
||||
processed_publishers = []
|
||||
if os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
|
||||
):
|
||||
publishers = pd.read_csv(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"
|
||||
)
|
||||
if publishers.empty:
|
||||
print(colorText("The publishers list is empty.", "red"))
|
||||
else:
|
||||
processed_publishers = (
|
||||
publishers[publishers["publisher"] != "Not Signed"]["publisher"]
|
||||
.drop_duplicates()
|
||||
.tolist()
|
||||
)
|
||||
for publisher in processed_publishers:
|
||||
print(publisher)
|
||||
|
||||
print(colorText("These hashes would be added to:", "yellow"))
|
||||
print(destination_allowlist)
|
||||
|
||||
processed_hashes = hashes["sha256"].unique().tolist()
|
||||
print_x_wide(processed_hashes, 3)
|
||||
|
||||
return processed_paths, processed_hashes, processed_publishers
|
||||
|
||||
|
||||
def menu_policy_enforce(
|
||||
api: AirlockAPIWrapper,
|
||||
): # TODO Need to clean up 6 and 7 into functions
|
||||
selected_policies = []
|
||||
destination_policy = []
|
||||
destination_allowlist = []
|
||||
processed_paths = []
|
||||
processed_hashes = []
|
||||
processed_publishers = []
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
|
||||
while True:
|
||||
printEnforceChecklist(
|
||||
selected_policies, destination_policy, destination_allowlist
|
||||
)
|
||||
choice = get_sanitized_input("\nEnter your choice: ")
|
||||
|
||||
if choice == "1":
|
||||
clear_screen()
|
||||
selected_policies = selectPolicies(api, True)
|
||||
|
||||
elif choice == "2":
|
||||
clear_screen()
|
||||
print(
|
||||
colorText(
|
||||
"Please choose destination_name Policy for Path Exclusions", "white"
|
||||
)
|
||||
)
|
||||
|
||||
destination_policy = selectPolicies(api, False)
|
||||
|
||||
print(colorText("Please choose Allowlist for Hashes", "white"))
|
||||
|
||||
destination_allowlist = selectAllowlists(api, destination_policy, False)
|
||||
|
||||
elif choice == "3":
|
||||
clear_screen()
|
||||
sortHashes(
|
||||
api,
|
||||
selected_policies,
|
||||
type=[1, 2, 6, 7],
|
||||
)
|
||||
|
||||
elif choice == "4":
|
||||
clear_screen()
|
||||
if os.path.exists(
|
||||
f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"
|
||||
):
|
||||
buildPathsandPublishers(selected_policies, False)
|
||||
else:
|
||||
print(
|
||||
"File not found. Please make sure it's saved correctly and try again."
|
||||
)
|
||||
|
||||
elif choice == "5":
|
||||
clear_screen()
|
||||
if os.path.exists(
|
||||
f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
|
||||
) and os.path.exists(
|
||||
f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
|
||||
):
|
||||
buildPreflights(selected_policies)
|
||||
else:
|
||||
print(
|
||||
"File not found. Please make sure it's saved correctly and try again."
|
||||
)
|
||||
|
||||
elif choice == "6":
|
||||
clear_screen()
|
||||
if (
|
||||
os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
|
||||
)
|
||||
and os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
|
||||
)
|
||||
and destination_policy
|
||||
and destination_allowlist
|
||||
):
|
||||
processed_paths, processed_hashes, processed_publishers = testChange(
|
||||
selected_policies, destination_policy, destination_allowlist
|
||||
)
|
||||
else:
|
||||
# Log which condition(s) failed
|
||||
missing_items = []
|
||||
if not os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"
|
||||
):
|
||||
missing_items.append("approved_paths.csv not found")
|
||||
if not os.path.exists(
|
||||
f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"
|
||||
):
|
||||
missing_items.append("approved_hashes.csv not found")
|
||||
if not destination_policy:
|
||||
missing_items.append("destination_policy is empty or None")
|
||||
if not destination_allowlist:
|
||||
missing_items.append("destination_allowlist is empty or None")
|
||||
|
||||
logger.error("Preflight check failed due to the following:")
|
||||
for item in missing_items:
|
||||
logger.error(f" - {item}")
|
||||
|
||||
elif choice == "7":
|
||||
clear_screen()
|
||||
areYouSure()
|
||||
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
if (
|
||||
processed_paths
|
||||
and processed_hashes
|
||||
and processed_publishers
|
||||
and destination_policy
|
||||
and destination_allowlist
|
||||
and confirmation.strip() == "I AGREE"
|
||||
):
|
||||
print(colorText("Proceeding with the code...", "yellow"))
|
||||
api.hash_add_to_allowlist(
|
||||
destination_allowlist[0].applicationid, processed_hashes
|
||||
)
|
||||
api.policy_add_path_exclusions(
|
||||
destination_policy[0].groupid, processed_paths
|
||||
)
|
||||
if processed_publishers:
|
||||
api.policy_add_publishers(
|
||||
destination_policy[0].groupid, processed_publishers
|
||||
)
|
||||
|
||||
locked()
|
||||
|
||||
else:
|
||||
logger.error("Confirmation block failed. Reasons:")
|
||||
if not processed_publishers or processed_hashes or processed_paths:
|
||||
logger.error(" - Test not performed.")
|
||||
if not destination_policy:
|
||||
logger.error(" - `destination_policy` is missing or invalid.")
|
||||
if not destination_allowlist:
|
||||
logger.error(" - `destination_allowlist` is missing or invalid.")
|
||||
if confirmation.strip() != "I AGREE":
|
||||
logger.error(
|
||||
" - User did not confirm with 'I AGREE'. Received: '%s'",
|
||||
confirmation.strip(),
|
||||
)
|
||||
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
elif choice.upper() == "B":
|
||||
break
|
||||
|
||||
else:
|
||||
print(colorText("Invalid choice. Please try again.", "red"))
|
||||
|
||||
|
||||
def section_header(title):
|
||||
print(
|
||||
colorText(
|
||||
"\n --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(f" ------------- {title} -------------", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" --------------------------------------------------------------------",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
section_header("Prepare to Enforce Policy ")
|
||||
print(
|
||||
colorText(
|
||||
"\nSequentially follow these steps to prepare a policy for enforcement:",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
|
||||
# Step 1: Originating Policies
|
||||
print(
|
||||
colorText(
|
||||
"\n1. Choose which policy or policies to gather execution info from", "cyan"
|
||||
)
|
||||
)
|
||||
if not selected_policies:
|
||||
print(colorText(" [✗] No policies have been chosen", "red"))
|
||||
else:
|
||||
print(colorText("The following policies have been chosen:", "green"))
|
||||
for policy in selected_policies:
|
||||
print(colorText(f" [✓] {policy.name}", "green"))
|
||||
|
||||
# Step 2: Destination Policy and Allowlist
|
||||
print(
|
||||
colorText("2. Choose the destination policy and associated allowlist", "cyan")
|
||||
)
|
||||
if destination_policy:
|
||||
print(
|
||||
colorText(
|
||||
f" [✓] {destination_policy[0].name} has been selected as the destination policy",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(colorText(" [✗] No destination policy has been chosen", "red"))
|
||||
|
||||
if destination_allowlist:
|
||||
print(
|
||||
colorText(
|
||||
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(colorText(" [✗] No allowlist has been chosen", "red"))
|
||||
|
||||
# Step 3: Data Preparation
|
||||
print(
|
||||
colorText(
|
||||
f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
if selected_policies:
|
||||
policy_id = selected_policies[0].name
|
||||
review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv"
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Data has been fetched"
|
||||
if os.path.exists(review_path)
|
||||
else " [✗] Data has not been fetched"
|
||||
),
|
||||
"green" if os.path.exists(review_path) else "red",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check data fetch status", "red"
|
||||
)
|
||||
)
|
||||
|
||||
# Step 4: Manual Review
|
||||
print(colorText("4. Manually review the files:", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing hashes you do not approve of", "cyan"
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" This will start the process to generate possible filepath approvals",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
if selected_policies:
|
||||
policy_id = selected_policies[0].name
|
||||
approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv"
|
||||
second_review_path = (
|
||||
f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv"
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Reviewed hashes have been loaded"
|
||||
if os.path.exists(approved_path)
|
||||
else " [✗] Reviewed hashes have not been loaded"
|
||||
),
|
||||
"green" if os.path.exists(approved_path) else "red",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Path review list created"
|
||||
if os.path.exists(second_review_path)
|
||||
else " [✗] Path review list has not been created"
|
||||
),
|
||||
"green" if os.path.exists(second_review_path) else "red",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check reviewed hashes or path list",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
|
||||
# Step 5: Path Review
|
||||
print(
|
||||
colorText(
|
||||
f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Remove the rows containing path exclusions or publishers you do not approve of.",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
f" When complete, save the files to {working_dir}\\data\\Approved",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(" Choose this option when done to build your preflights", "cyan")
|
||||
)
|
||||
|
||||
if selected_policies:
|
||||
policy_id = selected_policies[0].name
|
||||
reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv"
|
||||
preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv"
|
||||
preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv"
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Reviewed path list detected"
|
||||
if os.path.exists(reviewed_path)
|
||||
else " [✗] Path review list has not been detected"
|
||||
),
|
||||
"green" if os.path.exists(reviewed_path) else "red",
|
||||
)
|
||||
)
|
||||
preflight_ready = os.path.exists(preflight_paths) and os.path.exists(
|
||||
preflight_hashes
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Preflight Path Exclusion List has been generated"
|
||||
if preflight_ready
|
||||
else " [✗] Preflight Path Exclusion List has not been generated"
|
||||
),
|
||||
"green" if preflight_ready else "red",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check preflight status", "red"
|
||||
)
|
||||
)
|
||||
|
||||
# Final Steps
|
||||
print(
|
||||
colorText(
|
||||
"6. Test ------------------------------------------------------", "cyan"
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Prints to console the changes that would be made, must be done to proceed. ",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
|
||||
print(
|
||||
colorText(
|
||||
"7. Liftoff ------------------------------------------------------", "cyan"
|
||||
)
|
||||
)
|
||||
print(
|
||||
colorText(
|
||||
" Apply path exclusions and approved publishers to selected policy",
|
||||
"cyan",
|
||||
)
|
||||
)
|
||||
print(colorText(" Apply approved hashes to allowlist", "cyan"))
|
||||
|
||||
# Utility Options
|
||||
print(colorText("F. Open Working Directory", "cyan"))
|
||||
print(colorText("B. Back", "cyan"))
|
||||
@@ -1,3 +1,18 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import logging
|
||||
|
||||
from textual.containers import Horizontal, Vertical
|
||||
@@ -59,15 +74,6 @@ class ResultsDisplay(Widget):
|
||||
margin-top: 1;
|
||||
width: 100%;
|
||||
}
|
||||
|
||||
#button_row {
|
||||
height: auto;
|
||||
margin: 1 0 0 0;
|
||||
}
|
||||
|
||||
#back_button {
|
||||
width: 1fr;
|
||||
}
|
||||
"""
|
||||
|
||||
class CopySuccess(Message):
|
||||
@@ -95,9 +101,9 @@ class ResultsDisplay(Widget):
|
||||
|
||||
def compose(self):
|
||||
with Vertical(id="results_screen"):
|
||||
yield Header(show_clock=True, icon="⚙")
|
||||
yield Header(show_clock=True, icon="⚙️")
|
||||
# Title
|
||||
title = Static(f"📊 {self.operation} - Results", id="results_title")
|
||||
title = Static(f"{self.operation} - Results", id="results_title")
|
||||
yield title
|
||||
|
||||
# Two-column layout
|
||||
@@ -107,7 +113,7 @@ class ResultsDisplay(Widget):
|
||||
yield Static("✅ Successful", id="success_label")
|
||||
yield Static(self.successful_results, id="success_results")
|
||||
yield Button(
|
||||
"📋✅ Copy Success List",
|
||||
"Copy Success List",
|
||||
id="copy_success",
|
||||
classes="copy_button",
|
||||
)
|
||||
@@ -117,15 +123,11 @@ class ResultsDisplay(Widget):
|
||||
yield Static("❌ Failed", id="failure_label")
|
||||
yield Static(self.unsuccessful_results, id="failure_results")
|
||||
yield Button(
|
||||
"📋❌ Copy Failure List",
|
||||
"Copy Failure List",
|
||||
id="copy_failure",
|
||||
classes="copy_button",
|
||||
)
|
||||
|
||||
# Back Button
|
||||
with Horizontal(id="button_row"):
|
||||
back_button = Button("← Back", id="back_button")
|
||||
yield back_button
|
||||
yield Footer()
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
@@ -138,18 +140,18 @@ class ResultsDisplay(Widget):
|
||||
|
||||
pyperclip.copy(str(success_widget.renderable))
|
||||
self.app.notify(
|
||||
"✅ Success list copied to clipboard!",
|
||||
"Success list copied to clipboard!",
|
||||
severity="information",
|
||||
timeout=2,
|
||||
)
|
||||
self.post_message(self.CopySuccess())
|
||||
except ImportError:
|
||||
self.app.notify(
|
||||
"âš ï¸ pyperclip not installed. Run: pip install pyperclip",
|
||||
"❌ pyperclip not installed. Run: pip install pyperclip",
|
||||
severity="warning",
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(f"⌠Failed to copy: {str(e)}", severity="error")
|
||||
self.app.notify(f"¢ Failed to copy: {str(e)}", severity="error")
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "copy_failure":
|
||||
@@ -159,20 +161,16 @@ class ResultsDisplay(Widget):
|
||||
|
||||
pyperclip.copy(str(failure_widget.renderable))
|
||||
self.app.notify(
|
||||
"✅ Failure list copied to clipboard!",
|
||||
"Failure list copied to clipboard!",
|
||||
severity="information",
|
||||
timeout=2,
|
||||
)
|
||||
self.post_message(self.CopyFailure())
|
||||
except ImportError:
|
||||
self.app.notify(
|
||||
"âš ï¸ pyperclip not installed. Run: pip install pyperclip",
|
||||
"❌ pyperclip not installed. Run: pip install pyperclip",
|
||||
severity="warning",
|
||||
)
|
||||
except Exception as e:
|
||||
self.app.notify(f"⌠Failed to copy: {str(e)}", severity="error")
|
||||
event.stop()
|
||||
|
||||
elif btn_id == "back_button":
|
||||
self.app.pop_screen()
|
||||
self.app.notify(f"❌ Failed to copy: {str(e)}", severity="error")
|
||||
event.stop()
|
||||
@@ -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)
|
||||
)
|
||||
@@ -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."""
|
||||
Generated
+34
-1
@@ -26,9 +26,10 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "airlock_libs"
|
||||
version = "5.0.0"
|
||||
version = "5.1.0"
|
||||
dependencies = [
|
||||
"chrono",
|
||||
"crossbeam",
|
||||
"indicatif",
|
||||
"mongodb",
|
||||
"opentelemetry 0.18.0",
|
||||
@@ -504,6 +505,19 @@ version = "1.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b"
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam"
|
||||
version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1137cd7e7fc0fb5d3c5a8678be38ec56e819125d8d7907411fe24ccb943faca8"
|
||||
dependencies = [
|
||||
"crossbeam-channel",
|
||||
"crossbeam-deque",
|
||||
"crossbeam-epoch",
|
||||
"crossbeam-queue",
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-channel"
|
||||
version = "0.5.15"
|
||||
@@ -513,6 +527,16 @@ dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-deque"
|
||||
version = "0.8.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51"
|
||||
dependencies = [
|
||||
"crossbeam-epoch",
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-epoch"
|
||||
version = "0.9.18"
|
||||
@@ -522,6 +546,15 @@ dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-queue"
|
||||
version = "0.3.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115"
|
||||
dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-utils"
|
||||
version = "0.8.21"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "airlock_libs"
|
||||
version = "5.0.0"
|
||||
version = "5.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[lib]
|
||||
@@ -25,6 +25,7 @@ tracing = "0.1.41"
|
||||
tracing-subscriber = "0.3.20"
|
||||
tracing-opentelemetry = "0.32.0"
|
||||
pyo3-async-runtimes = { version = "0.27.0", features = ["async-std", "tokio"] }
|
||||
crossbeam = "0.8.4"
|
||||
|
||||
[package.metadata.maturin]
|
||||
generate-abi-stubs = true
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "airlock_libs"
|
||||
version = "5.0.0"
|
||||
version = "5.1.0"
|
||||
description = "Airlock Digital API Wrapper"
|
||||
readme = "README.md"
|
||||
license = { text = "AGPL-3.0-only" }
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use pyo3::prelude::*;
|
||||
mod services;
|
||||
pub mod modules;
|
||||
pub mod prelude;
|
||||
pub mod services;
|
||||
#[pymodule]
|
||||
fn airlock_libs(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
|
||||
m.add_function(wrap_pyfunction!(services::pull_policy_exec_histories, py)?)?;
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
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 enum ExtractedValues {
|
||||
Headers(reqwest::header::HeaderMap),
|
||||
BaseUrl(String),
|
||||
}
|
||||
|
||||
pub trait Converter {
|
||||
fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues;
|
||||
}
|
||||
|
||||
pub struct PyData;
|
||||
|
||||
impl Converter for PyData {
|
||||
fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues {
|
||||
if extract_headers {
|
||||
let headers = py_self.getattr(py, "headers").unwrap().to_string();
|
||||
let headers_replace = headers.replace('\'', "\"");
|
||||
let parsed: Value = serde_json::from_str(headers_replace.as_str()).unwrap();
|
||||
let mut header_map = HeaderMap::new();
|
||||
if let Some(obj) = parsed.as_object() {
|
||||
for (_key, value) in obj {
|
||||
if let Some(v) = value.as_str() {
|
||||
let val = HeaderValue::from_str(v).unwrap();
|
||||
header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val);
|
||||
}
|
||||
}
|
||||
}
|
||||
ExtractedValues::Headers(header_map)
|
||||
} else {
|
||||
let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
|
||||
ExtractedValues::BaseUrl(base_url)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SkipBack;
|
||||
|
||||
impl SkipBack {
|
||||
pub fn find_checkpoint(days: i64) -> ObjectId {
|
||||
let date_days_ago = Local::now() - Duration::days(days);
|
||||
let timestamp = date_days_ago.timestamp() as u32;
|
||||
let mut hex_timestamp = String::new();
|
||||
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap();
|
||||
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
|
||||
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod datatypes;
|
||||
@@ -0,0 +1,27 @@
|
||||
pub use chrono::{Duration, Local, NaiveDate};
|
||||
pub use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
|
||||
pub use mongodb::bson::oid::ObjectId;
|
||||
pub use opentelemetry::global::shutdown_tracer_provider;
|
||||
pub use opentelemetry::sdk::Resource;
|
||||
pub use opentelemetry::trace::noop::NoopTracerProvider;
|
||||
pub use opentelemetry::trace::{Status, TraceContextExt, TraceError};
|
||||
pub use opentelemetry::{Context, KeyValue, sdk::trace as sdktrace, trace::Tracer};
|
||||
pub use opentelemetry::{Key, global};
|
||||
pub use opentelemetry_otlp::WithExportConfig;
|
||||
pub use pyo3::{prelude::*, types::PyString};
|
||||
pub use pyo3_async_runtimes::async_std;
|
||||
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,
|
||||
};
|
||||
+50
-189
@@ -1,128 +1,9 @@
|
||||
use chrono::{Duration, Local, NaiveDate};
|
||||
use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
|
||||
use mongodb::bson::oid::ObjectId;
|
||||
use opentelemetry::global::shutdown_tracer_provider;
|
||||
use opentelemetry::sdk::Resource;
|
||||
use opentelemetry::trace::noop::NoopTracerProvider;
|
||||
use opentelemetry::trace::{Status, TraceContextExt, TraceError};
|
||||
use opentelemetry::{Context, KeyValue, sdk::trace as sdktrace, trace::Tracer};
|
||||
use opentelemetry::{Key, global};
|
||||
use opentelemetry_otlp::WithExportConfig;
|
||||
use pyo3::{prelude::*, types::PyString};
|
||||
use pyo3_async_runtimes::async_std;
|
||||
use reqwest::{
|
||||
Client,
|
||||
header::{HeaderMap, HeaderName, HeaderValue},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
env,
|
||||
fmt::Write,
|
||||
fs::{self, File},
|
||||
io::{Read, Seek, SeekFrom},
|
||||
path::PathBuf,
|
||||
str::FromStr,
|
||||
};
|
||||
use std::thread;
|
||||
|
||||
#[allow(non_snake_case)]
|
||||
#[derive(Deserialize, Debug)]
|
||||
struct TelemetryConfig {
|
||||
TELEMETRY: bool,
|
||||
TELEM_URL: Option<String>,
|
||||
}
|
||||
|
||||
impl TelemetryConfig {
|
||||
pub fn load() -> Self {
|
||||
let cfg_path = get_base_directory().join("config\\user_config.json");
|
||||
if !cfg_path.exists() {
|
||||
return Self {
|
||||
TELEMETRY: false,
|
||||
TELEM_URL: None,
|
||||
};
|
||||
}
|
||||
match fs::read_to_string(&cfg_path) {
|
||||
Ok(contents) => serde_json::from_str::<Self>(&contents).unwrap_or(Self {
|
||||
TELEMETRY: false,
|
||||
TELEM_URL: None,
|
||||
}),
|
||||
Err(_) => Self {
|
||||
TELEMETRY: false,
|
||||
TELEM_URL: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct ApiResponse {
|
||||
error: String,
|
||||
response: ExecHistories,
|
||||
}
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct ExecHistories {
|
||||
exechistories: Vec<Group>,
|
||||
}
|
||||
#[derive(Debug, Deserialize, Serialize, Clone)]
|
||||
struct Group {
|
||||
checkpoint: String,
|
||||
#[serde(rename = "type")]
|
||||
exectype: u8,
|
||||
username: String,
|
||||
hostname: String,
|
||||
netdomain: String,
|
||||
filename: String,
|
||||
ppolicy: String,
|
||||
policyname: String,
|
||||
policyver: String,
|
||||
commandline: String,
|
||||
publisher: String,
|
||||
pprocess: String,
|
||||
gprocess: String,
|
||||
sha256: String,
|
||||
datetime: String,
|
||||
md5: String,
|
||||
sha128: String,
|
||||
sha384: String,
|
||||
sha512: String,
|
||||
ip: String,
|
||||
localip: String,
|
||||
}
|
||||
|
||||
enum ExtractedValues {
|
||||
Headers(reqwest::header::HeaderMap),
|
||||
BaseUrl(String),
|
||||
}
|
||||
|
||||
trait Converter {
|
||||
fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues;
|
||||
}
|
||||
|
||||
struct PyData;
|
||||
|
||||
impl Converter for PyData {
|
||||
fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues {
|
||||
if extract_headers {
|
||||
let headers = py_self.getattr(py, "headers").unwrap().to_string();
|
||||
let headers_replace = headers.replace('\'', "\"");
|
||||
let parsed: Value = serde_json::from_str(headers_replace.as_str()).unwrap();
|
||||
let mut header_map = HeaderMap::new();
|
||||
if let Some(obj) = parsed.as_object() {
|
||||
for (_key, value) in obj {
|
||||
if let Some(v) = value.as_str() {
|
||||
let val = HeaderValue::from_str(v).unwrap();
|
||||
header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val);
|
||||
}
|
||||
}
|
||||
}
|
||||
ExtractedValues::Headers(header_map)
|
||||
} else {
|
||||
let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
|
||||
ExtractedValues::BaseUrl(base_url)
|
||||
}
|
||||
}
|
||||
}
|
||||
use crossbeam::channel::unbounded;
|
||||
|
||||
use crate::modules::datatypes::*;
|
||||
use crate::prelude::*;
|
||||
#[pyfunction]
|
||||
pub fn pull_policy_exec_histories(
|
||||
py: Python<'_>,
|
||||
@@ -157,7 +38,6 @@ pub fn pull_policy_exec_histories(
|
||||
get_base_directory().display()
|
||||
)
|
||||
.into();
|
||||
let writeable_filepath = file_path.clone();
|
||||
if !&file_path.exists() {
|
||||
if let Some(parent_dir) = &file_path.parent()
|
||||
&& !parent_dir.exists()
|
||||
@@ -184,6 +64,7 @@ pub fn pull_policy_exec_histories(
|
||||
exechistories: vec![],
|
||||
},
|
||||
};
|
||||
let writeable_filepath = file_path.clone();
|
||||
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
|
||||
match fs::write(writeable_filepath.clone(), data_write) {
|
||||
Ok(_) => {}
|
||||
@@ -192,13 +73,13 @@ pub fn pull_policy_exec_histories(
|
||||
std::process::abort();
|
||||
}
|
||||
}
|
||||
let mut checkpoint_number: String = skipback(days).to_string();
|
||||
let mut checkpoint_number: String = SkipBack::find_checkpoint(days).to_string();
|
||||
let multi_progress = MultiProgress::new();
|
||||
multi_progress.set_draw_target(ProgressDrawTarget::stdout());
|
||||
let progress_bar = multi_progress.add(ProgressBar::new(100));
|
||||
progress_bar.set_style(
|
||||
ProgressStyle::default_bar()
|
||||
.template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
|
||||
.template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len} {message}")
|
||||
.unwrap(),
|
||||
);
|
||||
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
|
||||
@@ -231,55 +112,19 @@ pub fn pull_policy_exec_histories(
|
||||
}
|
||||
});
|
||||
let cutoff = Local::now().naive_local() - Duration::days(days);
|
||||
let mut f = match File::open(&writeable_filepath) {
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
println!("Failed to Access {:?}: {}", &writeable_filepath, e);
|
||||
std::process::abort();
|
||||
}
|
||||
};
|
||||
tracer.in_span("Airlock Data Retreival", |cx| {
|
||||
let span = cx.span();
|
||||
span.set_attribute(Key::new("Days").string(days.to_string()));
|
||||
span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
|
||||
loop {
|
||||
match f.seek(SeekFrom::Start(0)) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("Failed to seek start of {:?}: {}", f, e);
|
||||
std::process::abort();
|
||||
}
|
||||
}
|
||||
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 = execution_histories.response.exechistories;
|
||||
if parsed_responses.is_empty() {
|
||||
break;
|
||||
}
|
||||
let mut seen: HashMap<(String, String, String), Group> =
|
||||
if writeable_filepath.exists() {
|
||||
let mut contents = String::new();
|
||||
f.read_to_string(&mut contents).unwrap();
|
||||
let existing_data: ApiResponse =
|
||||
let (tx, rx) = unbounded::<Vec<Group>>();
|
||||
thread::spawn(move || {
|
||||
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
|
||||
{
|
||||
let contents = fs::read_to_string(&writeable_filepath).unwrap_or_default();
|
||||
let existing: ApiResponse =
|
||||
serde_json::from_str(&contents).unwrap_or(ApiResponse {
|
||||
error: "Success".to_string(),
|
||||
response: ExecHistories {
|
||||
exechistories: vec![],
|
||||
},
|
||||
});
|
||||
existing_data
|
||||
existing
|
||||
.response
|
||||
.exechistories
|
||||
.into_iter()
|
||||
@@ -297,14 +142,11 @@ pub fn pull_policy_exec_histories(
|
||||
} else {
|
||||
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() {
|
||||
continue;
|
||||
}
|
||||
if index == parsed_responses.len() - 1 {
|
||||
checkpoint_number = executions.checkpoint.clone();
|
||||
break;
|
||||
}
|
||||
let history_date = match NaiveDate::parse_from_str(
|
||||
&executions.datetime.replace(" +0000 UTC", ""),
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
@@ -334,7 +176,34 @@ pub fn pull_policy_exec_histories(
|
||||
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
|
||||
}
|
||||
}
|
||||
if let Some(last_item) = &final_response.response.exechistories.last()
|
||||
}
|
||||
});
|
||||
tracer.in_span("Airlock Data Retreival", |cx| {
|
||||
let span = cx.span();
|
||||
span.set_attribute(Key::new("Days").string(days.to_string()));
|
||||
span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
|
||||
loop {
|
||||
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 = 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(
|
||||
&last_item.datetime.replace(" +0000 UTC", ""),
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
@@ -342,20 +211,21 @@ pub fn pull_policy_exec_histories(
|
||||
{
|
||||
let date_diff = Local::now().naive_local().date() - last_date;
|
||||
let percentage_diff =
|
||||
(days - date_diff.num_days()) as f64 / days as f64 * 100.0;
|
||||
progress_bar.set_position(percentage_diff.round() as u64);
|
||||
((days - date_diff.num_days()) as f64 / days as f64 * 100.0).round() as u64;
|
||||
progress_bar.set_position(percentage_diff);
|
||||
progress_bar.set_message("Total Percent Complete");
|
||||
}
|
||||
}
|
||||
});
|
||||
progress_bar.finish_with_message("All Checkpoints Complete");
|
||||
let return_data = match fs::read_to_string(&writeable_filepath) {
|
||||
let return_data = match fs::read_to_string(file_path.clone()) {
|
||||
Ok(return_data) => return_data,
|
||||
Err(e) => {
|
||||
println!("Failed to read data from: {:?}: {}", &writeable_filepath, e);
|
||||
println!("Failed to read data from: {:?}: {}", &file_path, e);
|
||||
std::process::abort();
|
||||
}
|
||||
};
|
||||
drop(tx);
|
||||
shutdown_tracer_provider();
|
||||
return_data.to_string()
|
||||
});
|
||||
@@ -410,7 +280,7 @@ async fn history_logging(
|
||||
}
|
||||
}
|
||||
|
||||
fn get_base_directory() -> PathBuf {
|
||||
pub fn get_base_directory() -> PathBuf {
|
||||
let home = env::var_os("HOME")
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| env::var_os("USERPROFILE").map(PathBuf::from))
|
||||
@@ -427,15 +297,6 @@ fn get_base_directory() -> PathBuf {
|
||||
}
|
||||
}
|
||||
|
||||
fn skipback(days: i64) -> ObjectId {
|
||||
let date_days_ago = Local::now() - Duration::days(days);
|
||||
let timestamp = date_days_ago.timestamp() as u32;
|
||||
let mut hex_timestamp = String::new();
|
||||
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap();
|
||||
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
|
||||
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
|
||||
}
|
||||
|
||||
fn init_tracer() -> Result<Option<sdktrace::Tracer>, TraceError> {
|
||||
let cfg = TelemetryConfig::load();
|
||||
if !cfg.TELEMETRY {
|
||||
|
||||
+37
-28
@@ -1,6 +1,17 @@
|
||||
"""
|
||||
This module handles the creation of local approval requests.
|
||||
"""
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import logging
|
||||
import os
|
||||
@@ -51,7 +62,7 @@ class LocalApprovalRequestor:
|
||||
batch_id = int(time.time())
|
||||
|
||||
purpose = (
|
||||
f"🎫 Local Approval 🎫 - {duration_minutes} mins - "
|
||||
f" Local Approval - {duration_minutes} mins - "
|
||||
f"batch:{batch_id} Client:{agent_id} User:{self.username}"
|
||||
)
|
||||
|
||||
@@ -103,11 +114,9 @@ class LocalApprovalRequestor:
|
||||
success_count = 0
|
||||
failure_count = 0
|
||||
|
||||
print(colorText(f"\n📦 Processing batch {batch_id}...", "cyan"))
|
||||
print(colorText(f"👤 Requested by: {self.username}", "cyan"))
|
||||
print(
|
||||
colorText(f"📊 Moving {len(agents)} agent(s) to local approval\n", "cyan")
|
||||
)
|
||||
print(colorText(f"\n Processing batch {batch_id}...", "cyan"))
|
||||
print(colorText(f" Requested by: {self.username}", "cyan"))
|
||||
print(colorText(f" Moving {len(agents)} agent(s) to local approval\n", "cyan"))
|
||||
|
||||
for agent in agents:
|
||||
try:
|
||||
@@ -125,11 +134,11 @@ class LocalApprovalRequestor:
|
||||
if not move_success:
|
||||
raise Exception("Failed to move to audit policy")
|
||||
|
||||
print(colorText(f"✓ {agent.hostname}", "green"))
|
||||
print(colorText(f" {agent.hostname}", "green"))
|
||||
success_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(colorText(f"✗ {agent.hostname}: {e}", "red"))
|
||||
print(colorText(f" {agent.hostname}: {e}", "red"))
|
||||
logger.error(f"Error processing agent {agent.hostname}: {e}")
|
||||
failure_count += 1
|
||||
|
||||
@@ -152,7 +161,7 @@ class LocalApprovalRequestor:
|
||||
]
|
||||
|
||||
# Display duration options
|
||||
print(colorText("\nâ±ï¸ Select Local Approval Duration:", "white"))
|
||||
print(colorText("\n Select Local Approval Duration:", "white"))
|
||||
print(colorText("=" * 50, "white"))
|
||||
|
||||
for i, (minutes, label) in enumerate(duration_options, start=1):
|
||||
@@ -166,36 +175,36 @@ class LocalApprovalRequestor:
|
||||
|
||||
if 1 <= choice <= len(duration_options):
|
||||
duration_minutes, duration_label = duration_options[choice - 1]
|
||||
print(colorText(f"✓ Selected: {duration_label}", "green"))
|
||||
print(colorText(f" Selected: {duration_label}", "green"))
|
||||
logger.info(f"User selected duration: {duration_minutes} minutes")
|
||||
else:
|
||||
print(colorText("⌠Invalid choice.", "red"))
|
||||
print(colorText("❌ Invalid choice.", "red"))
|
||||
logger.warning("Invalid duration choice")
|
||||
return
|
||||
|
||||
except ValueError:
|
||||
print(colorText("⌠Invalid input. Please enter a number.", "red"))
|
||||
print(colorText("❌ Invalid input. Please enter a number.", "red"))
|
||||
logger.warning("Invalid input for duration selection")
|
||||
return
|
||||
|
||||
# Select agents
|
||||
print(colorText("\n🎯 Select Agents for Local Approval:", "white"))
|
||||
print(colorText("\nSelect Agents for Local Approval:", "white"))
|
||||
agents = selectAgents(self.api)
|
||||
|
||||
if not agents:
|
||||
print(colorText("⌠No agents found or error retrieving agents.", "red"))
|
||||
print(colorText("❌ No agents found or error retrieving agents.", "red"))
|
||||
logger.warning("No agents selected or error retrieving agents")
|
||||
return
|
||||
|
||||
# Confirm with user
|
||||
print(colorText("\n📋 Summary:", "cyan"))
|
||||
print(colorText("\nSummary:", "cyan"))
|
||||
print(colorText(f" Duration: {duration_label}", "white"))
|
||||
print(colorText(f" Agents: {len(agents)}", "white"))
|
||||
|
||||
confirm = get_sanitized_input("\nProceed? (y/n): ").lower()
|
||||
|
||||
if confirm != "y":
|
||||
print(colorText("⌠Operation cancelled.", "yellow"))
|
||||
print(colorText("❌ Operation cancelled.", "yellow"))
|
||||
return
|
||||
|
||||
# Process the batch
|
||||
@@ -219,24 +228,24 @@ class LocalApprovalRequestor:
|
||||
failure_count: Number of failed operations
|
||||
"""
|
||||
print(colorText(f"\n{'=' * 60}", "white"))
|
||||
print(colorText("📊 Local Approval Summary", "cyan"))
|
||||
print(colorText(" Local Approval Summary", "cyan"))
|
||||
print(colorText("=" * 60, "white"))
|
||||
|
||||
print(colorText(f"✓ Successfully processed: {success_count}", "green"))
|
||||
print(colorText(f" Successfully processed: {success_count}", "green"))
|
||||
|
||||
if failure_count > 0:
|
||||
print(colorText(f"✗ Failed: {failure_count}", "red"))
|
||||
print(colorText(f" Failed: {failure_count}", "red"))
|
||||
|
||||
print(colorText(f"\n📦 Batch ID: {batch_id}", "cyan"))
|
||||
print(colorText(f"â±ï¸ Duration: {duration_label}", "cyan"))
|
||||
print(colorText(f"\n Batch ID: {batch_id}", "cyan"))
|
||||
print(colorText(f" Duration: {duration_label}", "cyan"))
|
||||
|
||||
print(colorText("=" * 60, "white"))
|
||||
print(colorText("\n💡 Next Steps:", "yellow"))
|
||||
print(colorText(" • Agents have been moved to audit policies", "white"))
|
||||
print(colorText(" • Local approvals are active", "white"))
|
||||
print(colorText("\n Next Steps:", "yellow"))
|
||||
print(colorText(" ✅ Agents have been moved to audit policies", "white"))
|
||||
print(colorText(" ✅ Local approvals are active", "white"))
|
||||
print(
|
||||
colorText(
|
||||
f" • Agents will return to enforcement after {duration_label}",
|
||||
f" ✅ Agents will return to enforcement after {duration_label}",
|
||||
"white",
|
||||
)
|
||||
)
|
||||
|
||||
+23
-23
@@ -88,7 +88,7 @@ def sortHashes(
|
||||
):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
prompt="Enter how many days of history to pull (1-150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
@@ -655,7 +655,7 @@ def section_header(title):
|
||||
|
||||
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
section_header("ðŸ› ï¸ ðŸ”’ Prepare to Enforce Policy ðŸ› ï¸ ðŸ”’")
|
||||
section_header("Prepare to Enforce Policy")
|
||||
print(
|
||||
colorText(
|
||||
"\nSequentially follow these steps to prepare a policy for enforcement:",
|
||||
@@ -670,11 +670,11 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
)
|
||||
)
|
||||
if not selected_policies:
|
||||
print(colorText(" [✗] No policies have been chosen", "red"))
|
||||
print(colorText(" [❌] No policies have been chosen", "red"))
|
||||
else:
|
||||
print(colorText("The following policies have been chosen:", "green"))
|
||||
for policy in selected_policies:
|
||||
print(colorText(f" [✓] {policy.name}", "green"))
|
||||
print(colorText(f" [✅] {policy.name}", "green"))
|
||||
|
||||
# Step 2: Destination Policy and Allowlist
|
||||
print(
|
||||
@@ -683,22 +683,22 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
if destination_policy:
|
||||
print(
|
||||
colorText(
|
||||
f" [✓] {destination_policy[0].name} has been selected as the destination policy",
|
||||
f" [✅] {destination_policy[0].name} has been selected as the destination policy",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(colorText(" [✗] No destination policy has been chosen", "red"))
|
||||
print(colorText(" [❌] No destination policy has been chosen", "red"))
|
||||
|
||||
if destination_allowlist:
|
||||
print(
|
||||
colorText(
|
||||
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
|
||||
f" [✅] {destination_allowlist[0].name} has been selected as allowlist",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
print(colorText(" [✗] No allowlist has been chosen", "red"))
|
||||
print(colorText(" [❌] No allowlist has been chosen", "red"))
|
||||
|
||||
# Step 3: Data Preparation
|
||||
print(
|
||||
@@ -713,9 +713,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Data has been fetched"
|
||||
" [✅] Data has been fetched"
|
||||
if os.path.exists(review_path)
|
||||
else " [✗] Data has not been fetched"
|
||||
else " [❌] Data has not been fetched"
|
||||
),
|
||||
"green" if os.path.exists(review_path) else "red",
|
||||
)
|
||||
@@ -723,7 +723,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check data fetch status", "red"
|
||||
" [❌] No policies selected, cannot check data fetch status", "red"
|
||||
)
|
||||
)
|
||||
|
||||
@@ -756,9 +756,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Reviewed hashes have been loaded"
|
||||
" [✅] Reviewed hashes have been loaded"
|
||||
if os.path.exists(approved_path)
|
||||
else " [✗] Reviewed hashes have not been loaded"
|
||||
else " [❌] Reviewed hashes have not been loaded"
|
||||
),
|
||||
"green" if os.path.exists(approved_path) else "red",
|
||||
)
|
||||
@@ -766,9 +766,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Path review list created"
|
||||
" [✅] Path review list created"
|
||||
if os.path.exists(second_review_path)
|
||||
else " [✗] Path review list has not been created"
|
||||
else " [❌] Path review list has not been created"
|
||||
),
|
||||
"green" if os.path.exists(second_review_path) else "red",
|
||||
)
|
||||
@@ -776,7 +776,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check reviewed hashes or path list",
|
||||
" [❌] No policies selected, cannot check reviewed hashes or path list",
|
||||
"red",
|
||||
)
|
||||
)
|
||||
@@ -812,9 +812,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Reviewed path list detected"
|
||||
" [✅] Reviewed path list detected"
|
||||
if os.path.exists(reviewed_path)
|
||||
else " [✗] Path review list has not been detected"
|
||||
else " [❌] Path review list has not been detected"
|
||||
),
|
||||
"green" if os.path.exists(reviewed_path) else "red",
|
||||
)
|
||||
@@ -825,9 +825,9 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(
|
||||
colorText(
|
||||
(
|
||||
" [✓] Preflight Path Exclusion List has been generated"
|
||||
" [✅] Preflight Path Exclusion List has been generated"
|
||||
if preflight_ready
|
||||
else " [✗] Preflight Path Exclusion List has not been generated"
|
||||
else " [❌] Preflight Path Exclusion List has not been generated"
|
||||
),
|
||||
"green" if preflight_ready else "red",
|
||||
)
|
||||
@@ -835,7 +835,7 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
" [✗] No policies selected, cannot check preflight status", "red"
|
||||
" [❌] No policies selected, cannot check preflight status", "red"
|
||||
)
|
||||
)
|
||||
|
||||
@@ -866,5 +866,5 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all
|
||||
print(colorText(" Apply approved hashes to allowlist", "cyan"))
|
||||
|
||||
# Utility Options
|
||||
print(colorText("F. 📂 - Open Working Directory", "cyan"))
|
||||
print(colorText("B. 🔚 - Back", "cyan"))
|
||||
print(colorText("F. Open Working Directory", "cyan"))
|
||||
print(colorText("B. Back", "cyan"))
|
||||
|
||||
+4
-4
@@ -145,13 +145,13 @@ class Hash:
|
||||
approved_count += 1
|
||||
except (ValueError, TypeError):
|
||||
logger.debug(
|
||||
"Needs Review: Scannermatch score is missing or invalid. — {e}"
|
||||
"Needs Review: Scannermatch score is missing or invalid. — {e}"
|
||||
)
|
||||
hash_obj.at_decision = "needs_review"
|
||||
needs_review_count += 1
|
||||
|
||||
logger.debug(
|
||||
f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}"
|
||||
f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}"
|
||||
)
|
||||
return hashes
|
||||
|
||||
@@ -443,13 +443,13 @@ class ExecutionHistoryRecord:
|
||||
approved_count += 1
|
||||
except (ValueError, TypeError) as e:
|
||||
logger.debug(
|
||||
f"Needs Review: Scannermatch score is missing or invalid. — {e}"
|
||||
f"Needs Review: Scannermatch score is missing or invalid. — {e}"
|
||||
)
|
||||
hash_obj.at_decision = "needs_review"
|
||||
needs_review_count += 1
|
||||
|
||||
logger.debug(
|
||||
f"Final counts — Needs Review: {needs_review_count}, "
|
||||
f"Final counts — Needs Review: {needs_review_count}, "
|
||||
f"Approved: {approved_count}, Unapproved: {unapproved_count}"
|
||||
)
|
||||
|
||||
|
||||
+1
-1
@@ -11,4 +11,4 @@ urllib3==2.5.0
|
||||
pyperclip==1.11.0
|
||||
|
||||
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
|
||||
airlock_libs==5.0.0
|
||||
airlock_libs==5.1.0
|
||||
+14
-14
@@ -38,7 +38,7 @@ logger = logging.getLogger(__name__)
|
||||
def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||||
agents = selectAgents(api)
|
||||
history_days = Selector.select_value(
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
prompt="Enter how many days of history to pull (1–150): ",
|
||||
value_type=int,
|
||||
valid_range=(1, 150),
|
||||
)
|
||||
@@ -60,7 +60,7 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
|
||||
except Exception as e:
|
||||
print(
|
||||
colorText(
|
||||
f"⌠Error retrieving history for {agent.hostname}: {e}", "red"
|
||||
f"❌ Error retrieving history for {agent.hostname}: {e}", "red"
|
||||
)
|
||||
)
|
||||
continue
|
||||
@@ -139,7 +139,7 @@ def findAgents(api, return_dataframe):
|
||||
|
||||
print(
|
||||
colorText(
|
||||
f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
|
||||
f"\n✓ Matched devices exported to: {working_dir}\\{filename}",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
@@ -148,7 +148,7 @@ def findAgents(api, return_dataframe):
|
||||
|
||||
|
||||
def collect_device_names() -> List[str]:
|
||||
print(colorText("🔠Device Search", "cyan"))
|
||||
print(colorText("🖥�� Device Search", "cyan"))
|
||||
print(
|
||||
colorText(
|
||||
"Enter the device hostnames you'd like to search for, one per line.", "cyan"
|
||||
@@ -185,7 +185,7 @@ def collect_device_names() -> List[str]:
|
||||
else:
|
||||
print(
|
||||
colorText(
|
||||
f"âš ï¸ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
|
||||
f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.",
|
||||
"yellow",
|
||||
)
|
||||
)
|
||||
@@ -235,8 +235,8 @@ def show_unmatched(
|
||||
]
|
||||
|
||||
if unmatched:
|
||||
logger.debug(f"âš ï¸ No matches for: {', '.join(unmatched)}")
|
||||
print(colorText(f"âš ï¸ No matches for: {', '.join(unmatched)}", "yellow"))
|
||||
logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
|
||||
print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
|
||||
|
||||
|
||||
def enrich_agents(agents: List["Agent"], policies: List["Policy"]):
|
||||
@@ -248,7 +248,7 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
||||
device_names = collect_device_names()
|
||||
if not device_names:
|
||||
logger.debug("No device names entered")
|
||||
print(colorText("âš ï¸ No device names entered.", "red"))
|
||||
print(colorText("⚠️ No device names entered.", "red"))
|
||||
return []
|
||||
|
||||
use_exact = choose_match_type()
|
||||
@@ -261,11 +261,11 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
||||
show_unmatched(device_names, matched_agents, use_exact)
|
||||
|
||||
if not matched_agents:
|
||||
logger.debug("⌠No matching devices found.")
|
||||
print(colorText("⌠No matching devices found.", "red"))
|
||||
logger.debug("❌ No matching devices found.")
|
||||
print(colorText("❌ No matching devices found.", "red"))
|
||||
return []
|
||||
|
||||
print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green"))
|
||||
print(colorText(f"✓ Found {len(matched_agents)} matching device(s).", "green"))
|
||||
logger.info("Matched agent hostnames:")
|
||||
rows = (len(matched_agents) + 2) // 3 # 3 columns
|
||||
for row in range(rows):
|
||||
@@ -283,8 +283,8 @@ def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]:
|
||||
)
|
||||
|
||||
if not matched_agents:
|
||||
logger.debug("⌠No matching devices remain after refinement.")
|
||||
print(colorText("⌠No matching devices remain after refinement.", "red"))
|
||||
logger.debug("❌ No matching devices remain after refinement.")
|
||||
print(colorText("❌ No matching devices remain after refinement.", "red"))
|
||||
return []
|
||||
|
||||
enrich_agents(matched_agents, policies)
|
||||
@@ -302,7 +302,7 @@ def moveAgentToRelatedPolicy(
|
||||
Args:
|
||||
api: AirlockAPIWrapper instance.
|
||||
agent: Agent object.
|
||||
policy_relationship_map: Dict mapping enforcement → audit.
|
||||
policy_relationship_map: Dict mapping enforcement â–€ –€™ audit.
|
||||
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
|
||||
"""
|
||||
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
|
||||
@@ -27,9 +27,8 @@ import tqdm
|
||||
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_system_json
|
||||
from utils.setup import get_base_directory
|
||||
from utils.utils import areYouSure, colorText, get_sanitized_input
|
||||
from utils.utils import colorText
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -233,17 +232,3 @@ def skipback(days):
|
||||
hex_timestamp = format(timestamp, "08x")
|
||||
objectid_hex = hex_timestamp + "0000000000000000"
|
||||
return ObjectId(objectid_hex)
|
||||
|
||||
|
||||
def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
|
||||
policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
for enforcement_policy, audit_policy in policy_relationship_map.items():
|
||||
api.policy_clone(enforcement_policy, audit_policy)
|
||||
api.policy_set_auditmode(audit_policy, "1")
|
||||
|
||||
|
||||
def confirmUpdateAfromE(api: AirlockAPIWrapper):
|
||||
areYouSure()
|
||||
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
if confirmation.strip() == "I AGREE":
|
||||
updateAuditPoliciesFromEnforcementPolices(api)
|
||||
|
||||
+8
-4
@@ -184,7 +184,7 @@ class Selector:
|
||||
print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
|
||||
return selected
|
||||
elif mode == "e":
|
||||
print(colorText(f"🚫 Excluded {len(selected)} item(s).", "yellow"))
|
||||
print(colorText(f"👫 Excluded {len(selected)} item(s).", "yellow"))
|
||||
return [item for item in items if item not in selected]
|
||||
else:
|
||||
print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
|
||||
@@ -279,7 +279,9 @@ class Selector:
|
||||
df = df[columns]
|
||||
|
||||
items = [row for _, row in df.iterrows()]
|
||||
label_func = lambda row: str(row.to_dict())
|
||||
|
||||
def label_func(row):
|
||||
return str(row.to_dict())
|
||||
|
||||
result = Selector._select_from_list(
|
||||
items,
|
||||
@@ -311,7 +313,9 @@ class Selector:
|
||||
df = df[columns]
|
||||
|
||||
items = df.to_dict("records")
|
||||
label_func = lambda row: " | ".join(str(row[col]) for col in df.columns)
|
||||
|
||||
def label_func(row):
|
||||
return " | ".join(str(row[col]) for col in df.columns)
|
||||
|
||||
# Show rows first
|
||||
print(colorText(header, "cyan"))
|
||||
@@ -346,7 +350,7 @@ class Selector:
|
||||
print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
|
||||
return [pd.Series(row) for row in selected]
|
||||
elif mode == "e":
|
||||
print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow"))
|
||||
print(colorText(f"👫 Excluded {len(selected)} row(s).", "yellow"))
|
||||
return [pd.Series(row) for row in items if row not in selected]
|
||||
else:
|
||||
print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
|
||||
|
||||
+3
-3
@@ -118,14 +118,14 @@ def setup():
|
||||
logging.debug(f"{name.capitalize()} directory ensured at: {path}")
|
||||
|
||||
# Load system config (immutable)
|
||||
system_config = load_system_config()
|
||||
load_system_config()
|
||||
|
||||
# Configure logging with system-defined log level
|
||||
log_level = get_system_value("LOG_LEVEL", str, "INFO")
|
||||
configure_logging(dirs["logs"], log_level)
|
||||
|
||||
# Load user config (mutable)
|
||||
user_config = load_user_config(dirs["config"])
|
||||
load_user_config(dirs["config"])
|
||||
|
||||
# Set up .env file - ONLY for WORKING_DIR (runtime-configurable value)
|
||||
env_path = base_dir / ".env"
|
||||
@@ -155,6 +155,6 @@ def setup():
|
||||
for subfolder in subfolders:
|
||||
subfolder_path = folder_path / subfolder
|
||||
subfolder_path.mkdir(parents=True, exist_ok=True)
|
||||
logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
|
||||
logging.debug(f"'{subfolder}' subfolder created at: {subfolder_path}")
|
||||
|
||||
logging.info("✅ Setup complete")
|
||||
|
||||
Reference in New Issue
Block a user