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