Implement Add/Remove Cmdlets for VMSS diagnostics extension, add rela…

src/ResourceManager/Compute/Commands.Compute.Test/Commands.Compute.Test.csproj

@@ -232,6 +232,12 @@
     <Content Include="ConfigFiles\DiagnosticsExtensionConfig.json">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </Content>
+    <Content Include="ConfigFiles\DiagnosticsExtensionPrivateConfig.json">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </Content>
+    <Content Include="ConfigFiles\DiagnosticsExtensionPublicConfig.json">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </Content>
     <None Include="ScenarioTests\AddVhdTests.ps1">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </None>
@@ -332,6 +338,9 @@
     <None Include="SessionRecords\Microsoft.Azure.Commands.Compute.Test.ScenarioTests.DiagnosticsExtensionTests\TestDiagnosticsExtensionSupportJsonConfig.json">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </None>
+    <None Include="SessionRecords\Microsoft.Azure.Commands.Compute.Test.ScenarioTests.DiagnosticsExtensionTests\TestVmssDiagnosticsExtension.json">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </None>
     <None Include="SessionRecords\Microsoft.Azure.Commands.Compute.Test.ScenarioTests.DscExtensionTests\TestGetAzureRmVMDscExtension.json">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </None>

src/ResourceManager/Compute/Commands.Compute.Test/ConfigFiles/DiagnosticsExtensionPrivateConfig.json

@@ -0,0 +1,3 @@
+{
+  "storageAccountName": "definedinconfigstorage"
+}

src/ResourceManager/Compute/Commands.Compute.Test/ConfigFiles/DiagnosticsExtensionPublicConfig.json

@@ -0,0 +1,101 @@
+{
+  "WadCfg": {
+    "DiagnosticMonitorConfiguration": {
+      "DiagnosticInfrastructureLogs": {
+        "scheduledTransferLogLevelFilter": "Error"
+      },
+      "Directories": {
+        "IISLogs": {
+          "containerName": "wad-iis-logfiles"
+        },
+        "FailedRequestLogs": {
+          "containerName": "wad-failedrequestlogs"
+        },
+        "scheduledTransferPeriod": "PT1M"
+      },
+      "PerformanceCounters": {
+        "PerformanceCounterConfiguration": [
+          {
+            "annotation": [ ],
+            "counterSpecifier": "\\Memory\\Available MBytes",
+            "sampleRate": "PT3M"
+          },
+          {
+            "annotation": [ ],
+            "counterSpecifier": "\\Web Service(_Total)\\ISAPI Extension Requests/sec",
+            "sampleRate": "PT3M"
+          },
+          {
+            "annotation": [ ],
+            "counterSpecifier": "\\Web Service(_Total)\\Bytes Total/Sec",
+            "sampleRate": "PT3M"
+          },
+          {
+            "annotation": [ ],
+            "counterSpecifier": "\\ASP.NET Applications(__Total__)\\Requests/Sec",
+            "sampleRate": "PT3M"
+          },
+          {
+            "annotation": [ ],
+            "counterSpecifier": "\\ASP.NET Applications(__Total__)\\Errors Total/Sec",
+            "sampleRate": "PT3M"
+          },
+          {
+            "annotation": [ ],
+            "counterSpecifier": "\\ASP.NET\\Requests Queued",
+            "sampleRate": "PT3M"
+          },
+          {
+            "annotation": [ ],
+            "counterSpecifier": "\\ASP.NET\\Requests Rejected",
+            "sampleRate": "PT3M"
+          },
+          {
+            "annotation": [ ],
+            "counterSpecifier": "\\Processor(_Total)\\% Processor Time",
+            "sampleRate": "PT3M"
+          }
+        ],
+        "scheduledTransferPeriod": "PT1M"
+      },
+      "WindowsEventLog": {
+        "DataSource": [
+          {
+            "name": "Application!*[System[(Level=1 or Level=2 or Level=3)]]"
+          },
+          {
+            "name": "Windows Azure!*[System[(Level=1 or Level=2 or Level=3 or Level=4)]]"
+          }
+        ],
+        "scheduledTransferPeriod": "PT1M"
+      },
+      "CrashDumps": {
+        "CrashDumpConfiguration": [
+          {
+            "processName": "WaIISHost.exe"
+          },
+          {
+            "processName": "WaWorkerHost.exe"
+          },
+          {
+            "processName": "w3wp.exe"
+          }
+        ]
+      },
+      "Logs": {
+        "scheduledTransferLogLevelFilter": "Error",
+        "scheduledTransferPeriod": "PT1M"
+      },
+      "Metrics": {
+        "resourceId": "dummy",
+        "MetricAggregation": [
+          {
+            "scheduledTransferPeriod": "PT1M"
+          }
+        ]
+      },
+      "overallQuotaInMB": 4096
+    }
+  },
+  "StorageAccount": "definedinconfigstorage"
+}

