command: add TF_CLI_ARGS to specify additional CLI args

main.go

@@ -7,17 +7,24 @@ import (
 	"log"
 	"os"
 	"runtime"
+	"strings"
 	"sync"
 
 	"github.com/hashicorp/go-plugin"
 	"github.com/hashicorp/terraform/helper/logging"
 	"github.com/hashicorp/terraform/terraform"
 	"github.com/mattn/go-colorable"
+	"github.com/mattn/go-shellwords"
 	"github.com/mitchellh/cli"
 	"github.com/mitchellh/panicwrap"
 	"github.com/mitchellh/prefixedio"
 )
 
+const (
+	// EnvCLI is the environment variable name to set additional CLI args.
+	EnvCLI = "TF_CLI_ARGS"
+)
+
 func main() {
 	// Override global prefix set by go-dynect during init()
 	log.SetPrefix("")
@@ -129,9 +136,35 @@ func wrappedMain() int {
 	// Make sure we clean up any managed plugins at the end of this
 	defer plugin.CleanupClients()
 
-	// Get the command line args. We shortcut "--version" and "-v" to
-	// just show the version.
+	// Get the command line args.
 	args := os.Args[1:]
+
+	// Build the CLI so far, we do this so we can query the subcommand.
+	cliRunner := &cli.CLI{
+		Args:       args,
+		Commands:   Commands,
+		HelpFunc:   helpFunc,
+		HelpWriter: os.Stdout,
+	}
+
+	// Prefix the args with any args from the EnvCLI
+	args, err = mergeEnvArgs(EnvCLI, cliRunner.Subcommand(), args)
+	if err != nil {
+		Ui.Error(err.Error())
+		return 1
+	}
+
+	// Prefix the args with any args from the EnvCLI targeting this command
+	suffix := strings.Replace(strings.Replace(
+		cliRunner.Subcommand(), "-", "_", -1), " ", "_", -1)
+	args, err = mergeEnvArgs(
+		fmt.Sprintf("%s_%s", EnvCLI, suffix), cliRunner.Subcommand(), args)
+	if err != nil {
+		Ui.Error(err.Error())
+		return 1
+	}
+
+	// We shortcut "--version" and "-v" to just show the version
 	for _, arg := range args {
 		if arg == "-v" || arg == "-version" || arg == "--version" {
 			newArgs := make([]string, len(args)+1)
@@ -142,7 +175,9 @@ func wrappedMain() int {
 		}
 	}
 
-	cli := &cli.CLI{
+	// Rebuild the CLI with any modified args.
+	log.Printf("[INFO] CLI command args: %#v", args)
+	cliRunner = &cli.CLI{
 		Args:       args,
 		Commands:   Commands,
 		HelpFunc:   helpFunc,
@@ -153,7 +188,7 @@ func wrappedMain() int {
 	ContextOpts.Providers = config.ProviderFactories()
 	ContextOpts.Provisioners = config.ProvisionerFactories()
 
-	exitCode, err := cli.Run()
+	exitCode, err := cliRunner.Run()
 	if err != nil {
 		Ui.Error(fmt.Sprintf("Error executing CLI: %s", err.Error()))
 		return 1
@@ -241,3 +276,47 @@ func copyOutput(r io.Reader, doneCh chan<- struct{}) {
 
 	wg.Wait()
 }
+
+func mergeEnvArgs(envName string, cmd string, args []string) ([]string, error) {
+	v := os.Getenv(envName)
+	if v == "" {
+		return args, nil
+	}
+
+	log.Printf("[INFO] %s value: %q", envName, v)
+	extra, err := shellwords.Parse(v)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"Error parsing extra CLI args from %s: %s",
+			envName, err)
+	}
+
+	// Find the command to look for in the args. If there is a space,
+	// we need to find the last part.
+	search := cmd
+	if idx := strings.LastIndex(search, " "); idx >= 0 {
+		search = cmd[idx+1:]
+	}
+
+	// Find the index to place the flags. We put them exactly
+	// after the first non-flag arg.
+	idx := -1
+	for i, v := range args {
+		if v == search {
+			idx = i
+			break
+		}
+	}
+
+	// idx points to the exact arg that isn't a flag. We increment
+	// by one so that all the copying below expects idx to be the
+	// insertion point.
+	idx++
+
+	// Copy the args
+	newArgs := make([]string, len(args)+len(extra))
+	copy(newArgs, args[:idx])
+	copy(newArgs[idx:], extra)
+	copy(newArgs[len(extra)+idx:], args[idx:])
+	return newArgs, nil
+}

main_test.go

@@ -0,0 +1,257 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"reflect"
+	"testing"
+
+	"github.com/mitchellh/cli"
+)
+
+func TestMain_cliArgsFromEnv(t *testing.T) {
+	// Setup the state. This test really messes with the environment and
+	// global state so we set things up to be restored.
+
+	// Restore original CLI args
+	oldArgs := os.Args
+	defer func() { os.Args = oldArgs }()
+
+	// Setup test command and restore that
+	testCommandName := "unit-test-cli-args"
+	testCommand := &testCommandCLI{}
+	defer func() { delete(Commands, testCommandName) }()
+	Commands[testCommandName] = func() (cli.Command, error) {
+		return testCommand, nil
+	}
+
+	cases := []struct {
+		Name     string
+		Args     []string
+		Value    string
+		Expected []string
+		Err      bool
+	}{
+		{
+			"no env",
+			[]string{testCommandName, "foo", "bar"},
+			"",
+			[]string{"foo", "bar"},
+			false,
+		},
+
+		{
+			"both env var and CLI",
+			[]string{testCommandName, "foo", "bar"},
+			"-foo bar",
+			[]string{"-foo", "bar", "foo", "bar"},
+			false,
+		},
+
+		{
+			"only env var",
+			[]string{testCommandName},
+			"-foo bar",
+			[]string{"-foo", "bar"},
+			false,
+		},
+
+		{
+			"cli string has blank values",
+			[]string{testCommandName, "bar", "", "baz"},
+			"-foo bar",
+			[]string{"-foo", "bar", "bar", "", "baz"},
+			false,
+		},
+
+		{
+			"cli string has blank values before the command",
+			[]string{"", testCommandName, "bar"},
+			"-foo bar",
+			[]string{"-foo", "bar", "bar"},
+			false,
+		},
+
+		{
+			// this should fail gracefully, this is just testing
+			// that we don't panic with our slice arithmetic
+			"no command",
+			[]string{},
+			"-foo bar",
+			nil,
+			true,
+		},
+
+		{
+			"single quoted strings",
+			[]string{testCommandName, "foo"},
+			"-foo 'bar baz'",
+			[]string{"-foo", "bar baz", "foo"},
+			false,
+		},
+
+		{
+			"double quoted strings",
+			[]string{testCommandName, "foo"},
+			`-foo "bar baz"`,
+			[]string{"-foo", "bar baz", "foo"},
+			false,
+		},
+
+		{
+			"double quoted single quoted strings",
+			[]string{testCommandName, "foo"},
+			`-foo "'bar baz'"`,
+			[]string{"-foo", "'bar baz'", "foo"},
+			false,
+		},
+	}
+
+	for i, tc := range cases {
+		t.Run(fmt.Sprintf("%d-%s", i, tc.Name), func(t *testing.T) {
+			os.Unsetenv(EnvCLI)
+			defer os.Unsetenv(EnvCLI)
+
+			// Set the env var value
+			if tc.Value != "" {
+				if err := os.Setenv(EnvCLI, tc.Value); err != nil {
+					t.Fatalf("err: %s", err)
+				}
+			}
+
+			// Setup the args
+			args := make([]string, len(tc.Args)+1)
+			args[0] = oldArgs[0] // process name
+			copy(args[1:], tc.Args)
+
+			// Run it!
+			os.Args = args
+			testCommand.Args = nil
+			exit := wrappedMain()
+			if (exit != 0) != tc.Err {
+				t.Fatalf("bad: %d", exit)
+			}
+			if tc.Err {
+				return
+			}
+
+			// Verify
+			if !reflect.DeepEqual(testCommand.Args, tc.Expected) {
+				t.Fatalf("bad: %#v", testCommand.Args)
+			}
+		})
+	}
+}
+
+// This test just has more options than the test above. Use this for
+// more control over behavior at the expense of more complex test structures.
+func TestMain_cliArgsFromEnvAdvanced(t *testing.T) {
+	// Restore original CLI args
+	oldArgs := os.Args
+	defer func() { os.Args = oldArgs }()
+
+	cases := []struct {
+		Name     string
+		Command  string
+		EnvVar   string
+		Args     []string
+		Value    string
+		Expected []string
+		Err      bool
+	}{
+		{
+			"targeted to another command",
+			"command",
+			EnvCLI + "_foo",
+			[]string{"command", "foo", "bar"},
+			"-flag",
+			[]string{"foo", "bar"},
+			false,
+		},
+
+		{
+			"targeted to this command",
+			"command",
+			EnvCLI + "_command",
+			[]string{"command", "foo", "bar"},
+			"-flag",
+			[]string{"-flag", "foo", "bar"},
+			false,
+		},
+
+		{
+			"targeted to a command with a hyphen",
+			"command-name",
+			EnvCLI + "_command_name",
+			[]string{"command-name", "foo", "bar"},
+			"-flag",
+			[]string{"-flag", "foo", "bar"},
+			false,
+		},
+
+		{
+			"targeted to a command with a space",
+			"command name",
+			EnvCLI + "_command_name",
+			[]string{"command", "name", "foo", "bar"},
+			"-flag",
+			[]string{"-flag", "foo", "bar"},
+			false,
+		},
+	}
+
+	for i, tc := range cases {
+		t.Run(fmt.Sprintf("%d-%s", i, tc.Name), func(t *testing.T) {
+			// Setup test command and restore that
+			testCommandName := tc.Command
+			testCommand := &testCommandCLI{}
+			defer func() { delete(Commands, testCommandName) }()
+			Commands[testCommandName] = func() (cli.Command, error) {
+				return testCommand, nil
+			}
+
+			os.Unsetenv(tc.EnvVar)
+			defer os.Unsetenv(tc.EnvVar)
+
+			// Set the env var value
+			if tc.Value != "" {
+				if err := os.Setenv(tc.EnvVar, tc.Value); err != nil {
+					t.Fatalf("err: %s", err)
+				}
+			}
+
+			// Setup the args
+			args := make([]string, len(tc.Args)+1)
+			args[0] = oldArgs[0] // process name
+			copy(args[1:], tc.Args)
+
+			// Run it!
+			os.Args = args
+			testCommand.Args = nil
+			exit := wrappedMain()
+			if (exit != 0) != tc.Err {
+				t.Fatalf("bad: %d", exit)
+			}
+			if tc.Err {
+				return
+			}
+
+			// Verify
+			if !reflect.DeepEqual(testCommand.Args, tc.Expected) {
+				t.Fatalf("bad: %#v", testCommand.Args)
+			}
+		})
+	}
+}
+
+type testCommandCLI struct {
+	Args []string
+}
+
+func (c *testCommandCLI) Run(args []string) int {
+	c.Args = args
+	return 0
+}
+
+func (c *testCommandCLI) Synopsis() string { return "" }
+func (c *testCommandCLI) Help() string     { return "" }