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

package main

import (
	"sort"
	"strings"

	"github.com/spf13/cobra"

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

func sftpCmd() *cobra.Command {
	return transferCmd(internalssh.SFTP, "sftp [--dry-run] [--jump HOST|--no-jump] [sftp options] HOST[:PATH]",
		"Open an SFTP session to a NetBox host",
		`Resolve HOST like netssh does for ssh (cache, shortcuts, instance/host,
SSH options and rules, jump host) and run sftp. sftp options are passed
through; note that sftp uses -P for the port.

  netssh sftp fsnbx1
  netssh sftp -P 2222 prod/web01:/var/log`)
}

func scpCmd() *cobra.Command {
	return transferCmd(internalssh.SCP, "scp [--dry-run] [--jump HOST|--no-jump] [scp options] SOURCE... TARGET",
		"Copy files to or from a NetBox host",
		`Resolve the host in HOST:PATH operands like netssh does for ssh and run scp.
All remote operands must name the same host.

  netssh scp ./app.tar fsnbx1:/tmp/
  netssh scp -r prod/web01:/etc/nginx ./nginx`)
}

func transferCmd(client internalssh.Client, use, short, long string) *cobra.Command {
	return &cobra.Command{
		Use:                use,
		Short:              short,
		Long:               long,
		DisableFlagParsing: true,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) == 0 || args[0] == "--help" || args[0] == "-h" {
				_ = cmd.Help()
				return
			}
			flags, args, err := client.ExtractWrapperFlags(args)
			if err != nil {
				fatalf("%v", err)
			}
			if len(flags.Forwards) > 0 {
				fatalf("--forward is only available for ssh and netssh tunnel")
			}
			runClient(client, flags, args)
		},
		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
			if client.Name == internalssh.SFTP.Name && client.Parse(args) != nil {
				return nil, cobra.ShellCompDirectiveNoFileComp
			}
			if strings.HasPrefix(toComplete, "-") {
				return nil, cobra.ShellCompDirectiveNoFileComp
			}
			if client.Name == internalssh.SCP.Name {
				if strings.Contains(toComplete, ":") {
					return nil, cobra.ShellCompDirectiveNoFileComp // remote paths are not completed
				}
				// HOST: for remote operands; without a matching host, local files.
				hosts := completeHosts(toComplete)
				if len(hosts) == 0 {
					return nil, cobra.ShellCompDirectiveDefault
				}
				for i := range hosts {
					hosts[i] += ":"
				}
				return hosts, cobra.ShellCompDirectiveNoSpace | cobra.ShellCompDirectiveNoFileComp
			}
			return completeHosts(toComplete), cobra.ShellCompDirectiveNoFileComp
		},
	}
}

func tunnelCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "tunnel [--dry-run] [--jump HOST|--no-jump] [ssh options] HOST FORWARD...",
		Short: "Open port forwards to a NetBox host without a shell",
		Long: `Run ssh -N with the named forwards configured under ssh.forwards (globally,
per NetBox instance, or per rule). A FORWARD that is no configured name but a
port spec such as 8080:localhost:80 is used as a local forward (-L).

  netssh tunnel fsnbx1 web socks
  netssh tunnel prod/db01 5432:localhost:5432`,
		DisableFlagParsing: true,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) == 0 || args[0] == "--help" || args[0] == "-h" {
				_ = cmd.Help()
				return
			}
			flags, args, err := internalssh.SSH.ExtractWrapperFlags(args)
			if err != nil {
				fatalf("%v", err)
			}
			parsed := internalssh.SSH.Parse(args)
			if parsed == nil {
				fatalf("usage: netssh tunnel HOST FORWARD...")
			}
			// Everything after the host names forwards instead of a remote command.
			flags.Forwards = append(flags.Forwards, args[parsed.DestIdx+1:]...)
			flags.Tunnel = true
			runClient(internalssh.SSH, flags, args[:parsed.DestIdx+1])
		},
		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
			parsed := internalssh.SSH.Parse(args)
			if parsed == nil {
				return completeHosts(toComplete), cobra.ShellCompDirectiveNoFileComp
			}
			cfg, err := config.Load()
			if err != nil {
				return nil, cobra.ShellCompDirectiveError
			}
			var names []string
			if entry, ok := cachedHost(cfg, parsed.Host); ok {
				names = forwardNames(hostOptions(cfg, resolvedHost{Name: entry.Name, Kind: entry.Kind, Tags: entry.Tags}))
			} else {
				names = allForwardNames(cfg)
			}
			return filterPrefix(names, toComplete, args[parsed.DestIdx+1:]), cobra.ShellCompDirectiveNoFileComp
		},
	}
}

// completeHosts returns the cached hosts for a completion request.
func completeHosts(toComplete string) []string {
	cfg, err := config.Load()
	if err != nil {
		return nil
	}
	c := newCache(cfg)
	_ = c.Load()
	hosts := hostCompletions(c.All(), toComplete, cfg.Shortcuts)
	if limit := cfg.Completion.MaxResults; limit > 0 && len(hosts) > limit {
		hosts = hosts[:limit]
	}
	return hosts
}

// cachedHost finds a host for completion by name or shortcut, without NetBox.
func cachedHost(cfg *config.Config, host string) (cache.Entry, bool) {
	c := newCache(cfg)
	if err := c.Load(); err != nil {
		return cache.Entry{}, false
	}
	if entry, _ := c.Get(host); entry.Name != "" {
		return entry, true
	}
	if !cfg.Shortcuts.Enabled() {
		return cache.Entry{}, false
	}
	entry, found, _, err := c.GetByShortcut(host, normalizeTarget(cfg.Shortcuts, strings.Contains(host, "/")))
	return entry, found && err == nil
}

// allForwardNames lists every configured forward name, for when the host is unknown.
func allForwardNames(cfg *config.Config) []string {
	seen := map[string]bool{}
	add := func(o config.SSHOptions) {
		for name := range o.Forwards {
			seen[strings.ToLower(name)] = true
		}
	}
	add(cfg.SSH.SSHOptions)
	add(cfg.NetBox.SSH)
	for _, source := range cfg.NetBoxes {
		add(source.SSH)
	}
	for _, rule := range cfg.SSH.Rules {
		add(rule.SSHOptions)
	}
	names := make([]string, 0, len(seen))
	for name := range seen {
		names = append(names, name)
	}
	sort.Strings(names)
	return names
}

// filterPrefix keeps the names starting with prefix that were not given yet.
func filterPrefix(names []string, prefix string, given []string) []string {
	used := map[string]bool{}
	for _, name := range given {
		used[strings.ToLower(name)] = true
	}
	var out []string
	for _, name := range names {
		if strings.HasPrefix(name, strings.ToLower(prefix)) && !used[name] {
			out = append(out, name)
		}
	}
	return out
}