add tenant 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.79
 ++++++
 * Fix #11586: `az login` is not recorded in server telemetry

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]))
+                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()

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)

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
 ++++++
 

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/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],