From fbfe69172600936798387e2b5c9c1c9887bb38f4 Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Wed, 8 Jan 2020 13:43:41 +0800 Subject: [PATCH 1/9] add tenant to get_raw_token --- src/azure-cli-core/HISTORY.rst | 2 + src/azure-cli-core/azure/cli/core/_profile.py | 31 ++++++--- .../azure/cli/core/tests/test_profile.py | 66 ++++++++++++++++++- src/azure-cli/HISTORY.rst | 4 ++ .../cli/command_modules/profile/__init__.py | 1 + .../cli/command_modules/profile/custom.py | 8 +-- 6 files changed, 97 insertions(+), 15 deletions(-) diff --git a/src/azure-cli-core/HISTORY.rst b/src/azure-cli-core/HISTORY.rst index 5dfcf1adbd3..2d07c89e432 100644 --- a/src/azure-cli-core/HISTORY.rst +++ b/src/azure-cli-core/HISTORY.rst @@ -3,6 +3,8 @@ Release History =============== +* `get_raw_token`: Add `tenant` parameter to acquire token for the tenant directly, needless to specify a subscription + 2.0.79 ++++++ * Fix #11586: `az login` is not recorded in server telemetry diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 0e5b9e58601..d3a0b3fba35 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -587,7 +587,9 @@ def get_refresh_token(self, resource=None, sp_secret = self._creds_cache.retrieve_secret_of_service_principal(username_or_sp_id) return username_or_sp_id, sp_secret, None, str(account[_TENANT_ID]) - def get_raw_token(self, resource=None, subscription=None): + def get_raw_token(self, resource=None, subscription=None, tenant=None): + if subscription and tenant: + raise CLIError("Please specify only one of subscription and tenant, not both") account = self.get_subscription(subscription) user_type = account[_USER_ENTITY][_USER_TYPE] username_or_sp_id = account[_USER_ENTITY][_USER_NAME] @@ -595,23 +597,32 @@ def get_raw_token(self, resource=None, subscription=None): identity_type, identity_id = Profile._try_parse_msi_account_name(account) if identity_type: + # MSI + if tenant: + raise CLIError("Tenant shouldn't be specified for MSI account") msi_creds = MsiAccountTypes.msi_auth_factory(identity_type, identity_id, resource) msi_creds.set_token() token_entry = msi_creds.token creds = (token_entry['token_type'], token_entry['access_token'], token_entry) elif in_cloud_console() and account[_USER_ENTITY].get(_CLOUD_SHELL_ID): + # Cloud Shell + if tenant: + raise CLIError("Tenant shouldn't be specified for Cloud Shell account") creds = self._get_token_from_cloud_shell(resource) - - elif user_type == _USER: - creds = self._creds_cache.retrieve_token_for_user(username_or_sp_id, - account[_TENANT_ID], resource) else: - creds = self._creds_cache.retrieve_token_for_service_principal(username_or_sp_id, - resource, - account[_TENANT_ID]) + tenant_dest = tenant if tenant else account[_TENANT_ID] + if user_type == _USER: + # User + creds = self._creds_cache.retrieve_token_for_user(username_or_sp_id, + tenant_dest, resource) + else: + # Service Principal + creds = self._creds_cache.retrieve_token_for_service_principal(username_or_sp_id, + resource, + tenant_dest) return (creds, - str(account[_SUBSCRIPTION_ID]), - str(account[_TENANT_ID])) + str('N/A' if tenant else account[_SUBSCRIPTION_ID]), + str(tenant if tenant else account[_TENANT_ID])) def refresh_accounts(self, subscription_finder=None): subscriptions = self.load_cached_subscriptions() 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 65e8dd8c789..8b766fe5b57 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 @@ -716,6 +716,17 @@ def test_get_raw_token(self, mock_get_token, mock_read_cred_file): self.assertEqual(sub, '1') self.assertEqual(tenant, self.tenant_id) + # Test get_raw_token with tenant + creds, sub, tenant = profile.get_raw_token(resource='https://foo', tenant=self.tenant_id) + + self.assertEqual(creds[0], self.token_entry1['tokenType']) + self.assertEqual(creds[1], self.raw_token1) + self.assertEqual(creds[2]['expiresOn'], self.token_entry1['expiresOn']) + mock_get_token.assert_called_with(mock.ANY, self.user1, self.tenant_id, 'https://foo') + self.assertEqual(mock_get_token.call_count, 2) + self.assertEqual(sub, 'N/A') + self.assertEqual(tenant, self.tenant_id) + @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) @mock.patch('azure.cli.core._profile.CredsCache.retrieve_token_for_service_principal', autospec=True) def test_get_raw_token_for_sp(self, mock_get_token, mock_read_cred_file): @@ -744,6 +755,17 @@ def test_get_raw_token_for_sp(self, mock_get_token, mock_read_cred_file): self.assertEqual(sub, '1') self.assertEqual(tenant, self.tenant_id) + # Test get_raw_token with tenant + creds, sub, tenant = profile.get_raw_token(resource='https://foo', tenant=self.tenant_id) + + self.assertEqual(creds[0], self.token_entry1['tokenType']) + self.assertEqual(creds[1], self.raw_token1) + self.assertEqual(creds[2]['expiresOn'], self.token_entry1['expiresOn']) + mock_get_token.assert_called_with(mock.ANY, 'sp1', 'https://foo', self.tenant_id) + self.assertEqual(mock_get_token.call_count, 2) + self.assertEqual(sub, 'N/A') + 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) def test_get_raw_token_msi_system_assigned(self, mock_msi_auth, mock_read_cred_file): @@ -765,12 +787,54 @@ def test_get_raw_token_msi_system_assigned(self, mock_msi_auth, mock_read_cred_f mock_msi_auth.side_effect = MSRestAzureAuthStub # action - cred, subscription_id, _ = profile.get_raw_token(resource='http://test_resource') + cred, subscription_id, tenant_id = profile.get_raw_token(resource='http://test_resource') + + # assert + self.assertEqual(subscription_id, test_subscription_id) + self.assertEqual(cred[0], 'Bearer') + self.assertEqual(cred[1], TestProfile.test_msi_access_token) + self.assertEqual(subscription_id, test_subscription_id) + self.assertEqual(tenant_id, test_tenant_id) + + # verify tenant shouldn't be specified for MSI account + with self.assertRaisesRegex(CLIError, "MSI"): + cred, subscription_id, _ = profile.get_raw_token(resource='http://test_resource', tenant=self.tenant_id) + + @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) + 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 + + # setup an existing msi subscription + profile = Profile(cli_ctx=DummyCli(), storage={'subscriptions': None}, use_global_creds_cache=False, + async_persist=False) + test_subscription_id = '12345678-1bf0-4dda-aec3-cb9272f09590' + test_tenant_id = '12345678-38d6-4fb2-bad9-b7b93a3e1234' + msi_subscription = SubscriptionStub('/subscriptions/' + test_subscription_id, + self.display_name1, self.state1, test_tenant_id) + consolidated = profile._normalize_properties(self.user1, + [msi_subscription], + True) + consolidated[0]['user']['cloudShellID'] = True + profile._set_subscriptions(consolidated) + + mock_msi_auth.side_effect = MSRestAzureAuthStub + + # action + cred, subscription_id, tenant_id = profile.get_raw_token(resource='http://test_resource') # assert self.assertEqual(subscription_id, test_subscription_id) self.assertEqual(cred[0], 'Bearer') self.assertEqual(cred[1], TestProfile.test_msi_access_token) + self.assertEqual(subscription_id, test_subscription_id) + self.assertEqual(tenant_id, test_tenant_id) + + # verify tenant shouldn't be specified for Cloud Shell account + with self.assertRaisesRegex(CLIError, 'Cloud Shell'): + cred, subscription_id, _ = profile.get_raw_token(resource='http://test_resource', tenant=self.tenant_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) @mock.patch('azure.cli.core._profile.CredsCache.retrieve_token_for_user', autospec=True) diff --git a/src/azure-cli/HISTORY.rst b/src/azure-cli/HISTORY.rst index 9b5d88c1d25..8f2044bf4d4 100644 --- a/src/azure-cli/HISTORY.rst +++ b/src/azure-cli/HISTORY.rst @@ -3,6 +3,10 @@ Release History =============== +**Profile** + +* `az account get-access-token`: Add `--tenant` parameter to acquire token for the tenant directly, needless to specify a subscription + 2.0.79 ++++++ diff --git a/src/azure-cli/azure/cli/command_modules/profile/__init__.py b/src/azure-cli/azure/cli/command_modules/profile/__init__.py index 1a804788281..2259a9edda1 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/__init__.py +++ b/src/azure-cli/azure/cli/command_modules/profile/__init__.py @@ -75,6 +75,7 @@ def load_arguments(self, command): with self.argument_context('account get-access-token') as c: c.argument('resource_type', get_enum_type(cloud_resource_types), options_list=['--resource-type'], arg_group='', help='Type of well-known resource.') + c.argument('tenant', options_list=['--tenant', '-t'], is_preview=True, help='Tenant ID for which the token is acquired. Only available for user and service principal account, not for MSI or Cloud Shell account') COMMAND_LOADER_CLS = ProfileCommandsLoader diff --git a/src/azure-cli/azure/cli/command_modules/profile/custom.py b/src/azure-cli/azure/cli/command_modules/profile/custom.py index bf09876d794..9c5319d330b 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/custom.py +++ b/src/azure-cli/azure/cli/command_modules/profile/custom.py @@ -59,20 +59,20 @@ def show_subscription(cmd, subscription=None, show_auth_for_sdk=None): print(json.dumps(profile.get_sp_auth_info(subscription), indent=2)) -def get_access_token(cmd, subscription=None, resource=None, resource_type=None): - ''' +def get_access_token(cmd, subscription=None, resource=None, resource_type=None, tenant=None): + """ get AAD token to access to a specified resource :param resource: Azure resource endpoints. Default to Azure Resource Manager :param resource-type: Name of Azure resource endpoints. Can be used instead of resource. Use 'az cloud show' command for other Azure resources - ''' + """ if resource is None and resource_type is not None: endpoints_attr_name = cloud_resource_type_mappings[resource_type] resource = getattr(cmd.cli_ctx.cloud.endpoints, endpoints_attr_name) else: resource = (resource or cmd.cli_ctx.cloud.endpoints.active_directory_resource_id) profile = Profile(cli_ctx=cmd.cli_ctx) - creds, subscription, tenant = profile.get_raw_token(subscription=subscription, resource=resource) + creds, subscription, tenant = profile.get_raw_token(subscription=subscription, resource=resource, tenant=tenant) return { 'tokenType': creds[0], 'accessToken': creds[1], From 7722eb781ea9fc5686eafde31c14eb68c7e53502 Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Thu, 9 Jan 2020 11:28:36 +0800 Subject: [PATCH 2/9] use assertRaisesRegexp to pass py27 CI --- src/azure-cli-core/azure/cli/core/tests/test_profile.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 8b766fe5b57..e951e37878a 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 @@ -797,7 +797,7 @@ def test_get_raw_token_msi_system_assigned(self, mock_msi_auth, mock_read_cred_f self.assertEqual(tenant_id, test_tenant_id) # verify tenant shouldn't be specified for MSI account - with self.assertRaisesRegex(CLIError, "MSI"): + with self.assertRaisesRegexp(CLIError, "MSI"): cred, subscription_id, _ = profile.get_raw_token(resource='http://test_resource', tenant=self.tenant_id) @mock.patch('azure.cli.core._profile.in_cloud_console', autospec=True) @@ -833,7 +833,7 @@ def test_get_raw_token_in_cloud_console(self, mock_msi_auth, mock_read_cred_file self.assertEqual(tenant_id, test_tenant_id) # verify tenant shouldn't be specified for Cloud Shell account - with self.assertRaisesRegex(CLIError, 'Cloud Shell'): + with self.assertRaisesRegexp(CLIError, 'Cloud Shell'): cred, subscription_id, _ = profile.get_raw_token(resource='http://test_resource', tenant=self.tenant_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) From cd790da40c2acf5285ee276334e67d1d0f783e46 Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Thu, 9 Jan 2020 12:45:05 +0800 Subject: [PATCH 3/9] fix test --- .../profile/tests/latest/test_profile_custom.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py b/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py index 6c0dce4927e..b60cb78192f 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py +++ b/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py @@ -30,18 +30,18 @@ def test_list_only_enabled_one(self, logger_mock, load_subscription_mock): logger_mock.warning.assert_called_once_with(mock.ANY) @mock.patch('azure.cli.core._profile.Profile.get_raw_token', autospec=True) - def test_get_row_token(self, get_raw_token_mcok): + def test_get_raw_token(self, get_raw_token_mock): cmd = mock.MagicMock() cmd.cli_ctx = DummyCli() # arrange - get_raw_token_mcok.return_value = (['bearer', 'token123', {'expiresOn': '2100-01-01'}], 'sub123', 'tenant123') + get_raw_token_mock.return_value = (['bearer', 'token123', {'expiresOn': '2100-01-01'}], 'sub123', 'tenant123') # action result = get_access_token(cmd) # assert - get_raw_token_mcok.assert_called_with(mock.ANY, 'https://management.core.windows.net/', None) + get_raw_token_mock.assert_called_with(mock.ANY, 'https://management.core.windows.net/', None, None) expected_result = { 'tokenType': 'bearer', 'accessToken': 'token123', @@ -53,7 +53,12 @@ def test_get_row_token(self, get_raw_token_mcok): # assert it takes customized resource, subscription get_access_token(cmd, subscription='foosub', resource='foores') - get_raw_token_mcok.assert_called_with(mock.ANY, 'foores', 'foosub') + get_raw_token_mock.assert_called_with(mock.ANY, 'foores', 'foosub', None) + + # test get token with tenant + tenant_id = '00000000-0000-0000-0000-000000000000' + get_access_token(cmd, tenant=tenant_id) + get_raw_token_mock.assert_called_with(mock.ANY, 'https://management.core.windows.net/', None, tenant_id) @mock.patch('azure.cli.command_modules.profile.custom.Profile', autospec=True) def test_get_login(self, profile_mock): From 0a61ed45c583b0f4472df45db658522a26b8a59b Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Fri, 17 Jan 2020 11:15:38 +0800 Subject: [PATCH 4/9] auto-switch tenant for sp --- src/azure-cli-core/azure/cli/core/_profile.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index d3a0b3fba35..408b1246e0b 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -937,12 +937,21 @@ def retrieve_token_for_user(self, username, tenant, resource): def retrieve_token_for_service_principal(self, sp_id, resource, tenant, use_cert_sn_issuer=False): self.load_adal_token_cache() - matched = [x for x in self._service_principal_creds if sp_id == x[_SERVICE_PRINCIPAL_ID] and - tenant == x[_SERVICE_PRINCIPAL_TENANT]] + matched = [x for x in self._service_principal_creds if sp_id == x[_SERVICE_PRINCIPAL_ID]] if not matched: - raise CLIError("Please run 'az account set' to select active account.") - cred = matched[0] - context = self._auth_ctx_factory(self._ctx, cred[_SERVICE_PRINCIPAL_TENANT], None) + raise CLIError("Could not retrieve credential from local cache for service principal {}. " + "Please run 'az login' for this service principal." + .format(sp_id)) + matched_with_tenant = [x for x in matched if tenant == x[_SERVICE_PRINCIPAL_TENANT]] + if matched_with_tenant: + cred = matched_with_tenant[0] + else: + logger.warning("Could not retrieve credential from local cache for service principal {} under tenant {}. " + "Trying credential under tenant {}, assuming that is an app credential." + .format(sp_id, tenant, matched[0][_SERVICE_PRINCIPAL_TENANT])) + cred = matched[0] + + context = self._auth_ctx_factory(self._ctx, tenant, None) sp_auth = ServicePrincipalAuth(cred.get(_ACCESS_TOKEN, None) or cred.get(_SERVICE_PRINCIPAL_CERT_FILE, None), use_cert_sn_issuer) From 4243726ba8cd905b9f02b8579259c882e9f4654d Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Fri, 17 Jan 2020 11:47:46 +0800 Subject: [PATCH 5/9] fix linter --- 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 408b1246e0b..fd1d15764f5 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -946,9 +946,9 @@ def retrieve_token_for_service_principal(self, sp_id, resource, tenant, use_cert if matched_with_tenant: cred = matched_with_tenant[0] else: - logger.warning("Could not retrieve credential from local cache for service principal {} under tenant {}. " - "Trying credential under tenant {}, assuming that is an app credential." - .format(sp_id, tenant, matched[0][_SERVICE_PRINCIPAL_TENANT])) + logger.warning("Could not retrieve credential from local cache for service principal %s under tenant %s. " + "Trying credential under tenant %s, assuming that is an app credential.", + sp_id, tenant, matched[0][_SERVICE_PRINCIPAL_TENANT]) cred = matched[0] context = self._auth_ctx_factory(self._ctx, tenant, None) From 9f52ccd2bb0b491d19bb9e15ac97a77b9bd8381b Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Fri, 17 Jan 2020 12:23:52 +0800 Subject: [PATCH 6/9] add examples --- .../azure/cli/command_modules/profile/_help.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/azure-cli/azure/cli/command_modules/profile/_help.py b/src/azure-cli/azure/cli/command_modules/profile/_help.py index 820001a9760..afb27069bf2 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/_help.py +++ b/src/azure-cli/azure/cli/command_modules/profile/_help.py @@ -75,6 +75,19 @@ long-summary: > The token will be valid for at least 5 minutes with the maximum at 60 minutes. If the subscription argument isn't specified, the current account is used. + examples: + - name: Get an access token for the current account + text: > + az account get-access-token + - name: Get an access token for a specific subscription + text: > + az account get-access-token --subscription 00000000-0000-0000-0000-000000000000 + - name: Get an access token for a specific tenant + text: > + az account get-access-token --tenant 00000000-0000-0000-0000-000000000000 + - name: Get an access token to use with MS Graph API + text: > + az account get-access-token --resource-type ms-graph """ helps['self-test'] = """ From 15e8062973db625876a58744fcd15f9e39e4530b Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Fri, 17 Jan 2020 12:53:30 +0800 Subject: [PATCH 7/9] fix linter --- src/azure-cli/azure/cli/command_modules/profile/_help.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/azure-cli/azure/cli/command_modules/profile/_help.py b/src/azure-cli/azure/cli/command_modules/profile/_help.py index afb27069bf2..0f635743008 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/_help.py +++ b/src/azure-cli/azure/cli/command_modules/profile/_help.py @@ -76,7 +76,7 @@ The token will be valid for at least 5 minutes with the maximum at 60 minutes. If the subscription argument isn't specified, the current account is used. examples: - - name: Get an access token for the current account + - name: Get an access token for the current account text: > az account get-access-token - name: Get an access token for a specific subscription From 16d2ae8d83fd117e845bcdb557b7b38aaf4c2a19 Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Mon, 20 Jan 2020 11:00:04 +0800 Subject: [PATCH 8/9] add test for tenant not matching scenario --- src/azure-cli-core/azure/cli/core/tests/test_profile.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) 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 e951e37878a..0f6e2c4c652 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 @@ -1605,7 +1605,14 @@ def test_credscache_match_service_principal_correctly(self, _, mock_open_for_wri # action and verify(we plant an exception to throw after the SP was found; so if the exception is thrown, # we know the matching did go through) - self.assertRaises(ValueError, creds_cache.retrieve_token_for_service_principal, 'myapp', 'resource1', 'mytenant', False) + self.assertRaises(ValueError, creds_cache.retrieve_token_for_service_principal, + 'myapp', 'resource1', 'mytenant', False) + + # tenant doesn't exactly match, but it still succeeds + # before fully migrating to pytest and utilizing capsys fixture, use `pytest -o log_cli=True` to manually + # verify the warning log + self.assertRaises(ValueError, creds_cache.retrieve_token_for_service_principal, + 'myapp', 'resource1', 'mytenant2', False) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) @mock.patch('os.fdopen', autospec=True) From bc4dff24e551b36a8782ac084355944397d072ca Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Mon, 20 Jan 2020 11:32:22 +0800 Subject: [PATCH 9/9] for tenant token do not add subscription --- src/azure-cli-core/azure/cli/core/_profile.py | 2 +- .../azure/cli/core/tests/test_profile.py | 4 +-- .../cli/command_modules/profile/custom.py | 7 +++-- .../tests/latest/test_profile_custom.py | 26 ++++++++++++++++--- 4 files changed, 31 insertions(+), 8 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index fd1d15764f5..45a9a9089d7 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -621,7 +621,7 @@ def get_raw_token(self, resource=None, subscription=None, tenant=None): resource, tenant_dest) return (creds, - str('N/A' if tenant else account[_SUBSCRIPTION_ID]), + None if tenant else str(account[_SUBSCRIPTION_ID]), str(tenant if tenant else account[_TENANT_ID])) def refresh_accounts(self, subscription_finder=None): 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 0f6e2c4c652..ace6f490aa6 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 @@ -724,7 +724,7 @@ def test_get_raw_token(self, mock_get_token, mock_read_cred_file): self.assertEqual(creds[2]['expiresOn'], self.token_entry1['expiresOn']) mock_get_token.assert_called_with(mock.ANY, self.user1, self.tenant_id, 'https://foo') self.assertEqual(mock_get_token.call_count, 2) - self.assertEqual(sub, 'N/A') + self.assertIsNone(sub) self.assertEqual(tenant, self.tenant_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) @@ -763,7 +763,7 @@ def test_get_raw_token_for_sp(self, mock_get_token, mock_read_cred_file): self.assertEqual(creds[2]['expiresOn'], self.token_entry1['expiresOn']) mock_get_token.assert_called_with(mock.ANY, 'sp1', 'https://foo', self.tenant_id) self.assertEqual(mock_get_token.call_count, 2) - self.assertEqual(sub, 'N/A') + self.assertIsNone(sub) self.assertEqual(tenant, self.tenant_id) @mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True) diff --git a/src/azure-cli/azure/cli/command_modules/profile/custom.py b/src/azure-cli/azure/cli/command_modules/profile/custom.py index 9c5319d330b..4bea450f2be 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/custom.py +++ b/src/azure-cli/azure/cli/command_modules/profile/custom.py @@ -73,13 +73,16 @@ def get_access_token(cmd, subscription=None, resource=None, resource_type=None, resource = (resource or cmd.cli_ctx.cloud.endpoints.active_directory_resource_id) profile = Profile(cli_ctx=cmd.cli_ctx) creds, subscription, tenant = profile.get_raw_token(subscription=subscription, resource=resource, tenant=tenant) - return { + + result = { 'tokenType': creds[0], 'accessToken': creds[1], 'expiresOn': creds[2].get('expiresOn', 'N/A'), - 'subscription': subscription, 'tenant': tenant } + if subscription: + result['subscription'] = subscription + return result def set_active_subscription(cmd, subscription): diff --git a/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py b/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py index b60cb78192f..07bee1a8b74 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py +++ b/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py @@ -52,13 +52,33 @@ def test_get_raw_token(self, get_raw_token_mock): self.assertEqual(result, expected_result) # assert it takes customized resource, subscription - get_access_token(cmd, subscription='foosub', resource='foores') - get_raw_token_mock.assert_called_with(mock.ANY, 'foores', 'foosub', None) + resource = 'https://graph.microsoft.com/' + subscription_id = '00000001-0000-0000-0000-000000000000' + get_raw_token_mock.return_value = (['bearer', 'token123', {'expiresOn': '2100-01-01'}], subscription_id, + 'tenant123') + result = get_access_token(cmd, subscription=subscription_id, resource=resource) + get_raw_token_mock.assert_called_with(mock.ANY, resource, subscription_id, None) + expected_result = { + 'tokenType': 'bearer', + 'accessToken': 'token123', + 'expiresOn': '2100-01-01', + 'subscription': subscription_id, + 'tenant': 'tenant123' + } + self.assertEqual(result, expected_result) # test get token with tenant tenant_id = '00000000-0000-0000-0000-000000000000' - get_access_token(cmd, tenant=tenant_id) + get_raw_token_mock.return_value = (['bearer', 'token123', {'expiresOn': '2100-01-01'}], None, tenant_id) + result = get_access_token(cmd, tenant=tenant_id) get_raw_token_mock.assert_called_with(mock.ANY, 'https://management.core.windows.net/', None, tenant_id) + expected_result = { + 'tokenType': 'bearer', + 'accessToken': 'token123', + 'expiresOn': '2100-01-01', # subscription shouldn't be present + 'tenant': tenant_id + } + self.assertEqual(result, expected_result) @mock.patch('azure.cli.command_modules.profile.custom.Profile', autospec=True) def test_get_login(self, profile_mock):