felhom-agent/internal/wgtunnel/manager.go

package wgtunnel

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/netip"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"gitea.dooplex.hu/admin/felhom-agent/internal/hub"
	"gitea.dooplex.hu/admin/felhom-agent/internal/proxmox"
)

const (
	// Iface is the managed interface/unit name: wg-quick@wg-felhom.
	Iface = "wg-felhom"
	// confDest is the installed conf path — must match the FELHOM_WG sudoers entry EXACTLY.
	confDest = "/etc/wireguard/wg-felhom.conf"
	// unit is the systemd unit the sudoers allowlists.
	unit = "wg-quick@wg-felhom"

	markerName      = "registered.json"
	lastAppliedName = "last-applied.sha"
	stagedConfName  = "wg-felhom.conf"

	// registration backoff bounds (Scenario D: no hot loop).
	regBackoffMin = time.Minute
	regBackoffMax = 15 * time.Minute
)

// Registrar is the hub seam (satisfied by *hub.Client; tests inject a fake).
type Registrar interface {
	RegisterWG(ctx context.Context, pubkey string) (*hub.WGRegisterResponse, error)
}

// marker is the local registration record (<StateDir>/wg/registered.json). Its EXISTENCE is the
// registration gate: present → never register again except the pubkey-mismatch re-key path.
// KEPT on revocation (revoked stays revoked — doc 06 §3.5 completion).
type marker struct {
	Pubkey     string `json:"pubkey"`
	AssignedIP string `json:"assigned_ip"`
	Generation int64  `json:"generation"`
}

// Manager renders + applies the wg-felhom conf and drives wg-quick@wg-felhom through the
// narrow-sudoers runner (the lanresolver shape). All mutations go through `runner`; the only
// unprivileged execs are `systemctl is-active` (world-readable state) via isActive.
type Manager struct {
	runner   proxmox.Runner
	hub      Registrar
	stateDir string
	logger   *slog.Logger

	// injectable for tests
	isActive  func(ctx context.Context) bool
	wgPresent func() bool
	now       func() time.Time

	mu              sync.Mutex
	nextRegAt       time.Time
	regBackoff      time.Duration
	keyBroken       bool // corrupt key / partial state — loop idles until operator resolves
	brokenAnnounced bool
}

// NewManager builds a Manager. stateDir is the agent state dir (default /var/lib/felhom-agent —
// note the FELHOM_WG sudoers install entry hard-codes the staged path under it).
func NewManager(runner proxmox.Runner, registrar Registrar, stateDir string, logger *slog.Logger) *Manager {
	if logger == nil {
		logger = slog.Default()
	}
	return &Manager{
		runner:   runner,
		hub:      registrar,
		stateDir: stateDir,
		logger:   logger,
		isActive: func(ctx context.Context) bool {
			out, _ := exec.CommandContext(ctx, "systemctl", "is-active", unit).Output()
			return strings.TrimSpace(string(out)) == "active"
		},
		wgPresent: func() bool {
			_, err := os.Stat("/usr/bin/wg")
			return err == nil
		},
		now: time.Now,
	}
}

func (m *Manager) wgDir() string           { return filepath.Join(m.stateDir, "wg") }
func (m *Manager) markerPath() string      { return filepath.Join(m.wgDir(), markerName) }
func (m *Manager) lastAppliedPath() string { return filepath.Join(m.wgDir(), lastAppliedName) }
func (m *Manager) stagedConfPath() string  { return filepath.Join(m.wgDir(), stagedConfName) }

func (m *Manager) loadMarker() *marker {
	raw, err := os.ReadFile(m.markerPath())
	if err != nil {
		return nil
	}
	var mk marker
	if json.Unmarshal(raw, &mk) != nil || mk.Pubkey == "" {
		return nil
	}
	return &mk
}

func (m *Manager) writeMarker(mk marker) error {
	if err := os.MkdirAll(m.wgDir(), 0o700); err != nil {
		return err
	}
	raw, _ := json.Marshal(mk)
	return os.WriteFile(m.markerPath(), raw, 0o600)
}

// --- conf rendering (pure; every value strictly validated before it reaches the file) ---

var dnsNameRe = regexp.MustCompile(`^[A-Za-z0-9.-]+$`)

func validKeyB64(s string) error {
	if len(s) != 44 {
		return fmt.Errorf("key must be 44 base64 chars, got %d", len(s))
	}
	raw, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		return fmt.Errorf("key is not valid base64")
	}
	if len(raw) != 32 {
		return fmt.Errorf("key decodes to %d bytes, want 32", len(raw))
	}
	return nil
}

