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 (or runs it as a child to
// colourise json/yaml output on a terminal):
//
//	vctx --context sydney kv put kv/foo/bar secret=baz
//	vctx --context staging/sydney token lookup
package main

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"

	"git.unkin.net/unkin/golib/colour"
	"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

// osExit is swapped in tests to observe the child's exit code.
var osExit = os.Exit

// uncaptured commands are always exec'd: buffering their stdout would hide
// streamed output or prompts (unseal keys, MFA passcodes; write and patch can
// prompt for MFA) until exit, or break the TTY, for ssh.
var uncaptured = map[string]bool{"monitor": true, "debug": true, "events": true, "ssh": true, "agent": true, "proxy": true, "server": true, "login": true, "operator": true, "write": true, "patch": true}

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

func newRootCmd() *cobra.Command {
	var context, format, color 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, color, 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.Flags().StringVar(&color, "color", "auto", "Colourise json/yaml output: auto, always or never")
	_ = root.RegisterFlagCompletionFunc("color", cobra.FixedCompletions(
		[]string{"auto", "always", "never"}, 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), or runs it as a child and colourises its
// stdout when colour is enabled and the output format is json or yaml.
func runVctx(context, format, color string, args []string) error {
	if context == "" {
		return fmt.Errorf("--context is required")
	}
	mode, ok := map[string]colour.Mode{"auto": colour.Auto, "always": colour.Always, "never": colour.Never}[color]
	if !ok {
		return fmt.Errorf("--color must be auto, always or never, got %q", color)
	}
	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
	}
	if f := colourFormat(args, env); f != "" && colour.Enabled(mode, os.Stdout) {
		code, err := runChild(bin, argv, env, f, os.Stdout)
		if err != nil {
			return err
		}
		osExit(code)
		return nil
	}
	// 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 ""
}

// colourFormat returns the effective vault output format ("json" or "yaml")
// when it is one we colourise, else "". A -format arg beats VAULT_FORMAT, as in
// vault itself.
func colourFormat(args, env []string) string {
	f := ""
	var pos []string
	for _, e := range env {
		if v, ok := strings.CutPrefix(e, "VAULT_FORMAT="); ok {
			f = v
		}
	}
	for i := 0; i < len(args); i++ {
		a := args[i]
		if a == "--" {
			pos = append(pos, args[i+1:]...)
			break
		}
		if !strings.HasPrefix(a, "-") {
			pos = append(pos, a)
			continue
		}
		name, val, hasVal := strings.Cut(strings.TrimLeft(a, "-"), "=")
		if name != "format" {
			continue
		}
		if !hasVal && i+1 < len(args) {
			i++
			val = args[i]
		}
		f = val
	}
	f = strings.ToLower(f)
	if len(pos) == 0 || uncaptured[pos[0]] || (f != "json" && f != "yaml") {
		return ""
	}
	return f
}

// runChild runs vault with stdin/stderr passed through, forwards SIGTERM and
// SIGHUP, then writes its captured stdout colourised to w. It returns vault's
// exit code, or 128+n if it died from signal n.
func runChild(bin string, argv, env []string, format string, w io.Writer) (int, error) {
	var out bytes.Buffer
	cmd := exec.Command(bin, argv[1:]...)
	cmd.Env, cmd.Stdin, cmd.Stdout, cmd.Stderr = env, os.Stdin, &out, os.Stderr

	sigs := make(chan os.Signal, 1)
	// SIGINT is caught but not forwarded: vault is in our process group and already gets ^C.
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
	defer signal.Stop(sigs)
	if err := cmd.Start(); err != nil {
		return 0, fmt.Errorf("run vault: %w", err)
	}
	done := make(chan struct{})
	defer close(done)
	go func() {
		for {
			select {
			case s := <-sigs:
				if s != syscall.SIGINT {
					_ = cmd.Process.Signal(s)
				}
			case <-done:
				return
			}
		}
	}()

	err := cmd.Wait()
	var ee *exec.ExitError
	if err != nil && !errors.As(err, &ee) {
		return 0, fmt.Errorf("run vault: %w", err)
	}
	paint := colour.JSON
	if format == "yaml" {
		paint = colour.YAML
	}
	if err := paint(w, out.Bytes()); err != nil {
		return 0, err
	}
	if ws, ok := cmd.ProcessState.Sys().(syscall.WaitStatus); ok && ws.Signaled() {
		return 128 + int(ws.Signal()), nil
	}
	return cmd.ProcessState.ExitCode(), nil
}