diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 42bcda1e6f0..23a4ad0b35e 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -22,6 +22,8 @@ is_windows, is_wsl from azure.cli.core.cloud import get_active_cloud, set_cloud_subscription +from .adal_authentication import MSIAuthenticationWrapper + from knack.log import get_logger from knack.util import CLIError @@ -307,18 +309,17 @@ def find_subscriptions_in_vm_with_msi(self, identity_id=None, allow_no_subscript import jwt from requests import HTTPError - from msrestazure.azure_active_directory import MSIAuthentication from msrestazure.tools import is_valid_resource_id resource = self.cli_ctx.cloud.endpoints.active_directory_resource_id if identity_id: if is_valid_resource_id(identity_id): - msi_creds = MSIAuthentication(resource=resource, msi_res_id=identity_id) + msi_creds = MSIAuthenticationWrapper(resource=resource, msi_res_id=identity_id) identity_type = MsiAccountTypes.user_assigned_resource_id else: authenticated = False try: - msi_creds = MSIAuthentication(resource=resource, client_id=identity_id) + msi_creds = MSIAuthenticationWrapper(resource=resource, client_id=identity_id) identity_type = MsiAccountTypes.user_assigned_client_id authenticated = True except HTTPError as ex: @@ -330,7 +331,7 @@ def find_subscriptions_in_vm_with_msi(self, identity_id=None, allow_no_subscript if not authenticated: try: identity_type = MsiAccountTypes.user_assigned_object_id - msi_creds = MSIAuthentication(resource=resource, object_id=identity_id) + msi_creds = MSIAuthenticationWrapper(resource=resource, object_id=identity_id) authenticated = True except HTTPError as ex: if ex.response.reason == 'Bad Request' and ex.response.status == 400: @@ -343,7 +344,7 @@ def find_subscriptions_in_vm_with_msi(self, identity_id=None, allow_no_subscript else: identity_type = MsiAccountTypes.system_assigned - msi_creds = MSIAuthentication(resource=resource) + msi_creds = MSIAuthenticationWrapper(resource=resource) token_entry = msi_creds.token token = token_entry['access_token'] @@ -388,8 +389,7 @@ def find_subscriptions_in_cloud_console(self): return deepcopy(consolidated) def _get_token_from_cloud_shell(self, resource): # pylint: disable=no-self-use - from msrestazure.azure_active_directory import MSIAuthentication - auth = MSIAuthentication(resource=resource) + auth = MSIAuthenticationWrapper(resource=resource) auth.set_token() token_entry = auth.token return (token_entry['token_type'], token_entry['access_token'], token_entry) @@ -774,15 +774,14 @@ def valid_msi_account_types(): @staticmethod def msi_auth_factory(cli_account_name, identity, resource): - from msrestazure.azure_active_directory import MSIAuthentication if cli_account_name == MsiAccountTypes.system_assigned: - return MSIAuthentication(resource=resource) + return MSIAuthenticationWrapper(resource=resource) if cli_account_name == MsiAccountTypes.user_assigned_client_id: - return MSIAuthentication(resource=resource, client_id=identity) + return MSIAuthenticationWrapper(resource=resource, client_id=identity) if cli_account_name == MsiAccountTypes.user_assigned_object_id: - return MSIAuthentication(resource=resource, object_id=identity) + return MSIAuthenticationWrapper(resource=resource, object_id=identity) if cli_account_name == MsiAccountTypes.user_assigned_resource_id: - return MSIAuthentication(resource=resource, msi_res_id=identity) + return MSIAuthenticationWrapper(resource=resource, msi_res_id=identity) raise ValueError("unrecognized msi account name '{}'".format(cli_account_name)) 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 9b9d0db6db9..19cfd10b56c 100644 --- a/src/azure-cli-core/azure/cli/core/adal_authentication.py +++ b/src/azure-cli-core/azure/cli/core/adal_authentication.py @@ -8,6 +8,7 @@ import adal from msrest.authentication import Authentication +from msrestazure.azure_active_directory import MSIAuthentication from azure.core.credentials import AccessToken from azure.cli.core.util import in_cloud_console @@ -60,8 +61,10 @@ def _get_token(self): # This method is exposed for Azure Core. def get_token(self, *scopes, **kwargs): # pylint:disable=unused-argument _, token, full_token, _ = self._get_token() - - return AccessToken(token, int(full_token['expiresIn'] + time.time())) + try: + return AccessToken(token, int(full_token['expiresIn'] + time.time())) + except KeyError: # needed to deal with differing unserialized MSI token payload + return AccessToken(token, int(full_token['expires_on'])) # This method is exposed for msrest. def signed_session(self, session=None): # pylint: disable=arguments-differ @@ -83,3 +86,10 @@ def _log_hostname(): logger = get_logger(__name__) logger.warning("A Cloud Shell credential problem occurred. When you report the issue with the error " "below, please mention the hostname '%s'", socket.gethostname()) + + +class MSIAuthenticationWrapper(MSIAuthentication): + # This method is exposed for Azure Core. + def get_token(self): + self.set_token() + return AccessToken(self.token['access_token'], int(self.token['expires_on'])) 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 080be1e0996..2909a05e870 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 @@ -645,7 +645,7 @@ def test_get_login_credentials_aux_tenants(self, mock_get_token, mock_read_cred_ aux_tenants=[test_tenant_id2]) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_login_credentials_msi_system_assigned(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -676,7 +676,7 @@ def test_get_login_credentials_msi_system_assigned(self, mock_msi_auth, mock_rea self.assertTrue(cred.token_read_count) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_login_credentials_msi_user_assigned_with_client_id(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -707,7 +707,7 @@ def test_get_login_credentials_msi_user_assigned_with_client_id(self, mock_msi_a self.assertTrue(cred.client_id, test_client_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_login_credentials_msi_user_assigned_with_object_id(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -738,7 +738,7 @@ def test_get_login_credentials_msi_user_assigned_with_object_id(self, mock_msi_a self.assertTrue(cred.object_id, test_object_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_login_credentials_msi_user_assigned_with_res_id(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -849,7 +849,7 @@ def test_get_raw_token_for_sp(self, mock_get_token, mock_read_cred_file): self.assertEqual(tenant, self.tenant_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_raw_token_msi_system_assigned(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -884,7 +884,7 @@ def test_get_raw_token_msi_system_assigned(self, mock_msi_auth, mock_read_cred_f @mock.patch('azure.cli.core._profile.in_cloud_console', autospec=True) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_raw_token_in_cloud_console(self, mock_msi_auth, mock_read_cred_file, mock_in_cloud_console): mock_read_cred_file.return_value = [] mock_in_cloud_console.return_value = True @@ -1037,7 +1037,7 @@ def test_find_subscriptions_thru_username_non_password(self, mock_auth_context): # assert self.assertEqual([], subs) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.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) @mock.patch('azure.cli.core._profile.SubscriptionFinder', autospec=True) @@ -1186,7 +1186,7 @@ def __init__(self, *args, **kwargs): self.assertEqual(s['id'], self.id1.split('/')[-1]) self.assertEqual(s['tenantId'], '54826b22-38d6-4fb2-bad9-b7b93a3e9c5a') - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) @mock.patch('azure.cli.core.profiles._shared.get_client_class', autospec=True) @mock.patch('azure.cli.core._profile.SubscriptionFinder', autospec=True) def test_find_subscriptions_in_vm_with_msi_user_assigned_with_object_id(self, mock_subscription_finder, mock_get_client_class, diff --git a/src/azure-cli-core/azure/cli/core/tests/test_profile_v2016_06_01.py b/src/azure-cli-core/azure/cli/core/tests/test_profile_v2016_06_01.py index e9c16209562..5f84985d241 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_profile_v2016_06_01.py +++ b/src/azure-cli-core/azure/cli/core/tests/test_profile_v2016_06_01.py @@ -571,7 +571,7 @@ def test_get_login_credentials_aux_subscriptions(self, mock_get_token, mock_read self.assertEqual(mock_get_token.call_count, 2) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_login_credentials_msi_system_assigned(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -602,7 +602,7 @@ def test_get_login_credentials_msi_system_assigned(self, mock_msi_auth, mock_rea self.assertTrue(cred.token_read_count) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_login_credentials_msi_user_assigned_with_client_id(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -633,7 +633,7 @@ def test_get_login_credentials_msi_user_assigned_with_client_id(self, mock_msi_a self.assertTrue(cred.client_id, test_client_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_login_credentials_msi_user_assigned_with_object_id(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -664,7 +664,7 @@ def test_get_login_credentials_msi_user_assigned_with_object_id(self, mock_msi_a self.assertTrue(cred.object_id, test_object_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_login_credentials_msi_user_assigned_with_res_id(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -753,7 +753,7 @@ def test_get_raw_token_for_sp(self, mock_get_token, mock_read_cred_file): self.assertEqual(tenant, self.tenant_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) def test_get_raw_token_msi_system_assigned(self, mock_msi_auth, mock_read_cred_file): mock_read_cred_file.return_value = [] @@ -899,7 +899,7 @@ def test_find_subscriptions_thru_username_non_password(self, mock_auth_context): # assert self.assertEqual([], subs) - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.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) @mock.patch('azure.cli.core._profile.SubscriptionFinder', autospec=True) @@ -1048,7 +1048,7 @@ def __init__(self, *args, **kwargs): self.assertEqual(s['id'], self.id1.split('/')[-1]) self.assertEqual(s['tenantId'], '54826b22-38d6-4fb2-bad9-b7b93a3e9c5a') - @mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True) + @mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True) @mock.patch('azure.cli.core.profiles._shared.get_client_class', autospec=True) @mock.patch('azure.cli.core._profile.SubscriptionFinder', autospec=True) def test_find_subscriptions_in_vm_with_msi_user_assigned_with_object_id(self, mock_subscription_finder, mock_get_client_class,