// renderConf builds the wg-felhom.conf content from the hub block + the local private key.
// Client-side constants per doc 06 §4: MTU 1420, AllowedIPs = pbs_tunnel_ip/32 (the tunnel
// carries ONLY box→PBS traffic), PersistentKeepalive 25. All inputs validated — nothing
// user-controlled is interpolatable (strict charsets, netip parses).
func renderConf(block *hub.WireWireguard, privB64 string) (string, error) {
	if err := validKeyB64(privB64); err != nil {
		return "", fmt.Errorf("wgtunnel: private key: %w", err)
	}
	if err := validKeyB64(block.Endpoint.ServerPubkey); err != nil {
		return "", fmt.Errorf("wgtunnel: server_pubkey: %w", err)
	}
	addr, err := netip.ParsePrefix(block.AssignedIP)
	if err != nil || addr.Bits() != 32 {
		return "", fmt.Errorf("wgtunnel: assigned_ip %q is not an ip/32", block.AssignedIP)
	}
	pbsIP, err := netip.ParseAddr(block.Endpoint.PBSTunnelIP)
	if err != nil {
		return "", fmt.Errorf("wgtunnel: pbs_tunnel_ip %q is not an address", block.Endpoint.PBSTunnelIP)
	}
	if !dnsNameRe.MatchString(block.Endpoint.DNSName) {
		return "", fmt.Errorf("wgtunnel: endpoint dns_name has invalid characters")
	}
	if block.Endpoint.WGPort < 1 || block.Endpoint.WGPort > 65535 {
		return "", fmt.Errorf("wgtunnel: wg_port %d out of range", block.Endpoint.WGPort)
	}

	var b strings.Builder
	b.WriteString("# felhom offsite tunnel — agent-managed (S3); DO NOT EDIT\n")
	b.WriteString("[Interface]\n")
	fmt.Fprintf(&b, "PrivateKey = %s\n", privB64)
	fmt.Fprintf(&b, "Address = %s\n", block.AssignedIP)
	b.WriteString("MTU = 1420\n\n")
	b.WriteString("[Peer]\n")
	fmt.Fprintf(&b, "PublicKey = %s\n", block.Endpoint.ServerPubkey)
	fmt.Fprintf(&b, "Endpoint = %s:%d\n", block.Endpoint.DNSName, block.Endpoint.WGPort)
	fmt.Fprintf(&b, "AllowedIPs = %s/32\n", pbsIP)
	b.WriteString("PersistentKeepalive = 25\n")
	return b.String(), nil
}

// --- the state machine (doc 06 §3.3/§3.5 + spec §7/§8) ---

// Apply reconciles local reality toward (fetched, block). fetched=false means "no desired-state
// data seen yet this process" — NEVER a teardown signal (teardown only on a PRESENT desired-state
// without the block).
func (m *Manager) Apply(ctx context.Context, fetched bool, block *hub.WireWireguard) {
	m.mu.Lock()
	defer m.mu.Unlock()

	keyExists := false
	if _, err := os.Stat(KeyFilePath(m.stateDir)); err == nil {
		keyExists = true
	}
	mk := m.loadMarker()

	// Partial state: marker without key. Never guess — a fresh keygen here would silently orphan
	// the hub-registered identity. Operator repairs (delete marker → fresh keygen+register).
	if mk != nil && !keyExists {
		m.announceBroken("wgtunnel: registration marker exists but the key file is missing — idling until the operator resolves (delete the marker for a fresh keygen+register)")
		return
	}

	pub, created, err := EnsureKey(m.stateDir)
	if err != nil {
		m.announceBroken("wgtunnel: " + err.Error())
		return
	}
	m.keyBroken, m.brokenAnnounced = false, false
	if created {
		m.logger.Info("wgtunnel: generated new WG keypair", "pubkey", pub)
	}

	if mk == nil {
		// Not registered (by local knowledge). Adopt if the hub already knows this exact key —
		// the lost-marker-kept-key case; no re-register.
		if fetched && block != nil && block.Pubkey == pub {
			if err := m.writeMarker(marker{Pubkey: pub, AssignedIP: block.AssignedIP}); err != nil {
				m.logger.Error("wgtunnel: writing adoption marker", "err", err)
				return
			}
			m.logger.Info("wgtunnel: adopted existing hub registration (marker restored)", "pubkey", pub)
			mk = m.loadMarker()
		} else {
			// One-shot registration (backoff-gated). Registration is gated ONLY on marker
			// absence: a revoked box (marker present, block gone) never lands here.
			m.registerLocked(ctx, pub)
			return
		}
	}

	if !fetched {
		return // no desired data yet — keep last-applied state untouched
	}

	if block == nil {
		// PRESENT desired-state without the block = the hub revoked us. Stop + disable, KEEP the
		// marker (revoked stays revoked; the operator re-adds the peer using the reported pubkey).
		m.teardownLocked(ctx)
		return
	}

	if block.Pubkey != pub {
		// The hub's block is for another key (DR from escrowed key on a re-provisioned box, or
		// stale hub state) — the S2 re-key-in-place path. Bounded by the same backoff.
		m.logger.Warn("wgtunnel: desired block pubkey differs from local key — re-registering (re-key-in-place)",
			"local", pub, "desired", block.Pubkey)
		m.registerLocked(ctx, pub)
		return // the bumped generation brings a corrected block on the next fetch
	}

	m.ensureTunnelLocked(ctx, block)
}

