ssh-netbox-wrapper/internal/setup/wizard.go

package setup

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"charm.land/huh/v2"

	"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/hook"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/storage"
)

// RunWizard runs the interactive configurator, pre-filled with any existing cfg values.
// It saves the result to the config file on success.
func RunWizard(cfg *config.Config) error {
	sources, err := configureInstances(cfg)
	if err != nil {
		return cancelled(err)
	}
	work := prepare(*cfg, sources)
	if err := runMenu(&work); err != nil {
		return cancelled(err)
	}

	if err := save(work); err != nil {
		return err
	}
	loaded, err := config.Load()
	if err != nil {
		return err
	}
	c := cache.New(loaded.CachePath(), loaded.Cache.TTL)
	if err := c.Load(); err != nil {
		return fmt.Errorf("updating aliases: %w", err)
	}
	if err := hook.WriteAliasFiles(hook.BuildAliasEntries(c.All(), *loaded)); err != nil {
		return err
	}
	if loaded.Shell.Alias != "" {
		shell := filepath.Base(os.Getenv("SHELL"))
		profile, err := hook.ProfilePath(shell)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Run 'netssh hook install --shell bash|zsh|fish' to activate your alias.\n")
			return nil
		}
		if _, err := hook.Install(profile, shell); err != nil {
			return err
		}
		fmt.Fprintln(os.Stderr, hook.ReloadNote(profile))
	}
	return nil
}

func cancelled(err error) error {
	if errors.Is(err, huh.ErrUserAborted) {
		fmt.Fprintln(os.Stderr, "Setup cancelled.")
		os.Exit(0)
	}
	return err
}

// prepare starts from the loaded config so settings without a menu entry survive
// a rewrite, migrates a legacy single instance, and fills in usable defaults.
func prepare(cfg config.Config, sources []config.NetBoxConfig) config.Config {
	cfg.NetBox = config.NetBoxConfig{}
	cfg.NetBoxes = sources
	if len(cfg.Resolver.Strategies) == 0 {
		cfg.Resolver.Strategies = []string{"primary_ip"}
	}
	if cfg.Cache.TTL == 0 {
		cfg.Cache.TTL = 3600
	}
	return cfg
}

func runMenu(cfg *config.Config) error {
	for {
		action := "save"
		err := huh.NewForm(huh.NewGroup(huh.NewSelect[string]().
			Title("netssh configuration").
			Description("Choose a section to edit, then save.").
			Options(
				huh.NewOption("SSH defaults — user, port", "defaults"),
				huh.NewOption("SSH options — jump host, host keys, target, -o", "ssh"),
				huh.NewOption(fmt.Sprintf("SSH rules — per tag, instance, kind, name (%d)", len(cfg.SSH.Rules)), "rules"),
				huh.NewOption("Resolver & cache", "resolver"),
				huh.NewOption("Shortcuts & command alias", "shortcuts"),
				huh.NewOption("Completion", "completion"),
				huh.NewOption("Save and exit", "save"),
			).
			Value(&action))).Run()
		if err != nil {
			return err
		}
		switch action {
		case "defaults":
			err = editSSHDefaults(cfg)
		case "ssh":
			cfg.SSH.SSHOptions, err = editSSHOptions("SSH options (all hosts)", cfg.SSH.SSHOptions)
		case "rules":
			var rules []config.SSHRule
			if rules, err = configureRules(cfg.SSH.Rules, cfg.NetBoxes); err == nil {
				cfg.SSH.Rules = rules
			}
		case "resolver":
			err = editResolverCache(cfg)
		case "shortcuts":
			err = editShortcuts(cfg)
		case "completion":
			err = editCompletion(cfg)
		case "save":
			return nil
		}
		if err != nil {
			return err
		}
	}
}

func editSSHDefaults(cfg *config.Config) error {
	defaultUser := cfg.SSH.DefaultUser
	defaultPort := ""
	if cfg.SSH.DefaultPort > 0 {
		defaultPort = strconv.Itoa(cfg.SSH.DefaultPort)
	}
	err := huh.NewForm(huh.NewGroup(
		huh.NewInput().
			Title("Default SSH user").
			Description("Leave empty to use your system user ($USER).").
			Value(&defaultUser),
		huh.NewInput().
			Title("Default SSH port").
			Description("Leave empty to use the standard port (22).").
			Placeholder("22").
			Value(&defaultPort).
			Validate(func(s string) error {
				if strings.TrimSpace(s) == "" {
					return nil
				}
				n, err := strconv.Atoi(strings.TrimSpace(s))
				if err != nil || n < 1 || n > 65535 {
					return errors.New("must be a port number between 1 and 65535")
				}
				return nil
			}),
	).Title("SSH defaults")).Run()
	if err != nil {
		return err
	}
	cfg.SSH.DefaultUser = strings.TrimSpace(defaultUser)
	cfg.SSH.DefaultPort, _ = strconv.Atoi(strings.TrimSpace(defaultPort))
	return nil
}

