From 2eca97a78de5219b439492e0527cda51f4efe634 Mon Sep 17 00:00:00 2001
From: Florin Andrei <901867+FlorinAndrei@users.noreply.github.com>
Date: Thu, 4 Dec 2025 18:40:59 -0800
Subject: [PATCH] allow switch back to base account

---
 package-lock.json                   |  4 ++--
 src/js/content.js                   | 10 +++++++++
 src/js/lib/create_role_list_item.js |  3 +++
 src/js/lib/current_context.js       |  5 +++++
 src/js/lib/target_profiles.js       | 34 ++++++++++++++++++++++++++---
 5 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index d794ddcf..55611d31 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "aws-extend-switch-roles",
-  "version": "6.0.1",
+  "version": "6.0.3",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
     "": {
       "name": "aws-extend-switch-roles",
-      "version": "6.0.1",
+      "version": "6.0.3",
       "license": "MIT",
       "dependencies": {
         "aesr-config": "^0.6.0"
diff --git a/src/js/content.js b/src/js/content.js
index 739a9b1d..9087357d 100644
--- a/src/js/content.js
+++ b/src/js/content.js
@@ -159,6 +159,16 @@ function doSwitch(data, cb) {
 
   const form = document.getElementById('AESR_form');
   form.setAttribute('action', formActionUrl);
+
+  // Handle switch-back to base account
+  if (data.isswitchback === 'true') {
+    form.action.value = 'switchToBasis';
+    form.redirect_uri.value = data.redirecturi;
+    cb({ prism: session.prismModeEnabled, switchBack: true });
+    form.submit();
+    return false;
+  }
+
   form.account.value = data.account;
   form.color.value = data.color;
   form.roleName.value = data.rolename;
diff --git a/src/js/lib/create_role_list_item.js b/src/js/lib/create_role_list_item.js
index adf36e59..10887cd3 100644
--- a/src/js/lib/create_role_list_item.js
+++ b/src/js/lib/create_role_list_item.js
@@ -22,6 +22,9 @@ export function createRoleListItem(document, item, url, region, { hidesAccountId
   anchor.dataset.color = item.color || 'aaaaaa';
   anchor.dataset.redirecturi = createRedirectUri(url, region, item.region);
   anchor.dataset.search = item.name.toLowerCase() + ' ' + item.aws_account_id;
+  if (item.isSwitchBack) {
+    anchor.dataset.isswitchback = 'true';
+  }
 
   anchor.appendChild(headSquare);
   anchor.appendChild(document.createTextNode(item.name));
diff --git a/src/js/lib/current_context.js b/src/js/lib/current_context.js
index 262f40f9..9e545595 100644
--- a/src/js/lib/current_context.js
+++ b/src/js/lib/current_context.js
@@ -16,6 +16,11 @@ export class CurrentContext {
       return extractLoginRole(loginDisplayNameUser.split("/", 2)[0]);
     })();
     this.filterByTargetRole = showOnlyMatchingRoles ? (roleDisplayNameUser || this.loginRole) : null;
+
+    // Track if we're in a delegated role and what the source account is
+    this.isInDelegatedRole = !!(roleDisplayNameAccount || (prism && roleDisplayNameUser));
+    this.sourceAccount = this.isInDelegatedRole ? brushAccountId(loginDisplayNameAccount) : null;
+    this.sourceRole = this.isInDelegatedRole ? extractLoginRole(loginDisplayNameUser.split("/", 2)[0]) : null;
   }
 }
 
diff --git a/src/js/lib/target_profiles.js b/src/js/lib/target_profiles.js
index c74bb043..19d13fd4 100644
--- a/src/js/lib/target_profiles.js
+++ b/src/js/lib/target_profiles.js
@@ -13,7 +13,7 @@ export async function findTargetProfiles(ctx) {
 }
 
 async function retrieveTargetProfilesFromDB(ctx) {
-  const { baseAccount, loginRole, filterByTargetRole } = ctx;
+  const { baseAccount, loginRole, filterByTargetRole, isInDelegatedRole, sourceAccount, sourceRole } = ctx;
   const dbManager = new DBManager('aesr');
   await dbManager.open();
 
@@ -32,6 +32,20 @@ async function retrieveTargetProfilesFromDB(ctx) {
       }
       results.push(...targets)
     }
+
+    // If we're in a delegated role, add the base account as a switchable target
+    if (isInDelegatedRole && sourceAccount) {
+      const baseProfile = complexSrcItems.find(it => matchSourceProfile(it, sourceAccount, sourceRole));
+      if (baseProfile) {
+        const switchBackProfile = {
+          ...baseProfile,
+          isSwitchBack: true,
+          profile: `← ${baseProfile.name}`,
+        };
+        // Add it at the beginning for easy access
+        results.unshift(switchBackProfile);
+      }
+    }
   }, 'readonly');
 
   await dbManager.close();
@@ -58,8 +72,8 @@ function convertComplexTarget(item, baseProfile) {
 }
 
 async function retrieveTargetProfilesFromLztext(ctx) {
-  const { baseAccount, loginRole, filterByTargetRole } = ctx;
-  
+  const { baseAccount, loginRole, filterByTargetRole, isInDelegatedRole, sourceAccount, sourceRole } = ctx;
+
   const localRepo = StorageProvider.getLocalRepository();
   const cfgText = await loadConfigIni(localRepo);
   if (!cfgText) return [];
@@ -77,6 +91,20 @@ async function retrieveTargetProfilesFromLztext(ctx) {
     results.push(...targets)
   }
 
+  // If we're in a delegated role, add the base account as a switchable target
+  if (isInDelegatedRole && sourceAccount) {
+    const baseProfile = profileSet.complexes.find(it => matchSourceProfile(it, sourceAccount, sourceRole));
+    if (baseProfile) {
+      const switchBackProfile = {
+        ...baseProfile,
+        isSwitchBack: true,
+        profile: `← ${baseProfile.name}`,
+      };
+      // Add it at the beginning for easy access
+      results.unshift(switchBackProfile);
+    }
+  }
+
   return results;
 }