[Batch] Upgrade api-version to v2022-06-01, adding 'batch account network-profile' and 'batch account network-profile network-rule'

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

@@ -74,6 +74,42 @@
 long-summary: List the endpoints that a Batch Compute Node under this Batch Account may call as part of Batch service administration. If you are deploying a Pool inside of a virtual network that you specify, you must make sure your network allows outbound access to these endpoints. Failure to allow access to these endpoints may cause Batch to mark the affected nodes as unusable. For more information about creating a pool inside of a virtual network, see https://docs.microsoft.com/azure/batch/batch-virtual-network."
 """
 
+helps['batch account network-profile'] = """
+type: group
+short-summary: Manage Batch account Network profiles.
+"""
+
+helps['batch account network-profile show'] = """
+type: command
+short-summary: Get information about the Network profile for Batch account.
+"""
+
+helps['batch account network-profile set'] = """
+type: command
+short-summary: Set the Network profile for Batch account.
+"""
+
+helps['batch account network-profile network-rule'] = """
+type: group
+short-summary: Manage Batch account Network rules in Network Profile.
+"""
+
+helps['batch account network-profile network-rule list'] = """
+type: command
+short-summary: List the Network rules from a Network Profile.
+"""
+
+helps['batch account network-profile network-rule add'] = """
+type: command
+short-summary: Add a Network rule from a Network Profile.
+"""
+
+helps['batch account network-profile network-rule delete'] = """
+type: command
+short-summary: Delete a Network rule from a Network Profile.
+"""
+
+
 helps['batch application'] = """
 type: group
 short-summary: Manage Batch applications.

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

@@ -11,7 +11,8 @@
     AccountKeyType,
     KeySource,
     PublicNetworkAccessType,
