feat: enable 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")

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,31 @@ 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 GetRoleFromSAML(resp *saml.Response, profileARN string) (string, string, error) {
+
+	roles, err := GetAssumableRolesFromSAML(resp)
+	if err != nil {
+		return "", "", err
+	}
+	role, err := GetRole(roles, profileARN)
+	if err != nil {
+		return "", "", err
+	}
+	return role.Principal, role.Role, nil
+}
+
+func GetAssumableRolesFromSAML(resp *saml.Response) (saml.AssumableRoles, error) {
+	roleList := []saml.AssumableRole{}
+
 	for _, a := range resp.Assertion.AttributeStatement.Attributes {
 		if strings.HasSuffix(a.Name, "SAML/Attributes/Role") {
 			for _, v := range a.AttributeValues {
@@ -32,20 +47,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")
+	}
+
+	// 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
 	}
 
-	return "", "", fmt.Errorf("Role '%s' not authorized by Okta.  Contact Okta admin to make sure that the AWS app is configured properly.", profileARN)
+	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) {