[PSCmdAssistant] test changes

src/Compute/Compute.Test/ScenarioTests/VirtualMachineTests.cs

@@ -703,5 +703,12 @@ public void TestVirtualMachineGalleryApplicationFlags()
         {
             TestRunner.RunTestScript("Test-VirtualMachineGalleryApplicationFlags");
         }
-    }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void testgensetazdisksecurityprofile()
+        {
+            TestRunner.RunTestScript("TestGen-setazdisksecurityprofile");
+        }
+        }
 }

src/Compute/Compute.Test/ScenarioTests/VirtualMachineTests.ps1

@@ -8115,4 +8115,53 @@ function Test-VirtualMachineGalleryApplicationFlags
     finally {
         Clean-ResourceGroup $resourceGroupName
     }
-}
+}
+function TestGen-setazdisksecurityprofile
+{
+    # Setup
+    $rgname = Get-ComputeTestResourceName;
+    $loc = Get-Location;
+
+    try
+    {
+        New-AzResourceGroup -Name $rgname -Location $loc -Force;
+
+        # Disk Security Profile Test with new Shield parameter
+        $diskName = "diskWithShield";
+        $diskconfig = New-AzDiskConfig -AccountType Premium_LRS -OsType Linux -CreateOption "FromImage" -Location $loc;
+        $diskconfig = Set-AzDiskImageReference -Disk $diskconfig -Id "/Subscriptions/e37510d7-33b6-4676-886f-ee75bcc01871/Providers/Microsoft.Compute/Locations/northeurope/Publishers/MicrosoftWindowsServer/ArtifactTypes/VMImage/Offers/windows-cvm/Skus/2019-datacenter-cvm/Versions/latest";
+
+        # Test with ShieldOn
+        $diskconfig = Set-AzDiskSecurityProfile -Disk $diskconfig -Shield "ShieldOn";
+        New-AzDisk -ResourceGroupName $rgname -DiskName $diskName -Disk $diskconfig;
+        $disk = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskName;
+
+        Assert-AreEqual $disk.SecurityProfile.Shield "ShieldOn";
+
+        # Test with ShieldGone
+        $diskName = "diskWithShieldGone";
+        $diskconfig = New-AzDiskConfig -AccountType Premium_LRS -OsType Linux -CreateOption "FromImage" -Location $loc;
+        $diskconfig = Set-AzDiskImageReference -Disk $diskconfig -Id "/Subscriptions/e37510d7-33b6-4676-886f-ee75bcc01871/Providers/Microsoft.Compute/Locations/northeurope/Publishers/MicrosoftWindowsServer/ArtifactTypes/VMImage/Offers/windows-cvm/Skus/2019-datacenter-cvm/Versions/latest";
+        $diskconfig = Set-AzDiskSecurityProfile -Disk $diskconfig -Shield "ShieldGone";
+        New-AzDisk -ResourceGroupName $rgname -DiskName $diskName -Disk $diskconfig;
+        $disk2 = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskName;
+
+        Assert-AreEqual $disk2.SecurityProfile.Shield "ShieldGone";
+
+        # Test with ShieldDown
+        $diskName = "diskWithShieldDown";
+        $diskconfig = New-AzDiskConfig -AccountType Premium_LRS -OsType Linux -CreateOption "FromImage" -Location $loc;
+        $diskconfig = Set-AzDiskImageReference -Disk $diskconfig -Id "/Subscriptions/e37510d7-33b6-4676-886f-ee75bcc01871/Providers/Microsoft.Compute/Locations/northeurope/Publishers/MicrosoftWindowsServer/ArtifactTypes/VMImage/Offers/windows-cvm/Skus/2019-datacenter-cvm/Versions/latest";
+        $diskconfig = Set-AzDiskSecurityProfile -Disk $diskconfig -Shield "ShieldDown";
+        New-AzDisk -ResourceGroupName $rgname -DiskName $diskName -Disk $diskconfig;
+        $disk3 = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskName;
+
+        Assert-AreEqual $disk3.SecurityProfile.Shield "ShieldDown";
+
+    }
+    finally
+    {
+        # Cleanup
+        Remove-AzResourceGroup -Name $rgname -Force -ErrorAction SilentlyContinue;
+    }
+}

src/Compute/Compute/ChangeLog.md

@@ -20,6 +20,7 @@
 
 -->
 ## Upcoming Release
