From a4d6884344cfe25ab73e7f13b865f6ab84d43fa7 Mon Sep 17 00:00:00 2001 From: Zella Henderson Date: Fri, 24 Sep 2021 18:00:23 -0700 Subject: [PATCH 1/2] Add identity to cosmosdb-preview extension --- azure-cli-extensions.pyproj | 117 +++++++++++++++++ src/cosmosdb-preview/HISTORY.rst | 4 + src/cosmosdb-preview/README.md | 1 + .../azext_cosmosdb_preview/_help.py | 34 +++++ .../azext_cosmosdb_preview/_params.py | 7 + .../azext_cosmosdb_preview/commands.py | 5 + .../azext_cosmosdb_preview/custom.py | 121 +++++++++++++++++ .../latest/test_cosmosdb-identity_scenario.py | 123 ++++++++++++++++++ src/cosmosdb-preview/setup.py | 2 +- 9 files changed, 413 insertions(+), 1 deletion(-) create mode 100644 src/cosmosdb-preview/azext_cosmosdb_preview/tests/latest/test_cosmosdb-identity_scenario.py diff --git a/azure-cli-extensions.pyproj b/azure-cli-extensions.pyproj index c877de0a89c..2dbd7a29457 100644 --- a/azure-cli-extensions.pyproj +++ b/azure-cli-extensions.pyproj @@ -743,6 +743,103 @@ + + + + + Code + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -4668,6 +4765,15 @@ + + + + + + + + + @@ -4791,6 +4897,17 @@ + + + + + + + + + + + diff --git a/src/cosmosdb-preview/HISTORY.rst b/src/cosmosdb-preview/HISTORY.rst index c97fa6774f4..b96d6379c3a 100644 --- a/src/cosmosdb-preview/HISTORY.rst +++ b/src/cosmosdb-preview/HISTORY.rst @@ -2,6 +2,10 @@ Release History =============== +0.11.0 +++++++ +* Adding support for commands to manage identity. + 0.10.0 ++++++ * Adding support for Services APIs and Graph Resources. diff --git a/src/cosmosdb-preview/README.md b/src/cosmosdb-preview/README.md index f81aa322012..5aff6b61a1f 100644 --- a/src/cosmosdb-preview/README.md +++ b/src/cosmosdb-preview/README.md @@ -7,6 +7,7 @@ This package provides commands to - List the different versions of databases and collections that were modified - Trigger a point in time restore on the Azure CosmosDB continuous mode backup accounts - Update the backup interval and backup retention of periodic mode backup accounts +- Update the identities associated with a Cosmos DB account ## How to use ## diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py b/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py index b656d4d79a6..1f8074aebb2 100644 --- a/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py +++ b/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py @@ -213,3 +213,37 @@ type: command short-summary: Return if the given cosmosdb graph resource exist. """ + +helps['cosmosdb identity'] = """ +type: group +short-summary: Manage Azure Cosmos DB managed service identities. +""" + +helps['cosmosdb identity show'] = """ +type: command +short-summary: Show the identities for a Azure Cosmos DB database account. +examples: + - name: Show the identities for a Azure Cosmos DB database account. + text: az cosmosdb identity show --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup +""" + +helps['cosmosdb identity assign'] = """ +type: command +short-summary: Assign SystemAssigned identity for a Azure Cosmos DB database account. +examples: + - name: Assign SystemAssigned identity for a Azure Cosmos DB database account. + text: az cosmosdb identity assign --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup + - name: Assign one UserAssigned identity for a Azure Cosmos DB database account. + text: az cosmosdb identity assign --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup --identities /subscriptions/00000000-0000-0000-0000-00000000/resourcegroups/MyRG/providers/Microsoft.ManagedIdentity/userAssignedIdentities/MyID +""" + +helps['cosmosdb identity remove'] = """ +type: command +short-summary: Remove SystemAssigned identity for a Azure Cosmos DB database account. +examples: + - name: Remove SystemAssigned identity for a Azure Cosmos DB database account. + text: az cosmosdb identity remove --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup + - name: Remove a UserAssigned identity for a Azure Cosmos DB database account. + text: az cosmosdb identity remove --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup --identities /subscriptions/00000000-0000-0000-0000-00000000/resourcegroups/MyRG/providers/Microsoft.ManagedIdentity/userAssignedIdentities/MyID + +""" \ No newline at end of file diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py b/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py index fd031571481..ee9c0c08995 100644 --- a/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py +++ b/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py @@ -93,3 +93,10 @@ def load_arguments(self, _): c.argument('service_name', options_list=['--name', '-n'], help="Service Name.") c.argument('instance_count', options_list=['--count', '-c'], help="Instance Count.") c.argument('instance_size', options_list=['--size'], help="Instance Size. Possible values are: Cosmos.D4s, Cosmos.D8s, Cosmos.D16s etc") + + # Identity + with self.argument_context('cosmosdb identity assign') as c: + c.argument('identities', options_list=['--identities'], nargs='*', help="Space-separated identities to assign. Use '[system]' to refer to the system assigned identity. Default: '[system]'") + + with self.argument_context('cosmosdb identity remove') as c: + c.argument('identities', options_list=['--identities'], nargs='*', help="Space-separated identities to remove. Use '[system]' to refer to the system assigned identity. Default: '[system]'") diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py b/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py index 5aef8206d16..2e0efccc368 100644 --- a/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py +++ b/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py @@ -62,3 +62,8 @@ def load_command_table(self, _): g.command('list', 'list') g.show_command('show', 'get') g.command('delete', 'begin_delete', confirmation=True, supports_no_wait=True) + + with self.command_group('cosmosdb identity', client_factory=cf_db_accounts, is_preview=True) as g: + g.custom_show_command('show', 'cli_cosmosdb_identity_show') + g.custom_command('assign', 'cli_cosmosdb_identity_assign') + g.custom_command('remove', 'cli_cosmosdb_identity_remove') \ No newline at end of file diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/custom.py b/src/cosmosdb-preview/azext_cosmosdb_preview/custom.py index d5f673934f5..30aaeb114b8 100644 --- a/src/cosmosdb-preview/azext_cosmosdb_preview/custom.py +++ b/src/cosmosdb-preview/azext_cosmosdb_preview/custom.py @@ -216,6 +216,127 @@ def cli_cosmosdb_managed_cassandra_datacenter_update(client, resource_group_name ) return client.begin_create_update(resource_group_name, cluster_name, data_center_name, data_center_resource) + + +def cli_cosmosdb_identity_show(client, resource_group_name, account_name): + """ Show the identity associated with a Cosmos DB account """ + + cosmos_db_account = client.get(resource_group_name, account_name) + return cosmos_db_account.identity + + +def cli_cosmosdb_identity_assign(client, + resource_group_name, + account_name, + identities=None): + """ Update the identities associated with a Cosmos DB account """ + + existing = client.get(resource_group_name, account_name) + + SYSTEM_ID = '[system]' + enable_system = identities is None or SYSTEM_ID in identities + new_user_identities = [] + if identities is not None: + new_user_identities = [x for x in identities if x != SYSTEM_ID] + + only_enabling_system = enable_system and len(new_user_identities) == 0 + system_already_added = existing.identity.type == ResourceIdentityType.system_assigned or existing.identity.type == ResourceIdentityType.system_assigned_user_assigned + all_new_users_already_added = new_user_identities and existing.identity and existing.identity.user_assigned_identities and all(x in existing.identity.user_assigned_identities for x in new_user_identities) + if only_enabling_system and system_already_added: + return existing.identity + if (not enable_system) and all_new_users_already_added: + return existing.identity + if enable_system and system_already_added and all_new_users_already_added: + return existing.identity + + if existing.identity and existing.identity.type == ResourceIdentityType.system_assigned_user_assigned: + identity_type = ResourceIdentityType.system_assigned_user_assigned + elif existing.identity and existing.identity.type == ResourceIdentityType.system_assigned and new_user_identities: + identity_type = ResourceIdentityType.system_assigned_user_assigned + elif existing.identity and existing.identity.type == ResourceIdentityType.user_assigned and enable_system: + identity_type = ResourceIdentityType.system_assigned_user_assigned + elif new_user_identities and enable_system: + identity_type = ResourceIdentityType.system_assigned_user_assigned + elif new_user_identities: + identity_type = ResourceIdentityType.user_assigned + else: + identity_type = ResourceIdentityType.system_assigned + + if identity_type in [ResourceIdentityType.system_assigned, ResourceIdentityType.none]: + new_identity = ManagedServiceIdentity(type=identity_type.value) + else: + new_assigned_identities = existing.identity.user_assigned_identities or {} + for identity in new_user_identities: + new_assigned_identities[identity] = Components1Jq1T4ISchemasManagedserviceidentityPropertiesUserassignedidentitiesAdditionalproperties() + + new_identity = ManagedServiceIdentity(type=identity_type.value, user_assigned_identities=new_assigned_identities) + + params = DatabaseAccountUpdateParameters(identity=new_identity) + async_cosmos_db_update = client.begin_update(resource_group_name, account_name, params) + cosmos_db_account = async_cosmos_db_update.result() + return cosmos_db_account.identity + + +def cli_cosmosdb_identity_remove(client, + resource_group_name, + account_name, + identities=None): + """ Remove the identities associated with a Cosmos DB account """ + + existing = client.get(resource_group_name, account_name) + + SYSTEM_ID = '[system]' + remove_system_assigned_identity = False + if not identities: + remove_system_assigned_identity = True + elif SYSTEM_ID in identities: + remove_system_assigned_identity = True + identities.remove(SYSTEM_ID) + + if existing.identity is None: + return ManagedServiceIdentity(type=ResourceIdentityType.none.value) + if existing.identity.user_assigned_identities: + existing_identities = existing.identity.user_assigned_identities.keys() + else: + existing_identities = [] + if identities: + identities_to_remove = identities + else: + identities_to_remove = [] + non_existing = [x for x in identities_to_remove if x not in set(existing_identities)] + + if non_existing: + raise CLIError("'{}' are not associated with '{}'".format(','.join(non_existing), account_name)) + identities_remaining = [x for x in existing_identities if x not in set(identities_to_remove)] + if remove_system_assigned_identity and ((not existing.identity) or (existing.identity and existing.identity.type in [ResourceIdentityType.none, ResourceIdentityType.user_assigned])): + raise CLIError("System-assigned identity is not associated with '{}'".format(account_name)) + + if identities_remaining and not remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.system_assigned_user_assigned: + set_type = ResourceIdentityType.system_assigned_user_assigned + elif identities_remaining and remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.system_assigned_user_assigned: + set_type = ResourceIdentityType.user_assigned + elif identities_remaining and not remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.user_assigned: + set_type = ResourceIdentityType.user_assigned + elif not identities_remaining and not remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.system_assigned_user_assigned: + set_type = ResourceIdentityType.system_assigned + elif not identities_remaining and not remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.system_assigned: + set_type = ResourceIdentityType.system_assigned + else: + set_type = ResourceIdentityType.none + + new_user_identities = {} + for identity in identities_remaining: + new_user_identities[identity] = Components1Jq1T4ISchemasManagedserviceidentityPropertiesUserassignedidentitiesAdditionalproperties() + if set_type in [ResourceIdentityType.system_assigned_user_assigned, ResourceIdentityType.user_assigned]: + for removed_identity in identities_to_remove: + new_user_identities[removed_identity] = None + if not new_user_identities: + new_user_identities = None + + params = DatabaseAccountUpdateParameters(identity=ManagedServiceIdentity(type=set_type, user_assigned_identities=new_user_identities)) + async_cosmos_db_update = client.begin_update(resource_group_name, account_name, params) + cosmos_db_account = async_cosmos_db_update.result() + return cosmos_db_account.identity def _handle_exists_exception(http_response_error): diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/tests/latest/test_cosmosdb-identity_scenario.py b/src/cosmosdb-preview/azext_cosmosdb_preview/tests/latest/test_cosmosdb-identity_scenario.py new file mode 100644 index 00000000000..dd7f01cf662 --- /dev/null +++ b/src/cosmosdb-preview/azext_cosmosdb_preview/tests/latest/test_cosmosdb-identity_scenario.py @@ -0,0 +1,123 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import os +from unittest import mock + + +TEST_DIR = os.path.abspath(os.path.join(os.path.abspath(__file__), '..')) + +class Cosmosdb_previewIdentityTest(ScenarioTest): + + @ResourceGroupPreparer(name_prefix='cli_test_cosmosdb_managed_service_identity') + @KeyVaultPreparer(name_prefix='cli', name_len=15, location='eastus2', additional_params='--enable-purge-protection') + def test_cosmosdb_preview_identity(self, resource_group, key_vault): + key_name = self.create_random_name(prefix='cli', length=15) + key_uri = "https://{}.vault.azure.net/keys/{}".format(key_vault, key_name) + + self.kwargs.update({ + 'acc': self.create_random_name(prefix='cli', length=15), + 'acc2': self.create_random_name(prefix='cli', length=15), + 'kv_name': key_vault, + 'key_name': key_name, + 'key_uri': key_uri, + 'location': "eastus2", + 'id1': self.create_random_name(prefix='cli', length=15), + 'id2': self.create_random_name(prefix='cli', length=15) + }) + + self.cmd('az keyvault set-policy -n {kv_name} -g {rg} --spn a232010e-820c-4083-83bb-3ace5fc29d0b --key-permissions get unwrapKey wrapKey') + self.cmd('az keyvault key create -n {key_name} --kty RSA --size 3072 --vault-name {kv_name}') + + cmk_account = self.cmd('az cosmosdb create -n {acc} -g {rg} --locations regionName={location} failoverPriority=0 --key-uri {key_uri} --assign-identity [system] --default-identity FirstPartyIdentity').get_output_in_json() + + assert cmk_account["keyVaultKeyUri"] == key_uri + assert cmk_account["defaultIdentity"] == 'FirstPartyIdentity' + assert cmk_account["identity"]['type'] == 'SystemAssigned' + + identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg}').get_output_in_json() + assert identity_output["type"] == "None" + + identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg}').get_output_in_json() + assert identity_output["type"] == "SystemAssigned" + + identity_principal_id = identity_output["principalId"] + self.kwargs.update({ + 'identity_principal_id': identity_principal_id + }) + self.cmd('az keyvault set-policy -n {kv_name} -g {rg} --object-id {identity_principal_id} --key-permissions get unwrapKey wrapKey') + + # System assigned identity tests + cmk_account = self.cmd('az cosmosdb update -n {acc} -g {rg} --default-identity SystemAssignedIdentity').get_output_in_json() + assert cmk_account["defaultIdentity"] == 'SystemAssignedIdentity' + + identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg}').get_output_in_json() + assert identity_output["type"] == "None" + + identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg}').get_output_in_json() + assert identity_output["type"] == "SystemAssigned" + + # User assigned identity tests + user_identity1 = self.cmd('az identity create -n {id1} -g {rg}').get_output_in_json() + user_identity2 = self.cmd('az identity create -n {id2} -g {rg}').get_output_in_json() + id1 = user_identity1["id"] + id1principal = user_identity1["principalId"] + id2 = user_identity2["id"] + self.kwargs.update({ + 'id1': id1, + 'id2': id2, + 'id1principal': id1principal + }) + + identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id1}').get_output_in_json() + assert identity_output["type"] == "SystemAssigned,UserAssigned" + assert list(identity_output["userAssignedIdentities"])[0] == id1 + assert len(identity_output["userAssignedIdentities"]) == 1 + + identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id2}').get_output_in_json() + assert identity_output["type"] == "SystemAssigned,UserAssigned" + assert (list(identity_output["userAssignedIdentities"])[0] == id2 or list(identity_output["userAssignedIdentities"])[1] == id2) + assert len(identity_output["userAssignedIdentities"]) == 2 + + identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg}').get_output_in_json() + assert identity_output["type"] == "UserAssigned" + assert len(identity_output["userAssignedIdentities"]) == 2 + + identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id2}').get_output_in_json() + assert identity_output["type"] == "UserAssigned" + assert list(identity_output["userAssignedIdentities"])[0] == id1 + assert len(identity_output["userAssignedIdentities"]) == 1 + + identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id2}').get_output_in_json() + assert identity_output["type"] == "UserAssigned" + assert (list(identity_output["userAssignedIdentities"])[0] == id2 or list(identity_output["userAssignedIdentities"])[1] == id2) + assert len(identity_output["userAssignedIdentities"]) == 2 + + identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id1} {id2}').get_output_in_json() + assert identity_output["type"] == "None" + + identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id1} {id2} [system]').get_output_in_json() + assert identity_output["type"] == "SystemAssigned,UserAssigned" + assert len(identity_output["userAssignedIdentities"]) == 2 + + identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id2}').get_output_in_json() + assert identity_output["type"] == "SystemAssigned,UserAssigned" + assert len(identity_output["userAssignedIdentities"]) == 1 + + identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id1}').get_output_in_json() + assert identity_output["type"] == "SystemAssigned" + + identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id1} {id2} [system]').get_output_in_json() + assert identity_output["type"] == "SystemAssigned,UserAssigned" + assert len(identity_output["userAssignedIdentities"]) == 2 + identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id1} {id2} [system]').get_output_in_json() + assert identity_output["type"] == "None" + + # Default identity tests + self.cmd('az keyvault set-policy --name {kv_name} --object-id {id1principal} --key-permissions get unwrapKey wrapKey') + default_id_acct = self.cmd('az cosmosdb create -n {acc2} -g {rg} --locations regionName={location} failoverPriority=0 --key-uri {key_uri} --assign-identity {id1} --default-identity "UserAssignedIdentity={id1}"').get_output_in_json() + assert default_id_acct["identity"]["type"] == "UserAssigned" + assert list(default_id_acct["identity"]["userAssignedIdentities"])[0] == id1 + assert default_id_acct["defaultIdentity"] == "UserAssignedIdentity=" + id1 \ No newline at end of file diff --git a/src/cosmosdb-preview/setup.py b/src/cosmosdb-preview/setup.py index 9ffdda9ef03..86aa239c392 100644 --- a/src/cosmosdb-preview/setup.py +++ b/src/cosmosdb-preview/setup.py @@ -16,7 +16,7 @@ # TODO: Confirm this is the right version number you want and it matches your # HISTORY.rst entry. -VERSION = '0.10.0' +VERSION = '0.11.0' # The full list of classifiers is available at # https://pypi.python.org/pypi?%3Aaction=list_classifiers From b17055ea510de4e4f20ad97674e027649b0a707e Mon Sep 17 00:00:00 2001 From: Zella Henderson Date: Fri, 24 Sep 2021 18:22:46 -0700 Subject: [PATCH 2/2] Update --- src/cosmosdb-preview/azext_cosmosdb_preview/_help.py | 2 +- src/cosmosdb-preview/azext_cosmosdb_preview/commands.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py b/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py index 1f8074aebb2..dd291b33671 100644 --- a/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py +++ b/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py @@ -246,4 +246,4 @@ - name: Remove a UserAssigned identity for a Azure Cosmos DB database account. text: az cosmosdb identity remove --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup --identities /subscriptions/00000000-0000-0000-0000-00000000/resourcegroups/MyRG/providers/Microsoft.ManagedIdentity/userAssignedIdentities/MyID -""" \ No newline at end of file +""" diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py b/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py index 2e0efccc368..f540421b3e4 100644 --- a/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py +++ b/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py @@ -63,7 +63,7 @@ def load_command_table(self, _): g.show_command('show', 'get') g.command('delete', 'begin_delete', confirmation=True, supports_no_wait=True) - with self.command_group('cosmosdb identity', client_factory=cf_db_accounts, is_preview=True) as g: + with self.command_group('cosmosdb identity', client_factory=cf_service, is_preview=True) as g: g.custom_show_command('show', 'cli_cosmosdb_identity_show') g.custom_command('assign', 'cli_cosmosdb_identity_assign') - g.custom_command('remove', 'cli_cosmosdb_identity_remove') \ No newline at end of file + g.custom_command('remove', 'cli_cosmosdb_identity_remove')