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

package main

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"text/tabwriter"
	"time"

	tea "github.com/charmbracelet/bubbletea"
	"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/hook"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/netbox"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/resolver"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/setup"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/shortcuts"
	internalssh "git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/ssh"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/tui"
)

// managedSubcommands are dispatched to cobra; everything else is treated as SSH wrapper mode.
var managedSubcommands = map[string]bool{
	"configure":  true,
	"search":     true,
	"cache":      true,
	"alias":      true,
	"hook":       true,
	"shell-init": true,
	"completion": true,
	"__complete": true,
	"help":       true,
}

func main() {
	args := os.Args[1:]

	if len(args) == 0 {
		runTUI()
		return
	}

	if managedSubcommands[args[0]] {
		if err := rootCmd().Execute(); err != nil {
			os.Exit(1)
		}
		return
	}

	runSSHWrapper(args)
}

// loadConfigOrSetup loads the config and runs the setup wizard if NetBox is not configured.
func loadConfigOrSetup() *config.Config {
	cfg, err := config.Load()
	if err != nil {
		fatalf("config: %v", err)
	}
	if cfg.NetBox.URL == "" {
		fmt.Fprintln(os.Stderr, "No configuration found. Starting setup…")
		if err := setup.RunWizard(cfg); err != nil {
			fatalf("setup: %v", err)
		}
		cfg, err = config.Load()
		if err != nil {
			fatalf("config: %v", err)
		}
	}
	return cfg
}

// runSSHWrapper resolves the target hostname via NetBox and execs the native ssh binary.
func runSSHWrapper(args []string) {
	cfg := loadConfigOrSetup()

	parsed := internalssh.Parse(args)
	if parsed == nil {
		// No destination found — pass arguments straight through to ssh.
		if err := internalssh.Exec(args); err != nil {
			fatalf("%v", err)
		}
		return
	}

	// Inject the configured default user if none was given on the command line.
	if cfg.SSH.DefaultUser != "" && parsed.User == "" && !internalssh.HasUserFlag(args) {
		args = append([]string{"-l", cfg.SSH.DefaultUser}, args...)
		parsed.DestIdx += 2
	}
	// Inject the configured default port if none was given on the command line.
	if cfg.SSH.DefaultPort > 0 && !internalssh.HasPortFlag(args) {
		args = append([]string{"-p", strconv.Itoa(cfg.SSH.DefaultPort)}, args...)
		parsed.DestIdx += 2
	}

	c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
	_ = c.Load()
	maybeBackgroundRefresh(cfg, c)

	// Cache hit with a fresh TTL — connect directly without querying NetBox.
	if entry, fresh := c.Get(parsed.Host); fresh {
		c.MarkUsed(parsed.Host)
		_ = c.Save()
		connect(entry.IP, parsed, args)
		return
	}

	// Shortcut cache lookup: scan entries whose normalized name matches the input.
	// If a stale match is found, use the canonical name for the NetBox re-fetch below.
	lookupHost := parsed.Host
	shortcutsEnabled := len(cfg.Shortcuts.Domains) > 0 || cfg.Shortcuts.StripHyphens
	if shortcutsEnabled {
		normalize := shortcuts.MakeNormalizer(cfg.Shortcuts)
		if entry, found, fresh := c.GetByShortcut(lookupHost, normalize); found {
			if fresh {
				c.MarkUsed(entry.Name)
				_ = c.Save()
				connect(entry.IP, parsed, args)
				return
			}
			// Stale shortcut match — use canonical name for NetBox re-fetch.
			lookupHost = entry.Name
		}
	}

	if cfg.NetBox.URL == "" {
		fatalf("netbox.url is not configured (~/.config/netssh.yaml)")
	}
	nbClient := netbox.NewClient(cfg.NetBox.URL, cfg.NetBox.Token, cfg.NetBox.TokenVersion)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	entries, err := nbClient.Search(ctx, lookupHost, netbox.SearchOptions{})
	if err != nil {
		fatalf("NetBox search failed: %v", err)
	}

	var target *netbox.HostEntry
	for i, e := range entries {
		if strings.EqualFold(e.Name, lookupHost) {
			target = &entries[i]
			break
		}
	}

	// Domain expansion: if still no match and the input has no dots, try appending
	// each configured domain suffix and re-querying NetBox.
	if target == nil && !strings.Contains(parsed.Host, ".") && len(cfg.Shortcuts.Domains) > 0 {
		for _, domain := range cfg.Shortcuts.Domains {
			expanded := parsed.Host + domain
			expandedEntries, searchErr := nbClient.Search(ctx, expanded, netbox.SearchOptions{})
			if searchErr != nil {
				continue
			}
			for i, e := range expandedEntries {
				if strings.EqualFold(e.Name, expanded) {
					target = &expandedEntries[i]
					break
				}
			}
			if target != nil {
				break
			}
		}
	}

	if target == nil {
		fatalf("host %q not found in NetBox", parsed.Host)
	}

	chain, err := resolver.New(cfg.Resolver)
	if err != nil {
		fatalf("resolver: %v", err)
	}

	ip, err := chain.Resolve(ctx, target, nbClient)
	if err != nil {
		fatalf("IP resolution for %q: %v", parsed.Host, err)
	}

	c.Upsert(cache.Entry{Name: target.Name, IP: ip, Kind: target.Kind, Tags: target.Tags})
	c.MarkUsed(target.Name)
	_ = c.Save()

	connect(ip, parsed, args)
}

