diff --git a/AirlockTools_Server.py b/AirlockTools_Server.py deleted file mode 100644 index 183635c..0000000 --- a/AirlockTools_Server.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -# TODO Add CSV injection prevention -# TODO Continue OTP and Local approval rewrites -# TODO Explore pywin32 -# TODO Fix Requirements.txt -# TODO Create Generic system_config.json for gitea - - -import logging -import os - -import dotenv -import urllib3 - -import flows.localApproval as la -from Server.scheduler_async import ( - recurring_job, - register_function, - reload_jobs, - start_scheduler, -) -from services.API import AirlockAPIWrapper -from services.policyhandler import updateAuditPoliciesFromEnforcementPolices -from services.security import getAPI -from utils.setup import setup - -urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) - - -def main(): - - # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored - - working_dir = setup() - - logger = logging.getLogger(__name__) - - dotenv.load_dotenv(dotenv_path=working_dir / ".env") - - try: - url = os.getenv("URL") - username = os.getenv("USERNAME") - - if not url: - raise ValueError("Missing URL in environment variables.") - if not username: - raise ValueError("Missing USERNAME in environment variables.") - - logger.debug(f"Retrieved URL: {url}") - logger.debug(f"Retrieved Username: {username}") - - except ValueError as e: - logger.error(f"Configuration error: {e}", exc_info=True) - raise - - api = AirlockAPIWrapper( - base_url=str(os.getenv("URL")), - api_key=getAPI(username, "AirlockTools"), - ) - - logger.info("Running non-interactively to start monitoring Airlock Changes") - - register_function("monitorLA", la.scheduleAddingLAHashes) - register_function("updateAuditPolicies", updateAuditPoliciesFromEnforcementPolices) - - if not os.path.exists("scheduling\\jobs.json"): - recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[api]) - recurring_job( - "updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api] - ) - else: - reload_jobs() - - start_scheduler() - - -if __name__ == "__main__": - main() diff --git a/IRT_icon_32-512.ico b/IRT_icon_32-512.ico deleted file mode 100644 index 103b72e..0000000 Binary files a/IRT_icon_32-512.ico and /dev/null differ diff --git a/Server/scheduler_async.py b/Server/scheduler_async.py deleted file mode 100644 index 87e8c80..0000000 --- a/Server/scheduler_async.py +++ /dev/null @@ -1,215 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import asyncio -import json -import logging -import os -from typing import Any, Callable, Dict, List - -logger = logging.getLogger(__name__) - -# Registry of functions that can be scheduled -FUNCTION_MAP: Dict[str, Callable] = {} - -# Dictionary to manually track scheduled jobs by ID -scheduled_jobs: Dict[str, asyncio.TimerHandle] = {} - -# Path to the JSON file for job persistence TODO - pin this to the correct place -JOBS_FILE = os.path.join(os.getcwd(), "jobs.json") - - -def register_function(name: str, func: Callable): - """ - Register a function so it can be called by name later. - Example: - register_function("say_hello", say_hello) - """ - FUNCTION_MAP[name] = func - - -def load_jobs() -> List[Dict[str, Any]]: - """ - Load jobs from the JSON file, or return [] if none exist. - """ - if not os.path.exists(JOBS_FILE): - return [] - with open(JOBS_FILE, "r") as f: - return json.load(f) - - -def save_jobs(jobs: List[Dict[str, Any]]): - """ - Save jobs to the JSON file (overwrite). - """ - with open(JOBS_FILE, "w") as f: - json.dump(jobs, f, indent=4) - - -def cancel_job(job_id: str): - """ - Cancel a scheduled job by ID and remove it from the registry and persistence. - """ - handle = scheduled_jobs.pop(job_id, None) - if handle: - handle.cancel() - logger.info(f"Cancelled job '{job_id}'") - - jobs = [j for j in load_jobs() if j.get("id") != job_id] - save_jobs(jobs) - - -def run_once_job( - job_id: str, - func_name: str, - delay_seconds: float, - args=None, - kwargs=None, - persist=True, -): - """ - Schedule a job to run once after a delay (in seconds). - """ - args = args or [] - kwargs = kwargs or {} - - def job_wrapper(): - func = FUNCTION_MAP.get(func_name) - if func is None: - logger.error(f"Function '{func_name}' is not registered.") - return - func(*args, **kwargs) - cancel_job(job_id) - - loop = asyncio.get_event_loop() - handle = loop.call_later(delay_seconds, job_wrapper) - scheduled_jobs[job_id] = handle - - if persist: - jobs = [j for j in load_jobs() if j.get("id") != job_id] - jobs.append( - { - "id": job_id, - "type": "once", - "delay": delay_seconds, - "function": func_name, - "args": args, - "kwargs": kwargs, - } - ) - save_jobs(jobs) - logger.info( - f"Scheduled one-time job '{job_id}' to run in {delay_seconds} seconds." - ) - - -def recurring_job( - job_id: str, func_name: str, interval: float, args=None, kwargs=None, persist=True -): - """ - Schedule a recurring job. - """ - args = args or [] - kwargs = kwargs or {} - - def job_wrapper(): - func = FUNCTION_MAP.get(func_name) - if func is None: - logger.error(f"Function '{func_name}' is not registered.") - return - func(*args, **kwargs) - # Reschedule the job - handle = asyncio.get_event_loop().call_later(interval, job_wrapper) - scheduled_jobs[job_id] = handle - - cancel_job(job_id) - handle = asyncio.get_event_loop().call_later(interval, job_wrapper) - scheduled_jobs[job_id] = handle - - if persist: - jobs = [j for j in load_jobs() if j.get("id") != job_id] - jobs.append( - { - "id": job_id, - "type": "recurring", - "interval": interval, - "function": func_name, - "args": args, - "kwargs": kwargs, - } - ) - save_jobs(jobs) - logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.") - - -def reload_jobs(): - """ - Reload jobs from JSON and reschedule them. - """ - jobs = load_jobs() - for job in jobs: - if job["type"] == "once": - run_once_job( - job["id"], - job["function"], - job["delay"], - job.get("args"), - job.get("kwargs"), - persist=False, - ) - elif job["type"] == "recurring": - recurring_job( - job["id"], - job["function"], - job["interval"], - job.get("args"), - job.get("kwargs"), - persist=False, - ) - - -async def start_scheduler(): - """ - Start the asynchronous scheduler loop. - - This function is a placeholder to keep the event loop alive. - Jobs are scheduled using asyncio.call_later and do not require polling. - """ - try: - await asyncio.Event().wait() - except asyncio.CancelledError: - logger.critical("Scheduler stopped.") - - """ - Start the asynchronous scheduler loop. - - This function is a placeholder for compatibility. Since we use asyncio.call_later, - jobs are scheduled directly on the event loop and no polling is required. - - Usage: - # In an async app (e.g., Textual) - asyncio.create_task(start_scheduler()) - - # Or in a standalone script - async def main(): - await start_scheduler() - - asyncio.run(main()) - """ - try: - while True: - await asyncio.sleep(3600) # Sleep indefinitely; jobs run via call_later - except asyncio.CancelledError: - logger.critical("Scheduler stopped.") diff --git a/README.md b/docs/README.md similarity index 100% rename from README.md rename to docs/README.md diff --git a/flows/localApproval.py b/flows/localApproval.py index 5cc69a2..da89c8c 100644 --- a/flows/localApproval.py +++ b/flows/localApproval.py @@ -1,311 +1,242 @@ -# 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 . +""" +This module handles the creation of local approval requests. +""" - -import datetime import logging import os -import re import time - -import dotenv -import numpy as np -import pandas as pd +from typing import List, Optional from models.agent import Agent -from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents +from services.agenthandler import moveAgentToRelatedPolicy, selectAgents from services.API import AirlockAPIWrapper -from utils.configmanager import get_protected_json, load_env, load_env_json -from utils.setup import get_base_directory +from utils.configmanager import get_protected_json from utils.utils import colorText, get_sanitized_input logger = logging.getLogger(__name__) -dotenv.load_dotenv() +class LocalApprovalRequestor: + """Handles creation of local approval requests in Loxide.""" -def getLocalApprovals(api: AirlockAPIWrapper): - base_dir = get_base_directory - result = api.otp_find_awaiting() - local_approval = pd.DataFrame(result["response"]["otpusage"]) - if os.path.exists(f"{base_dir}\\cache\\newest_local_approval.parquet"): - previous_run = pd.read_parquet( - f"{base_dir}\\cache\\newest_local_approval.parquet" - ) - previous_run.to_parquet( - f"{base_dir}\\cache\\last_local_approval.parquet", index=False - ) - os.remove(f"{base_dir}\\cache\\newest_local_approval.parquet") + def __init__(self, api: AirlockAPIWrapper, username: str = None): + """ + Initialize the local approval requestor. - # Only keep rows presumably created by the generate local approval function - local_approval = local_approval[ - local_approval["purpose"].str.startswith("🎫 Local Approval 🎫") - ] - - local_approval["batchid"] = local_approval["purpose"].apply( - lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1) - ) - - if not local_approval.empty: - local_approval.to_parquet( - f"{base_dir}\\cache\\newest_local_approval.parquet", index=False + Args: + api: AirlockAPIWrapper instance + username: Username creating the approvals (for tracking) + """ + self.api = api + self.policy_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") + self.username = ( + username or os.getenv("USERNAME") or os.getenv("USER") or "unknown" ) - return local_approval + def create_local_approval( + self, agent_id: str, duration_minutes: int, batch_id: Optional[int] = None + ) -> bool: + """ + Create a single local approval request. + Args: + agent_id: Agent ID to create approval for + duration_minutes: Duration of approval in minutes + batch_id: Optional batch identifier (defaults to timestamp) -def scheduleAddingLAHashes(api: AirlockAPIWrapper): + Returns: + True if successful, False otherwise + """ + if batch_id is None: + batch_id = int(time.time()) - policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") - bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]") - pups = load_env_json("PUPS", "[]") - threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type=int) + purpose = ( + f"🎫 Local Approval 🎫 - {duration_minutes} mins - " + f"batch:{batch_id} Client:{agent_id} User:{self.username}" + ) - try: - register_function("add_hash", returnFromLocalApproval) - register_function("move_device", moveAgentToRelatedPolicy) - except Exception as e: - logger.warning(f"Failed to register functions: {e}") - return - - try: - approvals_df = getNewLocalApprovals(api) - if approvals_df.empty: - logger.debug("No new local approvals found. Nothing to schedule.") - return - batches = approvals_df.groupby("batchid") - except Exception as e: - logger.warning(f"Failed to retrieve or group local approvals: {e}") - return - - for batchid, batch_df in batches: try: - duration_minutes = int(batch_df["duration"].iloc[0]) - start_time = datetime.datetime.now() - run_time = start_time + datetime.timedelta(minutes=duration_minutes) - early_time = start_time + datetime.timedelta( - minutes=np.floor(duration_minutes * 0.95) + self.api.otp_generate(agent_id, duration_minutes, purpose) + logger.info( + f"Generated local approval for {agent_id}, batch {batch_id}, by {self.username}" ) - - early_timestamp = early_time.timestamp() - run_timestamp = run_time.timestamp() - - # Schedule add_hash job - try: - run_once_job( - f"add_hash_{batchid}", - "add_hash", - early_timestamp, - [ - api, - batch_df, - policy_relationship_map, - bad_publisher_list, - pups, - threat_tolerance_constant, - ], - None, - ) - logger.debug(f"Scheduled add_hash for batch {batchid} at {early_time}") - except Exception: - logger.debug("Failed to schedule add_hash for batch {batchid}: {e}") - - # Schedule move_device jobs - devices = batch_df["agentid"].drop_duplicates().tolist() - agents = [] - - for device in devices: - rows = api.agent_find_by_hostname(device).iterrows() - agents += [Agent(**row["data"]) for _, row in rows] - - for agent in agents: - try: - run_once_job( - f"move_device_{agent.hostame}_{batchid}", - "move_device", - run_timestamp, - [api, agent, policy_relationship_map], - "enforcement", - ) - - print( - f"Scheduled move_device for device {agent.hostname} in batch {batchid} at {run_time}" - ) - except Exception as e: - print( - f"Failed to schedule move_device for device {agent.hostname} in batch {batchid}: {e}" - ) - + return True except Exception as e: - logger.warning(f"Failed to process batch {batchid}: {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. -def returnFromLocalApproval( - api, - device_df, - policy_relationship_map, - bad_publisher_list, - pups, - threat_tolerance_constant, -): - """ - # Get unique policy names from device list - policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist()) + Args: + agent: Agent object to move - # Create inverse map to go from Audit to Enforcement - inverse_map = {v: k for k, v in policy_relationship_map.items()} - - # Fetch all policies - all_policies = [Policy(row['groupid'], row['hidden'], row['name'], row['parent']) for _, row in api.policy_find_all().iterrows()] - - # Define policy types - policy_types = [1, 2, 6, 7] - - #TODO finish logic for adding hashes - """ - working_dir = load_env("WORKING_DIR") - policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") - bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]") - pups = load_env_json("PUPS", "[]") - threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE") - print( - f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}" - ) - - -def moveToLocalApproval(api: AirlockAPIWrapper): - possible_durations = [15, 60, 360, 1440, 10080] - duration_selected = None - - print(colorText("Please select a duration:", "white")) - for i, option in enumerate(possible_durations, start=1): - print(f"{i}. {option}") - - try: - - choice = int(get_sanitized_input("Enter the number of your choice:")) - if 1 <= choice <= len(possible_durations): - duration_selected = possible_durations[choice - 1] - print(colorText(f"You selected: {duration_selected}", "yellow")) - logger.debug(f"You selected: {duration_selected}") - else: - print(colorText("❌ Invalid choice.", "red")) - logger.debug("Invalid Input") - return - except ValueError: - print(colorText("❌ Invalid input. Please enter a number.", "red")) - logger.debug("Invalid Input") - return - - agents = selectAgents(api) - batch = int(time.time()) - - if not agents: - print(colorText("❌ No agents found or error retrieving agents.", "red")) - logger.debug("No agents found or error retrieving agents") - return - - for agent in agents: + Returns: + True if successful, False otherwise + """ try: - addLocalApproval(api, batch, duration_selected, agent.agentid) - moveAgentToRelatedPolicy(api, agent, "audit") + moveAgentToRelatedPolicy(self.api, agent, "audit") + logger.info(f"Moved {agent.hostname} to audit policy") + return True except Exception as e: - print(colorText(f"❌ Error processing agent {agent.hostname}: {e}", "red")) + 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. -def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid): + Args: + agents: List of Agent objects + duration_minutes: Duration of approval in minutes + db_path: Optional path to database for history tracking - purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}" - api.otp_generate(agentid, duration_selected, purpose) + 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") + ) -def monitorAuditStatus(api: AirlockAPIWrapper): - current_agents = findAllAgents(api) - last_agents = [] - if not last_agents: - last_agents = current_agents - policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") + for agent in agents: + try: + # Create local approval + approval_success = self.create_local_approval( + agent.agentid, duration_minutes, batch_id + ) - # Reverse map for audit β†’ enforcement - reverse_policy_map = {v: k for k, v in policy_relationship_map.items()} - known_transitions = set(policy_relationship_map.items()) | set( - reverse_policy_map.items() - ) + if not approval_success: + raise Exception("Failed to create local approval") - # Index last_agents by hostname for quick lookup - last_agent_map = {agent.hostname: agent for agent in last_agents} + # Move to audit policy + move_success = self.move_agent_to_audit(agent) - # Result buckets - newly_added = [] - same_policy = [] - moved_to_audit = [] - moved_to_enforcement = [] - unusual_move = [] + if not move_success: + raise Exception("Failed to move to audit policy") - for current in current_agents: - previous = last_agent_map.get(current.hostname) + print(colorText(f"βœ“ {agent.hostname}", "green")) + success_count += 1 - if not previous: - newly_added.append(current) - continue + except Exception as e: + print(colorText(f"βœ— {agent.hostname}: {e}", "red")) + logger.error(f"Error processing agent {agent.hostname}: {e}") + failure_count += 1 - if current.groupid == previous.groupid: - same_policy.append(current) - elif (previous.groupid, current.groupid) in known_transitions: - moved_to_audit.append(current) - elif (current.groupid, previous.groupid) in known_transitions: - moved_to_enforcement.append(current) - else: - unusual_move.append(current) + return batch_id, success_count, failure_count - # Return all five DataFrames - return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move + 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"), + ] -def getNewLocalApprovals(api: AirlockAPIWrapper): + # Display duration options + print(colorText("\n⏱️ Select Local Approval Duration:", "white")) + print(colorText("=" * 50, "white")) - working_dir = load_env("WORKING_DIR") - current_la = getLocalApprovals(api) + for i, (minutes, label) in enumerate(duration_options, start=1): + print(f" {i}. {label} ({minutes} minutes)") - # Load old approval list - old_la_path = f"{working_dir}\\Scheduling\\last_local_approval.parquet" - if os.path.exists(old_la_path): - old_la = pd.read_parquet(old_la_path) - else: - old_la = pd.DataFrame(columns=current_la.columns) + print(colorText("=" * 50, "white")) - # Create composite keys - current_la["key"] = ( - current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str) - ) - old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str) + # Get user selection + try: + choice = int(get_sanitized_input("\nEnter the number of your choice: ")) - # Find new entries - new_entries = current_la[~current_la["key"].isin(old_la["key"])] + 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 - # Convert 'granted' to datetime and filter by last 10 minutes - new_entries["granted"] = pd.to_datetime( - new_entries["granted"], utc=True, errors="coerce" - ) - ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( - minutes=10 - ) - recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago] + except ValueError: + print(colorText("❌ Invalid input. Please enter a number.", "red")) + logger.warning("Invalid input for duration selection") + return - # Save current approvals for next run - current_la.drop(columns=["key"], inplace=True) - current_la.to_parquet(old_la_path, index=False) + # Select agents + print(colorText("\n🎯 Select Agents for Local Approval:", "white")) + agents = selectAgents(self.api) - return recent_entries + 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("\nπŸ“‹ Summary:", "cyan")) + print(colorText(f" Duration: {duration_label}", "white")) + print(colorText(f" Agents: {len(agents)}", "white")) + + confirm = get_sanitized_input("\nProceed? (y/n): ").lower() + + 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/screens/moveagentworkflowscreen.py b/screens/moveagentworkflowscreen.py new file mode 100644 index 0000000..7199f55 --- /dev/null +++ b/screens/moveagentworkflowscreen.py @@ -0,0 +1,61 @@ +from typing import List + +from textual.app import ComposeResult +from textual.screen import Screen + +from models.agent import Agent +from widgets.agentmoveoperations import AgentMoveOperations +from widgets.multiagentselector import MultiAgentSelector +from widgets.resultsdisplay import ResultsDisplay + + +class MoveAgentWorkflowScreen(Screen): + """Screen that handles the agent movement workflow.""" + + def __init__(self, all_agents: List[Agent]): + super().__init__() + self.all_agents = all_agents + self.selected_agents = None + + def compose(self) -> ComposeResult: + """Start with the multi-agent selector.""" + yield MultiAgentSelector(self.all_agents) + + def on_multi_agent_selector_agents_selected( + self, message: MultiAgentSelector.AgentsSelected + ) -> None: + """Handle selected agents - switch to operations screen.""" + self.selected_agents = message.selected_agents + + # Remove the MultiAgentSelector + selector = self.query_one(MultiAgentSelector) + selector.remove() + + # Mount the AgentMoveOperations with the selected Agent objects + self.mount(AgentMoveOperations(self.selected_agents)) + + def on_agent_move_operations_operation_complete( + self, message: AgentMoveOperations.OperationComplete + ) -> None: + """Handle completion of move operation - transition to results screen.""" + # Format successful results + success_lines = [] + for agent, result in message.successful: + success_lines.append(f"βœ“ {agent.hostname}") + + # Format unsuccessful results + failure_lines = [] + for agent, error in message.unsuccessful: + failure_lines.append(f"βœ— {agent.hostname}: {error}") + + successful_text = "\n".join(success_lines) if success_lines else "(none)" + unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)" + + # Remove the operations widget + ops_widget = self.query_one(AgentMoveOperations) + ops_widget.remove() + + # Mount the results display + self.mount( + ResultsDisplay(message.operation, successful_text, unsuccessful_text) + ) diff --git a/screens/policyselectorscreen.py b/screens/policyselectorscreen.py new file mode 100644 index 0000000..5937860 --- /dev/null +++ b/screens/policyselectorscreen.py @@ -0,0 +1,91 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +""" +Policy Selector Screen Module + +Provides a Textual Screen wrapper for the PolicySelector widget that manages +the policy selection workflow. +""" + +import logging + +from textual.app import ComposeResult +from textual.screen import Screen + +from widgets.policyselector import PolicySelector + +logger = logging.getLogger(__name__) + + +class PolicySelectorScreen(Screen): + """ + A Textual Screen for policy selection in agent move operations. + + This screen wraps the PolicySelector widget and manages the workflow + of selecting a target policy for bulk agent movements. + + Attributes: + policies: List of available policies (Policy objects or DataFrame). + agent_move_operations: Reference to the parent AgentMoveOperations widget. + """ + + CSS = """ + Screen { + layout: vertical; + background: $surface; + } + """ + + def __init__( + self, + policies, + agent_move_operations=None, + ): + """ + Initialize the PolicySelectorScreen. + + Args: + policies: List of available policies to display. + agent_move_operations: Reference to parent AgentMoveOperations widget. + Used to call back when policy selection is confirmed. + """ + super().__init__() + self.policies = policies + self.agent_move_operations = agent_move_operations + + def compose(self) -> ComposeResult: + """Create the PolicySelector widget.""" + yield PolicySelector(self.policies) + + def on_policy_selector_policy_selected( + self, message: PolicySelector.PolicySelected + ) -> None: + """ + Handle policy selection from the PolicySelector widget. + + When a policy is selected, this handler: + 1. Closes the selector screen + 2. Calls the parent AgentMoveOperations to execute the move + + Args: + message (PolicySelector.PolicySelected): Contains the selected policy. + """ + # Pop this screen to return to AgentMoveOperations + self.app.pop_screen() + + # Call parent widget's method to execute the move + if self.agent_move_operations: + self.agent_move_operations._execute_move_to_policy(message.policy) diff --git a/utils/selector.py b/utils/selector.py index 0a0f498..2cac574 100644 --- a/utils/selector.py +++ b/utils/selector.py @@ -33,7 +33,7 @@ class Selector: def _display_choices( items: List[Any], label_func: Callable[[Any], str], - num_columns: int = 4, + num_columns: int = 3, header: str = "Available Choices:", ) -> None: # Force single column if items are DataFrame rows @@ -54,7 +54,7 @@ class Selector: @staticmethod def _display_selected_items( - selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 4 + selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 3 ) -> None: print(colorText("\nCurrent selections:", "cyan")) if not selected: @@ -93,7 +93,7 @@ class Selector: allow_multiple: bool = False, prompt_each: bool = False, header: str = "Available Choices:", - num_columns: int = 4, + num_columns: int = 3, ) -> Union[Optional[Any], List[Any]]: if not items: logger.warning("No items available for selection.") diff --git a/utils/tui.py b/utils/tui.py index 6eb070d..9fb1c54 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -23,6 +23,7 @@ from flows.prepPolicy import menu_policy_enforce from flows.quietAgent import findQuietAgents from models.agent import Agent from models.policy import Policy +from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen from screens.otpworkflowscreen import OTPWorkflowScreen from services.agenthandler import findAgents, moveAgents, toggleEnforcement from services.API import AirlockAPIWrapper @@ -30,9 +31,12 @@ from services.policyhandler import confirmUpdateAfromE from utils.configmanager import load_env from utils.setup import get_base_directory, load_user_config from utils.utils import open_directory +from widgets.agentmoveoperations import AgentMoveOperations from widgets.multiagentselector import MultiAgentSelector from widgets.OTP_generate import OTPGenerator from widgets.policytreewidget import PolicyTreeWidget +from widgets.resultsdisplay import ResultsDisplay +from widgets.retro_terminal_theme import get_retro_terminal_theme from widgets.themeselector import ThemeSelector dotenv.load_dotenv() @@ -106,6 +110,7 @@ class MainMenuScreen(Screen): ("πŸ”‡ - Find Quiet Hosts", "find_quiet_button"), ], "move": [ + ("πŸ”„ - Move Agent Workflow", "move_agent_workflow_button"), ("βœ… - Move to local approval", "move_local_button"), ("πŸ”„ - Move to Audit/Enforcement", "move_audit_button"), ("πŸ”€ - Move - Other", "move_other_button"), @@ -261,6 +266,47 @@ class MainMenuScreen(Screen): self.app.exit() + def on_agent_move_operations_operation_complete( + self, message: AgentMoveOperations.OperationComplete + ) -> None: + """Handle completion of agent move operation - show results.""" + logger.info( + "Agent move operation completed: %s, %d successful, %d unsuccessful", + message.operation, + len(message.successful), + len(message.unsuccessful), + ) + + # Format results for display + successful_text = "\n".join( + [f"{agent.hostname}" for agent, _ in message.successful] + ) + unsuccessful_text = "\n".join( + [f"{agent.hostname}: {error}" for agent, error in message.unsuccessful] + ) + + # Remove the operations widget + try: + ops_widget = self.query_one(AgentMoveOperations) + ops_widget.remove() + except Exception: + pass + + # Show results + self.query_one("#content", Vertical).mount( + ResultsDisplay(message.operation, successful_text, unsuccessful_text) + ) + + def on_results_display_go_back(self, message: ResultsDisplay.GoBack) -> None: + """Handle back button from results display.""" + try: + results_widget = self.query_one(ResultsDisplay) + results_widget.remove() + except Exception: + pass + # Return to main menu + self.app.pop_screen() + def on_directory_tree_file_selected( self, event: DirectoryTree.FileSelected ) -> None: @@ -282,6 +328,11 @@ class MainMenuScreen(Screen): _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {}) case "find_quiet_button": _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {}) + case "move_agent_workflow_button": + # Push Move Agent workflow screen + self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices)) + event.stop() + return # Don't exit the app case "move_local_button": _PENDING_JOB = ( "legacy", @@ -349,12 +400,21 @@ class Loxide(App): self.devices = [ Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows() ] + + # Enrich agents with policy information + if self.policies and self.devices: + for agent in self.devices: + agent.enrich_with_policies(self.policies) + logger.debug( + f"Enriched {len(self.devices)} agents with policy information" + ) except Exception as exc: logger.error("Failed to load policies/devices: %s", exc) self.policies = None self.devices = None def on_mount(self, api: AirlockAPIWrapper) -> None: + self.register_theme(get_retro_terminal_theme()) self.theme = self._textual_theme self.push_screen(MainMenuScreen(api)) diff --git a/utils/utils.py b/utils/utils.py index 26b2023..b6ed034 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -151,44 +151,6 @@ def irtang(): ) -def displayIntro(): - - print( - colorText( - r""" - _____ .__ .__ __ ___________ .__ - / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______ - / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/ -/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \ -\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ > - \/ \/ \/ \/ -""", - "cyan", - ) - ) - - -def welcome(): - print( - colorText( - "=================================================================================", - "cyan", - ) - ) - print( - colorText( - "======================== Welcome to the Airlock API Tool ========================", - "cyan", - ) - ) - print( - colorText( - "=================================================================================", - "cyan", - ) - ) - - def section_header(title): print( colorText( diff --git a/widgets/OTP_generate.py b/widgets/OTP_generate.py index 6aff859..a3476c6 100644 --- a/widgets/OTP_generate.py +++ b/widgets/OTP_generate.py @@ -6,7 +6,16 @@ from textual.css.query import NoMatches from textual.message import Message from textual.reactive import reactive from textual.widget import Widget -from textual.widgets import Button, Input, RadioButton, RadioSet, Static, TextArea +from textual.widgets import ( + Button, + Footer, + Header, + Input, + RadioButton, + RadioSet, + Static, + TextArea, +) from models.agent import Agent @@ -70,6 +79,7 @@ class OTPGenerator(Widget): pass def compose(self): + yield Header(show_clock=True, icon="βš™") title_text = Static( f"🎫 Generate One Time Passes for {len(self.devices)} device(s)", id="otpgen_title", @@ -165,6 +175,7 @@ class OTPGenerator(Widget): copy_button.styles.margin = (1, 0, 0, 0) copy_button.styles.display = "none" yield copy_button + yield Footer() def on_mount(self) -> None: """Set initial button state.""" diff --git a/widgets/agentmoveoperations.py b/widgets/agentmoveoperations.py new file mode 100644 index 0000000..1e55fdf --- /dev/null +++ b/widgets/agentmoveoperations.py @@ -0,0 +1,704 @@ +""" +Agent Move Operations Widget Module + +This module provides a Textual-based UI widget for performing bulk operations on +agent devices in the Airlock system. It allows users to: +- View selected agents and their current policy assignments +- Move agents to local approval mode with OTP enforcement +- Toggle agents between audit and enforcement policy modes +- Select and move agents to alternate policies (future implementation) + +The widget tracks operation state, manages button availability, and displays +results with success/failure summaries that can be copied to clipboard. + +Dependencies: + - textual: TUI framework for building the widget and UI components + - models.agent: Agent model class + - services.agenthandler: Core agent operation functions + - flows.localApproval: Local approval workflow handling +""" + +import logging +from typing import List + +from textual.containers import Horizontal, Vertical +from textual.css.query import NoMatches +from textual.message import Message +from textual.reactive import reactive +from textual.widget import Widget +from textual.widgets import Button, DataTable, Header, Static, TextArea + +from models.agent import Agent +from screens.policyselectorscreen import PolicySelectorScreen + +logger = logging.getLogger(__name__) + + +class AgentMoveOperations(Widget): + """ + A Textual widget for managing bulk agent operations and policy migrations. + + This widget provides a comprehensive UI for performing operations on multiple + selected agents. It displays the list of target agents and provides buttons to + trigger various bulk operations like toggling policy modes or enabling local approval. + + The widget manages its own state through reactive properties and provides real-time + feedback on operation progress and results. Operations are executed sequentially + per agent with error handling that tracks both successful and failed operations. + + Attributes: + operation_in_progress (reactive[bool]): Tracks whether an operation is currently + executing. Used to disable buttons during execution. + selected_operation (reactive[str]): Tracks which operation type is currently + selected or in progress (e.g., "local_approval", "toggle_enforcement"). + + Example: + ```python + agents = [agent1, agent2, agent3] + widget = AgentMoveOperations(agents) + ``` + """ + + # Reactive property to track if an operation is in progress + operation_in_progress = reactive(False) + # Tracks the currently selected operation type + selected_operation = reactive("") + + class OperationComplete(Message): + """ + Message posted when a bulk operation completes. + + This message is broadcast to parent widgets/screens to notify them of + operation completion along with detailed results. It contains the list + of agents that were processed and the outcome for each. + + Attributes: + operation (str): Name of the operation that completed (e.g., "Local Approval Mode"). + agents (List[Agent]): List of all agents that were targeted by the operation. + successful (List[tuple]): List of (Agent, result_data) tuples for successfully + processed agents. Result data varies by operation type. + unsuccessful (List[tuple]): List of (Agent, error_message) tuples for agents + where the operation failed. Error message is a string explaining the failure. + """ + + def __init__( + self, + operation: str, + agents: List[Agent], + successful: List[tuple], + unsuccessful: List[tuple], + ): + super().__init__() + self.operation = operation + self.agents = agents + self.successful = successful # List of (agent, result) tuples + self.unsuccessful = unsuccessful # List of (agent, error) tuples + + def __init__(self, agents: List[Agent]): + """ + Initialize the AgentMoveOperations widget. + + Args: + agents (List[Agent]): List of Agent objects to perform operations on. + These agents will be displayed in the widget's agent table. + """ + super().__init__() + self.agents = agents + + def watch_operation_in_progress(self, old_value: bool, new_value: bool) -> None: + """ + React to changes in the operation_in_progress reactive property. + + This is called automatically by Textual when operation_in_progress changes. + It updates the button states to reflect whether an operation is running. + + Args: + old_value (bool): Previous value of operation_in_progress. + new_value (bool): New value of operation_in_progress. + """ + self._update_button_states() + + def _update_button_states(self) -> None: + """ + Update the enabled/disabled state of operation buttons based on current status. + + This method implements the following logic: + - If an operation is in progress: disable all buttons + - If an operation is selected: disable only that operation's button + - If no operation is selected: enable all buttons + + The state transitions prevent users from starting multiple operations + simultaneously and provide visual feedback on which operation is active. + + Handles NoMatches exceptions gracefully in case buttons are not yet rendered. + """ + try: + local_approval_btn = self.query_one("#local_approval_btn", Button) + toggle_enforcement_btn = self.query_one("#toggle_enforcement_btn", Button) + other_policy_btn = self.query_one("#other_policy_btn", Button) + + # If operation in progress, disable all + if self.operation_in_progress: + local_approval_btn.disabled = True + toggle_enforcement_btn.disabled = True + other_policy_btn.disabled = True + else: + # If an operation was selected, keep it disabled, enable others + if self.selected_operation: + local_approval_btn.disabled = ( + self.selected_operation == "local_approval" + ) + toggle_enforcement_btn.disabled = ( + self.selected_operation == "toggle_enforcement" + ) + other_policy_btn.disabled = ( + self.selected_operation == "other_policy" + ) + else: + # Enable all buttons + local_approval_btn.disabled = False + toggle_enforcement_btn.disabled = False + other_policy_btn.disabled = False + + except NoMatches: + pass + + def _display_results( + self, operation_name: str, successful: list, unsuccessful: list + ) -> None: + """ + Display operation results in the results text area. + + Formats the results into a human-readable summary including: + - Operation name and separator + - List of successful operations with agent hostnames + - List of failed operations with agent hostnames and error messages + - Summary statistics (total successful/failed count) + + The results are displayed in the results_text TextArea widget and the + results container is made visible after being initially hidden. + + Args: + operation_name (str): Human-readable name of the operation (e.g., "Local Approval Mode"). + successful (list): List of (Agent, result_data) tuples for successful operations. + unsuccessful (list): List of (Agent, error_message) tuples for failed operations. + """ + try: + # Build results text + results_lines = [ + f"Operation: {operation_name}", + f"{'=' * 50}", + "", + f"Òœ… Successful ({len(successful)}):", + ] + + if successful: + for agent, result in successful: + results_lines.append(f" Ò€’ {agent.hostname}") + else: + results_lines.append(" (none)") + + results_lines.append("") + results_lines.append(f"ҝŒ Failed ({len(unsuccessful)}):") + + if unsuccessful: + for agent, error in unsuccessful: + results_lines.append(f" Ò€’ {agent.hostname}: {error}") + else: + results_lines.append(" (none)") + + results_lines.append("") + results_lines.append(f"{'=' * 50}") + results_lines.append( + f"Total: {len(successful)} successful, {len(unsuccessful)} failed" + ) + + results_text_widget = self.query_one("#results_text", TextArea) + results_text_widget.text = "\n".join(results_lines) + + # Show results container + results_container = self.query_one("#results_container", Vertical) + results_container.styles.display = "block" + + except Exception as e: + logger.error(f"Error displaying results: {e}") + + def compose(self): + """ + Build the UI layout for the AgentMoveOperations widget. + + This method is called by Textual to create the widget's UI structure. + It builds a two-column layout with: + - Left side: Agent table showing selected agents and their current policies + - Right side: Operation buttons and results display area + - Bottom: Navigation buttons (Back, Reset) + + The layout is responsive with: + - Agent table: 2/3 width + - Operations panel: 1/3 width + - Results area: Initially hidden, shown after operation completion + """ + yield Header(show_clock=True, icon="βš™") + title_text = Static( + f"↔️ Move Agent Operations - {len(self.agents)} device(s) selected", + id="move_ops_title", + ) + title_text.styles.margin = (0, 0, 1, 0) + yield title_text + + with Horizontal() as main_layout: + main_layout.styles.height = "auto" + + # Left side - Agent list + with Vertical() as left_side: + left_side.styles.width = "2fr" + left_side.styles.height = "auto" + + agents_label = Static("Selected Agents:") + agents_label.styles.margin = (0, 0, 0, 0) + yield agents_label + + # Create a DataTable to show agents with their current policies + agent_table = DataTable(id="agent_table") + agent_table.styles.height = "1fr" + agent_table.styles.margin = (1, 0, 1, 0) + yield agent_table + + # Right side - Operation buttons + with Vertical() as right_side: + right_side.styles.width = "1fr" + right_side.styles.height = "auto" + + operations_label = Static("Operations:") + operations_label.styles.margin = (0, 0, 1, 0) + yield operations_label + + # Operation buttons + local_approval_btn = Button( + "βœ… Local Approval Mode", id="local_approval_btn" + ) + local_approval_btn.styles.width = "100%" + local_approval_btn.styles.margin = (0, 0, 1, 0) + yield local_approval_btn + + toggle_enforcement_btn = Button( + "πŸ”„ Toggle Audit/Enforcement", id="toggle_enforcement_btn" + ) + toggle_enforcement_btn.styles.width = "100%" + toggle_enforcement_btn.styles.margin = (0, 0, 1, 0) + yield toggle_enforcement_btn + + other_policy_btn = Button( + "πŸ”€ Move to Other Policy", id="other_policy_btn" + ) + other_policy_btn.styles.width = "100%" + other_policy_btn.styles.margin = (0, 0, 1, 0) + yield other_policy_btn + + # Status label + status_label = Static("", id="status_label") + status_label.styles.margin = (2, 0, 0, 0) + yield status_label + + # Results display area (initially hidden) + with Vertical(id="results_container") as results_container: + results_container.styles.height = "auto" + results_container.styles.margin = (1, 0, 0, 0) + results_container.styles.display = "none" + + results_label = Static("Γ°ΕΈβ€œΕ  Results:", id="results_label") + results_label.styles.margin = (0, 0, 0, 0) + yield results_label + + results_text = TextArea(id="results_text", read_only=True) + results_text.styles.height = 15 + results_text.styles.margin = (0, 0, 1, 0) + yield results_text + + copy_results_btn = Button( + "Γ°ΕΈβ€œβ€Ή Copy Results to Clipboard", id="copy_results_btn" + ) + copy_results_btn.styles.width = "100%" + yield copy_results_btn + + # Bottom buttons + with Horizontal() as button_row: + button_row.styles.height = "auto" + button_row.styles.margin = (1, 0, 0, 0) + + back_button = Button("Ò† Back", id="back_button") + back_button.styles.width = "1fr" + yield back_button + + reset_button = Button("Γ°ΕΈβ€β€ž Reset Selection", id="reset_button") + reset_button.styles.width = "1fr" + yield reset_button + + def on_mount(self) -> None: + """ + Initialize widget after it has been mounted on the screen. + + This Textual lifecycle method is called after the widget is added to the DOM. + It performs initialization tasks: + - Populates the agent table with columns for Hostname, Policy, and Status + - Adds rows to the table for each agent in self.agents + - Initializes button states based on current widget state + + The agent table displays agent.hostname, agent.groupname (or "Unknown"), + and agent.status_text (or "Unknown") for each agent. + """ + table = self.query_one("#agent_table", DataTable) + table.add_columns("Hostname", "Current Policy", "Status") + + for agent in self.agents: + table.add_row( + agent.hostname, + agent.groupname or "Unknown", + agent.status_text or "Unknown", + ) + + self._update_button_states() + + def on_button_pressed(self, event: Button.Pressed): + """ + Handle button press events from the widget. + + This Textual event handler routes button presses to appropriate actions: + - back_button: Pop this screen (return to parent) + - reset_button: Clear operation state and hide results + - copy_results_btn: Copy results text to clipboard (requires pyperclip) + - local_approval_btn: Start local approval operation + - toggle_enforcement_btn: Start toggle audit/enforcement operation + - other_policy_btn: Start move to other policy operation + + After handling, event.stop() is called to prevent event propagation. + + Args: + event (Button.Pressed): The button press event containing the button reference. + """ + + btn_id = event.button.id + + if btn_id == "back_button": + self.app.pop_screen() + event.stop() + + elif btn_id == "reset_button": + # Reset operation selection + self.selected_operation = "" + self.operation_in_progress = False + status_label = self.query_one("#status_label", Static) + status_label.update("") + # Hide results + try: + results_container = self.query_one("#results_container", Vertical) + results_container.styles.display = "none" + except NoMatches: + pass + event.stop() + + elif btn_id == "copy_results_btn": + try: + results_text = self.query_one("#results_text", TextArea) + import pyperclip + + pyperclip.copy(results_text.text) + self.app.notify( + "Òœ… Results copied to clipboard!", + severity="information", + timeout=2, + ) + except ImportError: + self.app.notify( + "Òő ï¸ pyperclip not installed. Run: pip install pyperclip", + severity="warning", + ) + except Exception as e: + self.app.notify(f"ҝŒ Failed to copy: {str(e)}", severity="error") + event.stop() + + elif btn_id == "local_approval_btn": + self._start_local_approval_operation() + event.stop() + + elif btn_id == "toggle_enforcement_btn": + self._start_toggle_enforcement_operation() + event.stop() + + elif btn_id == "other_policy_btn": + self._start_other_policy_operation() + event.stop() + + def _start_local_approval_operation(self) -> None: + """ + Execute the local approval mode operation on all selected agents. + + This operation performs the following steps for each agent: + 1. Generate a unique batch ID (current Unix timestamp) + 2. Create a local approval OTP with default duration of 360 minutes (6 hours) + 3. Move the agent to its related audit policy mode + + The operation: + - Sets operation state flags (selected_operation, operation_in_progress) + - Updates the status label with progress indicator + - Iterates through all agents, tracking successful and unsuccessful operations + - Displays formatted results via _display_results() + - Posts an OperationComplete message for parent widget handling + + Agents that fail are logged and added to the unsuccessful list with error details. + The operation completes and returns to a non-busy state regardless of individual + agent success/failure. + + Note: The OTP duration (360 minutes) is currently hardcoded and could be + made configurable in future versions. + """ + self.selected_operation = "local_approval" + self.operation_in_progress = True + + status_label = self.query_one("#status_label", Static) + status_label.update("ҏ³ Moving agents to local approval...") + + # Get API from app + api = self.app.api + + successful = [] + unsuccessful = [] + + try: + import time + + from services.agenthandler import moveAgentToRelatedPolicy + + # Generate batch ID + batch = int(time.time()) + duration = 360 # Default 6 hours, could make this configurable + + for agent in self.agents: + try: + # Add local approval OTP + addLocalApproval(api, batch, duration, agent.agentid) + # Move to audit mode + result = moveAgentToRelatedPolicy(api, agent, "audit") + successful.append((agent, result)) + logger.info( + f"Successfully moved {agent.hostname} to local approval" + ) + except Exception as e: + unsuccessful.append((agent, str(e))) + logger.error( + f"Failed to move {agent.hostname} to local approval: {e}" + ) + + except Exception as e: + logger.error(f"Error during local approval operation: {e}") + status_label.update(f"ҝŒ Error: {str(e)}") + self.operation_in_progress = False + return + + self.operation_in_progress = False + status_label.update("Òœ… Operation complete!") + + # Display results in the widget + self._display_results("Local Approval Mode", successful, unsuccessful) + + # Also post message for potential parent handling + self.post_message( + self.OperationComplete( + "Local Approval Mode", self.agents, successful, unsuccessful + ) + ) + + def _start_toggle_enforcement_operation(self) -> None: + """ + Toggle agents between enforcement and audit policy modes. + + This operation intelligently switches each agent between enforcement and + audit modes based on its current state: + - If agent.groupid is in POLICY_MAP_ENF_AUD: currently enforcing Ò†' move to audit + - Otherwise: currently in audit Ò†' move to enforcement + + The operation: + - Retrieves the enforcement/audit policy relationship map from protected config + - Sets operation state flags and updates status label + - Iterates through agents, determining current mode and toggling to opposite + - Tracks successful toggles with the new mode in the result message + - Logs both successes and failures + - Displays results and posts OperationComplete message + + The policy relationship map (POLICY_MAP_ENF_AUD) must be present in protected + configuration and maps enforcement policy IDs to audit policy IDs. If the map + is empty or not found, all agents are assumed to be in audit mode and will + be moved to enforcement. + + Returns to a non-busy state after completion regardless of individual results. + """ + self.selected_operation = "toggle_enforcement" + self.operation_in_progress = True + + status_label = self.query_one("#status_label", Static) + status_label.update("ҏ³ Toggling enforcement mode...") + + # Get API from app + api = self.app.api + + successful = [] + unsuccessful = [] + + try: + from services.agenthandler import moveAgentToRelatedPolicy + from utils.configmanager import get_protected_json + + policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") + + for agent in self.agents: + try: + # Determine current mode and toggle + if agent.groupid in policy_relationship_map: + # Currently in enforcement, move to audit + result = moveAgentToRelatedPolicy(api, agent, "audit") + mode = "audit" + else: + # Currently in audit, move to enforcement + result = moveAgentToRelatedPolicy(api, agent, "enforcement") + mode = "enforcement" + + successful.append((agent, f"Moved to {mode}: {result}")) + logger.info(f"Successfully toggled {agent.hostname} to {mode}") + except Exception as e: + unsuccessful.append((agent, str(e))) + logger.error(f"Failed to toggle {agent.hostname}: {e}") + + except Exception as e: + logger.error(f"Error during toggle enforcement operation: {e}") + status_label.update(f"ҝŒ Error: {str(e)}") + self.operation_in_progress = False + return + + self.operation_in_progress = False + status_label.update("Òœ… Operation complete!") + + # Display results in the widget + self._display_results("Toggle Audit/Enforcement", successful, unsuccessful) + + # Also post message for potential parent handling + self.post_message( + self.OperationComplete( + "Toggle Audit/Enforcement", self.agents, successful, unsuccessful + ) + ) + + def _start_other_policy_operation(self) -> None: + """ + Move agents to a user-selected policy (currently unimplemented). + + This operation is intended to allow bulk movement of selected agents to any + alternative policy via a policy selection dialog. Currently, this feature + is not fully implemented. + + Planned Implementation: + 1. Push a new policy selector screen (TUI modal/overlay) + 2. Allow user to choose target policy from available options + 3. Move all selected agents to the chosen policy + 4. Display results like other operations + + Current Behavior: + - Sets selected_operation to "other_policy" + - Displays "Policy selection not yet implemented" status message + - Clears selected_operation without performing any action + + TODO: Complete implementation by: + - Creating a policy selector screen component + - Implementing the policy selection logic + - Integrating with moveAgentToPolicy API call + - Adding proper result tracking and display + """ + self.selected_operation = "other_policy" + self.operation_in_progress = True + + status_label = self.query_one("#status_label", Static) + status_label.update("Loading available policies...") + + try: + # Fetch all policies from API + api = self.app.api + + # Fetch all available policies + all_policies_df = api.policy_find_all() + + if all_policies_df.empty: + status_label.update("No policies available") + self.operation_in_progress = False + self.selected_operation = "" + return + + # Create and push the policy selector screen + policy_selector_screen = PolicySelectorScreen( + policies=all_policies_df, + agent_move_operations=self, + ) + self.app.push_screen(policy_selector_screen) + + except Exception as e: + logger.error(f"Error loading policies: {e}") + status_label.update(f"Error: {str(e)}") + self.operation_in_progress = False + self.selected_operation = "" + self.app.notify(f"Failed to load policies: {str(e)}", severity="error") + + def _execute_move_to_policy(self, target_policy) -> None: + """ + Execute the actual move of agents to the selected policy. + + Moves each agent sequentially to the target policy, tracking success/failure. + Updates the status label and displays results upon completion. + + Args: + target_policy: The Policy object selected by the user. + """ + status_label = self.query_one("#status_label", Static) + status_label.update(f"Moving agents to {target_policy.name}...") + + api = self.app.api + successful = [] + unsuccessful = [] + + try: + for agent in self.agents: + try: + # Move agent to target policy + result = api.agent_move(agent.agentid, target_policy.groupid) + successful.append((agent, f"Moved to {target_policy.name}")) + logger.info( + f"Successfully moved {agent.hostname} to policy {target_policy.name}" + ) + except Exception as e: + unsuccessful.append((agent, str(e))) + logger.error( + f"Failed to move {agent.hostname} to policy {target_policy.name}: {e}" + ) + + except Exception as e: + logger.error(f"Error during move to policy operation: {e}") + status_label.update(f"Error: {str(e)}") + self.operation_in_progress = False + return + + self.operation_in_progress = False + status_label.update("Operation complete!") + + # Display results in the widget + self._display_results( + f"Move to {target_policy.name}", + successful, + unsuccessful, + ) + + # Also post message for potential parent handling + self.post_message( + self.OperationComplete( + f"Move to {target_policy.name}", + self.agents, + successful, + unsuccessful, + ) + ) diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py index e46c636..dd0fe1b 100644 --- a/widgets/multiagentselector.py +++ b/widgets/multiagentselector.py @@ -6,7 +6,15 @@ from textual.containers import Horizontal, Vertical from textual.css.query import NoMatches from textual.message import Message from textual.widget import Widget -from textual.widgets import Button, SelectionList, Static, Switch, TextArea +from textual.widgets import ( + Button, + Footer, + Header, + SelectionList, + Static, + Switch, + TextArea, +) from models.agent import Agent @@ -31,6 +39,7 @@ class MultiAgentSelector(Widget): self._match_type = value def compose(self): + yield Header(show_clock=True, icon="βš™") title_text = Static("πŸ–§ Multi-Agent Selector", id="selector_title") title_text.styles.margin = (0, 0, 0, 1) yield title_text @@ -98,6 +107,7 @@ class MultiAgentSelector(Widget): right_pane.styles.width = "2fr" yield SelectionList(id="match_results") yield Static(id="unmatched_label") + yield Footer() def on_switch_changed(self, event: Switch.Changed): self.match_type = "fuzzy" if event.value else "exact" diff --git a/widgets/policyselector.py b/widgets/policyselector.py new file mode 100644 index 0000000..1796968 --- /dev/null +++ b/widgets/policyselector.py @@ -0,0 +1,505 @@ +""" +Policy Selector Widget Module + +Provides a Textual widget for selecting target policies for bulk agent operations. +Allows users to browse available policies and select one as the destination for +moving agents. Automatically excludes parent/logical policies. +""" + +import logging +import re +from typing import Optional + +import pandas as pd +from textual.containers import Horizontal, Vertical +from textual.message import Message +from textual.widget import Widget +from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea + +from models.policy import Policy + +logger = logging.getLogger(__name__) + + +class PolicySelector(Widget): + """ + A Textual widget for selecting a target policy for agent operations. + + This widget displays available policies in a table and allows users to select + one policy as the destination for bulk agent movements. It automatically excludes: + - Parent/logical policies (where parent == "global-policy-settings") + - Specified policy IDs (e.g., the current policy) + + Features: + - Wildcard filtering (* and ?) + - Interactive table for policy browsing + - Explicit confirm button for selection + - Cancel/back button to dismiss + + Attributes: + policies (list[Policy]): List of available Policy objects to display. + excluded_policy_ids (set[str]): Set of policy IDs to exclude from selection. + selected_policy (Optional[Policy]): The currently selected policy (if any). + + Automatically Filtered Out: + - Policies with parent == "global-policy-settings" (parent policies for organization) + - Any policies in excluded_policy_ids set + + Example: + ```python + policies = [policy1, policy2, policy3] + widget = PolicySelector(policies, excluded_policy_ids={current_policy.groupid}) + ``` + """ + + class PolicySelected(Message): + """ + Message posted when a policy is selected. + + Attributes: + policy (Policy): The selected policy object. + """ + + def __init__(self, policy: Policy): + super().__init__() + self.policy = policy + + def __init__(self, policies: list): + """ + Initialize the PolicySelector widget. + + Args: + policies (list): List of Policy objects or DataFrame rows to display. + Can be a list of Policy objects or a pandas DataFrame of policy data. + """ + super().__init__() + self.policies = policies + self.selected_policy: Optional[Policy] = None + self._filtered_policies = [] + self._displayed_policies = [] # Track what's currently shown in the table + self._filter_text = "" + + def compose(self): + """ + Build the UI layout for the PolicySelector widget. + + The layout includes: + - Title indicating policy selection + - Search/filter text area with wildcard support + - Filter help text showing wildcard options + - Apply Filter button + - Clear Filter button + - Confirm Selection button + - Policy table displaying available policies + - Back and Continue buttons for navigation + """ + yield Header(show_clock=True, icon="βš™") + title_text = Static( + "🎯 Select Target Policy", + id="policy_selector_title", + ) + title_text.styles.margin = (0, 0, 1, 0) + yield title_text + + with Horizontal() as main_layout: + main_layout.styles.height = "auto" + + # Left side - Filter and controls + with Vertical() as left_side: + left_side.styles.width = "1fr" + left_side.styles.height = "auto" + + filter_label = Static("Filter Policies:") + filter_label.styles.margin = (0, 0, 0, 0) + yield filter_label + + filter_input = TextArea( + id="policy_filter", + text="", + ) + filter_input.styles.height = 3 + filter_input.styles.margin = (0, 0, 1, 0) + yield filter_input + + filter_help = Static("(Use * and ? for wildcards)", id="filter_help") + filter_help.styles.margin = (0, 0, 1, 0) + yield filter_help + + apply_button = Button("βœ“ Apply Filter", id="filter_button") + apply_button.styles.width = "100%" + apply_button.styles.margin = (0, 0, 1, 0) + yield apply_button + + clear_button = Button("πŸ—‘οΈ Clear Filter", id="clear_filter_button") + clear_button.styles.width = "100%" + clear_button.styles.margin = (0, 0, 1, 0) + yield clear_button + + confirm_button = Button("βœ… Confirm Selection", id="confirm_button") + confirm_button.styles.width = "100%" + confirm_button.styles.margin = (1, 0, 1, 0) + yield confirm_button + + selected_label = Static("", id="selected_policy_label") + selected_label.styles.margin = (2, 0, 1, 0) + yield selected_label + + # Right side - Policy table + with Vertical() as right_side: + right_side.styles.width = "2fr" + right_side.styles.height = "auto" + + table_label = Static("Available Policies:") + table_label.styles.margin = (0, 0, 0, 0) + yield table_label + + policy_table = DataTable(id="policy_table", cursor_type="row") + policy_table.styles.height = "1fr" + policy_table.styles.margin = (1, 0, 1, 0) + yield policy_table + + # Bottom buttons + with Horizontal() as button_row: + button_row.styles.height = "auto" + button_row.styles.margin = (1, 0, 0, 0) + + cancel_button = Button("βœ• Cancel", id="back_button", variant="error") + cancel_button.styles.width = "1fr" + yield cancel_button + + continue_button = Button( + "β–Ά Continue", + id="continue_button", + variant="primary", + ) + continue_button.styles.width = "1fr" + continue_button.styles.margin = (0, 0, 0, 1) + yield continue_button + yield Footer() + + def on_mount(self) -> None: + """ + Initialize the policy table when the widget is mounted. + + Populates the table with column (Policy Name) and rows for each + available policy (excluding those in excluded_policy_ids and parent policies). + Sets up event handlers for table row selection. + + Filters out: + - Parent policies (where parent == "global-policy-settings") + """ + table = self.query_one("#policy_table", DataTable) + + # Configure table for row selection + table.cursor_type = "row" + table.zebra_stripes = True + + # Only add Policy Name column + table.add_columns("Policy Name") + + # Filter out excluded policies and convert to list if DataFrame + if isinstance(self.policies, pd.DataFrame): + policies_list = self.policies.to_dict("records") + else: + policies_list = self.policies + + self._filtered_policies = [] + self._displayed_policies = [] # Initialize displayed list + + for policy_data in policies_list: + # Handle both Policy objects and dict/DataFrame rows + if isinstance(policy_data, Policy): + policy_id = policy_data.groupid + policy_name = policy_data.name + parent = policy_data.parent + else: + policy_id = policy_data.get("groupid", "Unknown") + policy_name = policy_data.get("name", "Unknown") + parent = policy_data.get("parent", None) + + # Skip parent policies (logical policies that shouldn't have devices) + if parent == "global-policy-settings": + logger.debug(f"Skipping parent policy: {policy_name}") + continue + + self._filtered_policies.append(policy_data) + self._displayed_policies.append(policy_data) # Add to displayed list + + table.add_row( + policy_name, + key=policy_id, + ) + + def on_button_pressed(self, event: Button.Pressed): + """ + Handle button press events from the widget. + + Routes to: + - back_button (Cancel): Pop screen without selecting + - filter_button (Apply Filter): Filter policies with wildcard support + - clear_filter_button: Clear filter and show all policies + - confirm_button: Confirm selection and post message + - continue_button: Continue without posting message + + Args: + event (Button.Pressed): The button press event. + """ + btn_id = event.button.id + + if btn_id == "back_button": + self.app.pop_screen() + event.stop() + + elif btn_id == "filter_button": + self._apply_filter() + event.stop() + + elif btn_id == "clear_filter_button": + self._clear_filter() + event.stop() + + elif btn_id == "confirm_button": + self._confirm_selection() + event.stop() + + elif btn_id == "continue_button": + self.app.pop_screen() + event.stop() + + def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None: + """ + Handle row selection in the policy table. + + Updates the selected_policy and displays the selection in the UI. + + Args: + event: DataTable.RowSelected event containing the selected row data. + """ + try: + # Get the row key from the event + row_key = event.row_key + if row_key is None: + return + + # Find the policy with matching groupid + for policy_data in self._displayed_policies: + if isinstance(policy_data, Policy): + if policy_data.groupid == row_key.value: + self.selected_policy = policy_data + break + else: + if policy_data.get("groupid") == row_key.value: + self.selected_policy = Policy( + groupid=policy_data.get("groupid"), + hidden=policy_data.get("hidden", False), + name=policy_data.get("name"), + parent=policy_data.get("parent"), + ) + break + + if self.selected_policy: + # Update selection display + label = self.query_one("#selected_policy_label", Static) + label.update(f"βœ“ Selected: {self.selected_policy.name}") + + # Log for debugging + logger.debug( + f"Selected policy: {self.selected_policy.name} (ID: {self.selected_policy.groupid})" + ) + self.app.notify( + f"Selected: {self.selected_policy.name}", + severity="information", + timeout=1, + ) + + except Exception as e: + logger.error(f"Error handling row selection: {e}") + self.app.notify(f"Selection error: {str(e)}", severity="error") + + def on_data_table_row_highlighted(self, event: DataTable.RowHighlighted) -> None: + """ + Handle row highlighting (cursor movement) in the table. + + This provides immediate visual feedback when navigating rows. + """ + try: + # Get the row key from the event + row_key = event.row_key + if row_key is None: + return + + # Find the highlighted policy + highlighted_name = None + for policy_data in self._displayed_policies: + if isinstance(policy_data, Policy): + if policy_data.groupid == row_key.value: + highlighted_name = policy_data.name + break + else: + if policy_data.get("groupid") == row_key.value: + highlighted_name = policy_data.get("name") + break + + if highlighted_name: + label = self.query_one("#selected_policy_label", Static) + label.update(f"β†’ Highlighting: {highlighted_name}") + + except Exception as e: + logger.error(f"Error handling row highlight: {e}") + + def _apply_filter(self) -> None: + """ + Apply filter text to policy list with wildcard support. + + Supports wildcards: + - * matches any sequence of characters + - ? matches a single character + + Examples: + - "policy*" matches "policy_prod", "policy_dev", etc. + - "policy?" matches "policy1", "policy2", etc. + - "*audit*" matches anything containing "audit" + - "*test*" matches "AT Testing", "test_policy", etc. + + Filters policies by name or ID (case-insensitive) and refreshes the table display + with only matching policies. Only filters from already-filtered list + (which excludes parent policies and excluded IDs). + """ + try: + filter_input = self.query_one("#policy_filter", TextArea) + filter_text = filter_input.text.strip() + + table = self.query_one("#policy_table", DataTable) + table.clear() + + # Clear the displayed policies list + self._displayed_policies = [] + + # Compile wildcard pattern if filter text is provided + pattern = None + if filter_text: + # Escape special regex chars but preserve wildcards + pattern_text = re.escape(filter_text.lower()) + pattern_text = pattern_text.replace(r"\*", ".*").replace(r"\?", ".") + # Use search() for partial matching + pattern = re.compile(pattern_text, re.IGNORECASE) + + # Filter policies based on search text + for policy_data in self._filtered_policies: + # Handle both Policy objects and dict/DataFrame rows + if isinstance(policy_data, Policy): + policy_name = policy_data.name.lower() + policy_id = policy_data.groupid.lower() + display_name = policy_data.name + key_id = policy_data.groupid + else: + policy_name = str(policy_data.get("name", "")).lower() + policy_id = str(policy_data.get("groupid", "Unknown")).lower() + display_name = policy_data.get("name") + key_id = policy_data.get("groupid") + + # Match against filter text with wildcard support + if pattern: + # Use search() for partial matching + matches = pattern.search(policy_name) or pattern.search(policy_id) + else: + matches = True + + if matches: + # Add to displayed policies list + self._displayed_policies.append(policy_data) + + # Add row to table + table.add_row( + display_name, + key=key_id, + ) + + displayed_count = len(self._displayed_policies) + status_text = f"πŸ“Š Showing {displayed_count} of {len(self._filtered_policies)} policies" + self.app.notify(status_text, severity="information", timeout=2) + + # Clear selection when filter is applied + self.selected_policy = None + label = self.query_one("#selected_policy_label", Static) + label.update("") + + except Exception as e: + logger.error(f"Error applying filter: {e}") + self.app.notify(f"❌ Filter error: {str(e)}", severity="error") + + def _clear_filter(self) -> None: + """ + Clear the filter and display all available policies. + + Resets the filter text and refreshes the table to show all policies + (already excluding parent policies and excluded IDs). + """ + try: + filter_input = self.query_one("#policy_filter", TextArea) + filter_input.text = "" + + table = self.query_one("#policy_table", DataTable) + table.clear() + + # Reset displayed policies to all filtered policies + self._displayed_policies = list(self._filtered_policies) + + # Reload all policies + for policy_data in self._filtered_policies: + if isinstance(policy_data, Policy): + policy_id = policy_data.groupid + policy_name = policy_data.name + else: + policy_id = policy_data.get("groupid", "Unknown") + policy_name = policy_data.get("name", "Unknown") + + # Add row with only policy name + table.add_row( + policy_name, + key=policy_id, + ) + + self.selected_policy = None + label = self.query_one("#selected_policy_label", Static) + label.update("") + + except Exception as e: + logger.error(f"Error clearing filter: {e}") + + def on_text_area_changed(self, event) -> None: + """ + Handle TextArea change events - specifically for Enter key in filter. + + When the user types in the filter TextArea and the text ends with a newline, + treat it as pressing Enter and apply the filter. + """ + if event.text_area.id == "policy_filter": + # Check if the text ends with a newline (Enter was pressed) + if event.text_area.text.endswith("\n"): + # Remove the newline that was added + event.text_area.text = event.text_area.text.rstrip("\n") + # Apply the filter + self._apply_filter() + + def _confirm_selection(self) -> None: + """ + Confirm the selected policy and post selection message. + + Posts a PolicySelected message to the parent widget/screen with the + selected policy. If no policy is selected, displays an error notification. + """ + if self.selected_policy is None: + self.app.notify( + "⚠️ Please select a policy first by clicking on a row in the table", + severity="warning", + timeout=3, + ) + return + + # Log confirmation for debugging + logger.info(f"Confirming selection of policy: {self.selected_policy.name}") + self.app.notify( + f"βœ… Confirmed: {self.selected_policy.name}", severity="success", timeout=2 + ) + self.post_message(self.PolicySelected(self.selected_policy)) diff --git a/widgets/resultsdisplay.py b/widgets/resultsdisplay.py new file mode 100644 index 0000000..ebabce7 --- /dev/null +++ b/widgets/resultsdisplay.py @@ -0,0 +1,178 @@ +import logging + +from textual.containers import Horizontal, Vertical +from textual.message import Message +from textual.widget import Widget +from textual.widgets import Button, Footer, Header, Static + +logger = logging.getLogger(__name__) + + +class ResultsDisplay(Widget): + """Widget for displaying operation results in a two-column layout.""" + + CSS = """ + ResultsDisplay { + height: 100%; + } + + #results_screen { + height: 100%; + } + + #results_title { + text-align: center; + margin: 1 0; + text-style: bold; + } + + #results_layout { + height: 1fr; + margin: 1 0; + } + + #left_column, #right_column { + width: 1fr; + height: 100%; + border: solid green; + padding: 1; + } + + #right_column { + border: solid red; + } + + #success_label, #failure_label { + text-style: bold; + margin-bottom: 1; + } + + #success_results, #failure_results { + height: 1fr; + overflow-y: auto; + background: $surface; + border: round $primary; + padding: 1; + } + + .copy_button { + margin-top: 1; + width: 100%; + } + + #button_row { + height: auto; + margin: 1 0 0 0; + } + + #back_button { + width: 1fr; + } + """ + + class CopySuccess(Message): + """Posted when success results are copied.""" + + pass + + class CopyFailure(Message): + """Posted when failure results are copied.""" + + pass + + class GoBack(Message): + """Posted when back button is pressed.""" + + pass + + def __init__( + self, operation: str, successful_results: str, unsuccessful_results: str + ) -> None: + super().__init__() + self.operation = operation + self.successful_results = successful_results + self.unsuccessful_results = unsuccessful_results + + def compose(self): + with Vertical(id="results_screen"): + yield Header(show_clock=True, icon="βš™") + # Title + title = Static(f"πŸ“Š {self.operation} - Results", id="results_title") + yield title + + # Two-column layout + with Horizontal(id="results_layout"): + # Left Column - Success + with Vertical(id="left_column"): + yield Static("βœ… Successful", id="success_label") + yield Static(self.successful_results, id="success_results") + yield Button( + "πŸ“‹βœ… Copy Success List", + id="copy_success", + classes="copy_button", + ) + + # Right Column - Failure + with Vertical(id="right_column"): + yield Static("❌ Failed", id="failure_label") + yield Static(self.unsuccessful_results, id="failure_results") + yield Button( + "πŸ“‹βŒ Copy Failure List", + id="copy_failure", + classes="copy_button", + ) + + # Back Button + with Horizontal(id="button_row"): + back_button = Button("← Back", id="back_button") + yield back_button + yield Footer() + + def on_button_pressed(self, event: Button.Pressed) -> None: + btn_id = event.button.id + + if btn_id == "copy_success": + success_widget = self.query_one("#success_results", Static) + try: + import pyperclip + + pyperclip.copy(str(success_widget.renderable)) + self.app.notify( + "Òœ… Success list copied to clipboard!", + severity="information", + timeout=2, + ) + self.post_message(self.CopySuccess()) + except ImportError: + self.app.notify( + "Òő ï¸ pyperclip not installed. Run: pip install pyperclip", + severity="warning", + ) + except Exception as e: + self.app.notify(f"ҝŒ Failed to copy: {str(e)}", severity="error") + event.stop() + + elif btn_id == "copy_failure": + failure_widget = self.query_one("#failure_results", Static) + try: + import pyperclip + + pyperclip.copy(str(failure_widget.renderable)) + self.app.notify( + "Òœ… Failure list copied to clipboard!", + severity="information", + timeout=2, + ) + self.post_message(self.CopyFailure()) + except ImportError: + self.app.notify( + "Òő ï¸ pyperclip not installed. Run: pip install pyperclip", + severity="warning", + ) + except Exception as e: + self.app.notify(f"ҝŒ Failed to copy: {str(e)}", severity="error") + event.stop() + + elif btn_id == "back_button": + self.app.pop_screen() + event.stop() diff --git a/widgets/retro_terminal_theme.py b/widgets/retro_terminal_theme.py new file mode 100644 index 0000000..3ff0d02 --- /dev/null +++ b/widgets/retro_terminal_theme.py @@ -0,0 +1,38 @@ +from textual.color import Color + + +def get_retro_terminal_theme(): + from textual.theme import Theme + + return Theme( + name="retro-terminal", + background=Color.parse("#000000"), + primary=Color.parse("#00ff00"), + secondary=Color.parse("#00aa00"), + success=Color.parse("#00ff00"), + warning=Color.parse("#ffff00"), + error=Color.parse("#ff0000"), + surface=Color.parse("#111111"), + ) + + +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/widgets/themeselector.py b/widgets/themeselector.py index fda7cd1..6a423c2 100644 --- a/widgets/themeselector.py +++ b/widgets/themeselector.py @@ -15,26 +15,25 @@ class ThemeSelector(Widget): self.theme_name = theme_name AVAILABLE_THEMES = [ - ("textual-dark", "textual-dark"), - ("textual-light", "textual-light"), - ("nord", "nord"), - ("gruvbox", "gruvbox"), - ("catppuccin-mocha", "catppuccin-mocha"), - ("dracula", "dracula"), - ("tokyo-night", "tokyo-night"), - ("monokai", "monokai"), - ("flexoki", "flexoki"), - ("catppuccin-latte", "catppuccin-latte"), - ("solarized-light", "solarized-light"), + ("Textual Dark", "textual-dark"), + ("Textual Light", "textual-light"), + ("Nord", "nord"), + ("Gruvbox", "gruvbox"), + ("Catppuccin Mocha", "catppuccin-mocha"), + ("Dracula", "dracula"), + ("Tokyo Night", "tokyo-night"), + ("Monokai", "monokai"), + ("Flexoki", "flexoki"), + ("Catppuccin Latte", "catppuccin-latte"), + ("Solarized Light", "solarized-light"), + ("Retro Terminal", "retro-terminal"), # your custom theme ] def compose(self): yield Static("Theme Options", id="theme_title") - with Vertical() as column: column.styles.width = "1fr" column.styles.height = "auto" - for label, btn_id in self.AVAILABLE_THEMES: yield Button(label, id=f"set_theme_{btn_id}", compact=True)