-    ResourceIdentityType)
+    ResourceIdentityType,
+    EndpointAccessDefaultAction)
 from azure.batch.models import ComputeNodeDeallocationOption
 
 from azure.cli.core.commands.parameters import (
@@ -71,15 +72,44 @@ def load_arguments(self, _):
         c.argument('encryption_key_source', help='Part of the encryption configuration for the Batch account. Type of the key source. Can be either Microsoft.Batch or Microsoft.KeyVault', arg_type=get_enum_type(KeySource))
         c.argument('encryption_key_identifier', help='Part of the encryption configuration for the Batch account. '
                                                      'Full path to the versioned secret. Example https://mykeyvault.vault.azure.net/keys/testkey/6e34a81fef704045975661e297a4c053.')
-        c.argument('identity_type', help="The type of identity used for the Batch account. Possible values include: 'SystemAssigned', 'None'.", arg_type=get_enum_type(ResourceIdentityType))
+        c.argument('identity_type', help="The type of identity used for the Batch account.", arg_type=get_enum_type(ResourceIdentityType))
+        c.argument('mi_user_assigned', help="The managed User identity fully qualified resource Id. Must be in conjunction with identity_type=UserAssigned")
         c.ignore('keyvault_url')
 
     with self.argument_context('batch account set') as c:
         c.argument('tags', tags_type)
         c.argument('storage_account', help='The storage account name or resource ID to be used for auto storage.', validator=storage_account_id)
         c.argument('encryption_key_source', help='Part of the encryption configuration for the Batch account. Type of the key source. Can be either Microsoft.Batch or Microsoft.KeyVault')
+        c.argument('public_network_access', help="The network access type for accessing Azure Batch account. Values can either be enabled or disabled.", arg_type=get_enum_type(PublicNetworkAccessType))
         c.argument('encryption_key_identifier', help='Part of the encryption configuration for the Batch account. Full path to the versioned secret. Example https://mykeyvault.vault.azure.net/keys/testkey/6e34a81fef704045975661e297a4c053.')
-        c.argument('identity_type', help="The type of identity used for the Batch account. Possible values include: 'SystemAssigned', 'None'.", arg_type=get_enum_type(ResourceIdentityType))
+        c.argument('identity_type', help="The type of identity used for the Batch account.", arg_type=get_enum_type(ResourceIdentityType))
+        c.argument('mi_user_assigned', help="The managed User identity fully qualified resource Id. Must be in conjunction with identity_type=UserAssigned")
+
+    with self.argument_context('batch account network-profile show') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.', required=False)
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.', required=False)
+
+    with self.argument_context('batch account network-profile set') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.', required=False)
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.', required=False)
+        c.argument('profile', help="Network profile to set. Allowed values include: 'BatchAccount', 'NodeManagement'")
+        c.argument('default_action', help="Default action for endpoint access. It is only applicable when publicNetworkAccess is enabled. Possible values include: 'Allow', 'Deny'", arg_type=get_enum_type(EndpointAccessDefaultAction))
+
+    with self.argument_context('batch account network-profile network-rule list') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.', required=False)
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.', required=False)
+
+    with self.argument_context('batch account network-profile network-rule add') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.', required=False)
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.', required=False)
+        c.argument('profile', help="Network profile to set. Allowed values include: 'BatchAccount', 'NodeManagement'")
+        c.argument('ip_address', help='IPv4 address or CIDR range.')
+
+    with self.argument_context('batch account network-profile network-rule delete') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.', required=False)
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.', required=False)
+        c.argument('profile', help="Network profile to set. Allowed values include: 'BatchAccount', 'NodeManagement'")
+        c.argument('ip_address', help='IPv4 address or CIDR range.')
 
     with self.argument_context('batch account keys renew') as c:
         c.argument('resource_group_name', resource_group_name_type,

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

@@ -62,6 +62,15 @@ def get_data_factory(name):
         g.custom_command('keys renew', 'renew_accounts_keys', table_transformer=account_keys_renew_table_format)
         g.command('outbound-endpoints', 'list_outbound_network_dependencies_endpoints')
 
+    with self.command_group('batch account network-profile', get_mgmt_type('batch_account'), client_factory=get_mgmt_factory('batch_account')) as g:
+        g.custom_show_command('show', 'get_network_profile')
+        g.custom_command('set', 'update_network_profile')
+
+    with self.command_group('batch account network-profile network-rule', get_mgmt_type('batch_account'), client_factory=get_mgmt_factory('batch_account')) as g:
+        g.custom_show_command('list', 'list_network_rules')
+        g.custom_command('add', 'add_network_rule')
+        g.custom_command('delete', 'delete_network_rule', confirmation=True)
+
     with self.command_group('batch application', get_mgmt_type('application'), client_factory=get_mgmt_factory('application')) as g:
         g.command('list', 'list', table_transformer=application_list_table_format)
         g.show_command('show', 'get')

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

@@ -16,7 +16,9 @@
                                      AutoStorageBaseProperties, ActivateApplicationPackageParameters,
                                      Application, EncryptionProperties,
                                      KeyVaultProperties, BatchAccountIdentity,
-                                     BatchAccountRegenerateKeyParameters)
+                                     BatchAccountRegenerateKeyParameters, PublicNetworkAccessType, BatchAccount,
+                                     NetworkProfile, EndpointAccessProfile, IPRule, EndpointAccessDefaultAction,
+                                     ResourceIdentityType, UserAssignedIdentities)
 from azure.mgmt.batch.operations import (ApplicationPackageOperations)
 
 from azure.batch.models import (CertificateAddParameter, PoolStopResizeOptions, PoolResizeParameter,
@@ -77,10 +79,20 @@ def get_account(cmd, client, resource_group_name=None, account_name=None):
 def create_account(client,
                    resource_group_name, account_name, location, tags=None, storage_account=None,
                    keyvault=None, keyvault_url=None, no_wait=False, public_network_access=None,
-                   encryption_key_source=None, encryption_key_identifier=None, identity_type=None):
+                   encryption_key_source=None, encryption_key_identifier=None, identity_type=None,
+                   mi_user_assigned=None):
     properties = AutoStorageBaseProperties(storage_account_id=storage_account) \
         if storage_account else None
