[Core][Profile] Add tenant parameter to get_raw_token

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.80
 ++++++
 * No changes

src/azure-cli-core/azure/cli/core/_profile.py

@@ -587,31 +587,42 @@ 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]
         resource = resource or self.cli_ctx.cloud.endpoints.active_directory_resource_id
 
         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]))
+                None if tenant else str(account[_SUBSCRIPTION_ID]),
+                str(tenant if tenant else account[_TENANT_ID]))
 
     def refresh_accounts(self, subscription_finder=None):
         subscriptions = self.load_cached_subscriptions()
@@ -926,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 %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)
         sp_auth = ServicePrincipalAuth(cred.get(_ACCESS_TOKEN, None) or
                                        cred.get(_SERVICE_PRINCIPAL_CERT_FILE, None),
                                        use_cert_sn_issuer)

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.assertIsNone(sub)
+        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.assertIsNone(sub)
+        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.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)
+    @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.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)
     @mock.patch('azure.cli.core._profile.CredsCache.retrieve_token_for_user', autospec=True)
@@ -1541,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)

src/azure-cli/HISTORY.rst

@@ -52,6 +52,10 @@ Release History
 
 * Fix #2092: az network dns record-set add/remove: add warning when record-set is not found. In the future, an extra argument will be supported to confirm this auto creation.
 
+**Profile**
+
+* `az account get-access-token`: Add `--tenant` parameter to acquire token for the tenant directly, needless to specify a subscription
+
 **Security**
 
 * Added new commands `az atp show` and `az atp update` to view and manage advanced threat protection settings for storage accounts.

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

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'] = """

src/azure-cli/azure/cli/command_modules/profile/custom.py

@@ -59,27 +59,30 @@ 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)
-    return {
+    creds, subscription, tenant = profile.get_raw_token(subscription=subscription, resource=resource, tenant=tenant)
+
+    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):

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',
@@ -52,8 +52,33 @@ def test_get_row_token(self, get_raw_token_mcok):
         self.assertEqual(result, expected_result)
 
         # 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')
+        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_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):