Add identity to cosmosdb-preview extension

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.

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

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

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

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_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')

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):

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

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