[Az.Batch] Adds encryption configuration (#20687)

* Add BatchAccount encryption

Update help

Updated changelog

* Added unit test

* whitespace

* Add BatchAccount encryption

Update help

Updated changelog

* Added unit test

* whitespace

* Flattened encryption

* Fixed identity type, unit test, and context

Fixed unit test and context

Updated help

* typo

---------

Co-authored-by: Yabo Hu <[email protected]>

Author: paterasMSFT

src/Batch/Batch.Test/BatchAccounts/NewBatchAccountCommandTests.cs

@@ -144,5 +144,44 @@ public void CanCreateUserSubscriptionBatchAccount()
             Assert.Equal(keyVaultId, actualCreateParameters.KeyVaultId);
             Assert.Equal(keyVaultUrl, actualCreateParameters.KeyVaultUrl);
         }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void NewBatchWithCustomerEncryptionKeyAccountTest()
+        {
+            string accountName = "account01";
+            string resourceGroup = "resourceGroup";
+            string location = "location";
+            string identityId = "subscriptions/0000/resourceGroups/resourceGroup/providers/Microsoft.ManagedIdentity/userAssignedIdentities/encryptionIdentity";
+            string encryptionKeyVaultId = "subscriptions/0000/resourceGroups/resourceGroup/providers/Microsoft.KeyVault/vaults/encryptionVault";
+            AccountCreateParameters actualCreateParameters = null;
+
+            // Setup the mock client to return a fake response and capture the account create parameters
+            BatchAccount accountResource = BatchTestHelpers.CreateAccountResource(accountName, resourceGroup, location);
+            BatchAccountContext fakeResponse = BatchAccountContext.ConvertAccountResourceToNewAccountContext(accountResource, null);
+
+            batchClientMock.Setup(b => b.CreateAccount(It.IsAny<AccountCreateParameters>()))
+                .Returns(fakeResponse)
+                .Callback((AccountCreateParameters p) => actualCreateParameters = p);
+
+            // Setup and run the cmdlet
+            cmdlet.AccountName = accountName;
+            cmdlet.ResourceGroupName = resourceGroup;
+            cmdlet.Location = location;
+            cmdlet.IdentityType = ResourceIdentityType.UserAssigned;
+            cmdlet.IdentityId = new string[] { identityId };
+            cmdlet.EncryptionKeySource = KeySource.MicrosoftKeyVault;
+            cmdlet.EncryptionKeyIdentifier = encryptionKeyVaultId;
+            cmdlet.ExecuteCmdlet();
+
+            // Verify the fake response was written to the pipeline and that the captured account create
+            // parameters matched expectations.
+            commandRuntimeMock.Verify(r => r.WriteObject(fakeResponse), Times.Once());
+            Assert.Equal(accountName, actualCreateParameters.BatchAccount);
+            Assert.Equal(resourceGroup, actualCreateParameters.ResourceGroup);
+            Assert.Equal(location, actualCreateParameters.Location);
+            Assert.Equal(KeySource.MicrosoftKeyVault, actualCreateParameters.Encryption.KeySource);
+            Assert.Equal(encryptionKeyVaultId, actualCreateParameters.Encryption.KeyVaultProperties.KeyIdentifier);
+        }
     }
 }

src/Batch/Batch/BatchAccounts/NewBatchAccountCommand.cs

@@ -69,6 +69,12 @@ public class NewBatchAccountCommand : BatchCmdletBase
         [Parameter(Mandatory = false, HelpMessage = "An array containing user assigned identities associated with the BatchAccount. This parameter is only used when IdentityType is set to UserAssigned.")]
         public string[] IdentityId { get; set; }
 
+        [Parameter(Mandatory = false, HelpMessage = "Configures how customer data is encrypted inside the Batch account.\r\nBy default, accounts are encrypted using a Microsoft managed key.\r\nFor additional control, a customer-managed key can be used instead.")]
+        public KeySource EncryptionKeySource { get; set; }
+
+        [Parameter(Mandatory = false, HelpMessage = "The Key Identifier for customer-based encryption.")]
+        public string EncryptionKeyIdentifier { get; set; }
+
         protected override void ExecuteCmdletImpl()
         {
             Dictionary<string, UserAssignedIdentities> identityDictionary = null;
@@ -82,6 +88,22 @@ protected override void ExecuteCmdletImpl()
                 identityDictionary = IdentityId.ToDictionary(i => i, i => new UserAssignedIdentities());
             }
 
+            EncryptionProperties encryption = null;
+            if (EncryptionKeySource == KeySource.MicrosoftKeyVault)
+            {
+                if (IdentityType != ResourceIdentityType.UserAssigned)
+                {
+                    throw new PSArgumentException("If EncryptionKeySource is set to 'MicrosoftKeyVault', the Batch Account identity must be set to `UserAssigned`.");
+                }
+
+                if (EncryptionKeyIdentifier == null)
+                {
+                    throw new PSArgumentNullException("EncryptionKeyIdentifier", "If EncryptionKeySource is set to 'MicrosoftKeyVault', a valid Key Identifier must also be supplied as parameter 'EncryptionKeyIdentifier'.");
+                }
+
+                encryption = new EncryptionProperties(EncryptionKeySource, new KeyVaultProperties(EncryptionKeyIdentifier));
+            }
+
             AccountCreateParameters parameters = new AccountCreateParameters(this.ResourceGroupName, this.AccountName, this.Location)
             {
                 AutoStorageAccountId = this.AutoStorageAccountId,
@@ -90,7 +112,8 @@ protected override void ExecuteCmdletImpl()
                 KeyVaultUrl = this.KeyVaultUrl,
                 Tags = this.Tag,
                 PublicNetworkAccess = this.PublicNetworkAccess,
-                Identity = new BatchAccountIdentity(IdentityType, null, null, identityDictionary)
+                Identity = new BatchAccountIdentity(IdentityType, null, null, identityDictionary),
+                Encryption = encryption
             };
 
             BatchAccountContext context = BatchClient.CreateAccount(parameters);

