AirlockTools/flows/localApproval.py

"""
This module handles the creation of local approval requests.
"""

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("\n🎯 Select 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("\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"))