func connect(ip string, parsed *internalssh.ParsedArgs, originalArgs []string) {
	newArgs := internalssh.ReplaceHost(originalArgs, parsed.DestIdx, ip)
	if err := internalssh.Exec(newArgs); err != nil {
		fatalf("%v", err)
	}
}

// runTUI starts the interactive host search.
func runTUI() {
	cfg := loadConfigOrSetup()

	c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
	_ = c.Load()
	maybeBackgroundRefresh(cfg, c)

	var nbClient *netbox.Client
	if cfg.NetBox.URL != "" {
		nbClient = netbox.NewClient(cfg.NetBox.URL, cfg.NetBox.Token, cfg.NetBox.TokenVersion)
	}

	m := tui.New(nbClient, c, cfg.SSH.DefaultUser)
	p := tea.NewProgram(m, tea.WithAltScreen())
	final, err := p.Run()
	if err != nil {
		fatalf("TUI: %v", err)
	}

	tuiModel, ok := final.(*tui.Model)
	if !ok {
		return
	}
	host := tuiModel.Selected()
	if host == nil {
		return
	}

	fmt.Fprintf(os.Stderr, "Connecting to %s (%s)…\n", host.Name, host.IP)

	var sshArgs []string

	// User override from TUI edit mode takes priority, then config default.
	user := host.User
	if user == "" {
		user = cfg.SSH.DefaultUser
	}
	if user != "" {
		sshArgs = append(sshArgs, "-l", user)
	}

	port := host.Port
	if port == "" && cfg.SSH.DefaultPort > 0 {
		port = strconv.Itoa(cfg.SSH.DefaultPort)
	}
	if port != "" {
		sshArgs = append(sshArgs, "-p", port)
	}

	sshArgs = append(sshArgs, host.IP)

	c.MarkUsed(host.Name)
	_ = c.Save()

	if err := internalssh.Exec(sshArgs); err != nil {
		fatalf("%v", err)
	}
}

// --- cobra root + subcommands ---

func rootCmd() *cobra.Command {
	root := &cobra.Command{
		Use:                "netssh",
		Short:              "SSH wrapper with NetBox hostname resolution",
		DisableFlagParsing: true,
		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
			cfg, err := config.Load()
			if err != nil {
				return nil, cobra.ShellCompDirectiveError
			}
			c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
			_ = c.Load()

			var completions []cobra.Completion
			for _, e := range c.Search(toComplete) {
				completions = append(completions, cobra.Completion(e.Name))
			}
			// Subcommands at the end, after all hostnames.
			for _, sub := range cmd.Commands() {
				if sub.IsAvailableCommand() && strings.HasPrefix(sub.Name(), toComplete) {
					completions = append(completions, cobra.Completion(sub.Name()+"\t"+sub.Short))
				}
			}
			return completions, cobra.ShellCompDirectiveNoFileComp | cobra.ShellCompDirectiveKeepOrder
		},
	}

	root.AddCommand(configureCmd(), searchCmd(), cacheCmd(), aliasCmd(), hookCmd(), shellInitCmd())

	// cobra builds the "completion" command lazily; force init so we can extend it.
	root.InitDefaultCompletionCmd()
	for _, cmd := range root.Commands() {
		if cmd.Name() == "completion" {
			cmd.AddCommand(completionInstallCmd(root))
			break
		}
	}

	return root
}

