From 7eb77b66df727bcfd5c41382de9c952f3ef2c4e6 Mon Sep 17 00:00:00 2001
From: Darren Tu <darrentu@users.noreply.github.com>
Date: Mon, 30 Sep 2024 13:27:06 -0700
Subject: [PATCH] Update PSCloudShellStartup.ps1 (#488)

---
 linux/powershell/PSCloudShellStartup.ps1 | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/linux/powershell/PSCloudShellStartup.ps1 b/linux/powershell/PSCloudShellStartup.ps1
index f7709979..c5a45099 100644
--- a/linux/powershell/PSCloudShellStartup.ps1
+++ b/linux/powershell/PSCloudShellStartup.ps1
@@ -13,7 +13,8 @@ $script:SkipMSIAuth = [System.Environment]::GetEnvironmentVariable('SkipMSIAuth'
 Microsoft.PowerShell.Core\Import-Module -Name PSCloudShellUtility
 $script:PSCloudShellUtilityModuleInfo = Microsoft.PowerShell.Core\Get-Module PSCloudShellUtility
 
-$script:CloudEnvironmentMap = @{
+# This Cloud env map is for Connect-AzureAD
+$script:CloudEnvironmentMapAzureAD = @{
     PROD = 'AzureCloud'; 
     Fairfax = 'AzureUSGovernment'; 
     Mooncake = 'AzureChinaCloud';
@@ -23,6 +24,16 @@ $script:CloudEnvironmentMap = @{
     USSec = 'AzureUSGovernment3'
 }
 
+# This Cloud Env Map is for Connect-AzAccount
+# We need to have two because the mapping is slightly different for the two commands
+$script:CloudEnvironmentMapAzAccount = @{
+    PROD = 'AzureCloud'; 
+    Fairfax = 'AzureUSGovernment'; 
+    Mooncake = 'AzureChinaCloud';
+    BlackForest = 'AzureGermanCloud';
+    dogfood = 'dogfood';
+}
+
 # For the Az.Tools.Predictor
 PSReadline\Set-PSReadLineOption -Colors @{ InLinePrediction = '#8d8d8d'}
 Microsoft.PowerShell.Core\Import-Module Az.Tools.Predictor -Force
@@ -128,9 +139,9 @@ function Connect-AzService
         # the authentication will be successful.
 
         $envName = $env:ACC_CLOUD
-        if ($CloudEnvironmentMap.ContainsKey($env:ACC_CLOUD))
+        if ($CloudEnvironmentMapAzAccount.ContainsKey($env:ACC_CLOUD))
         {
-            $envName = $script:CloudEnvironmentMap[$env:ACC_CLOUD]
+            $envName = $script:CloudEnvironmentMapAzAccount[$env:ACC_CLOUD]
         }
 
         $addAzAccountParameters = @{'Identity' = $true; 'TenantId' = $env:ACC_TID; 'EnvironmentName' = $envName}
@@ -175,9 +186,9 @@ function Connect-AzureAD
     try
     {  
         $envName = $env:ACC_CLOUD
-        if ($CloudEnvironmentMap.ContainsKey($env:ACC_CLOUD))
+        if ($CloudEnvironmentMapAzureAD.ContainsKey($env:ACC_CLOUD))
         {
-            $envName = $script:CloudEnvironmentMap[$env:ACC_CLOUD]
+            $envName = $script:CloudEnvironmentMapAzureAD[$env:ACC_CLOUD]
         }
 
         # Remove AccountId from parameters since it's missing for some users; Plus, it doesn't affect the authorization.