func editResolverCache(cfg *config.Config) error {
	strategiesRaw := strings.Join(cfg.Resolver.Strategies, ", ")
	subnets := strings.Join(cfg.Resolver.ManagementSubnets, ", ")
	interfaceName := cfg.Resolver.InterfaceName
	cacheTTL := strconv.Itoa(cfg.Cache.TTL)
	cachePath := cfg.Cache.Path
	err := huh.NewForm(huh.NewGroup(
		huh.NewInput().
			Title("Resolver strategies").
			Description("Comma-separated, in priority order. First match wins.\nAvailable: primary_ip, management_subnet, interface_name").
			Placeholder("primary_ip, management_subnet").
			Value(&strategiesRaw).
			Validate(validateStrategies),
		huh.NewInput().
			Title("Management subnets").
			Description("Comma-separated CIDRs, e.g. 10.0.0.0/8, 192.168.0.0/16\nOnly used when management_subnet strategy is active.").
			Value(&subnets),
		huh.NewInput().
			Title("Interface name").
			Description("Only used when interface_name strategy is active.").
			Placeholder("eth0").
			Value(&interfaceName),
		huh.NewInput().
			Title("Cache TTL (seconds)").
			Value(&cacheTTL).
			Validate(func(s string) error {
				if _, err := strconv.Atoi(s); err != nil {
					return errors.New("must be a number")
				}
				return nil
			}),
		huh.NewInput().
			Title("Cache file").
			Description("Where resolved hosts are stored. Leave empty for the default location.").
			Value(&cachePath),
	).Title("Resolver & cache")).Run()
	if err != nil {
		return err
	}
	cfg.Resolver.Strategies = parseStrategies(strategiesRaw)
	cfg.Resolver.ManagementSubnets = splitList(subnets)
	cfg.Resolver.InterfaceName = strings.TrimSpace(interfaceName)
	cfg.Cache.TTL, _ = strconv.Atoi(cacheTTL)
	cfg.Cache.Path = strings.TrimSpace(cachePath)
	return nil
}

func editShortcuts(cfg *config.Config) error {
	shortcutDomains := strings.Join(cfg.Shortcuts.Domains, ", ")
	stripHyphens := cfg.Shortcuts.StripHyphens
	rewrites := formatRewrites(cfg.Shortcuts.Rewrites)
	commandAlias := cfg.Shell.Alias
	err := huh.NewForm(huh.NewGroup(
		huh.NewInput().
			Title("Domain suffixes").
			Description("Comma-separated suffixes stripped for shortcuts, e.g. .example.com, .example.de\nAllows typing 'web01' instead of 'web01.example.com'.").
			Placeholder(".example.com").
			Value(&shortcutDomains),
		huh.NewConfirm().
			Title("Strip hyphens").
			Description("When enabled, fsn1-web01.example.com can be accessed as fsn1web01.\nOnly works for hosts already in the cache.").
			Value(&stripHyphens),
		huh.NewText().
			Title("Rewrite rules").
			Description("One rule per line: REGEX => REPLACEMENT, applied after the domain suffix is removed.\n"+
				"Example: ^([a-z]+)[0-9]+- => $1  turns fsn1-bx1.example.com into fsnbx1.\n"+
				"Check rules with: netssh alias --test HOSTNAME").
			Lines(3).
			Value(&rewrites).
			Validate(func(s string) error {
				_, err := parseRewrites(s)
				return err
			}),
		huh.NewInput().
			Title("Command alias").
			Description("Optional alias for netssh itself, e.g. j. Leave empty to disable.").
			Placeholder("j").
			Value(&commandAlias).
			Validate(func(s string) error { return config.ValidateCommandAlias(strings.TrimSpace(s)) }),
	).Title("Shortcuts")).Run()
	if err != nil {
		return err
	}
	cfg.Shortcuts.Domains = nil
	for _, s := range splitList(shortcutDomains) {
		if !strings.HasPrefix(s, ".") {
			s = "." + s
		}
		cfg.Shortcuts.Domains = append(cfg.Shortcuts.Domains, s)
	}
	cfg.Shortcuts.StripHyphens = stripHyphens
	cfg.Shortcuts.Rewrites, _ = parseRewrites(rewrites)
	cfg.Shell.Alias = strings.TrimSpace(commandAlias)
	return nil
}

