From e9d8e9f5374b694bec8e98111c9ba263d1bb5fc7 Mon Sep 17 00:00:00 2001
From: Jack Pearkes <jackpearkes@gmail.com>
Date: Thu, 27 Jun 2013 14:39:16 +0200
Subject: [PATCH] provisioner/shell: add support for environment variables to
 be injected

---
 provisioner/shell/provisioner.go      | 33 ++++++++++++++++++++++++---
 provisioner/shell/provisioner_test.go | 29 +++++++++++++++++++++++
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/provisioner/shell/provisioner.go b/provisioner/shell/provisioner.go
index c5b3ee7c5..c931b1111 100644
--- a/provisioner/shell/provisioner.go
+++ b/provisioner/shell/provisioner.go
@@ -31,12 +31,17 @@ type config struct {
 	// An array of multiple scripts to run.
 	Scripts []string
 
+	// An array of environment variables that will be injected before
+	// your command(s) are executed.
+	Vars []string `mapstructure:"environment_vars"`
+
 	// The remote path where the local shell script will be uploaded to.
 	// This should be set to a writable file that is in a pre-existing directory.
 	RemotePath string `mapstructure:"remote_path"`
 
 	// The command used to execute the script. The '{{ .Path }}' variable
-	// should be used to specify where the script goes.
+	// should be used to specify where the script goes, {{ .Vars }}
+	// can be used to inject the environment_vars into the environment.
 	ExecuteCommand string `mapstructure:"execute_command"`
 }
 
@@ -45,6 +50,7 @@ type Provisioner struct {
 }
 
 type ExecuteCommandTemplate struct {
+	Vars string
 	Path string
 }
 
@@ -56,7 +62,13 @@ func (p *Provisioner) Prepare(raws ...interface{}) error {
 	}
 
 	if p.config.ExecuteCommand == "" {
-		p.config.ExecuteCommand = "sh {{.Path}}"
+		// Don't leave an empty space if Vars isn't configured
+		if p.config.Vars == nil {
+			p.config.ExecuteCommand = "sh {{.Path}}"
+		} else {
+			p.config.ExecuteCommand = "{{.Vars}} sh {{.Path}}"
+		}
+
 	}
 
 	if p.config.Inline != nil && len(p.config.Inline) == 0 {
@@ -71,6 +83,10 @@ func (p *Provisioner) Prepare(raws ...interface{}) error {
 		p.config.Scripts = make([]string, 0)
 	}
 
+	if p.config.Vars == nil {
+		p.config.Vars = make([]string, 0)
+	}
+
 	errs := make([]error, 0)
 
 	if p.config.Path != "" && len(p.config.Scripts) > 0 {
@@ -93,6 +109,14 @@ func (p *Provisioner) Prepare(raws ...interface{}) error {
 		}
 	}
 
+	// Do a check for bad environment variables, such as '=foo', 'foobar'
+	for _, kv := range p.config.Vars {
+		vs := strings.Split(kv, "=")
+		if len(vs) != 2 || vs[0] == "" {
+			errs = append(errs, fmt.Errorf("Environment variable not in format 'key=value': %s", kv))
+		}
+	}
+
 	if len(errs) > 0 {
 		return &packer.MultiError{errs}
 	}
@@ -146,10 +170,13 @@ func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
 			return fmt.Errorf("Error uploading shell script: %s", err)
 		}
 
+		// Flatten the environment variables
+		flattendVars := strings.Join(p.config.Vars, " ")
+
 		// Compile the command
 		var command bytes.Buffer
 		t := template.Must(template.New("command").Parse(p.config.ExecuteCommand))
-		t.Execute(&command, &ExecuteCommandTemplate{p.config.RemotePath})
+		t.Execute(&command, &ExecuteCommandTemplate{flattendVars, p.config.RemotePath})
 
 		// Setup the remote command
 		stdout_r, stdout_w := io.Pipe()
diff --git a/provisioner/shell/provisioner_test.go b/provisioner/shell/provisioner_test.go
index bb4209051..1ac35253a 100644
--- a/provisioner/shell/provisioner_test.go
+++ b/provisioner/shell/provisioner_test.go
@@ -131,3 +131,32 @@ func TestProvisionerPrepare_Scripts(t *testing.T) {
 		t.Fatalf("should not have error: %s", err)
 	}
 }
+
+func TestProvisionerPrepare_EnvironmentVars(t *testing.T) {
+	config := testConfig()
+
+	// Test with a bad case
+	config["environment_vars"] = []string{"badvar", "good=var"}
+	p := new(Provisioner)
+	err := p.Prepare(config)
+	if err == nil {
+		t.Fatal("should have error")
+	}
+
+	// Test with a trickier case
+	config["environment_vars"] = []string{"=bad"}
+	p = new(Provisioner)
+	err = p.Prepare(config)
+	if err == nil {
+		t.Fatal("should have error")
+	}
+
+	// Test with a good case
+	// Note: baz= is a real env variable, just empty
+	config["environment_vars"] = []string{"FOO=bar", "baz="}
+	p = new(Provisioner)
+	err = p.Prepare(config)
+	if err != nil {
+		t.Fatalf("should not have error: %s", err)
+	}
+}