rmm2/agent-linux/ws/handler.go

package ws

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

type Handler struct {
	client *Client
}

func NewHandler(client *Client) *Handler {
	return &Handler{client: client}
}

func (h *Handler) HandleCommand(msg Message) {
	log.Printf("Command erhalten: %s (ID: %s)", msg.Command, msg.ID)

	var response Message
	response.Type = "response"
	response.ID = msg.ID

	switch msg.Command {
	case "exec":
		response = h.handleExec(msg)
	case "reboot":
		response = h.handleReboot(msg)
	case "tunnel_connect":
		response = h.handleTunnelConnect(msg)
	case "tunnel_disconnect":
		h.handleTunnelDisconnect(msg)
		return
	case "agent_update":
		response = h.handleAgentUpdate(msg)
	case "update_check":
		response = h.handleUpdateCheck(msg)
	case "update":
		response = h.handleUpdate(msg)
	case "zpoolscrub":
		response = h.handleZpoolScrub(msg)
	case "pty_start":
		response = h.handlePTYStart(msg)
	case "pty_stop":
		h.handlePTYStop(msg)
		return
	case "pty_resize":
		h.handlePTYResize(msg)
		return
	default:
		response.Status = "error"
		response.Error = fmt.Sprintf("Unbekanntes Command: %s", msg.Command)
	}

	if err := h.client.SendMessage(response); err != nil {
		log.Printf("Response senden fehlgeschlagen: %v", err)
	}
}

// handleExec fuehrt einen beliebigen Befehl aus
func (h *Handler) handleExec(msg Message) Message {
	response := Message{Type: "response", ID: msg.ID}

	command, ok := msg.Params["command"].(string)
	if !ok || command == "" {
		response.Status = "error"
		response.Error = "Parameter 'command' erforderlich"
		return response
	}

	timeoutSec := 30
	if t, ok := msg.Params["timeout"]; ok {
		if ts, ok := t.(float64); ok {
			timeoutSec = int(ts)
		} else if ts, ok := t.(string); ok {
			if parsed, err := strconv.Atoi(ts); err == nil {
				timeoutSec = parsed
			}
		}
	}

	output, err := h.runCommand(command, timeoutSec)

	if err != nil {
		response.Status = "error"
		response.Error = err.Error()
		response.Data = map[string]interface{}{"output": output}
	} else {
		response.Status = "ok"
		response.Data = map[string]interface{}{"output": output}
	}

	return response
}

// handleReboot fuehrt einen Reboot durch
func (h *Handler) handleReboot(msg Message) Message {
	response := Message{Type: "response", ID: msg.ID}

	delaySec := 5
	if d, ok := msg.Params["delay"].(float64); ok && d > 0 {
		delaySec = int(d)
	}

	log.Printf("Reboot angefordert (Verzoegerung: %ds)", delaySec)

	response.Status = "ok"
	response.Data = map[string]interface{}{
		"message": fmt.Sprintf("Reboot in %d Sekunden", delaySec),
	}

	go func() {
		time.Sleep(time.Duration(delaySec) * time.Second)
		log.Println("Reboot wird ausgefuehrt...")
		exec.Command("/sbin/shutdown", "-r", "now").Run()
	}()

	return response
}

