Fix setting of AWS variables in env subcommand

miquella · miquella · commit f9cb9eaa2e29 · 2017-04-13T17:23:42.000-06:00
diff --git a/env.go b/env.go
@@ -20,22 +20,30 @@ func (e *Env) Run(steward Steward) error {
 	}
 
 	usageHint := ""
+	unsetVar := ""
 	setVar := ""
 	quoteReplacement := "\""
 	switch e.Shell {
 	case "fish":
 		usageHint = "# To load these variables into your shell, execute:\n#   eval (%s)\n"
+		unsetVar = "set -e %s;\n"
 		setVar = "set -x %s \"%s\";\n"
 		quoteReplacement = "\\\""
 	default:
 		usageHint = "# To load these variables into your shell, execute:\n#   eval $(%s)\n"
+		unsetVar = "unset %s\n"
 		setVar = "export %s=\"%s\"\n"
 		quoteReplacement = "\\\""
 	}
 
-	// sort the vars
+	vars := env.Variables()
+
+	// sort the vars to unset
+	sort.Strings(vars.Unset)
+
+	// sort the vars to set
 	var keys []string
-	for key := range env.Vars {
+	for key := range vars.Set {
 		keys = append(keys, key)
 	}
 	sort.Strings(keys)
@@ -45,8 +53,12 @@ func (e *Env) Run(steward Steward) error {
 		fmt.Printf(usageHint, strings.Join(os.Args, " "))
 	}
 
+	for _, key := range vars.Unset {
+		fmt.Printf(unsetVar, key)
+	}
+
 	for _, key := range keys {
-		fmt.Printf(setVar, key, strings.Replace(env.Vars[key], "\"", quoteReplacement, -1))
+		fmt.Printf(setVar, key, strings.Replace(vars.Set[key], "\"", quoteReplacement, -1))
 	}
 
 	return nil
diff --git a/env_test.go b/env_test.go
@@ -17,6 +17,14 @@ set -x TWO "222";
 set -x ONE "111111";
 set -x THREE "333";
 set -x TWO "222";
+`
+	envFishOutputWithPermCreds = `set -e AWS_SECURITY_TOKEN;
+set -e AWS_SESSION_TOKEN;
+set -x AWS_ACCESS_KEY_ID "aws-key-id";
+set -x AWS_SECRET_ACCESS_KEY "aws-secret-key";
+set -x ONE "111111";
+set -x THREE "333";
+set -x TWO "222";
 `
 
 	envShOutput = `export ONE="111111"
@@ -28,6 +36,14 @@ export TWO="222"
 export ONE="111111"
 export THREE="333"
 export TWO="222"
+`
+	envShOutputWithPermCreds = `unset AWS_SECURITY_TOKEN
+unset AWS_SESSION_TOKEN
+export AWS_ACCESS_KEY_ID="aws-key-id"
+export AWS_SECRET_ACCESS_KEY="aws-secret-key"
+export ONE="111111"
+export THREE="333"
+export TWO="222"
 `
 )
 
@@ -106,4 +122,45 @@ func TestEng(t *testing.T) {
 	if string(output) != envShOutputWithHint {
 		t.Fatalf("Incorrect output: %s", output)
 	}
+
+	// cached environment
+	steward.Environments["one"] = &vaulted.Environment{
+		AWSCreds: &vaulted.AWSCredentials{
+			ID:     "aws-key-id",
+			Secret: "aws-secret-key",
+		},
+		Vars: map[string]string{
+			"TWO":   "222",
+			"ONE":   "111111",
+			"THREE": "333",
+		},
+	}
+
+	output = CaptureStdout(func() {
+		e := Env{
+			VaultName: "one",
+			Shell:     "fish",
+		}
+		err := e.Run(steward)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+	if string(output) != envFishOutputWithPermCreds {
+		t.Fatalf("Incorrect output: %s", output)
+	}
+
+	output = CaptureStdout(func() {
+		e := Env{
+			VaultName: "one",
+			Shell:     "sh",
+		}
+		err := e.Run(steward)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+	if string(output) != envShOutputWithPermCreds {
+		t.Fatalf("Incorrect output: %s", output)
+	}
 }
diff --git a/lib/environment.go b/lib/environment.go
@@ -132,40 +132,61 @@ func (e *Environment) startProxyKeyring() (string, error) {
 	return sock, err
 }
 
-func (e *Environment) buildEnviron(extraVars map[string]string) []string {
-	// load the current environ
-	env := make(map[string]string)
-	for _, envVar := range os.Environ() {
-		parts := strings.SplitN(envVar, "=", 2)
-		env[parts[0]] = parts[1]
+func (e *Environment) Variables() *Variables {
+	vars := Variables{
+		Set: make(map[string]string),
 	}
 
-	// merge the vars
 	for key, value := range e.Vars {
-		env[key] = value
-	}
-	for key, value := range extraVars {
-		env[key] = value
+		vars.Set[key] = value
 	}
 
 	if e.AWSCreds != nil {
-		delete(env, "AWS_ACCESS_KEY_ID")
-		delete(env, "AWS_SECRET_ACCESS_KEY")
-		delete(env, "AWS_SESSION_TOKEN")
-		delete(env, "AWS_SECURITY_TOKEN")
+		vars.Set["AWS_ACCESS_KEY_ID"] = e.AWSCreds.ID
+		vars.Set["AWS_SECRET_ACCESS_KEY"] = e.AWSCreds.Secret
 
-		env["AWS_ACCESS_KEY_ID"] = e.AWSCreds.ID
-		env["AWS_SECRET_ACCESS_KEY"] = e.AWSCreds.Secret
 		if e.AWSCreds.Token != "" {
-			env["AWS_SESSION_TOKEN"] = e.AWSCreds.Token
-			env["AWS_SECURITY_TOKEN"] = e.AWSCreds.Token
+			vars.Set["AWS_SESSION_TOKEN"] = e.AWSCreds.Token
+			vars.Set["AWS_SECURITY_TOKEN"] = e.AWSCreds.Token
+		} else {
+			vars.Unset = append(
+				vars.Unset,
+				"AWS_SESSION_TOKEN",
+				"AWS_SECURITY_TOKEN",
+			)
 		}
 	}
 
-	// recombine into environ
-	environ := make([]string, 0, len(env))
-	for key, value := range env {
+	return &vars
+}
+
+func (e *Environment) buildEnviron(extraVars map[string]string) []string {
+	vars := make(map[string]string)
+	for _, v := range os.Environ() {
+		parts := strings.SplitN(v, "=", 2)
+		vars[parts[0]] = parts[1]
+	}
+
+	v := e.Variables()
+	for _, key := range v.Unset {
+		delete(vars, key)
+	}
+	for key, value := range v.Set {
+		vars[key] = value
+	}
+
+	for key, value := range extraVars {
+		vars[key] = value
+	}
+
+	environ := make([]string, 0, len(vars))
+	for key, value := range vars {
 		environ = append(environ, fmt.Sprintf("%s=%s", key, value))
 	}
 	return environ
 }
+
+type Variables struct {
+	Set   map[string]string
+	Unset []string
+}
diff --git a/steward_test.go b/steward_test.go
@@ -59,14 +59,16 @@ func WriteStdin(b []byte, f func()) {
 
 func NewTestSteward() *TestSteward {
 	return &TestSteward{
-		Passwords: make(map[string]string),
-		Vaults:    make(map[string]*vaulted.Vault),
+		Passwords:    make(map[string]string),
+		Vaults:       make(map[string]*vaulted.Vault),
+		Environments: make(map[string]*vaulted.Environment),
 	}
 }
 
 type TestSteward struct {
-	Passwords map[string]string
-	Vaults    map[string]*vaulted.Vault
+	Passwords    map[string]string
+	Vaults       map[string]*vaulted.Vault
+	Environments map[string]*vaulted.Environment
 
 	LegacyPassword     string
 	LegacyEnvironments map[string]legacy.Environment
@@ -139,22 +141,38 @@ func (ts TestSteward) GetEnvironment(name string, password *string) (string, *va
 		}
 	}
 
-	vault := ts.Vaults[name]
-
-	env := &vaulted.Environment{
+	e := &vaulted.Environment{
 		Vars:    make(map[string]string),
 		SSHKeys: make(map[string]string),
 	}
+	if _, exists := ts.Environments[name]; exists {
+		env := ts.Environments[name]
 
-	for key, value := range vault.Vars {
-		env.Vars[key] = value
-	}
+		e.Expiration = env.Expiration
 
-	for key, value := range vault.SSHKeys {
-		env.SSHKeys[key] = value
+		creds := *env.AWSCreds
+		e.AWSCreds = &creds
+
+		for key, value := range env.Vars {
+			e.Vars[key] = value
+		}
+
+		for key, value := range env.SSHKeys {
+			e.SSHKeys[key] = value
+		}
+	} else {
+		vault := ts.Vaults[name]
+
+		for key, value := range vault.Vars {
+			e.Vars[key] = value
+		}
+
+		for key, value := range vault.SSHKeys {
+			e.SSHKeys[key] = value
+		}
 	}
 
-	return ts.Passwords[name], env, nil
+	return ts.Passwords[name], e, nil
 }
 
 func (ts TestSteward) OpenLegacyVault() (password string, environments map[string]legacy.Environment, err error) {
