Add support for querying encryption status from the AzureDiskEncryptionForLinux extension

src/ResourceManager/Compute/Commands.Compute/Extension/AzureDiskEncryption/AzureDiskEncryptionExtensionConstants.cs

@@ -23,6 +23,9 @@ public static class AzureDiskEncryptionExtensionConstants
         public const string aadClientSecretParameterSet = "AAD Client Secret Parameters";
         public const string enableEncryptionOperation = "EnableEncryption";
         public const string disableEncryptionOperation = "DisableEncryption";
+        public const string queryEncryptionStatusOperation = "QueryEncryptionStatus";
+        public const string encryptionResultOsKey = "os";
+        public const string encryptionResultDataKey = "data";
         public const string aadClientIDKey = "AADClientID";
         public const string aadClientSecretKey = "AADClientSecret";
         public const string aadClientCertThumbprintKey = "AADClientCertThumbprint";

src/ResourceManager/Compute/Commands.Compute/Extension/AzureDiskEncryption/AzureDiskEncryptionExtensionContext.cs

@@ -94,6 +94,7 @@ public AzureDiskEncryptionExtensionContext(PSVirtualMachineExtension psExt)
             ProtectedSettings = psExt.ProtectedSettings;
             ProvisioningState = psExt.ProvisioningState;
             Statuses = psExt.Statuses;
+            SubStatuses = psExt.SubStatuses;
 
             InitializeAzureDiskEncryptionMembers(psExt);
         }

src/ResourceManager/Compute/Commands.Compute/Extension/AzureDiskEncryption/DisableAzureDiskEncryption.cs

@@ -231,6 +231,20 @@ public override void ExecuteCmdlet()
 
                 currentOSType = virtualMachineResponse.StorageProfile.OsDisk.OsType;
 
+                if (OperatingSystemTypes.Linux.Equals(currentOSType) &&
+                    !AzureDiskEncryptionExtensionContext.VolumeTypeData.Equals(VolumeType, StringComparison.InvariantCultureIgnoreCase))
+                {
+                    ThrowTerminatingError(
+                        new ErrorRecord(
+                            new ArgumentException(
+                                string.Format(
+                                    CultureInfo.CurrentUICulture,
+                                    "Disabling encryption is only allowed on Data volumes for Linux VMs.")),
+                            "InvalidType",
+                            ErrorCategory.NotImplemented,
+                            null));
+                }
+
                 if (this.ShouldProcess(VMName, Properties.Resources.DisableDiskEncryptionAction)
                     && (this.Force.IsPresent ||
                     this.ShouldContinue(Properties.Resources.DisableAzureDiskEncryptionConfirmation, Properties.Resources.DisableAzureDiskEncryptionCaption)))

src/ResourceManager/Compute/Commands.Compute/Extension/AzureDiskEncryption/GetAzureDiskEncryptionStatus.cs

@@ -16,7 +16,11 @@
 using Microsoft.Azure.Commands.Compute.Models;
 using Microsoft.Azure.Management.Compute;
 using Microsoft.Azure.Management.Compute.Models;
+using Microsoft.Rest.Azure;
+using Newtonsoft.Json;
 using System;
+using System.Collections;
+using System.Collections.Generic;
 using System.Globalization;
 using System.Management.Automation;
 
@@ -45,6 +49,140 @@ public class GetAzureDiskEncryptionStatusCommand : VirtualMachineExtensionBaseCm
         [ValidateNotNullOrEmpty]
         public string VMName { get; set; }
 
