llmgw/internal/cli/wipe.go

package cli

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/spf13/cobra"
)

// wipeUnits are the systemd units llmgw installs, stopped first so nothing
// writes while the rest goes.
var wipeUnits = []string{
	"llmgw-server.service", "llmgw-agent.service", "llmgw-mirror.service",
	firewallSyncPathUnit, "llmgw-firewall-sync.service", "llmgw-firewall.service",
	"llmgw-backup.timer", "llmgw-backup.service",
}

// wiper removes llmgw from a machine. root prefixes every path, for tests;
// run runs a command.
type wiper struct {
	root   string
	run    func(name string, args ...string) error
	store  bool
	models bool
	all    bool
}

// wipeStep is one thing to remove, said before it is done.
type wipeStep struct {
	what string
	do   func() error
}

func (w *wiper) path(p string) string { return filepath.Join(w.root, p) }

func (w *wiper) exists(p string) bool {
	_, err := os.Lstat(w.path(p))
	return err == nil
}

func (w *wiper) glob(pattern string) []string {
	m, _ := filepath.Glob(w.path(pattern))
	out := make([]string, len(m))
	for i, p := range m {
		out[i] = strings.TrimPrefix(p, w.root)
	}
	return out
}

// mentions is whether a file is there and says llmgw: a file of a name
// another tool also uses (keepalived's, etcd's) is removed only when llmgw
// wrote it.
func (w *wiper) mentions(p string) bool {
	b, err := os.ReadFile(w.path(p))
	return err == nil && strings.Contains(string(b), "llmgw")
}

func (w *wiper) remove(paths ...string) func() error {
	return func() error {
		var errs []error
		for _, p := range paths {
			if err := os.RemoveAll(w.path(p)); err != nil {
				errs = append(errs, err)
			}
		}
		return errors.Join(errs...)
	}
}

// plan is every step, in order, for what is on this machine.
func (w *wiper) plan() []wipeStep {
	var steps []wipeStep
	add := func(what string, do func() error) { steps = append(steps, wipeStep{what, do}) }
	files := func(what string, paths ...string) {
		var there []string
		for _, p := range paths {
			if w.exists(p) {
				there = append(there, p)
			}
		}
		if len(there) > 0 {
			add(what+": "+strings.Join(there, ", "), w.remove(there...))
		}
	}

	// The services, stopped and gone.
	var units []string
	for _, u := range wipeUnits {
		if w.exists("/etc/systemd/system/" + u) {
			units = append(units, u)
		}
	}
	if len(units) > 0 {
		add("stop and remove "+strings.Join(units, ", "), func() error {
			for _, u := range units {
				_ = w.run("systemctl", "disable", "--now", u)
				if err := os.Remove(w.path("/etc/systemd/system/" + u)); err != nil && !os.IsNotExist(err) {
					return err
				}
			}
			return w.run("systemctl", "daemon-reload")
		})
	}
	// Engines run as their own user, kept alive across the agent's restarts:
	// stopped by user, or they would keep their GPUs.
	if w.userExists("llmgw-engine") {
		add("stop the engines (every process of llmgw-engine)", func() error {
			_ = w.run("pkill", "-KILL", "-u", "llmgw-engine")
			return nil
		})
	}
	add("take llmgw's table out of the firewall (nft inet llmgw), if loaded", func() error {
		_ = w.run("nft", "delete", "table", "inet", "llmgw")
		return nil
	})
	if w.mentions("/etc/keepalived/keepalived.conf") {
		add("stop keepalived and remove llmgw's floating address (/etc/keepalived/keepalived.conf)", func() error {
			_ = w.run("systemctl", "disable", "--now", "keepalived")
			return w.remove("/etc/keepalived/keepalived.conf")()
		})
	}

	files("configuration, certificates, the authority and secrets", "/etc/llmgw")
	files("state: the admin token, a SQLite store, what the agent kept",
		"/var/lib/llmgw", "/var/lib/llmgw-agent")
	files("the sudo grant setup left", "/etc/sudoers.d/llmgw-setup")
	if w.exists("/usr/local/share/ca-certificates/llmgw-ca.crt") {
		add("the cluster's authority from the system's trust (llmgw-ca.crt)", func() error {
			if err := w.remove("/usr/local/share/ca-certificates/llmgw-ca.crt")(); err != nil {
				return err
			}
			_ = w.run("update-ca-certificates", "--fresh")
			return nil
		})
	}
	if w.exists("/etc/pki/ca-trust/source/anchors/llmgw-ca.pem") {
		add("the cluster's authority from the system's trust (llmgw-ca.pem)", func() error {
			if err := w.remove("/etc/pki/ca-trust/source/anchors/llmgw-ca.pem")(); err != nil {
				return err
			}
			_ = w.run("update-ca-trust")
			return nil
		})
	}

	if w.store || w.all {
		w.planStore(add, files)
	}
	if w.models || w.all {
		files("model weights and the mirror", "/srv/models", "/var/lib/llmgw-weights", "/var/lib/llmgw-models")
		files("installed engines", "/opt/llmgw")
	}
	if w.all {
		files("backups", DefaultBackupDir)
	}

	for _, u := range []string{"llmgw-engine", "llmgw"} {
		if w.userExists(u) {
			add("the user "+u, func() error { return w.run("userdel", u) })
		}
	}
	if w.all {
		files("the llmgw binary", "/usr/local/bin/llmgw", "/usr/local/bin/llmgw.prev")
	}
	return steps
}

