diff --git a/Loxide.py b/Loxide.py index 748dffe..07d1ec6 100644 --- a/Loxide.py +++ b/Loxide.py @@ -30,10 +30,10 @@ from typing import Optional from bson import ObjectId import dotenv +import plotext as plt from textual.app import App, ComposeResult -from textual.containers import Horizontal, Vertical +from textual.containers import Horizontal, Vertical, VerticalScroll from textual.message import Message -from textual.reactive import reactive from textual.screen import Screen from textual.widgets import ( Button, @@ -51,7 +51,6 @@ import airlock_libs from models.agent import Agent from models.policy import Policy from services.API import AirlockAPIWrapper -from services.security import getAPI from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen from TUI.Screens.moveagentworkflowscreen import MoveAgentWorkflowScreen from TUI.Screens.otpactivityscreen import OTPActivitiesScreen @@ -59,19 +58,19 @@ from TUI.Screens.otprevokescreen import OTPRevokeScreen from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen from TUI.Screens.policyprepworkflowscreen import PolicyPrepWorkflowScreen from TUI.Screens.quietagentworkflowscreen import QuietAgentWorkflowScreen -from TUI.Themes.theme_amber_terminal import get_amber_terminal_theme -from TUI.Themes.theme_retro_terminal import get_retro_terminal_theme -from TUI.Themes.themeselector import ThemeSelector from TUI.Widgets.agentmoveoperations import AgentMoveOperations +from TUI.Widgets.multiagentselector import MultiAgentSelector from TUI.Widgets.policytreewidget import PolicyTreeWidget from TUI.Widgets.resultsdisplay import ResultsDisplay from TUI.Widgets.serverlogwidget import ServerLogWidget +from TUI.Widgets.settingswidget import SettingsWidget from utils.configmanager import ( get_system_value, get_user_value, load_env, save_user_config, ) +from utils.security import getAPI from utils.setup import get_base_directory, setup from utils.utils import irtang, open_directory @@ -87,6 +86,81 @@ _APP_RESTART_REASON = None logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Helper for plotext charts +# --------------------------------------------------------------------------- +def create_plotext_chart( + chart_type: str, + labels: list, + values: list, + height: int = 10, + width: int = 60, + title: str = "", + color: str = None, +) -> Static: + """Create a plotext chart and return it as a Static widget with the rendered output.""" + # Completely reset plotext state + plt.clear_data() + plt.clear_figure() + plt.clear_color() + plt.clear_terminal() + plt.clf() + + # Set theme and canvas size + plt.theme("clear") + plt.plotsize(width, height) + + if chart_type == "simple_bar": + # Simple horizontal bar - clean text-based look + if color: + plt.simple_bar(labels, values, width=width, color=color) + else: + plt.simple_bar(labels, values, width=width) + elif chart_type == "bar_h": + # Calculate integer ticks for bar charts + max_val = max(values) if values else 0 + if max_val <= 5: + ticks = list(range(0, max_val + 1)) + elif max_val <= 10: + ticks = list(range(0, max_val + 1, 2)) + else: + step = max(1, int(max_val / 5)) + ticks = list(range(0, max_val + step, step)) + plt.bar(labels, values, orientation="h", width=0.5) + plt.xlabel("Count") + plt.xticks(ticks) + plt.grid(False, False) + elif chart_type == "bar_v": + max_val = max(values) if values else 0 + if max_val <= 5: + ticks = list(range(0, max_val + 1)) + elif max_val <= 10: + ticks = list(range(0, max_val + 1, 2)) + else: + step = max(1, int(max_val / 5)) + ticks = list(range(0, max_val + step, step)) + plt.bar(labels, values, orientation="v", width=0.5) + plt.ylabel("Count") + plt.yticks(ticks) + plt.grid(False, False) + + if title: + plt.title(title) + + # Use build() to get the chart string + chart_output = plt.build() + + # Clear state after building + plt.clf() + plt.clear_data() + + # Create Static widget with the chart - disable selection + chart_widget = Static(chart_output, classes="chart_display") + chart_widget.can_focus = False + + return chart_widget + + # --------------------------------------------------------------------------- # helper to persist TEXTUAL_THEME to *user* config and mirror to .env # --------------------------------------------------------------------------- @@ -110,31 +184,30 @@ def _persist_user_theme(theme_name: str) -> None: # --------------------------------------------------------------------------- class MainMenuScreen(Screen): api: AirlockAPIWrapper - current_tab = reactive("") BUTTON_DEFS = { "agent_actions": [ { - "label": "🖥️ - Multi-Agent Operations", + "label": "🔧 - Multi-Agent Operations", "id": "move_agent_workflow_button", - "description": "Select agents to: Move policies, Generate OTPs, Toggle audit/enforcement, View history, Export data", + "description": "Select agents to:\n • Move policies\n • Generate OTPs\n • Toggle audit/enforcement\n • View history\n • Export data", }, { - "label": "🎫 - Review and approve OTP Activities", + "label": "🎫 - Review and approve OTP Activities", "id": "otp_activities_button", }, { - "label": "🛑 - Revoke Active OTP Session", + "label": "🛑 - Revoke Active OTP Session", "id": "otp_revoke_button", }, ], "policy": [ { - "label": "⚖️ - Prepare Policy For Enforcement", + "label": "⚙️ - Prepare Policy For Enforcement", "id": "policy_prep_button", }, { - "label": "🔕 - Find and Move Quiet Hosts to Enforcement", + "label": "🔍 - Find and Move Quiet Hosts to Enforcement", "id": "find_quiet_button", }, ], @@ -154,35 +227,81 @@ class MainMenuScreen(Screen): def _make_buttons_for(self, tab_id: str) -> Vertical: defs = self.BUTTON_DEFS.get(tab_id, []) + + # Special layout for agent_actions tab + if tab_id == "agent_actions": + left_col = Vertical() + left_col.styles.width = "1fr" + right_col = Vertical() + right_col.styles.width = "1fr" + + for item in defs: + if isinstance(item, dict): + label = item["label"] + btn_id = item["id"] + description = item.get("description") + else: + label, btn_id = item + description = None + + btn = Button(label, id=btn_id) + btn.styles.width = "100%" + btn.styles.margin = (0, 1, 1, 0) + + btn_container = Vertical() + btn_container.styles.height = "auto" + btn_container.compose_add_child(btn) + + if description: + desc_text = Static(description, classes="button_description") + desc_text.styles.width = "100%" + desc_text.styles.color = "ansi_bright_black" + desc_text.styles.text_align = "left" + desc_text.styles.margin = (0, 0, 1, 0) + btn_container.compose_add_child(desc_text) + + # Multi-Agent on left, OTP buttons on right + if "otp" in btn_id: + right_col.compose_add_child(btn_container) + else: + left_col.compose_add_child(btn_container) + + row = Horizontal() + row.styles.width = "100%" + row.styles.height = "auto" + row.compose_add_child(left_col) + row.compose_add_child(right_col) + + return Vertical(row) + + # Default single-column layout for other tabs widgets = [] for item in defs: - # Support both old tuple format and new dict format if isinstance(item, dict): label = item["label"] btn_id = item["id"] description = item.get("description") else: - # Old tuple format: (label, id) label, btn_id = item description = None btn = Button(label, id=btn_id) btn.styles.width = "100%" + btn.styles.margin = (0, 1, 1, 0) widgets.append(btn) - # Add description text if provided if description: desc_text = Static(description, classes="button_description") desc_text.styles.width = "100%" desc_text.styles.color = "ansi_bright_black" - desc_text.styles.text_align = "center" + desc_text.styles.text_align = "left" desc_text.styles.margin = (0, 0, 1, 0) widgets.append(desc_text) return Vertical(*widgets) def compose(self) -> ComposeResult: - yield Header(show_clock=True, icon="âš™") + yield Header(show_clock=True, icon="⚙") tabs = [ Tab("Agents", id="agent_actions"), @@ -203,6 +322,50 @@ class MainMenuScreen(Screen): def on_mount(self) -> None: self.switch_tab("agent_actions") + # Delay version check to ensure UI is fully ready for notifications + if hasattr(self.app, "_version_checker"): + self.set_timer(1.0, self._do_version_check) + + def _do_version_check(self) -> None: + """Perform version check and show notification if update available.""" + logger.info("_do_version_check called") + + if not hasattr(self.app, "_version_checker"): + logger.warning( + "No _version_checker on app - was create_update_notifier called?" + ) + return + + checker = self.app._version_checker + logger.info(f"Checker exists: {checker}") + logger.info("Calling check_now(force=True)...") + + try: + result = checker.check_now(force=True) + logger.info( + f"Version check result: update_available={result.update_available}, latest={result.latest_version}, current={result.current_version}, error={result.error}" + ) + + if result.error: + logger.warning(f"Version check returned error: {result.error}") + return + + if result.update_available: + # Always nag on startup - don't check dismissed version here + # User can dismiss from settings but we still want startup reminder + logger.info(f"Showing toast for update {result.latest_version}") + msg = f"🆕 Update available: {result.latest_version}\nGo to Settings to download" + self.app.notify( + msg, title="Loxide Update Available", severity="warning", timeout=15 + ) + logger.info("Toast notification called") + else: + logger.info( + f"No update available - current {result.current_version} is latest" + ) + except Exception as e: + logger.error(f"Version check failed: {e}", exc_info=True) + def on_key(self, event) -> None: """Handle up/down arrow keys for button navigation.""" if event.key == "down": @@ -252,7 +415,6 @@ class MainMenuScreen(Screen): buttons[new_index].focus() def switch_tab(self, tab_id: str) -> None: - self.current_tab = tab_id content = self.query_one("#content", Vertical) content.remove_children() @@ -267,14 +429,13 @@ class MainMenuScreen(Screen): elif tab_id == "statistics": content.mount(self._create_statistics_widget()) elif tab_id == "settings": - content.mount(ThemeSelector()) + content.mount(SettingsWidget()) else: content.mount(Static(f"Unknown tab: {tab_id}")) def _create_statistics_widget(self) -> Vertical: """Create the enhanced statistics display widget with time period selection.""" stats_container = Vertical(id="statistics_container") - # Build all widgets first, then mount them all at once widgets_to_mount = [] @@ -306,41 +467,56 @@ class MainMenuScreen(Screen): ) time_select.styles.width = 20 - # Refresh button - refresh_btn = Button("🔄 Refresh", id="stats_refresh_btn", variant="primary") - refresh_btn.styles.margin = (0, 0, 0, 2) + # Fetch/Refresh button - starts as "Fetch" + if self.current_stats_days in self.statistics_cache: + button_label = "🔍„ Refresh" + else: + button_label = "📊 Fetch Statistics" + + fetch_btn = Button(button_label, id="stats_refresh_btn", variant="primary") + fetch_btn.styles.margin = (0, 0, 0, 2) # Compose controls controls.compose_add_child(time_label) controls.compose_add_child(time_select) - controls.compose_add_child(refresh_btn) + controls.compose_add_child(fetch_btn) widgets_to_mount.append(controls) # Timestamp and status area status_container = Vertical(id="stats_status_area") - status_container.styles.margin = (0, 2, 1, 2) + status_container.styles.margin = (0, 2, 0, 2) widgets_to_mount.append(status_container) - # Data display area - using horizontal layout - data_container = Horizontal(id="stats_data_area") + # Data display area - scrollable vertical layout + data_container = VerticalScroll(id="stats_data_area") data_container.styles.margin = (0, 2, 0, 2) - data_container.styles.height = "auto" widgets_to_mount.append(data_container) # Compose all widgets into the container for widget in widgets_to_mount: stats_container.compose_add_child(widget) - # Schedule data fetch after the widget is mounted - if self.current_stats_days not in self.statistics_cache: - self.call_after_refresh( - lambda: self._fetch_execution_statistics(self.current_stats_days) - ) - else: + # If cached, display it. Otherwise show prompt to fetch + if self.current_stats_days in self.statistics_cache: self.call_after_refresh(self._update_statistics_display) + else: + self.call_after_refresh(self._show_fetch_prompt) return stats_container + def _show_fetch_prompt(self) -> None: + """Show a prompt to click Fetch to load statistics.""" + try: + status_area = self.query_one("#stats_status_area", Vertical) + status_area.remove_children() + + prompt = Static("Click 'Fetch Statistics' to load execution data.") + prompt.styles.text_style = "italic" + prompt.styles.text_align = "center" + status_area.mount(prompt) + except Exception as e: + logger.warning(f"Could not show fetch prompt: {e}") + def _skipback(self, days: int) -> ObjectId: """Generate a MongoDB ObjectId for a given number of days ago.""" date_days_ago = datetime.now(UTC) - timedelta(days=days) @@ -372,10 +548,9 @@ class MainMenuScreen(Screen): ) logger.debug(f"Calling with exec_types={exec_types}, days={days}") - # Use pull_policy_exec_histories - must pass all 4 args in order: (api, policy_name, type, days) - # Empty string for policy_name means all policies + # Use pull_policy_exec_histories - pass None for policy_name to get all policies execs = airlock_libs.pull_policy_exec_histories( - self.app.api, None, str(exec_types), days # Empty string = all policies + self.app.api, None, str(exec_types), days # None = all policies ) logger.debug(f"Received response, length: {len(execs) if execs else 0}") @@ -449,24 +624,41 @@ class MainMenuScreen(Screen): 18: "Trusted Browser Metadata Execution", } + # Categorize execution types + # "Untrusted" types: contain "Untrusted" in name + untrusted_types = {2, 3, 13, 14} + # "Block" types: contain "Block" in name (Blocked, Blocklist) + block_types = {1, 6, 7, 12, 15, 16} + # Count execution types type_counter = Counter() - policy_counter = Counter() - hostname_counter = Counter() + + # Separate counters for untrusted vs block + policy_counter_untrusted = Counter() + policy_counter_block = Counter() + hostname_counter_untrusted = Counter() + hostname_counter_block = Counter() for exec_record in executions: exec_type = exec_record.get("type", -1) type_counter[exec_type] += 1 policy_name = exec_record.get("policyname", "Unknown") - policy_counter[policy_name] += 1 - hostname = exec_record.get("hostname", "Unknown") - hostname_counter[hostname] += 1 - # Get top 5 policies and hostnames - top_policies = policy_counter.most_common(5) - top_hostnames = hostname_counter.most_common(5) + # Categorize by type + if exec_type in untrusted_types: + policy_counter_untrusted[policy_name] += 1 + hostname_counter_untrusted[hostname] += 1 + elif exec_type in block_types: + policy_counter_block[policy_name] += 1 + hostname_counter_block[hostname] += 1 + + # Get top 5 for each category + top_policies_untrusted = policy_counter_untrusted.most_common(5) + top_policies_block = policy_counter_block.most_common(5) + top_hostnames_untrusted = hostname_counter_untrusted.most_common(5) + top_hostnames_block = hostname_counter_block.most_common(5) # Format execution type counts type_counts = [] @@ -477,22 +669,36 @@ class MainMenuScreen(Screen): return { "total_executions": len(executions), "type_counts": type_counts, - "top_policies": top_policies, - "top_hostnames": top_hostnames, + "top_policies_untrusted": top_policies_untrusted, + "top_policies_block": top_policies_block, + "top_hostnames_untrusted": top_hostnames_untrusted, + "top_hostnames_block": top_hostnames_block, } def _update_statistics_display(self) -> None: - """Update the statistics display with cached data using horizontal layout.""" + """Update the statistics display with cached data using vertical scrollable layout.""" try: - data_area = self.query_one("#stats_data_area", Horizontal) + data_area = self.query_one("#stats_data_area", VerticalScroll) status_area = self.query_one("#stats_status_area", Vertical) except Exception as e: logger.warning(f"Could not find statistics display areas: {e}") return + # Clear any existing plotext state + plt.clear_data() + plt.clear_figure() + data_area.remove_children() status_area.remove_children() + # Get theme color for charts + theme_name = ( + self.app._textual_theme + if hasattr(self.app, "_textual_theme") + else "textual-dark" + ) + chart_color = None # use plotext default + # Get cached data cache_entry = self.statistics_cache.get(self.current_stats_days) if not cache_entry: @@ -511,109 +717,187 @@ class MainMenuScreen(Screen): timestamp_label.styles.margin = (0, 0, 1, 0) status_area.mount(timestamp_label) - # Create three columns for horizontal layout - # Column 1: Summary and Execution Types - col1 = Vertical() - col1.styles.width = "1fr" - col1.styles.padding = (0, 1, 0, 0) + # Create single scrollable column for better chart display + main_col = Vertical() + main_col.styles.width = "100%" + main_col.styles.height = "auto" - # Column 2: Top Policies - col2 = Vertical() - col2.styles.width = "1fr" - col2.styles.padding = (0, 1) - - # Column 3: Top Machines and System Overview - col3 = Vertical() - col3.styles.width = "1fr" - col3.styles.padding = (0, 0, 0, 1) - - # === COLUMN 1: Summary and Execution Types === - total = Static(f"📊 Total Non-Trusted Executions: {stats['total_executions']}") - total.styles.text_style = "bold" - total.styles.margin = (0, 0, 1, 0) - col1.compose_add_child(total) - - # Execution type breakdown - if stats["type_counts"]: - type_header = Static("📋 Execution Types") - type_header.styles.text_style = "bold" - type_header.styles.margin = (1, 0, 0, 0) - col1.compose_add_child(type_header) - - for type_name, count in stats["type_counts"]: - # Abbreviate long type names - short_name = type_name.replace("Execution", "Exec").replace( - "[Audit]", "[Aud]" - ) - type_line = Static(f" {short_name}: {count}") - type_line.styles.margin = (0, 0, 0, 1) - col1.compose_add_child(type_line) - - # === COLUMN 2: Top Policies === - if stats["top_policies"]: - policy_header = Static("⚖️ Top 5 Policies") - policy_header.styles.text_style = "bold" - policy_header.styles.margin = (0, 0, 0, 0) - col2.compose_add_child(policy_header) - - for i, (policy_name, count) in enumerate(stats["top_policies"], 1): - policy_line = Static(f"{i}. {policy_name}: {count}") - policy_line.styles.margin = (0, 0, 0, 1) - col2.compose_add_child(policy_line) - - # === COLUMN 3: Top Machines and System Overview === - if stats["top_hostnames"]: - host_header = Static("🖥️ Top 5 Machines") - host_header.styles.text_style = "bold" - host_header.styles.margin = (0, 0, 0, 0) - col3.compose_add_child(host_header) - - for i, (hostname, count) in enumerate(stats["top_hostnames"], 1): - host_line = Static(f"{i}. {hostname}: {count}") - host_line.styles.margin = (0, 0, 0, 1) - col3.compose_add_child(host_line) - - # System Overview in column 3 + # === SYSTEM OVERVIEW (First) === total_agents = len(self.app.devices) if self.app.devices else 0 total_policies = len(self.app.policies) if self.app.policies else 0 if self.app.devices: - enforced_agents = sum( - 1 for agent in self.app.devices if getattr(agent, "enforcement", False) - ) - audit_agents = sum( - 1 - for agent in self.app.devices - if not getattr(agent, "enforcement", False) - ) + # status: 0=Offline, 1=Online, 2=Hidden, 3=Safemode online_agents = sum( - 1 for agent in self.app.devices if getattr(agent, "online", False) + 1 for agent in self.app.devices if getattr(agent, "status", 0) == 1 + ) + offline_agents = sum( + 1 for agent in self.app.devices if getattr(agent, "status", 0) == 0 + ) + hidden_agents = sum( + 1 for agent in self.app.devices if getattr(agent, "status", 0) == 2 + ) + safemode_agents = sum( + 1 for agent in self.app.devices if getattr(agent, "status", 0) == 3 ) - offline_agents = total_agents - online_agents else: - enforced_agents = audit_agents = online_agents = offline_agents = 0 + online_agents = offline_agents = hidden_agents = safemode_agents = 0 system_header = Static("ℹ️ System Overview") system_header.styles.text_style = "bold" - system_header.styles.margin = (2, 0, 0, 0) - col3.compose_add_child(system_header) + system_header.styles.margin = (0, 0, 1, 0) + main_col.compose_add_child(system_header) - system_data = [ - f"Agents: {total_agents} ({online_agents} online)", - f"Policies: {total_policies}", - f"Enforced: {enforced_agents}", - f"Audit: {audit_agents}", + # Total counts + totals_line = Static( + f"Total Agents: {total_agents} Total Policies: {total_policies}" + ) + totals_line.styles.margin = (0, 0, 1, 1) + main_col.compose_add_child(totals_line) + + # Agent status with plotext bar chart and percentages + status_data = [ + ("Online", online_agents), + ("Offline", offline_agents), + ("Hidden", hidden_agents), + ("Safemode", safemode_agents), ] - for line in system_data: - line_widget = Static(f" {line}") - line_widget.styles.margin = (0, 0, 0, 1) - col3.compose_add_child(line_widget) + labels = [] + counts = [] + for status_name, count in status_data: + percentage = int((count / total_agents) * 100) if total_agents > 0 else 0 + labels.append(f"{status_name} ({percentage}%)") + counts.append(count) - # Mount all columns - data_area.mount(col1) - data_area.mount(col2) - data_area.mount(col3) + chart = create_plotext_chart( + "simple_bar", labels, counts, height=15, width=60, color=chart_color + ) + chart.styles.margin = (0, 0, 1, 0) + main_col.compose_add_child(chart) + + # === EXECUTION SUMMARY === + total_execs = stats["total_executions"] + total = Static(f"📊 Total Untrusted Executions: {total_execs}") + total.styles.text_style = "bold" + total.styles.margin = (1, 0, 1, 0) + main_col.compose_add_child(total) + + # Execution type breakdown with percentages + if stats["type_counts"]: + type_header = Static("📋 Execution Types") + type_header.styles.text_style = "bold" + type_header.styles.margin = (1, 0, 0, 0) + main_col.compose_add_child(type_header) + + labels = [] + counts = [] + for type_name, count in stats["type_counts"]: + short_name = type_name.replace("Execution", "Exec").replace( + "[Audit]", "[Aud]" + ) + if len(short_name) > 18: + short_name = short_name[:15] + "..." + percentage = int((count / total_execs) * 100) if total_execs > 0 else 0 + labels.append(f"{short_name} ({percentage}%)") + counts.append(count) + + chart = create_plotext_chart( + "simple_bar", labels, counts, height=15, width=60, color=chart_color + ) + chart.styles.margin = (0, 0, 1, 0) + main_col.compose_add_child(chart) + + # === TOP POLICIES - Untrusted === + if stats["top_policies_untrusted"]: + policy_header = Static("⚖️ Most Active Policies (Untrusted)") + policy_header.styles.text_style = "bold" + policy_header.styles.margin = (1, 0, 0, 0) + main_col.compose_add_child(policy_header) + + policy_total = sum(count for _, count in stats["top_policies_untrusted"]) + labels = [] + values = [] + for policy_name, count in stats["top_policies_untrusted"]: + short_name = policy_name[:18] if len(policy_name) > 18 else policy_name + percentage = ( + int((count / policy_total) * 100) if policy_total > 0 else 0 + ) + labels.append(f"{short_name} ({percentage}%)") + values.append(count) + chart = create_plotext_chart( + "simple_bar", labels, values, height=12, width=60, color=chart_color + ) + chart.styles.margin = (0, 0, 1, 0) + main_col.compose_add_child(chart) + + # === TOP POLICIES - Blocked === + if stats["top_policies_block"]: + policy_header = Static("⚖️ Most Active Policies (Blocked)") + policy_header.styles.text_style = "bold" + policy_header.styles.margin = (1, 0, 0, 0) + main_col.compose_add_child(policy_header) + + policy_total = sum(count for _, count in stats["top_policies_block"]) + labels = [] + values = [] + for policy_name, count in stats["top_policies_block"]: + short_name = policy_name[:18] if len(policy_name) > 18 else policy_name + percentage = ( + int((count / policy_total) * 100) if policy_total > 0 else 0 + ) + labels.append(f"{short_name} ({percentage}%)") + values.append(count) + chart = create_plotext_chart( + "simple_bar", labels, values, height=12, width=60, color=chart_color + ) + chart.styles.margin = (0, 0, 1, 0) + main_col.compose_add_child(chart) + + # === TOP MACHINES - Untrusted === + if stats["top_hostnames_untrusted"]: + host_header = Static("🖥️ Most Active Machines (Untrusted)") + host_header.styles.text_style = "bold" + host_header.styles.margin = (1, 0, 0, 0) + main_col.compose_add_child(host_header) + + host_total = sum(count for _, count in stats["top_hostnames_untrusted"]) + labels = [] + values = [] + for hostname, count in stats["top_hostnames_untrusted"]: + short_name = hostname[:18] if len(hostname) > 18 else hostname + percentage = int((count / host_total) * 100) if host_total > 0 else 0 + labels.append(f"{short_name} ({percentage}%)") + values.append(count) + chart = create_plotext_chart( + "simple_bar", labels, values, height=12, width=60, color=chart_color + ) + chart.styles.margin = (0, 0, 1, 0) + main_col.compose_add_child(chart) + + # === TOP MACHINES - Blocked === + if stats["top_hostnames_block"]: + host_header = Static("🖥️ Most Active Machines (Blocked)") + host_header.styles.text_style = "bold" + host_header.styles.margin = (1, 0, 0, 0) + main_col.compose_add_child(host_header) + + host_total = sum(count for _, count in stats["top_hostnames_block"]) + labels = [] + values = [] + for hostname, count in stats["top_hostnames_block"]: + short_name = hostname[:18] if len(hostname) > 18 else hostname + percentage = int((count / host_total) * 100) if host_total > 0 else 0 + labels.append(f"{short_name} ({percentage}%)") + values.append(count) + chart = create_plotext_chart( + "simple_bar", labels, values, height=12, width=60, color=chart_color + ) + chart.styles.margin = (0, 0, 1, 0) + main_col.compose_add_child(chart) + + # Mount the main scrollable column + data_area.mount(main_col) def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None: self.switch_tab(event.tab.id) @@ -625,13 +909,23 @@ class MainMenuScreen(Screen): if new_days != self.current_stats_days: self.current_stats_days = new_days - # Fetch if not cached - if new_days not in self.statistics_cache: - self._fetch_execution_statistics(new_days) - else: - self._update_statistics_display() + # Update button label based on cache status + try: + button = self.query_one("#stats_refresh_btn", Button) + if new_days in self.statistics_cache: + button.label = "🔍„ Refresh" + # Display cached data + self._update_statistics_display() + else: + button.label = "📊 Fetch Statistics" + # Show fetch prompt + self._show_fetch_prompt() + except Exception as e: + logger.warning(f"Could not update button or display: {e}") - def on_multi_agent_selector_agents_selected() -> None: + def on_multi_agent_selector_agents_selected( + self, message: MultiAgentSelector.AgentsSelected + ) -> None: """Handle selected agents from AgentSelector.""" global _APP_RESTART_REASON selected_agents = message.selected_agents @@ -640,10 +934,10 @@ class MainMenuScreen(Screen): _APP_RESTART_REASON = ("multi_agent_action", selected_agents) self.app.exit() - def on_theme_selector_theme_selected( - self, message: ThemeSelector.ThemeSelected + def on_settings_widget_theme_selected( + self, message: SettingsWidget.ThemeSelected ) -> None: - """Handle theme selection from ThemeSelector.""" + """Handle theme selection from SettingsWidget.""" global _APP_RESTART_REASON _persist_user_theme(message.theme_name) _APP_RESTART_REASON = ("restart",) @@ -713,7 +1007,7 @@ class MainMenuScreen(Screen): logger.info("Toggling enforcement for device: %s", message.device.hostname) try: - from services.agenthandler import moveAgentToRelatedPolicy + from TUI.Widgets.agentmoveoperations import moveAgentToRelatedPolicy from utils.configmanager import get_system_json policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}") @@ -768,6 +1062,12 @@ class MainMenuScreen(Screen): # Handle statistics refresh button if button_id == "stats_refresh_btn": self._fetch_execution_statistics(self.current_stats_days) + # Update button label to "Refresh" after first fetch + try: + button = event.button + button.label = "🔍„ Refresh" + except Exception as e: + logger.warning(f"Could not update button label: {e}") event.stop() return @@ -826,6 +1126,22 @@ class Loxide(App[Message]): content-align: center middle; text-align: center; } + .chart_display { + overflow: hidden; + } + .chart_display:hover { + background: transparent; + } + #statistics_container { + height: 1fr; + } + #stats_status_area { + height: auto; + } + #stats_data_area { + height: 1fr; + scrollbar-gutter: stable; + } """ BINDINGS = [ ("q", "quit", "Quit"), @@ -867,8 +1183,6 @@ class Loxide(App[Message]): self.devices = None def on_mount(self, api: AirlockAPIWrapper) -> None: - self.register_theme(get_retro_terminal_theme()) - self.register_theme(get_amber_terminal_theme()) self.theme = self._textual_theme self.push_screen(MainMenuScreen()) diff --git a/TUI/Screens/quietagentworkflowscreen.py b/TUI/Screens/quietagentworkflowscreen.py index bca4b31..fcac2dc 100644 --- a/TUI/Screens/quietagentworkflowscreen.py +++ b/TUI/Screens/quietagentworkflowscreen.py @@ -38,9 +38,9 @@ from textual.widgets import Button, DataTable, Footer, Header, Input, Static from models.policy import Policy from services.API import AirlockAPIWrapper -from services.policyhandler import getPolicyInfo from TUI.Widgets.policyselector import PolicySelector from utils.configmanager import load_env +from utils.executionfetcher import getExecutions logger = logging.getLogger(__name__) @@ -429,7 +429,7 @@ class QuietAgentWorkflowScreen(Screen): return # Get execution history (this shows progress bars in terminal via airlock_libs) - policy_exec_history = getPolicyInfo( + policy_exec_history = getExecutions( self.api, self.selected_policy, [1, 2, 6, 7], self.history_days ) diff --git a/TUI/Themes/theme_amber_terminal.py b/TUI/Themes/theme_amber_terminal.py deleted file mode 100644 index 738be38..0000000 --- a/TUI/Themes/theme_amber_terminal.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -from textual.color import Color -from textual.theme import Theme - - -def get_amber_terminal_theme(): - """Amber CRT theme with compensated brightness for blending.""" - return Theme( - name="amber-terminal", - background=Color.parse("#000000"), # pure black - primary=Color.parse("#ffb733"), # bright amber - secondary=Color.parse("#e69500"), # strong amber - success=Color.parse("#ffb733"), - warning=Color.parse("#ffff66"), - error=Color.parse("#ff3300"), - surface=Color.parse("#49331a"), # brighter brown for blending - ) - - -AMBER_TERMINAL_CSS = """ -Screen { - align: center middle; - background: #000000; /* force black */ - color: #ffb733; /* force amber text */ -} - -.widget { - border: tall #ffb733; /* force amber border */ - background: #3a1f00; /* compensated surface */ - width: 80%; -} - -* { - font-family: "Courier New", monospace; -} -""" diff --git a/TUI/Themes/theme_retro_terminal.py b/TUI/Themes/theme_retro_terminal.py deleted file mode 100644 index 1de2669..0000000 --- a/TUI/Themes/theme_retro_terminal.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -from textual.color import Color - - -def get_retro_terminal_theme(): - from textual.theme import Theme - - return Theme( - name="retro-terminal", - background=Color.parse("#000000"), - primary=Color.parse("#00ff00"), - secondary=Color.parse("#00aa00"), - success=Color.parse("#00ff00"), - warning=Color.parse("#ffff00"), - error=Color.parse("#ff0000"), - surface=Color.parse("#071802"), - ) - - -RETRO_TERMINAL_CSS = """ -/* Retro terminal CRT effect */ -Screen { - align: center middle; - background: $background; - color: $text; -} - -/* Blocky, pixelated widgets */ -.widget { - border: tall $primary; - background: $surface; - width: 80%; -} - -/* Monospaced font */ -* { - font-family: "Courier New", monospace; -} -""" diff --git a/TUI/Widgets/OTP_generate.py b/TUI/Widgets/OTP_generate.py index 4dacdf6..e34ddfa 100644 --- a/TUI/Widgets/OTP_generate.py +++ b/TUI/Widgets/OTP_generate.py @@ -43,7 +43,6 @@ class OTPGenerator(Widget): # Reactive properties to track form completion requestor_filled = reactive(False) reasoning_filled = reactive(False) - duration_selected = reactive(True) # Default is selected otp_generated = reactive(False) class OTPInfo(Message): diff --git a/TUI/Widgets/agentmoveoperations.py b/TUI/Widgets/agentmoveoperations.py index 15cfc96..412a200 100644 --- a/TUI/Widgets/agentmoveoperations.py +++ b/TUI/Widgets/agentmoveoperations.py @@ -29,14 +29,71 @@ from textual.widget import Widget from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea from models.agent import Agent +from services.API import AirlockAPIWrapper from TUI.Screens.executionhistoryscreen import ExecutionHistoryScreen from TUI.Screens.otpworkflowscreen import OTPWorkflowScreen from TUI.Screens.policyselectorscreen import PolicySelectorScreen from TUI.Widgets.OTP_generate import OTPGenerator +from utils.configmanager import get_system_json logger = logging.getLogger(__name__) +def moveAgentToRelatedPolicy( + api: AirlockAPIWrapper, + agent: Agent, + mode: str = "audit", +): + """ + Moves an agent between audit and enforcement policies based on the mode. + + Args: + api: AirlockAPIWrapper instance. + agent: Agent object. + mode: 'audit' to move to audit, 'enforcement' to move to enforcement. + """ + policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}") + + if mode == "audit": + if agent.groupid in policy_relationship_map: + target_policy = policy_relationship_map[agent.groupid] + elif agent.groupid in policy_relationship_map.values(): + logger.debug( + f"Agent {agent.hostname} is already in an audit group. No action needed." + ) + print( + f"Agent {agent.hostname} is already in an audit group. No action needed." + ) + return + else: + logger.warning( + f"Error: No corresponding audit policy found for groupid: {agent.groupid}." + ) + return + + elif mode == "enforcement": + inverse_map = {v: k for k, v in policy_relationship_map.items()} + if agent.groupid in inverse_map: + target_policy = inverse_map[agent.groupid] + elif agent.groupid in inverse_map.values(): + logger.info( + f"Agent {agent.hostname} is already in an enforcement group. No action needed." + ) + return + else: + logger.warning( + f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}." + ) + return + + else: + logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.") + return + + result = api.agent_move(agent.agentid, target_policy) + return result + + class AgentMoveOperations(Widget): """ A Textual widget for managing bulk agent operations and policy migrations. @@ -216,11 +273,11 @@ class AgentMoveOperations(Widget): 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)") @@ -255,9 +312,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) @@ -292,39 +349,39 @@ class AgentMoveOperations(Widget): yield operations_label # Operation buttons - export_csv_btn = Button("📄 Export CSV", id="export_csv_btn") + export_csv_btn = Button("📄 Export CSV", id="export_csv_btn") export_csv_btn.styles.width = "100%" export_csv_btn.styles.margin = (0, 0, 1, 0) yield export_csv_btn local_approval_btn = Button( - "✔️ Local Approval Mode", id="local_approval_btn" + "✔❗ Local Approval Mode", id="local_approval_btn" ) local_approval_btn.styles.width = "100%" local_approval_btn.styles.margin = (0, 0, 1, 0) yield local_approval_btn - otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn") + otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn") otp_gen_btn.styles.width = "100%" otp_gen_btn.styles.margin = (0, 0, 1, 0) yield otp_gen_btn toggle_enforcement_btn = Button( - "🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn" + "🔄 Toggle Audit/Enforcement", id="toggle_enforcement_btn" ) toggle_enforcement_btn.styles.width = "100%" toggle_enforcement_btn.styles.margin = (0, 0, 1, 0) yield toggle_enforcement_btn other_policy_btn = Button( - "🔀 Move to Other Policy", id="other_policy_btn" + "🔀 Move to Other Policy", id="other_policy_btn" ) other_policy_btn.styles.width = "100%" other_policy_btn.styles.margin = (0, 0, 1, 0) yield other_policy_btn exec_history_btn = Button( - "📊 View Execution History", id="exec_history_btn" + "📊 View Execution History", id="exec_history_btn" ) exec_history_btn.styles.width = "100%" exec_history_btn.styles.margin = (0, 0, 1, 0) @@ -391,17 +448,17 @@ class AgentMoveOperations(Widget): 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() @@ -452,7 +509,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 @@ -463,8 +520,6 @@ class AgentMoveOperations(Widget): try: import time - from services.agenthandler import moveAgentToRelatedPolicy - # Generate batch ID batch = int(time.time()) duration = 360 # Default 6 hours, could make this configurable @@ -487,7 +542,7 @@ 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 @@ -537,7 +592,7 @@ class AgentMoveOperations(Widget): successful.append(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 @@ -581,7 +636,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 @@ -590,9 +645,6 @@ class AgentMoveOperations(Widget): unsuccessful = [] try: - from services.agenthandler import moveAgentToRelatedPolicy - from utils.configmanager import get_system_json - policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}") for agent in self.agents: @@ -617,7 +669,7 @@ 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 @@ -687,7 +739,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") @@ -722,7 +774,7 @@ class AgentMoveOperations(Widget): ) except Exception as e: logger.error(f"Failed to open execution history viewer: {e}") - status_label.update(f"❌ Error: {str(e)}") + status_label.update(f"❌ Error: {str(e)}") self.app.notify( f"Failed to open execution history: {str(e)}", severity="error" ) diff --git a/TUI/Widgets/prepPolicy.py b/TUI/Widgets/prepPolicy.py deleted file mode 100644 index 87e2a41..0000000 --- a/TUI/Widgets/prepPolicy.py +++ /dev/null @@ -1,870 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import logging -import os -import os.path -import re -from typing import List - -import dotenv -import pandas as pd - -from models.execution import ExecutionHistoryRecord -from models.policy import Allowlist, Policy -from services.API import AirlockAPIWrapper -from utils.configmanager import get_system_list, get_system_value, load_env -from utils.selector import Selector -from utils.utils import ( - areYouSure, - clear_screen, - colorText, - formatHTML, - get_sanitized_input, - locked, - open_directory, - print_x_wide, - regulator, -) - -logger = logging.getLogger(__name__) - -dotenv.load_dotenv() - - -def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]: - - policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()] - logger.debug("Prompting for Policies") - print(colorText("Please select policy/policies", "white")) - selected = Selector.select_objects(policies, allow_multiple, prompt_each=True) - - if selected is None: - return [] - - # Normalize to always return a list - logger.debug("Returning {selected.dict}") - return selected if isinstance(selected, list) else [selected] - - -def selectAllowlists( - api: AirlockAPIWrapper, policy=all, allow_multiple=True -) -> List[Allowlist]: - if policy == "all": - allowlists = [ - Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows() - ] - else: - allowlists = [ - Allowlist(**row.to_dict()) - for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows() - ] - logger.debug("Prompting for Allowlist(s)") - print(colorText("Please select allowlist(s)", "white")) - selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True) - - if selected is None: - return [] - - # Normalize to always return a list - logger.debug(f"Returning {selected}") - return selected if isinstance(selected, list) else [selected] - - -def sortHashes( - api: AirlockAPIWrapper, selected_policies: List[Policy], type=[1, 2, 6, 7] -): - working_dir = load_env("WORKING_DIR") - history_days = Selector.select_value( - prompt="Enter how many days of history to pull (1-365): ", - value_type=int, - valid_range=(1, 365), - ) - - logger.debug(f"{history_days} day selected for history") - - if history_days is None: - logging.warning("No history range selected. Aborting.") - return - - policy_executions = ExecutionHistoryRecord.from_policies( - api, selected_policies, type_=type, history_days=history_days - ) - - logger.debug(f"Executions contains {policy_executions}") - - enriched_executions = ExecutionHistoryRecord.enrich_with_hashes( - api, policy_executions - ) - categorized_executions = ( - ExecutionHistoryRecord.categorize_executions_by_hash_decision( - enriched_executions - ) - ) - approved, unapproved, needs_review, unknown = ( - ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions) - ) - - categories = { - "needs_review": needs_review, - "approved": approved, - "unapproved": unapproved, - "leftover": unknown, - } - - for label, records in categories.items(): - if not records: - continue # Skip empty or falsy categories - - csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv" - html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html" - - # Convert ExecutionHistoryRecord objects to dictionaries - df = pd.DataFrame([r.__dict__ for r in records]) - - # Optional: flatten hash_obj if needed - if not df.empty and "hash_obj" in df.columns: - hash_df = df["hash_obj"].apply(lambda h: h.to_dict() if h else {}) - df = pd.concat([df.drop(columns=["hash_obj"]), hash_df], axis=1) - - # Save to CSV - df.to_csv(csv_path, index=False) - logger.info(f"Saved {label} executions to {csv_path}") - - # Generate HTML - formatHTML(df, html_path) - logger.info(f"Generated HTML report at {html_path}") - - -def buildPathsandPublishers(selected_policies: List[Policy], split): - working_dir = load_env("WORKING_DIR") - df1 = pd.DataFrame() - df2 = pd.DataFrame() - all_approved_hashes = pd.DataFrame() - path1 = ( - f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv" - ) - path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv" - path_exclusion_constant = get_system_value("PATH_EXCLUSION_CONST", cast_type=int) - - if os.path.exists(path1): - df1 = pd.read_csv(path1) - else: - logger.warning(f"File not found: {path1}") - - if os.path.exists(path2): - df2 = pd.read_csv(path2) - else: - logger.warning(f"File not found: {path2}") - - if df1.empty and df2.empty: - logger.warning("Both DataFrames are empty. Skipping sort.") - all_approved_hashes = pd.DataFrame() - logger.debug(all_approved_hashes.head) - else: - all_approved_hashes = pd.concat([df1, df2], ignore_index=True) - if "filename" in all_approved_hashes.columns: - all_approved_hashes = all_approved_hashes.sort_values(by="filename") - else: - logger.warning( - "Warning: 'filename' column not found in concatenated DataFrame." - ) - - if not all_approved_hashes.empty and path_exclusion_constant: - - primary_path_exclusions = calculatePath( - all_approved_hashes, - path_exclusion_constant, - split, - ) - remaining_hashes = all_approved_hashes[ - ~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"]) - ] - secondary_path_exclusions = calculatePath( - remaining_hashes, (path_exclusion_constant - 1), split - ) - remaining_hashes = remaining_hashes[ - ~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"]) - ] - dataframes = { - "all_approved_hashes": all_approved_hashes, - "primary_Paths": primary_path_exclusions, - "secondary_Paths": secondary_path_exclusions, - "hashes_not_approvable_by_path": remaining_hashes, - } - logger.debug("Preparing to sort dataframes") - for name, df in dataframes.items(): - logger.debug(f" DataFrame headers: {list(df.columns)}") - if "hashes" in name: - df.sort_values(by="filename", inplace=True) - else: - df.sort_values(by="longestcfp", inplace=True) - - df.to_csv( - f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv", - index=False, - ) - formatHTML( - df, - f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html", - ) - - if not all_approved_hashes.empty: - # Drop all not signed, only keep unique values - publist = all_approved_hashes[ - all_approved_hashes["publisher"] != "Not Signed" - ].drop_duplicates(subset=["publisher"]) - # Remove Bad publisher if somehow they made it this far - pattern = regulator(get_system_list("BAD_PUBLISHERS")) - publist = publist[~publist["publisher"].str.contains(pattern, na=False)] - publist = publist[["publisher"]] - publist.sort_values(by="publisher", inplace=True) - publist.to_csv( - f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv", - index=False, - ) - else: - logger.debug("Approved Hashes list appears empty") - - -def buildPreflights(selected_policies: List[Policy]): - working_dir = load_env("WORKING_DIR") - - df1 = pd.DataFrame() - df2 = pd.DataFrame() - approved_hashes = pd.DataFrame() - approved_publishers = pd.DataFrame() - - hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv" - path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv" - path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv" - publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv" - - # Read in and combine the two path generations - if os.path.exists(path1): - df1 = pd.read_csv(path1) - else: - logger.warning(f"File not found: {path1}") - - if os.path.exists(path2): - df2 = pd.read_csv(path2) - else: - logger.warning(f"File not found: {path2}") - - if df1.empty and df2.empty: - logger.warning("Both DataFrames are empty. Skipping sort.") - approved_paths = pd.DataFrame() - else: - approved_paths = pd.concat([df1, df2], ignore_index=True) - - approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first") - - # We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions. - if os.path.exists(hash): - hashes = pd.read_csv(hash) - approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])] - - approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first") - - else: - logger.warning(f"File not found: {hash}") - - if os.path.exists(publishers): - approved_publishers = pd.read_csv(publishers) - - else: - logger.warning(f"File not found: {publishers}") - - dataframes = { - "approved_paths": approved_paths, - "approved_hashes": approved_hashes, - "approved_publishers": approved_publishers, - } - - for name, df in dataframes.items(): - logger.debug(f" DataFrame headers: {list(df.columns)}") - if name == "approved_paths": - df.sort_values(by="longestcfp", inplace=True) - elif name == "approved_hashes": - df.sort_values(by="filename", inplace=True) - elif name == "approved_publishers": - df.sort_values(by="publisher", inplace=True) - - df.to_csv( - f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv", - index=False, - ) - formatHTML( - df, - f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html", - ) - - -def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): - min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int) - - def clean_split(path): - if not isinstance(path, (str, bytes, os.PathLike)): - return [] - parts = str(os.path.normpath(path)).split(os.sep) - parts = [p for p in parts if p] # Remove empty strings - return parts - - # Diagnostic: log any non-string entries - non_string_entries = df[ - ~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike))) - ] - if not non_string_entries.empty: - print(f"[WARNING] Non-string entries found in column '{col}':") - print(non_string_entries) - - df = df.copy() - split_paths = df[col].apply(clean_split) - - if min_files_for_path is not None: - df = df[ - split_paths.apply(lambda parts: len(parts) >= min_files_for_path) - ].copy() - split_paths = split_paths[df.index] - - df["group_key"] = split_paths.apply( - lambda parts: os.sep.join(parts[:path_exclusion_constant]) - ) - grouped = df.groupby("group_key") - new_rows = [] - - for _, group_df in grouped: - paths = group_df[col].tolist() - split_parts = [clean_split(p) for p in paths] - - def longest_common_prefix(paths): - if not paths: - return [] - prefix = paths[0] - for path in paths[1:]: - prefix = [a for a, b in zip(prefix, path) if a == b] - if not prefix: - break - return prefix - - common_prefix = longest_common_prefix(split_parts) - prefix_str = os.sep.join(common_prefix) - - for i, parts in enumerate(split_parts): - filename = parts[-1] - middle = ( - os.sep.join(parts[len(common_prefix) : -1]) - if len(parts) > len(common_prefix) + 1 - else "" - ) - row = group_df.iloc[i].copy() - row["longestcfp"] = prefix_str - row["middle"] = middle - row["filename_only"] = filename - row["file_extension"] = os.path.splitext(filename)[1].lower() - new_rows.append(row) - - return pd.DataFrame(new_rows).drop(columns=["group_key"]) - - -def calculatePath(approved_hashes, path_exclusion_constant, split): - if split: - dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")] - else: - dfs_by_policy = [approved_hashes] - - badpathparts = get_system_list("BAD_PATH_PARTS") - min_files_for_path = get_system_value("MIN_FILES_FOR_PATH", cast_type=int) - - processed_dfs = [] - - for df in dfs_by_policy: - haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename") - haslcp = haslcp.drop_duplicates() - - forbidden = regulator(badpathparts, True) - forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False) - - logger.debug("Removing forbidden filepaths for path exceptions") - print(colorText("Removing forbidden filepaths for path exceptions", "green")) - lcp_not_forbidden = haslcp[~forbidden_lcfp].copy() - - lcp_not_forbidden_review = lcp_not_forbidden[ - [ - "policyname", - "longestcfp", - "middle", - "filename_only", - "file_extension", - "sha256", - ] - ] - - unique_sha_counts = ( - lcp_not_forbidden_review.groupby("longestcfp")["sha256"] - .nunique() - .reset_index() - ) - unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"] - - lcp_not_forbidden_review = lcp_not_forbidden_review.merge( - unique_sha_counts, on="longestcfp", how="left" - ) - lcp_not_forbidden_review = lcp_not_forbidden_review[ - lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path - ] - processed_dfs.append(lcp_not_forbidden_review) - - pathExclusions = pd.concat(processed_dfs, ignore_index=True) - - return pathExclusions - - -def testChange(selected_policies, destination_policy, destination_allowlist): - working_dir = load_env("WORKING_DIR") - - logger.info("These path exclusions would be added to:") - logger.info(destination_policy) - - pathexclusions = pd.read_csv( - f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv" - ) - hashes = pd.read_csv( - f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv" - ) - - unique_combinations = pathexclusions[ - ["longestcfp", "file_extension"] - ].drop_duplicates() - - drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\") - processed_paths = [ - (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}" - for path, ext in unique_combinations.itertuples(index=False, name=None) - ] - - for path in processed_paths: - logger.info(path) - - print(colorText("These publishers would added", "yellow")) - processed_publishers = [] - if os.path.exists( - f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv" - ): - publishers = pd.read_csv( - f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv" - ) - if publishers.empty: - print(colorText("The publishers list is empty.", "red")) - else: - processed_publishers = ( - publishers[publishers["publisher"] != "Not Signed"]["publisher"] - .drop_duplicates() - .tolist() - ) - for publisher in processed_publishers: - print(publisher) - - print(colorText("These hashes would be added to:", "yellow")) - print(destination_allowlist) - - processed_hashes = hashes["sha256"].unique().tolist() - print_x_wide(processed_hashes, 3) - - return processed_paths, processed_hashes, processed_publishers - - -def menu_policy_enforce( - api: AirlockAPIWrapper, -): # TODO Need to clean up 6 and 7 into functions - selected_policies = [] - destination_policy = [] - destination_allowlist = [] - processed_paths = [] - processed_hashes = [] - processed_publishers = [] - working_dir = load_env("WORKING_DIR") - - while True: - printEnforceChecklist( - selected_policies, destination_policy, destination_allowlist - ) - choice = get_sanitized_input("\nEnter your choice: ") - - if choice == "1": - clear_screen() - selected_policies = selectPolicies(api, True) - - elif choice == "2": - clear_screen() - print( - colorText( - "Please choose destination_name Policy for Path Exclusions", "white" - ) - ) - - destination_policy = selectPolicies(api, False) - - print(colorText("Please choose Allowlist for Hashes", "white")) - - destination_allowlist = selectAllowlists(api, destination_policy, False) - - elif choice == "3": - clear_screen() - sortHashes( - api, - selected_policies, - type=[1, 2, 6, 7], - ) - - elif choice == "4": - clear_screen() - if os.path.exists( - f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv" - ): - buildPathsandPublishers(selected_policies, False) - else: - print( - "File not found. Please make sure it's saved correctly and try again." - ) - - elif choice == "5": - clear_screen() - if os.path.exists( - f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv" - ) and os.path.exists( - f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv" - ): - buildPreflights(selected_policies) - else: - print( - "File not found. Please make sure it's saved correctly and try again." - ) - - elif choice == "6": - clear_screen() - if ( - os.path.exists( - f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv" - ) - and os.path.exists( - f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv" - ) - and destination_policy - and destination_allowlist - ): - processed_paths, processed_hashes, processed_publishers = testChange( - selected_policies, destination_policy, destination_allowlist - ) - else: - # Log which condition(s) failed - missing_items = [] - if not os.path.exists( - f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv" - ): - missing_items.append("approved_paths.csv not found") - if not os.path.exists( - f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv" - ): - missing_items.append("approved_hashes.csv not found") - if not destination_policy: - missing_items.append("destination_policy is empty or None") - if not destination_allowlist: - missing_items.append("destination_allowlist is empty or None") - - logger.error("Preflight check failed due to the following:") - for item in missing_items: - logger.error(f" - {item}") - - elif choice == "7": - clear_screen() - areYouSure() - confirmation = get_sanitized_input("Type 'I AGREE' to continue: ") - if ( - processed_paths - and processed_hashes - and processed_publishers - and destination_policy - and destination_allowlist - and confirmation.strip() == "I AGREE" - ): - print(colorText("Proceeding with the code...", "yellow")) - api.hash_add_to_allowlist( - destination_allowlist[0].applicationid, processed_hashes - ) - api.policy_add_path_exclusions( - destination_policy[0].groupid, processed_paths - ) - if processed_publishers: - api.policy_add_publishers( - destination_policy[0].groupid, processed_publishers - ) - - locked() - - else: - logger.error("Confirmation block failed. Reasons:") - if not processed_publishers or processed_hashes or processed_paths: - logger.error(" - Test not performed.") - if not destination_policy: - logger.error(" - `destination_policy` is missing or invalid.") - if not destination_allowlist: - logger.error(" - `destination_allowlist` is missing or invalid.") - if confirmation.strip() != "I AGREE": - logger.error( - " - User did not confirm with 'I AGREE'. Received: '%s'", - confirmation.strip(), - ) - - elif choice.upper() == "F": - open_directory(working_dir) - elif choice.upper() == "B": - break - - else: - print(colorText("Invalid choice. Please try again.", "red")) - - -def section_header(title): - print( - colorText( - "\n --------------------------------------------------------------------", - "cyan", - ) - ) - print(colorText(f" ------------- {title} -------------", "cyan")) - print( - colorText( - " --------------------------------------------------------------------", - "cyan", - ) - ) - - -def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist): - working_dir = load_env("WORKING_DIR") - section_header("Prepare to Enforce Policy") - print( - colorText( - "\nSequentially follow these steps to prepare a policy for enforcement:", - "white", - ) - ) - - # Step 1: Originating Policies - print( - colorText( - "\n1. Choose which policy or policies to gather execution info from", "cyan" - ) - ) - if not selected_policies: - print(colorText(" [❌] No policies have been chosen", "red")) - else: - print(colorText("The following policies have been chosen:", "green")) - for policy in selected_policies: - print(colorText(f" [✅] {policy.name}", "green")) - - # Step 2: Destination Policy and Allowlist - print( - colorText("2. Choose the destination policy and associated allowlist", "cyan") - ) - if destination_policy: - print( - colorText( - f" [✅] {destination_policy[0].name} has been selected as the destination policy", - "green", - ) - ) - else: - print(colorText(" [❌] No destination policy has been chosen", "red")) - - if destination_allowlist: - print( - colorText( - f" [✅] {destination_allowlist[0].name} has been selected as allowlist", - "green", - ) - ) - else: - print(colorText(" [❌] No allowlist has been chosen", "red")) - - # Step 3: Data Preparation - print( - colorText( - f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review", - "cyan", - ) - ) - if selected_policies: - policy_id = selected_policies[0].name - review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv" - print( - colorText( - ( - " [✅] Data has been fetched" - if os.path.exists(review_path) - else " [❌] Data has not been fetched" - ), - "green" if os.path.exists(review_path) else "red", - ) - ) - else: - print( - colorText( - " [❌] No policies selected, cannot check data fetch status", "red" - ) - ) - - # Step 4: Manual Review - print(colorText("4. Manually review the files:", "cyan")) - print( - colorText( - " Remove the rows containing hashes you do not approve of", "cyan" - ) - ) - print( - colorText( - f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.", - "cyan", - ) - ) - print( - colorText( - " This will start the process to generate possible filepath approvals", - "cyan", - ) - ) - - if selected_policies: - policy_id = selected_policies[0].name - approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv" - second_review_path = ( - f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv" - ) - print( - colorText( - ( - " [✅] Reviewed hashes have been loaded" - if os.path.exists(approved_path) - else " [❌] Reviewed hashes have not been loaded" - ), - "green" if os.path.exists(approved_path) else "red", - ) - ) - print( - colorText( - ( - " [✅] Path review list created" - if os.path.exists(second_review_path) - else " [❌] Path review list has not been created" - ), - "green" if os.path.exists(second_review_path) else "red", - ) - ) - else: - print( - colorText( - " [❌] No policies selected, cannot check reviewed hashes or path list", - "red", - ) - ) - - # Step 5: Path Review - print( - colorText( - f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\", - "cyan", - ) - ) - print( - colorText( - " Remove the rows containing path exclusions or publishers you do not approve of.", - "cyan", - ) - ) - print( - colorText( - f" When complete, save the files to {working_dir}\\data\\Approved", - "cyan", - ) - ) - print( - colorText(" Choose this option when done to build your preflights", "cyan") - ) - - if selected_policies: - policy_id = selected_policies[0].name - reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv" - preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv" - preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv" - print( - colorText( - ( - " [✅] Reviewed path list detected" - if os.path.exists(reviewed_path) - else " [❌] Path review list has not been detected" - ), - "green" if os.path.exists(reviewed_path) else "red", - ) - ) - preflight_ready = os.path.exists(preflight_paths) and os.path.exists( - preflight_hashes - ) - print( - colorText( - ( - " [✅] Preflight Path Exclusion List has been generated" - if preflight_ready - else " [❌] Preflight Path Exclusion List has not been generated" - ), - "green" if preflight_ready else "red", - ) - ) - else: - print( - colorText( - " [❌] No policies selected, cannot check preflight status", "red" - ) - ) - - # Final Steps - print( - colorText( - "6. Test ------------------------------------------------------", "cyan" - ) - ) - print( - colorText( - " Prints to console the changes that would be made, must be done to proceed. ", - "cyan", - ) - ) - - print( - colorText( - "7. Liftoff ------------------------------------------------------", "cyan" - ) - ) - print( - colorText( - " Apply path exclusions and approved publishers to selected policy", - "cyan", - ) - ) - print(colorText(" Apply approved hashes to allowlist", "cyan")) - - # Utility Options - print(colorText("F. Open Working Directory", "cyan")) - print(colorText("B. Back", "cyan")) diff --git a/TUI/Widgets/settingswidget.py b/TUI/Widgets/settingswidget.py new file mode 100644 index 0000000..1c97ee8 --- /dev/null +++ b/TUI/Widgets/settingswidget.py @@ -0,0 +1,567 @@ +# 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 . + +""" +Settings widget combining theme selection and update checking. +""" + +import logging +import webbrowser + +from textual.containers import Horizontal, Vertical, VerticalScroll +from textual.message import Message +from textual.widget import Widget +from textual.widgets import Button, Rule, Static + +from utils.versionchecker import ( + RELEASES_PAGE_URL, + UpdateCheckResult, + check_for_updates, + get_current_version, + get_version_checker, +) + +logger = logging.getLogger(__name__) + + +class SettingsWidget(Widget): + """Widget for application settings including themes and updates.""" + + DEFAULT_CSS = """ + SettingsWidget { + height: 1fr; + } + + /* Update section buttons - add margin between them */ + #update_buttons Button { + margin-right: 1; + } + + /* Theme buttons - consistent width within columns, slightly smaller */ + .theme_btn { + width: 100%; + margin-bottom: 1; + } + + /* Column headers */ + .theme_column_header { + text-align: center; + text-style: bold; + margin-bottom: 1; + } + + /* Section titles */ + .settings_section_title { + text-style: bold; + margin-bottom: 1; + } + + /* Theme columns - reduce overall width */ + #theme_columns { + width: 80%; + } + + /* Theme columns spacing */ + #dark_themes_col1, #dark_themes_col2 { + margin-right: 1; + } + + #light_themes_col { + margin-left: 1; + } + """ + + class ThemeSelected(Message): + """Message posted when a theme is selected.""" + + def __init__(self, theme_name: str): + super().__init__() + self.theme_name = theme_name + + # Dark themes - Column 1 + DARK_THEMES_COL1 = [ + ("Textual Dark", "textual-dark"), + ("Nord", "nord"), + ("Gruvbox", "gruvbox"), + ("Dracula", "dracula"), + ] + + # Dark themes - Column 2 + DARK_THEMES_COL2 = [ + ("Catppuccin Mocha", "catppuccin-mocha"), + ("Tokyo Night", "tokyo-night"), + ("Monokai", "monokai"), + ] + + # Light themes (third column) + LIGHT_THEMES = [ + ("Textual Light", "textual-light"), + ("Flexoki", "flexoki"), + ("Catppuccin Latte", "catppuccin-latte"), + ("Solarized Light", "solarized-light"), + ] + + # Combined for backward compatibility + DARK_THEMES = DARK_THEMES_COL1 + DARK_THEMES_COL2 + AVAILABLE_THEMES = DARK_THEMES + LIGHT_THEMES + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._update_result: UpdateCheckResult | None = None + self._checking = False + + def compose(self): + # Wrap everything in a scrollable container with auto height children + with VerticalScroll(id="settings_scroll"): + # Version & Updates Section + with Vertical(id="updates_section") as updates: + updates.styles.height = "auto" + yield Static( + "📦 Version & Updates", + id="updates_title", + classes="settings_section_title", + ) + + version_text = f"Current Version: v{get_current_version()}" + yield Static(version_text, id="current_version") + + with Horizontal(id="update_buttons") as btn_row: + btn_row.styles.height = "auto" + yield Button("🔍 Check for Updates", id="check_updates_btn") + yield Button("📥 View Releases", id="view_releases_btn") + + yield Static("", id="update_status") + + yield Rule() + + # Theme Section - Three columns: Dark 1, Dark 2, Light + with Vertical(id="themes_section") as themes: + themes.styles.height = "auto" + yield Static( + "🎨 Theme Options", + id="theme_title", + classes="settings_section_title", + ) + + with Horizontal(id="theme_columns") as cols: + cols.styles.height = "auto" + + # Dark themes section (2 columns under one header) + with Vertical(id="dark_themes_section") as dark_section: + dark_section.styles.width = "2fr" + dark_section.styles.height = "auto" + yield Static( + "🌙 Dark Themes", + classes="theme_column_header", + id="dark_header", + ) + + with Horizontal(id="dark_columns") as dark_cols: + dark_cols.styles.height = "auto" + + # Dark themes column 1 + with Vertical(id="dark_themes_col1") as dark_col1: + dark_col1.styles.width = "1fr" + dark_col1.styles.height = "auto" + for label, btn_id in self.DARK_THEMES_COL1: + yield Button( + label, + id=f"set_theme_{btn_id}", + classes="theme_btn", + ) + + # Dark themes column 2 + with Vertical(id="dark_themes_col2") as dark_col2: + dark_col2.styles.width = "1fr" + dark_col2.styles.height = "auto" + for label, btn_id in self.DARK_THEMES_COL2: + yield Button( + label, + id=f"set_theme_{btn_id}", + classes="theme_btn", + ) + + # Light themes column + with Vertical(id="light_themes_col") as light_col: + light_col.styles.width = "1fr" + light_col.styles.height = "auto" + yield Static("☀️ Light Themes", classes="theme_column_header") + for label, btn_id in self.LIGHT_THEMES: + yield Button( + label, id=f"set_theme_{btn_id}", classes="theme_btn" + ) + + def on_mount(self) -> None: + """Check for cached update result on mount.""" + checker = get_version_checker() + cached_result = checker.get_last_result() + if cached_result and cached_result.update_available: + self._update_result = cached_result + self._show_update_available(cached_result) + + def on_button_pressed(self, event: Button.Pressed) -> None: + button_id = event.button.id + + if button_id == "check_updates_btn": + self._check_for_updates() + event.stop() + elif button_id == "view_releases_btn": + self._open_releases_page() + event.stop() + elif button_id == "download_update_btn": + self._download_update() + event.stop() + elif button_id == "dismiss_update_btn": + self._dismiss_update() + event.stop() + elif button_id and button_id.startswith("set_theme_"): + theme_name = button_id.replace("set_theme_", "") + self.post_message(self.ThemeSelected(theme_name)) + event.stop() + + def _check_for_updates(self) -> None: + """Check for updates and update UI.""" + if self._checking: + return + + self._checking = True + status = self.query_one("#update_status", Static) + check_btn = self.query_one("#check_updates_btn", Button) + + # Show checking status + check_btn.disabled = True + check_btn.label = "⏳ Checking..." + status.update("🔄 Checking for updates...") + + # Run check in worker to avoid blocking UI + self.run_worker(self._do_update_check, exclusive=True) + + async def _do_update_check(self) -> None: + """Worker to perform update check.""" + try: + result = check_for_updates() + self._update_result = result + + # Since we're in an async worker (not a thread), we can call directly + self._update_check_complete(result) + except Exception as e: + logger.error(f"Update check failed: {e}") + self._update_check_failed(str(e)) + finally: + self._checking = False + + def _update_check_complete(self, result: UpdateCheckResult) -> None: + """Handle completed update check.""" + check_btn = self.query_one("#check_updates_btn", Button) + check_btn.disabled = False + check_btn.label = "🔍 Check for Updates" + + if result.error: + self._update_check_failed(result.error) + return + + if result.update_available: + self._show_update_available(result) + self.app.notify( + f"🆕 Update available: {result.latest_version}", + title="Update Available", + severity="information", + timeout=8, + ) + else: + status = self.query_one("#update_status", Static) + status.update(f"✅ Loxide is up to date (v{result.current_version})") + self.app.notify( + "✅ Loxide is up to date!", + severity="information", + timeout=5, + ) + + def _update_check_failed(self, error: str) -> None: + """Handle failed update check.""" + check_btn = self.query_one("#check_updates_btn", Button) + check_btn.disabled = False + check_btn.label = "🔍 Check for Updates" + + status = self.query_one("#update_status", Static) + status.update(f"⚠️ Could not check for updates: {error}") + + def _show_update_available(self, result: UpdateCheckResult) -> None: + """Show update available UI with release notes.""" + status = self.query_one("#update_status", Static) + + msg = f"🆕 New version available: {result.latest_version}\n" + msg += f" Current: v{result.current_version}" + + if result.release_info and result.release_info.body: + # Show release notes (truncate if very long) + notes = result.release_info.body.strip() + # Limit to ~500 chars to avoid overwhelming the UI + if len(notes) > 500: + notes = notes[:500] + "\n..." + msg += f"\n\n📋 Release Notes:\n{notes}" + + status.update(msg) + + # Add download/dismiss buttons if not already there + try: + self.query_one("#download_update_btn") + except Exception: + # Buttons don't exist, add them + button_container = self.query_one("#update_buttons", Horizontal) + download_btn = Button( + "📥 Download Update", id="download_update_btn", variant="success" + ) + dismiss_btn = Button( + "✖ Dismiss", id="dismiss_update_btn", variant="default" + ) + button_container.mount(download_btn) + button_container.mount(dismiss_btn) + + def _open_releases_page(self) -> None: + """Open the releases page in browser.""" + try: + webbrowser.open(RELEASES_PAGE_URL) + self.app.notify("📂 Opened releases page in browser", timeout=3) + except Exception as e: + logger.error(f"Could not open browser: {e}") + self.app.notify(f"⚠️ Could not open browser: {e}", severity="warning") + + def _download_update(self) -> None: + """Download the update exe file.""" + import os + from pathlib import Path + + if not self._update_result or not self._update_result.release_info: + self.app.notify("⚠️ No update information available", severity="warning") + return + + download_url = self._update_result.release_info.download_url + if not download_url: + # Fall back to opening the release page + url = self._update_result.release_info.html_url + try: + webbrowser.open(url) + self.app.notify( + "📥 Opened download page in browser (no direct download available)", + timeout=5, + ) + except Exception as e: + logger.error(f"Could not open browser: {e}") + self.app.notify(f"⚠️ Could not open browser: {e}", severity="warning") + return + + # Determine destination path + if os.name == "nt": # Windows + downloads_dir = Path.home() / "Downloads" + else: + downloads_dir = Path.home() / "Downloads" + if not downloads_dir.exists(): + downloads_dir = Path.home() + + # Extract filename from URL + filename = download_url.split("/")[-1] + if not filename.endswith(".exe"): + filename = f"Loxide_{self._update_result.latest_version}.exe" + + dest_path = downloads_dir / filename + + # Disable the button while downloading + try: + btn = self.query_one("#download_update_btn", Button) + btn.disabled = True + btn.label = "⏳ Downloading..." + except Exception: + pass + + self.app.notify( + f"📥 Downloading to:\n{dest_path}", title="Download Starting", timeout=5 + ) + + # Small delay so user sees the "downloading to" toast before download completes + self.set_timer( + 0.5, + lambda: self.run_worker( + self._do_download(download_url, dest_path), exclusive=True + ), + ) + + async def _do_download(self, download_url: str, dest_path) -> None: + """Worker to download the update file.""" + try: + # Download the file + import requests + + response = requests.get(download_url, stream=True, timeout=60) + response.raise_for_status() + + with open(dest_path, "wb") as f: + for chunk in response.iter_content(chunk_size=8192): + if chunk: + f.write(chunk) + + # Success + self.app.notify( + f"✅ Downloaded to:\n{dest_path}", + title="Download Complete", + severity="information", + timeout=10, + ) + logger.info(f"Update downloaded to {dest_path}") + + # Re-enable button + try: + btn = self.query_one("#download_update_btn", Button) + btn.disabled = False + btn.label = "📥 Download Again" + except Exception: + pass + + except Exception as e: + logger.error(f"Download failed: {e}") + self.app.notify(f"❌ Download failed: {e}", severity="error", timeout=10) + + # Re-enable button + try: + btn = self.query_one("#download_update_btn", Button) + btn.disabled = False + btn.label = "📥 Download Update" + except Exception: + pass + + def _dismiss_update(self) -> None: + """Dismiss the current update notification.""" + if self._update_result and self._update_result.latest_version: + checker = get_version_checker() + checker.dismiss_update(self._update_result.latest_version) + + # Remove the extra buttons + try: + self.query_one("#download_update_btn").remove() + self.query_one("#dismiss_update_btn").remove() + except Exception: + pass + + status = self.query_one("#update_status", Static) + status.update(f"✓ Dismissed update {self._update_result.latest_version}") + self._update_result = None + + +# Keep ThemeSelector as a standalone for backward compatibility +class ThemeSelector(Widget): + """Widget for selecting and applying Textual themes. + + DEPRECATED: Use SettingsWidget instead for combined settings UI. + """ + + DEFAULT_CSS = """ + ThemeSelector { + height: 1fr; + } + + /* Theme buttons - consistent width within columns */ + .theme_btn { + width: 100%; + margin-bottom: 1; + } + + /* Column headers */ + .theme_column_header { + text-align: center; + text-style: bold; + margin-bottom: 1; + } + + /* Theme columns - reduce overall width */ + #theme_columns { + width: 80%; + } + + /* Theme columns spacing */ + #dark_themes_col1, #dark_themes_col2 { + margin-right: 1; + } + + #light_themes_col { + margin-left: 1; + } + """ + + class ThemeSelected(Message): + """Message posted when a theme is selected.""" + + def __init__(self, theme_name: str): + super().__init__() + self.theme_name = theme_name + + DARK_THEMES_COL1 = SettingsWidget.DARK_THEMES_COL1 + DARK_THEMES_COL2 = SettingsWidget.DARK_THEMES_COL2 + DARK_THEMES = SettingsWidget.DARK_THEMES + LIGHT_THEMES = SettingsWidget.LIGHT_THEMES + AVAILABLE_THEMES = SettingsWidget.AVAILABLE_THEMES + + def compose(self): + with VerticalScroll(id="theme_scroll"): + yield Static("Theme Options", id="theme_title") + + with Horizontal(id="theme_columns") as cols: + cols.styles.height = "auto" + + # Dark themes section (2 columns under one header) + with Vertical(id="dark_themes_section") as dark_section: + dark_section.styles.width = "2fr" + dark_section.styles.height = "auto" + yield Static( + "🌙 Dark Themes", + classes="theme_column_header", + id="dark_header", + ) + + with Horizontal(id="dark_columns") as dark_cols: + dark_cols.styles.height = "auto" + + # Dark themes column 1 + with Vertical(id="dark_themes_col1") as dark_col1: + dark_col1.styles.width = "1fr" + dark_col1.styles.height = "auto" + for label, btn_id in self.DARK_THEMES_COL1: + yield Button( + label, id=f"set_theme_{btn_id}", classes="theme_btn" + ) + + # Dark themes column 2 + with Vertical(id="dark_themes_col2") as dark_col2: + dark_col2.styles.width = "1fr" + dark_col2.styles.height = "auto" + for label, btn_id in self.DARK_THEMES_COL2: + yield Button( + label, id=f"set_theme_{btn_id}", classes="theme_btn" + ) + + # Light themes column + with Vertical(id="light_themes_col") as light_col: + light_col.styles.width = "1fr" + light_col.styles.height = "auto" + yield Static("☀️ Light Themes", classes="theme_column_header") + for label, btn_id in self.LIGHT_THEMES: + yield Button( + label, id=f"set_theme_{btn_id}", classes="theme_btn" + ) + + def on_button_pressed(self, event: Button.Pressed) -> None: + button_id = event.button.id + if button_id and button_id.startswith("set_theme_"): + theme_name = button_id.replace("set_theme_", "") + self.post_message(self.ThemeSelected(theme_name)) diff --git a/TUI/Themes/themeselector.py b/TUI/Widgets/themeselector.py similarity index 100% rename from TUI/Themes/themeselector.py rename to TUI/Widgets/themeselector.py diff --git a/airlock_libs-6.2.0-cp313-cp313-win_amd64.whl b/airlock_libs-6.2.0-cp313-cp313-win_amd64.whl deleted file mode 100644 index 755c56f..0000000 Binary files a/airlock_libs-6.2.0-cp313-cp313-win_amd64.whl and /dev/null differ diff --git a/docs/README.md b/docs/README.md index 8257c1c..a763c86 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,39 +1,47 @@ - # 🛡️ Loxide -Python/Rust/Oxide toolkit for secure, auditable, and automated airlock agent and policy management. Designed for enterprise environments, it supports advanced policy workflows, device tracking, and terminal-based interaction. +Python/Rust/Oxide toolkit for secure, auditable, and automated Airlock agent and policy management. Designed for enterprise environments, it supports advanced policy workflows, device tracking, and terminal-based interaction. + +![Python](https://img.shields.io/badge/python-3.10+-blue.svg) +![Rust](https://img.shields.io/badge/rust-backend-orange.svg) +![License](https://img.shields.io/badge/license-AGPL--3.0-green.svg) +![Textual](https://img.shields.io/badge/textual-6.5.0-purple.svg) --- - ## 🚀 Features + - 🔍 **Fuzzy Device Search** - Quickly locate devices using partial or approximate matches. + Quickly locate devices using partial or approximate matches with wildcard support (`*`, `?`). - 📦 **Batch Move Devices** - Move multiple devices between groups or policies easily. + Move multiple devices between groups or policies easily with multi-select and file import. - 🔄 **Toggle Enforcement/Audit Policies** Seamlessly switch devices between enforcement and audit modes. -- 🕵️‍♂️ **Device History Search** - Track agent executions. +- 🕵️ **Device History Search** + Track agent executions with configurable date ranges (1-365 days) and CSV export. - 🧰 **Prepare Policies for Enforcement** - Validate and stage policies before pushing them to enforcement. + Multi-step wizard to validate and stage policies before pushing to enforcement. - 💤 **Find Quiet Hosts** - Identify devices ready for enforcement. + Identify inactive devices ready for enforcement using Rust-powered analysis. -- 🎛️ **TUI** - Navigate with arrow keys and F-key shortcuts using a custom ANSI-colored terminal UI. +- 🔐 **OTP Management** + Generate, revoke, and monitor one-time passwords for agents. + +- 📋 **Server Log Viewer** + Real-time server activity monitoring with DataTable display. + +- 🎛️ **Modern TUI** + Full keyboard navigation with Textual framework, toast notifications, and color-coded results. --- ## 🧭 Roadmap -- ⚙️ **Rust-based Async API Calls** - Improve performance and concurrency with a Rust-powered backend. - ✅ **Carbon Black-style Local Approval** Enable local user approvals for policy exceptions and enforcement actions. @@ -43,9 +51,105 @@ Python/Rust/Oxide toolkit for secure, auditable, and automated airlock agent and --- -## 🧑‍💻 Requirements +## 📦 Installation -[airlock_libs](https://git.racooncity.org/brotoskyj/-/packages/pypi/airlock-libs/) +### Prerequisites +- Python 3.10+ +- Access to Airlock API +- Network access to private PyPI server + +### Install Dependencies + +```bash +pip install -r requirements.txt +``` + +### Requirements + +``` +# Core TUI +textual==6.5.0 + +# API and data handling +requests==2.32.5 +pandas==2.3.3 +numpy==2.3.4 + +# Database +pymongo==4.15.3 + +# Security +cryptography==46.0.3 +keyring==25.6.0 + +# Utilities +python-dotenv==1.2.1 +tqdm==4.67.1 +urllib3==2.5.0 +pyperclip==1.11.0 + +# Private package +--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/ +airlock_libs==5.2.1 +``` + +See also: [airlock_libs](https://git.racooncity.org/brotoskyj/-/packages/pypi/airlock-libs/) + +--- + +## 📖 Usage + +```bash +python Loxide.py +``` + +### Keyboard Shortcuts + +| Key | Action | +|-----|--------| +| `Tab` | Navigate between elements | +| `Enter` | Select/Confirm | +| `Escape` | Go back / Cancel | +| `←` `→` | Navigate tabs | +| `r` | Refresh (context-dependent) | +| `e` | Export to CSV | +| `q` | Quit application | + +--- + +## 🏗️ Architecture + +``` +Loxide/ +├── Loxide.py # Main entry point +├── API.py # AirlockAPIWrapper +├── setup.py # Initialization & logging +├── configmanager.py # Config management +│ +├── models/ # Data models +├── services/ # Business logic +├── utils/ # Helper functions +│ +└── TUI/ + ├── Screens/ # Workflow screens + │ ├── policyprepworkflowscreen.py + │ ├── quietagentworkflowscreen.py + │ ├── agentmoveoperations.py + │ ├── otprevokescreen.py + │ └── executionhistoryscreen.py + │ + └── Widgets/ # Reusable components + ├── policyselector.py + ├── multiagentselector.py + └── serverlogwidget.py +``` + +--- + +## 👥 Authors + +- **Brandon Wickline** - *Lead Python Developer* +- **James Brotosky** - *Lead Rust Developer* --- @@ -55,5 +159,15 @@ Python/Rust/Oxide toolkit for secure, auditable, and automated airlock agent and You may copy, distribute, and modify the software under the terms of the AGPL-3.0 license. +``` +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. +``` + See the [LICENSE](LICENSE.md) file for full details, or visit -[https://www.gnu.org/license/agpl-3.0.html](https://www.gnu.org/license/agpl-3.0.html) +[https://www.gnu.org/licenses/agpl-3.0.html](https://www.gnu.org/licenses/agpl-3.0.html) + diff --git a/flows/localApproval.py b/flows/localApproval.py deleted file mode 100644 index d6bda28..0000000 --- a/flows/localApproval.py +++ /dev/null @@ -1,252 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import logging -import os -import time -from typing import List, Optional - -from models.agent import Agent -from services.agenthandler import moveAgentToRelatedPolicy, selectAgents -from services.API import AirlockAPIWrapper -from utils.configmanager import get_system_json -from utils.utils import colorText, get_sanitized_input - -logger = logging.getLogger(__name__) - - -class LocalApprovalRequestor: - """Handles creation of local approval requests in Loxide.""" - - def __init__(self, api: AirlockAPIWrapper, username: str = None): - """ - Initialize the local approval requestor. - - Args: - api: AirlockAPIWrapper instance - username: Username creating the approvals (for tracking) - """ - self.api = api - self.policy_map = get_system_json("POLICY_MAP_ENF_AUD", "{}") - self.username = ( - username or os.getenv("USERNAME") or os.getenv("USER") or "unknown" - ) - - def create_local_approval( - self, agent_id: str, duration_minutes: int, batch_id: Optional[int] = None - ) -> bool: - """ - Create a single local approval request. - - Args: - agent_id: Agent ID to create approval for - duration_minutes: Duration of approval in minutes - batch_id: Optional batch identifier (defaults to timestamp) - - Returns: - True if successful, False otherwise - """ - if batch_id is None: - batch_id = int(time.time()) - - purpose = ( - f" Local Approval - {duration_minutes} mins - " - f"batch:{batch_id} Client:{agent_id} User:{self.username}" - ) - - try: - self.api.otp_generate(agent_id, duration_minutes, purpose) - logger.info( - f"Generated local approval for {agent_id}, batch {batch_id}, by {self.username}" - ) - return True - except Exception as e: - logger.error(f"Failed to generate local approval for {agent_id}: {e}") - return False - - def move_agent_to_audit(self, agent: Agent) -> bool: - """ - Move an agent to its corresponding audit policy. - - Args: - agent: Agent object to move - - Returns: - True if successful, False otherwise - """ - try: - moveAgentToRelatedPolicy(self.api, agent, "audit") - logger.info(f"Moved {agent.hostname} to audit policy") - return True - except Exception as e: - logger.error(f"Failed to move {agent.hostname} to audit: {e}") - return False - - def create_local_approval_batch( - self, - agents: List[Agent], - duration_minutes: int, - ) -> tuple[int, int, int]: - """ - Create local approvals for multiple agents and move them to audit. - - Args: - agents: List of Agent objects - duration_minutes: Duration of approval in minutes - db_path: Optional path to database for history tracking - - Returns: - Tuple of (batch_id, success_count, failure_count) - """ - batch_id = int(time.time()) - success_count = 0 - failure_count = 0 - - print(colorText(f"\n Processing batch {batch_id}...", "cyan")) - print(colorText(f" Requested by: {self.username}", "cyan")) - print(colorText(f" Moving {len(agents)} agent(s) to local approval\n", "cyan")) - - for agent in agents: - try: - # Create local approval - approval_success = self.create_local_approval( - agent.agentid, duration_minutes, batch_id - ) - - if not approval_success: - raise Exception("Failed to create local approval") - - # Move to audit policy - move_success = self.move_agent_to_audit(agent) - - if not move_success: - raise Exception("Failed to move to audit policy") - - print(colorText(f" {agent.hostname}", "green")) - success_count += 1 - - except Exception as e: - print(colorText(f" {agent.hostname}: {e}", "red")) - logger.error(f"Error processing agent {agent.hostname}: {e}") - failure_count += 1 - - return batch_id, success_count, failure_count - - def interactive_local_approval(self): - """ - Interactive workflow to create local approvals for selected agents. - - This prompts the user to select a duration and agents, then creates - the local approvals and moves agents to audit policies. - """ - # Duration options in minutes - duration_options = [ - (15, "15 minutes"), - (60, "1 hour"), - (360, "6 hours"), - (1440, "1 day"), - (10080, "1 week"), - ] - - # Display duration options - print(colorText("\n Select Local Approval Duration:", "white")) - print(colorText("=" * 50, "white")) - - for i, (minutes, label) in enumerate(duration_options, start=1): - print(f" {i}. {label} ({minutes} minutes)") - - print(colorText("=" * 50, "white")) - - # Get user selection - try: - choice = int(get_sanitized_input("\nEnter the number of your choice: ")) - - if 1 <= choice <= len(duration_options): - duration_minutes, duration_label = duration_options[choice - 1] - print(colorText(f" Selected: {duration_label}", "green")) - logger.info(f"User selected duration: {duration_minutes} minutes") - else: - print(colorText("❌ Invalid choice.", "red")) - logger.warning("Invalid duration choice") - return - - except ValueError: - print(colorText("❌ Invalid input. Please enter a number.", "red")) - logger.warning("Invalid input for duration selection") - return - - # Select agents - 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")) - logger.warning("No agents selected or error retrieving agents") - return - - # Confirm with user - 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")) - return - - # Process the batch - batch_id, success_count, failure_count = self.create_local_approval_batch( - agents, duration_minutes - ) - - # Display summary - self._display_summary(batch_id, duration_label, success_count, failure_count) - - def _display_summary( - self, batch_id: int, duration_label: str, success_count: int, failure_count: int - ): - """ - Display operation summary. - - Args: - batch_id: Batch identifier - duration_label: Human-readable duration - success_count: Number of successful operations - failure_count: Number of failed operations - """ - print(colorText(f"\n{'=' * 60}", "white")) - print(colorText(" Local Approval Summary", "cyan")) - print(colorText("=" * 60, "white")) - - print(colorText(f" Successfully processed: {success_count}", "green")) - - if failure_count > 0: - print(colorText(f" Failed: {failure_count}", "red")) - - print(colorText(f"\n Batch ID: {batch_id}", "cyan")) - print(colorText(f" Duration: {duration_label}", "cyan")) - - print(colorText("=" * 60, "white")) - print(colorText("\n Next Steps:", "yellow")) - print(colorText(" ✅ Agents have been moved to audit policies", "white")) - print(colorText(" ✅ Local approvals are active", "white")) - print( - colorText( - f" ✅ Agents will return to enforcement after {duration_label}", - "white", - ) - ) - print(colorText("=" * 60 + "\n", "white")) diff --git a/flows/otp.py b/flows/otp.py deleted file mode 100644 index ee07c98..0000000 --- a/flows/otp.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -import logging - -import pandas as pd - -from services.agenthandler import selectAgents -from services.API import AirlockAPIWrapper -from utils.selector import Selector -from utils.utils import get_sanitized_input - -logger = logging.getLogger(__name__) - - -def otp_revoke(api: AirlockAPIWrapper): - - activeagents = api.otp_find_active() - awaitingagents = api.otp_find_awaiting() - - activeagents["status"] = "active" - awaitingagents["status"] = "awaiting" - - combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True) - combined_agents = combined_agents.sort_values(by="otpid", ascending=False) - - # Combine all into one DataFrame - combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True) - combined_agents = combined_agents.sort_values(by="otpid", ascending=False) - - # Optionally, select specific hosts - user_input = ( - get_sanitized_input("\nWould you like to search for a specific device? (y/n): ") - .strip() - .lower() - ) - if user_input == "y": - agentnames = [] - agents = selectAgents(api) - for agent in agents: - agentnames.append(agent.hostname) - - combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)] - - # Present and select rows - selected_rows = Selector.select_dataframe_with_mode( - combined_agents, - columns=["otpid", "hostname", "status", "purpose", "granted"], - header="OTP Sessions", - ) - - for row in selected_rows: - otpid = row["otpid"] - hostname = row["hostname"] - result = api.otp_revoke(otpid) - logger.info(f"{hostname} (otpid: {otpid}):\n{result}") diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py deleted file mode 100644 index e6e6298..0000000 --- a/flows/prepPolicy.py +++ /dev/null @@ -1,870 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -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")) diff --git a/loxide_tight.ico b/loxide_tight.ico deleted file mode 100644 index 631fb96..0000000 Binary files a/loxide_tight.ico and /dev/null differ diff --git a/models/execution.py b/models/execution.py index 439515f..52dfb78 100644 --- a/models/execution.py +++ b/models/execution.py @@ -29,7 +29,7 @@ import pandas as pd import airlock_libs from services.API import AirlockAPIWrapper from utils.configmanager import get_system_list, get_system_value -from utils.utils import colorText, regulator +from utils.utils import regulator logger = logging.getLogger(__name__) @@ -293,12 +293,6 @@ class ExecutionHistoryRecord: logger.debug( f"Staging of Execution history for policy: {policy.name} is complete" ) - print( - colorText( - f"Staging of Execution history for policy: {policy.name} is complete", - "green", - ) - ) return executions diff --git a/requirements.txt b/requirements.txt index 30d5f35..7b901a7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,8 +17,8 @@ keyring==25.6.0 python-dotenv==1.2.1 # Utilities -tqdm==4.67.1 urllib3==2.5.0 +plotext==5.3.2 pyperclip==1.11.0 # Custom/Private packages diff --git a/services/agenthandler.py b/services/agenthandler.py deleted file mode 100644 index 2d50f4d..0000000 --- a/services/agenthandler.py +++ /dev/null @@ -1,388 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -from dataclasses import asdict -from datetime import datetime, timedelta -import json -import logging -import os -import re -from typing import List - -import pandas as pd - -from flows.prepPolicy import selectPolicies -from models.agent import Agent -from models.policy import Policy -from services.API import AirlockAPIWrapper -from utils.configmanager import get_system_json, load_env -from utils.selector import Selector -from utils.utils import colorText, get_sanitized_input - -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–365): ", - value_type=int, - valid_range=(1, 365), - ) - - if not agents or not history_days: - print(colorText("No agents selected or invalid history range.", "red")) - return - - historical_date = (datetime.now() - timedelta(days=history_days)).strftime( - "%Y-%m-%d" - ) - today = datetime.now().strftime("%Y-%m-%d") - - all_history = [] - - for agent in agents: - try: - exechistory = api.history_execution(today, historical_date, agent.hostname) - except Exception as e: - print( - colorText( - f"❌ Error retrieving history for {agent.hostname}: {e}", "red" - ) - ) - continue - - if isinstance(exechistory, list): - for block in exechistory: - record = { - "Command": block.get("commandline", "N/A"), - "Date": block.get("datetime", "N/A"), - "Filename": block.get("filename", "N/A"), - "Policy Name": block.get("policyname", "N/A"), - "Hostname": block.get("hostname", "N/A"), - "Hash": block.get("sha256", "N/A"), - } - all_history.append(record) - - if not outputjson: - for key, value in record.items(): - print(colorText(f"{key}: {value}", "green")) - print("\n") - else: - print( - colorText(f"No execution history found for {agent.hostname}.", "yellow") - ) - - if outputjson: - print(json.dumps(all_history, indent=2)) - - -def findAllAgents(api): - # Step 1: Load data from API - policies = [Policy(**row["data"]) for _, row in api.policy_find_all().iterrows()] - agents = [Agent(**row["data"]) for _, row in api.agent_find_all().iterrows()] - - for agent in agents: - agent.enrich_with_policies(policies) - - return agents - - -def findAgents(api, return_dataframe): - agents = selectAgents(api) - working_dir = load_env("WORKING_DIR") - - if not agents: - logging.warning("No agents or policies found.") - print("No agents matched the criteria.") - return - - # Convert enriched agents to DataFrame - agent_dicts = [asdict(agent) for agent in agents] - agent_df = pd.DataFrame(agent_dicts) - - if return_dataframe: - logging.debug("Returning DataFrame to caller.") - return agent_df - - # Otherwise, print and optionally export - print(agent_df) - logging.debug("Displayed DataFrame to console.") - - user_input = ( - get_sanitized_input( - "\nWould you like to export the results to a CSV file? (y/n): " - ) - .strip() - .lower() - ) - if user_input == "y": - timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - filename = f"agentsearch_{timestamp}.csv" - file_path = os.path.join(str(working_dir), filename) - - agent_df.to_csv(file_path, index=False) - logging.info(f"Exported DataFrame to {file_path}") - - print( - colorText( - f"\n✓ Matched devices exported to: {working_dir}\\{filename}", - "green", - ) - ) - else: - logging.debug("User declined to export the DataFrame.") - - -def collect_device_names() -> List[str]: - print(colorText("🖥��Â Device Search", "cyan")) - print( - colorText( - "Enter the device hostnames you'd like to search for, one per line.", "cyan" - ) - ) - print( - colorText( - "When you're done, press Enter twice (Three times if you have a single device).\n", - "cyan", - ) - ) - print(colorText("Example:", "cyan")) - print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan")) - print(colorText("Paste or type your device names below:", "white")) - - device_input_lines = [] - empty_line_count = 0 - valid_line_pattern = re.compile(r"^[a-zA-Z0-9_\- ]+$") - - while True: - line = get_sanitized_input("") - stripped_line = line.strip() - - if stripped_line == "": - empty_line_count += 1 - if empty_line_count == 2: - break - continue - else: - empty_line_count = 0 - - if valid_line_pattern.match(stripped_line): - device_input_lines.append(stripped_line) - else: - print( - colorText( - f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", - "yellow", - ) - ) - - return [name for name in device_input_lines if name] - - -def choose_match_type() -> bool: - print(colorText("Use exact match? (Y for exact, N for fuzzy):", "white")) - return get_sanitized_input("").strip().lower() in ["y", "yes"] - - -def match_agents( - device_names: List[str], agents: List["Agent"], use_exact: bool -) -> List["Agent"]: - if use_exact: - return [ - agent - for agent in agents - if agent.hostname.lower() in [name.lower() for name in device_names] - ] - else: - pattern = "|".join(map(re.escape, device_names)) - regex = re.compile(pattern, re.IGNORECASE) - return [agent for agent in agents if regex.search(agent.hostname)] - - -def show_unmatched( - device_names: List[str], matched_agents: List["Agent"], use_exact: bool -): - if use_exact: - unmatched = [ - name - for name in device_names - if not any( - agent.hostname.lower() == name.lower() for agent in matched_agents - ) - ] - else: - unmatched = [ - name - for name in device_names - if not any( - re.search(re.escape(name), agent.hostname, re.IGNORECASE) - for agent in matched_agents - ) - ] - - if unmatched: - 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"]): - for agent in agents: - agent.enrich_with_policies(policies) - - -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")) - return [] - - use_exact = choose_match_type() - - policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()] - agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()] - matched_agents = match_agents(device_names, agents, use_exact) - matched_agents.sort(key=lambda agent: agent.hostname.lower()) - - 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")) - return [] - - 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): - line = "" - for col in range(3): - idx = row + col * rows - if idx < len(matched_agents): - line += f"{matched_agents[idx].hostname:<30}" - logger.info(line) - - matched_agents = Selector.select_with_mode( - matched_agents, - label_func=lambda agent: agent.hostname, - header="Matched Devices:", - ) - - if not matched_agents: - logger.debug("❌ No matching devices remain after refinement.") - print(colorText("❌ No matching devices remain after refinement.", "red")) - return [] - - enrich_agents(matched_agents, policies) - return matched_agents - - -def moveAgentToRelatedPolicy( - api: AirlockAPIWrapper, - agent: Agent, - mode: str = "audit", -): - """ - Moves an agent between audit and enforcement policies based on the mode. - - Args: - api: AirlockAPIWrapper instance. - agent: Agent object. - 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", "{}") - - if mode == "audit": - if agent.groupid in policy_relationship_map: - target_policy = policy_relationship_map[agent.groupid] - elif agent.groupid in policy_relationship_map.values(): - logger.debug( - f"Agent {agent.hostname} is already in an audit group. No action needed." - ) - print( - f"Agent {agent.hostname} is already in an audit group. No action needed." - ) - return - else: - logger.warning( - f"Error: No corresponding audit policy found for groupid: {agent.groupid}." - ) - return - - elif mode == "enforcement": - inverse_map = {v: k for k, v in policy_relationship_map.items()} - if agent.groupid in inverse_map: - target_policy = inverse_map[agent.groupid] - elif agent.groupid in inverse_map.values(): - logger.info( - f"Agent {agent.hostname} is already in an enforcement group. No action needed." - ) - return - else: - logger.warning( - f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}." - ) - return - - else: - logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.") - return - - result = api.agent_move(agent.agentid, target_policy) - return result - - -def toggleEnforcement(api: AirlockAPIWrapper): - choices = ["Audit", "Enforcement", "Exit"] - print(colorText("Move devices to which state?:", "yellow")) - direction = Selector.select_string(choices, False, False) - if direction == "Exit": - pass - else: - devices = selectAgents(api) - for device in devices: - print(device.hostname) - confirm = Selector.confirm( - "Would you like to continue with these devices? Y/N: " - ) - if direction and devices and confirm: - for device in devices: - result = moveAgentToRelatedPolicy(api, device, str(direction).lower()) - logger.info(f"{device.hostname}: result: {result}") - get_sanitized_input("Press enter to continue") - - -def moveAgents(api: AirlockAPIWrapper): - devices = selectAgents(api) - for device in devices: - print(device.hostname) - confirm_devices = Selector.confirm( - "Would you like to continue with these devices? Y/N: " - ) - if devices and confirm_devices: - policies = selectPolicies(api, False) - confirm_move = Selector.confirm( - f"Would you like to move these devices to {policies[0].name}?" - ) - if confirm_move: - for device in devices: - result = api.agent_move(device.agentid, policies[0].groupid) - logger.info(f"{device.hostname}: result: {result}") - else: - logger.info("Exiting without change") - get_sanitized_input("Press enter to continue") diff --git a/services/policyhandler.py b/services/policyhandler.py deleted file mode 100644 index e36fe23..0000000 --- a/services/policyhandler.py +++ /dev/null @@ -1,234 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -import datetime -import gc -import json -import logging -import os -import sys - -from bson import ObjectId -import pandas as pd -import tqdm - -from models.policy import Policy -from services.API import AirlockAPIWrapper -from utils.setup import get_base_directory -from utils.utils import colorText - -logger = logging.getLogger(__name__) - - -def pullPolicyExechistories( - api: AirlockAPIWrapper, - policy: Policy, - type: list, - days, - outputjson: bool, -): - - file_path = f"{get_base_directory()}\\cache\\chunkinator.json" - - # Ensure the file exists - if not os.path.exists(file_path): - with open(file_path, "w") as file: - json.dump({"error": "Success", "response": {"exechistories": []}}, file) - logger.debug(f"File '{file_path}' has been created.") - else: - logger.debug(f"File '{file_path}' already exists.") - - checkpoint = str(skipback(days)) - json_output = {"error": "Success", "response": {"exechistories": []}} - - with tqdm.tqdm( - file=sys.stdout, - leave=True, - total=10000, - desc=f"Checkpoint Progress: {checkpoint}", - colour="blue", - initial=1, - ) as filebar: - with tqdm.tqdm( - file=sys.stdout, - leave=True, - total=100, - desc=f"Total of {policy} Complete: ", - ) as pbar: - while True: - histories = api.history_logging( - type=type, checkpoint=checkpoint, policy=[policy.name] - ) - - # Ensure histories is a list of dictionaries - if not isinstance(histories, list) or not all( - isinstance(h, dict) for h in histories - ): - logger.error( - "Unexpected response format from API. Expected list of dictionaries." - ) - break - - filebar.total = len(histories) - - if not histories: - break - - for index, history_item in enumerate(histories): - if ( - "checkpoint" not in history_item - or "datetime" not in history_item - ): - continue # Skip malformed entries - - # Update checkpoint on last item - if index == len(histories) - 1: - checkpoint = history_item[ - "checkpoint" - ] # pyright: ignore[reportArgumentType] - filebar.desc = f"Checkpoint Progress: {checkpoint}" - break - - try: - history_date = datetime.datetime.strptime( - history_item["datetime"].replace( - " +0000 UTC", "" - ), # pyright: ignore[reportArgumentType] - "%Y-%m-%dT%H:%M:%SZ", - ).date() - except ValueError: - continue # Skip if date format is invalid - - if ( - datetime.date.today() - datetime.timedelta(days=days) - ) <= history_date: - json_output["response"]["exechistories"].append(history_item) - - filebar.update(1) - filebar.refresh() - - # Deduplicate entries - seen = {} - if os.path.exists(file_path): - with open(file_path, "r") as file: - existing_data = json.load(file) - combined = ( - existing_data["response"]["exechistories"] - + json_output["response"]["exechistories"] - ) - else: - combined = json_output["response"]["exechistories"] - - for entry in combined: - key = ( - entry.get("sha256"), - entry.get("filename"), - entry.get("hostname"), - ) - seen[key] = entry - - deduplicated = list(seen.values()) - with open(file_path, "w") as file: - json.dump( - { - "error": "Success", - "response": {"exechistories": deduplicated}, - }, - file, - ) - - json_output["response"]["exechistories"].clear() - - # Update progress bar based on last valid item - try: - last_date = datetime.datetime.strptime( - history_item["datetime"].replace(" +0000 UTC", ""), # type: ignore - "%Y-%m-%dT%H:%M:%SZ", - ).date() - date_diff = datetime.date.today() - last_date - percentage_diff = ( - ((days + 10) - date_diff.days) / (days + 10) - ) * 100 - pbar.n = round(percentage_diff) - pbar.set_description_str(f"Total of {policy} Complete: ") - pbar.refresh() - except Exception: - pass - - filebar.n = 1 - - # Final output - with open(file_path, "r") as file: - final_output = json.load(file) - os.remove(file_path) - - return json.dumps(final_output) if outputjson else None - - -def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days): - import airlock_libs - - executionhist_policy = pd.DataFrame() - exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days) - if exehist is not None: - data = json.loads(exehist) - executionhist_policy = pd.DataFrame(data["response"]["exechistories"]) - if not executionhist_policy.empty: - executionhist_policy = executionhist_policy[ - [ - "datetime", - "sha256", - "publisher", - "filename", - "hostname", - "username", - "pprocess", - "gprocess", - "commandline", - ] - ] - executionhist_policy["policy"] = policy # Add policy column here - executionhist_policy = executionhist_policy.drop_duplicates( - subset=["sha256", "filename", "hostname"] - ) - executionhist_policy = executionhist_policy.sort_values( - by=["sha256", "filename"] - ) - logger.debug(f"Staging of Execution history for policy: {policy} is complete") - print( - colorText( - f"Staging of Execution history for policy: {policy} is complete", - "green", - ) - ) - del data - del exehist - gc.collect() - return executionhist_policy - - -def skipback(days): - """ - Generate a MongoDB ObjectId for a given number of days ago from today. - """ - adjusted_days = days - date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta( - days=adjusted_days - ) - timestamp = int(date_days_ago.timestamp()) - hex_timestamp = format(timestamp, "08x") - objectid_hex = hex_timestamp + "0000000000000000" - return ObjectId(objectid_hex) diff --git a/utils/configmanager.py b/utils/configmanager.py index 0d5391a..d455da1 100644 --- a/utils/configmanager.py +++ b/utils/configmanager.py @@ -336,10 +336,3 @@ def load_env_json(key: str, default: str = "[]") -> Any: except Exception as e: logger.error(f"Failed to parse {key}: {e}") return json.loads(default) - - -# Backwards compatibility aliases (deprecated - use get_system_value instead) -get_protected_value = get_system_value -get_protected_json = get_system_json -load_protected_config = load_system_config -PROTECTED_KEYS = SYSTEM_CONFIG_KEYS # For backwards compatibility diff --git a/utils/executionfetcher.py b/utils/executionfetcher.py new file mode 100644 index 0000000..87e454d --- /dev/null +++ b/utils/executionfetcher.py @@ -0,0 +1,78 @@ +# 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 . + + +import datetime +import gc +import json +import logging + +from bson import ObjectId +import pandas as pd + +from services.API import AirlockAPIWrapper + +logger = logging.getLogger(__name__) + + +def getExecutions(api: AirlockAPIWrapper, policy, type, days): + import airlock_libs + + executionhist_policy = pd.DataFrame() + exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days) + if exehist is not None: + data = json.loads(exehist) + executionhist_policy = pd.DataFrame(data["response"]["exechistories"]) + if not executionhist_policy.empty: + executionhist_policy = executionhist_policy[ + [ + "datetime", + "sha256", + "publisher", + "filename", + "hostname", + "username", + "pprocess", + "gprocess", + "commandline", + ] + ] + executionhist_policy["policy"] = policy # Add policy column here + executionhist_policy = executionhist_policy.drop_duplicates( + subset=["sha256", "filename", "hostname"] + ) + executionhist_policy = executionhist_policy.sort_values( + by=["sha256", "filename"] + ) + logger.debug(f"Staging of Execution history for policy: {policy} is complete") + + del data + del exehist + gc.collect() + return executionhist_policy + + +def skipback(days): + """ + Generate a MongoDB ObjectId for a given number of days ago from today. + """ + adjusted_days = days + date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta( + days=adjusted_days + ) + timestamp = int(date_days_ago.timestamp()) + hex_timestamp = format(timestamp, "08x") + objectid_hex = hex_timestamp + "0000000000000000" + return ObjectId(objectid_hex) diff --git a/services/security.py b/utils/security.py similarity index 93% rename from services/security.py rename to utils/security.py index 70229e2..59c2e76 100644 --- a/services/security.py +++ b/utils/security.py @@ -169,17 +169,3 @@ def getAPI(USERNAME, SERVICE_NAME): logging.warning( "Password does not meet complexity requirements. Try again." ) - - -class APIKeyManager: - _api_key = None - - @classmethod - def load(cls, service: str, username: str, password: str): - cls._api_key = retrieve_api_key(service, username, password) - - @classmethod - def get(cls) -> str: - if cls._api_key is None: - raise ValueError("API key not loaded. Call APIKeyManager.load() first.") - return cls._api_key diff --git a/utils/selector.py b/utils/selector.py deleted file mode 100644 index d43f6d8..0000000 --- a/utils/selector.py +++ /dev/null @@ -1,357 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -import logging -from typing import Any, Callable, List, Optional, Union - -import pandas as pd - -from utils.utils import colorText, get_sanitized_input - -logger = logging.getLogger(__name__) - - -class Selector: - @staticmethod - def _get_sorted_items( - items: List[Any], label_func: Callable[[Any], str] - ) -> List[Any]: - return sorted(items, key=lambda item: label_func(item).lower()) - - @staticmethod - def _display_choices( - items: List[Any], - label_func: Callable[[Any], str], - num_columns: int = 3, - header: str = "Available Choices:", - ) -> None: - # Force single column if items are DataFrame rows - - if items and isinstance(items[0], (pd.Series, dict)): - num_columns = 1 - - rows = (len(items) + num_columns - 1) // num_columns - print(f"\n{header}") - for row in range(rows): - line = "" - for col in range(num_columns): - idx = row + col * rows - if idx < len(items): - label = label_func(items[idx]) - line += f"{idx + 1}: {label:<30}" - print(line) - - @staticmethod - def _display_selected_items( - selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 3 - ) -> None: - print(colorText("\nCurrent selections:", "cyan")) - if not selected: - print(" (none)") - return - sorted_selected = sorted(selected, key=lambda item: label_func(item).lower()) - rows = (len(sorted_selected) + num_columns - 1) // num_columns - for row in range(rows): - line = "" - for col in range(num_columns): - idx = row + col * rows - if idx < len(sorted_selected): - label = label_func(sorted_selected[idx]) - line += f"{label:<30}" - print(line) - - @staticmethod - def _parse_selection_input(input_str: str, max_index: int) -> List[int]: - selections = [] - for part in input_str.split(","): - part = part.strip() - if "-" in part: - try: - start, end = map(int, part.split("-")) - selections.extend(range(start, end + 1)) - except ValueError: - continue - elif part.isdigit(): - selections.append(int(part)) - return [i for i in selections if 1 <= i <= max_index] - - @staticmethod - def _select_from_list( - items: List[Any], - label_func: Callable[[Any], str], - allow_multiple: bool = False, - prompt_each: bool = False, - header: str = "Available Choices:", - num_columns: int = 3, - ) -> Union[Optional[Any], List[Any]]: - if not items: - logger.warning("No items available for selection.") - return None - - full_sorted_items = Selector._get_sorted_items(items, label_func) - remaining_items = full_sorted_items.copy() - selected = [] - - if allow_multiple: - while True: - Selector._display_choices( - remaining_items, label_func, num_columns=num_columns, header=header - ) - Selector._display_selected_items( - selected, label_func, num_columns=num_columns - ) - choice = ( - get_sanitized_input( - "Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: " - ) - .strip() - .lower() - ) - if choice == "q": - break - elif choice == "r": - selected.clear() - remaining_items = full_sorted_items.copy() - print(colorText("🔄 Selections reset.", "yellow")) - continue - indices = Selector._parse_selection_input(choice, len(remaining_items)) - newly_selected = [] - for index in indices: - item = remaining_items[index - 1] - if item not in selected: - selected.append(item) - newly_selected.append(item) - if prompt_each: - logger.info(f"Selected: {label_func(item)}") - else: - logger.warning("Item already selected.") - remaining_items = [ - item for item in remaining_items if item not in newly_selected - ] - return selected if selected else None - else: - Selector._display_choices( - full_sorted_items, label_func, num_columns=num_columns, header=header - ) - try: - choice = int(get_sanitized_input("Select one item by number: ")) - if 1 <= choice <= len(full_sorted_items): - selected_item = full_sorted_items[choice - 1] - logger.info(f"Selected: {label_func(selected_item)}") - return selected_item - else: - logger.warning("Selection out of range.") - except ValueError: - logger.warning("Invalid input.") - return None - - @staticmethod - def select_with_mode( - items: List[Any], - label_func: Callable[[Any], str], - header: str = "Available Choices:", - ) -> List[Any]: - print( - colorText( - "Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", - "white", - ) - ) - mode = get_sanitized_input("").strip().lower() - if mode == "a": - return items - selected = Selector._select_from_list( - items, - label_func=label_func, - allow_multiple=True, - prompt_each=False, - header=header, - ) - if not selected: - return items - if mode == "i": - print(colorText(f"✅ Included {len(selected)} item(s).", "green")) - return selected - elif mode == "e": - 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")) - return items - - @staticmethod - def select_objects( - objects: List[Any], allow_multiple: bool = False, prompt_each: bool = False - ) -> Union[Optional[Any], List[Any]]: - return Selector._select_from_list( - objects, - label_func=lambda obj: getattr(obj, "name", str(obj)), - allow_multiple=allow_multiple, - prompt_each=prompt_each, - header="Available Objects:", - ) - - @staticmethod - def select_string( - options: List[str], allow_multiple: bool = False, prompt_each: bool = False - ) -> Union[Optional[str], List[str]]: - return Selector._select_from_list( - options, - label_func=str, - allow_multiple=allow_multiple, - prompt_each=prompt_each, - header="Available Options:", - ) - - @staticmethod - def select_int( - options: List[int], allow_multiple: bool = False, prompt_each: bool = False - ) -> Union[Optional[int], List[int]]: - return Selector._select_from_list( - options, - label_func=lambda x: str(x), - allow_multiple=allow_multiple, - prompt_each=prompt_each, - header="Available Integers:", - ) - - @staticmethod - def select_value( - prompt: str, - value_type: type = int, - valid_range: Optional[tuple] = None, - allow_quit: bool = False, - ) -> Optional[Any]: - while True: - user_input = get_sanitized_input(prompt).strip().lower() - if allow_quit and user_input == "q": - logger.info("User opted to quit value selection.") - return None - try: - value = value_type(user_input) - if valid_range: - min_val, max_val = valid_range - if not (min_val <= value <= max_val): - logger.warning(f"Value out of range ({min_val}–{max_val}).") - continue - logger.info(f"User selected value: {value}") - return value - except ValueError: - logger.warning(f"Invalid input. Expected a {value_type.__name__}.") - - @staticmethod - def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool: - while True: - response = get_sanitized_input(prompt).strip().lower() - if response in ["y", "yes"]: - logger.info("User confirmed action.") - return True - elif response in ["n", "no"]: - logger.info("User declined action.") - return False - else: - logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.") - - @staticmethod - def select_dataframe_rows( - df: pd.DataFrame, - columns: Optional[List[str]] = None, - allow_multiple: bool = False, - prompt_each: bool = False, - header: str = "Available Rows:", - ) -> List[pd.Series]: - if df.empty: - print("DataFrame is empty.") - return [] - - if columns: - df = df[columns] - - items = [row for _, row in df.iterrows()] - - def label_func(row): - return str(row.to_dict()) - - result = Selector._select_from_list( - items, - label_func=label_func, - allow_multiple=allow_multiple, - prompt_each=prompt_each, - header=header, - ) - - if isinstance(result, pd.Series): - return [result] - elif isinstance(result, list): - return result - else: - return [] - - @staticmethod - def select_dataframe_with_mode( - df: pd.DataFrame, - columns: Optional[List[str]] = None, - header: str = "Available Rows:", - ) -> List[pd.Series]: - if df.empty: - print("⚠️ DataFrame is empty.") - return [] - - # Filter columns if specified - if columns: - df = df[columns] - - items = df.to_dict("records") - - def label_func(row): - return " | ".join(str(row[col]) for col in df.columns) - - # Show rows first - print(colorText(header, "cyan")) - for i, row in enumerate(items): - print(f"{i}: {label_func(row)}") - - # Prompt for mode once - print( - colorText( - "\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", - "white", - ) - ) - mode = get_sanitized_input("").strip().lower() - - if mode == "a": - return [pd.Series(row) for row in items] - - # Prompt for selection only once - selected = Selector._select_from_list( - items, - label_func=label_func, - allow_multiple=True, - prompt_each=False, - header=header, - ) - - if not selected: - return [pd.Series(row) for row in items] - - if mode == "i": - 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")) - return [pd.Series(row) for row in items if row not in selected] - else: - print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow")) - return [] diff --git a/utils/setup.py b/utils/setup.py index 36cf927..32201a3 100644 --- a/utils/setup.py +++ b/utils/setup.py @@ -83,7 +83,18 @@ def get_base_directory() -> Path: return home / ".local" / "share" / "Loxide" -def configure_logging(log_dir: Path, log_level: str = "INFO"): +def configure_logging(log_dir: Path, cache_dir: Path, log_level: str = "INFO"): + """ + Configure logging and return function to attach notification handler. + + Args: + log_dir: Directory for log files + cache_dir: Directory for cache files (used by version checker) + log_level: Logging level string + + Returns: + Function to attach notification handler to Textual app + """ log_file = log_dir / "Loxide.log" config = { @@ -134,7 +145,8 @@ def configure_logging(log_dir: Path, log_level: str = "INFO"): # Return a function to attach the notification handler once the app is created def attach_notification_handler(app): - """Attach the Textual notification handler to the root logger.""" + """Attach the Textual notification handler and version checker to the app.""" + # Attach logging handler handler = TextualNotificationHandler(app) handler.setLevel(logging.WARNING) # Only WARNING and above formatter = logging.Formatter("%(name)s: %(message)s") @@ -142,6 +154,17 @@ def configure_logging(log_dir: Path, log_level: str = "INFO"): logging.getLogger().addHandler(handler) logging.getLogger().debug("✅ Textual notification handler attached.") + # Attach version checker (checks in background, notifies if update available) + try: + from utils.versionchecker import create_update_notifier + + create_update_notifier(app, cache_dir=cache_dir) + logging.getLogger().debug("✅ Version checker attached.") + except ImportError as e: + logging.getLogger().debug(f"Version checker not available: {e}") + except Exception as e: + logging.getLogger().warning(f"Could not initialize version checker: {e}") + return attach_notification_handler @@ -173,7 +196,7 @@ def setup(): # Configure logging with system-defined log level log_level = get_system_value("LOG_LEVEL", str, "INFO") - attach_handler = configure_logging(dirs["logs"], log_level) + attach_handler = configure_logging(dirs["logs"], dirs["cache"], log_level) # Load user config (mutable) load_user_config(dirs["config"]) diff --git a/utils/utils.py b/utils/utils.py index b6ed034..06c561d 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -19,81 +19,10 @@ import os import platform import re import subprocess -import tempfile -import tkinter as tk -from tkinter import filedialog - -import pandas as pd logger = logging.getLogger(__name__) -def import_to_dataframe(file_path: str) -> pd.DataFrame: - df = pd.DataFrame() - - try: - if not os.path.exists(file_path): - print(colorText(f"Error: File '{file_path}' does not exist.", "red")) - return df - - ext = os.path.splitext(file_path)[1].lower() - - if ext == ".csv": - df = pd.read_csv(file_path) - elif ext == ".parquet": - df = pd.read_parquet(file_path) - else: - print(colorText(f"Error: Unsupported file extension '{ext}'.", "red")) - return df - - if df.empty: - print(colorText("Error: File has headers but no data rows.", "red")) - else: - print(colorText(f"Data loaded successfully from {file_path}", "green")) - - return df - - except pd.errors.EmptyDataError: - print( - colorText( - "Notice: CSV file is completely empty, falling back to empty frame", - "white", - ) - ) - return pd.DataFrame() - - except Exception as e: - print(colorText(f"Error reading file: {e}", "red")) - return pd.DataFrame() - - -def choose_directory(): - root = tk.Tk() - root.withdraw() # Hide the main window - directory = filedialog.askdirectory(title="Select a Directory") - print("Selected directory:", directory) - return directory - - -def choose_file(initial_directory=None, required_substring=None): - """Open a file dialog and ensure the selected file contains a required substring.""" - while True: - root = tk.Tk() - root.withdraw() # Hide the main window - file_path = filedialog.askopenfilename(initialdir=initial_directory) - - if not file_path: - print("No file selected.") - return None - - if required_substring and required_substring not in file_path: - print( - f"The selected file must contain '{required_substring}' in its path or name. Please try again." - ) - else: - return file_path - - def get_sanitized_input(prompt: str) -> str: while True: user_input = input(prompt) @@ -151,235 +80,6 @@ def irtang(): ) -def section_header(title): - print( - colorText( - "\n --------------------------------------------------------------------", - "cyan", - ) - ) - print(colorText(f" ------------- {title} -------------", "cyan")) - print( - colorText( - " --------------------------------------------------------------------", - "cyan", - ) - ) - - -def areYouSure(): - print( - colorText( - "🛑****************************************************************************************************************************************🛑", - "red", - ) - ) - print( - colorText( - "⚠️=========================================================================================================================================⚠️", - "yellow", - ) - ) - print( - colorText( - "🛑========================================================================================================================================🛑", - "red", - ) - ) - print( - colorText( - "⚠️-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------⚠️", - "yellow", - ) - ) - print( - colorText( - "🛑========================================================================================================================================🛑", - "red", - ) - ) - print( - colorText( - "⚠️=========================================================================================================================================⚠️", - "yellow", - ) - ) - print( - colorText( - "🛑****************************************************************************************************************************************🛑", - "red", - ) - ) - - -def locked(): - print( - colorText( - r""" - ████████████████████████████████████████████████████████████████ - ███ ██ - ██ ██████ ███ - ██ ████████████ ███ - ██ ████ ███ ███ - ██ ███ ███ ███ - ██ ███ ███ ███ - ██ ▒████████████████████ ███ - ██ ██████████████████████ ███ - ██ ██████████████████████ ███ - ██ ██████████████████████ ███ - ██ ██████████████████████ ███ - ██ ██████████████████████ ███ - ██ ███ - ███ ███ - ████████████████████████████████████████████████████████████████████ - ▒██████████████████████████████████████████████████████████████████▒ - ▒████ - ▒████ - ▓██████████████████████████████████████████ - █████████████████████████████████████████████░ -""", - "yellow", - ) - ) - - -def printDeviceEnforceChecklist(): - print( - colorText( - "\n --------------------------------------------------------------------", - "cyan", - ) - ) - print( - colorText( - " ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------", - "cyan", - ) - ) - print( - colorText( - " --------------------------------------------------------------------", - "cyan", - ) - ) - print( - colorText( - "\nSequentually follow these steps to prepare a policy for enforcement:", - "white", - ) - ) - - print( - colorText( - "\n1. Choose which originating policy or policies to move to enforcement", - "cyan", - ) - ) - print( - colorText( - "2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes", - "cyan", - ) - ) - print(colorText("3. Manually review the files:", "cyan")) - print( - colorText( - " 'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'", - "cyan", - ) - ) - print( - colorText( - " Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.", - "cyan", - ) - ) - print( - colorText( - " If metarules need to be created, please make note of them, and remove the row from the csv.", - "cyan", - ) - ) - print( - colorText( - " When complete, save both csv files to the directory 'approved' and choose this option.", - "cyan", - ) - ) - print( - colorText( - " This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed", - "cyan", - ) - ) - print( - colorText( - "4. Manually review the file 'needs_approved\\paths_needing_review.csv'", - "cyan", - ) - ) - print( - colorText( - " Remove the rows containing path exclusions you do not approve of", - "cyan", - ) - ) - print( - colorText( - " When complete, save the csv file to the directory 'approved'", "cyan" - ) - ) - print( - colorText( - " Do the same process with the list of publishers forthe same directories", - "cyan", - ) - ) - print(colorText(" Preflight Lists will be generated", "cyan")) - - print( - colorText( - "5. Choose the destination policy and parent and child allow list", "cyan" - ) - ) - - print( - colorText( - "6. Test ------------------------------------------------------", "cyan" - ) - ) - print(colorText(" Print rather than apply selected data.", "cyan")) - - print( - colorText( - "7. Liftoff ------------------------------------------------------", "cyan" - ) - ) - print( - colorText( - " Apply path exclusions according to allowed and approved paths", - "cyan", - ) - ) - print( - colorText(" Apply signed or attested hashes to Parent Allow List", "cyan") - ) - print( - colorText( - " Apply approved, but unsigned hashes to the Child Allow List", "cyan" - ) - ) - - print( - colorText( - "R. Remove/Reset Generated data - will prompt to allow keeping execution history", - "cyan", - ) - ) - - print(colorText("B. Back", "cyan")) - - def colorText(text, color): colors = { "red": "\033[91m", @@ -394,133 +94,6 @@ def colorText(text, color): return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}" -def formatHTML(df, output_html_path=None, overwrite=True): - from datetime import datetime - - # Get current date and filename for subtitle - today = datetime.now().strftime("%d %B %Y") # Changed to "Day Month Year" - filename = output_html_path.replace(".html", "") if output_html_path else "Report" - - dark_css = """ - - """ - - header = f""" -
-

Airlock Tools

-

{filename} - {today}

-
- """ - - html_table = df.to_html(index=False, escape=False) - styled_html = ( - f"\n" - f"Airlock Tools Report\n" - f"\n" - f"{dark_css}\n" - f"{header}\n" - f"
\n" - f" {html_table}\n" - f"
\n" - f"\n" - f"" - ) - if output_html_path: - with open(output_html_path, "w", encoding="utf-8") as f: - f.write(styled_html) - print(f"✅ Styled table saved to '{output_html_path}'") - elif overwrite: - with tempfile.NamedTemporaryFile( - suffix=".html", delete=False, mode="w", encoding="utf-8" - ) as f: - f.write(styled_html) - temp_path = f.name - - print(f"✅ Styled table saved to temporary file: {temp_path}") - else: - return styled_html - - def open_directory(path): system = platform.system() @@ -530,13 +103,3 @@ def open_directory(path): subprocess.run(["xdg-open", path]) else: raise OSError(f"Unsupported operating system: {system}") - - -def print_x_wide(items: list, width: int): - for i in range(0, len(items), width): - row = items[i : i + width] - print(" | ".join(row)) - - -def clear_screen(): - os.system("cls" if os.name == "nt" else "clear") diff --git a/utils/versionchecker.py b/utils/versionchecker.py new file mode 100644 index 0000000..9b4bd24 --- /dev/null +++ b/utils/versionchecker.py @@ -0,0 +1,560 @@ +# 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 . + +""" +Version checking and update notification system for Loxide. + +Checks against Gitea releases at: +https://git.racooncity.org/brotoskyj/AirlockTools/releases +""" + +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +import json +import logging +from pathlib import Path +import re +import threading +from typing import Callable, Optional + +import requests + +logger = logging.getLogger(__name__) + +# Current application version - UPDATE THIS ON EACH RELEASE +__version__ = "0.7.0" + +# Gitea release API configuration +GITEA_API_BASE = "https://git.racooncity.org/api/v1" +REPO_OWNER = "brotoskyj" +REPO_NAME = "AirlockTools" +RELEASES_URL = f"{GITEA_API_BASE}/repos/{REPO_OWNER}/{REPO_NAME}/releases" +RELEASES_PAGE_URL = f"https://git.racooncity.org/{REPO_OWNER}/{REPO_NAME}/releases" + +# How often to check for updates (in hours) +CHECK_INTERVAL_HOURS = 24 + + +@dataclass +class ReleaseInfo: + """Information about a release.""" + + tag_name: str + version: tuple # Parsed semantic version (major, minor, patch) + name: str + body: str # Release notes + published_at: datetime + html_url: str + download_url: Optional[str] = None # URL to download the release asset + is_prerelease: bool = False + + +@dataclass +class UpdateCheckResult: + """Result of an update check.""" + + current_version: str + latest_version: Optional[str] + update_available: bool + release_info: Optional[ReleaseInfo] + error: Optional[str] = None + + +def parse_version(version_str: str) -> Optional[tuple]: + """ + Parse a version string into a comparable tuple. + Supports formats: v1.2.3, 1.2.3, v1.2, 1.2 + + Returns (major, minor, patch) tuple or None if parsing fails. + """ + if not version_str: + return None + + # Strip 'v' prefix if present + clean = version_str.lstrip("vV").strip() + + # Match semantic version pattern + match = re.match(r"^(\d+)(?:\.(\d+))?(?:\.(\d+))?", clean) + if not match: + return None + + major = int(match.group(1)) + minor = int(match.group(2)) if match.group(2) else 0 + patch = int(match.group(3)) if match.group(3) else 0 + + return (major, minor, patch) + + +def compare_versions(v1: tuple, v2: tuple) -> int: + """ + Compare two version tuples. + + Returns: + -1 if v1 < v2 + 0 if v1 == v2 + 1 if v1 > v2 + """ + for a, b in zip(v1, v2): + if a < b: + return -1 + if a > b: + return 1 + return 0 + + +def get_current_version() -> str: + """Get the current application version.""" + return __version__ + + +def _parse_release_response(release_data: dict) -> Optional[ReleaseInfo]: + """Parse a release from Gitea API response.""" + try: + tag_name = release_data.get("tag_name", "") + version = parse_version(tag_name) + if not version: + logger.debug(f"Could not parse version from tag: {tag_name}") + return None + + # Parse published date + published_str = release_data.get("published_at", "") + try: + published_at = datetime.fromisoformat(published_str.replace("Z", "+00:00")) + except (ValueError, AttributeError): + published_at = datetime.now(UTC) + + # Get download URL from assets if available + download_url = None + assets = release_data.get("assets", []) + for asset in assets: + # Prefer .exe or .zip files + name = asset.get("name", "").lower() + if name.endswith((".exe", ".zip", ".msi")): + download_url = asset.get("browser_download_url") + break + + return ReleaseInfo( + tag_name=tag_name, + version=version, + name=release_data.get("name", tag_name), + body=release_data.get("body", ""), + published_at=published_at, + html_url=release_data.get("html_url", RELEASES_PAGE_URL), + download_url=download_url, + is_prerelease=release_data.get("prerelease", False), + ) + except Exception as e: + logger.warning(f"Failed to parse release data: {e}") + return None + + +def fetch_latest_release( + include_prerelease: bool = False, timeout: int = 10 +) -> Optional[ReleaseInfo]: + """ + Fetch the latest release from Gitea. + + Args: + include_prerelease: Whether to include pre-release versions + timeout: Request timeout in seconds + + Returns: + ReleaseInfo for the latest release, or None if fetch fails + """ + try: + response = requests.get( + RELEASES_URL, + params={"limit": 10}, # Get last 10 releases to find latest stable + timeout=timeout, + headers={"Accept": "application/json"}, + ) + response.raise_for_status() + + releases = response.json() + if not releases: + logger.debug("No releases found") + return None + + # Find the latest release (first non-prerelease if we're excluding them) + for release_data in releases: + release_info = _parse_release_response(release_data) + if release_info is None: + continue + + if include_prerelease or not release_info.is_prerelease: + return release_info + + # If all are prereleases and we're excluding them, return the first one anyway + # but log a warning + if releases: + logger.debug("All releases are pre-releases") + return _parse_release_response(releases[0]) + + return None + + except requests.exceptions.Timeout: + logger.warning("Timeout fetching releases from Gitea") + return None + except requests.exceptions.RequestException as e: + logger.warning(f"Failed to fetch releases: {e}") + return None + except (json.JSONDecodeError, KeyError) as e: + logger.warning(f"Failed to parse release response: {e}") + return None + + +def check_for_updates(include_prerelease: bool = False) -> UpdateCheckResult: + """ + Check if a newer version is available. + + Args: + include_prerelease: Whether to consider pre-release versions + + Returns: + UpdateCheckResult with the check results + """ + current = get_current_version() + current_parsed = parse_version(current) + + if not current_parsed: + return UpdateCheckResult( + current_version=current, + latest_version=None, + update_available=False, + release_info=None, + error="Could not parse current version", + ) + + release_info = fetch_latest_release(include_prerelease=include_prerelease) + + if release_info is None: + return UpdateCheckResult( + current_version=current, + latest_version=None, + update_available=False, + release_info=None, + error="Could not fetch release information", + ) + + is_newer = compare_versions(release_info.version, current_parsed) > 0 + + return UpdateCheckResult( + current_version=current, + latest_version=release_info.tag_name, + update_available=is_newer, + release_info=release_info, + ) + + +class VersionChecker: + """ + Background version checker that periodically checks for updates + and can notify the application when updates are available. + """ + + def __init__( + self, + cache_dir: Optional[Path] = None, + check_interval_hours: int = CHECK_INTERVAL_HOURS, + on_update_available: Optional[Callable[[UpdateCheckResult], None]] = None, + ): + """ + Initialize the version checker. + + Args: + cache_dir: Directory to store last check timestamp + check_interval_hours: Hours between automatic checks + on_update_available: Callback when update is available + """ + self.cache_dir = cache_dir + self.check_interval = timedelta(hours=check_interval_hours) + self.on_update_available = on_update_available + self._last_check: Optional[datetime] = None + self._last_result: Optional[UpdateCheckResult] = None + self._check_thread: Optional[threading.Thread] = None + self._dismissed_version: Optional[str] = None + + # Load cached state + self._load_cache() + + @property + def cache_file(self) -> Optional[Path]: + if self.cache_dir: + return self.cache_dir / "version_check_cache.json" + return None + + def _load_cache(self) -> None: + """Load cached check state.""" + if not self.cache_file or not self.cache_file.exists(): + return + + try: + with open(self.cache_file, "r") as f: + data = json.load(f) + + if "last_check" in data: + self._last_check = datetime.fromisoformat(data["last_check"]) + # Don't load dismissed_version - dismiss is session-only + + except (json.JSONDecodeError, ValueError, OSError) as e: + logger.debug(f"Could not load version check cache: {e}") + + def _save_cache(self) -> None: + """Save check state to cache.""" + if not self.cache_file: + return + + try: + self.cache_file.parent.mkdir(parents=True, exist_ok=True) + data = {} + if self._last_check: + data["last_check"] = self._last_check.isoformat() + # Don't save dismissed_version - dismiss is session-only + + with open(self.cache_file, "w") as f: + json.dump(data, f) + + except OSError as e: + logger.debug(f"Could not save version check cache: {e}") + + def should_check(self) -> bool: + """Determine if enough time has passed to check again.""" + if self._last_check is None: + return True + + elapsed = datetime.now(UTC) - self._last_check + return elapsed >= self.check_interval + + def check_now( + self, force: bool = False, include_prerelease: bool = False + ) -> UpdateCheckResult: + """ + Check for updates immediately. + + Args: + force: Check even if recently checked + include_prerelease: Include pre-release versions + + Returns: + UpdateCheckResult + """ + if not force and not self.should_check() and self._last_result: + return self._last_result + + result = check_for_updates(include_prerelease=include_prerelease) + self._last_check = datetime.now(UTC) + self._last_result = result + self._save_cache() + + # Notify if update available and not dismissed + if ( + result.update_available + and self.on_update_available + and result.latest_version != self._dismissed_version + ): + self.on_update_available(result) + + return result + + def check_async( + self, force: bool = False, include_prerelease: bool = False + ) -> None: + """ + Check for updates in background thread. + + Args: + force: Check even if recently checked + include_prerelease: Include pre-release versions + """ + if self._check_thread and self._check_thread.is_alive(): + return # Already checking + + if not force and not self.should_check(): + return # Too soon to check again + + def _check(): + try: + self.check_now(force=True, include_prerelease=include_prerelease) + except Exception as e: + logger.debug(f"Background version check failed: {e}") + + self._check_thread = threading.Thread(target=_check, daemon=True) + self._check_thread.start() + + def dismiss_update(self, version: str) -> None: + """ + Dismiss update notification for a specific version. + Only lasts for the current session - will nag again on next startup. + + Args: + version: Version to dismiss (e.g., "v1.2.3") + """ + # Session-only dismiss - don't save to cache + self._dismissed_version = version + + def clear_dismissed(self) -> None: + """Clear the dismissed version so user will be nagged again.""" + self._dismissed_version = None + + def get_last_result(self) -> Optional[UpdateCheckResult]: + """Get the result of the last check.""" + return self._last_result + + +# Global instance for easy access +_global_checker: Optional[VersionChecker] = None + + +def get_version_checker( + cache_dir: Optional[Path] = None, + on_update_available: Optional[Callable[[UpdateCheckResult], None]] = None, +) -> VersionChecker: + """ + Get or create the global version checker instance. + + Args: + cache_dir: Directory for caching (only used on first call) + on_update_available: Callback for updates (only used on first call) + + Returns: + The global VersionChecker instance + """ + global _global_checker + + if _global_checker is None: + _global_checker = VersionChecker( + cache_dir=cache_dir, + on_update_available=on_update_available, + ) + + return _global_checker + + +def format_update_message(result: UpdateCheckResult, short: bool = False) -> str: + """ + Format a human-readable update message. + + Args: + result: The update check result + short: Whether to use a short format + + Returns: + Formatted message string + """ + if not result.update_available: + return f"✅ Loxide is up to date (v{result.current_version})" + + if short: + return f"🆕 Update available: {result.latest_version}" + + msg = f"🆕 Loxide {result.latest_version} is available! (current: v{result.current_version})" + + if result.release_info: + msg += f"\n📥 Download: {result.release_info.html_url}" + + # Include release notes preview if available + if result.release_info.body: + notes = result.release_info.body.strip() + # Truncate if too long + if len(notes) > 200: + notes = notes[:200] + "..." + msg += f"\n\n📋 Release Notes:\n{notes}" + + return msg + + +# --------------------------------------------------------------------------- +# Textual TUI Integration +# --------------------------------------------------------------------------- + + +def create_update_notifier( + app, cache_dir: Optional[Path] = None, nag_on_startup: bool = True +): + """ + Create a version checker that notifies via Textual toast notifications. + + This should be called after the Textual app is created. + + Args: + app: The Textual App instance + cache_dir: Directory for caching check state + nag_on_startup: Always show notification on startup if update available + + Returns: + The VersionChecker instance + """ + + def on_update_available(result: UpdateCheckResult): + """Callback when update is available - show toast notification.""" + try: + msg = f"🆕 Update available: {result.latest_version}\nGo to Settings to download" + try: + app.notify( + msg, title="Loxide Update Available", severity="warning", timeout=15 + ) + except RuntimeError: + app.call_from_thread( + app.notify, + msg, + title="Loxide Update Available", + severity="warning", + timeout=15, + ) + except Exception as e: + logger.debug(f"Could not show update notification: {e}") + + checker = get_version_checker( + cache_dir=cache_dir, + on_update_available=on_update_available, + ) + + # Store checker on app so Loxide.on_mount can use it + if nag_on_startup: + app._version_checker = checker + app._version_nag_shown = False + + return checker + + +def check_for_updates_startup( + cache_dir: Optional[Path] = None, +) -> Optional[UpdateCheckResult]: + """ + Check for updates during application startup. + + This performs a synchronous check but respects the cache interval, + so it will only actually query the network once per CHECK_INTERVAL_HOURS. + + Returns the result if an update is available, None otherwise. + + Example usage: + result = check_for_updates_startup(cache_dir) + if result and result.update_available: + print(format_update_message(result)) + """ + checker = get_version_checker(cache_dir=cache_dir) + + # Only check if enough time has passed (uses cache) + if not checker.should_check(): + result = checker.get_last_result() + if result and result.update_available: + return result + return None + + result = checker.check_now(force=False) + if result.update_available: + return result + return None