feat: enablet prompt if role not provided

If the `role_arn` is provided, attempt to assume that role. No change to
current behaviour.

If `role_arn` is empty or not present in the profile and there is more than one
role then display the full list of roles the user can assume and prompt
them to choose the role to assume. The choice will be sticky for the
rest of the session for that profile. Whne the session ends they will be
prompted again to choose a role.

lib/okta.go

@@ -189,6 +189,7 @@ func (o *OktaClient) AuthenticateProfile(profileARN string, duration time.Durati
 
 	// Attempt to reuse session cookie
 	var assertion SAMLAssertion
+
 	err := o.Get("GET", o.OktaAwsSAMLUrl, nil, &assertion, "saml")
 	if err != nil {
 		log.Debug("Failed to reuse session token, starting flow from start")
@@ -204,26 +205,29 @@ func (o *OktaClient) AuthenticateProfile(profileARN string, duration time.Durati
 			return sts.Credentials{}, "", err
 		}
 	}
-
-	principal, role, err := GetRoleFromSAML(assertion.Resp, profileARN)
+	roleList, err := GetAssumableRolesFromSAML(assertion.Resp)
 	if err != nil {
 		return sts.Credentials{}, "", err
 	}
 
+	assumableRole, err := GetRole(roleList, profileARN)
+	if err != nil {
+		return sts.Credentials{}, "", err
+	}
 	// Step 4 : Assume Role with SAML
 	samlSess := session.Must(session.NewSession())
 	svc := sts.New(samlSess)
 
 	samlParams := &sts.AssumeRoleWithSAMLInput{
-		PrincipalArn:    aws.String(principal),
-		RoleArn:         aws.String(role),
+		PrincipalArn:    aws.String(assumableRole.Principal),
+		RoleArn:         aws.String(assumableRole.Role),
 		SAMLAssertion:   aws.String(string(assertion.RawData)),
 		DurationSeconds: aws.Int64(int64(duration.Seconds())),
 	}
 
 	samlResp, err := svc.AssumeRoleWithSAML(samlParams)
 	if err != nil {
-		log.WithField("role", role).Errorf(
+		log.WithField("role", assumableRole.Role).Errorf(
 			"error assuming role with SAML: %s", err.Error())
 		return sts.Credentials{}, "", err
 	}

lib/provider.go

@@ -161,10 +161,7 @@ func (p *Provider) getSamlSessionCreds() (sts.Credentials, error) {
 	}
 	oktaSessionCookieKey := p.getOktaSessionCookieKey()
 
-	profileARN, ok := p.profiles[source]["role_arn"]
-	if !ok {
-		return sts.Credentials{}, errors.New("Source profile must provide `role_arn`")
-	}
+	profileARN := p.profiles[source]["role_arn"]
 
 	provider := OktaProvider{
 		MFAConfig:            p.ProviderOptions.MFAConfig,
@@ -192,10 +189,7 @@ func (p *Provider) GetSAMLLoginURL() (*url.URL, error) {
 	}
 	oktaSessionCookieKey := p.getOktaSessionCookieKey()
 
-	profileARN, ok := p.profiles[source]["role_arn"]
-	if !ok {
-		return &url.URL{}, errors.New("Source profile must provide `role_arn`")
-	}
+	profileARN := p.profiles[source]["role_arn"]
 
 	provider := OktaProvider{
 		MFAConfig:            p.ProviderOptions.MFAConfig,

lib/saml/struct.go

@@ -2,6 +2,13 @@ package saml
 
 import "encoding/xml"
 
+type AssumableRole struct {
+	Role      string
+	Principal string
+}
+
+type AssumableRoles []AssumableRole
+
 type Response struct {
 	XMLName      xml.Name
 	SAMLP        string `xml:"xmlns:samlp,attr"`

lib/utils.go

@@ -3,16 +3,18 @@ package lib
 import (
 	"encoding/base64"
 	"encoding/xml"
+	"errors"
 	"fmt"
+	"strconv"
 	"strings"
 
 	"github.com/segmentio/aws-okta/lib/saml"
 	"golang.org/x/net/html"
 )
 
-//TODO: Move those functions into saml package
+func GetAssumableRolesFromSAML(resp *saml.Response) (saml.AssumableRoles, error) {
+	roleList := []saml.AssumableRole{}
 
-func GetRoleFromSAML(resp *saml.Response, profileARN string) (string, string, error) {
 	for _, a := range resp.Assertion.AttributeStatement.Attributes {
 		if strings.HasSuffix(a.Name, "SAML/Attributes/Role") {
 			for _, v := range a.AttributeValues {
@@ -32,20 +34,67 @@ func GetRoleFromSAML(resp *saml.Response, profileARN string) (string, string, er
 				// In practice, though, Okta SAML integrations with AWS will succeed
 				// with either the role or principal ARN first, and these `if` statements
 				// allow that behavior in this program.
-				if tokens[0] == profileARN {
-					// if true, Role attribute is formatted like:
-					// arn:aws:iam::account:role/roleName,arn:aws:iam::ACCOUNT:saml-provider/provider
-					return tokens[1], tokens[0], nil
-				} else if tokens[1] == profileARN {
+				if strings.Contains(tokens[0], ":saml-provider/") {
 					// if true, Role attribute is formatted like:
 					// arn:aws:iam::ACCOUNT:saml-provider/provider,arn:aws:iam::account:role/roleName
-					return tokens[0], tokens[1], nil
+					roleList = append(roleList, saml.AssumableRole{Role: tokens[1],
+						Principal: tokens[0]})
+				} else if strings.Contains(tokens[1], ":saml-provider/") {
+					// if true, Role attribute is formatted like:
+					// arn:aws:iam::account:role/roleName,arn:aws:iam::ACCOUNT:saml-provider/provider
+					roleList = append(roleList, saml.AssumableRole{Role: tokens[0],
+						Principal: tokens[1]})
+				} else {
+					return saml.AssumableRoles{}, fmt.Errorf("Unable to get roles from %s", v.Value)
 				}
+
 			}
 		}
 	}
+	return roleList, nil
+}
+
+func GetRole(roleList saml.AssumableRoles, profileARN string) (saml.AssumableRole, error) {
+
+	// if the user doesn't have any roles they can assume return an error.
+	if len(roleList) == 0 {
+		return saml.AssumableRole{}, fmt.Errorf("There are no roles that can be assumed")
+	}
 
-	return "", "", fmt.Errorf("Role '%s' not authorized by Okta.  Contact Okta admin to make sure that the AWS app is configured properly.", profileARN)
+	// A role arn was provided as part of the profile, we will assume that role.
+	if profileARN != "" {
+		for _, arole := range roleList {
+			if profileARN == arole.Role {
+				return arole, nil
+			}
+		}
+		return saml.AssumableRole{}, fmt.Errorf("ARN isn't valid")
+	}
+
+	// if the user only has one role assume that role without prompting.
+	if len(roleList) == 1 {
+		return roleList[0], nil
+	}
+
+	for i, arole := range roleList {
+		fmt.Printf("%d - %s\n", i, arole.Role)
+	}
+
+	i, err := Prompt("Select Role to Assume", false)
+	if err != nil {
+		return saml.AssumableRole{}, err
+	}
+	if i == "" {
+		return saml.AssumableRole{}, errors.New("Invalid selection - Please use an option that is listed")
+	}
+	factorIdx, err := strconv.Atoi(i)
+	if err != nil {
+		return saml.AssumableRole{}, err
+	}
+	if factorIdx > (len(roleList) - 1) {
+		return saml.AssumableRole{}, errors.New("Invalid selection - Please use an option that is listed")
+	}
+	return roleList[factorIdx], nil
 }
 
 func ParseSAML(body []byte, resp *SAMLAssertion) (err error) {