image-copy: add support for setting target subscription

src/image-copy/azext_imagecopy/__init__.py

@@ -41,6 +41,8 @@ def load_arguments(self, _):
                        help='Include this switch to delete temporary resources upon completion')
             c.argument('target_name', options_list=['--target-name'],
                        help='Name of the final image that will be created')
+            c.argument('target_subscription', options_list=['--target-subscription'],
+                       help='Name or ID of the subscription where the final image should be created.')
 
 
 COMMAND_LOADER_CLS = ImageCopyCommandsLoader

src/image-copy/azext_imagecopy/cli_utils.py

@@ -25,8 +25,8 @@ def run_cli_command(cmd, return_as_json=False):
             if cmd_output:
                 json_output = json.loads(cmd_output)
                 return json_output
-            else:
-                raise CLIError("Command returned an unexpected empty string.")
+
+            raise CLIError("Command returned an unexpected empty string.")
         else:
             return cmd_output
     except CalledProcessError as ex:
@@ -38,14 +38,18 @@ def run_cli_command(cmd, return_as_json=False):
         raise
 
 
-def prepare_cli_command(cmd, output_as_json=True, tags=None):
+def prepare_cli_command(cmd, output_as_json=True, tags=None, subscription=None):
     full_cmd = [sys.executable, '-m', 'azure.cli'] + cmd
 
     if output_as_json:
         full_cmd += ['--output', 'json']
     else:
         full_cmd += ['--output', 'tsv']
 
+    # override the default subscription if needed
+    if subscription is not None:
+        full_cmd += ['--subscription', subscription]
+
     # tag newly created resources, containers don't have tags
     if 'create' in cmd and ('container' not in cmd):
         full_cmd += ['--tags', EXTENSION_TAG_STRING]

src/image-copy/azext_imagecopy/create_target.py

@@ -3,7 +3,6 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-import hashlib
 import datetime
 import time
 from azext_imagecopy.cli_utils import run_cli_command, prepare_cli_command
@@ -18,32 +17,29 @@
 # pylint: disable=too-many-locals
 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, tags, target_name):
+                        target_resource_group_name, azure_pool_frequency, tags, target_name, target_subscription):
 
-    subscription_id = get_subscription_id()
-
-    subscription_hash = hashlib.sha1(
-        subscription_id.encode("UTF-8")).hexdigest()
-    unique_subscription_string = subscription_hash[:(
-        STORAGE_ACCOUNT_NAME_LENGTH - len(location))]
+    random_string = get_random_string(STORAGE_ACCOUNT_NAME_LENGTH - len(location))
 
     # create the target storage account
     logger.warn(
         "%s - Creating target storage account (can be slow sometimes)", location)
-    target_storage_account_name = location + unique_subscription_string
+    target_storage_account_name = location + random_string
     cli_cmd = prepare_cli_command(['storage', 'account', 'create',
                                    '--name', target_storage_account_name,
                                    '--resource-group', transient_resource_group_name,
                                    '--location', location,
-                                   '--sku', 'Standard_LRS'])
+                                   '--sku', 'Standard_LRS'],
+                                  subscription=target_subscription)
 
     json_output = run_cli_command(cli_cmd, return_as_json=True)
     target_blob_endpoint = json_output['primaryEndpoints']['blob']
 
     # Setup the target storage account
     cli_cmd = prepare_cli_command(['storage', 'account', 'keys', 'list',
                                    '--account-name', target_storage_account_name,
-                                   '--resource-group', transient_resource_group_name])
+                                   '--resource-group', transient_resource_group_name],
+                                  subscription=target_subscription)
 
     json_output = run_cli_command(cli_cmd, return_as_json=True)
 
@@ -59,7 +55,8 @@ def create_target_image(location, transient_resource_group_name, source_type, so
                                    '--expiry', expiry.strftime(expiry_format),
                                    '--permissions', 'aclrpuw', '--resource-types',
                                    'sco', '--services', 'b', '--https-only'],
-                                  output_as_json=False)
+                                  output_as_json=False,
+                                  subscription=target_subscription)
 
     sas_token = run_cli_command(cli_cmd)
     sas_token = sas_token.rstrip("\n\r")  # STRANGE
@@ -71,7 +68,8 @@ def create_target_image(location, transient_resource_group_name, source_type, so
     target_container_name = 'snapshots'
     cli_cmd = prepare_cli_command(['storage', 'container', 'create',
                                    '--name', target_container_name,
-                                   '--account-name', target_storage_account_name])
+                                   '--account-name', target_storage_account_name],
+                                  subscription=target_subscription)
 
     run_cli_command(cli_cmd)
 
@@ -84,14 +82,15 @@ def create_target_image(location, transient_resource_group_name, source_type, so
                                    '--destination-blob', blob_name,
                                    '--destination-container', target_container_name,
                                    '--account-name', target_storage_account_name,
-                                   '--sas-token', sas_token])
+                                   '--sas-token', sas_token],
+                                  subscription=target_subscription)
 
     run_cli_command(cli_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)
+    wait_for_blob_copy_operation(blob_name, target_container_name, target_storage_account_name,
+                                 azure_pool_frequency, location, target_subscription)
     msg = "{0} - Copy time: {1}".format(
         location, datetime.datetime.now() - start_datetime)
     logger.warn(msg)
@@ -106,7 +105,8 @@ def create_target_image(location, transient_resource_group_name, source_type, so
                                    '--resource-group', transient_resource_group_name,
                                    '--name', target_snapshot_name,
                                    '--location', location,
-                                   '--source', target_blob_path])
+                                   '--source', target_blob_path],
+                                  subscription=target_subscription)
 
     json_output = run_cli_command(cli_cmd, return_as_json=True)
     target_snapshot_id = json_output['id']
