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

package main

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"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/netbox"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/resolver"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/setup"
	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,
	"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
	}

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

	// 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
	}

	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, parsed.Host, 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, parsed.Host) {
			target = &entries[i]
			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()

	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)
	}

	if host.Port != "" {
		sshArgs = append(sshArgs, "-p", host.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())

	// 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
			}
			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
}

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