src/ResourceManager/Compute/Commands.Compute.Test/ScenarioTests/ComputeTestCommon.ps1

@@ -576,3 +576,8 @@ function Get-SubscriptionIdFromResourceGroup
       $last = $rgid.IndexOf('/', $first + 1);
       return $rgid.Substring($first + 1, $last - $first - 1);
 }
+
+function Get-ComputeVmssLocation
+{
+      Get-ResourceProviderLocation "Microsoft.Compute/virtualMachineScaleSets"
+}

src/ResourceManager/Compute/Commands.Compute.Test/ScenarioTests/DiagnosticsExtensionTests.cs

@@ -51,5 +51,12 @@ public void TestDiagnosticsExtensionSupportJsonConfig()
         {
             ComputeTestController.NewInstance.RunPsTest("Test-DiagnosticsExtensionSupportJsonConfig");
         }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void TestVmssDiagnosticsExtension()
+        {
+            ComputeTestController.NewInstance.RunPsTest("Test-VmssDiagnosticsExtension");
+        }
     }
 }

src/ResourceManager/Compute/Commands.Compute.Test/ScenarioTests/DiagnosticsExtensionTests.ps1

@@ -183,4 +183,114 @@ function Test-DiagnosticsExtensionSupportJsonConfig
         # Cleanup
         Clean-ResourceGroup $rgname
     }