// announceBroken logs a broken-state error ONCE (then idles quietly until state changes).
func (m *Manager) announceBroken(msg string) {
	if !m.brokenAnnounced {
		m.logger.Error(msg)
		m.brokenAnnounced = true
	}
	m.keyBroken = true
}

// registerLocked runs one backoff-gated RegisterWG and writes the marker on success.
func (m *Manager) registerLocked(ctx context.Context, pub string) {
	if m.now().Before(m.nextRegAt) {
		return
	}
	resp, err := m.hub.RegisterWG(ctx, pub)
	if err != nil {
		if m.regBackoff == 0 {
			m.regBackoff = regBackoffMin
		} else if m.regBackoff *= 2; m.regBackoff > regBackoffMax {
			m.regBackoff = regBackoffMax
		}
		m.nextRegAt = m.now().Add(m.regBackoff)
		m.logger.Warn("wgtunnel: registration failed — will retry", "err", err, "backoff", m.regBackoff)
		return
	}
	m.regBackoff, m.nextRegAt = 0, time.Time{}
	if err := m.writeMarker(marker{Pubkey: pub, AssignedIP: resp.AssignedIP, Generation: resp.Generation}); err != nil {
		m.logger.Error("wgtunnel: registration succeeded but marker write failed", "err", err)
		return
	}
	m.logger.Info("wgtunnel: registered with hub", "pubkey", pub,
		"assigned_ip", resp.AssignedIP, "existed", resp.Existed, "generation", resp.Generation, "sync", resp.Sync)
}

// teardownLocked disables the tunnel once (idempotent via the last-applied hash file + a live
// service check). NO re-registration follows — the marker stays.
func (m *Manager) teardownLocked(ctx context.Context) {
	applied := false
	if _, err := os.Stat(m.lastAppliedPath()); err == nil {
		applied = true
	}
	if !applied && !m.isActive(ctx) {
		return // already torn down — nothing to exec (Scenario C's quiet ticks)
	}
	m.logger.Warn("wgtunnel: desired-state no longer carries the wireguard block — hub revoked this peer; disabling the tunnel (marker kept; NO re-registration)")
	if _, errOut, err := m.runner.Run(ctx, "systemctl", "disable", "--now", unit); err != nil {
		m.logger.Error("wgtunnel: disable failed", "err", err, "stderr", strings.TrimSpace(string(errOut)))
		return
	}
	os.Remove(m.lastAppliedPath())
}