// handleAgentUpdate aktualisiert den Agent
func (h *Handler) handleAgentUpdate(msg Message) Message {
	response := Message{Type: "response", ID: msg.ID}

	// Binary kommt als Base64 im Data-Feld
	data, ok := msg.Data.(map[string]interface{})
	if !ok {
		response.Status = "error"
		response.Error = "Ungueltige Daten"
		return response
	}

	binaryB64, _ := data["binary"].(string)
	expectedHash, _ := data["hash"].(string)
	newVersion, _ := data["version"].(string)

	if binaryB64 == "" || expectedHash == "" {
		response.Status = "error"
		response.Error = "Binary oder Hash fehlt"
		return response
	}

	log.Printf("Agent-Update empfangen: Version %s, Hash %s", newVersion, expectedHash[:16])

	// Base64 dekodieren
	binary, err := base64.StdEncoding.DecodeString(binaryB64)
	if err != nil {
		response.Status = "error"
		response.Error = fmt.Sprintf("Base64-Dekodierung fehlgeschlagen: %v", err)
		return response
	}

	// Hash pruefen
	hash := sha256.Sum256(binary)
	actualHash := hex.EncodeToString(hash[:])
	if actualHash != expectedHash {
		response.Status = "error"
		response.Error = fmt.Sprintf("Hash-Mismatch: erwartet %s, bekommen %s", expectedHash[:16], actualHash[:16])
		return response
	}

	// Aktuelles Binary finden
	binaryPath := "/usr/local/bin/rmm-agent-linux"
	if _, err := os.Stat(binaryPath); os.IsNotExist(err) {
		// Fallback: eigener Pfad
		binaryPath, _ = os.Executable()
	}

	// Backup vom alten Binary
	backupPath := binaryPath + ".old"
	if err := os.Rename(binaryPath, backupPath); err != nil {
		response.Status = "error"
		response.Error = fmt.Sprintf("Backup fehlgeschlagen: %v", err)
		return response
	}

	// Neues Binary schreiben
	if err := os.WriteFile(binaryPath, binary, 0755); err != nil {
		// Rollback
		os.Rename(backupPath, binaryPath)
		response.Status = "error"
		response.Error = fmt.Sprintf("Binary schreiben fehlgeschlagen: %v", err)
		return response
	}

	log.Printf("Neues Binary geschrieben: %s (%d bytes)", binaryPath, len(binary))

	response.Status = "ok"
	response.Data = map[string]interface{}{
		"message":     fmt.Sprintf("Update auf %s erfolgreich, Neustart...", newVersion),
		"old_version": data["old_version"],
		"new_version": newVersion,
	}

	// Response senden, dann Neustart
	go func() {
		time.Sleep(2 * time.Second)
		log.Printf("Agent-Neustart nach Update auf %s", newVersion)
		// Versuche systemd service restart, fallback zu kill + restart
		if err := exec.Command("systemctl", "restart", "rmm-agent-linux").Run(); err != nil {
			log.Printf("systemctl restart fehlgeschlagen, versuche manuelle Wiederbelebung: %v", err)
			// Als letzter Ausweg: sich selbst ersetzen
			os.Exit(0)
		}
	}()

	return response
}

// --- PTY Handlers ---

func (h *Handler) handlePTYStart(msg Message) Message {
	response := Message{Type: "response", ID: msg.ID}

	sessionID, _ := msg.Params["session_id"].(string)
	if sessionID == "" {
		response.Status = "error"
		response.Error = "session_id erforderlich"
		return response
	}

	cols := uint16(80)
	rows := uint16(24)
	if c, ok := msg.Params["cols"].(float64); ok && c > 0 {
		cols = uint16(c)
	}
	if r, ok := msg.Params["rows"].(float64); ok && r > 0 {
		rows = uint16(r)
	}

	if err := h.client.ptyManager.StartPTY(sessionID, cols, rows); err != nil {
		response.Status = "error"
		response.Error = err.Error()
		return response
	}

	response.Status = "ok"
	response.Data = map[string]interface{}{"session_id": sessionID}
	return response
}

func (h *Handler) handlePTYStop(msg Message) {
	sessionID, _ := msg.Params["session_id"].(string)
	if sessionID != "" {
		h.client.ptyManager.StopPTY(sessionID)
	}
}

func (h *Handler) handlePTYResize(msg Message) {
	sessionID, _ := msg.Params["session_id"].(string)
	cols := uint16(80)
	rows := uint16(24)
	if c, ok := msg.Params["cols"].(float64); ok && c > 0 {
		cols = uint16(c)
	}
	if r, ok := msg.Params["rows"].(float64); ok && r > 0 {
		rows = uint16(r)
	}
	if sessionID != "" {
		h.client.ptyManager.ResizePTY(sessionID, cols, rows)
	}
}

