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 vctl), 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"
	"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 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/vault (shared with vctl), loads its\n" +
			"cached token from ~/.cache/vault/<context>, sets VAULT_ADDR, VAULT_TOKEN and\n" +
			"VAULT_NAMESPACE for this invocation only, and execs `vault` with the\n" +
			"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, 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.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 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, 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());
// the Vault settings are appended so they override any ambient values (exec
// semantics: the last assignment of a variable wins).
func buildInvocation(context 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)
	}

	env = make([]string, len(baseEnv), len(baseEnv)+3)
	copy(env, baseEnv)
	env = append(env,
		"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 != "" {
		env = append(env, "VAULT_NAMESPACE="+ns)
	}

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