From 2ffe200f5874cf691d21512da435b1d020860036 Mon Sep 17 00:00:00 2001
From: Aaron Turner <synfinatic@gmail.com>
Date: Sat, 15 Jan 2022 16:36:12 -0800
Subject: [PATCH] Correctly set `AWS_SSO` env variable

We were incorrectly leaving this blank so it was being
unset in `eval` and set to an empty string in `exec`

Fixes: #251
---
 CHANGELOG.md         |  4 ++++
 cmd/exec_cmd.go      |  3 ++-
 sso/settings.go      | 29 ++++++++++++++++++++++-------
 sso/settings_test.go | 22 +++++++++++++++++++---
 4 files changed, 47 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e5ff04e6..f570a43f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [1.7.1] - Unreleased
 
+### Bug Fixes
+
+ * `AWS_SSO` env var is now set with the `eval` and `exec` command #251 
+
 ### Changes
 
  * `flush` now flushes the STS IAM Role credentials first by default #236
diff --git a/cmd/exec_cmd.go b/cmd/exec_cmd.go
index 824d3702..115203d3 100644
--- a/cmd/exec_cmd.go
+++ b/cmd/exec_cmd.go
@@ -132,6 +132,7 @@ func execShellEnvs(ctx *RunContext, awssso *sso.AWSSSO, accountid int64, role, r
 	credsPtr := GetRoleCredentials(ctx, awssso, accountid, role)
 	creds := *credsPtr
 
+	ssoName, _ := ctx.Settings.GetSelectedSSOName(ctx.Cli.SSO)
 	shellVars := map[string]string{
 		"AWS_ACCESS_KEY_ID":          creds.AccessKeyId,
 		"AWS_SECRET_ACCESS_KEY":      creds.SecretAccessKey,
@@ -140,7 +141,7 @@ func execShellEnvs(ctx *RunContext, awssso *sso.AWSSSO, accountid int64, role, r
 		"AWS_SSO_ROLE_NAME":          creds.RoleName,
 		"AWS_SSO_SESSION_EXPIRATION": creds.ExpireString(),
 		"AWS_SSO_ROLE_ARN":           utils.MakeRoleARN(creds.AccountId, creds.RoleName),
-		"AWS_SSO":                    ctx.Cli.SSO,
+		"AWS_SSO":                    ssoName,
 	}
 
 	if len(region) > 0 {
diff --git a/sso/settings.go b/sso/settings.go
index 338373e2..2bca9613 100644
--- a/sso/settings.go
+++ b/sso/settings.go
@@ -287,18 +287,33 @@ func (s *Settings) CreatedAt() int64 {
 // GetSelectedSSO returns a valid SSOConfig based on user intput, configured
 // value or our hardcoded 'Default' if it exists and name is empty String
 func (s *Settings) GetSelectedSSO(name string) (*SSOConfig, error) {
-	if c, ok := s.SSO[name]; ok {
-		return c, nil
+	n, err := s.GetSelectedSSOName(name)
+	if err != nil {
+		return &SSOConfig{}, err
+	}
+	return s.SSO[n], nil
+}
+
+// GetSelectedSSOName returns the name of the selected SSO name where
+// the input is the option passed in via the CLI (should be an empty string)
+// if user did not specify a value on the CLI
+func (s *Settings) GetSelectedSSOName(name string) (string, error) {
+	if name != "" {
+		if _, ok := s.SSO[name]; ok {
+			return name, nil
+		}
+
+		return "", fmt.Errorf("'%s' is not a valid AWS SSO Instance", name)
 	}
 
-	if c, ok := s.SSO[s.DefaultSSO]; ok && s.DefaultSSO != "Default" {
-		return c, nil
+	if _, ok := s.SSO[s.DefaultSSO]; ok {
+		return s.DefaultSSO, nil
 	}
 
-	if c, ok := s.SSO["Default"]; ok && name == "" {
-		return c, nil
+	if _, ok := s.SSO["Default"]; ok {
+		return "Default", nil
 	}
-	return &SSOConfig{}, fmt.Errorf("No available SSOConfig Provider")
+	return "", fmt.Errorf("No available AWS SSO Instance")
 }
 
 // Returns the Tag name => Environment variable name
diff --git a/sso/settings_test.go b/sso/settings_test.go
index 382cc5c2..81c2bd2d 100644
--- a/sso/settings_test.go
+++ b/sso/settings_test.go
@@ -65,18 +65,34 @@ func (suite *SettingsTestSuite) TestGetSelectedSSO() {
 	t := suite.T()
 
 	sso, err := suite.settings.GetSelectedSSO("Default")
-	assert.Nil(t, err)
+	assert.NoError(t, err)
 	assert.Equal(t, "https://d-754545454.awsapps.com/start", sso.StartUrl)
 
 	sso, err = suite.settings.GetSelectedSSO("Foobar")
-	assert.NotNil(t, err)
+	assert.Error(t, err)
 	assert.Equal(t, "", sso.StartUrl)
 
 	sso, err = suite.settings.GetSelectedSSO("")
-	assert.Nil(t, err)
+	assert.NoError(t, err)
 	assert.Equal(t, "https://d-754545454.awsapps.com/start", sso.StartUrl)
 }
 
+func (suite *SettingsTestSuite) TestGetSelectedSSOName() {
+	t := suite.T()
+
+	name, err := suite.settings.GetSelectedSSOName("Default")
+	assert.NoError(t, err)
+	assert.Equal(t, "Default", name)
+
+	name, err = suite.settings.GetSelectedSSOName("Foobar")
+	assert.Error(t, err)
+	assert.Equal(t, "", name)
+
+	name, err = suite.settings.GetSelectedSSOName("")
+	assert.NoError(t, err)
+	assert.Equal(t, "Default", name)
+}
+
 func (suite *SettingsTestSuite) TestCreatedAt() {
 	t := suite.T()
 	sso, _ := suite.settings.GetSelectedSSO("")