// --- Tunnel Handlers ---

func (h *Handler) handleTunnelConnect(msg Message) Message {
	response := Message{Type: "response", ID: msg.ID}

	sessionID, _ := msg.Params["session_id"].(string)
	tunnelID, _ := msg.Params["tunnel_id"].(string)
	targetHost, _ := msg.Params["target_host"].(string)
	targetPortFloat, _ := msg.Params["target_port"].(float64)
	targetPort := int(targetPortFloat)

	if sessionID == "" || targetHost == "" || targetPort <= 0 {
		response.Status = "error"
		response.Error = "session_id, target_host und target_port erforderlich"
		return response
	}

	if err := h.client.tunnelManager.ConnectSession(sessionID, tunnelID, targetHost, targetPort); err != nil {
		response.Status = "error"
		response.Error = fmt.Sprintf("Session-Connect fehlgeschlagen: %v", err)
		return response
	}

	response.Status = "ok"
	response.Data = map[string]interface{}{"session_id": sessionID}
	return response
}

func (h *Handler) handleTunnelDisconnect(msg Message) {
	sessionID, _ := msg.Params["session_id"].(string)
	if sessionID != "" {
		h.client.tunnelManager.DisconnectSession(sessionID)
	}
}

// SendSessionData schickt Tunnel-Daten als Binary-Message zurueck zum Backend
func (h *Handler) SendSessionData(sessionID string, data []byte) error {
	idBytes := []byte(sessionID)
	if len(idBytes) > 255 {
		return fmt.Errorf("session_id zu lang")
	}

	message := make([]byte, 1+len(idBytes)+len(data))
	message[0] = byte(len(idBytes))
	copy(message[1:], idBytes)
	copy(message[1+len(idBytes):], data)

	return h.client.SendBinaryData(message)
}

// runCommand fuehrt einen Shell-Befehl mit Timeout aus
func (h *Handler) runCommand(command string, timeoutSec int) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutSec)*time.Second)
	defer cancel()

	cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
	output, err := cmd.CombinedOutput()
	return string(output), err
}
// handleUpdateCheck: apt update + Liste verfügbarer Pakete zurückgeben
func (h *Handler) handleUpdateCheck(msg Message) Message {
	var response Message
	response.Type = "response"
	response.ID = msg.ID

	log.Println("Update-Check gestartet (apt update)...")

	// apt update (Paketlisten aktualisieren)
	_, err := h.runCommand("DEBIAN_FRONTEND=noninteractive apt-get update -q 2>&1", 120)
	if err != nil {
		log.Printf("apt update Fehler: %v", err)
		// Fehler ignorieren, apt update gibt oft non-zero bei Warnings zurück
	}

	// Verfügbare Updates abfragen
	output, _ := h.runCommand("LC_ALL=C apt list --upgradable 2>/dev/null", 60)

	// Security-Updates separat zählen
	secOutput, _ := h.runCommand("LC_ALL=C apt list --upgradable 2>/dev/null | grep -c security || true", 30)
	secCount := 0
	fmt.Sscanf(strings.TrimSpace(secOutput), "%d", &secCount)

	// Pakete zählen (Zeilen ohne "Listing..." und Leerzeilen)
	lines := strings.Split(output, "\n")
	pkgCount := 0
	for _, l := range lines {
		l = strings.TrimSpace(l)
		if l != "" && !strings.HasPrefix(l, "Listing") && !strings.HasPrefix(l, "WARNING") {
			pkgCount++
		}
	}

	response.Status = "ok"
	response.Data = map[string]interface{}{
		"update_available": pkgCount > 0,
		"pending_count":    pkgCount,
		"security_count":   secCount,
		"output":           output,
	}
	return response
}

