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 setup form, 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 {
		if errors.Is(err, huh.ErrUserAborted) {
			fmt.Fprintln(os.Stderr, "Setup cancelled.")
			os.Exit(0)
		}
		return err
	}
	defaultUser := cfg.SSH.DefaultUser
	defaultPort := ""
	if cfg.SSH.DefaultPort > 0 {
		defaultPort = strconv.Itoa(cfg.SSH.DefaultPort)
	}
	strategiesRaw := strings.Join(cfg.Resolver.Strategies, ", ")
	subnets := strings.Join(cfg.Resolver.ManagementSubnets, ", ")
	interfaceName := cfg.Resolver.InterfaceName
	cacheTTL := strconv.Itoa(cfg.Cache.TTL)
	shortcutDomains := strings.Join(cfg.Shortcuts.Domains, ", ")
	stripHyphens := cfg.Shortcuts.StripHyphens
	commandAlias := cfg.Shell.Alias

	if strategiesRaw == "" {
		strategiesRaw = "primary_ip"
	}
	if cacheTTL == "0" {
		cacheTTL = "3600"
	}

	form := 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"),

		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
				}),
		).Title("Resolver & cache"),

		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.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"),
	)

	if err := form.Run(); err != nil {
		if errors.Is(err, huh.ErrUserAborted) {
			fmt.Fprintln(os.Stderr, "Setup cancelled.")
			os.Exit(0)
		}
		return err
	}

	strategies := parseStrategies(strategiesRaw)

	ttl, _ := strconv.Atoi(cacheTTL)

	port := 0
	if p, err := strconv.Atoi(strings.TrimSpace(defaultPort)); err == nil {
		port = p
	}

	var subnetList []string
	for _, s := range strings.Split(subnets, ",") {
		if s = strings.TrimSpace(s); s != "" {
			subnetList = append(subnetList, s)
		}
	}

	var domainList []string
	for _, s := range strings.Split(shortcutDomains, ",") {
		if s = strings.TrimSpace(s); s != "" {
			if !strings.HasPrefix(s, ".") {
				s = "." + s
			}
			domainList = append(domainList, s)
		}
	}

	out := config.Config{
		Shell:    config.ShellConfig{Alias: strings.TrimSpace(commandAlias)},
		NetBoxes: sources,
		SSH: config.SSHConfig{
			DefaultUser: strings.TrimSpace(defaultUser),
			DefaultPort: port,
		},
		Resolver: config.ResolverConfig{
			Strategies:        strategies,
			ManagementSubnets: subnetList,
			InterfaceName:     strings.TrimSpace(interfaceName),
		},
		Cache: config.CacheConfig{
			TTL:  ttl,
			Path: cfg.Cache.Path,
		},
		Shortcuts: config.ShortcutsConfig{
			Domains:      domainList,
			StripHyphens: stripHyphens,
		},
	}

	if err := save(out); 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
}

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 {
	var out []string
	for _, part := range strings.Split(s, ",") {
		if name := strings.TrimSpace(part); name != "" {
			out = append(out, name)
		}
	}
	return out
}

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
	}
	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 len(cfg.Shortcuts.Domains) > 0 || cfg.Shortcuts.StripHyphens {
		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 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])
		}
	}
}