ssh-netbox-wrapper/cmd/netssh/multi.go

package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/cache"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/config"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/multiplex"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/netbox"
)

// maxHostsWithoutConfirm guards against opening hundreds of panes by accident.
const maxHostsWithoutConfirm = 24

// multiFlags are the options shared by netssh multi and the TUI.
type multiFlags struct {
	kind, user, jump, multiplexer string
	noJump, sync, dryRun, yes     bool
}

func multiCmd() *cobra.Command {
	var f multiFlags
	var useTmux, useZellij bool
	cmd := &cobra.Command{
		Use:   "multi [flags] HOST|tag:TAG|cluster:NAME|@GROUP...",
		Short: "Open several hosts side by side in tmux or zellij",
		Long: `Open one pane per host in a tmux window or zellij tab. Hosts are given as
names (shortcuts, instance/host, static hosts) or selectors:

  tag:TAG        all hosts with this NetBox tag, and static hosts with this tag
  cluster:NAME   devices of a NetBox virtualization cluster (--kind vm|all for VMs)
  @GROUP         a group from the config (groups:)

Broadcast typing to all panes of a window with --sync, or toggle it later with
netssh sync. Inside tmux or zellij the panes open in the current session.`,
		Example: `  netssh multi web01 web02 db01
  netssh multi cluster:pve-fsn1
  netssh multi --kind vm cluster:pve-fsn1
  netssh multi --sync tag:k8s-worker
  netssh multi @k8s`,
		Args: cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if useTmux && useZellij {
				return errors.New("--tmux and --zellij are mutually exclusive")
			}
			if useTmux {
				f.multiplexer = multiplex.Tmux
			} else if useZellij {
				f.multiplexer = multiplex.Zellij
			}
			return runMulti(loadConfigOrSetup(), args, f)
		},
		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
			cfg, err := config.Load()
			if err != nil {
				return nil, cobra.ShellCompDirectiveError
			}
			return multiCompletions(cfg, toComplete)
		},
	}
	cmd.Flags().StringVar(&f.kind, "kind", "", "Host kind for tag: and cluster: — device, vm, or all (default: all for tags, device for clusters)")
	cmd.Flags().BoolVar(&f.sync, "sync", false, "Start with typing broadcast to all panes")
	cmd.Flags().BoolVar(&useTmux, "tmux", false, "Use tmux")
	cmd.Flags().BoolVar(&useZellij, "zellij", false, "Use zellij")
	cmd.Flags().StringVarP(&f.user, "user", "l", "", "Login user for every host")
	cmd.Flags().StringVar(&f.jump, "jump", "", "Jump host for every host")
	cmd.Flags().BoolVar(&f.noJump, "no-jump", false, "Disable jump hosts")
	cmd.Flags().BoolVar(&f.dryRun, "dry-run", false, "Print the hosts and multiplexer commands")
	cmd.Flags().BoolVarP(&f.yes, "yes", "y", false, fmt.Sprintf("Open more than %d hosts without asking", maxHostsWithoutConfirm))
	_ = cmd.RegisterFlagCompletionFunc("kind", cobra.FixedCompletions([]cobra.Completion{"device", "vm", "all"}, cobra.ShellCompDirectiveNoFileComp))
	return cmd
}

func syncCmd() *cobra.Command {
	var pane bool
	cmd := &cobra.Command{
		Use:   "sync [--pane]",
		Short: "Toggle typing into all panes of the current tmux window or zellij tab",
		Long: `Toggle broadcast in the current tmux window or zellij tab. With --pane,
exclude the current tmux pane from broadcast, or include it again.

In windows opened by netssh multi, tmux also has keys for this:
prefix S (all panes) and prefix e (this pane); see multi.sync_key and
multi.solo_key. In zellij, use Ctrl+t then s.`,
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			state, err := multiplex.ToggleSync(multiplex.System{}, pane)
			if err != nil {
				return err
			}
			fmt.Fprintln(os.Stderr, "netssh: "+state)
			return nil
		},
	}
	cmd.Flags().BoolVar(&pane, "pane", false, "Exclude the current tmux pane from broadcast, or include it again")
	return cmd
}

