vantage-agent/internal/sync/sync.go

package agentsync

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"gitea.hostxtra.co.uk/vantage/vantage-agent/internal/config"
	agentexec "gitea.hostxtra.co.uk/vantage/vantage-agent/internal/exec"
	grpcclient "gitea.hostxtra.co.uk/vantage/vantage-agent/internal/grpc"
	"gitea.hostxtra.co.uk/vantage/vantage-agent/internal/inventory"
	"gitea.hostxtra.co.uk/vantage/vantage-agent/internal/keys"
	"gitea.hostxtra.co.uk/vantage/vantage-agent/internal/monitors"
	agentproxy "gitea.hostxtra.co.uk/vantage/vantage-agent/internal/proxy"
	"gitea.hostxtra.co.uk/vantage/vantage-agent/internal/updates"
	"gitea.hostxtra.co.uk/vantage/vantage-shared/grpc/pb"
)

func Run(ctx context.Context, cfg *config.Config, version string) error {
	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
	if err != nil {
		return fmt.Errorf("dial grpc: %w", err)
	}
	defer client.Close()

	if cfg.PreRegToken != "" {
		log.Println("registering with server...")
		hostname, _ := os.Hostname()
		ipAddress := localIP()
		osInfo := fmt.Sprintf("%s %s", runtime.GOOS, runtime.GOARCH)

		agentToken, err := client.Register(cfg.ServerID, cfg.PreRegToken, hostname, ipAddress, osInfo)
		if err != nil {
			return fmt.Errorf("registration failed: %w", err)
		}

		cfg.AgentToken = agentToken
		cfg.PreRegToken = ""
		if err := config.Save(cfg); err != nil {
			return fmt.Errorf("save config: %w", err)
		}
		log.Println("registration successful")

		client.Close()
		client, err = grpcclient.New(cfg.ServerURL, cfg.TLS)
		if err != nil {
			return fmt.Errorf("reconnect: %w", err)
		}
	}

	if cfg.AgentToken == "" {
		return fmt.Errorf("no agent token available registration required")
	}

	go runCommandStream(ctx, cfg)

	go runUpdateCheck(ctx, cfg)

	go runInventory(ctx, cfg)

	go runWorkloads(ctx, cfg)

	go monitors.Run(ctx, cfg)

	ticker := time.NewTicker(cfg.PollInterval)
	defer ticker.Stop()

	if err := poll(client, cfg, version); err != nil {
		log.Printf("poll error: %v", err)
	}

	for {
		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
			if err := poll(client, cfg, version); err != nil {
				log.Printf("poll error: %v", err)
			}
		}
	}
}

func poll(client *grpcclient.Client, cfg *config.Config, version string) error {
	resp, err := client.SyncKeys(cfg.ServerID, cfg.AgentToken, version)
	if err != nil {
		return fmt.Errorf("SyncKeys: %w", err)
	}

	// Stored atomically: the hourly package loop reads this from another
	// goroutine. Absent on the wire decodes as false, so an older server leaves
	// collection off rather than on.
	collectPackagesFlag.Store(resp.CollectPackages)
	markFirstPoll()

	desired := resp.PublicKeys

	if runtime.GOOS != "linux" {
		return nil
	}

	current, err := keys.ReadAuthorizedKeys()
	if err != nil {
		return fmt.Errorf("read authorized_keys: %w", err)
	}

	if !keys.StateChanged(current, desired) {
		log.Println("authorized_keys unchanged, skipping write")
		return nil
	}

	if err := keys.WriteAuthorizedKeys(desired); err != nil {
		return fmt.Errorf("write authorized_keys: %w", err)
	}
	log.Printf("authorized_keys updated (%d keys)", len(desired))
	return nil
}

// How long a command stream must survive before it counts as having worked.
// Past this, the next drop is treated as a fresh incident rather than as the
// continuation of a run of failures.
const streamHealthyAfter = time.Minute

// Stream staleness. The server beats every 20s, so 70s tolerates three missed
// beats before the stream is written off - high enough that a slow network or a
// briefly busy server does not cost a reconnect, low enough that an agent is
// not uncommandable for minutes after a control-plane restart.
const (
	streamStaleAfter = 70 * time.Second
	streamStaleCheck = 10 * time.Second

	// How often a healthy stream reports itself. Also the interval at which an
	// agent talking to a control plane too old to send heartbeats says so -
	// that agent is running without a watchdog, and the journal should not be
	// silent about it.
	pingSummaryInterval = 5 * time.Minute
)

