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
6# TODO refactor out _image_builder commands.
7# i.e something like image_builder/_client_factory image_builder/commands.py image_builder/_params.py
8import os
9import re
10import json
11import traceback
12from enum import Enum
13
14import requests
15
16try:
17    from urllib.parse import urlparse
18except ImportError:
19    from urlparse import urlparse  # pylint: disable=import-error
20
21from knack.util import CLIError
22from knack.log import get_logger
23
24from msrestazure.azure_exceptions import CloudError
25from msrestazure.tools import is_valid_resource_id, resource_id, parse_resource_id
26
27from azure.cli.core.commands import cached_get, cached_put
28from azure.cli.core.commands.client_factory import get_subscription_id
29from azure.cli.core.commands.validators import get_default_location_from_resource_group, validate_tags
30
31from azure.cli.command_modules.vm._client_factory import _compute_client_factory
32from azure.cli.command_modules.vm._validators import _get_resource_id
33
34logger = get_logger(__name__)
35
36
37class _SourceType(Enum):
38    PLATFORM_IMAGE = "PlatformImage"
39    ISO_URI = "ISO"
40    MANAGED_IMAGE = "ManagedImage"
41    SIG_VERSION = "SharedImageVersion"
42
43
44class _DestType(Enum):
45    MANAGED_IMAGE = 1
46    SHARED_IMAGE_GALLERY = 2
47
48
49class ScriptType(Enum):
50    SHELL = "shell"
51    POWERSHELL = "powershell"
52    WINDOWS_RESTART = "windows-restart"
53    WINDOWS_UPDATE = "windows-update"
54    FILE = "file"
55
56
57# region Client Factories
58
59def image_builder_client_factory(cli_ctx, _):
60    from azure.cli.core.commands.client_factory import get_mgmt_service_client
61    from azure.mgmt.imagebuilder import ImageBuilderClient
62    client = get_mgmt_service_client(cli_ctx, ImageBuilderClient)
63    return client
64
65
66def cf_img_bldr_image_templates(cli_ctx, _):
67    return image_builder_client_factory(cli_ctx, _).virtual_machine_image_templates
68
69# endregion
70
71
72def _no_white_space_or_err(words):
73    for char in words:
74        if char.isspace():
75            raise CLIError("Error: White space in {}".format(words))
76
77
78def _require_defer(cmd):
79    use_cache = cmd.cli_ctx.data.get('_cache', False)
80    if not use_cache:
81        raise CLIError("This command requires --defer")
82
83
84def _parse_script(script_str):
85    script_name = script_str
86    script = {"script": script_str, "name": script_name, "type": None}
87    if urlparse(script_str).scheme and "://" in script_str:
88        _, script_name = script_str.rsplit("/", 1)
89        script["name"] = script_name
90        script["is_url"] = True
91    else:
92        raise CLIError("Expected a url, got: {}".format(script_str))
93
94    if script_str.lower().endswith(".sh"):
95        script["type"] = ScriptType.SHELL
96    elif script_str.lower().endswith(".ps1"):
97        script["type"] = ScriptType.POWERSHELL
98
99    return script
100
101
102def _parse_image_destination(cmd, rg, destination, is_shared_image):
103
104    if any([not destination, "=" not in destination]):
105        raise CLIError("Invalid Format: the given image destination {} must contain the '=' delimiter."
106                       .format(destination))
107
108    rid, location = destination.rsplit("=", 1)
109    if not rid or not location:
110        raise CLIError("Invalid Format: destination {} should have format 'destination=location'.".format(destination))
111
112    _no_white_space_or_err(rid)
113
114    result = None
115    if is_shared_image:
116        if not is_valid_resource_id(rid):
117            if "/" not in rid:
118                raise CLIError("Invalid Format: {} must have a shared image gallery name and definition. "
119                               "They must be delimited by a '/'.".format(rid))
120
121            sig_name, sig_def = rid.rsplit("/", 1)
122
123            rid = resource_id(
124                subscription=get_subscription_id(cmd.cli_ctx), resource_group=rg,
125                namespace='Microsoft.Compute',
126                type='galleries', name=sig_name,
127                child_type_1='images', child_name_1=sig_def
128            )
129
130        result = rid, location.split(",")
131    else:
132        if not is_valid_resource_id(rid):
133            rid = resource_id(
134                subscription=get_subscription_id(cmd.cli_ctx),
135                resource_group=rg,
136                namespace='Microsoft.Compute', type='images',
137                name=rid
138            )
139
140        result = rid, location
141
142    return result
143
144
145def _validate_location(location, location_names, location_display_names):
146
147    if ' ' in location:
148        # if display name is provided, attempt to convert to short form name
149        location = next((name for name in location_display_names if name.lower() == location.lower()), location)
150
151    if location.lower() not in [location_name.lower() for location_name in location_names]:
152        raise CLIError("Location {} is not a valid subscription location. "
153                       "Use one from `az account list-locations`.".format(location))
154
155    return location
156
157
158def process_image_template_create_namespace(cmd, namespace):  # pylint: disable=too-many-locals, too-many-branches, too-many-statements
159    if namespace.image_template is not None:
160        return
161
162    from azure.cli.core.commands.parameters import get_subscription_locations
163
164    source = None
165    scripts = []
166
167    # default location to RG location.
168    if not namespace.location:
169        get_default_location_from_resource_group(cmd, namespace)
170
171    # validate tags.
172    validate_tags(namespace)
173
174    # Validate and parse scripts
175    if namespace.scripts:
176        for ns_script in namespace.scripts:
177            scripts.append(_parse_script(ns_script))
178
179    # Validate and parse destination and locations
180    destinations = []
181    subscription_locations = get_subscription_locations(cmd.cli_ctx)
182    location_names = [location.name for location in subscription_locations]
183    location_display_names = [location.display_name for location in subscription_locations]
184
185    if namespace.managed_image_destinations:
186        for dest in namespace.managed_image_destinations:
187            rid, location = _parse_image_destination(cmd, namespace.resource_group_name, dest, is_shared_image=False)
188            location = _validate_location(location, location_names, location_display_names)
189            destinations.append((_DestType.MANAGED_IMAGE, rid, location))
190
191    if namespace.shared_image_destinations:
192        for dest in namespace.shared_image_destinations:
193            rid, locations = _parse_image_destination(cmd, namespace.resource_group_name, dest, is_shared_image=True)
194            locations = [_validate_location(location, location_names, location_display_names) for location in locations]
195            destinations.append((_DestType.SHARED_IMAGE_GALLERY, rid, locations))
196
197    # Validate and parse source image
198    # 1 - check if source is a URN. A urn e.g "Canonical:UbuntuServer:18.04-LTS:latest"
199    urn_match = re.match('([^:]*):([^:]*):([^:]*):([^:]*)', namespace.source)
200    if urn_match:  # if platform image urn
201        source = {
202            'publisher': urn_match.group(1),
203            'offer': urn_match.group(2),
204            'sku': urn_match.group(3),
205            'version': urn_match.group(4),
206            'type': _SourceType.PLATFORM_IMAGE
207        }
208
209        likely_linux = bool("windows" not in source["offer"].lower() and "windows" not in source["sku"].lower())
210
211        logger.info("%s looks like a platform image URN", namespace.source)
212
213    # 2 - check if a fully-qualified ID (assumes it is an image ID)
214    elif is_valid_resource_id(namespace.source):
215
216        parsed = parse_resource_id(namespace.source)
217        image_type = parsed.get('type')
218        image_resource_type = parsed.get('type')
219
220        if not image_type:
221            pass
222
223        elif image_type.lower() == 'images':
224            source = {
225                'image_id': namespace.source,
226                'type': _SourceType.MANAGED_IMAGE
227            }
228            logger.info("%s looks like a managed image id.", namespace.source)
229
230        elif image_type == "galleries" and image_resource_type:
231            source = {
232                'image_version_id': namespace.source,
233                'type': _SourceType.SIG_VERSION
234            }
235            logger.info("%s looks like a shared image version id.", namespace.source)
236
237    # 3 - check if source is a Redhat iso uri. If so a checksum must be provided.
238    elif urlparse(namespace.source).scheme and "://" in namespace.source and ".iso" in namespace.source.lower():
239        if not namespace.checksum:
240            raise CLIError("Must provide a checksum for source uri: {}".format(namespace.source))
241        source = {
242            'source_uri': namespace.source,
243            'sha256_checksum': namespace.checksum,
244            'type': _SourceType.ISO_URI
245        }
246        likely_linux = True
247
248        logger.info("%s looks like a RedHat iso uri.", namespace.source)
249
250    # 4 - check if source is a urn alias from the vmImageAliasDoc endpoint. See "az cloud show"
251    if not source:
252        from azure.cli.command_modules.vm._actions import load_images_from_aliases_doc
253        images = load_images_from_aliases_doc(cmd.cli_ctx)
254        matched = next((x for x in images if x['urnAlias'].lower() == namespace.source.lower()), None)
255        if matched:
256            source = {
257                'publisher': matched['publisher'],
258                'offer': matched['offer'],
259                'sku': matched['sku'],
260                'version': matched['version'],
261                'type': _SourceType.PLATFORM_IMAGE
262            }
263
264        if "windows" not in source["offer"].lower() and "windows" not in source["sku"].lower():
265            likely_linux = True
266
267        logger.info("%s looks like a platform image alias.", namespace.source)
268
269    # 5 - check if source is an existing managed disk image resource
270    if not source:
271        compute_client = _compute_client_factory(cmd.cli_ctx)
272        try:
273            image_name = namespace.source
274            compute_client.images.get(namespace.resource_group_name, namespace.source)
275            namespace.source = _get_resource_id(cmd.cli_ctx, namespace.source, namespace.resource_group_name,
276                                                'images', 'Microsoft.Compute')
277            source = {
278                'image_id': namespace.source,
279                'type': _SourceType.MANAGED_IMAGE
280            }
281
282            logger.info("%s, looks like a managed image name. Using resource ID: %s", image_name, namespace.source)  # pylint: disable=line-too-long
283        except CloudError:
284            pass
285
286    if not source:
287        err = 'Invalid image "{}". Use a valid image URN, managed image name or ID, ISO URI, ' \
288              'or pick a platform image alias from {}.\nSee vm create -h for more information on specifying an image.'\
289            .format(namespace.source, ", ".join([x['urnAlias'] for x in images]))
290        raise CLIError(err)
291
292    for script in scripts:
293        if script["type"] is None:
294            try:
295                script["type"] = ScriptType.SHELL if likely_linux else ScriptType.POWERSHELL
296                logger.info("For script %s, likely linux is %s.", script["script"], likely_linux)
297            except NameError:
298                raise CLIError("Unable to infer the type of script {}.".format(script["script"]))
299
300    namespace.source_dict = source
301    namespace.scripts_list = scripts
302    namespace.destinations_lists = destinations
303
304
305# first argument is `cmd`, but it is unused. Feel free to substitute it in.
306def process_img_tmpl_customizer_add_namespace(cmd, namespace):  # pylint:disable=unused-argument
307
308    if namespace.customizer_type.lower() in [ScriptType.SHELL.value.lower(), ScriptType.POWERSHELL.value.lower()]:  # pylint:disable=no-member, line-too-long
309        if not (namespace.script_url or namespace.inline_script):
310            raise CLIError("A script must be provided if the customizer type is one of: {} {}"
311                           .format(ScriptType.SHELL.value, ScriptType.POWERSHELL.value))
312
313        if namespace.script_url and namespace.inline_script:
314            raise CLIError("Cannot supply both script url and inline script.")
315
316    elif namespace.customizer_type.lower() == ScriptType.WINDOWS_RESTART.value.lower():  # pylint:disable=no-member
317        if namespace.script_url or namespace.inline_script:
318            logger.warning("Ignoring the supplied script as scripts are not used for Windows Restart.")
319
320
321def process_img_tmpl_output_add_namespace(cmd, namespace):
322    from azure.cli.core.commands.parameters import get_subscription_locations
323
324    outputs = [output for output in [namespace.managed_image, namespace.gallery_image_definition, namespace.is_vhd] if output]  # pylint:disable=line-too-long
325
326    if len(outputs) != 1:
327        err = "Supplied outputs: {}".format(outputs)
328        logger.debug(err)
329        raise CLIError("Usage error: must supply exactly one destination type to add. Supplied {}".format(len(outputs)))
330
331    if namespace.managed_image:
332        if not is_valid_resource_id(namespace.managed_image):
333            namespace.managed_image = resource_id(
334                subscription=get_subscription_id(cmd.cli_ctx),
335                resource_group=namespace.resource_group_name,
336                namespace='Microsoft.Compute', type='images',
337                name=namespace.managed_image
338            )
339
340    if namespace.gallery_image_definition:
341        if not is_valid_resource_id(namespace.gallery_image_definition):
342            if not namespace.gallery_name:
343                raise CLIError("Usage error: gallery image definition is a name and not an ID.")
344
345            namespace.gallery_image_definition = resource_id(
346                subscription=get_subscription_id(cmd.cli_ctx), resource_group=namespace.resource_group_name,
347                namespace='Microsoft.Compute',
348                type='galleries', name=namespace.gallery_name,
349                child_type_1='images', child_name_1=namespace.gallery_image_definition
350            )
351
352    if namespace.is_vhd and not namespace.output_name:
353        raise CLIError("Usage error: If --is-vhd is used, a run output name must be provided via --output-name.")
354
355    subscription_locations = get_subscription_locations(cmd.cli_ctx)
356    location_names = [location.name for location in subscription_locations]
357    location_display_names = [location.display_name for location in subscription_locations]
358
359    if namespace.managed_image_location:
360        namespace.managed_image_location = _validate_location(namespace.managed_image_location,
361                                                              location_names, location_display_names)
362
363    if namespace.gallery_replication_regions:
364        processed_regions = []
365        for loc in namespace.gallery_replication_regions:
366            processed_regions.append(_validate_location(loc, location_names, location_display_names))
367        namespace.gallery_replication_regions = processed_regions
368
369    # get default location from resource group
370    if not any([namespace.managed_image_location, namespace.gallery_replication_regions]) and hasattr(namespace, 'location'):  # pylint: disable=line-too-long
371        # store location in namespace.location for use in custom method.
372        get_default_location_from_resource_group(cmd, namespace)
373
374    # validate tags.
375    validate_tags(namespace)
376
377
378# region Custom Commands
379
380def create_image_template(  # pylint: disable=too-many-locals, too-many-branches, too-many-statements, unused-argument
381        cmd, client, resource_group_name, image_template_name, location=None,
382        source_dict=None, scripts_list=None, destinations_lists=None, build_timeout=None, tags=None,
383        source=None, scripts=None, checksum=None, managed_image_destinations=None,
384        shared_image_destinations=None, no_wait=False, image_template=None, identity=None,
385        vm_size=None, os_disk_size=None, vnet=None, subnet=None):
386    from azure.mgmt.imagebuilder.models import (ImageTemplate, ImageTemplateSharedImageVersionSource,
387                                                ImageTemplatePlatformImageSource, ImageTemplateManagedImageSource,
388                                                ImageTemplateShellCustomizer, ImageTemplatePowerShellCustomizer,
389                                                ImageTemplateManagedImageDistributor,
390                                                ImageTemplateSharedImageDistributor, ImageTemplateIdentity,
391                                                ImageTemplateIdentityUserAssignedIdentitiesValue,
392                                                ImageTemplateVmProfile, VirtualNetworkConfig)
393
394    if image_template is not None:
395        logger.warning('You are using --image-template. All other parameters will be ignored.')
396        if os.path.exists(image_template):
397            # Local file
398            with open(image_template) as f:
399                content = f.read()
400        else:
401            # It should be an URL
402            msg = '\nusage error: --image-template is not a correct local path or URL'
403            try:
404                r = requests.get(image_template)
405            except Exception:
406                raise CLIError(traceback.format_exc() + msg)
407            if r.status_code != 200:
408                raise CLIError(traceback.format_exc() + msg)
409            content = r.content
410
411        try:
412            obj = json.loads(content)
413        except json.JSONDecodeError:
414            raise CLIError(traceback.format_exc() +
415                           '\nusage error: Content of --image-template is not a valid JSON string')
416        content = {}
417        if 'properties' in obj:
418            content = obj['properties']
419        if 'location' in obj:
420            content['location'] = obj['location']
421        if 'tags' in obj:
422            content['tags'] = obj['tags']
423        if 'identity' in obj:
424            content['identity'] = obj['identity']
425        return client.virtual_machine_image_templates.create_or_update(
426            parameters=content, resource_group_name=resource_group_name, image_template_name=image_template_name)
427
428    template_source, template_scripts, template_destinations = None, [], []
429
430    # create image template source settings
431    if source_dict['type'] == _SourceType.PLATFORM_IMAGE:
432        template_source = ImageTemplatePlatformImageSource(**source_dict)
433    elif source_dict['type'] == _SourceType.ISO_URI:
434        # It was supported before but is removed in the current service version.
435        raise CLIError('usage error: Source type ISO URI is not supported.')
436    elif source_dict['type'] == _SourceType.MANAGED_IMAGE:
437        template_source = ImageTemplateManagedImageSource(**source_dict)
438    elif source_dict['type'] == _SourceType.SIG_VERSION:
439        template_source = ImageTemplateSharedImageVersionSource(**source_dict)
440
441    # create image template customizer settings
442    # Script structure can be found in _parse_script's function definition
443    for script in scripts_list:
444        script.pop("is_url")
445        script["script_uri"] = script.pop("script")
446
447        if script["type"] == ScriptType.SHELL:
448            template_scripts.append(ImageTemplateShellCustomizer(**script))
449        elif script["type"] == ScriptType.POWERSHELL:
450            template_scripts.append(ImageTemplatePowerShellCustomizer(**script))
451        else:  # Should never happen
452            logger.debug("Script %s has type %s", script["script"], script["type"])
453            raise CLIError("Script {} has an invalid type.".format(script["script"]))
454
455    # create image template distribution / destination settings
456    for dest_type, rid, loc_info in destinations_lists:
457        parsed = parse_resource_id(rid)
458        if dest_type == _DestType.MANAGED_IMAGE:
459            template_destinations.append(ImageTemplateManagedImageDistributor(
460                image_id=rid, location=loc_info, run_output_name=parsed['name']))
461        elif dest_type == _DestType.SHARED_IMAGE_GALLERY:
462            template_destinations.append(ImageTemplateSharedImageDistributor(
463                gallery_image_id=rid, replication_regions=loc_info, run_output_name=parsed['child_name_1']))
464        else:
465            logger.info("No applicable destination found for destination %s", str(tuple([dest_type, rid, loc_info])))
466
467    # Identity
468    identity_body = None
469    if identity is not None:
470        subscription_id = get_subscription_id(cmd.cli_ctx)
471        user_assigned_identities = {}
472        for ide in identity:
473            if not is_valid_resource_id(ide):
474                ide = resource_id(subscription=subscription_id, resource_group=resource_group_name,
475                                  namespace='Microsoft.ManagedIdentity', type='userAssignedIdentities', name=ide)
476            user_assigned_identities[ide] = ImageTemplateIdentityUserAssignedIdentitiesValue()
477        identity_body = ImageTemplateIdentity(type='UserAssigned', user_assigned_identities=user_assigned_identities)
478
479    # VM profile
480    vnet_config = None
481    if vnet or subnet:
482        if not is_valid_resource_id(subnet):
483            subnet = resource_id(subscription=client.config.subscription_id, resource_group=resource_group_name,
484                                 namespace='Microsoft.Network', type='virtualNetworks', name=vnet,
485                                 child_type_1='subnets', child_name_1=subnet)
486        vnet_config = VirtualNetworkConfig(subnet_id=subnet)
487    vm_profile = ImageTemplateVmProfile(vm_size=vm_size, os_disk_size_gb=os_disk_size, vnet_config=vnet_config)
488
489    image_template = ImageTemplate(source=template_source, customize=template_scripts, distribute=template_destinations,
490                                   location=location, build_timeout_in_minutes=build_timeout, tags=(tags or {}),
491                                   identity=identity_body, vm_profile=vm_profile)
492
493    return cached_put(cmd, client.virtual_machine_image_templates.create_or_update, parameters=image_template,
494                      resource_group_name=resource_group_name, image_template_name=image_template_name)
495
496
497def list_image_templates(client, resource_group_name=None):
498    if resource_group_name:
499        return client.virtual_machine_image_templates.list_by_resource_group(resource_group_name)
500    return client.virtual_machine_image_templates.list()
501
502
503def show_build_output(client, resource_group_name, image_template_name, output_name=None):
504    if output_name:
505        return client.virtual_machine_image_templates.get_run_output(resource_group_name, image_template_name, output_name)  # pylint: disable=line-too-long
506    return client.virtual_machine_image_templates.list_run_outputs(resource_group_name, image_template_name)
507
508
509def add_template_output(cmd, client, resource_group_name, image_template_name, gallery_name=None, location=None,  # pylint: disable=line-too-long, unused-argument
510                        output_name=None, is_vhd=None, tags=None,
511                        gallery_image_definition=None, gallery_replication_regions=None,
512                        managed_image=None, managed_image_location=None):  # pylint: disable=line-too-long, unused-argument
513
514    _require_defer(cmd)
515
516    from azure.mgmt.imagebuilder.models import (ImageTemplateManagedImageDistributor, ImageTemplateVhdDistributor,
517                                                ImageTemplateSharedImageDistributor)
518    existing_image_template = cached_get(cmd, client.virtual_machine_image_templates.get,
519                                         resource_group_name=resource_group_name,
520                                         image_template_name=image_template_name)
521
522    distributor = None
523
524    if managed_image:
525        parsed = parse_resource_id(managed_image)
526        distributor = ImageTemplateManagedImageDistributor(
527            run_output_name=output_name or parsed['name'],
528            image_id=managed_image, location=managed_image_location or location)
529    elif gallery_image_definition:
530        parsed = parse_resource_id(gallery_image_definition)
531        distributor = ImageTemplateSharedImageDistributor(
532            run_output_name=output_name or parsed['child_name_1'], gallery_image_id=gallery_image_definition,
533            replication_regions=gallery_replication_regions or [location])
534    elif is_vhd:
535        distributor = ImageTemplateVhdDistributor(run_output_name=output_name)
536
537    if distributor:
538        distributor.artifact_tags = tags or {}
539
540    if existing_image_template.distribute is None:
541        existing_image_template.distribute = []
542    else:
543        for existing_distributor in existing_image_template.distribute:
544            if existing_distributor.run_output_name == distributor.run_output_name:
545                raise CLIError("Output with output name {} already exists in image template {}."
546                               .format(distributor.run_output_name.lower(), image_template_name))
547
548    existing_image_template.distribute.append(distributor)
549
550    return cached_put(cmd, client.virtual_machine_image_templates.create_or_update, parameters=existing_image_template,
551                      resource_group_name=resource_group_name, image_template_name=image_template_name)
552
553
554def remove_template_output(cmd, client, resource_group_name, image_template_name, output_name):
555    _require_defer(cmd)
556
557    existing_image_template = cached_get(cmd, client.virtual_machine_image_templates.get,
558                                         resource_group_name=resource_group_name,
559                                         image_template_name=image_template_name)
560    if not existing_image_template.distribute:
561        raise CLIError("No outputs to remove.")
562
563    new_distribute = []
564    for existing_distributor in existing_image_template.distribute:
565        if existing_distributor.run_output_name.lower() == output_name.lower():
566            continue
567        new_distribute.append(existing_distributor)
568
569    if len(new_distribute) == len(existing_image_template.distribute):
570        raise CLIError("Output with output name {} not in image template distribute list.".format(output_name))
571
572    existing_image_template.distribute = new_distribute
573
574    return cached_put(cmd, client.virtual_machine_image_templates.create_or_update, parameters=existing_image_template,
575                      resource_group_name=resource_group_name, image_template_name=image_template_name)
576
577
578def clear_template_output(cmd, client, resource_group_name, image_template_name):
579    _require_defer(cmd)
580
581    existing_image_template = cached_get(cmd, client.virtual_machine_image_templates.get,
582                                         resource_group_name=resource_group_name,
583                                         image_template_name=image_template_name)
584    if not existing_image_template.distribute:
585        raise CLIError("No outputs to remove.")
586
587    existing_image_template.distribute = []
588
589    return cached_put(cmd, client.virtual_machine_image_templates.create_or_update, parameters=existing_image_template,
590                      resource_group_name=resource_group_name, image_template_name=image_template_name)
591
592
593def add_template_customizer(cmd, client, resource_group_name, image_template_name, customizer_name, customizer_type,
594                            script_url=None, inline_script=None, valid_exit_codes=None,
595                            restart_command=None, restart_check_command=None, restart_timeout=None,
596                            file_source=None, dest_path=None, search_criteria=None, filters=None, update_limit=None):
597    _require_defer(cmd)
598
599    from azure.mgmt.imagebuilder.models import (ImageTemplateShellCustomizer, ImageTemplatePowerShellCustomizer,
600                                                ImageTemplateRestartCustomizer, ImageTemplateFileCustomizer,
601                                                ImageTemplateWindowsUpdateCustomizer)
602
603    existing_image_template = cached_get(cmd, client.virtual_machine_image_templates.get,
604                                         resource_group_name=resource_group_name,
605                                         image_template_name=image_template_name)
606
607    if existing_image_template.customize is None:
608        existing_image_template.customize = []
609    else:
610        for existing_customizer in existing_image_template.customize:
611            if existing_customizer.name == customizer_name:
612                raise CLIError("Output with output name {} already exists in image template {}."
613                               .format(customizer_name, image_template_name))
614
615    new_customizer = None
616
617    if customizer_type.lower() == ScriptType.SHELL.value.lower():  # pylint:disable=no-member
618        new_customizer = ImageTemplateShellCustomizer(name=customizer_name, script_uri=script_url, inline=inline_script)
619    elif customizer_type.lower() == ScriptType.POWERSHELL.value.lower():  # pylint:disable=no-member
620        new_customizer = ImageTemplatePowerShellCustomizer(name=customizer_name, script_uri=script_url,
621                                                           inline=inline_script, valid_exit_codes=valid_exit_codes)
622    elif customizer_type.lower() == ScriptType.WINDOWS_RESTART.value.lower():  # pylint:disable=no-member
623        new_customizer = ImageTemplateRestartCustomizer(name=customizer_name, restart_command=restart_command,
624                                                        restart_check_command=restart_check_command,
625                                                        restart_timeout=restart_timeout)
626    elif customizer_type.lower() == ScriptType.FILE.value.lower():  # pylint:disable=no-member
627        new_customizer = ImageTemplateFileCustomizer(name=customizer_name, source_uri=file_source,
628                                                     destination=dest_path)
629    elif customizer_type.lower() == ScriptType.WINDOWS_UPDATE.value.lower():
630        new_customizer = ImageTemplateWindowsUpdateCustomizer(name=customizer_name, search_criteria=search_criteria,
631                                                              filters=filters, update_limit=update_limit)
632
633    if not new_customizer:
634        raise CLIError("Cannot determine customizer from type {}.".format(customizer_type))
635
636    existing_image_template.customize.append(new_customizer)
637
638    return cached_put(cmd, client.virtual_machine_image_templates.create_or_update, parameters=existing_image_template,
639                      resource_group_name=resource_group_name, image_template_name=image_template_name)
640
641
642def remove_template_customizer(cmd, client, resource_group_name, image_template_name, customizer_name):
643    existing_image_template = cached_get(cmd, client.virtual_machine_image_templates.get,
644                                         resource_group_name=resource_group_name,
645                                         image_template_name=image_template_name)
646    _require_defer(cmd)
647
648    if not existing_image_template.customize:
649        raise CLIError("No customizers to remove.")
650
651    new_customize = []
652    for existing_customizer in existing_image_template.customize:
653        if existing_customizer.name == customizer_name:
654            continue
655        new_customize.append(existing_customizer)
656
657    if len(new_customize) == len(existing_image_template.customize):
658        raise CLIError("Customizer with name {} not in image template customizer list.".format(customizer_name))
659
660    existing_image_template.customize = new_customize
661
662    return cached_put(cmd, client.virtual_machine_image_templates.create_or_update, parameters=existing_image_template,
663                      resource_group_name=resource_group_name, image_template_name=image_template_name)
664
665
666def clear_template_customizer(cmd, client, resource_group_name, image_template_name):
667    _require_defer(cmd)
668
669    existing_image_template = cached_get(cmd, client.virtual_machine_image_templates.get,
670                                         resource_group_name=resource_group_name,
671                                         image_template_name=image_template_name)
672
673    if not existing_image_template.customize:
674        raise CLIError("No customizers to remove.")
675
676    existing_image_template.customize = []
677
678    return cached_put(cmd, client.virtual_machine_image_templates.create_or_update, parameters=existing_image_template,
679                      resource_group_name=resource_group_name, image_template_name=image_template_name)
680
681# endregion
682