// runMulti expands selectors and opens one pane per host.
func runMulti(cfg *config.Config, selectors []string, f multiFlags) error {
	if f.jump != "" && f.noJump {
		return errors.New("--jump and --no-jump are mutually exclusive")
	}
	switch f.kind {
	case "", "device", "vm", "all":
	default:
		return fmt.Errorf("--kind must be device, vm, or all")
	}
	c := newCache(cfg)
	_ = c.Load()
	targets, warnings, err := expandSelectors(cfg, c, selectors, f.kind)
	for _, warning := range warnings {
		fmt.Fprintln(os.Stderr, "netssh: warning: "+warning)
	}
	if err != nil {
		return err
	}
	if len(targets) > maxHostsWithoutConfirm && !f.yes && !f.dryRun {
		return fmt.Errorf("%d hosts selected; pass --yes to open them all", len(targets))
	}

	var runner multiplex.Runner = multiplex.System{}
	if f.dryRun {
		fmt.Printf("# %d hosts: %s\n", len(targets), strings.Join(targets, " "))
		runner = &multiplex.DryRun{W: os.Stdout}
	}
	requested := f.multiplexer
	if requested == "" {
		requested = cfg.Multi.Multiplexer
	}
	kind, err := multiplex.Detect(requested, runner)
	if err != nil {
		return err
	}
	if !f.dryRun {
		// Checked once here; the panes skip it (NETSSH_NO_REFRESH).
		maybeBackgroundRefresh(cfg, c)
	}

	panes := make([]multiplex.Pane, 0, len(targets))
	for _, target := range targets {
		panes = append(panes, multiplex.Pane{Title: target, Argv: paneArgv(target, f)})
	}
	layoutDir := filepath.Dir(cfg.Cache.Path)
	if err := os.MkdirAll(layoutDir, 0o700); err != nil {
		return err
	}
	syncKey, soloKey := cfg.Multi.Keys()
	return multiplex.Open(kind, panes, multiplex.Options{
		SyncKey:   syncKey,
		SoloKey:   soloKey,
		Label:     multiLabel(selectors),
		Sync:      f.sync,
		PerWindow: cfg.Multi.PanesPerWindowOrDefault(),
		LayoutDir: layoutDir,
		Hint:      func(s string) { fmt.Fprintln(os.Stderr, "netssh: "+s) },
	}, runner)
}

// paneArgv runs netssh for one host. Panes of an existing tmux or zellij
// session inherit that server's environment, so the paths that select the
// config and cache are passed explicitly; NETSSH_* values are not, because
// arguments are visible to other users (tokens).
func paneArgv(target string, f multiFlags) []string {
	self, err := os.Executable()
	if err != nil {
		self = "netssh"
	}
	env, err := exec.LookPath("env")
	if err != nil {
		env = "/usr/bin/env"
	}
	argv := []string{env, "NETSSH_NO_REFRESH=1"}
	for _, key := range []string{"HOME", "XDG_CONFIG_HOME", "XDG_CACHE_HOME"} {
		if value, ok := os.LookupEnv(key); ok {
			argv = append(argv, key+"="+value)
		}
	}
	argv = append(argv, self)
	switch {
	case f.jump != "":
		argv = append(argv, "--jump", f.jump)
	case f.noJump:
		argv = append(argv, "--no-jump")
	}
	if f.user != "" {
		argv = append(argv, "-l", f.user)
	}
	return append(argv, target)
}

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

// multiLabel names the window, session, or tab after a single selector.
func multiLabel(selectors []string) string {
	if len(selectors) != 1 {
		return "multi"
	}
	s := selectors[0]
	for _, prefix := range []string{"@", "tag:", "cluster:"} {
		s = strings.TrimPrefix(s, prefix)
	}
	if s = strings.Trim(unsafeLabel.ReplaceAllString(s, "-"), "-"); s == "" {
		return "multi"
	}
	return s
}

// expander resolves selectors into an ordered list of unique host targets.
type expander struct {
	cfg      *config.Config
	cache    *cache.Cache
	kind     string
	client   *netbox.Client
	seen     map[string]bool
	visiting map[string]bool
	targets  []string
	warnings []string
}

func expandSelectors(cfg *config.Config, c *cache.Cache, selectors []string, kind string) ([]string, []string, error) {
	e := &expander{cfg: cfg, cache: c, kind: kind, seen: map[string]bool{}, visiting: map[string]bool{}}
	for _, selector := range selectors {
		if err := e.expand(selector); err != nil {
			return nil, e.warnings, err
		}
	}
	if len(e.targets) == 0 {
		return nil, e.warnings, errors.New("no hosts selected")
	}
	return e.targets, e.warnings, nil
}

func (e *expander) add(target string) {
	key := strings.ToLower(target)
	if !e.seen[key] {
		e.seen[key] = true
		e.targets = append(e.targets, target)
	}
}

func (e *expander) expand(selector string) error {
	switch {
	case strings.HasPrefix(selector, "@"):
		return e.group(strings.ToLower(strings.TrimPrefix(selector, "@")))
	case strings.HasPrefix(selector, "tag:"):
		return e.query(selector, netbox.SearchOptions{Tag: strings.TrimPrefix(selector, "tag:")}, "all")
	case strings.HasPrefix(selector, "cluster:"):
		return e.query(selector, netbox.SearchOptions{Cluster: strings.TrimPrefix(selector, "cluster:")}, "device")
	case selector == "":
		return nil
	}
	e.add(selector)
	return nil
}

