diff --git a/internal/webssoauth/webssoauth.go b/internal/webssoauth/webssoauth.go index 52483fd..6ded0f0 100644 --- a/internal/webssoauth/webssoauth.go +++ b/internal/webssoauth/webssoauth.go @@ -479,63 +479,53 @@ func (w *WebSSOAuthentication) choiceFriendlyLabelRole(arn string, roles map[str } // promptForRole prompt operator for the AWS Role ARN given a slice of Role ARNs -func (w *WebSSOAuthentication) promptForRole(idp string, roleARNs []string) (roleARN string, err error) { - oktaConfig, err := config.OktaConfig() - var configRoles map[string]string - if err == nil { - configRoles = oktaConfig.AWSCLI.ROLES +func (w *WebSSOAuthentication) promptForRole(idp string, roleARNs []string, configRoles map[string]string) (roleARN string, err error) { + // roleLabels are the friendly names if configured or the ARNs themselves + roleLabels := make([]string, len(roleARNs)) + roleArnByLabel := map[string]string{} + for _, arn := range roleARNs { + roleLabel := w.choiceFriendlyLabelRole(arn, configRoles) + roleLabels = append(roleLabels, roleLabel) + roleArnByLabel[roleLabel] = arn } - if len(roleARNs) == 1 || w.config.AWSIAMRole() != "" { - roleARN = w.config.AWSIAMRole() - if len(roleARNs) == 1 { - roleARN = roleARNs[0] - } - roleLabel := w.choiceFriendlyLabelRole(roleARN, configRoles) - roleData := roleTemplateData{ - Role: roleLabel, - } - - // reverse case when friendly role name alias is given as the input value - // --aws-iam-role "OK S3 Read" - if roleLabel == roleARN { - for rARN, rLbl := range configRoles { - if roleARN == rLbl { - roleARN = rARN - break - } - } - } + var roleLabelChoice string - if !w.config.IsProcessCredentialsFormat() { - rich, _, err := core.RunTemplate(roleSelectedTemplate, roleData) - if err != nil { - return "", err - } - fmt.Fprintln(os.Stderr, rich) - } - return roleARN, nil + // There is only a single choice so go ahead and use its label + if len(roleARNs) == 1 { + rArn := roleARNs[0] + roleLabelChoice = w.choiceFriendlyLabelRole(rArn, configRoles) } - promptRoles := []string{} - labelsARNs := map[string]string{} - for _, arn := range roleARNs { - roleLabel := w.choiceFriendlyLabelRole(arn, configRoles) - promptRoles = append(promptRoles, roleLabel) - labelsARNs[roleLabel] = arn + // The user already provided their choice via config + if roleLabelChoice == "" && w.config.AWSIAMRole() != "" { + rArg := w.config.AWSIAMRole() + roleLabelChoice = w.choiceFriendlyLabelRole(rArg, configRoles) } - prompt := &survey.Select{ - Message: chooseRole, - Options: promptRoles, - } - var selected string - err = survey.AskOne(prompt, &selected, survey.WithValidator(survey.Required), stderrIsOutAskOpt) - if err != nil { - return "", fmt.Errorf(askRoleError, err) + // Prompt the user to choose + if roleLabelChoice == "" { + prompt := &survey.Select{ + Message: chooseRole, + Options: roleLabels, + } + err = survey.AskOne(prompt, &roleLabelChoice, survey.WithValidator(survey.Required), stderrIsOutAskOpt) + if err != nil { + return "", fmt.Errorf(askRoleError, err) + } + } else if !w.config.IsProcessCredentialsFormat() { + // The choice was determined without prompting the user so pretty print the role + // todo: explain why we check IsProcessCredentialsFormat? + rich, _, err := core.RunTemplate(roleSelectedTemplate, roleTemplateData{ + Role: roleLabelChoice, + }) + if err != nil { + return "", err + } + fmt.Fprintln(os.Stderr, rich) } - roleARN = labelsARNs[selected] + roleARN = roleArnByLabel[roleLabelChoice] if roleARN == "" { return "", fmt.Errorf(noRolesError, idp) } @@ -546,60 +536,61 @@ func (w *WebSSOAuthentication) promptForRole(idp string, roleARNs []string) (rol // promptForIDP prompt operator for the AWS IdP ARN given a slice of IdP ARNs. // If the fedApp has already been selected via an ask one survey we don't need // to pretty print out the IdP name again. -func (w *WebSSOAuthentication) promptForIDP(idpARNs []string) (idpARN string, err error) { - var configIDPs map[string]string - if oktaConfig, cErr := config.OktaConfig(); cErr == nil { - configIDPs = oktaConfig.AWSCLI.IDPS +func (w *WebSSOAuthentication) promptForIDP(idpARNs []string, configIDPs map[string]string) (idpArnChoice string, err error) { + if len(idpARNs) == 0 { + return "", errors.New(noIDPsError) } - if len(idpARNs) == 0 { - return idpARN, errors.New(noIDPsError) + // idpLabels are the friendly names if configured or the ARNs itself + idpLabels := make([]string, len(idpARNs)) + idpArnByLabel := make(map[string]string, len(idpARNs)) + for i, arn := range idpARNs { + idpLabel := w.choiceFriendlyLabelIDP(arn, arn, configIDPs) + idpArnByLabel[idpLabel] = arn + idpLabels[i] = idpLabel } - if len(idpARNs) == 1 || w.config.AWSIAMIdP() != "" { - idpARN = w.config.AWSIAMIdP() - if len(idpARNs) == 1 { - idpARN = idpARNs[0] - } - if w.fedAppAlreadySelected { - return idpARN, nil - } + var idpLabelChoice string + + // There is only a single choice so go ahead and use its label + if len(idpARNs) == 1 { + idpArn := idpARNs[0] + idpLabelChoice = w.choiceFriendlyLabelIDP(idpArn, idpArn, configIDPs) + } + + // The user already provided their choice via config + if idpLabelChoice == "" && w.config.AWSIAMIdP() != "" { + iArg := w.config.AWSIAMIdP() + idpLabelChoice = w.choiceFriendlyLabelIDP(iArg, iArg, configIDPs) + } - idpLabel := w.choiceFriendlyLabelIDP(idpARN, idpARN, configIDPs) - idpData := idpTemplateData{ - IDP: idpLabel, + // Prompt the user to choose + if idpLabelChoice == "" { + prompt := &survey.Select{ + Message: chooseIDP, + Options: idpLabels, + } + err = survey.AskOne(prompt, &idpLabelChoice, survey.WithValidator(survey.Required), stderrIsOutAskOpt) + if err != nil { + return "", fmt.Errorf(askIDPError, err) } - rich, _, err := core.RunTemplate(idpSelectedTemplate, idpData) + } else if !w.fedAppAlreadySelected { + // The choice was determined without prompting the user and the fedApp has not already been selected so pretty print the idp + rich, _, err := core.RunTemplate(idpSelectedTemplate, idpTemplateData{ + IDP: idpLabelChoice, + }) if err != nil { return "", err } fmt.Fprintln(os.Stderr, rich) - return idpARN, nil } - idpChoices := make(map[string]string, len(idpARNs)) - idpChoiceLabels := make([]string, len(idpARNs)) - for i, arn := range idpARNs { - idpLabel := w.choiceFriendlyLabelIDP(arn, arn, configIDPs) - idpChoices[idpLabel] = arn - idpChoiceLabels[i] = idpLabel + idpArnChoice = idpArnByLabel[idpLabelChoice] + if idpArnChoice == "" { + return idpArnChoice, errors.New(idpValueNotSelectedError) } - var idpChoice string - prompt := &survey.Select{ - Message: chooseIDP, - Options: idpChoiceLabels, - } - err = survey.AskOne(prompt, &idpChoice, survey.WithValidator(survey.Required), stderrIsOutAskOpt) - if err != nil { - return idpARN, fmt.Errorf(askIDPError, err) - } - idpARN = idpChoices[idpChoice] - if idpARN == "" { - return idpARN, errors.New(idpValueNotSelectedError) - } - - return idpARN, nil + return idpArnChoice, nil } // promptForIdpAndRole UX to prompt operator for the AWS role whose credentials @@ -609,13 +600,22 @@ func (w *WebSSOAuthentication) promptForIdpAndRole(idpRoles map[string][]string) for idp := range idpRoles { idps = append(idps, idp) } - idp, err := w.promptForIDP(idps) + + var configRoles map[string]string + var configIDPs map[string]string + + if oktaConfig, cErr := config.OktaConfig(); cErr == nil { + configRoles = oktaConfig.AWSCLI.ROLES + configIDPs = oktaConfig.AWSCLI.IDPS + } + + idp, err := w.promptForIDP(idps, configIDPs) if err != nil { return nil, err } roles := idpRoles[idp] - role, err := w.promptForRole(idp, roles) + role, err := w.promptForRole(idp, roles, configRoles) if err != nil { return nil, err } diff --git a/internal/webssoauth/webssoauth_test.go b/internal/webssoauth/webssoauth_test.go index 49c89a9..a62ebdf 100644 --- a/internal/webssoauth/webssoauth_test.go +++ b/internal/webssoauth/webssoauth_test.go @@ -265,3 +265,199 @@ func TestChoiceFriendlyLabelRole(t *testing.T) { }) } } +func TestPromptForRole(t *testing.T) { + testCases := []struct { + name string + idpARN string + configRoles map[string]string + roleARNs []string + roleArg string + expected string + }{ + { + name: "friendly label", + idpARN: "arn:aws:iam::123:role/rickrole", + roleArg: "Rock N Role", + roleARNs: []string{ + "arn:aws:iam::123:role/rocknrole", + "arn:aws:iam::123:role/rickrole", + }, + configRoles: map[string]string{ + "arn:aws:iam::123:role/rocknrole": "Rock N Role", + "arn:aws:iam::123:role/rickrole": "Rick Role", + "arn:aws:iam::.*:role/never": "Never Gonna Give You Up", + }, + expected: "arn:aws:iam::123:role/rocknrole", + }, + { + name: "friendly label configured but arn arg supplied", + idpARN: "arn:aws:iam::123:role/rickrole", + roleArg: "arn:aws:iam::123:role/rocknrole", + roleARNs: []string{ + "arn:aws:iam::123:role/rocknrole", + "arn:aws:iam::123:role/rickrole", + }, + configRoles: map[string]string{ + "arn:aws:iam::123:role/rocknrole": "Rock N Role", + "arn:aws:iam::123:role/rickrole": "Rick Role", + "arn:aws:iam::.*:role/never": "Never Gonna Give You Up", + }, + expected: "arn:aws:iam::123:role/rocknrole", + }, + { + name: "friendly label with wildcard", + idpARN: "arn:aws:iam::123:role/rickrole", + roleArg: "Never Gonna Give You Up", + roleARNs: []string{ + "arn:aws:iam::123:role/never", + "arn:aws:iam::123:role/rocknrole", + }, + configRoles: map[string]string{ + "arn:aws:iam::123:role/rocknrole": "Rock N Role", + "arn:aws:iam::123:role/rickrole": "Rick Role", + "arn:aws:iam::.*:role/never": "Never Gonna Give You Up", + }, + expected: "arn:aws:iam::123:role/never", + }, + { + name: "no friendly labels arn arg supplied", + idpARN: "arn:aws:iam::123:role/rickrole", + roleArg: "arn:aws:iam::123:role/rocknrole", + roleARNs: []string{ + "arn:aws:iam::123:role/never", + "arn:aws:iam::123:role/rocknrole", + }, + configRoles: nil, + expected: "arn:aws:iam::123:role/rocknrole", + }, + { + name: "single arn option no arg supplied", + idpARN: "arn:aws:iam::123:role/rickrole", + roleArg: "", + roleARNs: []string{ + "arn:aws:iam::123:role/never", + }, + configRoles: nil, + expected: "arn:aws:iam::123:role/never", + }, + } + t.Parallel() + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + cfg, err := config.NewConfig(&config.Attributes{ + AWSIAMRole: tc.roleArg, + }) + require.NoError(t, err) + + w, err := NewWebSSOAuthentication(cfg) + roleARn, err := w.promptForRole(tc.idpARN, tc.roleARNs, tc.configRoles) + if roleARn != tc.expected { + t.Errorf("expected %q, got %q", tc.expected, roleARn) + } + }) + } +} + +func TestPromptForIdp(t *testing.T) { + testCases := []struct { + name string + configIdps map[string]string + idpARNs []string + idpArg string + expected string + }{ + { + name: "friendly label", + idpArg: "My IdP", + idpARNs: []string{ + "arn:aws:iam::123:saml-provider/youridp", + "arn:aws:iam::123:saml-provider/myidp", + "arn:aws:iam::123:saml-provider/aidp", + }, + configIdps: map[string]string{ + "arn:aws:iam::123:saml-provider/youridp": "Your IdP", + "arn:aws:iam::123:saml-provider/myidp": "My IdP", + "arn:aws:iam::.*:saml-provider/aidp": "A IdP", + }, + expected: "arn:aws:iam::123:saml-provider/myidp", + }, + { + name: "friendly label configured but arn arg supplied", + idpArg: "arn:aws:iam::123:saml-provider/myidp", + idpARNs: []string{ + "arn:aws:iam::123:saml-provider/youridp", + "arn:aws:iam::123:saml-provider/myidp", + "arn:aws:iam::123:saml-provider/aidp", + }, + configIdps: map[string]string{ + "arn:aws:iam::123:saml-provider/youridp": "Your IdP", + "arn:aws:iam::123:saml-provider/myidp": "My IdP", + "arn:aws:iam::.*:saml-provider/aidp": "A IdP", + }, + expected: "arn:aws:iam::123:saml-provider/myidp", + }, + { + name: "friendly label with wildcard", + idpArg: "A IdP", + idpARNs: []string{ + "arn:aws:iam::123:saml-provider/youridp", + "arn:aws:iam::123:saml-provider/myidp", + "arn:aws:iam::123:saml-provider/aidp", + }, + configIdps: map[string]string{ + "arn:aws:iam::123:saml-provider/youridp": "Your IdP", + "arn:aws:iam::123:saml-provider/myidp": "My IdP", + "arn:aws:iam::.*:saml-provider/aidp": "A IdP", + }, + expected: "arn:aws:iam::123:saml-provider/aidp", + }, + { + name: "no friendly labels arn arg supplied", + idpArg: "arn:aws:iam::123:saml-provider/youridp", + idpARNs: []string{ + "arn:aws:iam::123:saml-provider/youridp", + "arn:aws:iam::123:saml-provider/myidp", + "arn:aws:iam::123:saml-provider/aidp", + }, + configIdps: nil, + expected: "arn:aws:iam::123:saml-provider/youridp", + }, + { + name: "single arn option no arg supplied", + idpArg: "", + idpARNs: []string{ + "arn:aws:iam::123:saml-provider/myidp", + }, + configIdps: nil, + expected: "arn:aws:iam::123:saml-provider/myidp", + }, + { + name: "single arn option no arg supplied with friendly label", + idpArg: "", + idpARNs: []string{ + "arn:aws:iam::123:saml-provider/myidp", + }, + configIdps: map[string]string{ + "arn:aws:iam::123:saml-provider/youridp": "Your IdP", + "arn:aws:iam::123:saml-provider/myidp": "My IdP", + "arn:aws:iam::.*:saml-provider/aidp": "A IdP", + }, + expected: "arn:aws:iam::123:saml-provider/myidp", + }, + } + t.Parallel() + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + cfg, err := config.NewConfig(&config.Attributes{ + AWSIAMIdP: tc.idpArg, + }) + require.NoError(t, err) + + w, err := NewWebSSOAuthentication(cfg) + roleARn, err := w.promptForIDP(tc.idpARNs, tc.configIdps) + if roleARn != tc.expected { + t.Errorf("expected %q, got %q", tc.expected, roleARn) + } + }) + } +}