+* Added parameter `Shield` to `Set-AzDiskSecurityProfile` cmdlet. The `Shield` parameter is a string with allowed values: `ShieldOn`, `ShieldGone`, `ShieldDown`.
 * Improved user experience and consistency. This may introduce breaking changes. Please refer to [here](https://go.microsoft.com/fwlink/?linkid=2340249).
 * Updated Azure.Core from 1.45.0 to 1.47.3
 * Added `-EnableAutomaticUpgrade` and `-TreatFailureAsDeploymentFailure` parameters (Bool) to `New-AzVmGalleryApplication` and `New-AzVmssGalleryApplication` cmdlets.

src/Compute/Compute/Generated/Disk/Config/SetAzDiskSecurityProfile.cs

@@ -1,4 +1,4 @@
-// ----------------------------------------------------------------------------------
+// ----------------------------------------------------------------------------------
 //
 // Copyright Microsoft Corporation
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -23,6 +23,7 @@
 using Microsoft.Azure.Commands.Compute.Automation.Models;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
 using Microsoft.WindowsAzure.Commands.Utilities.Common;
+using System.Reflection;
 
 namespace Microsoft.Azure.Commands.Compute
 {
@@ -54,6 +55,14 @@ public class SetAzDiskSecurityProfile : Microsoft.Azure.Commands.ResourceManager
            HelpMessage = "ResourceId of the disk encryption set to use for enabling encryption at rest.")]
         public string SecureVMDiskEncryptionSet { get; set; }
 
+        [Parameter(
+            Mandatory = false,
+            ValueFromPipelineByPropertyName = true,
+            HelpMessage = "Specifies the Shield setting for the disk security profile. Possible values include: ShieldOn, ShieldGone, ShieldDown")]
+        [PSArgumentCompleter("ShieldOn", "ShieldGone", "ShieldDown")]
+        [ValidateSet("ShieldOn", "ShieldGone", "ShieldDown", IgnoreCase = true)]
+        public string Shield { get; set; }
+
         protected override void ProcessRecord()
         {
             if (ShouldProcess("DiskSecurityProfile", "Set"))
@@ -94,6 +103,44 @@ private void Run()
                 this.Disk.SecurityProfile.SecureVMDiskEncryptionSetId = this.SecureVMDiskEncryptionSet;
             }
 
+            if (this.IsParameterBound(c => c.Shield))
+            {
+                if (this.Disk.SecurityProfile == null)
+                {
+                    this.Disk.SecurityProfile = new DiskSecurityProfile();
+                }
+
+                bool shieldSet = false;
+                Type dspType = this.Disk.SecurityProfile.GetType();
+                PropertyInfo shieldProperty = dspType.GetProperty("Shield", BindingFlags.Public | BindingFlags.Instance | BindingFlags.IgnoreCase);
+                if (shieldProperty != null && shieldProperty.CanWrite)
+                {
+                    shieldProperty.SetValue(this.Disk.SecurityProfile, this.Shield);
+                    shieldSet = true;
+                }
+                else
+                {
+                    // Attempt to use AdditionalProperties dictionary if available
+                    PropertyInfo additionalPropsProperty = dspType.GetProperty("AdditionalProperties", BindingFlags.Public | BindingFlags.Instance);
+                    if (additionalPropsProperty != null)
+                    {
+                        var dict = additionalPropsProperty.GetValue(this.Disk.SecurityProfile) as IDictionary<string, object>;
+                        if (dict == null)
+                        {
+                            dict = new Dictionary<string, object>(StringComparer.OrdinalIgnoreCase);
+                            additionalPropsProperty.SetValue(this.Disk.SecurityProfile, dict);
+                        }
+                        dict["shield"] = this.Shield;
+                        shieldSet = true;
+                    }
+                }
+
+                if (!shieldSet)
+                {
+                    WriteWarning("Shield parameter is specified but could not be applied because the current DiskSecurityProfile model does not expose a suitable property.");
+                }
+            }
+
             WriteObject(this.Disk);
         }
     }

src/Compute/Compute/Generated/Models/PSDisk.cs

@@ -28,6 +28,13 @@ namespace Microsoft.Azure.Commands.Compute.Automation.Models
 {
     public partial class PSDisk
     {
+        private static readonly HashSet<string> AllowedShieldValues = new HashSet<string>(StringComparer.OrdinalIgnoreCase)
+        {
+            "ShieldOn",
+            "ShieldGone",
+            "ShieldDown"
+        };
+
         // Gets or sets the property of 'ResourceGroupName'
         public string ResourceGroupName
         {
@@ -81,5 +88,21 @@ public string ResourceGroupName
         public bool? OptimizedForFrequentAttach { get; set; }
         public string SecurityDataUri { get; set; }
         public string SecurityMetadataUri { get; set; }
+
+        private string _shield;
+
+        // Gets or sets the Shield status for the disk. Allowed values are: ShieldOn, ShieldGone, ShieldDown.
+        public string Shield
+        {
+            get => _shield;
+            set
+            {
+                if (!string.IsNullOrEmpty(value) && !AllowedShieldValues.Contains(value))
+                {
+                    throw new ArgumentException("Invalid Shield value. Allowed values are: ShieldOn, ShieldGone, ShieldDown");
+                }
+                _shield = value;
+            }
+        }
     }
-}
+}