func (e *expander) group(name string) error {
	members, ok := e.cfg.Groups[name]
	if !ok {
		names := make([]string, 0, len(e.cfg.Groups))
		for group := range e.cfg.Groups {
			names = append(names, "@"+group)
		}
		sort.Strings(names)
		if len(names) == 0 {
			return fmt.Errorf("unknown group @%s; no groups are configured (groups:)", name)
		}
		return fmt.Errorf("unknown group @%s; available: %s", name, strings.Join(names, ", "))
	}
	if e.visiting[name] {
		return fmt.Errorf("group @%s contains itself", name)
	}
	e.visiting[name] = true
	defer delete(e.visiting, name)
	for _, member := range members {
		if err := e.expand(member); err != nil {
			return err
		}
	}
	return nil
}

// query expands tag: and cluster: through NetBox, plus static hosts for tags;
// if NetBox fails, the cache answers with a warning.
func (e *expander) query(selector string, opts netbox.SearchOptions, defaultKind string) error {
	kind := e.kind
	if kind == "" {
		kind = defaultKind
	}
	if kind != "all" {
		opts.Kind = kind
	}
	before := len(e.targets)
	if opts.Tag != "" && (kind == "all" || kind == config.StaticKind) {
		for _, host := range e.cfg.Hosts {
			if hasTag(host.Tags, opts.Tag) {
				e.add(host.Name)
			}
		}
	}
	if e.cfg.HasNetBox() {
		if e.client == nil {
			e.client = newNetBoxClient(e.cfg)
		}
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		entries, err := e.client.SearchAll(ctx, "", opts)
		cancel()
		if err != nil {
			e.warnings = append(e.warnings, fmt.Sprintf("%s: NetBox unavailable, using the cache: %v", selector, err))
			entries = e.fromCache(opts)
		} else {
			sort.Slice(entries, func(i, j int) bool { return entries[i].Target() < entries[j].Target() })
		}
		for _, entry := range entries {
			e.add(entry.Target())
		}
	}
	if len(e.targets) == before {
		hint := ""
		if opts.Cluster != "" && e.kind == "" {
			hint = " (only devices by default; try --kind vm or --kind all)"
		}
		return fmt.Errorf("%s matches no hosts%s", selector, hint)
	}
	return nil
}

func (e *expander) fromCache(opts netbox.SearchOptions) []netbox.HostEntry {
	var out []netbox.HostEntry
	for _, entry := range e.cache.All() {
		if entry.Static || opts.Kind != "" && entry.Kind != opts.Kind {
			continue
		}
		if opts.Tag != "" && !hasTag(entry.Tags, opts.Tag) || opts.Cluster != "" && !strings.EqualFold(entry.Cluster, opts.Cluster) {
			continue
		}
		source, name := netbox.SplitTarget(entry.Name)
		out = append(out, netbox.HostEntry{Source: source, Name: name, Kind: entry.Kind})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Target() < out[j].Target() })
	return out
}

// hasTag compares a tag by name or slug, like ssh.rules.
func hasTag(tags []string, want string) bool {
	for _, tag := range tags {
		if strings.EqualFold(tag, want) || strings.EqualFold(strings.Join(strings.Fields(strings.ToLower(tag)), "-"), want) {
			return true
		}
	}
	return false
}

// multiCompletions offers hosts, @groups, and tag:/cluster: values from the cache.
func multiCompletions(cfg *config.Config, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
	c := newCache(cfg)
	_ = c.Load()
	lower := strings.ToLower(toComplete)
	var out []string
	switch {
	case strings.HasPrefix(lower, "tag:"), strings.HasPrefix(lower, "cluster:"):
		prefix, value, _ := strings.Cut(lower, ":")
		values := map[string]bool{}
		for _, entry := range c.All() {
			if prefix == "tag" {
				for _, tag := range entry.Tags {
					values[strings.ToLower(tag)] = true
				}
			} else if entry.Cluster != "" {
				values[entry.Cluster] = true
			}
		}
		for v := range values {
			if strings.HasPrefix(strings.ToLower(v), value) {
				out = append(out, prefix+":"+v)
			}
		}
	case strings.HasPrefix(lower, "@"):
		for group := range cfg.Groups {
			if strings.HasPrefix("@"+group, lower) {
				out = append(out, "@"+group)
			}
		}
	default:
		out = completeHosts(toComplete)
		var prefixes []string
		for _, p := range []string{"tag:", "cluster:"} {
			if toComplete != "" && strings.HasPrefix(p, lower) {
				prefixes = append(prefixes, p)
			}
		}
		if len(out) == 0 && len(prefixes) > 0 {
			return prefixes, cobra.ShellCompDirectiveNoSpace | cobra.ShellCompDirectiveNoFileComp
		}
		for group := range cfg.Groups {
			if toComplete == "" {
				out = append(out, "@"+group)
			}
		}
	}
	sort.Strings(out)
	return out, cobra.ShellCompDirectiveNoFileComp
}