rmm2/agent-windows/main.go

//go:build windows

package main

import (
	"crypto/sha256"
	"crypto/tls"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/eventlog"
	"golang.org/x/sys/windows/svc/mgr"

	"github.com/cynfo/rmm-agent-windows/client"
	"github.com/cynfo/rmm-agent-windows/collector"
	"github.com/cynfo/rmm-agent-windows/config"
	"github.com/cynfo/rmm-agent-windows/ws"
)

var Version = "1.0.0"

const (
	ServiceName        = "RMMAgent"
	ServiceDisplayName = "RMM Agent"
	ServiceDescription = "RMM Agent - Remote Monitoring & Management"
	AgentIDFile        = "agent_id.txt"
)

// agentService implementiert das Windows Service Interface
type agentService struct {
	cfg     *config.Config
	cfgPath string
}

func (s *agentService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (bool, uint32) {
	changes <- svc.Status{State: svc.StartPending}
	stop := make(chan struct{})
	go runAgent(s.cfg, s.cfgPath, stop)
	changes <- svc.Status{State: svc.Running, Accepts: svc.AcceptStop | svc.AcceptShutdown}

	for c := range r {
		switch c.Cmd {
		case svc.Stop, svc.Shutdown:
			changes <- svc.Status{State: svc.StopPending}
			close(stop)
			return false, 0
		}
	}
	return false, 0
}

func main() {
	var (
		cfgPath   = flag.String("config", defaultConfigPath(), "Pfad zur Konfigurationsdatei")
		flagDebug = flag.Bool("debug", false, "Im Vordergrund ausfuehren (kein Dienst)")
		flagInst  = flag.Bool("install", false, "Als Windows-Dienst installieren")
		flagUninst = flag.Bool("uninstall", false, "Windows-Dienst deinstallieren")
	)
	flag.Parse()

	// Logging in Datei (neben Binary)
	logPath := filepath.Join(filepath.Dir(*cfgPath), "rmm-agent.log")
	if lf, err := os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644); err == nil {
		log.SetOutput(lf)
	}
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.Printf("RMM Windows Agent v%s startet", Version)

	cfg, err := config.Load(*cfgPath)
	if err != nil {
		log.Fatalf("Konfiguration laden fehlgeschlagen (%s): %v", *cfgPath, err)
	}

	if *flagInst {
		installService(cfg, *cfgPath)
		return
	}
	if *flagUninst {
		uninstallService()
		return
	}

	// Interaktiv (debug) oder als Dienst?
	isInteractive, err := svc.IsAnInteractiveSession()
	if err != nil {
		log.Printf("IsAnInteractiveSession: %v", err)
		isInteractive = true
	}

	if *flagDebug || isInteractive {
		log.Println("Starte im Debug-Modus (Vordergrund)")
		stop := make(chan struct{})
		runAgent(cfg, *cfgPath, stop)
	} else {
		if err := svc.Run(ServiceName, &agentService{cfg: cfg, cfgPath: *cfgPath}); err != nil {
			log.Fatalf("Dienst-Fehler: %v", err)
		}
	}
}

// runAgent ist die Hauptschleife des Agents
func runAgent(cfg *config.Config, cfgPath string, stop <-chan struct{}) {
	c := client.New(cfg.BackendURL, cfg.APIKey, cfg.Insecure)

	// Agent-ID laden oder registrieren
	agentID, err := loadOrRegister(c, cfg, cfgPath)
	if err != nil {
		log.Fatalf("Registrierung fehlgeschlagen: %v", err)
	}
	log.Printf("Agent-ID: %s", agentID)

	// WS-Handler und -Client starten
	handler := ws.NewHandler()
	wsClient := ws.NewClient(cfg.BackendURL, agentID, cfg.APIKey, cfg.Insecure, handler)
	go wsClient.Run(stop)

	// Heartbeat-Schleife
	interval := time.Duration(cfg.IntervalSeconds) * time.Second
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	// Ersten Heartbeat sofort
	sendHeartbeat(c, cfg, agentID)

	for {
		select {
		case <-stop:
			log.Println("Agent wird gestoppt")
			return
		case <-ticker.C:
			sendHeartbeat(c, cfg, agentID)
		}
	}
}

