Merge pull request #6939 from MSSedusch/preview

fix AEM cmdlets for other environments e.g. Azure China

Author: markcowl

src/ResourceManager/Compute/Commands.Compute/ChangeLog.md

@@ -20,6 +20,7 @@
 -->
 ## Current Release
 * Fixed issue with default resource groups not being set.
+* Fix AEM Extension cmdlets for other environments, for example Azure China
 
 ## Version 5.5.0
 * Updated to the latest version of the Azure ClientRuntime.

src/ResourceManager/Compute/Commands.Compute/Extension/AEM/AEMHelper.cs

@@ -46,16 +46,18 @@ internal class AEMHelper
         private Dictionary<string, string> _StorageKeyCache = new Dictionary<string, string>(StringComparer.InvariantCultureIgnoreCase);
         private StorageManagementClient _StorageClient;
         private IAzureSubscription _Subscription;
+        private string _StorageEndpoint;
 
         public AEMHelper(Action<ErrorRecord> errorAction, Action<string> verboseAction, Action<string> warningAction,
-            PSHostUserInterface ui, StorageManagementClient storageClient, IAzureSubscription subscription)
+            PSHostUserInterface ui, StorageManagementClient storageClient, IAzureSubscription subscription, String storageEndpoint)
         {
             this._ErrorAction = errorAction;
             this._VerboseAction = verboseAction;
             this._WarningAction = warningAction;
             this._UI = ui;
             this._StorageClient = storageClient;
             this._Subscription = subscription;
+            this._StorageEndpoint = storageEndpoint;
         }
 
         internal string GetStorageAccountFromUri(string uri)
@@ -140,7 +142,7 @@ internal bool IsPremiumStorageAccount(string accountName)
                     StorageCredentialsFactory storageCredentialsFactory = new StorageCredentialsFactory(resGroupName,
                         this._StorageClient, this._Subscription);
                     StorageCredentials sc = storageCredentialsFactory.Create(blobUri);
-                    CloudStorageAccount cloudStorageAccount = new CloudStorageAccount(sc, true);
+                    CloudStorageAccount cloudStorageAccount = new CloudStorageAccount(sc,  this._StorageEndpoint, true);
                     CloudBlobClient blobClient = cloudStorageAccount.CreateCloudBlobClient();
                     CloudBlobContainer blobContainer = blobClient.GetContainerReference(blobUri.BlobContainerName);
                     var cloudBlob = blobContainer.GetPageBlobReference(blobUri.BlobName);
@@ -729,7 +731,7 @@ internal ServiceProperties GetStorageAnalytics(string storageAccountName)
         {
             var key = this.GetAzureStorageKeyFromCache(storageAccountName);
             var credentials = new StorageCredentials(storageAccountName, key);
-            var cloudStorageAccount = new CloudStorageAccount(credentials, true);
+            var cloudStorageAccount = new CloudStorageAccount(credentials, this._StorageEndpoint, true);
             return cloudStorageAccount.CreateCloudBlobClient().GetServicePropertiesAsync().ConfigureAwait(false).GetAwaiter().GetResult();
         }
 

src/ResourceManager/Compute/Commands.Compute/Extension/AEM/RemoveAzureRmVMAEMExtension.cs

@@ -74,7 +74,8 @@ public override void ExecuteCmdlet()
             this._Helper = new AEMHelper((err) => this.WriteError(err), (msg) => this.WriteVerbose(msg), (msg) => this.WriteWarning(msg),
                 this.CommandRuntime.Host.UI,
                 AzureSession.Instance.ClientFactory.CreateArmClient<StorageManagementClient>(DefaultProfile.DefaultContext, AzureEnvironment.Endpoint.ResourceManager),
-                this.DefaultContext.Subscription);
+                this.DefaultContext.Subscription, 
+                this.DefaultContext.Environment.GetEndpoint(AzureEnvironment.Endpoint.StorageEndpointSuffix));
 
             base.ExecuteCmdlet();
 

src/ResourceManager/Compute/Commands.Compute/Extension/AEM/SetAzureRmVMAEMExtension.cs

@@ -38,6 +38,7 @@ namespace Microsoft.Azure.Commands.Compute
     [OutputType(typeof(PSAzureOperationResponse))]
     public class SetAzureRmVMAEMExtension : VirtualMachineExtensionBaseCmdlet
     {
+        private string _StorageEndpoint;
         private AEMHelper _Helper = null;
 
         [Parameter(
@@ -91,8 +92,13 @@ public SetAzureRmVMAEMExtension()
 
         public override void ExecuteCmdlet()
         {
+            this._StorageEndpoint = this.DefaultContext.Environment.GetEndpoint(AzureEnvironment.Endpoint.StorageEndpointSuffix);
             this._Helper = new AEMHelper((err) => this.WriteError(err), (msg) => this.WriteVerbose(msg), (msg) => this.WriteWarning(msg),
-                this.CommandRuntime.Host.UI, AzureSession.Instance.ClientFactory.CreateArmClient<StorageManagementClient>(DefaultProfile.DefaultContext, AzureEnvironment.Endpoint.ResourceManager), this.DefaultContext.Subscription);
+                this.CommandRuntime.Host.UI, 
+                AzureSession.Instance.ClientFactory.CreateArmClient<StorageManagementClient>(
+                    DefaultProfile.DefaultContext, AzureEnvironment.Endpoint.ResourceManager), 
+                this.DefaultContext.Subscription,
+                this._StorageEndpoint);
 
             base.ExecuteCmdlet();
 
@@ -521,7 +527,7 @@ private void SetStorageAnalytics(string storageAccountName)
 
             var key = this._Helper.GetAzureStorageKeyFromCache(storageAccountName);
             var credentials = new StorageCredentials(storageAccountName, key);
-            var cloudStorageAccount = new CloudStorageAccount(credentials, true);
+            var cloudStorageAccount = new CloudStorageAccount(credentials,  this._StorageEndpoint, true);
 
             cloudStorageAccount.CreateCloudBlobClient().SetServicePropertiesAsync(props)
                                                        .ConfigureAwait(false).GetAwaiter().GetResult();

src/ResourceManager/Compute/Commands.Compute/Extension/AEM/TestAzureRmVMAEMExtension.cs

@@ -89,7 +89,8 @@ public override void ExecuteCmdlet()
             this._Helper = new AEMHelper((err) => this.WriteError(err), (msg) => this.WriteVerbose(msg), (msg) => this.WriteWarning(msg),
                 this.CommandRuntime.Host.UI,
                 AzureSession.Instance.ClientFactory.CreateArmClient<StorageManagementClient>(DefaultProfile.DefaultContext, AzureEnvironment.Endpoint.ResourceManager),
-                this.DefaultContext.Subscription);
+                this.DefaultContext.Subscription, 
+                this.DefaultContext.Environment.GetEndpoint(AzureEnvironment.Endpoint.StorageEndpointSuffix));
 
             this._Helper.WriteVerbose("Starting TestAzureRmVMAEMExtension");