image-copy: first commit

src/image-copy/azext_imagecopy/__init__.py

@@ -0,0 +1,26 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core.help_files import helps
+from azure.cli.core.sdk.util import ParametersContext
+
+helps['image copy'] = """
+    type: command
+    short-summary: Allows to copy an image (or vm) to other regions. Keep in mind that it requires the source disk to be available.
+"""
+
+def load_params(_):
+    with ParametersContext('image copy') as c:
+        c.register('source_resource_group_name', '--source-resource-group-name', help='Name of the resource gorup of the source resource')
+        c.register('source_object_name', '--source-object-name', help='The name of the image or vm resource')
+        c.register('target_location', '--target-location', help='Comma seperated location list to create the image in')
+        c.register('source_type', '--source-type', help='image (default) or vm')
+        c.register('target_resource_group_name', '--target-resource-group-name', help='Name of the resource group to create images in')
+        c.register('parallel_degree', '--parallel-degree', help='Number of parallel copy operations')
+        c.register('cleanup', '--cleanup', help='Will delete temporary resources created in process (false by default)')
+
+def load_commands():
+    from azure.cli.core.commands import cli_command
+    cli_command(__name__, 'image copy', 'azext_imagecopy.custom#imagecopy')

src/image-copy/azext_imagecopy/azext_metadata.json

@@ -0,0 +1,3 @@
+{
+    "azext.minCliVersion": "2.0.12"
+}

src/image-copy/azext_imagecopy/cli_utils.py

@@ -0,0 +1,44 @@
+import sys
+import json
+
+from subprocess import check_output, STDOUT, CalledProcessError
+
+import azure.cli.core.azlogging as azlogging
+
+logger = azlogging.get_az_logger(__name__)
+
+
+def run_cli_command(cmd, return_as_json=False):
+    try:
+        cmd_output = check_output(cmd, stderr=STDOUT, universal_newlines=True)
+        logger.debug(cmd_output)
+
+        if return_as_json is True:
+            if cmd_output:
+                json_output = json.loads(cmd_output)
+                return json_output
+            else:
+                raise Exception("Command returned an unexpected empty string.")
+        else:
+            return cmd_output
+    except CalledProcessError as ex:
+        print('command failed: ', cmd)
+        print('output: ', ex.output)
+        raise ex
+    except:
+        print('command: ', cmd)
+        raise
+
+def prepare_cli_command(cmd, output_as_json=True):
+    full_cmd = [sys.executable, '-m', 'azure.cli'] + cmd
+
+    if output_as_json:
+        full_cmd += ['--output', 'json']
+    else:
+        full_cmd += ['--output', 'tsv']
+
+    # tag newly created resources
+    if 'create' in cmd and ('container' not in cmd):
+        full_cmd += ['--tags', 'created_by=image-copy-extension']
+
+    return full_cmd

src/image-copy/azext_imagecopy/create_target.py

