From 2938407c27daa2f03bfdb7e780c1f04a1fa38557 Mon Sep 17 00:00:00 2001 From: jiasli <4003950+jiasli@users.noreply.github.com> Date: Tue, 14 Jan 2025 17:01:47 +0800 Subject: [PATCH] managed-identity-id --- src/azure-cli-core/azure/cli/core/_profile.py | 35 +++++++++++++--- .../azure/cli/core/tests/test_profile.py | 40 ++++++++++++++++--- .../cli/command_modules/profile/__init__.py | 6 +++ .../cli/command_modules/profile/_help.py | 6 ++- .../cli/command_modules/profile/custom.py | 8 ++-- .../tests/latest/test_profile_custom.py | 6 +-- 6 files changed, 81 insertions(+), 20 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 7c09c1ac6d3..712fdf2074e 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -60,6 +60,11 @@ _AZ_LOGIN_MESSAGE = "Please run 'az login' to setup account." +MANAGED_IDENTITY_ID_WARNING = ( + "Passing the managed identity ID with --username is deprecated and will be removed in a future release. " + "Please use --client-id, --object-id or --resource-id instead." +) + def load_subscriptions(cli_ctx, all_clouds=False, refresh=False): profile = Profile(cli_ctx=cli_ctx) @@ -219,7 +224,8 @@ def login(self, self._set_subscriptions(consolidated) return deepcopy(consolidated) - def login_with_managed_identity(self, identity_id=None, allow_no_subscriptions=None): + def login_with_managed_identity(self, identity_id=None, client_id=None, object_id=None, resource_id=None, + allow_no_subscriptions=None): if _on_azure_arc(): return self.login_with_managed_identity_azure_arc( identity_id=identity_id, allow_no_subscriptions=allow_no_subscriptions) @@ -229,7 +235,28 @@ def login_with_managed_identity(self, identity_id=None, allow_no_subscriptions=N from azure.cli.core.auth.adal_authentication import MSIAuthenticationWrapper resource = self.cli_ctx.cloud.endpoints.active_directory_resource_id - if identity_id: + id_arg_count = len([arg for arg in (client_id, object_id, resource_id, identity_id) if arg]) + if id_arg_count > 1: + raise CLIError('Usage error: Provide only one of --client-id, --object-id, --resource-id, or --username.') + + if id_arg_count == 0: + identity_type = MsiAccountTypes.system_assigned + msi_creds = MSIAuthenticationWrapper(resource=resource) + elif client_id: + identity_type = MsiAccountTypes.user_assigned_client_id + identity_id = client_id + msi_creds = MSIAuthenticationWrapper(resource=resource, client_id=client_id) + elif object_id: + identity_type = MsiAccountTypes.user_assigned_object_id + identity_id = object_id + msi_creds = MSIAuthenticationWrapper(resource=resource, object_id=object_id) + elif resource_id: + identity_type = MsiAccountTypes.user_assigned_resource_id + identity_id = resource_id + msi_creds = MSIAuthenticationWrapper(resource=resource, msi_res_id=resource_id) + # The old way of re-using the same --username for 3 types of ID + elif identity_id: + logger.warning(MANAGED_IDENTITY_ID_WARNING) if is_valid_resource_id(identity_id): msi_creds = MSIAuthenticationWrapper(resource=resource, msi_res_id=identity_id) identity_type = MsiAccountTypes.user_assigned_resource_id @@ -260,10 +287,6 @@ def login_with_managed_identity(self, identity_id=None, allow_no_subscriptions=N if not authenticated: raise CLIError('Failed to connect to MSI, check your managed service identity id.') - else: - identity_type = MsiAccountTypes.system_assigned - msi_creds = MSIAuthenticationWrapper(resource=resource) - token_entry = msi_creds.token token = token_entry['access_token'] logger.info('MSI: token was retrieved. Now trying to initialize local accounts...') 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 425f968c083..5954ec8f283 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 @@ -498,7 +498,7 @@ def test_login_in_cloud_shell(self, cloud_shell_credential_mock, create_subscrip @mock.patch('requests.get', autospec=True) @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) - def test_find_subscriptions_in_vm_with_mi_system_assigned(self, create_subscription_client_mock, mock_get): + def test_login_with_mi_system_assigned(self, create_subscription_client_mock, mock_get): mock_subscription_client = mock.MagicMock() mock_subscription_client.subscriptions.list.return_value = [deepcopy(self.subscription1_raw)] create_subscription_client_mock.return_value = mock_subscription_client @@ -531,7 +531,7 @@ def test_find_subscriptions_in_vm_with_mi_system_assigned(self, create_subscript @mock.patch('requests.get', autospec=True) @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) - def test_find_subscriptions_in_vm_with_mi_no_subscriptions(self, create_subscription_client_mock, mock_get): + def test_login_with_mi_no_subscriptions(self, create_subscription_client_mock, mock_get): mock_subscription_client = mock.MagicMock() mock_subscription_client.subscriptions.list.return_value = [] create_subscription_client_mock.return_value = mock_subscription_client @@ -566,8 +566,7 @@ def test_find_subscriptions_in_vm_with_mi_no_subscriptions(self, create_subscrip @mock.patch('requests.get', autospec=True) @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) - def test_find_subscriptions_in_vm_with_mi_user_assigned_with_client_id(self, create_subscription_client_mock, - mock_get): + def test_login_with_mi_user_assigned_client_id(self, create_subscription_client_mock, mock_get): mock_subscription_client = mock.MagicMock() mock_subscription_client.subscriptions.list.return_value = [deepcopy(self.subscription1_raw)] create_subscription_client_mock.return_value = mock_subscription_client @@ -587,6 +586,19 @@ def test_find_subscriptions_in_vm_with_mi_user_assigned_with_client_id(self, cre good_response.content = encoded_test_token mock_get.return_value = good_response + subscriptions = profile.login_with_managed_identity(client_id=test_client_id) + + self.assertEqual(len(subscriptions), 1) + s = subscriptions[0] + self.assertEqual(s['name'], self.display_name1) + self.assertEqual(s['id'], self.id1.split('/')[-1]) + self.assertEqual(s['tenantId'], self.test_mi_tenant) + + self.assertEqual(s['user']['name'], 'userAssignedIdentity') + self.assertEqual(s['user']['type'], 'servicePrincipal') + self.assertEqual(s['user']['assignedIdentityInfo'], 'MSIClient-{}'.format(test_client_id)) + + # Old way of using identity_id subscriptions = profile.login_with_managed_identity(identity_id=test_client_id) self.assertEqual(len(subscriptions), 1) @@ -601,7 +613,7 @@ def test_find_subscriptions_in_vm_with_mi_user_assigned_with_client_id(self, cre @mock.patch('azure.cli.core.auth.adal_authentication.MSIAuthenticationWrapper', autospec=True) @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) - def test_find_subscriptions_in_vm_with_mi_user_assigned_with_object_id(self, create_subscription_client_mock, + def test_login_with_mi_user_assigned_object_id(self, create_subscription_client_mock, mock_msi_auth): mock_subscription_client = mock.MagicMock() mock_subscription_client.subscriptions.list.return_value = [deepcopy(self.subscription1_raw)] @@ -632,6 +644,14 @@ def set_token(self): mock_msi_auth.side_effect = AuthStub test_object_id = '54826b22-38d6-4fb2-bad9-b7b93a3e9999' + subscriptions = profile.login_with_managed_identity(object_id=test_object_id) + + s = subscriptions[0] + self.assertEqual(s['user']['name'], 'userAssignedIdentity') + self.assertEqual(s['user']['type'], 'servicePrincipal') + self.assertEqual(s['user']['assignedIdentityInfo'], 'MSIObject-{}'.format(test_object_id)) + + # Old way of using identity_id subscriptions = profile.login_with_managed_identity(identity_id=test_object_id) s = subscriptions[0] @@ -641,7 +661,7 @@ def set_token(self): @mock.patch('requests.get', autospec=True) @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) - def test_find_subscriptions_in_vm_with_mi_user_assigned_with_res_id(self, create_subscription_client_mock, + def test_login_with_mi_user_assigned_resource_id(self, create_subscription_client_mock, mock_get): mock_subscription_client = mock.MagicMock() @@ -665,6 +685,14 @@ def test_find_subscriptions_in_vm_with_mi_user_assigned_with_res_id(self, create good_response.content = encoded_test_token mock_get.return_value = good_response + subscriptions = profile.login_with_managed_identity(resource_id=test_res_id) + + s = subscriptions[0] + self.assertEqual(s['user']['name'], 'userAssignedIdentity') + self.assertEqual(s['user']['type'], 'servicePrincipal') + self.assertEqual(subscriptions[0]['user']['assignedIdentityInfo'], 'MSIResource-{}'.format(test_res_id)) + + # Old way of using identity_id subscriptions = profile.login_with_managed_identity(identity_id=test_res_id) s = subscriptions[0] 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 76cb7a3f42b..abbf17ac342 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,12 @@ def load_arguments(self, command): # Managed identity c.argument('identity', options_list=('-i', '--identity'), action='store_true', help="Log in using managed identity", arg_group='Managed Identity') + c.argument('client_id', + help="Client ID of the user-assigned managed identity", arg_group='Managed Identity') + c.argument('object_id', + help="Object ID of the user-assigned managed identity", arg_group='Managed Identity') + c.argument('resource_id', + help="Resource ID of the user-assigned managed identity", arg_group='Managed Identity') with self.argument_context('logout') as c: c.argument('username', help='account user, if missing, logout the current active account') 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 3f5c1ead13d..e61c8738c73 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/_help.py +++ b/src/azure-cli/azure/cli/command_modules/profile/_help.py @@ -43,8 +43,10 @@ text: az login --service-principal --username APP_ID --certificate /path/to/cert.pem --tenant TENANT_ID - name: Log in with a system-assigned managed identity. text: az login --identity - - name: Log in with a user-assigned managed identity. You must specify the client ID, object ID or resource ID of the user-assigned managed identity with --username. - text: az login --identity --username 00000000-0000-0000-0000-000000000000 + - name: Log in with a user-assigned managed identity's client ID. + text: az login --identity --client-id 00000000-0000-0000-0000-000000000000 + - name: Log in with a user-assigned managed identity's resource ID. + text: az login --identity --resource-id /subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/MyResourceGroup/providers/Microsoft.ManagedIdentity/userAssignedIdentities/MyIdentity """ helps['account'] = """ 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 89416732a73..9623657761a 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/custom.py +++ b/src/azure-cli/azure/cli/command_modules/profile/custom.py @@ -114,14 +114,14 @@ def account_clear(cmd): profile.logout_all() -# pylint: disable=inconsistent-return-statements, too-many-branches +# pylint: disable=too-many-branches, too-many-locals def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_subscriptions=False, # Device code flow use_device_code=False, # Service principal service_principal=None, certificate=None, use_cert_sn_issuer=None, client_assertion=None, # Managed identity - identity=False): + identity=False, client_id=None, object_id=None, resource_id=None): """Log in to access Azure subscriptions""" # quick argument usage check @@ -143,7 +143,9 @@ def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_ if identity: if in_cloud_console(): return profile.login_in_cloud_shell() - return profile.login_with_managed_identity(username, allow_no_subscriptions) + return profile.login_with_managed_identity( + identity_id=username, client_id=client_id, object_id=object_id, resource_id=resource_id, + allow_no_subscriptions=allow_no_subscriptions) if in_cloud_console(): # tell users they might not need login logger.warning(_CLOUD_CONSOLE_LOGIN_WARNING) 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 8388ff70463..74547381ccf 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 @@ -88,15 +88,15 @@ def test_get_raw_token(self, get_raw_token_mock): get_raw_token_mock.assert_called_with(mock.ANY, None, None, None, tenant_id) @mock.patch('azure.cli.command_modules.profile.custom.Profile', autospec=True) - def test_get_login(self, profile_mock): + def test_login_with_mi(self, profile_mock): invoked = [] - def test_login(msi_port, identity_id=None): + def login_with_managed_identity_mock(*args, **kwargs): invoked.append(True) # mock the instance profile_instance = mock.MagicMock() - profile_instance.login_with_managed_identity = test_login + profile_instance.login_with_managed_identity = login_with_managed_identity_mock # mock the constructor profile_mock.return_value = profile_instance