func runCommandStream(ctx context.Context, cfg *config.Config) {
	backoff := time.Second

	// Two minutes was the old ceiling, and it was reached far too easily. The
	// command stream is what makes this agent controllable at all: while it is
	// down, workflows and console sessions fail as "agent offline" even though
	// SyncKeys keeps polling happily and the fleet list still shows the server
	// active. A shorter ceiling costs a few reconnect attempts; the old one cost
	// two minutes of an agent that looks fine and answers nothing.
	const maxBackoff = 30 * time.Second

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		started := time.Now()
		err := connectAndHandleStream(ctx, cfg)
		if ctx.Err() != nil {
			return
		}

		// A stream that stayed up is evidence the control plane is reachable,
		// whatever ended it. Without this the backoff only ever climbed:
		// connectAndHandleStream returns an error on *every* stream end,
		// including a healthy one dropped by a routine deploy, so an agent
		// pinned itself at the ceiling after a handful of ordinary restarts and
		// stayed there for the rest of its life.
		if time.Since(started) >= streamHealthyAfter {
			backoff = time.Second
		}

		// The uptime is in the line because it is what distinguishes a stream
		// that never worked from one that ran for hours and was dropped by a
		// deploy - and it is the same measure that decides whether the backoff
		// resets, so a reader can see why the delay is what it is.
		up := time.Since(started).Truncate(time.Second)
		if err != nil {
			log.Printf("command stream error after %s: %v, reconnecting in %s", up, err, backoff)
		} else {
			log.Printf("command stream closed after %s, reconnecting in %s", up, backoff)
		}

		select {
		case <-ctx.Done():
			return
		case <-time.After(backoff):
		}

		if backoff < maxBackoff {
			backoff *= 2
			if backoff > maxBackoff {
				backoff = maxBackoff
			}
		}
	}
}