func editCompletion(cfg *config.Config) error {
	maxResults := ""
	if cfg.Completion.MaxResults > 0 {
		maxResults = strconv.Itoa(cfg.Completion.MaxResults)
	}
	err := huh.NewForm(huh.NewGroup(
		huh.NewInput().
			Title("Maximum completion results").
			Description("Limits the hostnames offered per Tab completion. Leave empty for unlimited.").
			Value(&maxResults).
			Validate(func(s string) error {
				if strings.TrimSpace(s) == "" {
					return nil
				}
				if n, err := strconv.Atoi(strings.TrimSpace(s)); err != nil || n < 0 {
					return errors.New("must be a non-negative number")
				}
				return nil
			}),
	).Title("Completion")).Run()
	if err != nil {
		return err
	}
	cfg.Completion.MaxResults, _ = strconv.Atoi(strings.TrimSpace(maxResults))
	return nil
}

// parseRewrites reads "REGEX => REPLACEMENT" lines; the replacement may be empty.
func parseRewrites(s string) ([]config.ShortcutRewrite, error) {
	var out []config.ShortcutRewrite
	for _, line := range parseLines(s) {
		pattern, replace, ok := strings.Cut(line, "=>")
		if !ok {
			return nil, fmt.Errorf("%q: expected REGEX => REPLACEMENT", line)
		}
		out = append(out, config.ShortcutRewrite{Pattern: strings.TrimSpace(pattern), Replace: strings.TrimSpace(replace)})
	}
	if err := (config.ShortcutsConfig{Rewrites: out}).Validate(); err != nil {
		return nil, err
	}
	return out, nil
}

func formatRewrites(rewrites []config.ShortcutRewrite) string {
	lines := make([]string, 0, len(rewrites))
	for _, r := range rewrites {
		lines = append(lines, r.Pattern+" => "+r.Replace)
	}
	return strings.Join(lines, "\n")
}

// splitList returns the non-empty, trimmed entries of a comma-separated list.
func splitList(s string) []string {
	var out []string
	for _, part := range strings.Split(s, ",") {
		if part = strings.TrimSpace(part); part != "" {
			out = append(out, part)
		}
	}
	return out
}

var knownStrategies = map[string]bool{
	"primary_ip":        true,
	"management_subnet": true,
	"interface_name":    true,
}

func validateStrategies(s string) error {
	if strings.TrimSpace(s) == "" {
		return errors.New("at least one strategy is required")
	}
	for _, name := range parseStrategies(s) {
		if !knownStrategies[name] {
			return fmt.Errorf("unknown strategy %q — available: primary_ip, management_subnet, interface_name", name)
		}
	}
	return nil
}

func parseStrategies(s string) []string { return splitList(s) }

