{Keyvault} Support ip rules for MHSM

src/azure-cli/azure/cli/command_modules/keyvault/_help.py

@@ -569,21 +569,30 @@
 
 helps['keyvault network-rule'] = """
 type: group
-short-summary: Manage vault network ACLs.
+short-summary: Manage network ACLs for vault or managed hsm.
 """
 
 helps['keyvault network-rule add'] = """
 type: command
-short-summary: Add a network rule to the network ACLs for a Key Vault.
+short-summary: Add a network rule to the network ACLs for a Key Vault or a Managed HSM.
 examples:
   - name: Add a specific IP address to the network access rules of a Key Vault in a specified resource group using this command. (autogenerated)
     text: |
         az keyvault network-rule add --ip-address 1.2.3.4 --name cli-mock-kv-nr --resource-group cli_mock_keyvault_network_rule
 """
 
+helps['keyvault network-rule add'] = """
+type: command
+short-summary: Remove a network rule to the network ACLs for a Key Vault or a Managed HSM.
+examples:
+  - name: Remove a specific IP address to the network access rules of a Key Vault in a specified resource group.
+    text: |
+        az keyvault network-rule remove --ip-address 1.2.3.4 --name cli-mock-kv-nr --resource-group cli_mock_keyvault_network_rule
+"""
+
 helps['keyvault network-rule wait'] = """
 type: command
-short-summary: Place the CLI in a waiting state until a condition of the vault is met.
+short-summary: Place the CLI in a waiting state until a condition of the vault or managed hsm is met.
 examples:
   - name: Pause CLI until the network ACLs are updated.
     text: |

src/azure-cli/azure/cli/command_modules/keyvault/_params.py

@@ -104,7 +104,7 @@ class CLISecurityDomainOperation(str, Enum):
     # region vault (management)
     with self.argument_context('keyvault') as c:
         c.argument('resource_group_name', resource_group_name_type, id_part=None, required=False,
-                   help='Proceed only if Key Vault belongs to the specified resource group.',
+                   help='Name of resource group.',
                    validator=validate_resource_group_name)
         c.argument('vault_name', vault_name_type, options_list=['--name', '-n'])
         c.argument('object_id', help='a GUID that identifies the principal that will receive permissions')
@@ -244,6 +244,7 @@ class CLISecurityDomainOperation(str, Enum):
                    help='Space-separated list of storage permissions to assign.')
 
     with self.argument_context('keyvault network-rule') as c:
+        c.argument('hsm_name', mgmt_plane_hsm_name_type)
         c.argument('ip_address', help='IPv4 address or CIDR range.')
         c.argument('subnet', help='Name or ID of subnet. If name is supplied, `--vnet-name` must be supplied.')
         c.argument('vnet_name', help='Name of a virtual network.', validator=validate_subnet)

src/azure-cli/azure/cli/command_modules/keyvault/commands.py

@@ -97,10 +97,10 @@ def load_command_table(self, _):
     with self.command_group('keyvault network-rule',
                             mgmt_vaults_entity.command_type,
                             client_factory=mgmt_vaults_entity.client_factory) as g:
-        g.custom_command('add', 'add_network_rule', supports_no_wait=True)
-        g.custom_command('remove', 'remove_network_rule', supports_no_wait=True)
+        g.custom_command('add', 'add_network_rule_for_vault_or_hsm', supports_no_wait=True)
+        g.custom_command('remove', 'remove_network_rule_for_vault_or_hsm', supports_no_wait=True)
         g.custom_command('list', 'list_network_rules')
-        g.wait_command('wait')
+        g.custom_wait_command('wait', 'wait_vault_or_hsm')
 
     with self.command_group('keyvault private-endpoint-connection',
                             mgmt_pec_entity.command_type,

src/azure-cli/azure/cli/command_modules/keyvault/custom.py

@@ -474,6 +474,8 @@ def create_vault_or_hsm(cmd, client,
                               sku=sku,
                               enable_purge_protection=enable_purge_protection,
                               retention_days=retention_days,
+                              network_acls_ips=network_acls_ips,
+                              public_network_access=public_network_access,
                               bypass=bypass,
                               default_action=default_action,
                               tags=tags,
@@ -493,6 +495,7 @@ def create_hsm(cmd, client,
                resource_group_name, hsm_name, administrators, location=None, sku=None,
                enable_purge_protection=None,
                retention_days=None,
+               network_acls_ips=None,
                public_network_access=None,
                bypass=None,
                default_action=None,
@@ -517,6 +520,14 @@ def create_hsm(cmd, client,
                                           operation_group='managed_hsms')
     ManagedHsmSku = cmd.get_models('ManagedHsmSku', resource_type=ResourceType.MGMT_KEYVAULT,
                                    operation_group='managed_hsms')
+    MHSMIPRule = cmd.get_models('MHSMIPRule', resource_type=ResourceType.MGMT_KEYVAULT,
+                                   operation_group='managed_hsms')
+
+    network_acls = _create_network_rule_set(cmd, bypass, default_action)
+    network_acls.ip_rules = []
+    if network_acls_ips:
+        for ip in network_acls_ips:
+            network_acls.ip_rules.append(MHSMIPRule(value=ip))
 
     profile = Profile(cli_ctx=cmd.cli_ctx)
     tenant_id = profile.get_subscription(subscription=cmd.cli_ctx.data.get('subscription_id', None))[_TENANT_ID]
@@ -525,7 +536,7 @@ def create_hsm(cmd, client,
                                       enable_purge_protection=enable_purge_protection,
                                       soft_delete_retention_in_days=retention_days,
                                       initial_admin_object_ids=administrators,
-                                      network_acls=_create_network_rule_set(cmd, bypass, default_action),
+                                      network_acls=network_acls,
                                       public_network_access=public_network_access)
     parameters = ManagedHsm(location=location,
                             tags=tags,
@@ -551,6 +562,13 @@ def wait_hsm(client, hsm_name, resource_group_name):
     return client.get(resource_group_name=resource_group_name, name=hsm_name)
 
 
+def wait_vault_or_hsm(cmd, client, resource_group_name, vault_name=None, hsm_name=None):
+    if vault_name:
+        return client.get(resource_group_name=resource_group_name, vault_name=vault_name)
+    hsm_client = get_client_factory(ResourceType.MGMT_KEYVAULT, Clients.managed_hsms)(cmd.cli_ctx, None)
+    return hsm_client.get(resource_group_name=resource_group_name, name=hsm_name)
+
+
 def create_vault(cmd, client,  # pylint: disable=too-many-locals, too-many-statements
                  resource_group_name, vault_name, location=None, sku=None,
                  enabled_for_deployment=None,
@@ -824,8 +842,19 @@ def set_policy(cmd, client, resource_group_name, vault_name,
                            tags=vault.tags,
                            properties=vault.properties))
 
+def add_network_rule_for_vault_or_hsm(cmd, client, resource_group_name, vault_name=None, hsm_name=None,
+                                      ip_address=None, subnet=None, vnet_name=None, no_wait=False):
+    if vault_name:
+        return add_vault_network_rule(cmd, client, resource_group_name, vault_name,
+                                      ip_address=ip_address, subnet=subnet, vnet_name=vnet_name, no_wait=no_wait)
+    if hsm_name:
+        hsm_client = get_client_factory(ResourceType.MGMT_KEYVAULT, Clients.managed_hsms)(cmd.cli_ctx, None)
+        return add_hsm_network_rule(cmd, hsm_client, resource_group_name, hsm_name,
+                                    ip_address=ip_address, subnet=subnet, vnet_name=vnet_name, no_wait=no_wait)
 
-def add_network_rule(cmd, client, resource_group_name, vault_name, ip_address=None, subnet=None,
+
+
+def add_vault_network_rule(cmd, client, resource_group_name, vault_name, ip_address=None, subnet=None,
                      vnet_name=None, no_wait=False):  # pylint: disable=unused-argument
     """ Add a network rule to the network ACLs for a Key Vault. """
 
@@ -884,8 +913,57 @@ def add_network_rule(cmd, client, resource_group_name, vault_name, ip_address=No
                            properties=vault.properties))
 
 
-def remove_network_rule(cmd, client, resource_group_name, vault_name, ip_address=None, subnet=None,
-                        vnet_name=None, no_wait=False):  # pylint: disable=unused-argument
+def add_hsm_network_rule(cmd, client, resource_group_name, hsm_name,
+                         ip_address=None, subnet=None, vnet_name=None, no_wait=False):
+    if subnet or vnet_name:
+        raise InvalidArgumentValueError("--subnet and --vnet-name are not supported for MHSM yet")
+
+    hsm = client.get(resource_group_name=resource_group_name, name=hsm_name)
+    hsm.properties.network_acls = hsm.properties.network_acls or _create_network_rule_set(cmd)
+    rules = hsm.properties.network_acls
+    IPRule = cmd.get_models('MHSMIPRule', resource_type=ResourceType.MGMT_KEYVAULT)
+
+    to_update = False
+    if ip_address:
+        rules.ip_rules = rules.ip_rules or []
+        # if the rule already exists, don't add again
+        for ip in ip_address:
+            to_modify = True
+            for x in rules.ip_rules:
+                existing_ip_network = ip_network(x.value)
+                new_ip_network = ip_network(ip)
+                if new_ip_network.overlaps(existing_ip_network):
+                    logger.warning("IP/CIDR %s overlaps with %s, which exists already. Not adding duplicates.",
+                                   ip, x.value)
+                    to_modify = False
+                    break
+            if to_modify:
+                rules.ip_rules.append(IPRule(value=ip))
+                to_update = True
+
+    if not to_update:
+        return hsm
+
+    client.begin_create_or_update(resource_group_name=resource_group_name,
+                                  name=hsm_name,
+                                  parameters=hsm,
+                                  polling=not no_wait).result()
+    if not no_wait:
+        return client.get(resource_group_name=resource_group_name, name=hsm_name)
+
+
+def remove_network_rule_for_vault_or_hsm(cmd, client, resource_group_name, vault_name=None, hsm_name=None,
+                                         ip_address=None, subnet=None, vnet_name=None, no_wait=False):
+    if vault_name:
+        return remove_vault_network_rule(cmd, client, resource_group_name, vault_name,
+                                         ip_address=ip_address, subnet=subnet, vnet_name=vnet_name, no_wait=no_wait)
+    hsm_client = get_client_factory(ResourceType.MGMT_KEYVAULT, Clients.managed_hsms)(cmd.cli_ctx, None)
+    return remove_hsm_network_rule(cmd, hsm_client, resource_group_name, hsm_name,
+                                   ip_address=ip_address, subnet=subnet, vnet_name=vnet_name, no_wait=no_wait)
+
+
+def remove_vault_network_rule(cmd, client, resource_group_name, vault_name, ip_address=None,
+                              subnet=None, vnet_name=None, no_wait=False):  # pylint: disable=unused-argument
     """ Remove a network rule from the network ACLs for a Key Vault. """
 
     VaultCreateOrUpdateParameters = cmd.get_models('VaultCreateOrUpdateParameters',
@@ -927,10 +1005,48 @@ def remove_network_rule(cmd, client, resource_group_name, vault_name, ip_address
                            properties=vault.properties))
 
 
-def list_network_rules(cmd, client, resource_group_name, vault_name):  # pylint: disable=unused-argument
-    """ List the network rules from the network ACLs for a Key Vault. """
-    vault = client.get(resource_group_name=resource_group_name, vault_name=vault_name)
-    return vault.properties.network_acls
+def remove_hsm_network_rule(cmd, client, resource_group_name, hsm_name,
+                            ip_address=None, subnet=None, vnet_name=None, no_wait=False):
+    if subnet or vnet_name:
+        raise InvalidArgumentValueError("--subnet and --vnet-name are not supported for MHSM yet")
+
+    hsm = client.get(resource_group_name=resource_group_name, name=hsm_name)
+
+    if not hsm.properties.network_acls:
+        return hsm
+
+    rules = hsm.properties.network_acls
+
+    to_modify = False
+
+    if ip_address and rules.ip_rules:
+        to_remove = [ip_network(x) for x in ip_address]
+        new_rules = list(filter(lambda x: all(ip_network(x.value) != i for i in to_remove), rules.ip_rules))
+        to_modify |= len(new_rules) != len(rules.ip_rules)
+        if to_modify:
+            rules.ip_rules = new_rules
+
+    # if we didn't modify the network rules just return the vault as is
+    if not to_modify:
+        return hsm
+
+    # otherwise update
+    client.begin_create_or_update(resource_group_name=resource_group_name,
+                                  name=hsm_name,
+                                  parameters=hsm,
+                                  polling=not no_wait).result()
+    if not no_wait:
+        return client.get(resource_group_name=resource_group_name, name=hsm_name)
+
+
+def list_network_rules(cmd, client, resource_group_name, vault_name=None, hsm_name=None):  # pylint: disable=unused-argument
+    """ List the network rules from the network ACLs for a Key Vault or a Managed HSM. """
+    if vault_name:
+        vault = client.get(resource_group_name=resource_group_name, vault_name=vault_name)
+        return vault.properties.network_acls
+    hsm_client = get_client_factory(ResourceType.MGMT_KEYVAULT, Clients.managed_hsms)(cmd.cli_ctx, None)
+    hsm = hsm_client.get(resource_group_name, name=hsm_name)
+    return hsm.properties.network_acls
 
 
 def delete_policy(cmd, client, resource_group_name, vault_name,

src/azure-cli/azure/cli/command_modules/keyvault/tests/latest/policy3.json

@@ -1 +1 @@
-{"attributes": {"created": "2025-04-21T06:25:00+00:00", "enabled": true, "expires": null, "notBefore": null, "recoveryLevel": null, "updated": "2025-04-21T06:25:00+00:00"}, "id": "https://cli-test-kv-ct-s3vygjfgb.vault.azure.net/certificates/cert2/policy", "issuerParameters": {"certificateTransparency": null, "certificateType": null, "name": "Self"}, "keyProperties": {"curve": null, "exportable": true, "keySize": 2048, "keyType": "RSA", "reuseKey": false}, "lifetimeActions": [{"action": {"actionType": "AutoRenew"}, "trigger": {"daysBeforeExpiry": null, "lifetimePercentage": 90}}], "secretProperties": {"contentType": "application/x-pkcs12"}, "x509CertificateProperties": {"ekus": ["1.3.6.1.5.5.7.3.1", "1.3.6.1.5.5.7.3.2"], "keyUsage": ["digitalSignature", "keyAgreement", "keyCertSign", "keyEncipherment", "nonRepudiation"], "subject": "C=US, ST=WA, L=Redmon, O=Test Noodle, OU=TestNugget, CN=www.mytestdomain.com", "subjectAlternativeNames": {"dnsNames": ["1", "2"], "emails": null, "upns": null}, "validityInMonths": 60}}
+{"attributes": {"created": "2025-09-08T08:20:00+00:00", "enabled": true, "expires": null, "notBefore": null, "recoveryLevel": null, "updated": "2025-09-08T08:20:00+00:00"}, "id": "https://cli-test-kv-ct-ee4rkkdkg.vault.azure.net/certificates/cert2/policy", "issuerParameters": {"certificateTransparency": null, "certificateType": null, "name": "Self"}, "keyProperties": {"curve": null, "exportable": true, "keySize": 2048, "keyType": "RSA", "reuseKey": false}, "lifetimeActions": [{"action": {"actionType": "AutoRenew"}, "trigger": {"daysBeforeExpiry": null, "lifetimePercentage": 90}}], "secretProperties": {"contentType": "application/x-pkcs12"}, "x509CertificateProperties": {"ekus": ["1.3.6.1.5.5.7.3.1", "1.3.6.1.5.5.7.3.2"], "keyUsage": ["digitalSignature", "keyAgreement", "keyCertSign", "keyEncipherment", "nonRepudiation"], "subject": "C=US, ST=WA, L=Redmon, O=Test Noodle, OU=TestNugget, CN=www.mytestdomain.com", "subjectAlternativeNames": {"dnsNames": ["1", "2"], "emails": null, "upns": null}, "validityInMonths": 60}}