vault-tools/vctx/main.go

// Command vctx is a thin, context-aware wrapper around the real `vault` CLI.
//
// It resolves a context (via the same config + token cache as vcm), sets
// VAULT_ADDR / VAULT_TOKEN / VAULT_NAMESPACE for that single invocation, and
// execs `vault` with the remaining arguments:
//
//	vctx --context sydney kv put kv/foo/bar secret=baz
//	vctx --context staging/sydney token lookup
package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"

	"git.unkin.net/unkin/vault-tools/shared"

	"github.com/spf13/cobra"
)

var version = "dev"

// execVault replaces the current process with the vault binary. It is a package
// variable so tests can substitute a fake in place of syscall.Exec.
var execVault = syscall.Exec

func main() {
	if err := newRootCmd().Execute(); err != nil {
		os.Exit(1)
	}
}

func newRootCmd() *cobra.Command {
	var context, format string

	root := &cobra.Command{
		Use:   "vctx --context <context> <vault args...>",
		Short: "Run the vault CLI against a named context.",
		Long: "vctx resolves a context from ~/.config/vcm (shared with vcm), loads its\n" +
			"cached token from ~/.cache/vcm/<context>, sets VAULT_ADDR, VAULT_TOKEN and\n" +
			"VAULT_NAMESPACE (and VAULT_FORMAT: yaml for kv get and */creds/*, table\n" +
			"for kv list, unless already chosen) for this invocation only, and execs\n" +
			"`vault` with the remaining arguments.\n\n" +
			"Example: vctx --context sydney kv put kv/foo/bar secret=baz",
		// Everything from the first non-flag argument on is the vault command
		// line; SetInterspersed(false) below stops flag parsing there so flags
		// meant for vault (e.g. `kv get -field=foo`) are passed through untouched.
		Args:         cobra.ArbitraryArgs,
		SilenceUsage: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runVctx(context, format, args)
		},
	}

	// Do not treat vault flags interspersed with args as vctx flags: only the
	// leading --context is ours; everything from the first positional on is the
	// vault command line.
	root.Flags().SetInterspersed(false)
	root.Flags().StringVar(&context, "context", "", "Vault context to target (required)")
	_ = root.RegisterFlagCompletionFunc("context", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		cfg, err := shared.Load()
		if err != nil {
			return nil, cobra.ShellCompDirectiveNoFileComp
		}
		return cfg.ContextNames(), cobra.ShellCompDirectiveNoFileComp
	})
	root.Flags().StringVar(&format, "format", "", "VAULT_FORMAT for this call (default: picked from the vault command)")
	_ = root.RegisterFlagCompletionFunc("format", cobra.FixedCompletions(
		[]string{"table", "json", "yaml", "pretty"}, cobra.ShellCompDirectiveNoFileComp))

	root.AddCommand(&cobra.Command{
		Use:          "version",
		Short:        "Print the version",
		Args:         cobra.NoArgs,
		Run:          func(cmd *cobra.Command, args []string) { fmt.Println(version) },
		SilenceUsage: true,
	})

	return root
}

// runVctx validates inputs, builds the vault invocation for the context, and
// execs it (replacing this process).
func runVctx(context, format string, args []string) error {
	if context == "" {
		return fmt.Errorf("--context is required")
	}
	if len(args) == 0 {
		return fmt.Errorf("no vault arguments given (e.g. vctx --context %s kv list kv/)", context)
	}

	bin, argv, env, err := buildInvocation(context, format, args, os.Environ())
	if err != nil {
		return err
	}
	// Replace this process with vault so its exit status, signals and TTY
	// behaviour pass straight through.
	if err := execVault(bin, argv, env); err != nil {
		return fmt.Errorf("exec vault: %w", err)
	}
	return nil
}

// buildInvocation resolves the context and its cached token, locates the vault
// binary, and returns the binary path, argv (vault + args) and the environment
// to exec with. baseEnv is the starting environment (normally os.Environ());
// ambient copies of the keys vctx sets are dropped first, since Go programs
// (vault included) honour the first occurrence of a duplicated variable.
// VAULT_NAMESPACE is always dropped so a context without one is not redirected.
// format, when non-empty, is set as VAULT_FORMAT; otherwise autoFormat picks one.
func buildInvocation(context, format string, args, baseEnv []string) (bin string, argv []string, env []string, err error) {
	cfg, err := shared.Load()
	if err != nil {
		return "", nil, nil, err
	}
	rc, err := cfg.Resolve(context)
	if err != nil {
		return "", nil, nil, err
	}
	tok, err := shared.LoadToken(context)
	if err != nil {
		return "", nil, nil, err
	}

	vaultBin, err := exec.LookPath("vault")
	if err != nil {
		return "", nil, nil, fmt.Errorf("vault CLI not found in PATH: %w", err)
	}

	set := map[string]string{
		"VAULT_ADDR":  rc.Address,
		"VAULT_TOKEN": tok.Token,
	}
	// Namespace comes from the context, falling back to what the token was
	// issued under; only set when non-empty.
	ns := rc.Namespace
	if ns == "" {
		ns = tok.Namespace
	}
	if ns != "" {
		set["VAULT_NAMESPACE"] = ns
	}
	if format == "" {
		format = autoFormat(args, baseEnv)
	}
	if format != "" {
		set["VAULT_FORMAT"] = format
	}

	for _, e := range baseEnv {
		k, _, _ := strings.Cut(e, "=")
		if _, ok := set[k]; ok || k == "VAULT_NAMESPACE" {
			continue
		}
		env = append(env, e)
	}
	for _, k := range []string{"VAULT_ADDR", "VAULT_TOKEN", "VAULT_NAMESPACE", "VAULT_FORMAT"} {
		if v, ok := set[k]; ok {
			env = append(env, k+"="+v)
		}
	}

	argv = append([]string{vaultBin}, args...)
	return vaultBin, argv, env, nil
}

// autoFormat picks VAULT_FORMAT from the vault command line: yaml for
// `kv get` and dynamic creds (`*/creds/*`), table for `kv list`, else "".
// It returns "" when the caller already chose (ambient VAULT_FORMAT, or a
// -format/-field flag, since vault formats -field values in non-table modes).
func autoFormat(args, baseEnv []string) string {
	for _, e := range baseEnv {
		if strings.HasPrefix(e, "VAULT_FORMAT=") && e != "VAULT_FORMAT=" {
			return ""
		}
	}
	var pos []string
	for i, a := range args {
		if a == "--" {
			pos = append(pos, args[i+1:]...)
			break
		}
		name, _, _ := strings.Cut(strings.TrimLeft(a, "-"), "=")
		switch {
		case !strings.HasPrefix(a, "-"):
			pos = append(pos, a)
		case name == "format" || name == "field":
			return ""
		}
	}
	isKV := func(sub string) bool { return len(pos) >= 2 && pos[0] == "kv" && pos[1] == sub }
	switch {
	case isKV("list"):
		return "table"
	case isKV("get"):
		return "yaml"
	}
	for _, p := range pos {
		if strings.Contains(p, "/creds/") {
			return "yaml"
		}
	}
	return ""
}