Merge pull request 'vctx: colourise json/yaml output' (#6) from benvin/vctx-colour into main

Reviewed-on: #6

docs/vctx.md

@@ -8,7 +8,7 @@ invocation only, and execs `vault` with the remaining arguments.
 ## Synopsis
 
 ```
-vctx --context <context> [--format <fmt>] <any vault args...>
+vctx --context <context> [--format <fmt>] [--color auto|always|never] <any vault args...>
 vctx version
 ```
 
@@ -30,7 +30,8 @@ vctx --context sydney kv put kv/foo/bar secret=baz
      [Output format](#output-format))
 4. `exec`s the `vault` binary (found on `PATH`) with everything after the
    context. Because it replaces the process, `vault`'s exit status, signals and
-   TTY behaviour pass straight through.
+   TTY behaviour pass straight through (json/yaml output on a terminal is
+   instead run as a child and colourised; see [Colour](#colour)).
 
 The ambient `VAULT_ADDR` / `VAULT_TOKEN` / `VAULT_NAMESPACE` in your shell are ignored for the call
 — `vctx` always targets the chosen context.
@@ -50,9 +51,34 @@ It is left alone when `VAULT_FORMAT` is already set in the environment or the
 vault args include `-format` or `-field`. `--format` (`table`, `json`, `yaml`,
 `pretty`; `raw` for `vault read` only) always sets it, overriding the environment.
 
+## Colour
+
+When the effective format is `json` or `yaml` (from the table above,
+`--format`, a `-format` vault arg, or `VAULT_FORMAT`), `vctx` colourises vault's
+output like `jq`/`yq`. To do so it runs `vault` as a child instead of exec'ing
+it: stdin and stderr pass straight through, stdout is buffered and re-emitted
+coloured, SIGTERM/SIGHUP are forwarded (SIGINT already reaches vault from the
+terminal), and vault's exit code is returned
+(`128+n` if killed by signal `n`).
+
+`--color` controls this:
+
+| Value    | Behaviour |
+|----------|-----------|
+| `auto`   | colour only when stdout is a terminal, `NO_COLOR` is unset/empty and `TERM` is not `dumb` (default) |
+| `always` | always colour |
+| `never`  | never colour |
+
+Without colour (piped output, `NO_COLOR`, `table`/`raw`/`pretty` formats)
+`vctx` execs `vault` exactly as before, so scripts see unchanged output.
+Streaming or interactive commands (`monitor`, `debug`, `events`, `ssh`,
+`agent`, `proxy`, `server`, `login`, `operator`, and `write`/`patch`, which
+may prompt for MFA) are always exec'd and never coloured. `JQ_COLORS`
+overrides the JSON palette as in `jq`.
+
 ## Flags
 
-`--context` and `--format` are the only flags `vctx` consumes; they must come
+`--context`, `--format` and `--color` are the only flags `vctx` consumes; they must come
 **before** the vault command. Everything from the first non-flag argument onward is handed to `vault`
 untouched, so vault's own flags work normally:
 

vctx/main.go

@@ -2,19 +2,25 @@
 //
 // 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:
+// 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"
@@ -26,6 +32,14 @@ var version = "dev"
 // 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)
@@ -33,7 +47,7 @@ func main() {
 }
 
 func newRootCmd() *cobra.Command {
-	var context, format string
+	var context, format, color string
 
 	root := &cobra.Command{
 		Use:   "vctx --context <context> <vault args...>",
@@ -50,7 +64,7 @@ func newRootCmd() *cobra.Command {
 		Args:         cobra.ArbitraryArgs,
 		SilenceUsage: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runVctx(context, format, args)
+			return runVctx(context, format, color, args)
 		},
 	}
 
@@ -69,6 +83,9 @@ func newRootCmd() *cobra.Command {
 	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",
@@ -82,11 +99,16 @@ func newRootCmd() *cobra.Command {
 }
 
 // runVctx validates inputs, builds the vault invocation for the context, and
-// execs it (replacing this process).
-func runVctx(context, format string, args []string) error {
+// 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)
 	}
@@ -95,6 +117,14 @@ func runVctx(context, format string, args []string) error {
 	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 {
@@ -204,3 +234,89 @@ func autoFormat(args, baseEnv []string) string {
 	}
 	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
+}

vctx/main_test.go

@@ -1,12 +1,15 @@
 package main
 
 import (
+	"bytes"
 	"errors"
 	"os"
 	"path/filepath"
 	"reflect"
 	"strings"
+	"syscall"
 	"testing"
+	"time"
 
 	"git.unkin.net/unkin/vault-tools/shared"
 )
@@ -152,10 +155,10 @@ func TestBuildInvocationUnknownContext(t *testing.T) {
 }
 
 func TestRunVctxValidation(t *testing.T) {
-	if err := runVctx("", "", []string{"status"}); err == nil || !strings.Contains(err.Error(), "--context is required") {
+	if err := runVctx("", "", "auto", []string{"status"}); err == nil || !strings.Contains(err.Error(), "--context is required") {
 		t.Errorf("empty context: got %v", err)
 	}
-	if err := runVctx("sydney", "", nil); err == nil || !strings.Contains(err.Error(), "no vault arguments") {
+	if err := runVctx("sydney", "", "auto", nil); err == nil || !strings.Contains(err.Error(), "no vault arguments") {
 		t.Errorf("no args: got %v", err)
 	}
 }
@@ -174,7 +177,7 @@ func TestRunVctxExecsWithBuiltInvocation(t *testing.T) {
 	}
 	t.Cleanup(func() { execVault = orig })
 
-	if err := runVctx("sydney", "", []string{"kv", "list", "kv/"}); err != nil {
+	if err := runVctx("sydney", "", "auto", []string{"kv", "list", "kv/"}); err != nil {
 		t.Fatalf("runVctx: %v", err)
 	}
 	if gotBin != vault {
@@ -196,7 +199,7 @@ func TestRunVctxSurfacesExecError(t *testing.T) {
 	execVault = func(bin string, argv, env []string) error { return errors.New("boom") }
 	t.Cleanup(func() { execVault = orig })
 
-	err := runVctx("sydney", "", []string{"status"})
+	err := runVctx("sydney", "", "auto", []string{"status"})
 	if err == nil || !strings.Contains(err.Error(), "exec vault") {
 		t.Errorf("expected wrapped exec error, got %v", err)
 	}
@@ -258,3 +261,164 @@ func TestBuildInvocationFormat(t *testing.T) {
 		}
 	}
 }
+
+func TestColourFormat(t *testing.T) {
+	cases := []struct {
+		args []string
+		env  []string
+		want string
+	}{
+		{[]string{"kv", "get", "kv/foo"}, []string{"VAULT_FORMAT=yaml"}, "yaml"},
+		{[]string{"read", "x"}, []string{"VAULT_FORMAT=JSON"}, "json"},
+		{[]string{"kv", "list", "kv/"}, []string{"VAULT_FORMAT=table"}, ""},
+		{[]string{"read", "x"}, nil, ""},
+		{[]string{"read", "x"}, []string{"VAULT_FORMAT=pretty"}, ""},
+		{[]string{"read", "-format=raw", "x"}, []string{"VAULT_FORMAT=json"}, ""},
+		{[]string{"read", "-format=json", "x"}, nil, "json"},
+		{[]string{"read", "--format", "yaml", "x"}, []string{"VAULT_FORMAT=table"}, "yaml"},
+		{[]string{"read", "--", "-format=json"}, nil, ""},
+		{[]string{"events", "subscribe", "kv*"}, []string{"VAULT_FORMAT=json"}, ""},
+		{[]string{"monitor"}, []string{"VAULT_FORMAT=json"}, ""},
+		{[]string{"ssh", "-format=json", "host"}, nil, ""},
+		{[]string{"login", "-method=userpass"}, []string{"VAULT_FORMAT=json"}, ""},
+		{[]string{"operator", "unseal"}, []string{"VAULT_FORMAT=json"}, ""},
+		{[]string{"-format=yaml", "operator", "generate-root"}, nil, ""},
+		{[]string{"operator", "rekey", "-format=json"}, nil, ""},
+		{[]string{"write", "auth/userpass/login/bob"}, []string{"VAULT_FORMAT=json"}, ""},
+		{[]string{"write", "-format=json", "ns1/auth/x/login/y"}, nil, ""},
+		{[]string{"write", "-format=json", "secret/foo", "a=b"}, nil, ""},
+		{[]string{"patch", "auth/x/config"}, []string{"VAULT_FORMAT=json"}, ""},
+		{[]string{"patch", "-format=yaml", "secret/foo", "a=b"}, nil, ""},
+	}
+	for _, c := range cases {
+		if got := colourFormat(c.args, c.env); got != c.want {
+			t.Errorf("colourFormat(%v, %v) = %q, want %q", c.args, c.env, got, c.want)
+		}
+	}
+}
+
+// fakeVault writes a shell script standing in for vault and returns its path.
+func fakeVault(t *testing.T, body string) string {
+	t.Helper()
+	p := filepath.Join(t.TempDir(), "vault")
+	if err := os.WriteFile(p, []byte("#!/bin/sh\n"+body+"\n"), 0o755); err != nil {
+		t.Fatal(err)
+	}
+	return p
+}
+
+func TestRunChild(t *testing.T) {
+	cases := []struct {
+		name, body, format string
+		code               int
+		want               string
+	}{
+		{"json exit 0", `echo '{"a":1}'`, "json", 0, "\x1b[1;34m\"a\""},
+		{"yaml exit 3", "echo 'a: b'; exit 3", "yaml", 3, "\x1b[36ma"},
+		{"invalid passes through", "echo 'not json {'; exit 2", "json", 2, "not json {\n"},
+		{"signal", "kill -TERM $$", "json", 128 + 15, ""},
+	}
+	for _, c := range cases {
+		bin := fakeVault(t, c.body)
+		var out bytes.Buffer
+		code, err := runChild(bin, []string{bin}, nil, c.format, &out)
+		if err != nil {
+			t.Fatalf("%s: %v", c.name, err)
+		}
+		if code != c.code {
+			t.Errorf("%s: code = %d, want %d", c.name, code, c.code)
+		}
+		if !strings.Contains(out.String(), c.want) {
+			t.Errorf("%s: output %q missing %q", c.name, out.String(), c.want)
+		}
+	}
+}
+
+func TestRunChildForwardsSignal(t *testing.T) {
+	dir := t.TempDir()
+	bin := fakeVault(t, `trap 'echo INT >> "$D/got"' INT
+trap 'echo TERM >> "$D/got"; trap - TERM; kill -TERM $$' TERM
+touch "$D/ready"
+i=0; while [ $i -lt 100 ]; do sleep 0.05; i=$((i+1)); done`)
+	type res struct {
+		code int
+		err  error
+	}
+	ch := make(chan res, 1)
+	go func() {
+		code, err := runChild(bin, []string{bin}, append(os.Environ(), "D="+dir), "json", &bytes.Buffer{})
+		ch <- res{code, err}
+	}()
+	for i := 0; i < 200; i++ {
+		if _, err := os.Stat(filepath.Join(dir, "ready")); err == nil {
+			break
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	if err := syscall.Kill(os.Getpid(), syscall.SIGINT); err != nil {
+		t.Fatal(err)
+	}
+	time.Sleep(200 * time.Millisecond)
+	if err := syscall.Kill(os.Getpid(), syscall.SIGTERM); err != nil {
+		t.Fatal(err)
+	}
+	r := <-ch
+	if r.err != nil {
+		t.Fatal(r.err)
+	}
+	if got, _ := os.ReadFile(filepath.Join(dir, "got")); string(got) != "TERM\n" {
+		t.Errorf("child signals = %q, want only SIGTERM", got)
+	}
+	if r.code != 128+15 {
+		t.Errorf("code = %d, want %d", r.code, 128+15)
+	}
+}
+
+func TestRunChildPassesArgsAndEnv(t *testing.T) {
+	bin := fakeVault(t, `echo "{\"args\":\"$*\",\"tok\":\"$VAULT_TOKEN\"}"`)
+	var out bytes.Buffer
+	if _, err := runChild(bin, []string{bin, "kv", "get", "x"}, []string{"VAULT_TOKEN=s.T"}, "json", &out); err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(out.String(), "kv get x") || !strings.Contains(out.String(), "s.T") {
+		t.Errorf("output %q", out.String())
+	}
+}
+
+func TestRunVctxColourDecision(t *testing.T) {
+	setupVctx(t, vctxConfig, map[string]shared.Token{
+		"sydney": {Address: "https://vault.syd1.au.unkin.net", Token: "s.X"},
+	})
+	t.Setenv("VAULT_FORMAT", "")
+	origExec, origExit := execVault, osExit
+	t.Cleanup(func() { execVault, osExit = origExec, origExit })
+
+	cases := []struct {
+		color string
+		args  []string
+		exec  bool
+	}{
+		{"never", []string{"kv", "get", "kv/foo"}, true},
+		{"auto", []string{"kv", "get", "kv/foo"}, true}, // stdout is not a TTY under go test
+		{"always", []string{"kv", "list", "kv/"}, true},
+		{"always", []string{"status"}, true},
+		{"always", []string{"kv", "get", "kv/foo"}, false},
+		{"always", []string{"read", "-format=json", "x"}, false},
+		{"always", []string{"operator", "unseal", "-format=json"}, true},
+		{"always", []string{"login", "-format=json"}, true},
+	}
+	for _, c := range cases {
+		execed, exited := false, -1
+		execVault = func(string, []string, []string) error { execed = true; return nil }
+		osExit = func(code int) { exited = code }
+		if err := runVctx("sydney", "", c.color, c.args); err != nil {
+			t.Fatalf("%s %v: %v", c.color, c.args, err)
+		}
+		if execed != c.exec || (exited == 0) == c.exec {
+			t.Errorf("%s %v: exec=%v exit=%d, want exec=%v", c.color, c.args, execed, exited, c.exec)
+		}
+	}
+	if err := runVctx("sydney", "", "rainbow", []string{"status"}); err == nil {
+		t.Error("expected error for invalid --color")
+	}
+}