@@ -127,20 +127,23 @@ def create_target_image(location, transient_resource_group_name, source_type, so
                                    '--location', location,
                                    '--source', target_blob_path,
                                    '--os-type', source_os_type,
-                                   '--source', target_snapshot_id], tags=tags)
+                                   '--source', target_snapshot_id],
+                                  tags=tags,
+                                  subscription=target_subscription)
 
     run_cli_command(cli_cmd)
 
 
 def wait_for_blob_copy_operation(blob_name, target_container_name, target_storage_account_name,
-                                 azure_pool_frequency, location):
+                                 azure_pool_frequency, location, subscription):
     copy_status = "pending"
     prev_progress = -1
     while copy_status == "pending":
         cli_cmd = prepare_cli_command(['storage', 'blob', 'show',
                                        '--name', blob_name,
                                        '--container-name', target_container_name,
-                                       '--account-name', target_storage_account_name])
+                                       '--account-name', target_storage_account_name],
+                                      subscription=subscription)
 
         json_output = run_cli_command(cli_cmd, return_as_json=True)
         copy_status = json_output["properties"]["copy"]["status"]
@@ -161,17 +164,13 @@ def wait_for_blob_copy_operation(blob_name, target_container_name, target_storag
         except KeyboardInterrupt:
             return
 
-    if copy_status == 'success':
-        return
-    else:
-        logger.error(
-            "The copy operation didn't succeed. Last status: %s", copy_status)
+    if copy_status != 'success':
+        logger.error("The copy operation didn't succeed. Last status: %s", copy_status)
         raise CLIError('Blob copy failed')
 
 
-def get_subscription_id():
-    cli_cmd = prepare_cli_command(['account', 'show'])
-    json_output = run_cli_command(cli_cmd, return_as_json=True)
-    subscription_id = json_output['id']
-
-    return subscription_id
+def get_random_string(length):
+    import string
+    import random
+    chars = string.ascii_lowercase + string.digits
+    return ''.join(random.choice(chars) for _ in range(length))

src/image-copy/azext_imagecopy/custom.py

@@ -14,9 +14,10 @@
 
 
 # pylint: disable=too-many-statements
+# pylint: disable=too-many-locals
 def imagecopy(source_resource_group_name, source_object_name, target_location,
               target_resource_group_name, source_type='image', cleanup='false',
-              parallel_degree=-1, tags=None, target_name=None):
+              parallel_degree=-1, tags=None, target_name=None, target_subscription=None):
 
     # get the os disk id from source vm/image
     logger.warn("Getting os disk id of the source vm/image")
@@ -95,40 +96,43 @@ def imagecopy(source_resource_group_name, source_object_name, target_location,
     # pick the first location for the temp group
     transient_resource_group_location = target_location[0].strip()
     create_resource_group(transient_resource_group_name,
-                          transient_resource_group_location)
+                          transient_resource_group_location,
+                          target_subscription)
 
     target_locations_count = len(target_location)
     logger.warn("Target location count: %s", target_locations_count)
 
     create_resource_group(target_resource_group_name,
-                          target_location[0].strip())
+                          target_location[0].strip(),
+                          target_subscription)
 
-    if parallel_degree == -1:
-        pool = Pool(target_locations_count)
-    else:
-        pool = Pool(min(parallel_degree, target_locations_count))
+    try:
 
-    # try to get a handle on arm's 409s
-    azure_pool_frequency = 5
-    if target_locations_count >= 5:
-        azure_pool_frequency = 15
-    elif target_locations_count >= 3:
-        azure_pool_frequency = 10
+        # try to get a handle on arm's 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:
-        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,
-                      tags, target_name))
+        if parallel_degree == -1:
+            pool = Pool(target_locations_count)
+        else:
+            pool = Pool(min(parallel_degree, target_locations_count))
 
-    logger.warn("Starting async process for all locations")
+        tasks = []
+        for location in target_location:
+            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,
+                          tags, target_name, target_subscription))
 
-    for task in tasks:
-        pool.apply_async(create_target_image, task)
+        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:
@@ -145,7 +149,8 @@ def imagecopy(source_resource_group_name, source_object_name, target_location,
 
         # Delete resource group
         cli_cmd = prepare_cli_command(['group', 'delete', '--no-wait', '--yes',
-                                       '--name', transient_resource_group_name])
+                                       '--name', transient_resource_group_name],
+                                      subscription=target_subscription)
         run_cli_command(cli_cmd)
 
         # Revoke sas for source snapshot
@@ -162,10 +167,12 @@ def imagecopy(source_resource_group_name, source_object_name, target_location,
         run_cli_command(cli_cmd)
 
 
-def create_resource_group(resource_group_name, location):
+def create_resource_group(resource_group_name, location, subscription=None):
     # check if target resource group exists
     cli_cmd = prepare_cli_command(['group', 'exists',
-                                   '--name', resource_group_name], output_as_json=False)
+                                   '--name', resource_group_name],
+                                  output_as_json=False,
+                                  subscription=subscription)
 
     cmd_output = run_cli_command(cli_cmd)
 
@@ -176,6 +183,7 @@ def create_resource_group(resource_group_name, location):
     logger.warn("Creating resource group: %s", resource_group_name)
     cli_cmd = prepare_cli_command(['group', 'create',
                                    '--name', resource_group_name,
-                                   '--location', location])
+                                   '--location', location],
+                                  subscription=subscription)
 
     run_cli_command(cli_cmd)

src/image-copy/setup.py

@@ -8,7 +8,7 @@
 from codecs import open
 from setuptools import setup, find_packages
 
-VERSION = "0.0.7"
+VERSION = "0.0.8"
 
 CLASSIFIERS = [
     'Development Status :: 4 - Beta',