func sendHeartbeat(c *client.Client, cfg *config.Config, agentID string) {
	sysData, err := collector.Collect()
	if err != nil {
		log.Printf("Collector-Fehler: %v", err)
	}

	var rawData json.RawMessage
	if sysData != nil {
		rawData, _ = sysData.ToJSON()
	}

	resp, err := c.Heartbeat(client.HeartbeatRequest{
		AgentID:      agentID,
		AgentVersion: Version,
		Platform:     "windows",
		SystemData:   rawData,
	})
	if err != nil {
		log.Printf("Heartbeat fehlgeschlagen: %v", err)
		return
	}
	log.Printf("Heartbeat OK")

	if resp.UpdateAvailable && resp.UpdateVersion != "" && resp.UpdateVersion != Version {
		log.Printf("Update verfuegbar: v%s → v%s", Version, resp.UpdateVersion)
		go doSelfUpdate(cfg, resp.UpdateVersion, resp.UpdateHash)
	}
}

// doSelfUpdate laedt das neue Binary herunter und startet einen
// PowerShell-Script der nach Dienst-Stop das Binary tauscht und neu startet.
func doSelfUpdate(cfg *config.Config, newVersion, expectedHash string) {
	log.Printf("OTA-Update: Starte Download v%s", newVersion)

	exe, err := os.Executable()
	if err != nil {
		log.Printf("OTA-Update: Executable-Pfad: %v", err)
		return
	}
	dir := filepath.Dir(exe)
	tmpBin := filepath.Join(dir, "rmm-agent-windows-update.exe")
	scriptPath := filepath.Join(dir, "rmm-update.ps1")

	// Binary herunterladen
	dlURL := fmt.Sprintf("%s/api/v1/firmware/download?platform=windows&api_key=%s",
		cfg.BackendURL, cfg.APIKey)

	tr := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: cfg.Insecure}}
	hc := &http.Client{Transport: tr, Timeout: 10 * time.Minute}

	req, err := http.NewRequest("GET", dlURL, nil)
	if err != nil {
		log.Printf("OTA-Update: Request: %v", err)
		return
	}
	req.Header.Set("X-API-Key", cfg.APIKey)

	resp, err := hc.Do(req)
	if err != nil {
		log.Printf("OTA-Update: Download: %v", err)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		log.Printf("OTA-Update: HTTP %d", resp.StatusCode)
		return
	}

	f, err := os.Create(tmpBin)
	if err != nil {
		log.Printf("OTA-Update: Temp-Datei: %v", err)
		return
	}

	h := sha256.New()
	if _, err := io.Copy(io.MultiWriter(f, h), resp.Body); err != nil {
		f.Close()
		os.Remove(tmpBin)
		log.Printf("OTA-Update: Schreiben: %v", err)
		return
	}
	f.Close()

	// Hash pruefen
	gotHash := hex.EncodeToString(h.Sum(nil))
	if expectedHash != "" && gotHash != expectedHash {
		os.Remove(tmpBin)
		log.Printf("OTA-Update: Hash-Mismatch: erwartet=%s, erhalten=%s", expectedHash, gotHash)
		return
	}

	log.Printf("OTA-Update: Download OK (hash=%s), starte Update-Script", gotHash[:12])

	// PowerShell-Script schreiben das nach Dienst-Stop das Binary tauscht
	psScript := fmt.Sprintf(`
Start-Sleep -Seconds 3
Stop-Service -Name RMMAgent -Force -ErrorAction SilentlyContinue
Start-Sleep -Seconds 2
Copy-Item -Path "%s" -Destination "%s" -Force
Remove-Item -Path "%s" -Force -ErrorAction SilentlyContinue
Remove-Item -Path "%s" -Force -ErrorAction SilentlyContinue
Start-Service -Name RMMAgent
`, tmpBin, exe, tmpBin, scriptPath)

	if err := os.WriteFile(scriptPath, []byte(psScript), 0644); err != nil {
		log.Printf("OTA-Update: Script schreiben: %v", err)
		return
	}

	// Script detached starten (laeuft weiter wenn dieser Prozess beendet wird)
	cmd := exec.Command("powershell", "-NonInteractive", "-ExecutionPolicy", "Bypass",
		"-WindowStyle", "Hidden", "-File", scriptPath)
	if err := cmd.Start(); err != nil {
		log.Printf("OTA-Update: Script starten: %v", err)
		return
	}
	// cmd.Wait() nicht aufrufen — Script laeuft unabhaengig weiter
	log.Printf("OTA-Update: Script gestartet (PID %d), Dienst wird gestoppt", cmd.Process.Pid)
}

