Use V2 Storage accounts for RM based Diagnostics Extenstion

Author: rahulpandit85

src/ResourceManager/Compute/Commands.Compute/Extension/Diagnostics/SetAzureVMDiagnosticsExtension.cs

@@ -21,7 +21,8 @@
 using Microsoft.Azure.Management.Compute.Models;
 using Microsoft.WindowsAzure.Commands.Common.Storage;
 using Microsoft.WindowsAzure.Commands.Utilities.Common;
-using Microsoft.WindowsAzure.Management.Storage;
+using Microsoft.Azure.Commands.Management.Storage;
+using Microsoft.Azure.Management.Storage;
 using Newtonsoft.Json;
 
 namespace Microsoft.Azure.Commands.Compute
@@ -37,7 +38,7 @@ public class SetAzureVMDiagnosticsExtensionCommand : VirtualMachineExtensionBase
         private const string VirtualMachineExtension = "Microsoft.Compute/virtualMachines/extensions";
         private const string IaaSDiagnosticsExtension = "IaaSDiagnostics";
         private const string ExtensionPublisher = "Microsoft.Azure.Diagnostics";
-        private StorageManagementClient storageClient;
+        private IStorageManagementClient storageClient;
 
         [Parameter(
             Mandatory = true,
@@ -154,14 +155,14 @@ public string PrivateConfiguration
             }
         }
 
-        public StorageManagementClient StorageClient
+        public IStorageManagementClient StorageClient
         {
             get
             {
                 if (this.storageClient == null)
                 {
                     this.storageClient = AzureSession.ClientFactory.CreateClient<StorageManagementClient>(
-                        DefaultProfile.Context, AzureEnvironment.Endpoint.ServiceManagement);
+                        DefaultProfile.Context, AzureEnvironment.Endpoint.ResourceManager);
                 }
 
                 return this.storageClient;
@@ -201,14 +202,12 @@ protected string GetStorageKey()
 
             if (!string.IsNullOrEmpty(StorageAccountName))
             {
-                var storageAccount = this.StorageClient.StorageAccounts.Get(StorageAccountName);
-                if (storageAccount != null)
+                var storageAccountKeys =
+                    this.StorageClient.StorageAccounts.ListKeys(this.ResourceGroupName,
+                        this.StorageContext.StorageAccountName);
+                if (storageAccountKeys != null)
                 {
-                    var keys = this.StorageClient.StorageAccounts.GetKeys(StorageAccountName);
-                    if (keys != null)
-                    {
-                        storageKey = !string.IsNullOrEmpty(keys.PrimaryKey) ? keys.PrimaryKey : keys.SecondaryKey;
-                    }
+                    storageKey = !string.IsNullOrEmpty(storageAccountKeys.StorageAccountKeys.Key1) ? storageAccountKeys.StorageAccountKeys.Key1 : storageAccountKeys.StorageAccountKeys.Key2;
                 }
             }