func completionInstallCmd(root *cobra.Command) *cobra.Command {
	var shell string
	cmd := &cobra.Command{
		Use:   "install",
		Short: "Install shell completion for the current user (no sudo required)",
		RunE: func(cmd *cobra.Command, args []string) error {
			if shell == "" {
				shell = filepath.Base(os.Getenv("SHELL"))
			}

			var (
				dir  string
				file string
				gen  func() ([]byte, error)
				note string
			)

			switch shell {
			case "bash":
				dir = filepath.Join(os.Getenv("HOME"), ".local", "share", "bash-completion", "completions")
				file = filepath.Join(dir, "netssh")
				gen = func() ([]byte, error) {
					var buf strings.Builder
					err := root.GenBashCompletionV2(&buf, true)
					return []byte(buf.String()), err
				}
				note = "Reload your shell or run: source " + file
			case "zsh":
				dir = filepath.Join(os.Getenv("HOME"), ".zfunc")
				file = filepath.Join(dir, "_netssh")
				gen = func() ([]byte, error) {
					var buf strings.Builder
					err := root.GenZshCompletion(&buf)
					return []byte(buf.String()), err
				}
				note = "Make sure ~/.zfunc is in your fpath:\n  fpath=(~/.zfunc $fpath)\n  autoload -Uz compinit && compinit"
			case "fish":
				configDir, _ := os.UserConfigDir()
				dir = filepath.Join(configDir, "fish", "completions")
				file = filepath.Join(dir, "netssh.fish")
				gen = func() ([]byte, error) {
					var buf strings.Builder
					err := root.GenFishCompletion(&buf, true)
					return []byte(buf.String()), err
				}
				note = "Reload your shell or start a new fish session."
			default:
				return fmt.Errorf("unsupported shell %q — use --shell bash|zsh|fish", shell)
			}

			script, err := gen()
			if err != nil {
				return fmt.Errorf("generating completion: %w", err)
			}
			if err := os.MkdirAll(dir, 0o755); err != nil {
				return fmt.Errorf("creating %s: %w", dir, err)
			}
			if err := os.WriteFile(file, script, 0o644); err != nil {
				return fmt.Errorf("writing %s: %w", file, err)
			}
			fmt.Printf("Completion installed → %s\n%s\n", file, note)
			return nil
		},
	}
	cmd.Flags().StringVar(&shell, "shell", "", "Shell to install for (default: $SHELL). Supported: bash, zsh, fish")
	return cmd
}

func configureCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "configure",
		Short: "Interactively configure netssh",
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg, _ := config.Load()
			if cfg == nil {
				cfg = &config.Config{}
			}
			return setup.RunWizard(cfg)
		},
	}
}

func searchCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "search <query>",
		Short: "Print hostnames from the cache (used for shell completion)",
		Args:  cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg, err := config.Load()
			if err != nil {
				return err
			}
			c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
			if err := c.Load(); err != nil {
				return err
			}
			q := ""
			if len(args) > 0 {
				q = args[0]
			}
			for _, e := range c.Search(q) {
				fmt.Println(e.Name)
			}
			return nil
		},
	}
}

func cacheCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "cache",
		Short: "Manage the local host cache",
	}
	cmd.AddCommand(cacheListCmd(), cacheClearCmd(), cacheRefreshCmd())
	return cmd
}

func cacheListCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "Show all cached entries",
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg, err := config.Load()
			if err != nil {
				return err
			}
			c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
			if err := c.Load(); err != nil {
				return err
			}
			entries := c.All()
			sort.Slice(entries, func(i, j int) bool { return entries[i].Name < entries[j].Name })

			w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
			fmt.Fprintln(w, "NAME\tIP\tKIND\tCACHED AT")
			for _, e := range entries {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", e.Name, e.IP, e.Kind, e.CachedAt.Format(time.RFC3339))
			}
			return w.Flush()
		},
	}
}

func cacheClearCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "clear",
		Short: "Remove all cached entries",
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg, err := config.Load()
			if err != nil {
				return err
			}
			c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
			c.Clear()
			if err := c.Save(); err != nil {
				return err
			}
			fmt.Println("Cache cleared.")
			return nil
		},
	}
}

