From 87e2ce06a9a294bc680ee569b517cd86762c29ce Mon Sep 17 00:00:00 2001
From: Mike Mondragon <mikemondragon@gmail.com>
Date: Wed, 11 Oct 2023 17:02:09 -0700
Subject: [PATCH] Improve All Profiles implementation.

---
 CHANGELOG.md                      |  4 ++--
 README.md                         | 13 ++++++++-----
 internal/webssoauth/webssoauth.go | 26 +++++++++++++++++++-------
 3 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 91cfd8d..31bb289 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -65,8 +65,8 @@ $ okta-aws-cli web \
     --all-profiles
 
 ? Choose an IdP: AWS Account Federation
-Updated profile "myorg-S3-read" in credentials file "/Users/me/.aws/credentials".
-Updated profile "myorg-S3-write" in credentials file "/Users/me/.aws/credentials".
+Updated profile "dev-org-s3-write" in credentials file "/Users/me/.aws/credentials".
+Updated profile "prod-org-Admin-containers" in credentials file "/Users/me/.aws/credentials".
 ```
 
 ### (expected) Alternate web browser open command
diff --git a/README.md b/README.md
index cd2955c..25ed029 100644
--- a/README.md
+++ b/README.md
@@ -390,7 +390,7 @@ These settings are all optional:
 | AWS IAM Identity Provider ARN | Preselects the IdP list to this preferred IAM Identity Provider. If there are other IdPs available they will not be listed. | `--aws-iam-idp [value]` | `OKTA_AWSCLI_IAM_IDP` |
 | Display QR Code | `true` if flag is present | `--qr-code` | `OKTA_AWSCLI_QR_CODE=true` |
 | Automatically open the activation URL with the system web browser | `true` if flag is present | `--open-browser` | `OKTA_AWSCLI_OPEN_BROWSER=true` |
-| Gather all profiles for a given IdP (implies aws-credentials file output format)) | `true` if flag is present | `--all-profiles` | `OKTA_AWSCLI_OPEN_BROWSER=true` |
+| Gather all profiles for all IdPs and Roles associated with an AWS Fed App (implies aws-credentials file output format)) | `true` if flag is present | `--all-profiles` | `OKTA_AWSCLI_OPEN_BROWSER=true` |
 
 #### Allowed Web SSO Client ID
 
@@ -745,8 +745,9 @@ Federation app (IdP) at once.  This is a feature specific to writing the
 aliases is available on the given role) then `-` then abbreviated role name.
 
 ```
-# AWS account alias "myorg", given IdP associated with "AWS Account Federation"
-# and an app associated with two roles.
+# Two Okta "AWS Account Federation" apps, one on AWS alias "prod-org", the other
+# on alias "dev-org". Includes short name for the actual IAM IdPs and IAM Roles
+# on the Fed App.
 
 $ okta-aws-cli web \
     --org-domain test.okta.com \
@@ -755,8 +756,10 @@ $ okta-aws-cli web \
     --all-profiles
 
 ? Choose an IdP: AWS Account Federation
-Updated profile "myorg-S3-read" in credentials file "/Users/me/.aws/credentials".
-Updated profile "myorg-S3-write" in credentials file "/Users/me/.aws/credentials".
+Updated profile "dev-org-s3ops-read" in credentials file "/Users/me/.aws/credentials".
+Updated profile "dev-org-s3ops-write" in credentials file "/Users/me/.aws/credentials".
+Updated profile "prod-org-containerops-ec2-full" in credentials file "/Users/me/.aws/credentials".
+Updated profile "prod-org-containerops-eks-full" in credentials file "/Users/me/.aws/credentials".
 ```
 
 ### Help
diff --git a/internal/webssoauth/webssoauth.go b/internal/webssoauth/webssoauth.go
index e8daba0..0b4a275 100644
--- a/internal/webssoauth/webssoauth.go
+++ b/internal/webssoauth/webssoauth.go
@@ -203,6 +203,16 @@ AWS Federation App with --aws-acct-fed-app-id FED_APP_ID
 	if len(apps) == 1 {
 		// only one app, we don't need to prompt selection of idp / fed app
 		fedAppID = apps[0].ID
+	} else if w.config.AllProfiles() {
+		// special case, we're going to run the table and get all profiles for all apps
+		errArr := []error{}
+		for _, app := range apps {
+			if err = w.establishTokenWithFedAppID(clientID, app.ID, at); err != nil {
+				errArr = append(errArr, err)
+			}
+		}
+
+		return errors.Join(errArr...)
 	} else {
 		// Here, we do want to prompt for selection of the Fed App.
 		// If the app is making use of "Role value pattern" on AWS settings we
@@ -367,15 +377,17 @@ func (w *WebSSOAuthentication) awsAssumeRoleWithSAML(iar *idpAndRole, assertion
 		SessionToken:    *svcResp.Credentials.SessionToken,
 		Expiration:      svcResp.Credentials.Expiration,
 	}
-	if w.config.Profile() != "" {
+	if !w.config.AllProfiles() && w.config.Profile() != "" {
 		cc.Profile = w.config.Profile()
 		return cc, nil
 	}
 
-	var profileName string
-	var roleName string
+	var profileName, idpName, roleName string
+	if _, after, found := strings.Cut(iar.idp, "/"); found {
+		idpName = after
+	}
 	if _, after, found := strings.Cut(iar.role, "/"); found {
-		roleName = "-" + after
+		roleName = after
 	}
 	sessCopy := sess.Copy(&aws.Config{
 		Credentials: credentials.NewStaticCredentials(
@@ -385,12 +397,12 @@ func (w *WebSSOAuthentication) awsAssumeRoleWithSAML(iar *idpAndRole, assertion
 		),
 	})
 	if p, err := w.fetchAWSAccountAlias(sessCopy); err != nil {
-		fmt.Fprintf(os.Stderr, "unable to determine account alias, setting profile name to %q\n", iar.idp)
-		profileName = iar.idp
+		fmt.Fprintf(os.Stderr, "unable to determine account alias, setting alias name to %q\n", "org")
+		profileName = "org"
 	} else {
 		profileName = p
 	}
-	cc.Profile = fmt.Sprintf("%s%s", profileName, roleName)
+	cc.Profile = fmt.Sprintf("%s-%s-%s", profileName, idpName, roleName)
 
 	return cc, nil
 }