Update policy-write command

command/policy_write.go

@@ -7,84 +7,125 @@ import (
 	"os"
 	"strings"
 
-	"github.com/hashicorp/vault/meta"
+	"github.com/mitchellh/cli"
+	"github.com/posener/complete"
 )
 
-// PolicyWriteCommand is a Command that enables a new endpoint.
+// Ensure we are implementing the right interfaces.
+var _ cli.Command = (*PolicyWriteCommand)(nil)
+var _ cli.CommandAutocomplete = (*PolicyWriteCommand)(nil)
+
+// PolicyWriteCommand is a Command uploads a policy
 type PolicyWriteCommand struct {
-	meta.Meta
+	*BaseCommand
+
+	testStdin io.Reader // for tests
+}
+
+func (c *PolicyWriteCommand) Synopsis() string {
+	return "Uploads a policy file"
+}
+
+func (c *PolicyWriteCommand) Help() string {
+	helpText := `
+Usage: vault policy-write [options] NAME PATH
+
+  Uploads a policy with the given name from the contents of a local file or
+  stdin. If the path is "-", the policy is read from stdin. Otherwise, it is
+  loaded from the file at the given path.
+
+  Upload a policy named "my-policy" from /tmp/policy.hcl on the local disk:
+
+      $ vault policy-write my-policy /tmp/policy.hcl
+
+  Upload a policy from stdin:
+
+      $ cat my-policy.hcl | vault policy-write my-policy -
+
+  For a full list of examples, please see the documentation.
+
+` + c.Flags().Help()
+
+	return strings.TrimSpace(helpText)
+}
+
+func (c *PolicyWriteCommand) Flags() *FlagSets {
+	return c.flagSet(FlagSetHTTP)
+}
+
+func (c *PolicyWriteCommand) AutocompleteArgs() complete.Predictor {
+	return complete.PredictFunc(func(args complete.Args) []string {
+		// Predict the LAST argument hcl files - we don't want to predict the
+		// name argument as a filepath.
+		if len(args.All) == 3 {
+			return complete.PredictFiles("*.hcl").Predict(args)
+		}
+		return nil
+	})
+}
+
+func (c *PolicyWriteCommand) AutocompleteFlags() complete.Flags {
+	return c.Flags().Completions()
 }
 
 func (c *PolicyWriteCommand) Run(args []string) int {
-	flags := c.Meta.FlagSet("policy-write", meta.FlagSetDefault)
-	flags.Usage = func() { c.Ui.Error(c.Help()) }
-	if err := flags.Parse(args); err != nil {
+	f := c.Flags()
+
+	if err := f.Parse(args); err != nil {
+		c.UI.Error(err.Error())
 		return 1
 	}
 
-	args = flags.Args()
-	if len(args) != 2 {
-		flags.Usage()
-		c.Ui.Error(fmt.Sprintf(
-			"\npolicy-write expects exactly two arguments"))
+	args = f.Args()
+	switch {
+	case len(args) < 2:
+		c.UI.Error(fmt.Sprintf("Not enough arguments (expected 2, got %d)", len(args)))
+		return 1
+	case len(args) > 2:
+		c.UI.Error(fmt.Sprintf("Too many arguments (expected 2, got %d)", len(args)))
 		return 1
 	}
 
 	client, err := c.Client()
 	if err != nil {
-		c.Ui.Error(fmt.Sprintf(
-			"Error initializing client: %s", err))
+		c.UI.Error(err.Error())
 		return 2
 	}
 
 	// Policies are normalized to lowercase
-	name := strings.ToLower(args[0])
-	path := args[1]
+	name := strings.TrimSpace(strings.ToLower(args[0]))
+	path := strings.TrimSpace(args[1])
 
-	// Read the policy
-	var f io.Reader = os.Stdin
-	if path != "-" {
+	// Get the policy contents, either from stdin of a file
+	var reader io.Reader
+	if path == "-" {
+		reader = os.Stdin
+		if c.testStdin != nil {
+			reader = c.testStdin
+		}
+	} else {
 		file, err := os.Open(path)
 		if err != nil {
-			c.Ui.Error(fmt.Sprintf(
-				"Error opening file: %s", err))
-			return 1
+			c.UI.Error(fmt.Sprintf("Error opening policy file: %s", err))
+			return 2
 		}
 		defer file.Close()
-		f = file
+		reader = file
 	}
+
+	// Read the policy
 	var buf bytes.Buffer
-	if _, err := io.Copy(&buf, f); err != nil {
-		c.Ui.Error(fmt.Sprintf(
-			"Error reading file: %s", err))
-		return 1
+	if _, err := io.Copy(&buf, reader); err != nil {
+		c.UI.Error(fmt.Sprintf("Error reading policy: %s", err))
+		return 2
 	}
 	rules := buf.String()
 
 	if err := client.Sys().PutPolicy(name, rules); err != nil {
-		c.Ui.Error(fmt.Sprintf(
-			"Error: %s", err))
-		return 1
+		c.UI.Error(fmt.Sprintf("Error uploading policy: %s", err))
+		return 2
 	}
 
-	c.Ui.Output(fmt.Sprintf("Policy '%s' written.", name))
+	c.UI.Output(fmt.Sprintf("Success! Uploaded policy: %s", name))
 	return 0
 }
-
-func (c *PolicyWriteCommand) Synopsis() string {
-	return "Write a policy to the server"
-}
-
-func (c *PolicyWriteCommand) Help() string {
-	helpText := `
-Usage: vault policy-write [options] name path
-
-  Write a policy with the given name from the contents of a file or stdin.
-
-  If the path is "-", the policy is read from stdin. Otherwise, it is
-  loaded from the file at the given path.
-
-General Options:
-` + meta.GeneralOptionsUsage()
-	return strings.TrimSpace(helpText)
-}