func cacheRefreshCmd() *cobra.Command {
	var filterTag, filterKind string
	cmd := &cobra.Command{
		Use:   "refresh",
		Short: "Re-fetch all known hosts from NetBox",
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg, err := config.Load()
			if err != nil {
				return err
			}
			if cfg.NetBox.URL == "" {
				return fmt.Errorf("netbox.url is not configured")
			}

			nbClient := netbox.NewClient(cfg.NetBox.URL, cfg.NetBox.Token, cfg.NetBox.TokenVersion)
			ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
			defer cancel()

			opts := netbox.SearchOptions{Tag: filterTag, Kind: filterKind}
			entries, err := nbClient.SearchAll(ctx, "", opts)
			if err != nil {
				return fmt.Errorf("NetBox: %w", err)
			}

			c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
			_ = c.Load()

			chain, _ := resolver.New(cfg.Resolver)
			for i := range entries {
				e := &entries[i]

				var ip string
				if chain != nil {
					if resolved, err := chain.Resolve(ctx, e, nbClient); err == nil {
						ip = resolved
					}
				}
				if ip == "" {
					ip = e.PrimaryIP4
				}
				if ip == "" {
					ip = e.PrimaryIP6
				}

				if ip != "" {
					c.Upsert(cache.Entry{Name: e.Name, IP: ip, Kind: e.Kind, Tags: e.Tags})
				}
			}

			if err := c.Save(); err != nil {
				return err
			}
			if err := c.SetRefreshed(); err != nil {
				fmt.Fprintf(os.Stderr, "warning: could not update refresh timestamp: %v\n", err)
			}
			if aliasEntries := buildAliasEntries(c, *cfg); len(aliasEntries) > 0 {
				if err := hook.WriteAliasFiles(aliasEntries); err != nil {
					fmt.Fprintf(os.Stderr, "warning: could not write alias files: %v\n", err)
				}
			}
			fmt.Printf("%d entries written to cache.\n", len(entries))
			return nil
		},
	}
	cmd.Flags().StringVar(&filterTag, "tag", "", "Filter by NetBox tag slug (e.g. prod)")
	cmd.Flags().StringVar(&filterKind, "kind", "", "Filter by kind: device or vm")
	return cmd
}

// buildAliasEntries computes the sorted, deduplicated alias list from the cache.
func buildAliasEntries(c *cache.Cache, cfg config.Config) []hook.AliasEntry {
	all := c.All()
	sort.Slice(all, func(i, j int) bool { return all[i].Name < all[j].Name })

	seen := make(map[string]string)
	var out []hook.AliasEntry
	for _, e := range all {
		name := shortcuts.AliasName(e.Name, cfg.Shortcuts)
		if name == "" {
			continue
		}
		if _, dup := seen[name]; dup {
			continue
		}
		seen[name] = e.Name
		out = append(out, hook.AliasEntry{Name: name, Host: e.Name})
	}
	return out
}

func aliasCmd() *cobra.Command {
	var shell string
	cmd := &cobra.Command{
		Use:   "alias",
		Short: "Print shell aliases for all cached hosts",
		Long: `Print shell alias definitions for all cached hosts.
The alias name is the shortened form derived from the configured shortcuts
(domain suffixes stripped, hyphens optionally stripped).

Source the output in your shell profile:
  bash/zsh:  eval "$(netssh alias)"
  fish:      netssh alias --shell fish | source

Or use in a script:
  source <(netssh alias)`,
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg, err := config.Load()
			if err != nil {
				return err
			}
			c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
			if err := c.Load(); err != nil {
				return err
			}
			if shell == "" {
				shell = filepath.Base(os.Getenv("SHELL"))
			}

			entries := buildAliasEntries(c, *cfg)
			if len(entries) == 0 {
				fmt.Fprintln(os.Stderr, "netssh: cache is empty — run 'netssh cache refresh' first")
				return nil
			}

			switch shell {
			case "fish":
				fmt.Printf("# netssh aliases (%d hosts) — source with: netssh alias --shell fish | source\n", len(entries))
				for _, e := range entries {
					fmt.Printf("alias %s 'netssh %s'\n", e.Name, e.Host)
				}
			default:
				fmt.Printf("# netssh aliases (%d hosts) — source with: eval \"$(netssh alias)\"\n", len(entries))
				for _, e := range entries {
					fmt.Printf("alias %s='netssh %s'\n", e.Name, e.Host)
				}
			}
			return nil
		},
	}
	cmd.Flags().StringVar(&shell, "shell", "", "Output format: bash, zsh, fish (default: $SHELL)")
	return cmd
}

