Fixing Set-AzNetworkRuleSet and adding TrustedServiceAccess

src/ServiceBus/ServiceBus.Test/ScenarioTests/NetworkRuleSetTests.ps1

@@ -20,6 +20,9 @@ Tests New Parameter for ServiceBus Namespace Create List Remove operations.
 function NetworkRuleSetTests {
     # Setup  
 
+    #Customer can use this cmdlet to 
+    #Set default action, trustedservice, publicnetworkaccess
+
     $location = Get-Location
     $resourceGroupName = getAssetName "RSG"
     $namespaceName = getAssetName "ServiceBus-Namespace-"
@@ -92,11 +95,19 @@ function NetworkRuleSetTests {
     Assert-AreEqual $setResult.IpRules.Count 3 "Set - IPRules count did not matched"
 
     # Set-AzServiceBusNetworkRuleSet with parameters
-    $setResult = Set-AzServiceBusNetworkRuleSet -ResourceGroup $resourceGroupName -Name $namespaceName2 -IPRule $setResult.IpRules -VirtualNetworkRule $setResult.VirtualNetworkRules -DefaultAction "Allow" -PublicNetworkAccess "Disabled"
+    $setResult = Set-AzServiceBusNetworkRuleSet -ResourceGroup $resourceGroupName -Name $namespaceName2 -DefaultAction "Allow" -PublicNetworkAccess "Disabled" -TrustedServiceAccessEnabled
+    Assert-AreEqual $setResult.VirtualNetworkRules.Count 3 "Set -VirtualNetworkRules count did not matched"
+    Assert-AreEqual $setResult.IpRules.Count 3 "Set - IPRules count did not matched"
+    Assert-AreEqual $setResult.PublicNetworkAccess "Disabled"
+    Assert-AreEqual $setResult.DefaultAction "Allow"
+    Assert-True {$setResult.TrustedServiceAccessEnabled}
+
+    $setResult = Set-AzServiceBusNetworkRuleSet -ResourceGroup $resourceGroupName -Name $namespaceName2 -IPRule $setResult.IpRules -VirtualNetworkRule $setResult.VirtualNetworkRules
     Assert-AreEqual $setResult.VirtualNetworkRules.Count 3 "Set -VirtualNetworkRules count did not matched"
     Assert-AreEqual $setResult.IpRules.Count 3 "Set - IPRules count did not matched"
     Assert-AreEqual $setResult.PublicNetworkAccess "Disabled"
     Assert-AreEqual $setResult.DefaultAction "Allow"
+    Assert-True {$setResult.TrustedServiceAccessEnabled}
 
     # Set-AzServiceBusNetworkRuleSet with Resource ID
     $setResult1 = Set-AzServiceBusNetworkRuleSet -ResourceGroup $resourceGroupName -Name $namespaceName2 -ResourceId $getResult.Id

src/ServiceBus/ServiceBus.Test/ScenarioTests/ServiceBusServiceTests.cs

@@ -41,14 +41,14 @@ public void ServiceBusNameSpaceAuth_CURD_Tests()
             TestRunner.RunTestScript("ServiceBusNameSpaceAuthTests");
         }
 
-        [Fact(Skip = "Skip as current test framework does not support recording generated cmdlets.")]
+        [Fact]
         [Trait(Category.AcceptanceType, Category.LiveOnly)]
         public void ServiceBusNameSpaceEncryption_CRUD()
         {
             TestRunner.RunTestScript("EncryptionTest");
         }
 
-        [Fact(Skip = "Skip as current test framework does not support recording generated cmdlets.")]
+        [Fact]
         [Trait(Category.AcceptanceType, Category.LiveOnly)]
         public void ServiceBusNameSpaceMSI()
         {

src/ServiceBus/ServiceBus.Test/ScenarioTests/ServiceBusServiceTests.ps1

@@ -249,9 +249,9 @@ function MSITest{
     $namespace2 = getAssetName "Namespace2-"
     try{
 
-        $uad1 = Get-AzUserAssignedIdentity -ResourceGroupName $resourceGroupName -Name $msi1
-        $uad2 = Get-AzUserAssignedIdentity -ResourceGroupName $resourceGroupName -Name $msi2
-        $uad3 = Get-AzUserAssignedIdentity -ResourceGroupName $resourceGroupName -Name $msi3
+        $uad1 = "/subscriptions/326100e2-f69d-4268-8503-075374f62b6e/resourceGroups/PS-Testing/providers/Microsoft.ManagedIdentity/userAssignedIdentities/PS-Testing-MSI1"
+        $uad2 = "/subscriptions/326100e2-f69d-4268-8503-075374f62b6e/resourceGroups/PS-Testing/providers/Microsoft.ManagedIdentity/userAssignedIdentities/PS-Testing-MSI2"
+        $uad3 = "/subscriptions/326100e2-f69d-4268-8503-075374f62b6e/resourceGroups/PS-Testing/providers/Microsoft.ManagedIdentity/userAssignedIdentities/PS-Testing-MSI3"
 
         $namespace = New-AzServiceBusNamespace -ResourceGroupName $resourceGroupName -Name $namespace1 -SkuName Standard -Location northeurope
         Assert-AreEqual $namespace.Name $namespace1
@@ -262,13 +262,13 @@ function MSITest{
         Assert-AreEqual $namespace.Sku.Name "Standard"
         Assert-AreEqual $namespace.IdentityType "SystemAssigned"
 
-        $namespace = Set-AzServiceBusNamespace -ResourceGroupName $resourceGroupName -Name $namespace1 -IdentityType "UserAssigned" -IdentityId $uad1.Id,$uad2.Id
+        $namespace = Set-AzServiceBusNamespace -ResourceGroupName $resourceGroupName -Name $namespace1 -IdentityType "UserAssigned" -IdentityId $uad1,$uad2
         Assert-AreEqual $namespace.Name $namespace1
         Assert-AreEqual $namespace.Sku.Name "Standard"
         Assert-AreEqual $namespace.IdentityType "UserAssigned"
         Assert-True { $namespace.IdentityId.Count -eq 2 }
 
-        $namespace = Set-AzServiceBusNamespace -ResourceGroupName $resourceGroupName -Name $namespace1 -IdentityType "SystemAssigned, UserAssigned" -IdentityId $uad1.Id,$uad2.Id
+        $namespace = Set-AzServiceBusNamespace -ResourceGroupName $resourceGroupName -Name $namespace1 -IdentityType "SystemAssigned, UserAssigned"
         Assert-AreEqual $namespace.Name $namespace1
         Assert-AreEqual $namespace.Sku.Name "Standard"
         Assert-AreEqual $namespace.IdentityType "SystemAssignedUserAssigned"
@@ -290,16 +290,16 @@ function EncryptionTest{
         $msi1 = "PS-Testing-MSI1"
         $msi2 = "PS-Testing-MSI2"
         $msi3 = "PS-Testing-MSI3"
-        $kv1 = "PS-Testing-kv1"
-        $kv2 = "PS-Testing-kv2"
-        $kv1uri = "https://ps-testing-kv1.vault.azure.net/"
-        $kv2uri = "https://ps-testing-kv2.vault.azure.net"
+        $kv1 = "PS-Test-kv1"
+        $kv2 = "PS-Test-kv2"
+        $kv1uri = "https://ps-test-kv1.vault.azure.net/"
+        $kv2uri = "https://ps-test-kv2.vault.azure.net"
         $namespace1 = getAssetName "Namespace1-"
         $namespace2 = getAssetName "Namespace2-"
 
-        $uad1 = Get-AzUserAssignedIdentity -ResourceGroupName $resourceGroupName -Name $msi1
-        $uad2 = Get-AzUserAssignedIdentity -ResourceGroupName $resourceGroupName -Name $msi2
-        $uad3 = Get-AzUserAssignedIdentity -ResourceGroupName $resourceGroupName -Name $msi3
+        $uad1 = "/subscriptions/326100e2-f69d-4268-8503-075374f62b6e/resourceGroups/PS-Testing/providers/Microsoft.ManagedIdentity/userAssignedIdentities/PS-Testing-MSI1"
+        $uad2 = "/subscriptions/326100e2-f69d-4268-8503-075374f62b6e/resourceGroups/PS-Testing/providers/Microsoft.ManagedIdentity/userAssignedIdentities/PS-Testing-MSI2"
+        $uad3 = "/subscriptions/326100e2-f69d-4268-8503-075374f62b6e/resourceGroups/PS-Testing/providers/Microsoft.ManagedIdentity/userAssignedIdentities/PS-Testing-MSI3"
 
         $namespace = New-AzServiceBusNamespace -ResourceGroupName $resourceGroupName -Name $namespace2 -SkuName Premium -Location northeurope -IdentityType SystemAssigned
         Assert-AreEqual $namespace.Name $namespace2
@@ -318,17 +318,17 @@ function EncryptionTest{
         Assert-True { $namespace.EncryptionConfig.Count -eq 2 }
 
 
-        $ec1 = New-AzServiceBusEncryptionConfig -KeyName key1 -KeyVaultUri $kv1uri -UserAssignedIdentity $uad1.Id
-        $ec2 = New-AzServiceBusEncryptionConfig -KeyName key2 -KeyVaultUri $kv1uri -UserAssignedIdentity $uad1.Id
+        $ec1 = New-AzServiceBusEncryptionConfig -KeyName key1 -KeyVaultUri $kv1uri -UserAssignedIdentity $uad1
+        $ec2 = New-AzServiceBusEncryptionConfig -KeyName key2 -KeyVaultUri $kv1uri -UserAssignedIdentity $uad1
 
-        $namespace = New-AzServiceBusNamespace -ResourceGroupName $resourceGroupName -Name $namespace1 -SkuName Premium -Location northeurope -IdentityType UserAssigned -IdentityId $uad1.Id,$uad2.Id -EncryptionConfig $ec1,$ec2
+        $namespace = New-AzServiceBusNamespace -ResourceGroupName $resourceGroupName -Name $namespace1 -SkuName Premium -Location northeurope -IdentityType UserAssigned -IdentityId $uad1,$uad2 -EncryptionConfig $ec1,$ec2
         Assert-AreEqual $namespace.Name $namespace1
         Assert-AreEqual $namespace.Sku.Name "Premium"
         Assert-AreEqual $namespace.IdentityType "UserAssigned"
         Assert-True { $namespace.IdentityId.Count -eq 2 }
         Assert-True { $namespace.EncryptionConfig.Count -eq 2 }
 
-        $ec3 = New-AzServiceBusEncryptionConfig -KeyName key1 -KeyVaultUri $kv2uri -UserAssignedIdentity $uad1.id
+        $ec3 = New-AzServiceBusEncryptionConfig -KeyName key1 -KeyVaultUri $kv2uri -UserAssignedIdentity $uad1
         $namespace.EncryptionConfig += $ec3
 
         $namespace = Set-AzServiceBusNamespace -ResourceGroupName $resourceGroupName -Name $namespace1 -EncryptionConfig $namespace.EncryptionConfig -Location northeurope

src/ServiceBus/ServiceBus/ChangeLog.md

@@ -20,6 +20,7 @@
 -->
 ## Upcoming Release
 * Fixed miscellaneous network rule set typos across module.
+* Add `TrustedServiceAccessEnabled` to `Set-AzServiceBusNetworkRuleSet`
 
 ## Version 1.8.1
 * Fixed that `New-AzServiceBusAuthorizationRuleSASToken` returns invalid token. [#12975]

src/ServiceBus/ServiceBus/Cmdlets/NetworkRuleSet/SetAzureServiceBusNetworkrule.cs

@@ -49,12 +49,13 @@ public class SetAzureServiceBusNetworkrule : AzureServiceBusCmdletBase
         [PSDefaultValue(Value = "Enabled")]
         public string PublicNetworkAccess { get; set; }
 
-        [Parameter(Mandatory = true, ParameterSetName = NetworkRuleSetPropertiesParameterSet, Position = 2, HelpMessage = "List of IPRuleSet")]
-        [ValidateNotNullOrEmpty]
+        [Parameter(Mandatory = false, ParameterSetName = NetworkRuleSetPropertiesParameterSet, HelpMessage = "Trusted Service Access for NetworkRuleSet")]
+        public SwitchParameter TrustedServiceAccessEnabled { get; set; }
+
+        [Parameter(Mandatory = false, ParameterSetName = NetworkRuleSetPropertiesParameterSet, Position = 2, HelpMessage = "List of IPRuleSet")]
         public PSNWRuleSetIpRulesAttributes[] IPRule { get; set; }
 
-        [Parameter(Mandatory = true, ParameterSetName = NetworkRuleSetPropertiesParameterSet, Position = 3, HelpMessage = "List of VirtualNetworkRules")]
-        [ValidateNotNullOrEmpty]
+        [Parameter(Mandatory = false, ParameterSetName = NetworkRuleSetPropertiesParameterSet, Position = 3, HelpMessage = "List of VirtualNetworkRules")]
         [Alias(AliasVirtualNetworkRule)]
         public PSNWRuleSetVirtualNetworkRulesAttributes[] VirtualNetworkRule { get; set; }
 
@@ -78,15 +79,13 @@ public override void ExecuteCmdlet()
 
                     if (ParameterSetName.Equals(NetworkRuleSetPropertiesParameterSet))
                     {
-                        PSNetworkRuleSetAttributes networkRuleSetAttributes = new PSNetworkRuleSetAttributes()
-                        {
-                            DefaultAction = DefaultAction,
-                            IpRules = IPRule.OfType<PSNWRuleSetIpRulesAttributes>().ToList(),
-                            VirtualNetworkRules = VirtualNetworkRule.OfType<PSNWRuleSetVirtualNetworkRulesAttributes>().ToList(),
-                            PublicNetworkAccess = PublicNetworkAccess
-                        };
-
-                        WriteObject(Client.CreateOrUpdateNetworkRuleSet(ResourceGroupName, Name, networkRuleSetAttributes));
+                        WriteObject(Client.UpdateNetworkRuleSet(resourceGroupName: ResourceGroupName, 
+                                                                namespaceName: Name, 
+                                                                publicNetworkAccess: PublicNetworkAccess,
+                                                                trustedServiceAccessEnabled: TrustedServiceAccessEnabled,
+                                                                defaultAction: DefaultAction,
+                                                                iPRule: IPRule,
+                                                                virtualNetworkRule: VirtualNetworkRule));
                     }
 
                     if (ParameterSetName.Equals(NetworkRuleSetInputObjectParameterSet))

src/ServiceBus/ServiceBus/Models/PSNetworkRuleSetAttributes.cs

@@ -51,6 +51,7 @@ public PSNetworkRuleSetAttributes(NetworkRuleSet networkRuleSet)
             Name = networkRuleSet.Name;
             Type = networkRuleSet.Type;
             PublicNetworkAccess = networkRuleSet.PublicNetworkAccess;
+            TrustedServiceAccessEnabled = networkRuleSet.TrustedServiceAccessEnabled;
         }
 
         /// <summary>
@@ -80,6 +81,11 @@ public PSNetworkRuleSetAttributes(NetworkRuleSet networkRuleSet)
         /// </summary>
         public string PublicNetworkAccess { get; set; }
 
+        /// <summary>
+        /// Value that indicates whether Trusted Service Access is Enabled or not.
+        /// </summary>
+        public bool? TrustedServiceAccessEnabled { get; set; }
+
         /// <summary>
         /// Gets or sets list VirtualNetwork Rules
         /// </summary>

src/ServiceBus/ServiceBus/ServiceBus.format.ps1xml

@@ -198,6 +198,10 @@
                 <Label>PublicNetworkAccess</Label>
                 <PropertyName>PublicNetworkAccess</PropertyName>
               </ListItem>
+			  <ListItem>
+			    <Label>TrustedServiceAccessEnabled</Label>
+				<PropertyName>TrustedServiceAccessEnabled</PropertyName>
+			  </ListItem>
               <ListItem>
                 <Label>Id</Label>
                 <PropertyName>Id</PropertyName>

src/ServiceBus/ServiceBus/Utilities/ServiceBusClient.cs

@@ -216,10 +216,10 @@ public PSNamespaceAttributes UpdateNamespace(string resourceGroupName, string na
 
                 parameter.Identity.Type = FindIdentity(identityType);
 
-                if (parameter.Identity.Type == ManagedServiceIdentityType.None || parameter.Identity.Type == ManagedServiceIdentityType.SystemAssigned)
+                /*if (parameter.Identity.Type == ManagedServiceIdentityType.None || parameter.Identity.Type == ManagedServiceIdentityType.SystemAssigned)
                 {
                     parameter.Identity.UserAssignedIdentities = null;
-                }
+                }*/
             }
 
             if (identityIds != null)
@@ -236,9 +236,13 @@ public PSNamespaceAttributes UpdateNamespace(string resourceGroupName, string na
                 {
                     parameter.Identity.UserAssignedIdentities = UserAssignedIdentities;
                 }
-                if (parameter.Identity.Type == ManagedServiceIdentityType.None || parameter.Identity.Type == ManagedServiceIdentityType.SystemAssigned)
+                if (identityIds.Length == 0)
                 {
-                    throw new Exception("Please change -IdentityType to 'UserAssigned' or 'SystemAssigned, UserAssigned' if you want to add User Assigned Identities");
+                    parameter.Identity.UserAssignedIdentities = null;
+                }
+                else if (parameter.Identity.Type == ManagedServiceIdentityType.None || parameter.Identity.Type == ManagedServiceIdentityType.SystemAssigned)
+                {
+                    throw new Exception("Please change -IdentityType to UserAssigned or 'SystemAssigned, UserAssigned' if you want to add User Assigned Identities");
                 }
             }
 
@@ -336,6 +340,11 @@ public PSNetworkRuleSetAttributes CreateOrUpdateNetworkRuleSet(string resourceGr
 
             networkRuleSet.PublicNetworkAccess = psNetworkRuleSetAttributes.PublicNetworkAccess;
 
+            if(psNetworkRuleSetAttributes.TrustedServiceAccessEnabled != null)
+            {
+                networkRuleSet.TrustedServiceAccessEnabled = psNetworkRuleSetAttributes.TrustedServiceAccessEnabled;
+            }
+
             foreach (PSNWRuleSetIpRulesAttributes psiprules in psNetworkRuleSetAttributes.IpRules)
             {
                 networkRuleSet.IpRules.Add(new NWRuleSetIpRules { Action = psiprules.Action, IpMask = psiprules.IpMask });
@@ -350,6 +359,55 @@ public PSNetworkRuleSetAttributes CreateOrUpdateNetworkRuleSet(string resourceGr
             return new PSNetworkRuleSetAttributes(response);
         }
 
+        public PSNetworkRuleSetAttributes UpdateNetworkRuleSet(string resourceGroupName, string namespaceName, string publicNetworkAccess, bool trustedServiceAccessEnabled, string defaultAction, PSNWRuleSetIpRulesAttributes[] iPRule, PSNWRuleSetVirtualNetworkRulesAttributes[] virtualNetworkRule)
+        {
+            NetworkRuleSet networkRuleSet = Client.Namespaces.GetNetworkRuleSet(resourceGroupName, namespaceName);
+
+            if(networkRuleSet == null)
+            {
+                networkRuleSet = new NetworkRuleSet();
+            }
+
+            if(defaultAction != null)
+            {
+                networkRuleSet.DefaultAction = defaultAction;
+            }
+
+            if (publicNetworkAccess != null)
+            {
+                networkRuleSet.PublicNetworkAccess = publicNetworkAccess;
+            }
+
+            if (trustedServiceAccessEnabled == true)
+            {
+                networkRuleSet.TrustedServiceAccessEnabled = trustedServiceAccessEnabled;
+            }
+
+            if (iPRule != null)
+            {
+                networkRuleSet.IpRules = new List<NWRuleSetIpRules>();
+
+                foreach (PSNWRuleSetIpRulesAttributes psiprules in iPRule)
+                {
+                    networkRuleSet.IpRules.Add(new NWRuleSetIpRules { Action = psiprules.Action, IpMask = psiprules.IpMask });
+                }
+            }
+
+            if(virtualNetworkRule != null)
+            {
+                networkRuleSet.VirtualNetworkRules = new List<NWRuleSetVirtualNetworkRules>();
+
+                foreach (PSNWRuleSetVirtualNetworkRulesAttributes psvisrtualnetworkrules in virtualNetworkRule)
+                {
+                    networkRuleSet.VirtualNetworkRules.Add(new NWRuleSetVirtualNetworkRules { Subnet = new Subnet { Id = psvisrtualnetworkrules.Subnet.Id }, IgnoreMissingVnetServiceEndpoint = psvisrtualnetworkrules.IgnoreMissingVnetServiceEndpoint });
+                }
+            }
+
+            var response = Client.Namespaces.CreateOrUpdateNetworkRuleSet(resourceGroupName, namespaceName, networkRuleSet);
+            return new PSNetworkRuleSetAttributes(response);
+
+        }
+
         #endregion
 
         #region NameSpace AuthorizationRules