func save(cfg config.Config) error {
	if err := config.ValidateCommandAlias(cfg.Shell.Alias); err != nil {
		return err
	}
	if err := cfg.ValidateNetBoxes(); err != nil {
		return err
	}
	if err := cfg.ValidateSSH(); err != nil {
		return err
	}
	if err := cfg.Shortcuts.Validate(); err != nil {
		return err
	}
	path := config.Path()
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return fmt.Errorf("creating config dir: %w", err)
	}

	var b strings.Builder
	if len(cfg.NetBoxes) > 0 {
		b.WriteString("netboxes:\n")
		for _, source := range cfg.NetBoxes {
			fmt.Fprintf(&b, "  - name: %q\n    url: %q\n    token: %q\n    token_version: %d\n", source.Name, source.URL, source.Token, source.TokenVersion)
			if !sshOptionsEmpty(source.SSH) {
				b.WriteString("    ssh:\n")
				writeSSHOptions(&b, "      ", source.SSH)
			}
		}
	} else {
		b.WriteString("netbox:\n")
		fmt.Fprintf(&b, "  url: %q\n  token: %q\n  token_version: %d\n", cfg.NetBox.URL, cfg.NetBox.Token, cfg.NetBox.TokenVersion)
		if !sshOptionsEmpty(cfg.NetBox.SSH) {
			b.WriteString("  ssh:\n")
			writeSSHOptions(&b, "    ", cfg.NetBox.SSH)
		}
	}

	b.WriteString("\nresolver:\n")
	b.WriteString("  strategies:\n")
	for _, s := range cfg.Resolver.Strategies {
		fmt.Fprintf(&b, "    - %s\n", s)
	}
	if len(cfg.Resolver.ManagementSubnets) > 0 {
		b.WriteString("  management_subnets:\n")
		for _, s := range cfg.Resolver.ManagementSubnets {
			fmt.Fprintf(&b, "    - %s\n", s)
		}
	}
	if cfg.Resolver.InterfaceName != "" {
		fmt.Fprintf(&b, "  interface_name: %q\n", cfg.Resolver.InterfaceName)
	}

	b.WriteString("\ncache:\n")
	fmt.Fprintf(&b, "  ttl: %d\n", cfg.Cache.TTL)
	if cfg.Cache.Path != "" {
		fmt.Fprintf(&b, "  path: %q\n", cfg.Cache.Path)
	}

	b.WriteString("\nssh:\n")
	if cfg.SSH.DefaultUser != "" {
		fmt.Fprintf(&b, "  default_user: %q\n", cfg.SSH.DefaultUser)
	}
	if cfg.SSH.DefaultPort > 0 {
		fmt.Fprintf(&b, "  default_port: %d\n", cfg.SSH.DefaultPort)
	}
	writeSSHOptions(&b, "  ", cfg.SSH.SSHOptions)
	if len(cfg.SSH.Rules) > 0 {
		b.WriteString("  rules:\n")
		for _, rule := range cfg.SSH.Rules {
			b.WriteString("    - match:\n")
			for _, field := range [][2]string{{"tag", rule.Match.Tag}, {"instance", rule.Match.Instance}, {"kind", rule.Match.Kind}, {"name", rule.Match.Name}} {
				if field[1] != "" {
					fmt.Fprintf(&b, "        %s: %q\n", field[0], field[1])
				}
			}
			writeSSHOptions(&b, "      ", rule.SSHOptions)
		}
	}

	if cfg.Completion.MaxResults > 0 {
		fmt.Fprintf(&b, "\ncompletion:\n  max_results: %d\n", cfg.Completion.MaxResults)
	}

	if cfg.Shortcuts.Enabled() {
		b.WriteString("\nshortcuts:\n")
		if len(cfg.Shortcuts.Domains) > 0 {
			b.WriteString("  domains:\n")
			for _, d := range cfg.Shortcuts.Domains {
				fmt.Fprintf(&b, "    - %s\n", d)
			}
		}
		if cfg.Shortcuts.StripHyphens {
			b.WriteString("  strip_hyphens: true\n")
		}
		if len(cfg.Shortcuts.Rewrites) > 0 {
			b.WriteString("  rewrites:\n")
			for _, r := range cfg.Shortcuts.Rewrites {
				fmt.Fprintf(&b, "    - pattern: %q\n      replace: %q\n", r.Pattern, r.Replace)
			}
		}
	}

	if cfg.Shell.Alias != "" {
		fmt.Fprintf(&b, "\nshell:\n  alias: %q\n", cfg.Shell.Alias)
	}

	if err := storage.WriteFile(path, []byte(b.String()), 0o600); err != nil {
		return fmt.Errorf("writing config: %w", err)
	}

	fmt.Printf("\nConfig saved → %s\n", path)
	return nil
}

func sshOptionsEmpty(o config.SSHOptions) bool {
	return o.ProxyJump == "" && o.HostKeyChecking == "" && o.Target == "" && o.AddressFamily == "" && len(o.ExtraArgs) == 0 && len(o.Options) == 0
}

// writeSSHOptions keeps settings the wizard does not ask for when it rewrites the config.
func writeSSHOptions(b *strings.Builder, indent string, o config.SSHOptions) {
	for _, field := range [][2]string{{"proxy_jump", o.ProxyJump}, {"host_key_checking", o.HostKeyChecking}, {"target", o.Target}, {"address_family", o.AddressFamily}} {
		if field[1] != "" {
			fmt.Fprintf(b, "%s%s: %q\n", indent, field[0], field[1])
		}
	}
	if len(o.ExtraArgs) > 0 {
		fmt.Fprintf(b, "%sextra_args:\n", indent)
		for _, arg := range o.ExtraArgs {
			fmt.Fprintf(b, "%s  - %q\n", indent, arg)
		}
	}
	if len(o.Options) > 0 {
		keys := make([]string, 0, len(o.Options))
		for key := range o.Options {
			keys = append(keys, key)
		}
		sort.Strings(keys)
		fmt.Fprintf(b, "%soptions:\n", indent)
		for _, key := range keys {
			fmt.Fprintf(b, "%s  %s: %q\n", indent, key, o.Options[key])
		}
	}
}