src/Batch/Batch/ChangeLog.md

@@ -18,6 +18,8 @@
         - Additional information about change #1
 -->
 ## Upcoming Release
+* Added new property `Encryption` of type `EncryptionProperties` to `AccountCreateParameters`.
+  - Configures how customer data is encrypted inside the Batch account.
 
 ## Version 3.3.0
 * Added new properties `CurrentNodeCommunicationMode` (read only) and `TargetCommunicationMode` of type `NodeCommunicationMode` to `PSCloudPool`.

src/Batch/Batch/Models/AccountCreateParameters.cs

@@ -91,5 +91,10 @@ public AccountCreateParameters(string resourceGroup, string batchAccount, string
         /// The identity of the Batch account.
         /// </summary>
         public BatchAccountIdentity Identity { get; set; }
+
+        /// <summary>
+        /// Gets the encryption configuration for the Batch account.
+        /// </summary>
+        public EncryptionProperties Encryption { get; set; }
     }
 }

src/Batch/Batch/Models/BatchAccountContext.cs

@@ -179,6 +179,16 @@ public AccountKeyType KeyInUse
             }
         }
 
+        /// <summary>
+        /// Gets the encryption configuration for the Batch account.
+        /// </summary>
+        /// <remarks>
+        /// Configures how customer data is encrypted inside the Batch account.
+        /// By default, accounts are encrypted using a Microsoft managed key.
+        /// For additional control, a customer-managed key can be used instead.
+        /// </remarks>
+        public EncryptionProperties Encryption { get; private set; }
+
         internal bool HasKeys
         {
             get { return !string.IsNullOrEmpty(PrimaryAccountKey) || !string.IsNullOrEmpty(SecondaryAccountKey);  }
@@ -242,6 +252,7 @@ internal void ConvertAccountResourceToAccountContext(BatchAccount resource)
             this.PoolAllocationMode = resource.PoolAllocationMode;
             this.PublicNetworkAccess = resource.PublicNetworkAccess;
             this.Identity = resource.Identity;
+            this.Encryption = resource.Encryption;
 
             if (resource.AutoStorage != null)
             {

src/Batch/Batch/Models/BatchClient.BatchAccounts.cs

@@ -64,7 +64,8 @@ public virtual BatchAccountContext CreateAccount(AccountCreateParameters paramet
                 PoolAllocationMode = parameters.PoolAllocationMode,
                 KeyVaultReference = keyVaultRef,
                 PublicNetworkAccess = parameters.PublicNetworkAccess,
-                Identity = parameters.Identity
+                Identity = parameters.Identity,
+                Encryption = parameters.Encryption
             });
 
             var context = BatchAccountContext.ConvertAccountResourceToNewAccountContext(response, this.azureContext);

src/Batch/Batch/help/Get-AzBatchAccountKey.md

@@ -32,6 +32,8 @@ This command gets the account details and stores it in a `$Context` object for u
 
 ### Example 2: Get batch account keys and display them
 <!-- Skip: Output cannot be splitted from code -->
+
+
 ```powershell
 $Context = Get-AzBatchAccountKey -AccountName myaccount
 $Context.PrimaryAccountKey

src/Batch/Batch/help/New-AzBatchAccount.md

@@ -17,8 +17,8 @@ Creates a Batch account.
 New-AzBatchAccount [-AccountName] <String> [-Location] <String> [-ResourceGroupName] <String>
  [[-AutoStorageAccountId] <String>] [-PoolAllocationMode <PoolAllocationMode>] [-KeyVaultId <String>]
  [-KeyVaultUrl <String>] [-Tag <Hashtable>] [-PublicNetworkAccess <PublicNetworkAccessType>]
- [-IdentityType <ResourceIdentityType>] [-IdentityId <String[]>] [-DefaultProfile <IAzureContextContainer>]
- [<CommonParameters>]
+ [-IdentityType <ResourceIdentityType>] [-IdentityId <String[]>] [-EncryptionKeySource <KeySource>]
+ [-EncryptionKeyIdentifier <String>] [-DefaultProfile <IAzureContextContainer>] [<CommonParameters>]
 ```
 
 ## DESCRIPTION
@@ -104,6 +104,39 @@ Accept pipeline input: False
 Accept wildcard characters: False
 ```
 
+### -EncryptionKeyIdentifier
+The Key Identifier for customer-based encryption.
+
+```yaml
+Type: System.String
+Parameter Sets: (All)
+Aliases:
+
+Required: False
+Position: Named
+Default value: None
+Accept pipeline input: False
+Accept wildcard characters: False
+```
+
+### -EncryptionKeySource
+Configures how customer data is encrypted inside the Batch account.
+By default, accounts are encrypted using a Microsoft managed key.
+For additional control, a customer-managed key can be used instead.
+
+```yaml
+Type: Microsoft.Azure.Management.Batch.Models.KeySource
+Parameter Sets: (All)
+Aliases:
+Accepted values: MicrosoftBatch, MicrosoftKeyVault
+
+Required: False
+Position: Named
+Default value: None
+Accept pipeline input: False
+Accept wildcard characters: False
+```
+
 ### -IdentityId
 The list of user assigned identities associated with the BatchAccount. This parameter is only used when IdentityType is set to UserAssigned.