Update unwrap command

command/unwrap.go

@@ -1,104 +1,107 @@
 package command
 
 import (
-	"flag"
 	"fmt"
 	"strings"
 
-	"github.com/hashicorp/vault/api"
+	"github.com/mitchellh/cli"
-	"github.com/hashicorp/vault/meta"
+	"github.com/posener/complete"
 )
 
-// UnwrapCommand is a Command that behaves like ReadCommand but specifically
+// Ensure we are implementing the right interfaces.
-// for unwrapping cubbyhole-wrapped secrets
+var _ cli.Command = (*UnwrapCommand)(nil)
+var _ cli.CommandAutocomplete = (*UnwrapCommand)(nil)
+
+// UnwrapCommand is a Command that behaves like ReadCommand but specifically for
+// unwrapping cubbyhole-wrapped secrets
 type UnwrapCommand struct {
-	meta.Meta
+	*BaseCommand
+}
+
+func (c *UnwrapCommand) Synopsis() string {
+	return "Unwraps a wrapped secret"
+}
+
+func (c *UnwrapCommand) Help() string {
+	helpText := `
+Usage: vault unwrap [options] [TOKEN]
+
+  Unwraps a wrapped secret from Vault by the given token. The result is the
+  same as the "vault read" operation on the non-wrapped secret. If no token
+  is given, the data in the currently authenticated token is unwrapped.
+
+  Unwrap the data in the cubbyhole backend for a token:
+
+      $ vault unwrap 3de9ece1-b347-e143-29b0-dc2dc31caafd
+
+  Unwrap the data in the active token:
+
+      $ vault auth 848f9ccf-7176-098c-5e2b-75a0689d41cd
+      $ vault unwrap # unwraps 848f9ccf...
+
+  For a full list of examples and paths, please see the online documentation.
+
+` + c.Flags().Help()
+
+	return strings.TrimSpace(helpText)
+}
+
+func (c *UnwrapCommand) Flags() *FlagSets {
+	return c.flagSet(FlagSetHTTP | FlagSetOutputField | FlagSetOutputFormat)
+}
+
+func (c *UnwrapCommand) AutocompleteArgs() complete.Predictor {
+	return c.PredictVaultFiles()
+}
+
+func (c *UnwrapCommand) AutocompleteFlags() complete.Flags {
+	return c.Flags().Completions()
 }
 
 func (c *UnwrapCommand) Run(args []string) int {
-	var format string
+	f := c.Flags()
-	var field string
+
-	var err error
+	if err := f.Parse(args); err != nil {
-	var secret *api.Secret
+		c.UI.Error(err.Error())
-	var flags *flag.FlagSet
-	flags = c.Meta.FlagSet("unwrap", meta.FlagSetDefault)
-	flags.StringVar(&format, "format", "table", "")
-	flags.StringVar(&field, "field", "", "")
-	flags.Usage = func() { c.Ui.Error(c.Help()) }
-	if err := flags.Parse(args); err != nil {
 		return 1
 	}
 
-	var tokenID string
+	args = f.Args()
-
+	token := ""
-	args = flags.Args()
 	switch len(args) {
 	case 0:
+		// Leave token as "", that will use the local token
 	case 1:
-		tokenID = args[0]
+		token = strings.TrimSpace(args[0])
 	default:
-		c.Ui.Error("unwrap expects zero or one argument (the ID of the wrapping token)")
+		c.UI.Error(fmt.Sprintf("Too many arguments (expected 0-1, got %d)", len(args)))
-		flags.Usage()
 		return 1
 	}
 
 	client, err := c.Client()
 	if err != nil {
-		c.Ui.Error(fmt.Sprintf(
+		c.UI.Error(err.Error())
-			"Error initializing client: %s", err))
 		return 2
 	}
 
-	secret, err = client.Logical().Unwrap(tokenID)
+	secret, err := client.Logical().Unwrap(token)
 	if err != nil {
-		c.Ui.Error(err.Error())
+		c.UI.Error(fmt.Sprintf("Error unwrapping: %s", err))
-		return 1
+		return 2
 	}
 	if secret == nil {
-		c.Ui.Error("Server gave empty response or secret returned was empty")
+		c.UI.Error("Could not find wrapped response")
-		return 1
+		return 2
 	}
 
 	// Handle single field output
-	if field != "" {
+	if c.flagField != "" {
-		return PrintRawField(c.Ui, secret, field)
+		return PrintRawField(c.UI, secret, c.flagField)
 	}
 
-	// Check if the original was a list response and format as a list if so
+	// Check if the original was a list response and format as a list
-	if secret.Data != nil &&
+	if _, ok := extractListData(secret); ok {
-		len(secret.Data) == 1 &&
+		return OutputList(c.UI, c.flagFormat, secret)
-		secret.Data["keys"] != nil {
-		_, ok := secret.Data["keys"].([]interface{})
-		if ok {
-			return OutputList(c.Ui, format, secret)
-		}
 	}
-	return OutputSecret(c.Ui, format, secret)
+	return OutputSecret(c.UI, c.flagFormat, secret)
-}
-
-func (c *UnwrapCommand) Synopsis() string {
-	return "Unwrap a wrapped secret"
-}
-
-func (c *UnwrapCommand) Help() string {
-	helpText := `
-Usage: vault unwrap [options] <wrapping token ID>
-
-  Unwrap a wrapped secret.
-
-  Unwraps the data wrapped by the given token ID. The returned result is the
-  same as a 'read' operation on a non-wrapped secret.
-
-General Options:
-` + meta.GeneralOptionsUsage() + `
-Read Options:
-
-  -format=table           The format for output. By default it is a whitespace-
-                          delimited table. This can also be json or yaml.
-
-  -field=field            If included, the raw value of the specified field
-                          will be output raw to stdout.
-
-`
-	return strings.TrimSpace(helpText)
 }

command/unwrap_test.go

@@ -4,104 +4,169 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/hashicorp/vault/http"
+	"github.com/hashicorp/vault/api"
-	"github.com/hashicorp/vault/meta"
-	"github.com/hashicorp/vault/vault"
 	"github.com/mitchellh/cli"
 )
 
-func TestUnwrap(t *testing.T) {
+func testUnwrapCommand(tb testing.TB) (*cli.MockUi, *UnwrapCommand) {
-	core, _, token := vault.TestCoreUnsealed(t)
+	tb.Helper()
-	ln, addr := http.TestServer(t, core)
-	defer ln.Close()
 
-	ui := new(cli.MockUi)
+	ui := cli.NewMockUi()
-	c := &UnwrapCommand{
+	return ui, &UnwrapCommand{
-		Meta: meta.Meta{
+		BaseCommand: &BaseCommand{
-			ClientToken: token,
+			UI: ui,
-			Ui:          ui,
+		},
+	}
+}
+
+func testUnwrapWrappedToken(tb testing.TB, client *api.Client, data map[string]interface{}) string {
+	tb.Helper()
+
+	wrapped, err := client.Logical().Write("sys/wrapping/wrap", data)
+	if err != nil {
+		tb.Fatal(err)
+	}
+	if wrapped == nil || wrapped.WrapInfo == nil || wrapped.WrapInfo.Token == "" {
+		tb.Fatalf("missing wrap info: %v", wrapped)
+	}
+	return wrapped.WrapInfo.Token
+}
+
+func TestUnwrapCommand_Run(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		name string
+		args []string
+		out  string
+		code int
+	}{
+		{
+			"too_many_args",
+			[]string{"foo", "bar"},
+			"Too many arguments",
+			1,
+		},
+		{
+			"default",
+			nil, // Token comes in the test func
+			"bar",
+			0,
+		},
+		{
+			"field",
+			[]string{"-field", "foo"},
+			"bar",
+			0,
+		},
+		{
+			"field_not_found",
+			[]string{"-field", "not-a-real-field"},
+			"not present in secret",
+			1,
+		},
+		{
+			"format",
+			[]string{"-format", "json"},
+			"{",
+			0,
+		},
+		{
+			"format_bad",
+			[]string{"-format", "nope-not-real"},
+			"Invalid output format",
+			1,
 		},
 	}
 
-	args := []string{
+	t.Run("validations", func(t *testing.T) {
-		"-address", addr,
+		t.Parallel()
-		"-field", "zip",
-	}
 
-	// Run once so the client is setup, ignore errors
+		for _, tc := range cases {
-	c.Run(args)
+			tc := tc
 
-	// Get the client so we can write data
+			t.Run(tc.name, func(t *testing.T) {
-	client, err := c.Client()
+				t.Parallel()
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
 
-	wrapLookupFunc := func(method, path string) string {
+				client, closer := testVaultServer(t)
-		if method == "GET" && path == "secret/foo" {
+				defer closer()
-			return "60s"
+
+				wrappedToken := testUnwrapWrappedToken(t, client, map[string]interface{}{
+					"foo": "bar",
+				})
+
+				ui, cmd := testUnwrapCommand(t)
+				cmd.client = client
+
+				tc.args = append(tc.args, wrappedToken)
+				code := cmd.Run(tc.args)
+				if code != tc.code {
+					t.Errorf("expected %d to be %d", code, tc.code)
+				}
+
+				combined := ui.OutputWriter.String() + ui.ErrorWriter.String()
+				if !strings.Contains(combined, tc.out) {
+					t.Errorf("expected %q to contain %q", combined, tc.out)
+				}
+			})
 		}
-		if method == "LIST" && path == "secret" {
+	})
-			return "60s"
+
+	t.Run("communication_failure", func(t *testing.T) {
+		t.Parallel()
+
+		client, closer := testVaultServerBad(t)
+		defer closer()
+
+		ui, cmd := testUnwrapCommand(t)
+		cmd.client = client
+
+		code := cmd.Run([]string{
+			"foo",
+		})
+		if exp := 2; code != exp {
+			t.Errorf("expected %d to be %d", code, exp)
 		}
-		return ""
-	}
-	client.SetWrappingLookupFunc(wrapLookupFunc)
 
-	data := map[string]interface{}{"zip": "zap"}
+		expected := "Error unwrapping: "
-	if _, err := client.Logical().Write("secret/foo", data); err != nil {
+		combined := ui.OutputWriter.String() + ui.ErrorWriter.String()
-		t.Fatalf("err: %s", err)
+		if !strings.Contains(combined, expected) {
-	}
+			t.Errorf("expected %q to contain %q", combined, expected)
+		}
+	})
 
-	outer, err := client.Logical().Read("secret/foo")
+	// This test needs its own client and server because it modifies the client
-	if err != nil {
+	// to the wrapping token
-		t.Fatalf("err: %s", err)
+	t.Run("local_token", func(t *testing.T) {
-	}
+		t.Parallel()
-	if outer == nil {
-		t.Fatal("outer response was nil")
-	}
-	if outer.WrapInfo == nil {
-		t.Fatalf("outer wrapinfo was nil, response was %#v", *outer)
-	}
 
-	args = append(args, outer.WrapInfo.Token)
+		client, closer := testVaultServer(t)
+		defer closer()
 
-	// Run the read
+		wrappedToken := testUnwrapWrappedToken(t, client, map[string]interface{}{
-	if code := c.Run(args); code != 0 {
+			"foo": "bar",
-		t.Fatalf("bad: %d\n\n%s", code, ui.ErrorWriter.String())
+		})
-	}
 
-	output := ui.OutputWriter.String()
+		ui, cmd := testUnwrapCommand(t)
-	if output != "zap\n" {
+		cmd.client = client
-		t.Fatalf("unexpectd output:\n%s", output)
+		cmd.client.SetToken(wrappedToken)
-	}
 
-	// Now test with list handling, specifically that it will be called with
+		// Intentionally don't pass the token here - it shoudl use the local token
-	// the list output formatter
+		code := cmd.Run([]string{})
-	ui.OutputWriter.Reset()
+		if exp := 0; code != exp {
+			t.Errorf("expected %d to be %d", code, exp)
+		}
 
-	outer, err = client.Logical().List("secret")
+		combined := ui.OutputWriter.String() + ui.ErrorWriter.String()
-	if err != nil {
+		if !strings.Contains(combined, "bar") {
-		t.Fatalf("err: %s", err)
+			t.Errorf("expected %q to contain %q", combined, "bar")
-	}
+		}
-	if outer == nil {
+	})
-		t.Fatal("outer response was nil")
-	}
-	if outer.WrapInfo == nil {
-		t.Fatalf("outer wrapinfo was nil, response was %#v", *outer)
-	}
 
-	args = []string{
+	t.Run("no_tabs", func(t *testing.T) {
-		"-address", addr,
+		t.Parallel()
-		outer.WrapInfo.Token,
-	}
-	// Run the read
-	if code := c.Run(args); code != 0 {
-		t.Fatalf("bad: %d\n\n%s", code, ui.ErrorWriter.String())
-	}
 
-	output = ui.OutputWriter.String()
+		_, cmd := testUnwrapCommand(t)
-	if strings.TrimSpace(output) != "Keys\n----\nfoo" {
+		assertNoTabs(t, cmd)
-		t.Fatalf("unexpected output:\n%s", output)
+	})
-	}
 }