func connectAndHandleStream(ctx context.Context, cfg *config.Config) error {
	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
	if err != nil {
		return fmt.Errorf("dial: %w", err)
	}
	defer client.Close()

	// Cancelling this context is what unblocks Recv when the stream has gone
	// quiet. Without it the watchdog below would have no way to interrupt a
	// read that is never going to return.
	streamCtx, abandon := context.WithCancel(ctx)
	defer abandon()

	stream, err := client.CommandStream(streamCtx)
	if err != nil {
		return fmt.Errorf("open stream: %w", err)
	}

	if err := stream.Send(&pb.AgentMessage{
		ServerId:   cfg.ServerID,
		AgentToken: cfg.AgentToken,
		Ready:      &pb.AgentReady{},
	}); err != nil {
		return fmt.Errorf("send auth: %w", err)
	}

	log.Printf("command stream connected to %s", cfg.ServerURL)

	var sendMu sync.Mutex
	send := func(msg *pb.AgentMessage) error {
		sendMu.Lock()
		defer sendMu.Unlock()
		return stream.Send(msg)
	}

	// Results that could not be sent on an earlier stream go out first. In
	// its own goroutine: a queued reboot re-check can take a while on
	// Windows, and the receive loop below must start promptly.
	go flushPendingResults(send)

	// Stream liveness, tracked here rather than left to gRPC keepalive.
	//
	// Keepalive operates on the transport, and behind an L7 proxy the transport
	// ends at the proxy: it answers pings whether or not the server behind it
	// is still running. A control-plane pod that dies therefore leaves this
	// agent blocked in Recv on a stream that will never deliver another message
	// and never error, while the control plane dispatches commands into it and
	// the operator watches nothing happen.
	//
	// The watchdog only arms once a ping has actually been seen. A server too
	// old to send them must not be treated as dead - that would put the agent
	// in a reconnect loop against a control plane that is working perfectly.
	var (
		lastMu   sync.Mutex
		lastRecv = time.Now()
		pinged   bool
		beats    int
	)
	markRecv := func(isPing bool) {
		lastMu.Lock()
		lastRecv = time.Now()
		if isPing {
			beats++
			// Logged once per stream, because it is the moment the agent starts
			// holding the control plane to account: before this the watchdog is
			// disarmed and a dead stream would go unnoticed indefinitely.
			if !pinged {
				pinged = true
				log.Printf("command stream heartbeat detected, watchdog armed (%s threshold)", streamStaleAfter)
			}
		}
		lastMu.Unlock()
	}

	go func() {
		t := time.NewTicker(streamStaleCheck)
		defer t.Stop()

		// Reported periodically rather than per beat: at one every 20s the
		// journal would be nothing else. The count is what makes a partial
		// failure visible - beats arriving but fewer than expected is a
		// different problem from beats stopping altogether.
		summary := time.NewTicker(pingSummaryInterval)
		defer summary.Stop()

		for {
			select {
			case <-streamCtx.Done():
				return
			case <-summary.C:
				lastMu.Lock()
				n, armed := beats, pinged
				beats = 0
				lastMu.Unlock()
				if armed {
					log.Printf("command stream healthy, %d heartbeats in the last %s", n, pingSummaryInterval)
				} else {
					log.Printf("command stream up but sending no heartbeats; "+
						"control plane predates them, watchdog stays disarmed (last message %s ago)",
						time.Since(lastRecv).Truncate(time.Second))
				}
			case <-t.C:
				lastMu.Lock()
				idle, armed := time.Since(lastRecv), pinged
				lastMu.Unlock()
				if armed && idle > streamStaleAfter {
					log.Printf("command stream silent for %s (threshold %s), assuming it is dead and reconnecting",
						idle.Truncate(time.Second), streamStaleAfter)
					abandon()
					return
				}
			}
		}
	}()

	for {
		cmd, err := stream.Recv()
		if err != nil {
			return fmt.Errorf("recv: %w", err)
		}
		markRecv(cmd.Ping != nil)

		// Pings carry nothing and are not acknowledged; being received is their
		// whole purpose.
		if cmd.Ping != nil {
			continue
		}

		if cmd.GenerateKey != nil {
			go handleGenerateKey(cfg, cmd)
		}
		if cmd.DeleteKey != nil {
			go handleDeleteKey(cmd)
		}
		if cmd.UpdateAgent != nil {
			go handleUpdateAgent(cmd)
		}
		if cmd.ApplyUpdates != nil {
			go handleApplyUpdates(send, cfg, cmd)
		}
		if cmd.CleanupWorkspace != nil {
			go handleCleanupWorkspace(cmd)
		}
		if cmd.OpenProxy != nil {
			go handleOpenProxy(ctx, cfg, cmd.OpenProxy)
		}
		if cmd.RefreshWorkloads != nil {
			go handleRefreshWorkloads(cfg)
		}
		if cmd.ControlWorkload != nil {
			go handleControlWorkload(send, cfg, cmd.CommandId, cmd.ControlWorkload)
		}
		if cmd.WorkloadLogs != nil {
			go handleWorkloadLogs(send, cfg, cmd.CommandId, cmd.WorkloadLogs)
		}
		if cmd.RunStep != nil {
			go func(rc *pb.RunStepCmd, cid string) {
				emit := func(seq uint64, data []byte) {
					_ = send(&pb.AgentMessage{
						ServerId:   cfg.ServerID,
						AgentToken: cfg.AgentToken,
						StepOutput: &pb.StepOutputChunk{CommandId: cid, Seq: seq, Data: data},
					})
				}
				res := agentexec.RunStep(rc, emit)
				res.CommandId = cid

				_ = send(&pb.AgentMessage{
					ServerId:   cfg.ServerID,
					AgentToken: cfg.AgentToken,
					StepOutput: &pb.StepOutputChunk{CommandId: cid, Eof: true},
				})
				_ = send(&pb.AgentMessage{
					ServerId:   cfg.ServerID,
					AgentToken: cfg.AgentToken,
					StepResult: res,
				})
			}(cmd.RunStep, cmd.CommandId)
			continue
		}
	}
}

