1# --------------------------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for license information.
4# --------------------------------------------------------------------------------------------
5
6from argcomplete.completers import FilesCompleter
7
8from knack.arguments import CLIArgumentType
9
10from azure.mgmt.batch.models import (
11    AccountKeyType,
12    KeySource,
13    PublicNetworkAccessType,
14    ResourceIdentityType)
15from azure.batch.models import ComputeNodeDeallocationOption
16
17from azure.cli.core.commands.parameters import (
18    tags_type,
19    get_location_type,
20    resource_group_name_type,
21    get_resource_name_completion_list,
22    file_type,
23    get_enum_type)
24
25from azure.cli.command_modules.batch._completers import load_supported_images
26from azure.cli.command_modules.batch._validators import (
27    application_enabled,
28    application_package_reference_format,
29    certificate_reference_format,
30    datetime_format,
31    disk_encryption_configuration_format,
32    environment_setting_format,
33    keyvault_id,
34    metadata_item_format,
35    resource_file_format,
36    storage_account_id,
37    validate_cert_file,
38    validate_cert_settings,
39    validate_client_parameters,
40    validate_json_file,
41    validate_pool_resize_parameters)
42
43
44# pylint: disable=line-too-long, too-many-statements
45def load_arguments(self, _):
46    batch_name_type = CLIArgumentType(
47        help='Name of the Batch account.',
48        options_list=('--account-name',),
49        completer=get_resource_name_completion_list('Microsoft.Batch/batchAccounts'),
50        id_part=None)
51
52    with self.argument_context('batch') as c:
53        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group', completer=None, validator=None, required=True)
54
55    with self.argument_context('batch account') as c:
56        c.argument('account_name', batch_name_type, options_list=('--name', '-n'))
57
58    with self.argument_context('batch account show') as c:
59        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)
60        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)
61
62    with self.argument_context('batch account list') as c:
63        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group', required=False)
64
65    with self.argument_context('batch account create') as c:
66        c.argument('location', get_location_type(self.cli_ctx), help='The region in which to create the account.')
67        c.argument('tags', tags_type, help="Space-separated tags in 'key[=value]' format.")
68        c.argument('storage_account', help='The storage account name or resource ID to be used for auto storage.', validator=storage_account_id)
69        c.argument('keyvault', help='The KeyVault name or resource ID to be used for an account with a pool allocation mode of \'User Subscription\'.', validator=keyvault_id)
70        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))
71        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))
72        c.argument('encryption_key_identifier', help='Part of the encryption configuration for the Batch account. '
73                                                     'Full path to the versioned secret. Example https://mykeyvault.vault.azure.net/keys/testkey/6e34a81fef704045975661e297a4c053.')
74        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))
75        c.ignore('keyvault_url')
76
77    with self.argument_context('batch account set') as c:
78        c.argument('tags', tags_type)
79        c.argument('storage_account', help='The storage account name or resource ID to be used for auto storage.', validator=storage_account_id)
80        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')
81        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.')
82        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))
83
84    with self.argument_context('batch account keys renew') as c:
85        c.argument('resource_group_name', resource_group_name_type,
86                   help='Name of the resource group. If not specified will display currently set account.',
87                   required=False)
88        c.argument('account_name', batch_name_type, options_list=('--name', '-n'),
89                   help='Name of the batch account to show. If not specified will display currently set account.',
90                   required=False)
91        c.argument('key_name', arg_type=get_enum_type(AccountKeyType), help='Name of the batch account key.')
92
93    with self.argument_context('batch account login') as c:
94        c.argument('shared_key_auth', action='store_true', help='Using Shared Key authentication, if not specified, it will use Azure Active Directory authentication.')
95        c.argument('show', action='store_true', help='Display the credential information for the Batch account.')
96
97    with self.argument_context('batch application set') as c:
98        c.argument('application_name', options_list=('--application-name',), help="The name of the application.")
99        c.argument('allow_updates', options_list=('--allow-updates',), help="Specify to indicate whether packages within the application may be overwritten using the same version string. Specify either 'true' or 'false' to update the property.")
100
101    with self.argument_context('batch application create') as c:
102        c.argument('allow_updates', options_list=('--allow-updates',), action="store_true", help="Specify to indicate whether packages within the application may be overwritten using the same version string. True if flag present.")
103
104    for command in ['create', 'activate']:
105        with self.argument_context('batch application package {}'.format(command)) as c:
106            c.argument('package_file', type=file_type, help='The path of the application package in zip format', completer=FilesCompleter())
107            c.argument('application_name', options_list=('--application-name',), help="The name of the application.")
108            c.argument('version_name', options_list=('--version-name',), help="The version name of the application.")
109            c.argument('f_ormat', options_list=('--format',), help="The format of the application package binary file.")
110
111    with self.argument_context('batch location quotas show') as c:
112        c.argument('location_name', get_location_type(self.cli_ctx), help='The region for which to display the Batch service quotas.')
113
114    with self.argument_context('batch location list-skus') as c:
115        c.argument('location_name', get_location_type(self.cli_ctx), help='The region for which to display the available Batch VM SKUs.')
116
117    for command in ['list', 'show', 'create', 'set', 'delete', 'package']:
118        with self.argument_context('batch application {}'.format(command)) as c:
119            c.argument('account_name', batch_name_type, options_list=('--name', '-n'), validator=application_enabled)
120
121    # TODO: Refactor so the help text can be extracted automatically
122    with self.argument_context('batch pool resize') as c:
123        c.argument('if_modified_since', help='The operation will be performed only if the resource has been modified since the specified timestamp.', type=datetime_format, arg_group='Pre-condition and Query')
124        c.argument('if_unmodified_since', help='The operation will not be performed only if the resource has been modified since the specified timestamp.', type=datetime_format, arg_group='Pre-condition and Query')
125        c.argument('if_match', help='The operation will be performed only if the resource\'s current ETag exactly matches the specified value.', arg_group='Pre-condition and Query')
126        c.argument('if_none_match', help='The operation will not be performed only if the resource\'s current ETag exactly matches the specified value.', arg_group='Pre-condition and Query')
127        c.argument('pool_id', help='The ID of the pool.')
128        c.argument('abort', action='store_true', help='Stop the pool resize operation.', validator=validate_pool_resize_parameters)
129        c.argument('node_deallocation_option', options_list=('--node-deallocation-option',), help='When nodes may be removed from the pool, if the pool size is decreasing.', arg_type=get_enum_type(ComputeNodeDeallocationOption))
130
131    # TODO: Refactor so the help text can be extracted automatically
132    with self.argument_context('batch pool reset') as c:
133        c.argument('json_file', type=file_type, help='The file containing pool update properties parameter specification in JSON(formatted to match REST API request body). If this parameter is specified, all \'Pool Update Properties Parameter Arguments\' are ignored.', validator=validate_json_file, completer=FilesCompleter())
134        c.argument('pool_id', help='The ID of the pool to update.')
135        c.argument('application_package_references', nargs='+', type=application_package_reference_format, arg_group='Pool')
136        c.argument('certificate_references', nargs='+', type=certificate_reference_format, arg_group='Pool')
137        c.argument('metadata', nargs='+', type=metadata_item_format, arg_group='Pool')
138        c.argument('start_task_command_line', arg_group='Pool: Start Task',
139                   help='The command line of the start task. The command line does not run under a shell, and therefore cannot take advantage of shell features such as environment variable expansion. If you want to take advantage of such features, you should invoke the shell in the command line, for example using "cmd /c MyCommand" in Windows or "/bin/sh -c MyCommand" in Linux.')
140        c.argument('start_task_wait_for_success', action='store_true', arg_group='Pool: Start Task',
141                   help='Whether the Batch service should wait for the start task to complete successfully (that is, to exit with exit code 0) before scheduling any tasks on the compute node. True if flag present, otherwise defaults to False.')
142        c.argument('start_task_max_task_retry_count', arg_group='Pool: Start Task',
143                   help='The maximum number of times the task may be retried.')
144        c.argument('start_task_environment_settings', nargs='+', type=environment_setting_format, arg_group='Pool: Start Task',
145                   help='A list of environment variable settings for the start task. Space-separated values in \'key=value\' format.')
146
147    with self.argument_context('batch job list') as c:
148        c.argument('filter', help=' An OData $filter clause.', arg_group='Pre-condition and Query')
149        c.argument('select', help=' An OData $select clause.', arg_group='Pre-condition and Query')
150        c.argument('expand', help=' An OData $expand clause.', arg_group='Pre-condition and Query')
151        c.argument('job_schedule_id', help='The ID of the job schedule from which you want to get a list of jobs. If omitted, lists all jobs in the account.')
152
153    for command in ['job create', 'job set', 'job reset', 'job-schedule create', 'job-schedule set', 'job-schedule reset']:
154        with self.argument_context('batch {}'.format(command)) as c:
155            c.argument('pool_id', options_list=('--pool-id',), help='The id of an existing pool. All the tasks of the job will run on the specified pool.')
156
157    with self.argument_context('batch pool create') as c:
158        c.argument('os_family', arg_type=get_enum_type(['2', '3', '4', '5', '6']))
159        c.argument('auto_scale_formula', help='A formula for the desired number of compute nodes in the pool. The formula is checked for validity before the pool is created. If the formula is not valid, the Batch service rejects the request with detailed error information. For more information about specifying this formula, see https://azure.microsoft.com/documentation/articles/batch-automatic-scaling/.')
160        c.extra('disk_encryption_targets',
161                arg_group="Pool: Virtual Machine Configuration",
162                help='A space separated list of DiskEncryptionTargets. current possible values include OsDisk and TemporaryDisk.', type=disk_encryption_configuration_format)
163        c.extra('image', completer=load_supported_images, arg_group="Pool: Virtual Machine Configuration",
164                help="OS image reference. This can be either 'publisher:offer:sku[:version]' format, or a fully qualified ARM image id of the form '/subscriptions/{subscriptionId}/resourceGroups/{resourceGroup}/providers/Microsoft.Compute/images/{imageName}'. If 'publisher:offer:sku[:version]' format, version is optional and if omitted latest will be used. Valid values can be retrieved via 'az batch pool supported-images list'. For example: 'MicrosoftWindowsServer:WindowsServer:2012-R2-Datacenter:latest'")
165
166    with self.argument_context('batch certificate') as c:
167        c.argument('thumbprint', help='The certificate thumbprint.')
168        c.argument('password', help='The password to access the certificate\'s private key.')
169        c.argument('certificate_file', type=file_type, help='The certificate file: cer file or pfx file.', validator=validate_cert_file, completer=FilesCompleter())
170        c.argument('abort', action='store_true', help='Cancel the failed certificate deletion operation.')
171
172    with self.argument_context('batch certificate show') as c:
173        c.argument('thumbprint', help='The certificate thumbprint.', validator=validate_cert_settings)
174
175    with self.argument_context('batch task create') as c:
176        c.argument('json_file', type=file_type, help='The file containing the task(s) to create in JSON(formatted to match REST API request body). When submitting multiple tasks, accepts either an array of tasks or a TaskAddCollectionParamater. If this parameter is specified, all other parameters are ignored.', validator=validate_json_file, completer=FilesCompleter())
177        c.argument('application_package_references', nargs='+', help='The space-separated list of IDs specifying the application packages to be installed. Space-separated application IDs with optional version in \'id[#version]\' format.', type=application_package_reference_format)
178        c.argument('job_id', help='The ID of the job containing the task.')
179        c.argument('task_id', help='The ID of the task.')
180        c.argument('command_line', help='The command line of the task. The command line does not run under a shell, and therefore cannot take advantage of shell features such as environment variable expansion. If you want to take advantage of such features, you should invoke the shell in the command line, for example using "cmd /c MyCommand" in Windows or "/bin/sh -c MyCommand" in Linux.')
181        c.argument('environment_settings', nargs='+', help='A list of environment variable settings for the task. Space-separated values in \'key=value\' format.', type=environment_setting_format)
182        c.argument('resource_files', nargs='+', help='A list of files that the Batch service will download to the compute node before running the command line. Space-separated resource references in filename=httpurl format, with httpurl being any HTTP url with public access or a SAS url with read access.', type=resource_file_format)
183
184    for item in ['batch certificate delete', 'batch certificate create', 'batch pool resize', 'batch pool reset', 'batch job list', 'batch task create']:
185        with self.argument_context(item) as c:
186            c.extra('account_name', arg_group='Batch Account', validator=validate_client_parameters,
187                    help='The Batch account name. Alternatively, set by environment variable: AZURE_BATCH_ACCOUNT')
188            c.extra('account_key', arg_group='Batch Account',
189                    help='The Batch account key. Alternatively, set by environment variable: AZURE_BATCH_ACCESS_KEY')
190            c.extra('account_endpoint', arg_group='Batch Account',
191                    help='Batch service endpoint. Alternatively, set by environment variable: AZURE_BATCH_ENDPOINT')
192