From 7f67b321691a110678b00e632e2423998fe43db5 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 13 Feb 2017 14:05:37 -0800
Subject: [PATCH] main: add TF_CLI_ARGS to specify additional CLI args

---
 main.go      |  47 ++++++++++++++++-
 main_test.go | 139 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 184 insertions(+), 2 deletions(-)
 create mode 100644 main_test.go

diff --git a/main.go b/main.go
index bcf6a3f58..f0c736acc 100644
--- a/main.go
+++ b/main.go
@@ -13,11 +13,17 @@ import (
 	"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 +135,45 @@ 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:]
+
+	// Prefix the args with any args from the EnvCLI
+	if v := os.Getenv(EnvCLI); v != "" {
+		log.Printf("[INFO] %s value: %q", EnvCLI, v)
+		extra, err := shellwords.Parse(v)
+		if err != nil {
+			Ui.Error(fmt.Sprintf(
+				"Error parsing extra CLI args from %s: %s",
+				EnvCLI, err))
+			return 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[0] != '-' {
+				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:])
+		args = newArgs
+
+	}
+
+	// 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,6 +184,7 @@ func wrappedMain() int {
 		}
 	}
 
+	log.Printf("[INFO] CLI command args: %#v", args)
 	cli := &cli.CLI{
 		Args:       args,
 		Commands:   Commands,
diff --git a/main_test.go b/main_test.go
new file mode 100644
index 000000000..4c6cd5816
--- /dev/null
+++ b/main_test.go
@@ -0,0 +1,139 @@
+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,
+		},
+
+		{
+			// 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)
+
+			// 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)
+			}
+		})
+	}
+}
+
+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 "" }