[Key Vault] Update tests to only record with Managed Identity (#35859)

Author: mccoyp

sdk/keyvault/azure-keyvault-administration/README.md

@@ -65,21 +65,30 @@ client = KeyVaultAccessControlClient(vault_url=MANAGED_HSM_URL, credential=crede
 > **NOTE:** For an asynchronous client, import `azure.keyvault.administration.aio`'s `KeyVaultAccessControlClient` instead.
 
 #### Create a KeyVaultBackupClient
-After configuring your environment for the [DefaultAzureCredential][default_cred_ref] to use a suitable method of authentication, you can do the following to create a backup client (replacing the value of `vault_url` with your Managed HSM's URL):
+After creating a user-assigned [managed identity][managed_identity] and
+[granting it access to your Managed HSM][managed_identity_backup_setup], you can do the following to create a backup
+client (setting the value of `CLIENT_ID` to your managed identity's client ID):
 
 <!-- SNIPPET:backup_restore_operations.create_a_backup_restore_client -->
 
 ```python
-from azure.identity import DefaultAzureCredential
+from azure.identity import ManagedIdentityCredential
 from azure.keyvault.administration import KeyVaultBackupClient
 
 MANAGED_HSM_URL = os.environ["MANAGED_HSM_URL"]
-credential = DefaultAzureCredential()
+MANAGED_IDENTITY_CLIENT_ID = os.environ["CLIENT_ID"]
+credential = ManagedIdentityCredential(client_id=MANAGED_IDENTITY_CLIENT_ID)
 client = KeyVaultBackupClient(vault_url=MANAGED_HSM_URL, credential=credential)
 ```
 
 <!-- END SNIPPET -->
 
+Using the `ManagedIdentityCredential` is preferred in order to enable authenticating backup and restore operations with
+Managed Identity. Any other `azure-identity` credential could be provided instead if SAS tokens are used in these
+operations.
+
+See [azure-identity][managed_identity_ref] documentation for more information on Managed Identity authentication.
+
 > **NOTE:** For an asynchronous client, import `azure.keyvault.administration.aio`'s `KeyVaultBackupClient` instead.
 
 #### Create a KeyVaultSettingsClient
@@ -265,7 +274,11 @@ client.delete_role_assignment(scope=scope, name=role_assignment.name)
 
 ### Perform a full key backup
 The `KeyVaultBackupClient` can be used to back up your entire collection of keys. The backing store for full key
-backups is a blob storage container using Shared Access Signature (SAS) authentication.
+backups is a blob storage container using either Managed Identity (which is preferred) or Shared Access Signature (SAS)
+authentication.
+
+If using Managed Identity, first make sure your user-assigned managed identity has the correct access to your Storage
+account and Managed HSM per [the service's guidance][managed_identity_backup_setup].
 
 For more details on creating a SAS token using a `BlobServiceClient` from [`azure-storage-blob`][storage_blob], refer
 to the library's [credential documentation][sas_docs]. Alternatively, it is possible to
@@ -275,9 +288,8 @@ to the library's [credential documentation][sas_docs]. Alternatively, it is poss
 
 ```python
 CONTAINER_URL = os.environ["CONTAINER_URL"]
-SAS_TOKEN = os.environ["SAS_TOKEN"]
 
-backup_result: KeyVaultBackupResult = client.begin_backup(CONTAINER_URL, sas_token=SAS_TOKEN).result()
+backup_result: KeyVaultBackupResult = client.begin_backup(CONTAINER_URL, use_managed_identity=True).result()
 print(f"Azure Storage Blob URL of the backup: {backup_result.folder_url}")
 ```
 
@@ -289,8 +301,12 @@ the operation is complete without returning an object.
 
 ### Perform a full key restore
 The `KeyVaultBackupClient` can be used to restore your entire collection of keys from a backup. The data source for a
-full key restore is a storage blob accessed using Shared Access Signature authentication. You will also need the URL of
-the backup (`KeyVaultBackupResult.folder_url`) from the [above snippet](#perform-a-full-key-backup).
+full key restore is a storage blob accessed using either Managed Identity (which is preferred) or Shared Access
+Signature (SAS) authentication. You will also need the URL of the backup (`KeyVaultBackupResult.folder_url`) from the
+[above snippet](#perform-a-full-key-backup).
+
+If using Managed Identity, first make sure your user-assigned managed identity has the correct access to your Storage
+account and Managed HSM per [the service's guidance][managed_identity_backup_setup].
 
 For more details on creating a SAS token using a `BlobServiceClient` from [`azure-storage-blob`][storage_blob], refer
 to the library's [credential documentation][sas_docs]. Alternatively, it is possible to
@@ -299,10 +315,8 @@ to the library's [credential documentation][sas_docs]. Alternatively, it is poss
 <!-- SNIPPET:backup_restore_operations.begin_restore -->
 
 ```python
-SAS_TOKEN = os.environ["SAS_TOKEN"]
-
 # `backup_result` is the KeyVaultBackupResult returned by `begin_backup`
-client.begin_restore(backup_result.folder_url, sas_token=SAS_TOKEN).wait()
+client.begin_restore(backup_result.folder_url, use_managed_identity=True).wait()
 print("Vault restored successfully.")
 ```
 
@@ -398,7 +412,9 @@ contact [email protected] with any additional questions or comments.
 
 [managed_hsm]: https://docs.microsoft.com/azure/key-vault/managed-hsm/overview
 [managed_hsm_cli]: https://docs.microsoft.com/azure/key-vault/managed-hsm/quick-create-cli
-[managed_identity]: https://docs.microsoft.com/azure/active-directory/managed-identities-azure-resources/overview
+[managed_identity]: https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/overview
+[managed_identity_backup_setup]: https://learn.microsoft.com/azure/key-vault/managed-hsm/backup-restore#prerequisites-if-backing-up-and-restoring-using-user-assigned-managed-identity
+[managed_identity_ref]: https://aka.ms/azsdk/python/identity/docs#azure.identity.ManagedIdentityCredential
 
 [pip]: https://pypi.org/project/pip/
 [pypi_package_administration]: https://pypi.org/project/azure-keyvault-administration

sdk/keyvault/azure-keyvault-administration/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/keyvault/azure-keyvault-administration",
-  "Tag": "python/keyvault/azure-keyvault-administration_3353e6edf1"
+  "Tag": "python/keyvault/azure-keyvault-administration_df93633a52"
 }

sdk/keyvault/azure-keyvault-administration/samples/backup_restore_operations.py

@@ -13,15 +13,16 @@
 # 2. azure-keyvault-administration and azure-identity libraries (pip install these)
 #
 # 3. Set environment variable MANAGED_HSM_URL with the URL of your managed HSM
-#    
-# 4. Set up your environment to use azure-identity's DefaultAzureCredential. For more information about how to configure
-#    the DefaultAzureCredential, refer to https://aka.ms/azsdk/python/identity/docs#azure.identity.DefaultAzureCredential
 #
-# 5. A storage account containing a blob storage container
+# 4. A user-assigned managed identity that has access to your managed HSM. For more information about how to create a
+#    user-assigned managed identity, refer to
+#    https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/overview
+#    
+# 5. A storage account, that your managed identity has access to, containing a blob storage container
 #    (See https://docs.microsoft.com/azure/storage/blobs/storage-blobs-introduction)
 #
-# 6. Set environment variables CONTAINER_URL and SAS_TOKEN corresponding to your blob container's URI and SAS
-#    (See https://docs.microsoft.com/azure/storage/common/storage-sas-overview)
+# 6. Set environment variables CONTAINER_URL and CLIENT_ID, corresponding to your blob container's URI and your
+#    user-assigned managed identity's client ID, respectively.
 #
 # For more details, see https://docs.microsoft.com/azure/key-vault/managed-hsm/backup-restore
 #
@@ -34,13 +35,15 @@
 # ----------------------------------------------------------------------------------------------------------
 
 # Instantiate a backup client that will be used to call the service.
-# Here we use the DefaultAzureCredential, but any azure-identity credential can be used.
+# Here we use the ManagedIdentityCredential to use Managed Identity, but any azure-identity credential can be used if
+# opting for SAS token authentication.
 # [START create_a_backup_restore_client]
-from azure.identity import DefaultAzureCredential
+from azure.identity import ManagedIdentityCredential
 from azure.keyvault.administration import KeyVaultBackupClient
 
 MANAGED_HSM_URL = os.environ["MANAGED_HSM_URL"]
-credential = DefaultAzureCredential()
+MANAGED_IDENTITY_CLIENT_ID = os.environ["CLIENT_ID"]
+credential = ManagedIdentityCredential(client_id=MANAGED_IDENTITY_CLIENT_ID)
 client = KeyVaultBackupClient(vault_url=MANAGED_HSM_URL, credential=credential)
 # [END create_a_backup_restore_client]
 
@@ -50,9 +53,8 @@
 print("\n.. Back up the vault")
 # [START begin_backup]
 CONTAINER_URL = os.environ["CONTAINER_URL"]
-SAS_TOKEN = os.environ["SAS_TOKEN"]
 
-backup_result: KeyVaultBackupResult = client.begin_backup(CONTAINER_URL, sas_token=SAS_TOKEN).result()
+backup_result: KeyVaultBackupResult = client.begin_backup(CONTAINER_URL, use_managed_identity=True).result()
 print(f"Azure Storage Blob URL of the backup: {backup_result.folder_url}")
 # [END begin_backup]
 print("Vault backed up successfully.")
@@ -63,9 +65,7 @@
 # To restore a single key from the backed up vault instead, pass the key_name keyword argument.
 print("\n.. Restore the full vault")
 # [START begin_restore]
-SAS_TOKEN = os.environ["SAS_TOKEN"]
-
 # `backup_result` is the KeyVaultBackupResult returned by `begin_backup`
-client.begin_restore(backup_result.folder_url, sas_token=SAS_TOKEN).wait()
+client.begin_restore(backup_result.folder_url, use_managed_identity=True).wait()
 print("Vault restored successfully.")
 # [END begin_restore]

sdk/keyvault/azure-keyvault-administration/samples/backup_restore_operations_async.py

@@ -6,7 +6,7 @@
 import os
 
 from azure.keyvault.administration.aio import KeyVaultBackupClient
-from azure.identity.aio import DefaultAzureCredential
+from azure.identity.aio import ManagedIdentityCredential
 
 # ----------------------------------------------------------------------------------------------------------
 # Prerequisites:
@@ -15,15 +15,16 @@
 # 2. azure-keyvault-administration and azure-identity libraries (pip install these)
 #
 # 3. Set environment variable MANAGED_HSM_URL with the URL of your managed HSM
-#    
-# 4. Set up your environment to use azure-identity's DefaultAzureCredential. For more information about how to configure
-#    the DefaultAzureCredential, refer to https://aka.ms/azsdk/python/identity/docs#azure.identity.DefaultAzureCredential
 #
-# 5. A storage account containing a blob storage container
+# 4. A user-assigned managed identity that has access to your managed HSM. For more information about how to create a
+#    user-assigned managed identity, refer to
+#    https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/overview
+#    
+# 5. A storage account, that your managed identity has access to, containing a blob storage container
 #    (See https://docs.microsoft.com/azure/storage/blobs/storage-blobs-introduction)
 #
-# 6. Set environment variables CONTAINER_URL and SAS_TOKEN corresponding to your blob container's URI and SAS
-#    (See https://docs.microsoft.com/azure/storage/common/storage-sas-overview)
+# 6. Set environment variables CONTAINER_URL and CLIENT_ID, corresponding to your blob container's URI and your
+#    user-assigned managed identity's client ID, respectively.
 #
 # For more details, see https://docs.microsoft.com/azure/key-vault/managed-hsm/backup-restore
 #
@@ -38,18 +39,18 @@
 async def run_sample():
     MANAGED_HSM_URL = os.environ["MANAGED_HSM_URL"]
     CONTAINER_URL = os.environ["CONTAINER_URL"]
-    SAS_TOKEN = os.environ["SAS_TOKEN"]
+    MANAGED_IDENTITY_CLIENT_ID = os.environ["CLIENT_ID"]
 
     # Instantiate a backup client that will be used to call the service.
     # Here we use the DefaultAzureCredential, but any azure-identity credential can be used.
-    credential = DefaultAzureCredential()
+    credential = ManagedIdentityCredential(client_id=MANAGED_IDENTITY_CLIENT_ID)
     client = KeyVaultBackupClient(vault_url=MANAGED_HSM_URL, credential=credential)
 
     # Let's back up the vault with begin_backup, which returns a poller. Calling result() on the poller will return
     # a KeyVaultBackupResult that contains the URL of the backup after the operation completes. Calling wait() on
     # the poller will wait until the operation is complete.
     print("\n.. Back up the vault")
-    backup_poller = await client.begin_backup(CONTAINER_URL, sas_token=SAS_TOKEN)
+    backup_poller = await client.begin_backup(CONTAINER_URL, use_managed_identity=True)
     backup_result = await backup_poller.result()
     print("Vault backed up successfully.")
     assert backup_result.folder_url
@@ -58,7 +59,7 @@ async def run_sample():
     # return None after the operation completes. Calling wait() on the poller will wait until the operation is
     # complete. To restore a single key from the backed up vault instead, pass the key_name keyword argument.
     print("\n.. Restore the full vault")
-    restore_poller = await client.begin_restore(backup_result.folder_url, sas_token=SAS_TOKEN)
+    restore_poller = await client.begin_restore(backup_result.folder_url, use_managed_identity=True)
     await restore_poller.wait()
     print("Vault restored successfully.")
 

sdk/keyvault/azure-keyvault-administration/tests/_async_test_case.py

@@ -7,6 +7,7 @@
 import pytest
 from azure.keyvault.administration import ApiVersion
 from azure.keyvault.administration._internal.client_base import DEFAULT_VERSION
+from azure.identity.aio import ManagedIdentityCredential
 from devtools_testutils import AzureRecordedTestCase
 
 
@@ -30,6 +31,7 @@ def __init__(self, **kwargs) -> None:
             self.container_uri = container_playback_uri
             self.sas_token = playback_sas_token
 
+        self.managed_identity_client_id = os.environ.get("MANAGED_IDENTITY_CLIENT_ID")
         use_pwsh = os.environ.get("AZURE_TEST_USE_PWSH_AUTH", "false")
         use_cli = os.environ.get("AZURE_TEST_USE_CLI_AUTH", "false")
         use_vscode = os.environ.get("AZURE_TEST_USE_VSCODE_AUTH", "false")
@@ -52,6 +54,30 @@ def _set_mgmt_settings_real_values(self):
 
 
 class KeyVaultBackupClientPreparer(BaseClientPreparer):
+    def __call__(self, fn):
+        async def _preparer(test_class, api_version, **kwargs):
+            self._skip_if_not_configured(api_version)
+            kwargs["container_uri"] = self.container_uri
+            kwargs["managed_hsm_url"] = self.managed_hsm_url
+            client = self.create_backup_client(self.managed_identity_client_id, api_version=api_version, **kwargs)
+
+            async with client:
+                await fn(test_class, client, **kwargs)
+        return _preparer
+
+    def create_backup_client(self, managed_identity_client_id, **kwargs):
+        from azure.keyvault.administration.aio import KeyVaultBackupClient
+
+        if self.is_live:
+            credential = ManagedIdentityCredential(client_id=managed_identity_client_id)
+        else:
+            credential = self.get_credential(KeyVaultBackupClient, is_async=True)
+        return self.create_client_from_credential(
+            KeyVaultBackupClient, credential=credential, vault_url=self.managed_hsm_url, **kwargs
+        )
+
+
+class KeyVaultBackupClientSasPreparer(BaseClientPreparer):
     def __call__(self, fn):
         async def _preparer(test_class, api_version, **kwargs):
             self._skip_if_not_configured(api_version)

sdk/keyvault/azure-keyvault-administration/tests/_test_case.py

@@ -8,6 +8,7 @@
 
 from azure.keyvault.administration import ApiVersion
 from azure.keyvault.administration._internal.client_base import DEFAULT_VERSION
+from azure.identity import ManagedIdentityCredential
 from devtools_testutils import AzureRecordedTestCase
 
 
@@ -31,6 +32,7 @@ def __init__(self, **kwargs) -> None:
             self.container_uri = container_playback_uri
             self.sas_token = playback_sas_token
 
+        self.managed_identity_client_id = os.environ.get("MANAGED_IDENTITY_CLIENT_ID")
         use_pwsh = os.environ.get("AZURE_TEST_USE_PWSH_AUTH", "false")
         use_cli = os.environ.get("AZURE_TEST_USE_CLI_AUTH", "false")
         use_vscode = os.environ.get("AZURE_TEST_USE_VSCODE_AUTH", "false")
@@ -56,6 +58,33 @@ class KeyVaultBackupClientPreparer(BaseClientPreparer):
     def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
 
+    def __call__(self, fn):
+        def _preparer(test_class, api_version, **kwargs):
+            self._skip_if_not_configured(api_version)
+            kwargs["container_uri"] = self.container_uri
+            kwargs["managed_hsm_url"] = self.managed_hsm_url
+            client = self.create_backup_client(self.managed_identity_client_id, api_version=api_version, **kwargs)
+
+            with client:
+                fn(test_class, client, **kwargs)
+        return _preparer
+
+    def create_backup_client(self, managed_identity_client_id, **kwargs):
+        from azure.keyvault.administration import KeyVaultBackupClient
+
+        if self.is_live:
+            credential = ManagedIdentityCredential(client_id=managed_identity_client_id)
+        else:
+            credential = self.get_credential(KeyVaultBackupClient)
+        return self.create_client_from_credential(
+            KeyVaultBackupClient, credential=credential, vault_url=self.managed_hsm_url, **kwargs
+        )
+
+
+class KeyVaultBackupClientSasPreparer(BaseClientPreparer):
+    def __init__(self, **kwargs) -> None:
+       super().__init__(**kwargs)
+
     def __call__(self, fn):
         def _preparer(test_class, api_version, **kwargs):
             self._skip_if_not_configured(api_version)