-    identity = BatchAccountIdentity(type=identity_type) if identity_type else None
+
+    identity = None
+    if identity_type and identity_type == ResourceIdentityType.SYSTEM_ASSIGNED:
+        identity = BatchAccountIdentity(type=identity_type)
+
+    if mi_user_assigned and identity_type and identity_type == ResourceIdentityType.USER_ASSIGNED:
+        useridentity = UserAssignedIdentities()
+        my_dict = {mi_user_assigned: useridentity}
+        identity = BatchAccountIdentity(type=identity_type, user_assigned_identities=my_dict)
+
     if (encryption_key_source and
             encryption_key_source.lower() == "microsoft.keyvault" and not encryption_key_identifier):
         raise ValueError("The --encryption-key-identifier property is required when "
@@ -106,8 +118,9 @@ def create_account(client,
 
 @transfer_doc(AutoStorageBaseProperties)
 def update_account(client, resource_group_name, account_name,
-                   tags=None, storage_account=None, encryption_key_source=None,
-                   encryption_key_identifier=None, identity_type=None):
+                   tags=None, storage_account=None, encryption_key_source=None, public_network_access=None,
+                   encryption_key_identifier=None, identity_type=None, mi_user_assigned=None):
+
     properties = AutoStorageBaseProperties(storage_account_id=storage_account) \
         if storage_account else None
     if (encryption_key_source and
@@ -120,12 +133,23 @@ def update_account(client, resource_group_name, account_name,
     encryption = EncryptionProperties(
         key_source=encryption_key_source,
         encryption_key_identifier=encryption_key_identifier) if encryption_key_source else None
-    identity = BatchAccountIdentity(type=identity_type) if identity_type else None
+
+    identity = None
+    if identity_type and identity_type == ResourceIdentityType.SYSTEM_ASSIGNED:
+        identity = BatchAccountIdentity(type=identity_type)
+
+    if mi_user_assigned and identity_type and identity_type == ResourceIdentityType.USER_ASSIGNED:
+        useridentity = UserAssignedIdentities()
+        my_dict = {mi_user_assigned: useridentity}
+        identity = BatchAccountIdentity(type=identity_type, user_assigned_identities=my_dict)
+
     parameters = BatchAccountUpdateParameters(
         tags=tags,
         encryption=encryption,
         identity=identity,
+        public_network_access=public_network_access,
         auto_storage=properties)
+
     return client.update(resource_group_name=resource_group_name,
                          account_name=account_name,
                          parameters=parameters)
@@ -173,12 +197,130 @@ def login_account(cmd, client, resource_group_name, account_name, shared_key_aut
 
 
 def renew_accounts_keys(client, resource_group_name, account_name, key_name=None):
+
     parameters = BatchAccountRegenerateKeyParameters(key_name=key_name)
 
     return client.regenerate_key(resource_group_name=resource_group_name,
                                  account_name=account_name, parameters=parameters)
 
 
+def get_network_profile(cmd, client, resource_group_name=None, account_name=None):
+
+    batch_account: BatchAccount = get_account(cmd, client, resource_group_name, account_name)
+
+    return batch_account.network_profile
+
+
+def update_network_profile(cmd, client, resource_group_name, account_name,
+                           profile=None, default_action=None):
+
+    batch_account: BatchAccount = get_account(cmd, client, resource_group_name, account_name)
+
+    # we want to use the existing network_profile if it exists
+    networkprofile = batch_account.network_profile
+    if networkprofile is None:
+        networkprofile = NetworkProfile()
+
+    if profile.lower() == "batchaccount":
+        if networkprofile.account_access is None:
+            networkprofile.account_access = EndpointAccessProfile(default_action=default_action)
+
+        networkprofile.account_access.default_action = default_action
+
+    if profile.lower() == "nodemanagement":
+        if networkprofile.node_management_access is None:
+            networkprofile.node_management_access = EndpointAccessProfile(default_action=default_action)
+
+        networkprofile.node_management_access.default_action = default_action
+
+    parameters = BatchAccountUpdateParameters(network_profile=networkprofile)
+
+    return client.update(resource_group_name=resource_group_name,
+                         account_name=account_name,
+                         parameters=parameters)
+
+
+def list_network_rules(cmd, client, resource_group_name, account_name):
+    batch_account: BatchAccount = get_account(cmd, client, resource_group_name, account_name)
+
+    return batch_account.network_profile
+
+
+def add_network_rule(cmd, client, resource_group_name, account_name, profile=None, ip_address=None):
+    batch_account: BatchAccount = get_account(cmd, client, resource_group_name, account_name)
+
+    # we want to use the existing network_profile if it exists, else build one up
+    networkprofile = batch_account.network_profile
+    if networkprofile is None:
+        networkprofile = NetworkProfile()
+
+    if profile.lower() == "batchaccount":
+        if networkprofile.account_access is None:
+            networkprofile.account_access = EndpointAccessProfile(
+                default_action=EndpointAccessDefaultAction.ALLOW)
+        if networkprofile.account_access.ip_rules is None:
+            networkprofile.account_access.ip_rules = []
+        networkprofile.account_access.ip_rules.append(IPRule(value=ip_address))
+
+    if profile.lower() == "nodemanagement":
+        if networkprofile.node_management_access is None:
+            networkprofile.node_management_access = EndpointAccessProfile(
+                default_action=EndpointAccessDefaultAction.ALLOW)
+        if networkprofile.node_management_access.ip_rules is None:
+            networkprofile.node_management_access.ip_rules = []
+        networkprofile.node_management_access.ip_rules.append(IPRule(value=ip_address))
+
+    # Not sure if i want to enable public_network_access as part of this
+    parameters = BatchAccountUpdateParameters(public_network_access=PublicNetworkAccessType.ENABLED,
+                                              network_profile=networkprofile)
+
+    client.update(resource_group_name=resource_group_name,
+                  account_name=account_name,
+                  parameters=parameters)
+
+    return list_network_rules(cmd, client, resource_group_name, account_name)
+
+
+def delete_network_rule(cmd, client, resource_group_name, account_name, profile=None, ip_address=None):
+    batch_account: BatchAccount = get_account(cmd, client, resource_group_name, account_name)
+
+    # we want to use the existing network_profile if it exists, else build one up
+    networkprofile = batch_account.network_profile
+    if networkprofile is None:
+        networkprofile = NetworkProfile()
+
+    if profile.lower() == "batchaccount":
+        if networkprofile.account_access is None:
+            networkprofile.account_access = EndpointAccessProfile(
+                default_action=EndpointAccessDefaultAction.ALLOW)
+        if networkprofile.account_access.ip_rules is None:
+            networkprofile.account_access.ip_rules = []
+
+        for iprule in networkprofile.account_access.ip_rules:
+            if iprule.value == ip_address:
+                networkprofile.account_access.ip_rules.remove(iprule)
+
+    if profile.lower() == "nodemanagement":
+        if networkprofile.node_management_access is None:
+            networkprofile.node_management_access = EndpointAccessProfile(
+                default_action=EndpointAccessDefaultAction.ALLOW)
+        if networkprofile.node_management_access.ip_rules is None:
+            networkprofile.node_management_access.ip_rules = []
+        for iprule in networkprofile.node_management_access.ip_rules:
+            if iprule.value == ip_address:
+                networkprofile.node_management_access.ip_rules.remove(iprule)
+
+    # Not sure if i want to enable public_network_access as part of this
+    parameters = BatchAccountUpdateParameters(public_network_access=PublicNetworkAccessType.ENABLED,
+                                              network_profile=networkprofile)
+
+    client.update(resource_group_name=resource_group_name,
+                  account_name=account_name,
+                  parameters=parameters)
+
+    return list_network_rules(cmd, client, resource_group_name, account_name)
+
+
 @transfer_doc(Application)
 def update_application(client,
                        resource_group_name, account_name, application_name, allow_updates=None,