MAINT: Use DefaultAzureCredential directly for blob storage auth (microsoft#1946)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Roman Lutz <romanlutz13@gmail.com>

Author: 3 people

pyrit/memory/storage/storage.py

@@ -15,6 +15,7 @@
 from pyrit.common.deprecation import print_deprecation_message
 
 if TYPE_CHECKING:
+    from azure.identity.aio import DefaultAzureCredential
     from azure.storage.blob.aio import ContainerClient as AsyncContainerClient
 
 logger = logging.getLogger(__name__)
@@ -264,33 +265,63 @@ def __init__(
         self._container_url: str = container_url
         self._sas_token = sas_token
         self._client_async: AsyncContainerClient | None = None
+        self._credential: DefaultAzureCredential | None = None
 
     async def _create_container_client_async(self) -> AsyncContainerClient:
         """
         Create an asynchronous ContainerClient for Azure Storage.
 
         If a SAS token is provided via the
         AZURE_STORAGE_ACCOUNT_SAS_TOKEN environment variable or the init sas_token parameter, it will be used
-        for authentication. Otherwise, a delegation SAS token will be created using Entra ID authentication.
+        for authentication. Otherwise, ``DefaultAzureCredential`` is used directly, which requires the caller
+        to hold a data-plane role such as Storage Blob Data Contributor on the storage account.
 
         Returns:
             AsyncContainerClient: The initialized container client.
+
+        Raises:
+            ValueError: If the container URL does not include a container name in its path.
         """
         from azure.storage.blob.aio import ContainerClient as AsyncContainerClient
 
-        from pyrit.auth import AzureStorageAuth
+        if self._sas_token:
+            self._client_async = AsyncContainerClient.from_container_url(
+                container_url=self._container_url,
+                credential=self._sas_token,
+            )
+            return self._client_async
 
-        sas_token = self._sas_token
-        if not self._sas_token:
-            logger.info("SAS token not provided. Creating a delegation SAS token using Entra ID authentication.")
-            sas_token = await AzureStorageAuth.get_sas_token_async(self._container_url)
+        from azure.identity.aio import DefaultAzureCredential
 
-        self._client_async = AsyncContainerClient.from_container_url(
-            container_url=self._container_url,
-            credential=sas_token,
+        logger.info("SAS token not provided. Using DefaultAzureCredential for direct Entra ID authentication.")
+        parsed_url = urlparse(self._container_url)
+        path_parts = [part for part in parsed_url.path.split("/") if part]
+        if not path_parts:
+            raise ValueError(
+                f"Invalid Azure Storage container URL '{self._container_url}': expected a container name in the "
+                "path, e.g. https://<account>.blob.core.windows.net/<container>."
+            )
+        account_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
+        container_name = path_parts[0]
+        self._credential = DefaultAzureCredential()
+        self._client_async = AsyncContainerClient(
+            account_url=account_url,
+            container_name=container_name,
+            credential=self._credential,
         )
         return self._client_async
 
+    async def _close_client_async(self) -> None:
+        """Close the container client and credential, resetting both to None."""
+        client, self._client_async = self._client_async, None
+        credential, self._credential = self._credential, None
+        try:
+            if client:
+                await client.close()  # type: ignore[no-untyped-call, unused-ignore]
+        finally:
+            if credential:
+                await credential.close()
+
     async def _upload_blob_async(self, file_name: str, data: bytes, content_type: str) -> None:
         """
         (Async) Handles uploading blob to given storage container.
@@ -321,10 +352,10 @@ async def _upload_blob_async(self, file_name: str, data: bytes, content_type: st
         except Exception as exc:
             if isinstance(exc, ClientAuthenticationError):
                 logger.exception(
-                    msg="Authentication failed. Please check that the container existence in the "
-                    "Azure Storage Account and ensure the validity of the provided SAS token. If you "
-                    "haven't set the SAS token as an environment variable use `az login` to "
-                    "enable delegation-based SAS authentication to connect to the storage account"
+                    msg="Authentication failed. Please check that the container exists in the "
+                    "Azure Storage Account. If using a SAS token, ensure it is valid. Otherwise, "
+                    "ensure you are logged in via `az login` and hold a data-plane role such as "
+                    "Storage Blob Data Contributor on the storage account."
                 )
                 raise
             logger.exception(msg=f"An unexpected error occurred: {exc}")
@@ -417,8 +448,7 @@ async def read_file_async(self, path: Path | str) -> bytes:
             logger.exception(f"Failed to read file at {blob_name}: {exc}")
             raise
         finally:
-            await self._client_async.close()
-            self._client_async = None
+            await self._close_client_async()
 
     async def write_file_async(self, path: Path | str, data: bytes) -> None:
         """
@@ -440,8 +470,7 @@ async def write_file_async(self, path: Path | str, data: bytes) -> None:
             logger.exception(f"Failed to write file at {blob_name}: {exc}")
             raise
         finally:
-            await self._client_async.close()
-            self._client_async = None
+            await self._close_client_async()
 
     async def path_exists_async(self, path: Path | str) -> bool:
         """
@@ -465,8 +494,7 @@ async def path_exists_async(self, path: Path | str) -> bool:
         except ResourceNotFoundError:
             return False
         finally:
-            await self._client_async.close()
-            self._client_async = None
+            await self._close_client_async()
 
     async def is_file_async(self, path: Path | str) -> bool:
         """
@@ -490,8 +518,7 @@ async def is_file_async(self, path: Path | str) -> bool:
         except ResourceNotFoundError:
             return False
         finally:
-            await self._client_async.close()
-            self._client_async = None
+            await self._close_client_async()
 
     async def create_directory_if_not_exists_async(self, directory_path: Path | str) -> None:  # type: ignore[ty:invalid-method-override]
         """

pyrit/prompt_target/azure_blob_storage_target.py

@@ -6,10 +6,10 @@
 from urllib.parse import urlparse
 
 from azure.core.exceptions import ClientAuthenticationError
+from azure.identity.aio import DefaultAzureCredential
 from azure.storage.blob import ContentSettings
 from azure.storage.blob.aio import ContainerClient as AsyncContainerClient
 
-from pyrit.auth import AzureStorageAuth
 from pyrit.common import default_values
 from pyrit.models import ComponentIdentifier, Message, construct_response_from_request
 from pyrit.prompt_target.common.prompt_target import PromptTarget
@@ -37,8 +37,9 @@ class AzureBlobStorageTarget(PromptTarget):
 
     Args:
         container_url (str): URL to the Azure Blob Storage Container.
-        sas_token (optional[str]): Optional Blob SAS token needed to authenticate blob operations. If not provided, a
-            delegation SAS token will be created using Entra ID authentication.
+        sas_token (optional[str]): Optional Blob SAS token needed to authenticate blob operations. If not provided,
+            ``DefaultAzureCredential`` is used directly, which requires the caller to hold a data-plane role such
+            as Storage Blob Data Contributor on the storage account.
         blob_content_type (SupportedContentType): Expected Content Type of the blob, chosen from the
             SupportedContentType enum. Set to PLAIN_TEXT by default.
         max_requests_per_minute (int, Optional): Number of requests the target can handle per
@@ -96,6 +97,7 @@ def __init__(
 
         self._sas_token: str | None = sas_token
         self._client_async: AsyncContainerClient | None = None
+        self._credential: DefaultAzureCredential | None = None
 
         super().__init__(
             endpoint=self._container_url,
@@ -121,22 +123,42 @@ async def _create_container_client_async(self) -> None:
         """
         Create an asynchronous ContainerClient for Azure Storage. If a SAS token is provided via the
         AZURE_STORAGE_ACCOUNT_SAS_TOKEN environment variable or the init sas_token parameter, it will be used
-        for authentication. Otherwise, a delegation SAS token will be created using Entra ID authentication.
+        for authentication. Otherwise, ``DefaultAzureCredential`` is used directly, which requires the caller
+        to hold a data-plane role such as Storage Blob Data Contributor on the storage account.
         """
         container_url, _ = self._parse_url()
         try:
             sas_token: str = default_values.get_required_value(
                 env_var_name=self.SAS_TOKEN_ENVIRONMENT_VARIABLE, passed_value=self._sas_token
             )
-            logger.info("Using SAS token from environment variable or passed parameter.")
         except ValueError:
-            logger.info("SAS token not provided. Creating a delegation SAS token using Entra ID authentication.")
-            sas_token = await AzureStorageAuth.get_sas_token_async(container_url)
+            logger.info("SAS token not provided. Using DefaultAzureCredential for direct Entra ID authentication.")
+            account_url, _, container_name = container_url.rpartition("/")
+            self._credential = DefaultAzureCredential()
+            self._client_async = AsyncContainerClient(
+                account_url=account_url,
+                container_name=container_name,
+                credential=self._credential,
+            )
+            return
+
+        logger.info("Using SAS token from environment variable or passed parameter.")
         self._client_async = AsyncContainerClient.from_container_url(
             container_url=container_url,
             credential=sas_token,
         )
 
+    async def _close_client_async(self) -> None:
+        """Close the container client and credential, resetting both to None."""
+        client, self._client_async = self._client_async, None
+        credential, self._credential = self._credential, None
+        try:
+            if client:
+                await client.close()
+        finally:
+            if credential:
+                await credential.close()
+
     async def _upload_blob_async(self, file_name: str, data: bytes, content_type: str) -> None:
         """
         (Async) Handles uploading blob to given storage container.
@@ -152,14 +174,14 @@ async def _upload_blob_async(self, file_name: str, data: bytes, content_type: st
         content_settings = ContentSettings(content_type=f"{content_type}")
         logger.info(msg="\nUploading to Azure Storage as blob:\n\t" + file_name)
 
-        if not self._client_async:
-            await self._create_container_client_async()
         # Parse the Azure Storage Blob URL to extract components
         _, blob_prefix = self._parse_url()
         # If a blob prefix is provided, prepend it to the file name.
         # If not, the file will be put in the root of the container.
         blob_path = f"{blob_prefix}/{file_name}" if blob_prefix else file_name
         try:
+            if not self._client_async:
+                await self._create_container_client_async()
             if self._client_async is None:
                 raise RuntimeError("Blob storage client not initialized")
             blob_client = self._client_async.get_blob_client(blob=blob_path)
@@ -171,24 +193,34 @@ async def _upload_blob_async(self, file_name: str, data: bytes, content_type: st
         except Exception as exc:
             if isinstance(exc, ClientAuthenticationError):
                 logger.exception(
-                    msg="Authentication failed. Please check that the container existence in the "
-                    "Azure Storage Account and ensure the validity of the provided SAS token. If you "
-                    "haven't set the SAS token as an environment variable use `az login` to "
-                    "enable delegation-based SAS authentication to connect to the storage account"
+                    msg="Authentication failed. Please check that the container exists in the "
+                    "Azure Storage Account. If using a SAS token, ensure it is valid. Otherwise, "
+                    "ensure you are logged in via `az login` and hold a data-plane role such as "
+                    "Storage Blob Data Contributor on the storage account."
                 )
                 raise
             logger.exception(msg=f"An unexpected error occurred: {exc}")
             raise
+        finally:
+            await self._close_client_async()
 
     def _parse_url(self) -> tuple[str, str]:
         """
         Parse the Azure Storage Blob URL to extract components.
 
         Returns:
             tuple: A tuple containing the container URL and blob prefix.
+
+        Raises:
+            ValueError: If the container URL does not include a container name in its path.
         """
         parsed_url = urlparse(self._container_url)
         path_parts = parsed_url.path.split("/")
+        if len(path_parts) < 2 or not path_parts[1]:
+            raise ValueError(
+                f"Invalid Azure Storage container URL '{self._container_url}': expected a container name in the "
+                "path, e.g. https://<account>.blob.core.windows.net/<container>."
+            )
         container_name = path_parts[1]
         blob_prefix = "/".join(path_parts[2:])
         container_url = f"https://{parsed_url.netloc}/{container_name}"

tests/unit/memory/storage/test_storage.py

@@ -169,20 +169,60 @@ async def test_azure_blob_storage_io_create_container_client_uses_explicit_sas_t
 
     mock_container_client = AsyncMock()
 
+    with patch(
+        "azure.storage.blob.aio.ContainerClient.from_container_url", return_value=mock_container_client
+    ) as mock_from_container_url:
+        await azure_blob_storage_io._create_container_client_async()
+
+    mock_from_container_url.assert_called_once_with(container_url=container_url, credential=sas_token)
+    assert azure_blob_storage_io._client_async is mock_container_client
+    assert azure_blob_storage_io._credential is None
+
+
+async def test_azure_blob_storage_io_create_container_client_uses_default_credential_when_no_sas_token():
+    container_url = "https://youraccount.blob.core.windows.net/yourcontainer"
+    azure_blob_storage_io = AzureBlobStorageIO(container_url=container_url)
+
+    mock_container_client = AsyncMock()
+    mock_credential = AsyncMock()
+
     with (
-        patch("pyrit.auth.AzureStorageAuth.get_sas_token_async", new_callable=AsyncMock) as mock_get_sas_token,
-        patch(
-            "azure.storage.blob.aio.ContainerClient.from_container_url", return_value=mock_container_client
-        ) as mock_from_container_url,
+        patch("azure.identity.aio.DefaultAzureCredential", return_value=mock_credential) as mock_credential_cls,
+        patch("azure.storage.blob.aio.ContainerClient", return_value=mock_container_client) as mock_container_cls,
     ):
         await azure_blob_storage_io._create_container_client_async()
 
-    mock_get_sas_token.assert_not_awaited()
-    mock_from_container_url.assert_called_once_with(container_url=container_url, credential=sas_token)
+    mock_credential_cls.assert_called_once()
+    mock_container_cls.assert_called_once_with(
+        account_url="https://youraccount.blob.core.windows.net",
+        container_name="yourcontainer",
+        credential=mock_credential,
+    )
     assert azure_blob_storage_io._client_async is mock_container_client
+    assert azure_blob_storage_io._credential is mock_credential
+
+
+async def test_azure_blob_storage_io_close_client_async_closes_credential_and_client():
+    azure_blob_storage_io = AzureBlobStorageIO(container_url="https://youraccount.blob.core.windows.net/yourcontainer")
+
+    mock_client = AsyncMock()
+    mock_credential = AsyncMock()
+    azure_blob_storage_io._client_async = mock_client
+    azure_blob_storage_io._credential = mock_credential
 
+    await azure_blob_storage_io._close_client_async()
 
-async def test_azure_storage_io_path_exists(azure_blob_storage_io):
+    mock_client.close.assert_awaited_once()
+    mock_credential.close.assert_awaited_once()
+    assert azure_blob_storage_io._client_async is None
+    assert azure_blob_storage_io._credential is None
+
+
+async def test_azure_blob_storage_io_create_container_client_raises_for_url_without_container():
+    azure_blob_storage_io = AzureBlobStorageIO(container_url="https://youraccount.blob.core.windows.net")
+
+    with pytest.raises(ValueError, match="expected a container name"):
+        await azure_blob_storage_io._create_container_client_async()
     azure_blob_storage_io._client_async = AsyncMock()
 
     mock_blob_client = AsyncMock()