@@ -0,0 +1,170 @@
+import hashlib
+import datetime
+import time
+from azext_imagecopy.cli_utils import run_cli_command, prepare_cli_command
+from azure.cli.core.application import APPLICATION
+import azure.cli.core.azlogging as azlogging
+
+logger = azlogging.get_az_logger(__name__)
+
+
+def create_target_image(location, transient_resource_group_name, source_type, source_object_name, \
+    source_os_disk_snapshot_name, source_os_disk_snapshot_url, source_os_type, \
+    target_resource_group_name, azure_pool_frequency):
+
+    #from azure.cli.core.commands.client_factory import get_subscription_id
+    subscription_id = get_subscription_id()
+
+    subscription_hash = hashlib.sha1(subscription_id.encode("UTF-8")).hexdigest()
+    unique_subscription_string = subscription_hash[:7]
+
+
+    # create the target storage account
+    logger.warn("{0} - Creating target storage account (can be slow sometimes)".format(location))
+    target_storage_account_name = location + unique_subscription_string
+    cmd = prepare_cli_command(['storage', 'account', 'create', \
+        '--name', target_storage_account_name, \
+        '--resource-group', transient_resource_group_name, \
+        '--location', location, \
+        '--sku', 'Standard_LRS'])
+
+    json_output = run_cli_command(cmd, True)
+    target_blob_endpoint = json_output['primaryEndpoints']['blob']
+
+
+    # Setup the target storage account
+    cmd = prepare_cli_command(['storage', 'account', 'keys', 'list', \
+        '--account-name', target_storage_account_name, \
+        '--resource-group', transient_resource_group_name])
+
+    json_output = run_cli_command(cmd, True)
+
+    target_storage_account_key = json_output[0]['value']
+    logger.debug(target_storage_account_key)
+
+    expiry_format = "%Y-%m-%dT%H:%MZ"
+    expiry = datetime.datetime.utcnow() + datetime.timedelta(hours=1)
+
+    cmd = prepare_cli_command(['storage', 'account', 'generate-sas', \
+        '--account-name', target_storage_account_name, \
+        '--account-key', target_storage_account_key, \
+        '--expiry', expiry.strftime(expiry_format), \
+        '--permissions', 'aclrpuw', '--resource-types', \
+        'sco', '--services', 'b', '--https-only'], \
+        False)
+
+    sas_token = run_cli_command(cmd)
+    sas_token = sas_token.rstrip("\n\r") #STRANGE
+    logger.debug("sas token: " + sas_token)
+
+
+    # create a container in the target blob storage account
+    logger.warn("{0} - Creating container in the target storage account".format(location))
+    target_container_name = 'snapshots'
+    cmd = prepare_cli_command(['storage', 'container', 'create', \
+        '--name', target_container_name, \
+        '--account-name', target_storage_account_name])
+
+    run_cli_command(cmd)
+
+
+    # Copy the snapshot to the target region using the SAS URL
+    blob_name = source_os_disk_snapshot_name + '.vhd'
+    logger.warn("{0} - Copying blob to target storage account".format(location))
+    cmd = prepare_cli_command(['storage', 'blob', 'copy', 'start', \
+        '--source-uri', source_os_disk_snapshot_url, \
+        '--destination-blob', blob_name, \
+        '--destination-container', target_container_name, \
+        '--account-name', target_storage_account_name, \
+        '--sas-token', sas_token])
+
+    run_cli_command(cmd)
+
+
+    # Wait for the copy to complete
+    start_datetime = datetime.datetime.now()
+    wait_for_blob_copy_operation(blob_name, target_container_name, \
+        target_storage_account_name, azure_pool_frequency, location)
+    msg = "{0} - Copy time: {1}".format(location, datetime.datetime.now()-start_datetime).ljust(40)
+    logger.warn(msg)
+
+
+    # Create the snapshot in the target region from the copied blob
+    logger.warn("{0} - Creating snapshot in target region from the copied blob".format(location))
+    target_blob_path = target_blob_endpoint + target_container_name + '/' + blob_name
+    target_snapshot_name = source_os_disk_snapshot_name + '-' + location
+    cmd = prepare_cli_command(['snapshot', 'create', \
+        '--resource-group', transient_resource_group_name, \
+        '--name', target_snapshot_name, \
+        '--location', location, \
+        '--source', target_blob_path])
+
+    json_output = run_cli_command(cmd, True)
+    target_snapshot_id = json_output['id']
+
+    # Create the final image
+    logger.warn("{0} - Creating final image".format(location))
+    target_image_name = source_object_name
+    if source_type != 'image':
+        target_image_name += '-image'
+    target_image_name += '-' + location
+
+    cmd = prepare_cli_command(['image', 'create', \
+        '--resource-group', target_resource_group_name, \
+        '--name', target_image_name, \
+        '--location', location, \
+        '--source', target_blob_path, \
+        '--os-type', source_os_type, \
+        '--source', target_snapshot_id])
+
+    run_cli_command(cmd)
+
+
+def wait_for_blob_copy_operation(blob_name, target_container_name, target_storage_account_name, azure_pool_frequency, location):
+    progress_controller = APPLICATION.get_progress_controller()
+    copy_status = "pending"
+    prev_progress = -1
+    while copy_status == "pending":
+        cmd = prepare_cli_command(['storage', 'blob', 'show', \
+            '--name', blob_name, \
+            '--container-name', target_container_name, \
+            '--account-name', target_storage_account_name])
+
+        json_output = run_cli_command(cmd, True)
+        copy_status = json_output["properties"]["copy"]["status"]
+        copy_progress_1, copy_progress_2 = json_output["properties"]["copy"]["progress"].split("/")
+        current_progress = round(int(copy_progress_1)/int(copy_progress_2), 1)
+
+        if current_progress != prev_progress:
+            msg = "{0} - copy progress: {1}%".format(location, str(current_progress)).ljust(40) #need to justify since message overide each other
+            progress_controller.add(message=msg)
+
+        prev_progress = current_progress
+
+        try:
+            time.sleep(azure_pool_frequency)
+        except KeyboardInterrupt:
+            print('xxx - child')
+            progress_controller.stop()
+            return
+
+
+    if copy_status == 'success':
+        progress_controller.stop()
+    else:
+        logger.error("The copy operation didn't succeed. Last status: " + copy_status)
+        raise Exception('Blob copy failed')
+
+
+def get_subscription_id():
+    logger.warn("Get subscription id")
+
+    cmd = prepare_cli_command(['account', 'show'])
+    json_output = run_cli_command(cmd, True)
+    subscription_id = json_output['id']
+
+    # from azure.cli.core._profile import Profile
+    # profile = Profile()
+    # _, subscription_id, _ = profile.get_login_credentials()
+
+    return subscription_id

src/image-copy/azext_imagecopy/custom.py