func runUpdateCheck(ctx context.Context, cfg *config.Config) {
	const interval = time.Hour

	doCheck := func() {
		pkgs, err := updates.CheckAvailable()
		if err != nil {
			log.Printf("update check error: %v", err)
			return
		}
		pbUpdates := make([]pb.PackageUpdate, len(pkgs))
		for i, p := range pkgs {
			pbUpdates[i] = pb.PackageUpdate{
				Name:           p.Name,
				CurrentVersion: p.CurrentVersion,
				NewVersion:     p.NewVersion,
			}
		}
		client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
		if err != nil {
			log.Printf("update report dial error: %v", err)
			return
		}
		defer client.Close()
		if err := client.ReportUpdates(cfg.ServerID, cfg.AgentToken, pbUpdates); err != nil {
			log.Printf("ReportUpdates error: %v", err)
			return
		}
		log.Printf("reported %d available OS updates", len(pkgs))

		// Same hourly cadence, same connection. A package set changes on
		// roughly the schedule available updates do, so this needs no timer of
		// its own.
		reportPackages(client, cfg)
	}

	// The boot round only: after this the flag has long been set, and every
	// later tick is an hour past a poll that runs every 30s.
	waitFirstPoll(ctx, firstPollWait)

	doCheck()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			doCheck()
		}
	}
}

func runInventory(ctx context.Context, cfg *config.Config) {
	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
	if err != nil {
		log.Printf("inventory dial error: %v", err)
		return
	}
	defer client.Close()

	report := func(static bool) error {
		r := inventory.Collect(static)
		r.ServerId = cfg.ServerID
		r.AgentToken = cfg.AgentToken
		// Static snapshots only - every 15 minutes, not every 30 seconds. On
		// Windows this spawns a PowerShell process, which is not something to
		// do twice a minute forever, and a host rebooted by hand clearing the
		// flag within a quarter of an hour is soon enough.
		//
		// Computed here rather than inside inventory.Collect so the inventory
		// package gains no dependency on updates.
		if static {
			r.RebootRequired = updates.RebootRequired()
		}
		if err := client.ReportInventory(r); err != nil {
			log.Printf("report inventory: %v", err)
			return err
		}
		return nil
	}

	// The startup static report carries the boot time the server uses to
	// prove a patch reboot happened, so a failure is retried on the next
	// ticks (every 30 seconds, up to startupStaticAttempts in total) instead
	// of waiting a quarter of an hour for the next static snapshot.
	const startupStaticAttempts = 10
	attempts := 1
	startupPending := report(true) != nil
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	tick := 0
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			tick++
			if startupPending && attempts < startupStaticAttempts {
				attempts++
				startupPending = report(true) != nil
				continue
			}
			startupPending = false
			report(tick%30 == 0)
		}
	}
}

func handleCleanupWorkspace(cmd *pb.ServerCommand) {
	id := cmd.CleanupWorkspace.WorkspaceId
	dir := agentexec.WorkspacePath(id)
	if err := os.RemoveAll(dir); err != nil {
		log.Printf("cleanup workspace %s failed (cmd=%s): %v", dir, cmd.CommandId, err)
		return
	}
	log.Printf("removed run workspace %s (cmd=%s)", dir, cmd.CommandId)
}

// handleOpenProxy relays one console connection. It uses its own gRPC
// connection so console traffic never shares a stream with commands, key sync
// or workflow output.
func handleOpenProxy(ctx context.Context, cfg *config.Config, cmd *pb.OpenProxyCmd) {
	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
	if err != nil {
		log.Printf("proxy %s: dial control plane: %v", cmd.ProxyId, err)
		return
	}
	defer client.Close()

	stream, err := client.ProxyStream(ctx)
	if err != nil {
		log.Printf("proxy %s: open stream: %v", cmd.ProxyId, err)
		return
	}

	log.Printf("proxy %s: relaying 127.0.0.1:%d", cmd.ProxyId, cmd.Port)
	if err := agentproxy.Open(stream, cfg.ServerID, cfg.AgentToken, cmd.ProxyId, cmd.Port); err != nil {
		log.Printf("proxy %s: %v", cmd.ProxyId, err)
	}
}

func handleDeleteKey(cmd *pb.ServerCommand) {
	label := cmd.DeleteKey.Label
	keyPath := fmt.Sprintf("/root/.ssh/vantage_%s", strings.ReplaceAll(label, " ", "_"))

	if err := keys.RemoveSSHIdentity(keyPath); err != nil {
		log.Printf("remove ssh identity failed (cmd=%s): %v", cmd.CommandId, err)
	}

	for _, path := range []string{keyPath, keyPath + ".pub"} {
		if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
			log.Printf("delete key file %s (cmd=%s): %v", path, cmd.CommandId, err)
		}
	}
	log.Printf("deleted local key files for %q (cmd=%s)", label, cmd.CommandId)
}