// planStore is the store, when it is on this machine: a member of the
// replicated one (Patroni, etcd), or llmgw's database in a local PostgreSQL.
func (w *wiper) planStore(add func(string, func() error), files func(string, ...string)) {
	patroni := w.mentions("/etc/patroni/config.yml")
	if patroni || w.mentions("/etc/default/etcd") {
		add("stop Patroni and etcd", func() error {
			_ = w.run("systemctl", "disable", "--now", "patroni")
			_ = w.run("systemctl", "disable", "--now", "etcd")
			return nil
		})
	}
	if patroni {
		data := append(w.glob("/var/lib/postgresql/*/llmgw"), w.glob("/var/lib/pgsql/*/llmgw")...)
		files("the store's database on this member, and Patroni's configuration",
			append([]string{"/etc/patroni", "/var/lib/postgresql/.pgpass_patroni", "/root/llmgw-pg"}, data...)...)
	} else if w.exists("/usr/bin/psql") {
		add("drop the llmgw database and role from the local PostgreSQL, if there", func() error {
			_ = w.run("runuser", "-u", "postgres", "--", "psql", "-q",
				"-c", "DROP DATABASE IF EXISTS llmgw", "-c", "DROP ROLE IF EXISTS llmgw")
			return nil
		})
	}
	if w.mentions("/etc/default/etcd") {
		files("etcd's data and configuration", "/var/lib/etcd/default", "/etc/etcd/tls", "/etc/default/etcd")
	}
}

func (w *wiper) userExists(name string) bool {
	b, err := os.ReadFile(w.path("/etc/passwd"))
	return err == nil && strings.Contains("\n"+string(b), "\n"+name+":")
}

func newWipeCmd() *cobra.Command {
	var w wiper
	var yes bool
	c := &cobra.Command{
		Use:   "wipe",
		Short: "Remove llmgw from this machine: its services, configuration and state",
		Long: "Stops and removes llmgw's services, takes its firewall table out, and\n" +
			"deletes its configuration, certificates, the authority, secrets, state\n" +
			"and users — what server init, node add, server add and mirror init left —\n" +
			"so the machine can be set up again from nothing. It says everything it\n" +
			"will remove first, and asks for the machine's name.\n\n" +
			"Kept unless asked: model weights, the mirror and engines (--models), the\n" +
			"store's database when it is on this machine (--store), backups and the\n" +
			"binary (--all, which is all of these). Packages llmgw installed —\n" +
			"PostgreSQL, Patroni, etcd, nftables, the node exporter — stay.\n\n" +
			"A node wiped is then forgotten on the control plane with\n" +
			"'llmgw node remove <name>'.",
		Example: "  sudo llmgw wipe\n" +
			"  sudo llmgw wipe --models\n" +
			"  sudo llmgw wipe --all --yes",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if runtime.GOOS != "linux" {
				return fmt.Errorf("wipe removes llmgw from a Linux machine; this is %s", runtime.GOOS)
			}
			if os.Geteuid() != 0 {
				return errors.New("wipe removes services, users and files under /etc: run it with sudo")
			}
			out := cmd.OutOrStdout()
			w.run = func(name string, args ...string) error {
				return exec.CommandContext(cmd.Context(), name, args...).Run() //nolint:gosec // fixed commands
			}
			steps := w.plan()
			host, _ := os.Hostname()
			fmt.Fprintf(out, "On %s, this removes:\n", host)
			for _, s := range steps {
				fmt.Fprintf(out, "  - %s\n", s.what)
			}
			if !yes {
				if !isTerminal(out) {
					return errors.New("nothing removed: give --yes to wipe without being asked")
				}
				fmt.Fprintf(out, "\nType this machine's name (%s) to remove all of it: ", host)
				answer, _ := bufio.NewReader(cmd.InOrStdin()).ReadString('\n')
				if strings.TrimSpace(answer) != host {
					return errors.New("nothing removed")
				}
			}
			var failed []string
			for _, s := range steps {
				if err := s.do(); err != nil {
					failed = append(failed, s.what+": "+err.Error())
					fmt.Fprintf(out, "  !!    %s: %v\n", s.what, err)
					continue
				}
				fmt.Fprintf(out, "  ok    %s\n", s.what)
			}
			if len(failed) > 0 {
				return fmt.Errorf("%d step(s) failed; running wipe again finishes them", len(failed))
			}
			fmt.Fprintln(out, "\nllmgw is gone from this machine. If it was a node, forget it on the control plane: "+
				"llmgw node remove "+host)
			return nil
		},
	}
	f := c.Flags()
	f.BoolVar(&yes, "yes", false, "do not ask")
	f.BoolVar(&w.models, "models", false, "also remove model weights, the mirror's files and installed engines")
	f.BoolVar(&w.store, "store", false, "also remove the store's database here: a Patroni member's data, etcd's, or llmgw's database in a local PostgreSQL")
	f.BoolVar(&w.all, "all", false, "everything: --models, --store, backups and the binary")
	return c
}