kuma-strapper/backend/services/discovery.py

from dataclasses import dataclass
from typing import Optional, Callable
import uuid

from services.ssh_manager import get_ssh_manager, CommandResult


# Built-in safe commands that never require approval
SAFE_DISCOVERY_COMMANDS = {
    "system_info": "uname -a",
    "os_release": "cat /etc/os-release 2>/dev/null || cat /etc/*-release 2>/dev/null | head -20",
    "docker_containers": "docker ps --format '{{.ID}}\\t{{.Names}}\\t{{.Image}}\\t{{.Status}}\\t{{.Ports}}' 2>/dev/null || echo 'Docker not available'",
    "systemd_services": "systemctl list-units --type=service --state=running --no-pager 2>/dev/null | head -50 || echo 'Systemd not available'",
    "disk_usage": "df -h 2>/dev/null | head -20",
    "memory_usage": "free -h 2>/dev/null || cat /proc/meminfo 2>/dev/null | head -10",
    "cpu_count": "nproc 2>/dev/null || grep -c processor /proc/cpuinfo 2>/dev/null || echo 'Unknown'",
    "open_ports": "ss -tlnp 2>/dev/null || netstat -tlnp 2>/dev/null | head -30 || echo 'Unable to list ports'",
}


@dataclass
class DiscoveryResult:
    """Results from a host discovery scan."""

    scan_id: str
    hostname: str
    username: str
    port: int
    connected: bool
    error: Optional[str] = None
    system_info: str = ""
    os_release: str = ""
    docker_containers: str = ""
    systemd_services: str = ""
    disk_usage: str = ""
    memory_usage: str = ""
    cpu_count: str = ""
    open_ports: str = ""

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "scan_id": self.scan_id,
            "hostname": self.hostname,
            "username": self.username,
            "port": self.port,
            "connected": self.connected,
            "error": self.error,
            "system_info": self.system_info,
            "os_release": self.os_release,
            "docker_containers": self.docker_containers,
            "systemd_services": self.systemd_services,
            "disk_usage": self.disk_usage,
            "memory_usage": self.memory_usage,
            "cpu_count": self.cpu_count,
            "open_ports": self.open_ports,
        }


class DiscoveryService:
    """Service for discovering services on remote hosts."""

    def __init__(self):
        self.active_scans: dict[str, DiscoveryResult] = {}

    def is_safe_command(self, command: str) -> bool:
        """Check if a command is in the safe built-in list."""
        # Check exact matches
        if command in SAFE_DISCOVERY_COMMANDS.values():
            return True

        # Check if it's a safe read-only command pattern
        safe_patterns = [
            "cat ",
            "head ",
            "tail ",
            "grep ",
            "ls ",
            "ps ",
            "df ",
            "free ",
            "uname ",
            "uptime",
            "hostname",
            "whoami",
            "id ",
            "docker ps",
            "docker inspect",
            "docker logs",
            "systemctl status",
            "systemctl list-",
            "journalctl",
            "ss ",
            "netstat ",
            "curl -s",
            "wget -q",
            "nproc",
        ]

        command_lower = command.lower().strip()
        for pattern in safe_patterns:
            if command_lower.startswith(pattern):
                return True

        return False

    def scan_host(
        self,
        hostname: str,
        username: str = "root",
        port: int = 22,
        on_progress: Optional[Callable[[str, str], None]] = None,
    ) -> DiscoveryResult:
        """
        Scan a host for services and system information.

        Args:
            hostname: Target hostname or IP
            username: SSH username
            port: SSH port
            on_progress: Callback for progress updates (command_name, status)

        Returns:
            DiscoveryResult with all gathered information
        """
        scan_id = str(uuid.uuid4())
        result = DiscoveryResult(
            scan_id=scan_id,
            hostname=hostname,
            username=username,
            port=port,
            connected=False,
        )
        self.active_scans[scan_id] = result

        ssh = get_ssh_manager()

        # Connect to host
        if on_progress:
            on_progress("connect", "connecting")

        try:
            ssh.connect(hostname, username, port)
            result.connected = True
            if on_progress:
                on_progress("connect", "connected")
        except Exception as e:
            result.error = str(e)
            if on_progress:
                on_progress("connect", f"failed: {str(e)}")
            return result

        # Run discovery commands
        for cmd_name, command in SAFE_DISCOVERY_COMMANDS.items():
            if on_progress:
                on_progress(cmd_name, "running")

            try:
                cmd_result = ssh.execute(hostname, command, username, port, timeout=30)
                output = cmd_result.stdout if cmd_result.success else cmd_result.stderr
                setattr(result, cmd_name, output.strip())

                if on_progress:
                    on_progress(cmd_name, "complete" if cmd_result.success else "failed")
            except Exception as e:
                setattr(result, cmd_name, f"Error: {str(e)}")
                if on_progress:
                    on_progress(cmd_name, f"error: {str(e)}")

        return result

    def run_additional_command(
        self,
        hostname: str,
        command: str,
        username: str = "root",
        port: int = 22,
    ) -> CommandResult:
        """
        Run an additional command on a host.

        This should only be called after approval if in dev mode.
        """
        ssh = get_ssh_manager()

        if not ssh.is_connected(hostname, username, port):
            ssh.connect(hostname, username, port)

        return ssh.execute(hostname, command, username, port, timeout=60)

    def get_scan(self, scan_id: str) -> Optional[DiscoveryResult]:
        """Get a scan result by ID."""
        return self.active_scans.get(scan_id)


# Global discovery service instance
_discovery_service: Optional[DiscoveryService] = None


def get_discovery_service() -> DiscoveryService:
    """Get the global discovery service instance."""
    global _discovery_service
    if _discovery_service is None:
        _discovery_service = DiscoveryService()
    return _discovery_service