We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Image Copy '--target-subscription' Argument Ignored
### Related command
<!--- Please provide the related command with az {command} if you can, so that we can quickly route to the related team for help. --->
```bash
az image copy --source-resource-group <source_rg_name> \
        --source-object-name <source_image_name> \
        --source-type image \
        --target-subscription <target_subscription_id> \
        --target-location <target_location> \
        --target-resource-group <target_rg_name>  \
        --target-name <target_image_name> \
        --temporary-resource-group-name <temp_rg_name>
```

### Extension name (the extension in question)
image-copy-extension

### Description of issue (in as much detail as possible)

My organization is using the `image-copy-extension` and we noticed on recent workflows that the extension is now ignoring the `--target-subscription` argument and seems to use the authentication context subscription as the target subscription for all temporary resource and final image creation instead.

We noticed this change because previously the temporary resource group we used was stored within the target subscription and this meant that the temporary storage account, blob, and snapshot all got created within this temporary resource group but now the extension seems to be looking for the temporary resource group in the default subscription associated with the authentication context instead of the target subscription provided to the `--target-subscription` argument.

My Python experience is still limited but I have looked at the [recent commit](https://github.com/Azure/azure-cli-extensions/commit/3a6621cc0c96aab5ddd2388f4a2e6a81f4ed2f0a) and I see that the `target_subscription` variable is now being set without checking if the argument has been passed and I'm not sure if this was changed for a reason or if this is just a bug that can be fixed by only setting the `target_subscription` variable if the argument is null.

![image](https://user-images.githubusercontent.com/20447447/209253966-f9129522-3988-4c73-9abc-d6fc0e7ef991.png)

Link to the line in question.

https://github.com/Azure/azure-cli-extensions/blob/42102c3a81ad23407e7d929521ea212805b4688d/src/image-copy/azext_imagecopy/custom.py#L37

There is also a debug line below the variable declaration that we've seen continues to display the subscription ID for the authentication context and not the target subscription we're providing as an argument.

-----



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/image-copy/azext_imagecopy/custom.py`
Content:
```
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 from multiprocessing import Pool
7 
8 from knack.util import CLIError
9 from knack.log import get_logger
10 
11 from azext_imagecopy.cli_utils import run_cli_command, prepare_cli_command, get_storage_account_id_from_blob_path
12 from azext_imagecopy.create_target import create_target_image
13 
14 logger = get_logger(__name__)
15 
16 
17 # pylint: disable=too-many-statements
18 # pylint: disable=too-many-locals
19 # pylint: disable=too-many-branches
20 def imagecopy(cmd, source_resource_group_name, source_object_name, target_location,
21               target_resource_group_name, temporary_resource_group_name='image-copy-rg',
22               source_type='image', cleanup=False, parallel_degree=-1, tags=None, target_name=None,
23               target_subscription=None, export_as_snapshot='false', timeout=3600):
24     only_show_errors = cmd.cli_ctx.only_show_errors
25     if cleanup:
26         # If --cleanup is set, forbid using an existing temporary resource group name.
27         # It is dangerous to clean up an existing resource group.
28         cli_cmd = prepare_cli_command(['group', 'exists', '-n', temporary_resource_group_name],
29                                       output_as_json=False,
30                                       only_show_errors=only_show_errors)
31         cmd_output = run_cli_command(cli_cmd)
32         if 'true' in cmd_output:
33             raise CLIError('Don\'t specify an existing resource group in --temporary-resource-group-name '
34                            'when --cleanup is set')
35 
36     from azure.cli.core.commands.client_factory import get_subscription_id
37     target_subscription = get_subscription_id(cmd.cli_ctx)
38     logger.debug('subscription id - %s', target_subscription)
39 
40     # get the os disk id from source vm/image
41     logger.warning("Getting OS disk ID of the source VM/image")
42     cli_cmd = prepare_cli_command([source_type, 'show',
43                                    '--name', source_object_name,
44                                    '--resource-group', source_resource_group_name],
45                                   only_show_errors=only_show_errors)
46 
47     json_cmd_output = run_cli_command(cli_cmd, return_as_json=True)
48 
49     if json_cmd_output['storageProfile']['dataDisks']:
50         logger.warning(
51             "Data disks in the source detected, but are ignored by this extension!")
52 
53     source_os_disk_id = None
54     source_os_disk_type = None
55 
56     try:
57         source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['managedDisk']['id']
58         if source_os_disk_id is None:
59             raise TypeError
60         source_os_disk_type = "DISK"
61         logger.debug("found %s: %s", source_os_disk_type, source_os_disk_id)
62     except TypeError:
63         try:
64             source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['blobUri']
65             if source_os_disk_id is None:
66                 raise TypeError
67             source_os_disk_type = "BLOB"
68             logger.debug("found %s: %s", source_os_disk_type,
69                          source_os_disk_id)
70         except TypeError:
71             try:  # images created by e.g. image-copy extension
72                 source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['snapshot']['id']
73                 if source_os_disk_id is None:
74                     raise TypeError
75                 source_os_disk_type = "SNAPSHOT"
76                 logger.debug("found %s: %s", source_os_disk_type,
77                              source_os_disk_id)
78             except TypeError:
79                 pass
80 
81     if source_os_disk_type is None or source_os_disk_id is None:
82         logger.error(
83             'Unable to locate a supported OS disk type in the provided source object')
84         raise CLIError('Invalid OS Disk Source Type')
85 
86     source_os_type = json_cmd_output['storageProfile']['osDisk']['osType']
87     logger.debug("source_os_disk_type: %s. source_os_disk_id: %s. source_os_type: %s",
88                  source_os_disk_type, source_os_disk_id, source_os_type)
89 
90     # create source snapshots
91     # TODO: skip creating another snapshot when the source is a snapshot
92     logger.warning("Creating source snapshot")
93     source_os_disk_snapshot_name = source_object_name + '_os_disk_snapshot'
94     snapshot_location = json_cmd_output['location']
95     hyper_v_generation = json_cmd_output['hyperVGeneration']
96     if source_os_disk_type == "BLOB":
97         source_storage_account_id = get_storage_account_id_from_blob_path(source_os_disk_id,
98                                                                           source_resource_group_name,
99                                                                           target_subscription)
100         cmd_content = ['snapshot', 'create',
101                        '--name', source_os_disk_snapshot_name,
102                        '--location', snapshot_location,
103                        '--resource-group', source_resource_group_name,
104                        '--source', source_os_disk_id,
105                        '--source-storage-account-id', source_storage_account_id]
106     else:
107         cmd_content = ['snapshot', 'create',
108                        '--name', source_os_disk_snapshot_name,
109                        '--location', snapshot_location,
110                        '--resource-group', source_resource_group_name,
111                        '--source', source_os_disk_id]
112     if hyper_v_generation:
113         cmd_content = cmd_content + ['--hyper-v-generation', hyper_v_generation]
114     cli_cmd = prepare_cli_command(cmd_content, only_show_errors=only_show_errors)
115     run_cli_command(cli_cmd)
116 
117     # Get SAS URL for the snapshotName
118     logger.warning(
119         "Getting sas url for the source snapshot with timeout: %d seconds", timeout)
120     if timeout < 3600:
121         logger.error("Timeout should be greater than 3600 seconds")
122         raise CLIError('Invalid Timeout')
123 
124     cli_cmd = prepare_cli_command(['snapshot', 'grant-access',
125                                    '--name', source_os_disk_snapshot_name,
126                                    '--resource-group', source_resource_group_name,
127                                    '--duration-in-seconds', str(timeout)],
128                                   only_show_errors=only_show_errors)
129 
130     json_output = run_cli_command(cli_cmd, return_as_json=True)
131 
132     source_os_disk_snapshot_url = json_output['accessSas']
133     logger.debug("source os disk snapshot url: %s",
134                  source_os_disk_snapshot_url)
135 
136     # Start processing in the target locations
137 
138     transient_resource_group_name = temporary_resource_group_name
139     # pick the first location for the temp group
140     transient_resource_group_location = target_location[0].strip()
141     create_resource_group(transient_resource_group_name,
142                           transient_resource_group_location,
143                           target_subscription,
144                           only_show_errors)
145 
146     target_locations_count = len(target_location)
147     logger.warning("Target location count: %s", target_locations_count)
148 
149     create_resource_group(target_resource_group_name,
150                           target_location[0].strip(),
151                           target_subscription,
152                           only_show_errors)
153 
154     pool = None
155     try:
156 
157         # try to get a handle on arm's 409s
158         azure_pool_frequency = 5
159         if target_locations_count >= 5:
160             azure_pool_frequency = 15
161         elif target_locations_count >= 3:
162             azure_pool_frequency = 10
163 
164         if (target_locations_count == 1) or (parallel_degree == 1):
165             # Going to copy to targets one-by-one
166             logger.debug("Starting sync process for all locations")
167             for location in target_location:
168                 location = location.strip()
169                 create_target_image(location, transient_resource_group_name, source_type,
170                                     source_object_name, source_os_disk_snapshot_name, source_os_disk_snapshot_url,
171                                     source_os_type, target_resource_group_name, azure_pool_frequency,
172                                     tags, target_name, target_subscription, export_as_snapshot, timeout,
173                                     hyper_v_generation, only_show_errors)
174         else:
175             if parallel_degree == -1:
176                 pool = Pool(target_locations_count)
177             else:
178                 pool = Pool(min(parallel_degree, target_locations_count))
179 
180             tasks = []
181             for location in target_location:
182                 location = location.strip()
183                 task_content = (location, transient_resource_group_name, source_type,
184                                 source_object_name, source_os_disk_snapshot_name, source_os_disk_snapshot_url,
185                                 source_os_type, target_resource_group_name, azure_pool_frequency,
186                                 tags, target_name, target_subscription, export_as_snapshot, timeout,
187                                 hyper_v_generation, only_show_errors)
188                 tasks.append(task_content)
189 
190             logger.warning("Starting async process for all locations")
191 
192             for task in tasks:
193                 pool.apply_async(create_target_image, task)
194 
195             pool.close()
196             pool.join()
197 
198     except KeyboardInterrupt:
199         logger.warning('User cancelled the operation')
200         if cleanup:
201             logger.warning('To cleanup temporary resources look for ones tagged with "image-copy-extension". \n'
202                            'You can use the following command: az resource list --tag created_by=image-copy-extension')
203         if pool is not None:
204             pool.terminate()
205         return
206 
207     # Cleanup
208     if cleanup:
209         logger.warning('Deleting transient resources')
210 
211         # Delete resource group
212         cli_cmd = prepare_cli_command(['group', 'delete', '--no-wait', '--yes',
213                                        '--name', transient_resource_group_name],
214                                       subscription=target_subscription,
215                                       only_show_errors=only_show_errors)
216         run_cli_command(cli_cmd)
217 
218         # Revoke sas for source snapshot
219         cli_cmd = prepare_cli_command(['snapshot', 'revoke-access',
220                                        '--name', source_os_disk_snapshot_name,
221                                        '--resource-group', source_resource_group_name],
222                                       only_show_errors=only_show_errors)
223         run_cli_command(cli_cmd)
224 
225         # Delete source snapshot
226         # TODO: skip this if source is snapshot and not creating a new one
227         cli_cmd = prepare_cli_command(['snapshot', 'delete',
228                                        '--name', source_os_disk_snapshot_name,
229                                        '--resource-group', source_resource_group_name],
230                                       only_show_errors=only_show_errors)
231         run_cli_command(cli_cmd)
232 
233     logger.warning('Image copy finished')
234 
235 
236 def create_resource_group(resource_group_name, location, subscription=None, only_show_errors=None):
237     # check if target resource group exists
238     cli_cmd = prepare_cli_command(['group', 'exists',
239                                    '--name', resource_group_name],
240                                   output_as_json=False,
241                                   subscription=subscription,
242                                   only_show_errors=only_show_errors)
243 
244     cmd_output = run_cli_command(cli_cmd)
245 
246     if 'true' in cmd_output:
247         return
248 
249     # create the target resource group
250     logger.warning("Creating resource group: %s", resource_group_name)
251     cli_cmd = prepare_cli_command(['group', 'create',
252                                    '--name', resource_group_name,
253                                    '--location', location],
254                                   subscription=subscription,
255                                   only_show_errors=only_show_errors)
256 
257     run_cli_command(cli_cmd)
258
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/image-copy/azext_imagecopy/custom.py b/src/image-copy/azext_imagecopy/custom.py
--- a/src/image-copy/azext_imagecopy/custom.py
+++ b/src/image-copy/azext_imagecopy/custom.py
@@ -33,8 +33,9 @@
             raise CLIError('Don\'t specify an existing resource group in --temporary-resource-group-name '
                            'when --cleanup is set')
 
-    from azure.cli.core.commands.client_factory import get_subscription_id
-    target_subscription = get_subscription_id(cmd.cli_ctx)
+    if not target_subscription:
+        from azure.cli.core.commands.client_factory import get_subscription_id
+        target_subscription = get_subscription_id(cmd.cli_ctx)
     logger.debug('subscription id - %s', target_subscription)
 
     # get the os disk id from source vm/image

{"golden_diff": "diff --git a/src/image-copy/azext_imagecopy/custom.py b/src/image-copy/azext_imagecopy/custom.py\n--- a/src/image-copy/azext_imagecopy/custom.py\n+++ b/src/image-copy/azext_imagecopy/custom.py\n@@ -33,8 +33,9 @@\n             raise CLIError('Don\\'t specify an existing resource group in --temporary-resource-group-name '\n                            'when --cleanup is set')\n \n-    from azure.cli.core.commands.client_factory import get_subscription_id\n-    target_subscription = get_subscription_id(cmd.cli_ctx)\n+    if not target_subscription:\n+        from azure.cli.core.commands.client_factory import get_subscription_id\n+        target_subscription = get_subscription_id(cmd.cli_ctx)\n     logger.debug('subscription id - %s', target_subscription)\n \n     # get the os disk id from source vm/image\n", "issue": "Image Copy '--target-subscription' Argument Ignored\n### Related command\r\n<!--- Please provide the related command with az {command} if you can, so that we can quickly route to the related team for help. --->\r\n```bash\r\naz image copy --source-resource-group <source_rg_name> \\\r\n        --source-object-name <source_image_name> \\\r\n        --source-type image \\\r\n        --target-subscription <target_subscription_id> \\\r\n        --target-location <target_location> \\\r\n        --target-resource-group <target_rg_name>  \\\r\n        --target-name <target_image_name> \\\r\n        --temporary-resource-group-name <temp_rg_name>\r\n```\r\n\r\n### Extension name (the extension in question)\r\nimage-copy-extension\r\n\r\n### Description of issue (in as much detail as possible)\r\n\r\nMy organization is using the `image-copy-extension` and we noticed on recent workflows that the extension is now ignoring the `--target-subscription` argument and seems to use the authentication context subscription as the target subscription for all temporary resource and final image creation instead.\r\n\r\nWe noticed this change because previously the temporary resource group we used was stored within the target subscription and this meant that the temporary storage account, blob, and snapshot all got created within this temporary resource group but now the extension seems to be looking for the temporary resource group in the default subscription associated with the authentication context instead of the target subscription provided to the `--target-subscription` argument.\r\n\r\nMy Python experience is still limited but I have looked at the [recent commit](https://github.com/Azure/azure-cli-extensions/commit/3a6621cc0c96aab5ddd2388f4a2e6a81f4ed2f0a) and I see that the `target_subscription` variable is now being set without checking if the argument has been passed and I'm not sure if this was changed for a reason or if this is just a bug that can be fixed by only setting the `target_subscription` variable if the argument is null.\r\n\r\n![image](https://user-images.githubusercontent.com/20447447/209253966-f9129522-3988-4c73-9abc-d6fc0e7ef991.png)\r\n\r\nLink to the line in question.\r\n\r\nhttps://github.com/Azure/azure-cli-extensions/blob/42102c3a81ad23407e7d929521ea212805b4688d/src/image-copy/azext_imagecopy/custom.py#L37\r\n\r\nThere is also a debug line below the variable declaration that we've seen continues to display the subscription ID for the authentication context and not the target subscription we're providing as an argument.\r\n\r\n-----\r\n\r\n\n", "before_files": [{"content": "# --------------------------------------------------------------------------------------------\n# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License. See License.txt in the project root for license information.\n# --------------------------------------------------------------------------------------------\n\nfrom multiprocessing import Pool\n\nfrom knack.util import CLIError\nfrom knack.log import get_logger\n\nfrom azext_imagecopy.cli_utils import run_cli_command, prepare_cli_command, get_storage_account_id_from_blob_path\nfrom azext_imagecopy.create_target import create_target_image\n\nlogger = get_logger(__name__)\n\n\n# pylint: disable=too-many-statements\n# pylint: disable=too-many-locals\n# pylint: disable=too-many-branches\ndef imagecopy(cmd, source_resource_group_name, source_object_name, target_location,\n              target_resource_group_name, temporary_resource_group_name='image-copy-rg',\n              source_type='image', cleanup=False, parallel_degree=-1, tags=None, target_name=None,\n              target_subscription=None, export_as_snapshot='false', timeout=3600):\n    only_show_errors = cmd.cli_ctx.only_show_errors\n    if cleanup:\n        # If --cleanup is set, forbid using an existing temporary resource group name.\n        # It is dangerous to clean up an existing resource group.\n        cli_cmd = prepare_cli_command(['group', 'exists', '-n', temporary_resource_group_name],\n                                      output_as_json=False,\n                                      only_show_errors=only_show_errors)\n        cmd_output = run_cli_command(cli_cmd)\n        if 'true' in cmd_output:\n            raise CLIError('Don\\'t specify an existing resource group in --temporary-resource-group-name '\n                           'when --cleanup is set')\n\n    from azure.cli.core.commands.client_factory import get_subscription_id\n    target_subscription = get_subscription_id(cmd.cli_ctx)\n    logger.debug('subscription id - %s', target_subscription)\n\n    # get the os disk id from source vm/image\n    logger.warning(\"Getting OS disk ID of the source VM/image\")\n    cli_cmd = prepare_cli_command([source_type, 'show',\n                                   '--name', source_object_name,\n                                   '--resource-group', source_resource_group_name],\n                                  only_show_errors=only_show_errors)\n\n    json_cmd_output = run_cli_command(cli_cmd, return_as_json=True)\n\n    if json_cmd_output['storageProfile']['dataDisks']:\n        logger.warning(\n            \"Data disks in the source detected, but are ignored by this extension!\")\n\n    source_os_disk_id = None\n    source_os_disk_type = None\n\n    try:\n        source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['managedDisk']['id']\n        if source_os_disk_id is None:\n            raise TypeError\n        source_os_disk_type = \"DISK\"\n        logger.debug(\"found %s: %s\", source_os_disk_type, source_os_disk_id)\n    except TypeError:\n        try:\n            source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['blobUri']\n            if source_os_disk_id is None:\n                raise TypeError\n            source_os_disk_type = \"BLOB\"\n            logger.debug(\"found %s: %s\", source_os_disk_type,\n                         source_os_disk_id)\n        except TypeError:\n            try:  # images created by e.g. image-copy extension\n                source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['snapshot']['id']\n                if source_os_disk_id is None:\n                    raise TypeError\n                source_os_disk_type = \"SNAPSHOT\"\n                logger.debug(\"found %s: %s\", source_os_disk_type,\n                             source_os_disk_id)\n            except TypeError:\n                pass\n\n    if source_os_disk_type is None or source_os_disk_id is None:\n        logger.error(\n            'Unable to locate a supported OS disk type in the provided source object')\n        raise CLIError('Invalid OS Disk Source Type')\n\n    source_os_type = json_cmd_output['storageProfile']['osDisk']['osType']\n    logger.debug(\"source_os_disk_type: %s. source_os_disk_id: %s. source_os_type: %s\",\n                 source_os_disk_type, source_os_disk_id, source_os_type)\n\n    # create source snapshots\n    # TODO: skip creating another snapshot when the source is a snapshot\n    logger.warning(\"Creating source snapshot\")\n    source_os_disk_snapshot_name = source_object_name + '_os_disk_snapshot'\n    snapshot_location = json_cmd_output['location']\n    hyper_v_generation = json_cmd_output['hyperVGeneration']\n    if source_os_disk_type == \"BLOB\":\n        source_storage_account_id = get_storage_account_id_from_blob_path(source_os_disk_id,\n                                                                          source_resource_group_name,\n                                                                          target_subscription)\n        cmd_content = ['snapshot', 'create',\n                       '--name', source_os_disk_snapshot_name,\n                       '--location', snapshot_location,\n                       '--resource-group', source_resource_group_name,\n                       '--source', source_os_disk_id,\n                       '--source-storage-account-id', source_storage_account_id]\n    else:\n        cmd_content = ['snapshot', 'create',\n                       '--name', source_os_disk_snapshot_name,\n                       '--location', snapshot_location,\n                       '--resource-group', source_resource_group_name,\n                       '--source', source_os_disk_id]\n    if hyper_v_generation:\n        cmd_content = cmd_content + ['--hyper-v-generation', hyper_v_generation]\n    cli_cmd = prepare_cli_command(cmd_content, only_show_errors=only_show_errors)\n    run_cli_command(cli_cmd)\n\n    # Get SAS URL for the snapshotName\n    logger.warning(\n        \"Getting sas url for the source snapshot with timeout: %d seconds\", timeout)\n    if timeout < 3600:\n        logger.error(\"Timeout should be greater than 3600 seconds\")\n        raise CLIError('Invalid Timeout')\n\n    cli_cmd = prepare_cli_command(['snapshot', 'grant-access',\n                                   '--name', source_os_disk_snapshot_name,\n                                   '--resource-group', source_resource_group_name,\n                                   '--duration-in-seconds', str(timeout)],\n                                  only_show_errors=only_show_errors)\n\n    json_output = run_cli_command(cli_cmd, return_as_json=True)\n\n    source_os_disk_snapshot_url = json_output['accessSas']\n    logger.debug(\"source os disk snapshot url: %s\",\n                 source_os_disk_snapshot_url)\n\n    # Start processing in the target locations\n\n    transient_resource_group_name = temporary_resource_group_name\n    # pick the first location for the temp group\n    transient_resource_group_location = target_location[0].strip()\n    create_resource_group(transient_resource_group_name,\n                          transient_resource_group_location,\n                          target_subscription,\n                          only_show_errors)\n\n    target_locations_count = len(target_location)\n    logger.warning(\"Target location count: %s\", target_locations_count)\n\n    create_resource_group(target_resource_group_name,\n                          target_location[0].strip(),\n                          target_subscription,\n                          only_show_errors)\n\n    pool = None\n    try:\n\n        # try to get a handle on arm's 409s\n        azure_pool_frequency = 5\n        if target_locations_count >= 5:\n            azure_pool_frequency = 15\n        elif target_locations_count >= 3:\n            azure_pool_frequency = 10\n\n        if (target_locations_count == 1) or (parallel_degree == 1):\n            # Going to copy to targets one-by-one\n            logger.debug(\"Starting sync process for all locations\")\n            for location in target_location:\n                location = location.strip()\n                create_target_image(location, transient_resource_group_name, source_type,\n                                    source_object_name, source_os_disk_snapshot_name, source_os_disk_snapshot_url,\n                                    source_os_type, target_resource_group_name, azure_pool_frequency,\n                                    tags, target_name, target_subscription, export_as_snapshot, timeout,\n                                    hyper_v_generation, only_show_errors)\n        else:\n            if parallel_degree == -1:\n                pool = Pool(target_locations_count)\n            else:\n                pool = Pool(min(parallel_degree, target_locations_count))\n\n            tasks = []\n            for location in target_location:\n                location = location.strip()\n                task_content = (location, transient_resource_group_name, source_type,\n                                source_object_name, source_os_disk_snapshot_name, source_os_disk_snapshot_url,\n                                source_os_type, target_resource_group_name, azure_pool_frequency,\n                                tags, target_name, target_subscription, export_as_snapshot, timeout,\n                                hyper_v_generation, only_show_errors)\n                tasks.append(task_content)\n\n            logger.warning(\"Starting async process for all locations\")\n\n            for task in tasks:\n                pool.apply_async(create_target_image, task)\n\n            pool.close()\n            pool.join()\n\n    except KeyboardInterrupt:\n        logger.warning('User cancelled the operation')\n        if cleanup:\n            logger.warning('To cleanup temporary resources look for ones tagged with \"image-copy-extension\". \\n'\n                           'You can use the following command: az resource list --tag created_by=image-copy-extension')\n        if pool is not None:\n            pool.terminate()\n        return\n\n    # Cleanup\n    if cleanup:\n        logger.warning('Deleting transient resources')\n\n        # Delete resource group\n        cli_cmd = prepare_cli_command(['group', 'delete', '--no-wait', '--yes',\n                                       '--name', transient_resource_group_name],\n                                      subscription=target_subscription,\n                                      only_show_errors=only_show_errors)\n        run_cli_command(cli_cmd)\n\n        # Revoke sas for source snapshot\n        cli_cmd = prepare_cli_command(['snapshot', 'revoke-access',\n                                       '--name', source_os_disk_snapshot_name,\n                                       '--resource-group', source_resource_group_name],\n                                      only_show_errors=only_show_errors)\n        run_cli_command(cli_cmd)\n\n        # Delete source snapshot\n        # TODO: skip this if source is snapshot and not creating a new one\n        cli_cmd = prepare_cli_command(['snapshot', 'delete',\n                                       '--name', source_os_disk_snapshot_name,\n                                       '--resource-group', source_resource_group_name],\n                                      only_show_errors=only_show_errors)\n        run_cli_command(cli_cmd)\n\n    logger.warning('Image copy finished')\n\n\ndef create_resource_group(resource_group_name, location, subscription=None, only_show_errors=None):\n    # check if target resource group exists\n    cli_cmd = prepare_cli_command(['group', 'exists',\n                                   '--name', resource_group_name],\n                                  output_as_json=False,\n                                  subscription=subscription,\n                                  only_show_errors=only_show_errors)\n\n    cmd_output = run_cli_command(cli_cmd)\n\n    if 'true' in cmd_output:\n        return\n\n    # create the target resource group\n    logger.warning(\"Creating resource group: %s\", resource_group_name)\n    cli_cmd = prepare_cli_command(['group', 'create',\n                                   '--name', resource_group_name,\n                                   '--location', location],\n                                  subscription=subscription,\n                                  only_show_errors=only_show_errors)\n\n    run_cli_command(cli_cmd)\n", "path": "src/image-copy/azext_imagecopy/custom.py"}], "after_files": [{"content": "# --------------------------------------------------------------------------------------------\n# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License. See License.txt in the project root for license information.\n# --------------------------------------------------------------------------------------------\n\nfrom multiprocessing import Pool\n\nfrom knack.util import CLIError\nfrom knack.log import get_logger\n\nfrom azext_imagecopy.cli_utils import run_cli_command, prepare_cli_command, get_storage_account_id_from_blob_path\nfrom azext_imagecopy.create_target import create_target_image\n\nlogger = get_logger(__name__)\n\n\n# pylint: disable=too-many-statements\n# pylint: disable=too-many-locals\n# pylint: disable=too-many-branches\ndef imagecopy(cmd, source_resource_group_name, source_object_name, target_location,\n              target_resource_group_name, temporary_resource_group_name='image-copy-rg',\n              source_type='image', cleanup=False, parallel_degree=-1, tags=None, target_name=None,\n              target_subscription=None, export_as_snapshot='false', timeout=3600):\n    only_show_errors = cmd.cli_ctx.only_show_errors\n    if cleanup:\n        # If --cleanup is set, forbid using an existing temporary resource group name.\n        # It is dangerous to clean up an existing resource group.\n        cli_cmd = prepare_cli_command(['group', 'exists', '-n', temporary_resource_group_name],\n                                      output_as_json=False,\n                                      only_show_errors=only_show_errors)\n        cmd_output = run_cli_command(cli_cmd)\n        if 'true' in cmd_output:\n            raise CLIError('Don\\'t specify an existing resource group in --temporary-resource-group-name '\n                           'when --cleanup is set')\n\n    if not target_subscription:\n        from azure.cli.core.commands.client_factory import get_subscription_id\n        target_subscription = get_subscription_id(cmd.cli_ctx)\n    logger.debug('subscription id - %s', target_subscription)\n\n    # get the os disk id from source vm/image\n    logger.warning(\"Getting OS disk ID of the source VM/image\")\n    cli_cmd = prepare_cli_command([source_type, 'show',\n                                   '--name', source_object_name,\n                                   '--resource-group', source_resource_group_name],\n                                  only_show_errors=only_show_errors)\n\n    json_cmd_output = run_cli_command(cli_cmd, return_as_json=True)\n\n    if json_cmd_output['storageProfile']['dataDisks']:\n        logger.warning(\n            \"Data disks in the source detected, but are ignored by this extension!\")\n\n    source_os_disk_id = None\n    source_os_disk_type = None\n\n    try:\n        source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['managedDisk']['id']\n        if source_os_disk_id is None:\n            raise TypeError\n        source_os_disk_type = \"DISK\"\n        logger.debug(\"found %s: %s\", source_os_disk_type, source_os_disk_id)\n    except TypeError:\n        try:\n            source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['blobUri']\n            if source_os_disk_id is None:\n                raise TypeError\n            source_os_disk_type = \"BLOB\"\n            logger.debug(\"found %s: %s\", source_os_disk_type,\n                         source_os_disk_id)\n        except TypeError:\n            try:  # images created by e.g. image-copy extension\n                source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['snapshot']['id']\n                if source_os_disk_id is None:\n                    raise TypeError\n                source_os_disk_type = \"SNAPSHOT\"\n                logger.debug(\"found %s: %s\", source_os_disk_type,\n                             source_os_disk_id)\n            except TypeError:\n                pass\n\n    if source_os_disk_type is None or source_os_disk_id is None:\n        logger.error(\n            'Unable to locate a supported OS disk type in the provided source object')\n        raise CLIError('Invalid OS Disk Source Type')\n\n    source_os_type = json_cmd_output['storageProfile']['osDisk']['osType']\n    logger.debug(\"source_os_disk_type: %s. source_os_disk_id: %s. source_os_type: %s\",\n                 source_os_disk_type, source_os_disk_id, source_os_type)\n\n    # create source snapshots\n    # TODO: skip creating another snapshot when the source is a snapshot\n    logger.warning(\"Creating source snapshot\")\n    source_os_disk_snapshot_name = source_object_name + '_os_disk_snapshot'\n    snapshot_location = json_cmd_output['location']\n    hyper_v_generation = json_cmd_output['hyperVGeneration']\n    if source_os_disk_type == \"BLOB\":\n        source_storage_account_id = get_storage_account_id_from_blob_path(source_os_disk_id,\n                                                                          source_resource_group_name,\n                                                                          target_subscription)\n        cmd_content = ['snapshot', 'create',\n                       '--name', source_os_disk_snapshot_name,\n                       '--location', snapshot_location,\n                       '--resource-group', source_resource_group_name,\n                       '--source', source_os_disk_id,\n                       '--source-storage-account-id', source_storage_account_id]\n    else:\n        cmd_content = ['snapshot', 'create',\n                       '--name', source_os_disk_snapshot_name,\n                       '--location', snapshot_location,\n                       '--resource-group', source_resource_group_name,\n                       '--source', source_os_disk_id]\n    if hyper_v_generation:\n        cmd_content = cmd_content + ['--hyper-v-generation', hyper_v_generation]\n    cli_cmd = prepare_cli_command(cmd_content, only_show_errors=only_show_errors)\n    run_cli_command(cli_cmd)\n\n    # Get SAS URL for the snapshotName\n    logger.warning(\n        \"Getting sas url for the source snapshot with timeout: %d seconds\", timeout)\n    if timeout < 3600:\n        logger.error(\"Timeout should be greater than 3600 seconds\")\n        raise CLIError('Invalid Timeout')\n\n    cli_cmd = prepare_cli_command(['snapshot', 'grant-access',\n                                   '--name', source_os_disk_snapshot_name,\n                                   '--resource-group', source_resource_group_name,\n                                   '--duration-in-seconds', str(timeout)],\n                                  only_show_errors=only_show_errors)\n\n    json_output = run_cli_command(cli_cmd, return_as_json=True)\n\n    source_os_disk_snapshot_url = json_output['accessSas']\n    logger.debug(\"source os disk snapshot url: %s\",\n                 source_os_disk_snapshot_url)\n\n    # Start processing in the target locations\n\n    transient_resource_group_name = temporary_resource_group_name\n    # pick the first location for the temp group\n    transient_resource_group_location = target_location[0].strip()\n    create_resource_group(transient_resource_group_name,\n                          transient_resource_group_location,\n                          target_subscription,\n                          only_show_errors)\n\n    target_locations_count = len(target_location)\n    logger.warning(\"Target location count: %s\", target_locations_count)\n\n    create_resource_group(target_resource_group_name,\n                          target_location[0].strip(),\n                          target_subscription,\n                          only_show_errors)\n\n    pool = None\n    try:\n\n        # try to get a handle on arm's 409s\n        azure_pool_frequency = 5\n        if target_locations_count >= 5:\n            azure_pool_frequency = 15\n        elif target_locations_count >= 3:\n            azure_pool_frequency = 10\n\n        if (target_locations_count == 1) or (parallel_degree == 1):\n            # Going to copy to targets one-by-one\n            logger.debug(\"Starting sync process for all locations\")\n            for location in target_location:\n                location = location.strip()\n                create_target_image(location, transient_resource_group_name, source_type,\n                                    source_object_name, source_os_disk_snapshot_name, source_os_disk_snapshot_url,\n                                    source_os_type, target_resource_group_name, azure_pool_frequency,\n                                    tags, target_name, target_subscription, export_as_snapshot, timeout,\n                                    hyper_v_generation, only_show_errors)\n        else:\n            if parallel_degree == -1:\n                pool = Pool(target_locations_count)\n            else:\n                pool = Pool(min(parallel_degree, target_locations_count))\n\n            tasks = []\n            for location in target_location:\n                location = location.strip()\n                task_content = (location, transient_resource_group_name, source_type,\n                                source_object_name, source_os_disk_snapshot_name, source_os_disk_snapshot_url,\n                                source_os_type, target_resource_group_name, azure_pool_frequency,\n                                tags, target_name, target_subscription, export_as_snapshot, timeout,\n                                hyper_v_generation, only_show_errors)\n                tasks.append(task_content)\n\n            logger.warning(\"Starting async process for all locations\")\n\n            for task in tasks:\n                pool.apply_async(create_target_image, task)\n\n            pool.close()\n            pool.join()\n\n    except KeyboardInterrupt:\n        logger.warning('User cancelled the operation')\n        if cleanup:\n            logger.warning('To cleanup temporary resources look for ones tagged with \"image-copy-extension\". \\n'\n                           'You can use the following command: az resource list --tag created_by=image-copy-extension')\n        if pool is not None:\n            pool.terminate()\n        return\n\n    # Cleanup\n    if cleanup:\n        logger.warning('Deleting transient resources')\n\n        # Delete resource group\n        cli_cmd = prepare_cli_command(['group', 'delete', '--no-wait', '--yes',\n                                       '--name', transient_resource_group_name],\n                                      subscription=target_subscription,\n                                      only_show_errors=only_show_errors)\n        run_cli_command(cli_cmd)\n\n        # Revoke sas for source snapshot\n        cli_cmd = prepare_cli_command(['snapshot', 'revoke-access',\n                                       '--name', source_os_disk_snapshot_name,\n                                       '--resource-group', source_resource_group_name],\n                                      only_show_errors=only_show_errors)\n        run_cli_command(cli_cmd)\n\n        # Delete source snapshot\n        # TODO: skip this if source is snapshot and not creating a new one\n        cli_cmd = prepare_cli_command(['snapshot', 'delete',\n                                       '--name', source_os_disk_snapshot_name,\n                                       '--resource-group', source_resource_group_name],\n                                      only_show_errors=only_show_errors)\n        run_cli_command(cli_cmd)\n\n    logger.warning('Image copy finished')\n\n\ndef create_resource_group(resource_group_name, location, subscription=None, only_show_errors=None):\n    # check if target resource group exists\n    cli_cmd = prepare_cli_command(['group', 'exists',\n                                   '--name', resource_group_name],\n                                  output_as_json=False,\n                                  subscription=subscription,\n                                  only_show_errors=only_show_errors)\n\n    cmd_output = run_cli_command(cli_cmd)\n\n    if 'true' in cmd_output:\n        return\n\n    # create the target resource group\n    logger.warning(\"Creating resource group: %s\", resource_group_name)\n    cli_cmd = prepare_cli_command(['group', 'create',\n                                   '--name', resource_group_name,\n                                   '--location', location],\n                                  subscription=subscription,\n                                  only_show_errors=only_show_errors)\n\n    run_cli_command(cli_cmd)\n", "path": "src/image-copy/azext_imagecopy/custom.py"}]}