// handleUpdate: apt-get upgrade ausführen
func (h *Handler) handleUpdate(msg Message) Message {
	var response Message
	response.Type = "response"
	response.ID = msg.ID

	reboot := false
	if r, ok := msg.Params["reboot"].(bool); ok {
		reboot = r
	}
	securityOnly := false
	if s, ok := msg.Params["security_only"].(bool); ok {
		securityOnly = s
	}

	log.Printf("Update gestartet (reboot=%v, security_only=%v)", reboot, securityOnly)

	// apt update
	h.runCommand("DEBIAN_FRONTEND=noninteractive apt-get update -q 2>&1", 120)

	// apt-get upgrade oder nur Security-Updates
	var upgradeCmd string
	if securityOnly {
		upgradeCmd = `DEBIAN_FRONTEND=noninteractive apt-get install -y --only-upgrade \
			$(LC_ALL=C apt list --upgradable 2>/dev/null | grep security | cut -d/ -f1 | tr '\n' ' ') 2>&1`
	} else {
		// dist-upgrade statt upgrade: löst auch neue Abhängigkeiten auf (nötig für Proxmox/PVE-Pakete)
		upgradeCmd = "DEBIAN_FRONTEND=noninteractive apt-get dist-upgrade -y 2>&1"
	}

	output, err := h.runCommand(upgradeCmd, 900)
	if err != nil {
		log.Printf("apt upgrade Fehler: %v", err)
		response.Status = "error"
		response.Error = fmt.Sprintf("Update fehlgeschlagen: %v", err)
		response.Data = map[string]interface{}{"output": output}
		return response
	}

	log.Println("Update abgeschlossen")
	response.Status = "ok"
	response.Data = map[string]interface{}{
		"message": "Update erfolgreich",
		"output":  output,
		"reboot":  reboot,
	}

	if reboot {
		log.Println("Reboot nach Update in 10 Sekunden...")
		go func() {
			time.Sleep(10 * time.Second)
			exec.Command("/sbin/shutdown", "-r", "now").Run()
		}()
	}

	return response
}

// handleZpoolScrub: alle ZFS-Pools scrubben
func (h *Handler) handleZpoolScrub(msg Message) Message {
	response := Message{Type: "response", ID: msg.ID}

	log.Println("ZFS Scrub gestartet (alle Pools)...")

	// Alle Pools ermitteln
	poolOutput, err := h.runCommand("zpool list -H -o name 2>&1", 30)
	if err != nil || strings.TrimSpace(poolOutput) == "" {
		response.Status = "error"
		response.Error = "Keine ZFS-Pools gefunden oder zpool nicht verfügbar"
		return response
	}

	pools := []string{}
	for _, line := range strings.Split(strings.TrimSpace(poolOutput), "\n") {
		name := strings.TrimSpace(line)
		if name != "" {
			pools = append(pools, name)
		}
	}

	results := map[string]string{}
	scrubErrors := 0

	// Jeden Pool scrubben (nicht-blockierend: scrub startet sofort und läuft im Hintergrund)
	for _, pool := range pools {
		log.Printf("ZFS Scrub starte: %s", pool)
		out, err := h.runCommand(fmt.Sprintf("zpool scrub %s 2>&1", pool), 60)
		if err != nil {
			results[pool] = fmt.Sprintf("Fehler: %v — %s", err, strings.TrimSpace(out))
			scrubErrors++
		} else {
			results[pool] = "Scrub gestartet"
		}
	}

	if scrubErrors > 0 {
		response.Status = "error"
		response.Error = fmt.Sprintf("%d Pool(s) konnten nicht gestartet werden", scrubErrors)
	} else {
		response.Status = "ok"
	}

	log.Printf("ZFS Scrub: %d Pool(s) gestartet, %d Fehler", len(pools), scrubErrors)
	response.Data = map[string]interface{}{
		"pools":   pools,
		"results": results,
		"message": fmt.Sprintf("Scrub auf %d Pool(s) gestartet", len(pools)-scrubErrors),
	}
	return response
}