// ensureTunnelLocked renders + applies the conf and keeps the service alive (Scenario B).
func (m *Manager) ensureTunnelLocked(ctx context.Context, block *hub.WireWireguard) {
	priv, err := readPrivateKeyB64(m.stateDir)
	if err != nil {
		m.announceBroken("wgtunnel: " + err.Error())
		return
	}
	conf, err := renderConf(block, priv)
	if err != nil {
		m.logger.Error("wgtunnel: refusing to apply invalid desired block", "err", err)
		return
	}
	sum := sha256.Sum256([]byte(conf))
	hash := hex.EncodeToString(sum[:])
	last, _ := os.ReadFile(m.lastAppliedPath())
	active := m.isActive(ctx)

	if string(last) == hash {
		if active {
			return // steady state: zero execs
		}
		// self-heal: conf is current but the service is down (crash/manual stop)
		m.logger.Info("wgtunnel: service inactive with current conf — re-enabling (self-heal)")
		if err := m.ensureToolsLocked(ctx); err != nil {
			return
		}
		if _, errOut, err := m.runner.Run(ctx, "systemctl", "enable", "--now", unit); err != nil {
			m.logger.Error("wgtunnel: enable failed", "err", err, "stderr", strings.TrimSpace(string(errOut)))
		}
		return
	}

	// Conf changed (first apply, endpoint re-key, ip change): stage → install → enable/restart.
	if err := m.ensureToolsLocked(ctx); err != nil {
		return
	}
	if err := os.MkdirAll(m.wgDir(), 0o700); err != nil {
		m.logger.Error("wgtunnel: state dir", "err", err)
		return
	}
	if err := os.WriteFile(m.stagedConfPath(), []byte(conf), 0o600); err != nil {
		m.logger.Error("wgtunnel: staging conf", "err", err)
		return
	}
	// argv matches the FELHOM_WG sudoers entry exactly (fixed source + dest).
	if _, errOut, err := m.runner.Run(ctx, "install", "-o", "root", "-g", "root", "-m", "0600", "--", m.stagedConfPath(), confDest); err != nil {
		m.logger.Error("wgtunnel: conf install failed", "err", err, "stderr", strings.TrimSpace(string(errOut)))
		return
	}
	verb := []string{"enable", "--now", unit}
	if active {
		verb = []string{"restart", unit} // conf change on a running tunnel → restart (never reload)
	}
	if _, errOut, err := m.runner.Run(ctx, "systemctl", verb...); err != nil {
		m.logger.Error("wgtunnel: systemctl "+verb[0]+" failed", "err", err, "stderr", strings.TrimSpace(string(errOut)))
		return
	}
	if err := os.WriteFile(m.lastAppliedPath(), []byte(hash), 0o600); err != nil {
		m.logger.Error("wgtunnel: recording applied hash", "err", err)
		return
	}
	m.logger.Info("wgtunnel: tunnel conf applied", "endpoint",
		fmt.Sprintf("%s:%d", block.Endpoint.DNSName, block.Endpoint.WGPort), "assigned_ip", block.AssignedIP, "action", verb[0])
}

// ensureToolsLocked installs wireguard-tools once when absent (dnsmasq-install precedent).
func (m *Manager) ensureToolsLocked(ctx context.Context) error {
	if m.wgPresent() {
		return nil
	}
	m.logger.Info("wgtunnel: wireguard-tools absent — installing")
	if out, errOut, err := m.runner.Run(ctx, "apt-get", "install", "-y", "-q", "wireguard-tools"); err != nil {
		m.logger.Error("wgtunnel: apt-get install wireguard-tools failed",
			"err", err, "out", strings.TrimSpace(string(errOut))+strings.TrimSpace(string(out)))
		return err
	}
	return nil
}

// Status builds the heartbeat report stanza (doc 06 §4.6). Read-only: it NEVER creates keys or
// registers. The handshake read is the package's single wg invocation — `latest-handshakes`
// only (NEVER `dump`, whose interface line carries the private key).
func (m *Manager) Status(ctx context.Context) *hub.WireguardStatus {
	m.mu.Lock()
	defer m.mu.Unlock()

	st := &hub.WireguardStatus{}
	if raw, err := os.ReadFile(KeyFilePath(m.stateDir)); err == nil {
		if priv, derr := decodeKey(raw); derr == nil {
			if pub, perr := derivePublic(priv); perr == nil {
				st.Pubkey = pub
			}
		}
	}
	if mk := m.loadMarker(); mk != nil {
		st.Registered = true
		st.AssignedIP = mk.AssignedIP
	}
	st.Active = m.isActive(ctx)
	if st.Active {
		if out, _, err := m.runner.Run(ctx, "wg", "show", Iface, "latest-handshakes"); err == nil {
			if age, ok := parseHandshakeAge(string(out), m.now()); ok {
				st.LastHandshakeAgeS = &age
			}
		}
	}
	if st.Pubkey == "" && !st.Registered && !st.Active {
		return st // still a valid (empty-ish) stanza; caller decides whether to attach
	}
	return st
}

// parseHandshakeAge parses `wg show <if> latest-handshakes` output ("<peer-pub>\t<epoch>").
// epoch 0 = no handshake yet → not ok (nil in the report; nil ≠ 0).
func parseHandshakeAge(out string, now time.Time) (int64, bool) {
	for _, line := range strings.Split(strings.TrimSpace(out), "\n") {
		f := strings.Fields(line)
		if len(f) != 2 {
			continue
		}
		epoch, err := strconv.ParseInt(f[1], 10, 64)
		if err != nil || epoch <= 0 {
			continue
		}
		age := now.Unix() - epoch
		if age < 0 {
			age = 0
		}
		return age, true
	}
	return 0, false
}