+        [Alias("ExtensionName")]
+        [Parameter(
+            Mandatory = false,
+            Position = 2,
+            ValueFromPipelineByPropertyName = true,
+            HelpMessage = "The extension name. If this parameter is not specified, default values used are AzureDiskEncryption for Windows VMs and AzureDiskEncryptionForLinux for Linux VMs")]
+        [ValidateNotNullOrEmpty]
+        public string Name { get; set; }
+
+        private VirtualMachineExtension GetVmExtensionParameters(VirtualMachine vmParameters, OSType currentOSType)
+        {
+            Hashtable publicSettings = new Hashtable();
+            Hashtable protectedSettings = new Hashtable();
+
+            publicSettings.Add(AzureDiskEncryptionExtensionConstants.encryptionOperationKey, AzureDiskEncryptionExtensionConstants.queryEncryptionStatusOperation);
+            publicSettings.Add(AzureDiskEncryptionExtensionConstants.sequenceVersionKey, Guid.NewGuid().ToString());
+
+            if (vmParameters == null)
+            {
+                ThrowTerminatingError(new ErrorRecord(new ApplicationException(string.Format(CultureInfo.CurrentUICulture, "Get-AzureDiskEncryptionExtension can enable encryption only on a VM that was already created ")),
+                                                      "InvalidResult",
+                                                      ErrorCategory.InvalidResult,
+                                                      null));
+            }
+
+            VirtualMachineExtension vmExtensionParameters = null;
+
+            if (OSType.Windows.Equals(currentOSType))
+            {
+                this.Name = this.Name ?? AzureDiskEncryptionExtensionContext.ExtensionDefaultName;
+                vmExtensionParameters = new VirtualMachineExtension
+                {
+                    Location = vmParameters.Location,
+                    Publisher = AzureDiskEncryptionExtensionContext.ExtensionDefaultPublisher,
+                    VirtualMachineExtensionType = AzureDiskEncryptionExtensionContext.ExtensionDefaultName,
+                    TypeHandlerVersion = AzureDiskEncryptionExtensionContext.ExtensionDefaultVersion,
+                    Settings = publicSettings,
+                    ProtectedSettings = protectedSettings
+                };
+            }
+            else if (OSType.Linux.Equals(currentOSType))
+            {
+                this.Name = this.Name ?? AzureDiskEncryptionExtensionContext.LinuxExtensionDefaultName;
+                vmExtensionParameters = new VirtualMachineExtension
+                {
+                    Location = vmParameters.Location,
+                    Publisher = AzureDiskEncryptionExtensionContext.LinuxExtensionDefaultPublisher,
+                    VirtualMachineExtensionType = AzureDiskEncryptionExtensionContext.LinuxExtensionDefaultName,
+                    TypeHandlerVersion = AzureDiskEncryptionExtensionContext.LinuxExtensionDefaultVersion,
+                    Settings = publicSettings,
+                    ProtectedSettings = protectedSettings
+                };
+            }
+
+            return vmExtensionParameters;
+        }
+
+        private string GetExtensionStatusMessage(OSType currentOSType, bool returnSubstatusMessage=false)
+        {
+            AzureOperationResponse<VirtualMachineExtension> extensionResult = this.VirtualMachineExtensionClient.GetWithInstanceView(this.ResourceGroupName, this.VMName, this.Name);
+            if (extensionResult == null)
+            {
+                ThrowTerminatingError(new ErrorRecord(new ApplicationFailedException(string.Format(CultureInfo.CurrentUICulture, "Failed to retrieve extension status")),
+                                                      "InvalidResult",
+                                                      ErrorCategory.InvalidResult,
+                                                      null));
+            }
+
+            PSVirtualMachineExtension returnedExtension = extensionResult.ToPSVirtualMachineExtension(
+                this.ResourceGroupName, this.VMName);
+
+            if ((returnedExtension == null) ||
+                (string.IsNullOrWhiteSpace(returnedExtension.Publisher)) ||
+                (string.IsNullOrWhiteSpace(returnedExtension.ExtensionType)))
+            {
+                ThrowTerminatingError(new ErrorRecord(new ApplicationFailedException(string.Format(CultureInfo.CurrentUICulture, "Missing extension publisher and type info")),
+                                                      "InvalidResult",
+                                                      ErrorCategory.InvalidResult,
+                                                      null));
+            }
+            bool publisherMatch = false;
+            if (OSType.Linux.Equals(currentOSType))
+            {
+                if (returnedExtension.Publisher.Equals(AzureDiskEncryptionExtensionContext.LinuxExtensionDefaultPublisher, StringComparison.InvariantCultureIgnoreCase) &&
+                    returnedExtension.ExtensionType.Equals(AzureDiskEncryptionExtensionContext.LinuxExtensionDefaultName, StringComparison.InvariantCultureIgnoreCase))
+                {
+                    publisherMatch = true;
+                }
+            }
+            else if (OSType.Windows.Equals(currentOSType))
+            {
+                if (returnedExtension.Publisher.Equals(AzureDiskEncryptionExtensionContext.ExtensionDefaultPublisher, StringComparison.InvariantCultureIgnoreCase) &&
+                    returnedExtension.ExtensionType.Equals(AzureDiskEncryptionExtensionContext.ExtensionDefaultName, StringComparison.InvariantCultureIgnoreCase))
+                {
+                    publisherMatch = true;
+                }
+            }
+            if (publisherMatch)
+            {
+                AzureDiskEncryptionExtensionContext context = new AzureDiskEncryptionExtensionContext(returnedExtension);
+                if ((context == null) ||
+                    (context.Statuses == null) ||
+                    (context.Statuses.Count < 1) ||
+                    (string.IsNullOrWhiteSpace(context.Statuses[0].Message)))
+                {
+                    throw new KeyNotFoundException(string.Format(CultureInfo.CurrentUICulture, "Invalid extension status"));
+                }
+
+                if (returnSubstatusMessage)
+                {
+                    if((context == null) ||
+                       (context.SubStatuses == null) ||
+                       (context.SubStatuses.Count < 1))
+                    {
+                        throw new KeyNotFoundException(string.Format(CultureInfo.CurrentUICulture, "Invalid extension substatus"));
+                    }
+                    else
+                    {
+                        return context.SubStatuses[0].Message;
+                    }
+                }
+
+                return context.Statuses[0].Message;
+            }
+            else
+            {
+                ThrowTerminatingError(new ErrorRecord(new ApplicationFailedException(string.Format(CultureInfo.CurrentUICulture, "Extension publisher and type mismatched")),
+                                                      "InvalidResult",
+                                                      ErrorCategory.InvalidResult,
+                                                      null));
+            }
+            return null;
+        }
+
         private OSType GetOSType(VirtualMachine vmParameters)
         {
             if (vmParameters == null || vmParameters.StorageProfile == null || vmParameters.StorageProfile.OsDisk == null)
@@ -209,17 +347,72 @@ public override void ExecuteCmdlet()
                 EncryptionStatus osVolumeEncrypted = IsOsVolumeEncrypted(vmParameters);
                 DiskEncryptionSettings osVolumeEncryptionSettings = GetOsVolumeEncryptionSettings(vmParameters);
                 EncryptionStatus dataVolumesEncrypted = AreDataVolumesEncrypted(vmParameters);
+                AzureDiskEncryptionStatusContext encryptionStatus = null;
+                string progressMessage = null;
 
                 OSType osType = GetOSType(vmParameters);
                 switch (osType)
                 {
                     case OSType.Windows:
-                    case OSType.Linux:
-                        AzureDiskEncryptionStatusContext encryptionStatus = new AzureDiskEncryptionStatusContext
+                        try
+                        {
+                            progressMessage = GetExtensionStatusMessage(osType);
+                        }
+                        catch(KeyNotFoundException)
+                        {
+                            progressMessage = string.Format(CultureInfo.CurrentUICulture, "Extension status not available on the VM");
+                        }
+
+                        encryptionStatus = new AzureDiskEncryptionStatusContext
                         {
                             OsVolumeEncrypted = osVolumeEncrypted,
                             DataVolumesEncrypted = dataVolumesEncrypted,
-                            OsVolumeEncryptionSettings = osVolumeEncryptionSettings
+                            OsVolumeEncryptionSettings = osVolumeEncryptionSettings,
+                            ProgressMessage = progressMessage
+                        };
+                        WriteObject(encryptionStatus);
+                        break;
+                    case OSType.Linux:
+                        VirtualMachine virtualMachineResponse = this.ComputeClient.ComputeManagementClient.VirtualMachines.GetWithInstanceView(
+                            this.ResourceGroupName, VMName).Body;
+                        VirtualMachineExtension parameters = GetVmExtensionParameters(virtualMachineResponse, osType);
+
+                        this.VirtualMachineExtensionClient.CreateOrUpdateWithHttpMessagesAsync(
+                            this.ResourceGroupName,
+                            this.VMName,
+                            this.Name,
+                            parameters).GetAwaiter().GetResult();
+
+                        Dictionary<string, string> encryptionStatusParsed = null;
+                        try
+                        {
+                            string encryptionStatusJson = GetExtensionStatusMessage(osType, returnSubstatusMessage: true);
+                            encryptionStatusParsed = JsonConvert.DeserializeObject<Dictionary<string, string>>(encryptionStatusJson);
+                        }
+                        catch(KeyNotFoundException)
+                        {
+                            encryptionStatusParsed = new Dictionary<string, string>()
+                            {
+                                { AzureDiskEncryptionExtensionConstants.encryptionResultOsKey, EncryptionStatus.Unknown.ToString() },
+                                { AzureDiskEncryptionExtensionConstants.encryptionResultDataKey, EncryptionStatus.Unknown.ToString() }
+                            };
+                        }
+
+                        try
+                        {
+                            progressMessage = GetExtensionStatusMessage(osType);
+                        }
+                        catch(KeyNotFoundException)
+                        {
+                            progressMessage = string.Format(CultureInfo.CurrentUICulture, "Extension status not available on the VM");
+                        }
+
+                        encryptionStatus = new AzureDiskEncryptionStatusContext
+                        {
+                            OsVolumeEncrypted = (EncryptionStatus)Enum.Parse(typeof(EncryptionStatus), encryptionStatusParsed[AzureDiskEncryptionExtensionConstants.encryptionResultOsKey]),
+                            DataVolumesEncrypted = (EncryptionStatus)Enum.Parse(typeof(EncryptionStatus), encryptionStatusParsed[AzureDiskEncryptionExtensionConstants.encryptionResultDataKey]),
+                            OsVolumeEncryptionSettings = osVolumeEncryptionSettings,
+                            ProgressMessage = progressMessage
                         };
                         WriteObject(encryptionStatus);
                         break;

src/ResourceManager/Compute/Commands.Compute/Extension/AzureDiskEncryption/SetAzureDiskEncryptionExtension.cs

@@ -430,20 +430,6 @@ public override void ExecuteCmdlet()
 
                     currentOSType = virtualMachineResponse.StorageProfile.OsDisk.OsType;
 
-                    if (OperatingSystemTypes.Linux.Equals(currentOSType) &&
-                        !AzureDiskEncryptionExtensionContext.VolumeTypeData.Equals(VolumeType, StringComparison.InvariantCultureIgnoreCase))
-                    {
-                        ThrowTerminatingError(
-                            new ErrorRecord(
-                                new ArgumentException(
-                                    string.Format(
-                                        CultureInfo.CurrentUICulture,
-                                        "Enabling encryption is only allowed on Data volumes for Linux VMs.")),
-                                "InvalidType",
-                                ErrorCategory.NotImplemented,
-                                null));
-                    }
-
                     if (OperatingSystemTypes.Linux.Equals(currentOSType))
                     {
                         CreateVMBackupForLinx();

src/ResourceManager/Compute/Commands.Compute/Microsoft.Azure.Commands.Compute.format.ps1xml

@@ -556,6 +556,10 @@
                 <Label>OsVolumeEncryptionSettings</Label>
                 <PropertyName>OsVolumeEncryptionSettings</PropertyName>
               </ListItem>
+              <ListItem>
+                <Label>ProgressMessage</Label>
+                <PropertyName>ProgressMessage</PropertyName>
+              </ListItem>
             </ListItems>
           </ListEntry>
         </ListEntries>

src/ResourceManager/Compute/Commands.Compute/Models/AzureDiskEncryptionStatusContext.cs

@@ -21,6 +21,9 @@ enum EncryptionStatus
     {
         Encrypted,
         NotEncrypted,
+        NotMounted,
+        EncryptionInProgress,
+        VMRestartPending,
         Unknown
     }
 
@@ -29,5 +32,6 @@ class AzureDiskEncryptionStatusContext
         public EncryptionStatus OsVolumeEncrypted { get; set; }
         public DiskEncryptionSettings OsVolumeEncryptionSettings { get; set; }
         public EncryptionStatus DataVolumesEncrypted { get; set; }
+        public string ProgressMessage { get; set; }
     }
 }