From c17cd75161ead9a8c84c583dcbc706a5b40514fc Mon Sep 17 00:00:00 2001 From: houk-ms Date: Wed, 10 Mar 2021 15:58:13 +0800 Subject: [PATCH 1/6] security-domain sync to async support --- .../v7_2/key_vault_client.py | 86 ++++++++++++++++++- 1 file changed, 84 insertions(+), 2 deletions(-) diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/vendored_sdks/azure_keyvault_t1/v7_2/key_vault_client.py b/src/azure-cli/azure/cli/command_modules/keyvault/vendored_sdks/azure_keyvault_t1/v7_2/key_vault_client.py index 7ad7a4b3e91..fffe115ef71 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/vendored_sdks/azure_keyvault_t1/v7_2/key_vault_client.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/vendored_sdks/azure_keyvault_t1/v7_2/key_vault_client.py @@ -11,6 +11,7 @@ # pylint: skip-file # flake8: noqa import json +import time from msrest.service_client import SDKClient from msrest import Serializer, Deserializer @@ -114,20 +115,101 @@ def download(self, vault_base_url, certificates, custom_headers=None, raw=False, response = self._client.send( request, header_parameters, body_content, stream=False, **operation_config) + # v7.2-preview and v7.2 will introduce a breaking change to make the operation change from Sync to Async + # 200: for compatability of response before the change (Sync Operation) + # 202: for the support of new response after the change (Async Operation) + if response.status_code not in [200, 202]: + raise models.KeyVaultErrorException(self._deserialize, response) + + deserialized = None + + # for both old response, and new response + deserialized = self._deserialize('SecurityDomainObject', response) + + # for new response + if response.status_code == 202: + polling_interval = int(response.headers.get('retry-after', 1)) + status = 'InProgress' + + # keep polling if status is 'InProgress' + while status == 'InProgress': + # flush + time.sleep(0.5*polling_interval) + print('\r - Running .. ', end='') + time.sleep(0.5*polling_interval) + print('\r ', end='') + + # polling + operation_status = self.download_pending( + vault_base_url, + custom_headers=custom_headers, + raw=False, + **operation_config + ) + status = operation_status.status + + # Response won't be returned if the deployment fails + if status != 'Success': + raise models.KeyVaultErrorException(self._deserialize, response) + + if raw: + client_raw_response = ClientRawResponse(deserialized, response) + return client_raw_response + + return deserialized + download.metadata = {'url': '/securitydomain/download'} + + def download_pending(self, vault_base_url, custom_headers=None, raw=False, **operation_config): + """Get Security domain upload operation status. + :param vault_base_url: The vault name, for example https://myvault.vault.azure.net. + :type vault_base_url: str + :keyword callable cls: A custom type or function that will be passed the direct response + :return: SecurityDomainOperationStatus, or the result of cls(response) + :rtype: ~key_vault_client.models.SecurityDomainOperationStatus + :raises: ~azure.core.exceptions.HttpResponseError + """ + + # Construct URL + url = self.upload_pending.metadata['url'] + path_format_arguments = { + 'vaultBaseUrl': self._serialize.url("vault_base_url", vault_base_url, 'str', skip_quote=True) + } + url = self._client.format_url(url, **path_format_arguments) + + # Construct parameters + query_parameters = {} + query_parameters['api-version'] = self._serialize.query("self.api_version", self.api_version, 'str') + + # Construct headers + header_parameters = {} + header_parameters['Content-Type'] = 'application/json; charset=utf-8' + if self.config.generate_client_request_id: + header_parameters['x-ms-client-request-id'] = str(uuid.uuid1()) + if custom_headers: + header_parameters.update(custom_headers) + if self.config.accept_language is not None: + header_parameters['accept-language'] = self._serialize.header("self.config.accept_language", + self.config.accept_language, 'str') + + # Construct and send request + request = self._client.get(url, query_parameters) + response = self._client.send( + request, header_parameters, stream=False, **operation_config) + if response.status_code not in [200]: raise models.KeyVaultErrorException(self._deserialize, response) deserialized = None if response.status_code == 200: - deserialized = self._deserialize('SecurityDomainObject', response) + deserialized = self._deserialize('SecurityDomainOperationStatus', response) if raw: client_raw_response = ClientRawResponse(deserialized, response) return client_raw_response return deserialized - download.metadata = {'url': '/securitydomain/download'} + download_pending.metadata = {'url': '/securitydomain/download/pending'} def transfer_key(self, vault_base_url, custom_headers=None, raw=False, **operation_config): """Retrieve security domain transfer key. From 2fce0baef9c1948268583d60d6b9bd48ca2bd6ca Mon Sep 17 00:00:00 2001 From: houk-ms Date: Fri, 12 Mar 2021 14:21:46 +0800 Subject: [PATCH 2/6] security-domain download async --- .../cli/command_modules/keyvault/_params.py | 6 +++ .../cli/command_modules/keyvault/commands.py | 2 +- .../cli/command_modules/keyvault/custom.py | 40 ++++++++++++++----- .../v7_2/key_vault_client.py | 31 ++------------ 4 files changed, 39 insertions(+), 40 deletions(-) diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/_params.py b/src/azure-cli/azure/cli/command_modules/keyvault/_params.py index 8da92e09dbd..1f5b9d550ca 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/_params.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/_params.py @@ -77,6 +77,10 @@ class CLIJsonWebKeyCurveName(str, Enum): p_384 = "P-384" #: The NIST P-384 elliptic curve, AKA SECG curve SECP384R1. p_521 = "P-521" #: The NIST P-521 elliptic curve, AKA SECG curve SECP521R1. + class CLISecurityDomainOperation(str, Enum): + download = "download" #: Download operation + upload = "upload" #: Upload operation + (KeyPermissions, SecretPermissions, CertificatePermissions, StoragePermissions, NetworkRuleBypassOptions, NetworkRuleAction) = self.get_models( 'KeyPermissions', 'SecretPermissions', 'CertificatePermissions', 'StoragePermissions', @@ -496,6 +500,8 @@ class CLIJsonWebKeyCurveName(str, Enum): c.argument('identifier', options_list=['--id'], validator=validate_vault_or_hsm, help='Id of the HSM.') c.argument('resource_group_name', options_list=['--resource-group', '-g'], help='Proceed only if HSM belongs to the specified resource group.') + c.argument('target_operation', arg_type=get_enum_type(CLISecurityDomainOperation), + help='Target operation that needs waiting.') # endregion # region keyvault backup/restore diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/commands.py b/src/azure-cli/azure/cli/command_modules/keyvault/commands.py index 0c11d80a4ab..877ad3a81a3 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/commands.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/commands.py @@ -139,7 +139,7 @@ def load_command_table(self, _): is_preview=True) as g: g.keyvault_custom('init-recovery', 'security_domain_init_recovery') g.keyvault_custom('upload', 'security_domain_upload', supports_no_wait=True) - g.keyvault_custom('download', 'security_domain_download') + g.keyvault_custom('download', 'security_domain_download', supports_no_wait=True) g.keyvault_custom('wait', '_wait_security_domain_operation') with self.command_group('keyvault key', data_entity.command_type) as g: diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py index e0bfee612a8..cefc3b90fed 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py @@ -2204,14 +2204,21 @@ def get_x5c_as_pem(): raise ex -def _wait_security_domain_operation(client, hsm_name, identifier=None): # pylint: disable=unused-argument +def _wait_security_domain_operation(client, hsm_name, target_operation='upload', identifier=None): # pylint: disable=unused-argument retries = 0 max_retries = 30 wait_second = 5 while retries < max_retries: try: - ret = client.upload_pending(vault_base_url=hsm_name) - if ret and getattr(ret, 'status', None) in ['Succeeded', 'Failed']: + ret = None + if target_operation == 'upload': + ret = client.upload_pending(vault_base_url=hsm_name) + elif target_operation == 'download': + ret = client.upload_pending(vault_base_url=hsm_name) + + # v7.2-preview and v7.2 will change the upload operation from Sync to Async + # due to service defects, it returns 'Succeeded' before the change and 'Success' after the change + if ret and getattr(ret, 'status', None) in ['Succeeded', 'Success', 'Failed']: return ret except: # pylint: disable=bare-except pass @@ -2356,14 +2363,14 @@ def security_domain_upload(cmd, client, hsm_name, sd_file, sd_exchange_key, sd_w if no_wait: return retval - new_retval = _wait_security_domain_operation(client, hsm_name) + new_retval = _wait_security_domain_operation(client, hsm_name, 'upload') if new_retval: return new_retval return retval def security_domain_download(cmd, client, hsm_name, sd_wrapping_keys, security_domain_file, sd_quorum, - identifier=None, vault_base_url=None): # pylint: disable=unused-argument + identifier=None, vault_base_url=None, no_wait=False): # pylint: disable=unused-argument if os.path.exists(security_domain_file): raise CLIError("File named '{}' already exists.".format(security_domain_file)) @@ -2406,15 +2413,26 @@ def security_domain_download(cmd, client, hsm_name, sd_wrapping_keys, security_d certificates.append(sd_jwk) + # save security-domain backup value to local file + def _save_to_local_file(file_path, security_domain): + try: + with open(file_path, 'w') as f: + f.write(security_domain.value) + except: # pylint: disable=bare-except + if os.path.isfile(file_path): + os.remove(file_path) + ret = client.download( vault_base_url=hsm_name or vault_base_url, certificates=CertificateSet(certificates=certificates, required=sd_quorum) ) - try: - with open(security_domain_file, 'w') as f: - f.write(ret.value) - except: # pylint: disable=bare-except - if os.path.isfile(security_domain_file): - os.remove(security_domain_file) + if not no_wait: + polling_ret = _wait_security_domain_operation(client, hsm_name, 'download') + # Due to service defect, status could be 'Success' or 'Succeeded' when it succeeded + if polling_ret and getattr(polling_ret, 'status', None) != 'Failed': + _save_to_local_file(security_domain_file, ret) + return polling_ret + + _save_to_local_file(security_domain_file, ret) # endregion diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/vendored_sdks/azure_keyvault_t1/v7_2/key_vault_client.py b/src/azure-cli/azure/cli/command_modules/keyvault/vendored_sdks/azure_keyvault_t1/v7_2/key_vault_client.py index fffe115ef71..a2496ee5236 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/vendored_sdks/azure_keyvault_t1/v7_2/key_vault_client.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/vendored_sdks/azure_keyvault_t1/v7_2/key_vault_client.py @@ -123,34 +123,9 @@ def download(self, vault_base_url, certificates, custom_headers=None, raw=False, deserialized = None - # for both old response, and new response - deserialized = self._deserialize('SecurityDomainObject', response) - - # for new response - if response.status_code == 202: - polling_interval = int(response.headers.get('retry-after', 1)) - status = 'InProgress' - - # keep polling if status is 'InProgress' - while status == 'InProgress': - # flush - time.sleep(0.5*polling_interval) - print('\r - Running .. ', end='') - time.sleep(0.5*polling_interval) - print('\r ', end='') - - # polling - operation_status = self.download_pending( - vault_base_url, - custom_headers=custom_headers, - raw=False, - **operation_config - ) - status = operation_status.status - - # Response won't be returned if the deployment fails - if status != 'Success': - raise models.KeyVaultErrorException(self._deserialize, response) + # for both old response and new response + if response.status_code in [200, 202]: + deserialized = self._deserialize('SecurityDomainObject', response) if raw: client_raw_response = ClientRawResponse(deserialized, response) From 8a042c1a39c98942c363a282fa8a4fa1bfd476d2 Mon Sep 17 00:00:00 2001 From: houk-ms Date: Fri, 12 Mar 2021 14:44:48 +0800 Subject: [PATCH 3/6] fix polling bug --- src/azure-cli/azure/cli/command_modules/keyvault/custom.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py index cefc3b90fed..25751033b97 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py @@ -2214,7 +2214,7 @@ def _wait_security_domain_operation(client, hsm_name, target_operation='upload', if target_operation == 'upload': ret = client.upload_pending(vault_base_url=hsm_name) elif target_operation == 'download': - ret = client.upload_pending(vault_base_url=hsm_name) + ret = client.download_pending(vault_base_url=hsm_name) # v7.2-preview and v7.2 will change the upload operation from Sync to Async # due to service defects, it returns 'Succeeded' before the change and 'Success' after the change From 0d9c3378d0c4c3ae7b87c4475c5a2ff9a13e5ca3 Mon Sep 17 00:00:00 2001 From: houk-ms Date: Fri, 12 Mar 2021 16:16:54 +0800 Subject: [PATCH 4/6] raise error when file operation fails --- src/azure-cli/azure/cli/command_modules/keyvault/custom.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py index 25751033b97..b0dc98e8c3a 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py @@ -2418,9 +2418,11 @@ def _save_to_local_file(file_path, security_domain): try: with open(file_path, 'w') as f: f.write(security_domain.value) - except: # pylint: disable=bare-except + except Exception as ex: # pylint: disable=bare-except if os.path.isfile(file_path): os.remove(file_path) + from azure.cli.core.azclierror import FileOperationError + raise FileOperationError(str(ex)) ret = client.download( vault_base_url=hsm_name or vault_base_url, From f186f43aee9fec2791c0a2845b90d1387261f241 Mon Sep 17 00:00:00 2001 From: houk-ms Date: Mon, 15 Mar 2021 10:19:42 +0800 Subject: [PATCH 5/6] support --id to specify url and fix test --- .../cli/command_modules/keyvault/_params.py | 3 ++- .../cli/command_modules/keyvault/custom.py | 19 ++++++++++--------- .../tests/latest/test_keyvault_commands.py | 2 +- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/_params.py b/src/azure-cli/azure/cli/command_modules/keyvault/_params.py index 1f5b9d550ca..3d0083604b7 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/_params.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/_params.py @@ -471,6 +471,7 @@ class CLISecurityDomainOperation(str, Enum): c.argument('hsm_name', hsm_url_type, required=False, help='Name of the HSM. Can be omitted if --id is specified.') c.extra('identifier', options_list=['--id'], validator=validate_vault_or_hsm, help='Id of the HSM.') + c.ignore('vault_base_url') with self.argument_context('keyvault security-domain init-recovery') as c: c.argument('sd_exchange_key', help='Local file path to store the exported key.') @@ -492,7 +493,6 @@ class CLISecurityDomainOperation(str, Enum): help='Path to a file where the JSON blob returned by this command is stored.') c.argument('sd_quorum', type=int, help='The minimum number of shares required to decrypt the security domain ' 'for recovery.') - c.ignore('vault_base_url') with self.argument_context('keyvault security-domain wait') as c: c.argument('hsm_name', hsm_url_type, help='Name of the HSM. Can be omitted if --id is specified.', @@ -502,6 +502,7 @@ class CLISecurityDomainOperation(str, Enum): help='Proceed only if HSM belongs to the specified resource group.') c.argument('target_operation', arg_type=get_enum_type(CLISecurityDomainOperation), help='Target operation that needs waiting.') + c.ignore('vault_base_url') # endregion # region keyvault backup/restore diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py index b0dc98e8c3a..fdfbae1634a 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py @@ -2173,11 +2173,11 @@ def full_restore(cmd, client, token, folder_to_restore, storage_resource_uri=Non # region security domain def security_domain_init_recovery(client, hsm_name, sd_exchange_key, - identifier=None): # pylint: disable=unused-argument + identifier=None, vault_base_url=None): # pylint: disable=unused-argument if os.path.exists(sd_exchange_key): raise CLIError("File named '{}' already exists.".format(sd_exchange_key)) - ret = client.transfer_key(vault_base_url=hsm_name) + ret = client.transfer_key(vault_base_url=hsm_name or vault_base_url) exchange_key = json.loads(json.loads(ret)['transfer_key']) def get_x5c_as_pem(): @@ -2204,7 +2204,8 @@ def get_x5c_as_pem(): raise ex -def _wait_security_domain_operation(client, hsm_name, target_operation='upload', identifier=None): # pylint: disable=unused-argument +def _wait_security_domain_operation(client, hsm_name, target_operation='upload', + identifier=None, vault_base_url=None): # pylint: disable=unused-argument retries = 0 max_retries = 30 wait_second = 5 @@ -2212,9 +2213,9 @@ def _wait_security_domain_operation(client, hsm_name, target_operation='upload', try: ret = None if target_operation == 'upload': - ret = client.upload_pending(vault_base_url=hsm_name) + ret = client.upload_pending(vault_base_url=hsm_name or vault_base_url) elif target_operation == 'download': - ret = client.download_pending(vault_base_url=hsm_name) + ret = client.download_pending(vault_base_url=hsm_name or vault_base_url) # v7.2-preview and v7.2 will change the upload operation from Sync to Async # due to service defects, it returns 'Succeeded' before the change and 'Success' after the change @@ -2319,7 +2320,7 @@ def _security_domain_gen_blob(sd_exchange_key, share_arrays, enc_data, required) def security_domain_upload(cmd, client, hsm_name, sd_file, sd_exchange_key, sd_wrapping_keys, passwords=None, - identifier=None, no_wait=False): # pylint: disable=unused-argument + identifier=None, vault_base_url=None, no_wait=False): # pylint: disable=unused-argument resource_paths = [sd_file, sd_exchange_key] for p in resource_paths: if not os.path.exists(p): @@ -2358,12 +2359,12 @@ def security_domain_upload(cmd, client, hsm_name, sd_file, sd_exchange_key, sd_w ) SecurityDomainObject = cmd.get_models('SecurityDomainObject', resource_type=ResourceType.DATA_PRIVATE_KEYVAULT) security_domain = SecurityDomainObject(value=restore_blob_value) - retval = client.upload(vault_base_url=hsm_name, security_domain=security_domain) + retval = client.upload(vault_base_url=hsm_name or vault_base_url, security_domain=security_domain) if no_wait: return retval - new_retval = _wait_security_domain_operation(client, hsm_name, 'upload') + new_retval = _wait_security_domain_operation(client, hsm_name, 'upload', vault_base_url=vault_base_url) if new_retval: return new_retval return retval @@ -2430,7 +2431,7 @@ def _save_to_local_file(file_path, security_domain): ) if not no_wait: - polling_ret = _wait_security_domain_operation(client, hsm_name, 'download') + polling_ret = _wait_security_domain_operation(client, hsm_name, 'download', vault_base_url=vault_base_url) # Due to service defect, status could be 'Success' or 'Succeeded' when it succeeded if polling_ret and getattr(polling_ret, 'status', None) != 'Failed': _save_to_local_file(security_domain_file, ret) diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/tests/latest/test_keyvault_commands.py b/src/azure-cli/azure/cli/command_modules/keyvault/tests/latest/test_keyvault_commands.py index 8ee5f4de4b1..1493aa50b41 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/tests/latest/test_keyvault_commands.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/tests/latest/test_keyvault_commands.py @@ -371,7 +371,7 @@ def test_keyvault_hsm_security_domain(self): # download SD self.cmd('az keyvault security-domain download --hsm-name {hsm_name} --security-domain-file "{sdfile}" ' - '--sd-quorum 2 --sd-wrapping-keys "{cer1_path}" "{cer2_path}" "{cer3_path}"') + '--sd-quorum 2 --sd-wrapping-keys "{cer1_path}" "{cer2_path}" "{cer3_path}" --no-wait') # delete the HSM self.cmd('az keyvault delete --hsm-name {hsm_name}') From 15cd31b3fcda1c50e7cee58bf6188f508a4da3a4 Mon Sep 17 00:00:00 2001 From: houk-ms Date: Wed, 17 Mar 2021 14:47:42 +0800 Subject: [PATCH 6/6] wait before polling --- src/azure-cli/azure/cli/command_modules/keyvault/custom.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py index fdfbae1634a..bae72cb9f7f 100644 --- a/src/azure-cli/azure/cli/command_modules/keyvault/custom.py +++ b/src/azure-cli/azure/cli/command_modules/keyvault/custom.py @@ -2364,6 +2364,8 @@ def security_domain_upload(cmd, client, hsm_name, sd_file, sd_exchange_key, sd_w if no_wait: return retval + wait_second = 5 + time.sleep(wait_second) new_retval = _wait_security_domain_operation(client, hsm_name, 'upload', vault_base_url=vault_base_url) if new_retval: return new_retval @@ -2431,6 +2433,8 @@ def _save_to_local_file(file_path, security_domain): ) if not no_wait: + wait_second = 5 + time.sleep(wait_second) polling_ret = _wait_security_domain_operation(client, hsm_name, 'download', vault_base_url=vault_base_url) # Due to service defect, status could be 'Success' or 'Succeeded' when it succeeded if polling_ret and getattr(polling_ret, 'status', None) != 'Failed':