@@ -0,0 +1,135 @@
+from multiprocessing import Pool
+
+from azext_imagecopy.cli_utils import run_cli_command, prepare_cli_command
+from azext_imagecopy.create_target import create_target_image
+import azure.cli.core.azlogging as azlogging
+
+logger = azlogging.get_az_logger(__name__)
+
+
+def imagecopy(source_resource_group_name, source_object_name, target_location, \
+    target_resource_group_name, source_type='image', cleanup='false', parallel_degree=-1):
+
+    # get the os disk id from source vm/image
+    logger.warn("Getting os disk id of the source vm/image")
+    cmd = prepare_cli_command([source_type, 'show', \
+        '--name', source_object_name, \
+        '--resource-group', source_resource_group_name])
+
+    json_cmd_output = run_cli_command(cmd, True)
+
+    source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['managedDisk']['id']
+    source_os_type = json_cmd_output['storageProfile']['osDisk']['osType']
+    logger.debug("source_os_disk_id: " + source_os_disk_id + " source_os_type: " + source_os_type)
+
+
+    # create source snapshots
+    logger.warn("Creating source snapshot")
+    source_os_disk_snapshot_name = source_object_name + '_os_disk_snapshot'
+    cmd = prepare_cli_command(['snapshot', 'create', \
+        '--name', source_os_disk_snapshot_name, \
+        '--resource-group', source_resource_group_name, \
+        '--source', source_os_disk_id])
+
+    run_cli_command(cmd)
+
+
+    # Get SAS URL for the snapshotName
+    logger.warn("Getting sas url for the source snapshot")
+    cmd = prepare_cli_command(['snapshot', 'grant-access', \
+        '--name', source_os_disk_snapshot_name, \
+        '--resource-group', source_resource_group_name, \
+        '--duration-in-seconds', '3600'])
+
+    json_output = run_cli_command(cmd, True)
+
+    source_os_disk_snapshot_url = json_output['accessSas']
+    logger.debug(source_os_disk_snapshot_url)
+
+
+    # Start processing in the target locations
+
+    transient_resource_group_name = 'image-copy-rg'
+    create_resource_group(transient_resource_group_name, 'eastus')
+
+    target_locations_list = target_location.split(',')
+    target_locations_count = len(target_locations_list)
+    logger.warn("Target location count: {0}".format(target_locations_count))
+
+    create_resource_group(target_resource_group_name, target_locations_list[0].strip())
+
+    if parallel_degree == -1:
+        pool = Pool(target_locations_count)
+    elif parallel_degree == 1:
+        pool = Pool(1)
+    else:
+        pool = Pool(min(parallel_degree, target_locations_count))
+
+    # try to get a handle on arm 409s
+    azure_pool_frequency = 5
+    if target_locations_count >= 5:
+        azure_pool_frequency = 15
+    elif  target_locations_count >= 3:
+        azure_pool_frequency = 10
+
+    tasks = []
+    for location in target_location.split(','):
+        location = location.strip()
+        tasks.append((location, transient_resource_group_name, source_type, \
+        source_object_name, source_os_disk_snapshot_name, source_os_disk_snapshot_url, \
+        source_os_type, target_resource_group_name, azure_pool_frequency))
+
+    logger.warn("Starting async process for all locations")
+
+    for task in tasks:
+        pool.apply_async(create_target_image, task)
+
+    try:
+        pool.close()
+        pool.join()
+    except KeyboardInterrupt:
+        print('xxx - parent')
+        logger.warn('User cancelled the operation')
+        if 'true' in cleanup:
+            logger.warn('To cleanup temporary resources look for ones tagged with "image-copy-extension"')
+        pool.terminate()
+        return
+
+
+    # Cleanup
+    if 'true' in cleanup:
+        logger.warn('Deleting transient resources')
+
+        # Delete resource group
+        cmd = prepare_cli_command(['group', 'delete', '--no-wait', '--yes', \
+            '--name', transient_resource_group_name])
+        run_cli_command(cmd)
+
+        # Revoke sas for source snapshot
+        cmd = prepare_cli_command(['snapshot', 'revoke-access', \
+            '--name', source_os_disk_snapshot_name, \
+            '--resource-group', source_resource_group_name])
+        run_cli_command(cmd)
+
+        # Delete source snapshot
+        cmd = prepare_cli_command(['snapshot', 'delete', \
+            '--name', source_os_disk_snapshot_name, \
+            '--resource-group', source_resource_group_name])
+        run_cli_command(cmd)
+
+
+def create_resource_group(resource_group_name, location):
+    # check if target resource group exists
+    cmd = prepare_cli_command(['group', 'exists', \
+        '--name', resource_group_name], False)
+
+    cmd_output = run_cli_command(cmd)
+
+    if 'false' in cmd_output:
+        # create the target resource group
+        logger.warn("Creating resource group")
+        cmd = prepare_cli_command(['group', 'create', \
+            '--name', resource_group_name, \
+            '--location', location])
+
+        run_cli_command(cmd)

src/image-copy/setup.cfg

@@ -0,0 +1,2 @@
+[bdist_wheel]
+universal=1