restic-manager/cmd/agent/main.go

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log/slog"
	"os"
	"os/signal"
	"syscall"
	"time"

	"gitea.dcglab.co.uk/steve/restic-manager/internal/agent/config"
	"gitea.dcglab.co.uk/steve/restic-manager/internal/agent/runner"
	"gitea.dcglab.co.uk/steve/restic-manager/internal/agent/sysinfo"
	"gitea.dcglab.co.uk/steve/restic-manager/internal/agent/wsclient"
	"gitea.dcglab.co.uk/steve/restic-manager/internal/api"
	"gitea.dcglab.co.uk/steve/restic-manager/internal/restic"
)

var version = "dev"

func main() {
	if err := run(); err != nil {
		slog.Error("agent fatal", "err", err)
		os.Exit(1)
	}
}

func run() error {
	configPath := flag.String("config", config.DefaultPath(), "path to agent.yaml")
	enrollServer := flag.String("enroll-server", "", "server URL (used with -enroll-token to perform first-run enrollment)")
	enrollToken := flag.String("enroll-token", "", "one-time enrollment token (operator copies this from the UI)")
	showVersion := flag.Bool("version", false, "print version and exit")
	flag.Parse()

	if *showVersion {
		fmt.Println("restic-manager-agent", version)
		return nil
	}

	logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
	slog.SetDefault(logger)

	cfg, err := config.Load(*configPath)
	if err != nil {
		return fmt.Errorf("config: %w", err)
	}

	if *enrollToken != "" {
		if *enrollServer == "" {
			return errors.New("enrollment: -enroll-server is required with -enroll-token")
		}
		return doEnroll(*enrollServer, *enrollToken, cfg, version)
	}

	if !cfg.Enrolled() {
		return fmt.Errorf("agent is not enrolled; run with -enroll-server and -enroll-token first (config %q)", *configPath)
	}

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	snap, err := sysinfo.Collect(ctx, cfg.ResticPath)
	if err != nil {
		return fmt.Errorf("sysinfo: %w", err)
	}
	slog.Info("agent starting",
		"version", version,
		"host_id", cfg.HostID,
		"server", cfg.ServerURL,
		"restic_version", snap.ResticVersion,
		"protocol_version", snap.ProtocolVersion,
	)

	resticBin, _ := restic.Locate(cfg.ResticPath) // empty is fine; commands fail with a clear error later

	wsCfg := wsclient.Config{
		ServerURL:     cfg.ServerURL,
		AgentToken:    cfg.AgentToken,
		HostID:        cfg.HostID,
		CertPinSHA256: cfg.CertPinSHA256,
		HelloPayload: api.HelloPayload{
			ProtocolVersion: snap.ProtocolVersion,
			AgentVersion:    version,
			ResticVersion:   snap.ResticVersion,
			Hostname:        snap.Hostname,
			OS:              snap.OS,
			Arch:            snap.Arch,
		},
	}

	d := &dispatcher{
		resticBin:    resticBin,
		repoURL:      cfg.RepoURL,
		repoPassword: cfg.RepoPassword,
	}
	if err := wsclient.Run(ctx, wsCfg, d.handle); err != nil {
		return fmt.Errorf("ws run: %w", err)
	}
	slog.Info("agent shutting down")
	return nil
}

// dispatcher closes over the long-lived agent settings (restic path,
// repo creds) so handle() can spawn the runner without re-loading
// config every time.
type dispatcher struct {
	resticBin    string
	repoURL      string
	repoPassword string
}

func (d *dispatcher) handle(ctx context.Context, env api.Envelope, tx wsclient.Sender) error {
	switch env.Type {
	case api.MsgCommandRun:
		var p api.CommandRunPayload
		if err := env.UnmarshalPayload(&p); err != nil {
			return fmt.Errorf("command.run: %w", err)
		}
		return d.runJob(ctx, p, tx)

	case api.MsgCommandCancel:
		// TODO(P2): cancellation requires keeping a job→cancelFunc map.
		slog.Info("ws agent: command.cancel received (cancellation lands in P2)", "id", env.ID)

	case api.MsgScheduleSet:
		// TODO(P2): apply the schedule.
		slog.Info("ws agent: schedule.set received (handled in P2)", "id", env.ID)

	case api.MsgConfigUpdate:
		var p api.ConfigUpdatePayload
		_ = env.UnmarshalPayload(&p)
		// In-memory only for now — restart loses these. Persistent
		// secret storage lands with P2's keyring work.
		if p.RepoURL != "" {
			d.repoURL = p.RepoURL
			slog.Info("ws agent: repo URL updated via config.update")
		}
		if p.RepoPassword != "" {
			d.repoPassword = p.RepoPassword
			slog.Info("ws agent: repo password updated via config.update")
		}

	case api.MsgAgentUpdateAvail:
		var p api.AgentUpdateAvailablePayload
		_ = env.UnmarshalPayload(&p)
		slog.Info("ws agent: update available", "version", p.LatestVersion, "url", p.PackageURL)

	default:
		slog.Debug("ws agent: ignored message", "type", env.Type)
	}
	return nil
}

// runJob spawns a runner for one job. We launch a goroutine so the
// WS read loop keeps draining messages while restic chugs along.
func (d *dispatcher) runJob(ctx context.Context, p api.CommandRunPayload, tx wsclient.Sender) error {
	if d.resticBin == "" {
		return fmt.Errorf("restic binary not located on this agent")
	}
	if d.repoURL == "" || d.repoPassword == "" {
		return fmt.Errorf("repo credentials not configured (set repo_url + repo_password in agent.yaml or push via config.update)")
	}
	r := runner.New(runner.Config{
		ResticBin:    d.resticBin,
		RepoURL:      d.repoURL,
		RepoPassword: d.repoPassword,
	}, tx, time.Second)

	switch p.Kind {
	case api.JobBackup:
		// Agent.Args carries [paths...]. Excludes/tags are not yet
		// surfaced over the wire; they come with P2 schedule support.
		go func() {
			if err := r.RunBackup(ctx, p.JobID, p.Args, nil, nil); err != nil {
				slog.Warn("agent: backup job failed", "job_id", p.JobID, "err", err)
			}
		}()
	default:
		return fmt.Errorf("kind %q not implemented yet (Phase 2 lands the rest)", p.Kind)
	}
	return nil
}

func doEnroll(serverURL, token string, cfg *config.Config, agentVersion string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	snap, err := sysinfo.Collect(ctx, cfg.ResticPath)
	if err != nil {
		return fmt.Errorf("sysinfo: %w", err)
	}
	res, err := wsclient.Enroll(ctx, serverURL, wsclient.EnrollRequest{
		Token:         token,
		HostName:      snap.Hostname,
		OS:            snap.OS,
		Arch:          snap.Arch,
		AgentVersion:  agentVersion,
		ResticVersion: snap.ResticVersion,
	})
	if err != nil {
		return fmt.Errorf("enroll: %w", err)
	}
	cfg.ServerURL = serverURL
	cfg.HostID = res.HostID
	cfg.AgentToken = res.AgentToken
	cfg.CertPinSHA256 = res.CertPinSHA256
	if err := cfg.Save(); err != nil {
		return fmt.Errorf("save config: %w", err)
	}
	fmt.Fprintf(os.Stderr, "enrolled as host %s on %s\n", res.HostID, serverURL)
	return nil
}