From 405d48048b959095570b587346580e709ed8a8ed Mon Sep 17 00:00:00 2001
From: "Eric J. Holmes" <eric@ejholmes.net>
Date: Thu, 7 Apr 2016 13:44:30 +0700
Subject: [PATCH] Clean AWS_ environment vars between evals.

---
 README.md |  8 ++++++++
 main.go   | 18 ++++++++++++++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 109f429..3d61ebb 100644
--- a/README.md
+++ b/README.md
@@ -2,6 +2,14 @@ This is a small utility that makes it easier to use the `aws sts assume-role` co
 
 ## Installation
 
+On OS X, the best way to get it is to use homebrew:
+
+```bash
+brew install remind101/formulae/assume-role
+```
+
+If you have a working Go 1.6 environment:
+
 ```bash
 $ go get -u github.com/remind101/assume-role
 ```
diff --git a/main.go b/main.go
index a0b750f..18d5e07 100644
--- a/main.go
+++ b/main.go
@@ -38,11 +38,17 @@ func main() {
 		must(fmt.Errorf("%s not in ~/.aws/roles", role))
 	}
 
+	if os.Getenv("ASSUMED_ROLE") != "" {
+		// Clear out any previously set AWS_ environment variables so
+		// they aren't used with the assumeRole command.
+		cleanEnv()
+	}
+
 	creds, err := assumeRole(roleConfig.Role, roleConfig.MFA)
 	must(err)
 
 	if len(args) == 0 {
-		printCredentials(creds)
+		printCredentials(role, creds)
 		return
 	}
 
@@ -50,6 +56,13 @@ func main() {
 	must(err)
 }
 
+func cleanEnv() {
+	os.Unsetenv("AWS_ACCESS_KEY_ID")
+	os.Unsetenv("AWS_SECRET_ACCESS_KEY")
+	os.Unsetenv("AWS_SESSION_TOKEN")
+	os.Unsetenv("AWS_SECURITY_TOKEN")
+}
+
 func execWithCredentials(argv []string, creds *credentials) error {
 	argv0, err := exec.LookPath(argv[0])
 	if err != nil {
@@ -73,11 +86,12 @@ type credentials struct {
 
 // printCredentials prints the credentials in a way that can easily be sourced
 // with bash.
-func printCredentials(creds *credentials) {
+func printCredentials(role string, creds *credentials) {
 	fmt.Printf("export AWS_ACCESS_KEY_ID=\"%s\"\n", creds.AccessKeyID)
 	fmt.Printf("export AWS_SECRET_ACCESS_KEY=\"%s\"\n", creds.SecretAccessKey)
 	fmt.Printf("export AWS_SESSION_TOKEN=\"%s\"\n", creds.SessionToken)
 	fmt.Printf("export AWS_SECURITY_TOKEN=\"%s\"\n", creds.SessionToken)
+	fmt.Printf("export ASSUMED_ROLE=\"%s\"\n", role)
 	fmt.Printf("# Run this to configure your shell:\n")
 	fmt.Printf("# eval $(%s)\n", strings.Join(os.Args, " "))
 }