-}
+}
+
+<#
+.SYNOPSIS
+Test the basic usage of the Add/Remove VMSS diagnostics extension command
+#>
+function Test-VmssDiagnosticsExtension
+{
+    $rgname = Get-ComputeTestResourceName
+    $loc = Get-ComputeVmssLocation
+
+    try
+    {
+        # Common
+        New-AzureRMResourceGroup -Name $rgname -Location $loc -Force;
+
+        # Create VMSS
+
+        # SRP
+        $stoname = 'sto' + $rgname;
+        $stotype = 'Standard_GRS';
+        New-AzureRMStorageAccount -ResourceGroupName $rgname -Name $stoname -Location $loc -Type $stotype;
+        $stoaccount = Get-AzureRMStorageAccount -ResourceGroupName $rgname -Name $stoname;
+
+        # NRP
+        $subnet = New-AzureRMVirtualNetworkSubnetConfig -Name ('subnet' + $rgname) -AddressPrefix "10.0.0.0/24";
+        $vnet = New-AzureRMVirtualNetwork -Force -Name ('vnet' + $rgname) -ResourceGroupName $rgname -Location $loc -AddressPrefix "10.0.0.0/16" -Subnet $subnet;
+        $vnet = Get-AzureRMVirtualNetwork -Name ('vnet' + $rgname) -ResourceGroupName $rgname;
+        $subnetId = $vnet.Subnets[0].Id;
+
+        # New VMSS Parameters
+        $vmssName = 'vmss' + $rgname;
+        $vmssType = 'Microsoft.Compute/virtualMachineScaleSets';
+
+        $adminUsername = 'Foo12';
+        $adminPassword = "BaR@123" + $rgname;
+
+        $imgRef = Get-DefaultCRPImage -loc $loc;
+        $vhdContainer = "https://" + $stoname + ".blob.core.windows.net/" + $vmssName;
+
+        $extname = 'diagextest';
+        $diagExtPublisher = 'Microsoft.Azure.Diagnostics';
+        $diagExtType = 'IaaSDiagnostics';
+
+        # This storage name will be used in command line directly when set diagnostics extension
+        $storagename = 'definedinconfigstorage';
+        $storagetype = 'Standard_GRS';
+        New-AzureRmStorageAccount -ResourceGroupName $rgname -Name $storagename -Location $loc -Type $storagetype;
+
+        $ipCfg = New-AzureRmVmssIPConfig -Name 'test' -SubnetId $subnetId;
+        $vmss = New-AzureRmVmssConfig -Location $loc -SkuCapacity 2 -SkuName 'Standard_A0' -UpgradePolicyMode 'automatic' -NetworkInterfaceConfiguration $netCfg `
+            | Add-AzureRmVmssNetworkInterfaceConfiguration -Name 'test' -Primary $true -IPConfiguration $ipCfg `
+            | Set-AzureRmVmssOSProfile -ComputerNamePrefix 'test' -AdminUsername $adminUsername -AdminPassword $adminPassword `
+            | Set-AzureRmVmssStorageProfile -Name 'test' -OsDiskCreateOption 'FromImage' -OsDiskCaching 'None' `
+            -ImageReferenceOffer $imgRef.Offer -ImageReferenceSku $imgRef.Skus -ImageReferenceVersion $imgRef.Version `
+            -ImageReferencePublisher $imgRef.PublisherName -VhdContainer $vhdContainer;
+
+        # Full parameter test
+        $version = '1.5';
+        $publicSettingFilePath = "$TestOutputRoot\ConfigFiles\DiagnosticsExtensionPublicConfig.json";
+        $privateSettingFilePath = "$TestOutputRoot\ConfigFiles\DiagnosticsExtensionPrivateConfig.json";
+        $vmss = Add-AzureRmVmssDiagnosticsExtension -VirtualMachineScaleSet $vmss -Name $extname -SettingFilePath $publicSettingFilePath `
+            -ProtectedSettingFilePath $privateSettingFilePath -TypeHandlerVersion $version -AutoUpgradeMinorVersion $false -Force;
+
+        $vmssDiagExtensions = $vmss.VirtualMachineProfile.ExtensionProfile.Extensions | Where-Object {$_.Publisher -eq $diagExtPublisher -and $_.Type -eq $diagExtType};
+        Assert-AreEqual 1 $vmssDiagExtensions.Count;
+        $vmssDiagExtension = $vmssDiagExtensions | Select-Object -first 1;
+
+        Assert-AreEqual $extname $vmssDiagExtension.Name;
+        Assert-AreEqual $version $vmssDiagExtension.TypeHandlerVersion;
+        Assert-AreEqual $false $vmssDiagExtension.AutoUpgradeMinorVersion;
+
+        # Storage account key should be filled up by the cmdlet.
+        $storageAccountKey = $vmssDiagExtension.ProtectedSettings['storageAccountKey'];
+        Assert-NotNull $storageAccountKey;
+        Assert-AreNotEqual '' $storageAccountKey;
+
+        # Remove without specifying extension name, diagnostic extension is expected to be removed.
+        $vmss = Remove-AzureRmVmssDiagnosticsExtension -VirtualMachineScaleSet $vmss;
+        $vmssDiagExtensions = $vmss.VirtualMachineProfile.ExtensionProfile.Extensions | Where-Object {$_.Publisher -eq $diagExtPublisher -and $_.Type -eq $diagExtType};
+
+        Assert-Null $vmssDiagExtensions;
+
+        $vmss = $vmss | Add-AzureRmVmssDiagnosticsExtension -Name $extname -SettingFilePath $publicSettingFilePath `
+            | New-AzureRmVmss -ResourceGroupName $rgname -Name $vmssName;
+
+        $vmss = Get-AzureRmVmss -ResourceGroupName $rgname -VMScaleSetName $vmssName;
+
+        $vmssDiagExtensions = $vmss.VirtualMachineProfile.ExtensionProfile.Extensions | Where-Object {$_.Publisher -eq $diagExtPublisher -and $_.Type -eq $diagExtType};
+        Assert-AreEqual 1 $vmssDiagExtensions.Count;
+
+        $vmssDiagExtension = $vmssDiagExtensions | Select-Object -first 1;
+        Assert-AreEqual $extname $vmssDiagExtension.Name;
+
+        $settings = $vmssDiagExtension.Settings;
+        Assert-AreEqual $storagename $settings.storageAccount.Value;
+
+        Remove-AzureRmVmssDiagnosticsExtension -VirtualMachineScaleSet $vmss -Name $extname;
+        Update-AzureRmVmss -ResourceGroupName $rgname -Name $vmssName -VirtualMachineScaleSet $vmss;
+
+        $vmss = Get-AzureRmVmss -ResourceGroupName $rgname -VMScaleSetName $vmssName;
+        $vmssDiagExtensions = $vmss.VirtualMachineProfile.ExtensionProfile.Extensions | Where-Object {$_.Publisher -eq $diagExtPublisher -and $_.Type -eq $diagExtType};
+
+        Assert-Null $vmssDiagExtensions;
+    }
+    finally
+    {
+        # Cleanup
+        Clean-ResourceGroup $rgname
+    }
+}

src/ResourceManager/Compute/Commands.Compute/Commands.Compute.csproj

@@ -247,9 +247,11 @@
     <Compile Include="Extension\CustomScript\RemoveAzureVMCustomScriptExtensionCommand.cs" />
     <Compile Include="Extension\CustomScript\SetAzureVMCustomScriptExtensionCommand.cs" />
     <Compile Include="Extension\CustomScript\VirtualMachineCustomScriptExtensionContext.cs" />
+    <Compile Include="Extension\Diagnostics\AddAzureRmVmssDiagnosticsExtension.cs" />
     <Compile Include="Extension\Diagnostics\DiagnosticsExtensionConstants.cs" />
     <Compile Include="Extension\Diagnostics\GetAzureRmVMDiagnosticsExtension.cs" />
     <Compile Include="Extension\Diagnostics\RemoveAzureRmVMDiagnosticsExtension.cs" />
+    <Compile Include="Extension\Diagnostics\RemoveAzureRmVmssDiagnosticsExtension.cs" />
     <Compile Include="Extension\Diagnostics\SetAzureRmVMDiagnosticsExtension.cs" />
     <Compile Include="Extension\Diagnostics\WADPrivateConfigSchema.cs" />
     <Compile Include="Extension\DSC\GetAzureVMDscExtensionStatusCommand.cs" />

src/ResourceManager/Compute/Commands.Compute/Common/ConstantStringTypes.cs

@@ -116,6 +116,8 @@ public static class ProfileNouns
         public const string VaultSecretGroup = "AzureRmVMSecret";
         public const string RemoteDesktopFile = "AzureRmRemoteDesktopFile";
 
+        public const string VirtualMachineScaleSetDiagnosticsExtension = "AzureRmVmssDiagnosticsExtension";
+
         //DSC
         public const string VirtualMachineDscExtension = "AzureRmVMDscExtension";
         public const string VirtualMachineDscConfiguration = "AzureRmVMDscConfiguration";

src/ResourceManager/Compute/Commands.Compute/Common/DiagnosticsHelper.cs

@@ -14,6 +14,7 @@
 
 using Microsoft.Azure.Commands.Common.Authentication.Models;
 using Microsoft.Azure.Commands.Management.Storage.Models;
+using Microsoft.Azure.Management.Compute.Models;
 using Microsoft.Azure.Management.Storage;
 using Microsoft.Azure.Management.Storage.Models;
 using Microsoft.WindowsAzure.Commands.Common.Storage;
@@ -195,8 +196,15 @@ private static Hashtable GetPublicConfigFromJsonFile(string configurationPath, s
             var wadCfgProperty = properties.FirstOrDefault(p => p.Equals(WadCfg, StringComparison.OrdinalIgnoreCase));
             var wadCfgBlobProperty = properties.FirstOrDefault(p => p.Equals(WadCfgBlob, StringComparison.OrdinalIgnoreCase));
             var xmlCfgProperty = properties.FirstOrDefault(p => p.Equals(EncodedXmlCfg, StringComparison.OrdinalIgnoreCase));
+            var storageAccountProperty = properties.FirstOrDefault(p => p.Equals(StorageAccount, StringComparison.OrdinalIgnoreCase));
 
             var hashTable = new Hashtable();
+
+            if (string.IsNullOrEmpty(storageAccountName) && storageAccountProperty != null)
+            {
+                storageAccountName = (string)publicConfig[storageAccountProperty];
+            }
+
             hashTable.Add(StorageAccount, storageAccountName);
 
             if (wadCfgProperty != null && publicConfig[wadCfgProperty] is JObject)
@@ -578,5 +586,79 @@ private static string GetEndpointFromStorageContext(AzureStorageContext context)
             var scheme = context.BlobEndPoint.StartsWith("https://", StringComparison.OrdinalIgnoreCase) ? "https://" : "http://";
             return scheme + context.EndPointSuffix;
         }
+
+        /// <summary>
+        /// Parse public and private configurations, and automatically resolve storage key for private configuration.
+        /// </summary>
+        /// <param name="publicConfigPath">Public configuration file path</param>
+        /// <param name="privateConfigPath">Private configuration file path, can be empty</param>
+        /// <param name="storageClient">Storage client</param>
+        /// <returns>A tuple with public configuration as first element and private configuration as second element</returns>
+        public static Tuple<Hashtable, Hashtable> GetConfigurationsFromFiles(string publicConfigPath, string privateConfigPath, string resourceId, Cmdlet cmdlet, IStorageManagementClient storageClient)
+        {
+            var publicConfig = GetPublicConfigFromJsonFile(publicConfigPath, null, resourceId, cmdlet);
+            var privateConfig = string.IsNullOrEmpty(privateConfigPath) ? new Hashtable() :
+                JsonConvert.DeserializeObject<Hashtable>(File.ReadAllText(privateConfigPath));
+
+            // Resolve storage account name
+            // Storage account name must be provided in public config
+            var storageAccountName = publicConfig[StorageAccount] as string;
+            if (string.IsNullOrEmpty(storageAccountName))
+            {
+                throw new ArgumentException(Properties.Resources.DiagnosticsExtensionNullStorageAccountName);
+            }
+
+            privateConfig[StorageAccountNameTag] = storageAccountName;
+
+            // Resolve storage account key
+            var storageAccountKey = InitializeStorageAccountKey(storageClient, storageAccountName, privateConfigPath);
+            if (string.IsNullOrEmpty(storageAccountKey))
+            {
+                storageAccountKey = privateConfig[StorageAccountKeyTag] as string;
+
+                if (string.IsNullOrEmpty(storageAccountKey))
+                {
+                    // Throw exception if no storage key provided in private config and cannot retrieve it from server
+                    throw new ArgumentException(Properties.Resources.DiagnosticsExtensionNullStorageAccountKey);
+                }
+            }
+            else
+            {
+                // If storage key can be retrieved, use that one.
+                privateConfig[StorageAccountKeyTag] = storageAccountKey;
+            }
+
+
+            // Resolve storage account endpoint
+            var storageAccountEndpoint = InitializeStorageAccountEndpoint(storageAccountName, storageAccountKey, storageClient);
+            if (string.IsNullOrEmpty(storageAccountEndpoint))
+            {
+                storageAccountEndpoint = privateConfig[StorageAccountEndPointTag] as string;
+
+                if (string.IsNullOrEmpty(storageAccountEndpoint))
+                {
+                    // Throw exception if no storage endpoint provided in private config and cannot retrieve it from server
+                    throw new ArgumentNullException(Properties.Resources.DiagnosticsExtensionNullStorageAccountEndpoint);
+                }
+            }
+            else
+            {
+                // If storage account endpoint can be retrieved, use that one.
+                privateConfig[StorageAccountEndPointTag] = storageAccountEndpoint;
+            }
+
+            return new Tuple<Hashtable, Hashtable>(publicConfig, privateConfig);
+        }
+
+        /// <summary>
+        /// Check if a VMSS extension is diagnostics extension.
+        /// </summary>
+        /// <param name="extension">VMSS extension</param>
+        /// <returns>Whether the VMSS extension is diagnostics extension</returns>
+        public static bool IsDiagnosticsExtension(VirtualMachineScaleSetExtension extension)
+        {
+            return extension.Publisher.Equals(DiagnosticsExtensionConstants.ExtensionPublisher, StringComparison.InvariantCultureIgnoreCase) &&
+                extension.Type.Equals(DiagnosticsExtensionConstants.ExtensionType, StringComparison.InvariantCultureIgnoreCase);
+        }
     }
 }