// shellInitCmd is called at shell startup to update alias files and trigger a
// background cache refresh when stale. Intentionally silent — never disrupts the prompt.
func shellInitCmd() *cobra.Command {
	return &cobra.Command{
		Use:    "shell-init",
		Short:  "Regenerate alias files and trigger a background cache refresh if stale",
		Hidden: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg, err := config.Load()
			if err != nil || cfg.NetBox.URL == "" {
				return nil
			}
			c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
			_ = c.Load()

			// Regenerate alias files from the current cache every shell start
			// so aliases are always up-to-date even before a refresh completes.
			if entries := buildAliasEntries(c, *cfg); len(entries) > 0 {
				_ = hook.WriteAliasFiles(entries)
			}

			maybeBackgroundRefresh(cfg, c)
			return nil
		},
	}
}

func hookCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "hook",
		Short: "Manage shell hooks for automatic cache refresh at shell startup",
	}
	cmd.AddCommand(hookInstallCmd(), hookUninstallCmd())
	return cmd
}

func hookInstallCmd() *cobra.Command {
	var shell string
	cmd := &cobra.Command{
		Use:   "install",
		Short: "Add netssh shell-init to your shell profile",
		Long: `Appends a single line to your shell profile that runs 'netssh shell-init'
on every new shell session. shell-init checks whether the cache is older
than 24 hours and, if so, starts a background refresh — no delay to your prompt.

After installation, reload your profile or open a new shell.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if shell == "" {
				shell = filepath.Base(os.Getenv("SHELL"))
			}
			profile, err := hook.ProfilePath(shell)
			if err != nil {
				return err
			}

			installed, err := hook.Install(profile)
			if err != nil {
				return err
			}
			// For bash/zsh: also add the source line so aliases are loaded on startup.
			// Fish uses conf.d auto-sourcing — no explicit source line needed.
			if shell != "fish" {
				if _, err := hook.InstallAliasesSource(profile); err != nil {
					return err
				}
			}

			if !installed {
				fmt.Printf("Hook already installed in %s\n", profile)
			} else {
				fmt.Printf("Hook installed → %s\n%s\n", profile, hook.ReloadNote(profile))
			}

			// Generate the aliases file immediately so it is ready on the next shell start.
			cfg, cfgErr := config.Load()
			if cfgErr == nil && cfg.NetBox.URL != "" {
				c := cache.New(cfg.Cache.Path, cfg.Cache.TTL)
				if _ = c.Load(); true {
					if entries := buildAliasEntries(c, *cfg); len(entries) > 0 {
						if err := hook.WriteAliasFiles(entries); err != nil {
							fmt.Fprintf(os.Stderr, "warning: could not write alias files: %v\n", err)
						} else {
							fmt.Printf("%d aliases written to %s\n", len(entries), hook.AliasesPath())
						}
					} else {
						fmt.Println("Cache is empty — run 'netssh cache refresh' to populate aliases.")
					}
				}
			}
			return nil
		},
	}
	cmd.Flags().StringVar(&shell, "shell", "", "Shell to install for (default: $SHELL). Supported: bash, zsh, fish")
	return cmd
}

func hookUninstallCmd() *cobra.Command {
	var shell string
	cmd := &cobra.Command{
		Use:   "uninstall",
		Short: "Remove netssh shell-init from your shell profile",
		RunE: func(cmd *cobra.Command, args []string) error {
			if shell == "" {
				shell = filepath.Base(os.Getenv("SHELL"))
			}
			profile, err := hook.ProfilePath(shell)
			if err != nil {
				return err
			}
			removed, err := hook.Uninstall(profile)
			if err != nil {
				return err
			}
			if !removed {
				fmt.Printf("No hook found in %s\n", profile)
				return nil
			}
			fmt.Printf("Hook removed from %s\n", profile)
			return nil
		},
	}
	cmd.Flags().StringVar(&shell, "shell", "", "Shell to uninstall from (default: $SHELL). Supported: bash, zsh, fish")
	return cmd
}

// maybeBackgroundRefresh starts a background `netssh cache refresh` if the cache
// has not been fully refreshed within the last 24 hours.
func maybeBackgroundRefresh(cfg *config.Config, c *cache.Cache) {
	if cfg.NetBox.URL == "" {
		return
	}
	if !c.NeedsRefresh(24 * time.Hour) {
		return
	}
	self, err := os.Executable()
	if err != nil {
		return
	}
	cmd := exec.Command(self, "cache", "refresh")
	cmd.Stdout = nil
	cmd.Stderr = nil
	cmd.Stdin = nil
	_ = cmd.Start() // fire and forget — becomes orphan after parent execs ssh
}

func fatalf(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "netssh: "+format+"\n", args...)
	os.Exit(1)
}