From bc516689bb0b6ab48b4d6d2da6c2e4b4ca70f719 Mon Sep 17 00:00:00 2001 From: evelyn-ys Date: Tue, 10 Nov 2020 18:01:42 +0800 Subject: [PATCH 1/3] Add error handling when finding subscriptions --- src/azure-cli-core/azure/cli/core/_profile.py | 41 +++++++++++++++---- .../azure/cli/core/adal_authentication.py | 10 +++-- .../azure/cli/core/tests/test_profile.py | 35 ++++++++++++++++ 3 files changed, 75 insertions(+), 11 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 266d451081c..c42855b4aa6 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -24,7 +24,6 @@ is_windows, is_wsl from azure.cli.core.cloud import get_active_cloud, set_cloud_subscription - logger = get_logger(__name__) # Names below are used by azure-xplat-cli to persist account information into @@ -825,7 +824,11 @@ def create_arm_client_factory(credentials): from azure.cli.core.profiles._shared import get_client_class from azure.cli.core.profiles import ResourceType, get_api_version from azure.cli.core.commands.client_factory import configure_common_settings + from azure.cli.core.azclierror import CLIInternalError client_type = get_client_class(ResourceType.MGMT_RESOURCE_SUBSCRIPTIONS) + if client_type is None: + raise CLIInternalError("Unable to get '{}' in profile '{}'" + .format(ResourceType.MGMT_RESOURCE_SUBSCRIPTIONS, cli_ctx.cloud.profile)) api_version = get_api_version(cli_ctx, ResourceType.MGMT_RESOURCE_SUBSCRIPTIONS) client = client_type(credentials, api_version=api_version, base_url=self.cli_ctx.cloud.endpoints.resource_manager) @@ -837,10 +840,14 @@ def create_arm_client_factory(credentials): def find_from_user_account(self, username, password, tenant, resource): context = self._create_auth_context(tenant) - if password: - token_entry = context.acquire_token_with_username_password(resource, username, password, _CLIENT_ID) - else: # when refresh account, we will leverage local cached tokens - token_entry = context.acquire_token(resource, username, _CLIENT_ID) + from requests.exceptions import InvalidURL + try: + if password: + token_entry = context.acquire_token_with_username_password(resource, username, password, _CLIENT_ID) + else: # when refresh account, we will leverage local cached tokens + token_entry = context.acquire_token(resource, username, _CLIENT_ID) + except InvalidURL as err: + _handle_invalid_url(err) if not token_entry: return [] @@ -861,8 +868,12 @@ def find_through_authorization_code_flow(self, tenant, resource, authority_url): # exchange the code for the token context = self._create_auth_context(tenant) - token_entry = context.acquire_token_with_authorization_code(results['code'], results['reply_url'], - resource, _CLIENT_ID, None) + from requests.exceptions import InvalidURL + try: + token_entry = context.acquire_token_with_authorization_code(results['code'], results['reply_url'], + resource, _CLIENT_ID, None) + except InvalidURL as err: + _handle_invalid_url(err) self.user_id = token_entry[_TOKEN_ENTRY_USER_ID] logger.warning("You have logged in. Now let us find all the subscriptions to which you have access...") if tenant is None: @@ -873,7 +884,11 @@ def find_through_authorization_code_flow(self, tenant, resource, authority_url): def find_through_interactive_flow(self, tenant, resource): context = self._create_auth_context(tenant) - code = context.acquire_user_code(resource, _CLIENT_ID) + from requests.exceptions import InvalidURL + try: + code = context.acquire_user_code(resource, _CLIENT_ID) + except InvalidURL as err: + _handle_invalid_url(err) logger.warning(code['message']) token_entry = context.acquire_token_with_device_code(resource, code, _CLIENT_ID) self.user_id = token_entry[_TOKEN_ENTRY_USER_ID] @@ -1331,3 +1346,13 @@ def _get_authorization_code(resource, authority_url): if results.get('no_browser'): raise RuntimeError() return results + + +def _handle_invalid_url(ex): + from requests.exceptions import InvalidURL + if isinstance(ex, InvalidURL): + import traceback + from azure.cli.core.azclierror import UnknownError + logger.debug('Invalid url when acquiring token\n%s', traceback.format_exc()) + raise UnknownError(error_msg='Invalid url when acquiring token', + recommendation='Please make sure the cloud is registered with valid url') diff --git a/src/azure-cli-core/azure/cli/core/adal_authentication.py b/src/azure-cli-core/azure/cli/core/adal_authentication.py index 4d42159d5f9..d5dd278628c 100644 --- a/src/azure-cli-core/azure/cli/core/adal_authentication.py +++ b/src/azure-cli-core/azure/cli/core/adal_authentication.py @@ -119,9 +119,13 @@ def set_token(self): except requests.exceptions.HTTPError as err: logger.debug('throw requests.exceptions.HTTPError when doing MSIAuthentication: \n%s', traceback.format_exc()) - raise AzureResponseError('Failed to connect to MSI. Please make sure MSI is configured correctly.\n' - 'Get Token request returned http error: {}, reason: {}' - .format(err.response.status, err.response.reason)) + try: + raise AzureResponseError('Failed to connect to MSI. Please make sure MSI is configured correctly.\n' + 'Get Token request returned http error: {}, reason: {}' + .format(err.response.status, err.response.reason)) + except AttributeError: + raise AzureResponseError('Failed to connect to MSI. Please make sure MSI is configured correctly.\n' + 'Get Token request returned: {}'.format(err.response)) except TimeoutError as err: logger.debug('throw TimeoutError when doing MSIAuthentication: \n%s', traceback.format_exc()) diff --git a/src/azure-cli-core/azure/cli/core/tests/test_profile.py b/src/azure-cli-core/azure/cli/core/tests/test_profile.py index b96af36a0ee..c3fc07e6fc2 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_profile.py +++ b/src/azure-cli-core/azure/cli/core/tests/test_profile.py @@ -316,6 +316,15 @@ def test_subscription_finder_constructor(self, get_api_mock): result = finder._arm_client_factory(mock.MagicMock()) self.assertEqual(result.config.base_url, 'http://foo_arm') + @mock.patch('azure.cli.core.profiles._shared.get_client_class', autospec=True) + def test_subscription_finder_fail_on_arm_client_factory(self, get_client_class_mock): + cli = DummyCli() + get_client_class_mock.return_value = None + finder = SubscriptionFinder(cli, None, None, arm_client_factory=None) + from azure.cli.core.azclierror import CLIInternalError + with self.assertRaisesRegexp(CLIInternalError, 'Unable to get'): + finder._arm_client_factory(mock.MagicMock()) + @mock.patch('adal.AuthenticationContext', autospec=True) def test_get_auth_info_for_logged_in_service_principal(self, mock_auth_context): cli = DummyCli() @@ -1037,6 +1046,32 @@ def test_find_subscriptions_thru_username_non_password(self, mock_auth_context): # assert self.assertEqual([], subs) + @mock.patch('adal.AuthenticationContext', autospec=True) + @mock.patch('azure.cli.core._profile._get_authorization_code', autospec=True) + def test_find_subscriptions_with_invalid_authority_url(self, _get_authorization_code_mock, mock_auth_context): + from requests.exceptions import InvalidURL + from azure.cli.core.azclierror import UnknownError + + def mock_acquire(*args, **kwargs): + raise InvalidURL(request='http://some.unknown.endpoints') + + cli = DummyCli() + mock_auth_context.acquire_token_with_username_password.side_effect = mock_acquire + mock_auth_context.acquire_token_with_authorization_code.side_effect = mock_acquire + mock_auth_context.acquire_user_code.side_effect = mock_acquire + _get_authorization_code_mock.return_value = { + 'code': 'code1', + 'reply_url': 'http://localhost:8888' + } + + finder = SubscriptionFinder(cli, lambda _, _1, _2: mock_auth_context, None, lambda _: None) + with self.assertRaisesRegexp(UnknownError, 'Invalid url when acquiring token'): + finder.find_from_user_account(self.user1, 'bar', None, 'http://goo-resource') + with self.assertRaisesRegexp(UnknownError, 'Invalid url when acquiring token'): + finder.find_through_authorization_code_flow(None, 'https://management.core.windows.net/', 'https:/some_aad_point/common') + with self.assertRaisesRegexp(UnknownError, 'Invalid url when acquiring token'): + finder.find_through_interactive_flow(None, 'https://management.core.windows.net/') + @mock.patch('azure.cli.core.adal_authentication.MSIAuthenticationWrapper', autospec=True) @mock.patch('azure.cli.core.profiles._shared.get_client_class', autospec=True) @mock.patch('azure.cli.core._profile._get_cloud_console_token_endpoint', autospec=True) From 029d501ee1e55b53047eafcd931d6788b970caca Mon Sep 17 00:00:00 2001 From: evelyn-ys Date: Wed, 11 Nov 2020 10:28:36 +0800 Subject: [PATCH 2/3] exception handler --- src/azure-cli-core/azure/cli/core/_profile.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index c42855b4aa6..7e7f49c9824 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -840,14 +840,13 @@ def create_arm_client_factory(credentials): def find_from_user_account(self, username, password, tenant, resource): context = self._create_auth_context(tenant) - from requests.exceptions import InvalidURL try: if password: token_entry = context.acquire_token_with_username_password(resource, username, password, _CLIENT_ID) else: # when refresh account, we will leverage local cached tokens token_entry = context.acquire_token(resource, username, _CLIENT_ID) - except InvalidURL as err: - _handle_invalid_url(err) + except Exception as err: + _login_exception_handler(err) if not token_entry: return [] @@ -868,12 +867,11 @@ def find_through_authorization_code_flow(self, tenant, resource, authority_url): # exchange the code for the token context = self._create_auth_context(tenant) - from requests.exceptions import InvalidURL try: token_entry = context.acquire_token_with_authorization_code(results['code'], results['reply_url'], resource, _CLIENT_ID, None) - except InvalidURL as err: - _handle_invalid_url(err) + except Exception as err: + _login_exception_handler(err) self.user_id = token_entry[_TOKEN_ENTRY_USER_ID] logger.warning("You have logged in. Now let us find all the subscriptions to which you have access...") if tenant is None: @@ -884,11 +882,10 @@ def find_through_authorization_code_flow(self, tenant, resource, authority_url): def find_through_interactive_flow(self, tenant, resource): context = self._create_auth_context(tenant) - from requests.exceptions import InvalidURL try: code = context.acquire_user_code(resource, _CLIENT_ID) - except InvalidURL as err: - _handle_invalid_url(err) + except Exception as err: + _login_exception_handler(err) logger.warning(code['message']) token_entry = context.acquire_token_with_device_code(resource, code, _CLIENT_ID) self.user_id = token_entry[_TOKEN_ENTRY_USER_ID] @@ -1348,7 +1345,7 @@ def _get_authorization_code(resource, authority_url): return results -def _handle_invalid_url(ex): +def _login_exception_handler(ex): from requests.exceptions import InvalidURL if isinstance(ex, InvalidURL): import traceback From 91579bb81576326f9a68bf2767a3c1ff163e3a53 Mon Sep 17 00:00:00 2001 From: evelyn-ys Date: Wed, 11 Nov 2020 10:47:06 +0800 Subject: [PATCH 3/3] pylint --- src/azure-cli-core/azure/cli/core/_profile.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 7e7f49c9824..004a720fce8 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -845,7 +845,7 @@ def find_from_user_account(self, username, password, tenant, resource): token_entry = context.acquire_token_with_username_password(resource, username, password, _CLIENT_ID) else: # when refresh account, we will leverage local cached tokens token_entry = context.acquire_token(resource, username, _CLIENT_ID) - except Exception as err: + except Exception as err: # pylint: disable=broad-except _login_exception_handler(err) if not token_entry: @@ -870,7 +870,7 @@ def find_through_authorization_code_flow(self, tenant, resource, authority_url): try: token_entry = context.acquire_token_with_authorization_code(results['code'], results['reply_url'], resource, _CLIENT_ID, None) - except Exception as err: + except Exception as err: # pylint: disable=broad-except _login_exception_handler(err) self.user_id = token_entry[_TOKEN_ENTRY_USER_ID] logger.warning("You have logged in. Now let us find all the subscriptions to which you have access...") @@ -884,7 +884,7 @@ def find_through_interactive_flow(self, tenant, resource): context = self._create_auth_context(tenant) try: code = context.acquire_user_code(resource, _CLIENT_ID) - except Exception as err: + except Exception as err: # pylint: disable=broad-except _login_exception_handler(err) logger.warning(code['message']) token_entry = context.acquire_token_with_device_code(resource, code, _CLIENT_ID)