selective restore key for a managed hsm backup (Azure#13627)

Author: BethanyZhou

src/KeyVault/KeyVault.Test/PesterTests/ManagedHsmDatePlaneTests.Tests.ps1

@@ -4,11 +4,12 @@ $sut = (Split-Path -Leaf $MyInvocation.MyCommand.Path) -replace '\.Tests\.', '.'
 
 . $PSScriptRoot/ManagedHsmDatePlaneTests.ps1
 # ImportModules
-$hsmName = 'yeminghsm02'
+$hsmName = 'bezmhsm'
 $signInName = '[email protected]'
-$storageAccount = 'yemingsa01'
-$containerName = 'hsmbackup'
-$sasToken = ConvertTo-SecureString -AsPlainText -Force 'insert sas token'
+$storageAccount = 'bezstorageaccount'
+$containerName = 'backup'
+$keyName = 'test'
+# $sasToken = ConvertTo-SecureString -AsPlainText -Force 'insert sas token'
 $certs = "D:\sd1.cer", "D:\sd2.cer", "D:\sd3.cer" # for security domain
 $certsKeys = @{PublicKey = "D:\sd1.cer"; PrivateKey = "D:\sd1.key" }, @{PublicKey = "D:\sd2.cer"; PrivateKey = "D:\sd2.key" }, @{PublicKey = "D:\sd3.cer"; PrivateKey = "D:\sd3.key" }
 
@@ -170,15 +171,24 @@ Describe "BackupAndRestoreAzManagedHsmKey" {
 Describe "BackupAndRestoreAzManagedHsm" {
     $script:backupUri = ''
     $containerUri = "https://$storageAccount.blob.core.windows.net/$containerName"
-
     It "Backup then restore a managed HSM" {
         $script:backupUri = Backup-AzKeyVault -HsmName $hsmName -StorageContainerUri $containerUri -SasToken $sasToken
         $script:backupUri | Should -Not -Be $null
     }
 
+    It "Selective restore a managed HSM"{
+        $script:backupUri = [System.Uri]::new($script:backupUri)
+        $backupFolder = $script:backupUri.Segments[$script:backupUri.Segments.Length - 1]
+        $restoreResult = Restore-AzKeyVault -HsmName $hsmName -KeyName $keyName -StorageContainerUri $containerUri -BackupFolder $backupFolder -SasToken $sasToken -PassThru
+        $restoreResult | Should -Be $True
+    }
+
     It "Restore a managed HSM" {
         $script:backupUri = [System.Uri]::new($script:backupUri)
         $backupFolder = $script:backupUri.Segments[$script:backupUri.Segments.Length - 1]
+        # Clean hsm
+        Get-AzKeyVaultKey -HsmName $hsmName | Remove-AzKeyVaultKey -Force
+        Get-AzKeyVaultKey -HsmName $hsmName -InRemovedState| Remove-AzKeyVaultKey -InRemovedState -Force
         $restoreResult = Restore-AzKeyVault -HsmName $hsmName -StorageContainerUri $containerUri -BackupFolder $backupFolder -SasToken $sasToken -PassThru
         $restoreResult | Should -Be $True
     }

src/KeyVault/KeyVault/ChangeLog.md

@@ -18,6 +18,7 @@
         - Additional information about change #1
 -->
 ## Upcoming Release
+* Supported selective restore a key from a managed HSM full backup [#13526]
 * Added missing return objects of `Get-Secret` in SecretManagement module
 
 ## Version 3.2.0

src/KeyVault/KeyVault/Commands/FullBackupRestore/FullBackupRestoreCmdletBase.cs

@@ -42,8 +42,8 @@ public abstract class FullBackupRestoreCmdletBase : KeyVaultCmdletBase
         [Parameter(Mandatory = true, HelpMessage = "The shared access signature (SAS) token to authenticate the storage account.")]
         public SecureString SasToken { get; set; }
 
-        [Parameter(ParameterSetName = InputObjectStorageUri, Mandatory = true, HelpMessage = "Managed HSM object")]
-        [Parameter(ParameterSetName = InputObjectStorageName, Mandatory = true, HelpMessage = "Managed HSM object")]
+        [Parameter(ParameterSetName = InputObjectStorageUri, Mandatory = true, ValueFromPipeline = true, HelpMessage = "Managed HSM object")]
+        [Parameter(ParameterSetName = InputObjectStorageName, Mandatory = true, ValueFromPipeline = true, HelpMessage = "Managed HSM object")]
         public PSManagedHsm HsmObject { get; set; }
 
         public override void ExecuteCmdlet()

src/KeyVault/KeyVault/Commands/FullBackupRestore/RestoreAzureManagedHsm.cs

@@ -14,6 +14,9 @@ public class RestoreAzureManagedHsm : FullBackupRestoreCmdletBase
         [Parameter(Mandatory = true, HelpMessage = "Folder name of the backup, e.g. 'mhsm-*-2020101309020403'.\nIt can also be nested such as 'backups/mhsm-*-2020101309020403'.")]
         public string BackupFolder { get; set; }
 
+        [Parameter(Mandatory = false, HelpMessage = "Key name to restore.")]
+        public string KeyName { get; set; }
+
         [Parameter(Mandatory = false, HelpMessage = "Return true when the HSM is restored.")]
         public SwitchParameter PassThru { get; set; }
 
@@ -23,14 +26,24 @@ public override void DoExecuteCmdlet()
                     string.Format(Resources.DoFullRestore, StorageContainerUri),
                     HsmName, () =>
                     {
+                        var errorMsg = string.Format(Resources.FullRestoreFailed, HsmName);
                         try
                         {
-                            Track2DataClient.RestoreHsm(HsmName, StorageContainerUri, SasToken.ConvertToString(), BackupFolder);
+                            if(KeyName == null)
+                            {
+                                Track2DataClient.RestoreHsm(HsmName, StorageContainerUri, SasToken.ConvertToString(), BackupFolder);
+                            }
+                            else
+                            {
+                                Track2DataClient.SelectiveRestoreHsm(HsmName, KeyName,StorageContainerUri, SasToken.ConvertToString(), BackupFolder);
+                                errorMsg = string.Format(Resources.SelectiveRestoreFailed, KeyName, HsmName);
+                            }
                         }
                         catch (Exception ex)
                         {
-                            throw new Exception(string.Format(Resources.FullRestoreFailed, HsmName), ex);
+                            throw new Exception(errorMsg, ex);                            
                         }
+                        
                         if (PassThru)
                         {
                             WriteObject(true);

src/KeyVault/KeyVault/Models/IKeyVaultDataServiceClient.cs

@@ -25,7 +25,9 @@ namespace Microsoft.Azure.Commands.KeyVault.Models
 {
     public interface IKeyVaultDataServiceClient
     {
+        #region Key actions
         PSKeyVaultKey CreateKey(string vaultName, string keyName, PSKeyVaultKeyAttributes keyAttributes, int? size, string curveName);
+       
         PSKeyVaultKey CreateManagedHsmKey(string managedHsmName, string keyName, PSKeyVaultKeyAttributes keyAttributes, int? size, string curveName);
 
         PSKeyVaultKey ImportKey(string vaultName, string keyName, PSKeyVaultKeyAttributes keyAttributes, JsonWebKey webKey, bool? importToHsm);
@@ -68,6 +70,16 @@ public interface IKeyVaultDataServiceClient
 
         PSKeyVaultKey RecoverManagedHsmKey(string managedHsmName, string keyName);
 
+        string BackupKey(string vaultName, string keyName, string outputBlobPath);
+
+        string BackupManagedHsmKey(string managedHsmName, string keyName, string outputBlobPath);
+
+        PSKeyVaultKey RestoreKey(string vaultName, string inputBlobPath);
+
+        PSKeyVaultKey RestoreManagedHsmKey(string managedHsmName, string inputBlobPath);
+        #endregion
+
+        #region Secret actions
         PSKeyVaultSecret SetSecret(string vaultName, string secretName, SecureString secretValue, PSKeyVaultSecretAttributes secretAttributes);
 
         PSKeyVaultSecret UpdateSecret(string vaultName, string secretName, string secretVersion, PSKeyVaultSecretAttributes secretAttributes);
@@ -88,17 +100,10 @@ public interface IKeyVaultDataServiceClient
 
         PSKeyVaultSecret RecoverSecret(string vaultName, string secretName);
 
-        string BackupKey(string vaultName, string keyName, string outputBlobPath);
-
-        string BackupManagedHsmKey(string managedHsmName, string keyName, string outputBlobPath);
-
-        PSKeyVaultKey RestoreKey(string vaultName, string inputBlobPath);
-
-        PSKeyVaultKey RestoreManagedHsmKey(string managedHsmName, string inputBlobPath);
-
         string BackupSecret(string vaultName, string secretName, string outputBlobPath);
 
         PSKeyVaultSecret RestoreSecret(string vaultName, string inputBlobPath);
+        #endregion
 
         #region Certificate actions
 
@@ -198,7 +203,10 @@ public interface IKeyVaultDataServiceClient
 
         #region Full backup restore
         Uri BackupHsm(string hsmName, Uri blobStorageUri, string sasToken);
+        
         void RestoreHsm(string hsmName, Uri backupLocation, string sasToken, string backupFolder);
+       
+        void SelectiveRestoreHsm(string hsmName, string keyName, Uri backupLocation, string sasToken, string backupFolder);
         #endregion
 
         #region RBAC

src/KeyVault/KeyVault/Models/KeyVaultDataServiceClient.cs

@@ -2019,6 +2019,7 @@ public PSKeyVaultKey CreateManagedHsmKey(string managedHsmName, string keyName,
             throw new NotImplementedException("Creating keys on managed HSM is only possible in track 2 SDK.");
         }
 
+        #region Full backup restore
         public Uri BackupHsm(string hsmName, Uri blobStorageUri, string sasToken)
         {
             throw new NotImplementedException();
@@ -2029,6 +2030,12 @@ public void RestoreHsm(string hsmName, Uri blobStorageUri, string sasToken, stri
             throw new NotImplementedException();
         }
 
+        public void SelectiveRestoreHsm(string hsmName, string keyName, Uri backupLocation, string sasToken, string backupFolder)
+        {
+            throw new NotImplementedException();
+        }
+        #endregion
+
         public PSKeyVaultRoleDefinition[] GetHsmRoleDefinitions(string name, string scope)
         {
             throw new NotImplementedException();
@@ -2103,6 +2110,8 @@ public void PurgeManagedHsmKey(string managedHsmName, string keyName)
             throw new NotImplementedException("Purging deleted keys on managed HSM is only possible in track 2 SDK.");
         }
 
+
+
         private VaultUriHelper vaultUriHelper;
         private KeyVaultClient keyVaultClient;
     }

src/KeyVault/KeyVault/Properties/Resources.Designer.cs

@@ -582,4 +582,7 @@ You can find the object ID using Azure Active Directory Module for Windows Power
   <data name="UpdateHsmShouldProcessMessage" xml:space="preserve">
     <value>Updating managed HSM '{0}' in resource group '{1}'.</value>
   </data>
+  <data name="SelectiveRestoreFailed" xml:space="preserve">
+    <value>Failed to selective restore key {0} of managed HSM {1}.</value>
+  </data>
 </root>

src/KeyVault/KeyVault/Properties/Resources.resx

@@ -187,6 +187,20 @@ public void RestoreHsm(string hsmName, Uri backupLocation, string sasToken, stri
                 throw;
             }
         }
+        
+        public void SelectiveRestoreHsm(string hsmName, string keyName, Uri backupLocation, string sasToken, string backupFolder)
+        {
+            var client = CreateBackupClient(hsmName);
+            try
+            {
+                client.StartSelectiveRestore(keyName, backupLocation, sasToken, backupFolder)
+                    .WaitForCompletionAsync().ConfigureAwait(false).GetAwaiter().GetResult();
+            }
+            catch
+            {
+                throw;
+            }
+        }
 
         public PSKeyVaultRoleDefinition[] GetHsmRoleDefinitions(string hsmName, string scope)
         {

src/KeyVault/KeyVault/Track2Models/Track2HsmClient.cs

@@ -388,6 +388,7 @@ public PSKeyVaultSecret UpdateSecret(string vaultName, string secretName, string
             throw new NotImplementedException();
         }
 
+        #region Full backup restore
         public Uri BackupHsm(string hsmName, Uri blobStorageUri, string sasToken)
         {
             return HsmClient.BackupHsm(hsmName, blobStorageUri, sasToken);
@@ -398,6 +399,12 @@ public void RestoreHsm(string hsmName, Uri backupLocation, string sasToken, stri
             HsmClient.RestoreHsm(hsmName, backupLocation, sasToken, backupFolder);
         }
 
+        public void SelectiveRestoreHsm(string hsmName, string keyName, Uri backupLocation, string sasToken, string backupFolder)
+        {
+            HsmClient.SelectiveRestoreHsm(hsmName, keyName, backupLocation, sasToken, backupFolder);
+        }
+        #endregion
+
         public PSKeyVaultRoleDefinition[] GetHsmRoleDefinitions(string hsmName, string scope)
         {
             return HsmClient.GetHsmRoleDefinitions(hsmName, scope);