// loadOrRegister liest die Agent-ID aus Datei oder registriert sich neu
func loadOrRegister(c *client.Client, cfg *config.Config, cfgPath string) (string, error) {
	idFile := filepath.Join(filepath.Dir(cfgPath), AgentIDFile)

	if data, err := os.ReadFile(idFile); err == nil {
		id := string(data)
		if len(id) > 0 {
			// Re-registrieren um sicherzustellen dass der Agent im Backend bekannt ist
			hostname, _ := os.Hostname()
			c.Register(client.RegisterRequest{
				AgentID:      id,
				Name:         cfg.AgentName,
				Hostname:     hostname,
				IP:           "",
				AgentVersion: Version,
				Platform:     "windows",
			})
			return id, nil
		}
	}

	// Neu registrieren
	hostname, _ := os.Hostname()
	name := cfg.AgentName
	if name == "" {
		name = hostname
	}
	id, err := c.Register(client.RegisterRequest{
		Name:         name,
		Hostname:     hostname,
		AgentVersion: Version,
		Platform:     "windows",
	})
	if err != nil {
		return "", fmt.Errorf("Registrierung: %w", err)
	}

	if err := os.WriteFile(idFile, []byte(id), 0644); err != nil {
		log.Printf("Agent-ID speichern fehlgeschlagen: %v", err)
	}
	log.Printf("Neu registriert als: %s", id)
	return id, nil
}

// defaultConfigPath gibt den Standardpfad zur Konfigurationsdatei zurück
func defaultConfigPath() string {
	exe, err := os.Executable()
	if err != nil {
		return `C:\Program Files\RMMAgent\config.yaml`
	}
	return filepath.Join(filepath.Dir(exe), "config.yaml")
}

// installService installiert den Agent als Windows-Dienst
func installService(cfg *config.Config, cfgPath string) {
	m, err := mgr.Connect()
	if err != nil {
		log.Fatalf("Service Manager: %v", err)
	}
	defer m.Disconnect()

	s, err := m.OpenService(ServiceName)
	if err == nil {
		s.Close()
		log.Fatalf("Dienst '%s' existiert bereits — zuerst deinstallieren", ServiceName)
	}

	exe, err := os.Executable()
	if err != nil {
		log.Fatalf("Executable-Pfad: %v", err)
	}

	s, err = m.CreateService(ServiceName, exe, mgr.Config{
		DisplayName:  ServiceDisplayName,
		Description:  ServiceDescription,
		StartType:    mgr.StartAutomatic,
		ErrorControl: mgr.ErrorNormal,
	}, "-config", cfgPath)
	if err != nil {
		log.Fatalf("Dienst erstellen: %v", err)
	}
	defer s.Close()

	// Event Log registrieren
	eventlog.InstallAsEventCreate(ServiceName, eventlog.Error|eventlog.Warning|eventlog.Info)

	log.Printf("Dienst '%s' installiert", ServiceName)
	log.Println("Starten mit: net start RMMAgent")
}

// uninstallService entfernt den Windows-Dienst
func uninstallService() {
	m, err := mgr.Connect()
	if err != nil {
		log.Fatalf("Service Manager: %v", err)
	}
	defer m.Disconnect()

	s, err := m.OpenService(ServiceName)
	if err != nil {
		log.Fatalf("Dienst nicht gefunden: %v", err)
	}
	defer s.Close()

	s.Control(svc.Stop)
	time.Sleep(2 * time.Second)

	if err := s.Delete(); err != nil {
		log.Fatalf("Dienst loeschen: %v", err)
	}

	eventlog.Remove(ServiceName)
	log.Printf("Dienst '%s' deinstalliert", ServiceName)
}