func handleUpdateAgent(cmd *pb.ServerCommand) {
	if runtime.GOOS == "windows" {
		handleUpdateAgentWindows(cmd)
		return
	}

	u := cmd.UpdateAgent
	arch := runtime.GOARCH
	tag := "agent%2Fv" + u.Version
	binaryURL := fmt.Sprintf("%s/vantage/vantage-agent/releases/download/%s/vantage-agent-linux-%s", u.GiteaBaseURL, tag, arch)
	checksumURL := fmt.Sprintf("%s/vantage/vantage-agent/releases/download/%s/checksums.txt", u.GiteaBaseURL, tag)

	log.Printf("updating agent to v%s from %s (cmd=%s)", u.Version, u.GiteaBaseURL, cmd.CommandId)

	tmpBin := "/tmp/vantage-agent-update"
	if err := downloadFile(binaryURL, tmpBin); err != nil {
		log.Printf("update download failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}

	checksumData, err := httpGetBytes(checksumURL)
	if err != nil {
		log.Printf("update checksum fetch failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}
	if err := verifyChecksum(tmpBin, fmt.Sprintf("vantage-agent-linux-%s", arch), checksumData); err != nil {
		log.Printf("update checksum mismatch (cmd=%s): %v", cmd.CommandId, err)
		os.Remove(tmpBin)
		return
	}

	if err := os.Chmod(tmpBin, 0755); err != nil {
		log.Printf("update chmod failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}
	if err := os.Rename(tmpBin, "/usr/local/bin/vantage-agent"); err != nil {
		log.Printf("update replace binary failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}

	log.Printf("agent binary replaced, restarting service (cmd=%s)", cmd.CommandId)
	exec.Command("systemctl", "restart", "vantage-agent").Run()
}

func handleUpdateAgentWindows(cmd *pb.ServerCommand) {
	u := cmd.UpdateAgent
	tag := "agent%2Fv" + u.Version
	msiURL := fmt.Sprintf("%s/vantage/vantage-agent/releases/download/%s/vantage-agent.msi", u.GiteaBaseURL, tag)
	checksumURL := fmt.Sprintf("%s/vantage/vantage-agent/releases/download/%s/checksums-msi.txt", u.GiteaBaseURL, tag)

	log.Printf("updating agent to v%s from %s (cmd=%s)", u.Version, u.GiteaBaseURL, cmd.CommandId)

	msiPath := filepath.Join(os.TempDir(), "vantage-agent-update.msi")
	if err := downloadFile(msiURL, msiPath); err != nil {
		log.Printf("update download failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}

	checksumData, err := httpGetBytes(checksumURL)
	if err != nil {
		log.Printf("update checksum fetch failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}
	if err := verifyChecksum(msiPath, "vantage-agent.msi", checksumData); err != nil {
		log.Printf("update checksum mismatch (cmd=%s): %v", cmd.CommandId, err)
		os.Remove(msiPath)
		return
	}

	logPath := filepath.Join(os.TempDir(), "vantage-agent-msi.log")

	// The MSI stops the vantage-agent service as part of the upgrade. Anything
	// descended from this process is killed with it, so msiexec must not be a
	// child: run it from a scheduled task, which is parented to the Task
	// Scheduler service instead.
	if err := launchDetachedUpdate(msiPath, logPath, cmd.CommandId); err != nil {
		log.Printf("failed to launch msiexec (cmd=%s): %v", cmd.CommandId, err)
		return
	}
	log.Printf("scheduled msiexec for upgrade to v%s (cmd=%s)", u.Version, cmd.CommandId)
}

const updateTaskName = "VantageAgentUpdate"

func launchDetachedUpdate(msiPath, logPath, commandID string) error {
	scriptPath := filepath.Join(os.TempDir(), "vantage-agent-update.cmd")
	script := fmt.Sprintf("@echo off\r\n"+
		"timeout /t 5 /nobreak >nul\r\n"+
		"msiexec /i \"%s\" /qn /norestart /l*v \"%s\"\r\n"+
		"schtasks /delete /tn %s /f >nul 2>&1\r\n"+
		"del /f /q \"%s\" >nul 2>&1\r\n"+
		"(goto) 2>nul & del /f /q \"%%~f0\"\r\n",
		msiPath, logPath, updateTaskName, msiPath)
	if err := os.WriteFile(scriptPath, []byte(script), 0o600); err != nil {
		return fmt.Errorf("write update script: %w", err)
	}

	// Stale task from a previous attempt would make /create fail even with /f
	// if it is still running, so tear it down first and ignore the result.
	exec.Command("schtasks", "/end", "/tn", updateTaskName).Run()
	exec.Command("schtasks", "/delete", "/tn", updateTaskName, "/f").Run()

	create := exec.Command("schtasks", "/create",
		"/tn", updateTaskName,
		"/tr", `"`+scriptPath+`"`,
		"/sc", "once",
		// Already in the past: the task never fires on its own, only via /run.
		"/st", "00:00",
		"/ru", "SYSTEM",
		"/rl", "HIGHEST",
		"/f")
	if out, err := create.CombinedOutput(); err != nil {
		return fmt.Errorf("schtasks create: %v: %s", err, strings.TrimSpace(string(out)))
	}

	if out, err := exec.Command("schtasks", "/run", "/tn", updateTaskName).CombinedOutput(); err != nil {
		return fmt.Errorf("schtasks run: %v: %s", err, strings.TrimSpace(string(out)))
	}
	return nil
}

func downloadFile(url, dest string) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP %d from %s", resp.StatusCode, url)
	}
	f, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, resp.Body)
	return err
}

func httpGetBytes(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("HTTP %d from %s", resp.StatusCode, url)
	}
	return io.ReadAll(resp.Body)
}

func verifyChecksum(filePath, filename string, checksumData []byte) error {
	f, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return err
	}
	actual := hex.EncodeToString(h.Sum(nil))

	for _, line := range strings.Split(string(checksumData), "\n") {
		fields := strings.Fields(line)
		if len(fields) == 2 && fields[1] == filename {
			if fields[0] != actual {
				return fmt.Errorf("expected %s got %s", fields[0], actual)
			}
			return nil
		}
	}
	return fmt.Errorf("no checksum entry found for %s", filename)
}

func handleGenerateKey(cfg *config.Config, cmd *pb.ServerCommand) {
	g := cmd.GenerateKey
	label := g.Label
	keyPath := fmt.Sprintf("/root/.ssh/vantage_%s", strings.ReplaceAll(label, " ", "_"))

	opts := keys.KeyGenOptions{
		KeyType:    g.KeyType,
		KeySize:    g.KeySize,
		Passphrase: g.Passphrase,
		Comment:    g.Comment,
	}
	pubKey, err := keys.GenerateKeyPair(keyPath, opts)
	if err != nil {
		log.Printf("key generation failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}

	privKeyData, err := os.ReadFile(keyPath)
	if err != nil {
		log.Printf("read private key failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}

	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
	if err != nil {
		log.Printf("dial for key upload failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}
	defer client.Close()

	keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, string(privKeyData), label)
	if err != nil {
		log.Printf("key upload failed (cmd=%s): %v", cmd.CommandId, err)
		return
	}

	if err := keys.AddSSHIdentity(keyPath); err != nil {
		log.Printf("add ssh identity failed (cmd=%s): %v", cmd.CommandId, err)
	}
	log.Printf("generated and uploaded key %q (key_id=%s, cmd=%s)", label, keyID, cmd.CommandId)
}

func localIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return ""
	}
	for _, addr := range addrs {
		if ipNet, ok := addr.(*net.IPNet); ok && !ipNet.IP.IsLoopback() {
			if ipNet.IP.To4() != nil {
				return ipNet.IP.String()
			}
		}
	}
	return ""
}

func GenerateAndUpload(cfg *config.Config, label string) error {
	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
	if err != nil {
		return err
	}
	defer client.Close()

	keyPath := fmt.Sprintf("/root/.ssh/vantage_%s", strings.ReplaceAll(label, " ", "_"))
	pubKey, err := keys.GenerateKeyPair(keyPath, keys.KeyGenOptions{Comment: label})
	if err != nil {
		return err
	}

	privKeyData, err := os.ReadFile(keyPath)
	if err != nil {
		return fmt.Errorf("read private key: %w", err)
	}

	keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, string(privKeyData), label)
	if err != nil {
		return err
	}

	if err := keys.AddSSHIdentity(keyPath); err != nil {
		log.Printf("add ssh identity: %v", err)
	}
	log.Printf("uploaded generated key %s (key_id=%s)", label, keyID)
	return nil
}