From 531e33aecf2127da6ddc100dfcaa321ca947c65d Mon Sep 17 00:00:00 2001 From: lxcxjxhx Date: Sun, 19 Jul 2026 13:21:40 +0800 Subject: [PATCH 1/3] feat: expand SupportedContentType enum with common media types Add support for image, audio, video, and document content types to address TODO comment in storage.py. This enables PyRIT to handle multimodal data (images, audio, video, PDFs) when uploading to Azure Blob Storage. Added types: - Text: HTML, JSON, XML, CSV, Markdown - Image: PNG, JPEG, GIF, WebP, SVG, BMP - Audio: WAV, MP3, OGG, FLAC, M4A - Video: MP4, WebM, OGG, AVI - Document: PDF, ZIP, TAR, GZIP Includes unit tests to verify all enum values. --- pyrit/memory/storage/storage.py | 46 ++++++- tests/unit/memory/storage/test_storage.py | 156 ++++++++++++++++++---- 2 files changed, 169 insertions(+), 33 deletions(-) diff --git a/pyrit/memory/storage/storage.py b/pyrit/memory/storage/storage.py index 1cd5bac2a8..a230710b09 100644 --- a/pyrit/memory/storage/storage.py +++ b/pyrit/memory/storage/storage.py @@ -25,8 +25,40 @@ class SupportedContentType(Enum): See all options here: https://www.iana.org/assignments/media-types/media-types.xhtml. """ - # TODO, add other media supported types + # Text types PLAIN_TEXT = "text/plain" + HTML = "text/html" + JSON = "application/json" + XML = "application/xml" + CSV = "text/csv" + MARKDOWN = "text/markdown" + + # Image types + PNG = "image/png" + JPEG = "image/jpeg" + GIF = "image/gif" + WEBP = "image/webp" + SVG = "image/svg+xml" + BMP = "image/bmp" + + # Audio types + WAV = "audio/wav" + MP3 = "audio/mpeg" + OGG = "audio/ogg" + FLAC = "audio/flac" + M4A = "audio/mp4" + + # Video types + MP4 = "video/mp4" + WEBM = "video/webm" + OGG_VIDEO = "video/ogg" + AVI = "video/x-msvideo" + + # Document types + PDF = "application/pdf" + ZIP = "application/zip" + TAR = "application/x-tar" + GZIP = "application/gzip" class StorageIO(ABC): @@ -211,7 +243,9 @@ async def _create_container_client_async(self) -> AsyncContainerClient: from azure.identity.aio import DefaultAzureCredential - logger.info("SAS token not provided. Using DefaultAzureCredential for direct Entra ID authentication.") + 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: @@ -240,7 +274,9 @@ async def _close_client_async(self) -> None: if credential: await credential.close() - async def _upload_blob_async(self, file_name: str, data: bytes, content_type: str) -> None: + async def _upload_blob_async( + self, file_name: str, data: bytes, content_type: str + ) -> None: """ (Async) Handles uploading blob to given storage container. @@ -383,7 +419,9 @@ async def write_file_async(self, path: Path | str, data: bytes) -> None: self._client_async = await self._create_container_client_async() blob_name = self._resolve_blob_name(path) try: - await self._upload_blob_async(file_name=blob_name, data=data, content_type=self._blob_content_type) + await self._upload_blob_async( + file_name=blob_name, data=data, content_type=self._blob_content_type + ) except Exception as exc: logger.exception(f"Failed to write file at {blob_name}: {exc}") raise diff --git a/tests/unit/memory/storage/test_storage.py b/tests/unit/memory/storage/test_storage.py index d7c9720885..6b9b7d10e6 100644 --- a/tests/unit/memory/storage/test_storage.py +++ b/tests/unit/memory/storage/test_storage.py @@ -71,19 +71,26 @@ async def test_disk_storage_io_create_directory_if_not_exists(): storage = DiskStorageIO() directory_path = "sample_dir" - with patch("pathlib.Path.mkdir") as mock_mkdir, patch("pathlib.Path.exists", return_value=False) as mock_exists: + with ( + patch("pathlib.Path.mkdir") as mock_mkdir, + patch("pathlib.Path.exists", return_value=False) as mock_exists, + ): await storage.create_directory_if_not_exists_async(directory_path) mock_exists.assert_called_once() mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) async def test_azure_blob_storage_io_read_file(azure_blob_storage_io): - azure_blob_storage_io._client_async = AsyncMock() # Use Mock since get_blob_client is sync + azure_blob_storage_io._client_async = ( + AsyncMock() + ) # Use Mock since get_blob_client is sync mock_blob_client = AsyncMock() mock_blob_stream = AsyncMock() - azure_blob_storage_io._client_async.get_blob_client = Mock(return_value=mock_blob_client) + azure_blob_storage_io._client_async.get_blob_client = Mock( + return_value=mock_blob_client + ) mock_blob_client.download_blob = AsyncMock(return_value=mock_blob_stream) mock_blob_stream.readall = AsyncMock(return_value=b"Test file content") azure_blob_storage_io._client_async.close = AsyncMock() @@ -95,7 +102,9 @@ async def test_azure_blob_storage_io_read_file(azure_blob_storage_io): assert result == b"Test file content" -async def test_azure_blob_storage_io_read_file_with_relative_path(azure_blob_storage_io): +async def test_azure_blob_storage_io_read_file_with_relative_path( + azure_blob_storage_io, +): mock_container_client = AsyncMock() azure_blob_storage_io._client_async = mock_container_client @@ -110,7 +119,9 @@ async def test_azure_blob_storage_io_read_file_with_relative_path(azure_blob_sto result = await azure_blob_storage_io.read_file_async("dir1/dir2/sample.png") assert result == b"Test file content" - mock_container_client.get_blob_client.assert_called_once_with(blob="dir1/dir2/sample.png") + mock_container_client.get_blob_client.assert_called_once_with( + blob="dir1/dir2/sample.png" + ) async def test_azure_blob_storage_io_write_file(): @@ -126,7 +137,9 @@ async def test_azure_blob_storage_io_write_file(): mock_container_client.get_blob_client.return_value = mock_blob_client - with patch.object(azure_blob_storage_io, "_create_container_client_async", return_value=None): + with patch.object( + azure_blob_storage_io, "_create_container_client_async", return_value=None + ): azure_blob_storage_io._client_async = mock_container_client azure_blob_storage_io._upload_blob_async = AsyncMock() @@ -136,7 +149,9 @@ async def test_azure_blob_storage_io_write_file(): await azure_blob_storage_io.write_file_async(path, data_to_write) azure_blob_storage_io._upload_blob_async.assert_awaited_with( - file_name="testfile.txt", data=data_to_write, content_type=SupportedContentType.PLAIN_TEXT.value + file_name="testfile.txt", + data=data_to_write, + content_type=SupportedContentType.PLAIN_TEXT.value, ) @@ -148,12 +163,16 @@ async def test_azure_blob_storage_io_write_file_with_relative_path(): mock_container_client = AsyncMock() - with patch.object(azure_blob_storage_io, "_create_container_client_async", return_value=None): + with patch.object( + azure_blob_storage_io, "_create_container_client_async", return_value=None + ): azure_blob_storage_io._client_async = mock_container_client azure_blob_storage_io._upload_blob_async = AsyncMock() data_to_write = b"Test data" - await azure_blob_storage_io.write_file_async("dir1/dir2/testfile.txt", data_to_write) + await azure_blob_storage_io.write_file_async( + "dir1/dir2/testfile.txt", data_to_write + ) azure_blob_storage_io._upload_blob_async.assert_awaited_with( file_name="dir1/dir2/testfile.txt", @@ -165,16 +184,21 @@ async def test_azure_blob_storage_io_write_file_with_relative_path(): async def test_azure_blob_storage_io_create_container_client_uses_explicit_sas_token(): container_url = "https://youraccount.blob.core.windows.net/yourcontainer" sas_token = "explicit-sas-token" - azure_blob_storage_io = AzureBlobStorageIO(container_url=container_url, sas_token=sas_token) + azure_blob_storage_io = AzureBlobStorageIO( + container_url=container_url, sas_token=sas_token + ) mock_container_client = AsyncMock() with patch( - "azure.storage.blob.aio.ContainerClient.from_container_url", return_value=mock_container_client + "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) + 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 @@ -187,8 +211,12 @@ async def test_azure_blob_storage_io_create_container_client_uses_default_creden mock_credential = AsyncMock() with ( - 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, + 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() @@ -203,7 +231,9 @@ async def test_azure_blob_storage_io_create_container_client_uses_default_creden 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") + azure_blob_storage_io = AzureBlobStorageIO( + container_url="https://youraccount.blob.core.windows.net/yourcontainer" + ) mock_client = AsyncMock() mock_credential = AsyncMock() @@ -219,7 +249,9 @@ async def test_azure_blob_storage_io_close_client_async_closes_credential_and_cl 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") + 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() @@ -227,10 +259,14 @@ async def test_azure_blob_storage_io_create_container_client_raises_for_url_with mock_blob_client = AsyncMock() - azure_blob_storage_io._client_async.get_blob_client = Mock(return_value=mock_blob_client) + azure_blob_storage_io._client_async.get_blob_client = Mock( + return_value=mock_blob_client + ) mock_blob_client.get_blob_properties = AsyncMock() azure_blob_storage_io._client_async.close = AsyncMock() - file_path = "https://example.blob.core.windows.net/container/dir1/dir2/blob_name.txt" + file_path = ( + "https://example.blob.core.windows.net/container/dir1/dir2/blob_name.txt" + ) exists = await azure_blob_storage_io.path_exists_async(file_path) assert exists is True @@ -248,7 +284,9 @@ async def test_azure_storage_io_path_exists_with_relative_path(azure_blob_storag exists = await azure_blob_storage_io.path_exists_async("dir1/dir2/blob_name.txt") assert exists is True - mock_container_client.get_blob_client.assert_called_once_with(blob="dir1/dir2/blob_name.txt") + mock_container_client.get_blob_client.assert_called_once_with( + blob="dir1/dir2/blob_name.txt" + ) async def test_azure_storage_io_is_file(azure_blob_storage_io): @@ -256,11 +294,15 @@ async def test_azure_storage_io_is_file(azure_blob_storage_io): mock_blob_client = AsyncMock() - azure_blob_storage_io._client_async.get_blob_client = Mock(return_value=mock_blob_client) + azure_blob_storage_io._client_async.get_blob_client = Mock( + return_value=mock_blob_client + ) mock_blob_properties = Mock(size=1024) mock_blob_client.get_blob_properties = AsyncMock(return_value=mock_blob_properties) azure_blob_storage_io._client_async.close = AsyncMock() - file_path = "https://example.blob.core.windows.net/container/dir1/dir2/blob_name.txt" + file_path = ( + "https://example.blob.core.windows.net/container/dir1/dir2/blob_name.txt" + ) is_file = await azure_blob_storage_io.is_file_async(file_path) assert is_file is True @@ -279,11 +321,15 @@ async def test_azure_storage_io_is_file_with_relative_path(azure_blob_storage_io is_file = await azure_blob_storage_io.is_file_async("dir1/dir2/blob_name.txt") assert is_file is True - mock_container_client.get_blob_client.assert_called_once_with(blob="dir1/dir2/blob_name.txt") + mock_container_client.get_blob_client.assert_called_once_with( + blob="dir1/dir2/blob_name.txt" + ) def test_azure_storage_io_parse_blob_url_valid(azure_blob_storage_io): - file_path = "https://example.blob.core.windows.net/container/dir1/dir2/blob_name.txt" + file_path = ( + "https://example.blob.core.windows.net/container/dir1/dir2/blob_name.txt" + ) container_name, blob_name = azure_blob_storage_io.parse_blob_url(file_path) assert container_name == "container" @@ -297,7 +343,9 @@ def test_azure_storage_io_parse_blob_url_invalid(azure_blob_storage_io): def test_azure_storage_io_parse_blob_url_without_scheme(azure_blob_storage_io): with pytest.raises(ValueError, match="Invalid blob URL"): - azure_blob_storage_io.parse_blob_url("example.blob.core.windows.net/container/dir1/blob_name.txt") + azure_blob_storage_io.parse_blob_url( + "example.blob.core.windows.net/container/dir1/blob_name.txt" + ) def test_azure_storage_io_parse_blob_url_without_netloc(azure_blob_storage_io): @@ -306,12 +354,17 @@ def test_azure_storage_io_parse_blob_url_without_netloc(azure_blob_storage_io): def test_resolve_blob_name_with_full_url(azure_blob_storage_io): - result = azure_blob_storage_io._resolve_blob_name("https://account.blob.core.windows.net/container/dir1/file.txt") + result = azure_blob_storage_io._resolve_blob_name( + "https://account.blob.core.windows.net/container/dir1/file.txt" + ) assert result == "dir1/file.txt" def test_resolve_blob_name_with_relative_path(azure_blob_storage_io): - assert azure_blob_storage_io._resolve_blob_name("dir1/dir2/file.txt") == "dir1/dir2/file.txt" + assert ( + azure_blob_storage_io._resolve_blob_name("dir1/dir2/file.txt") + == "dir1/dir2/file.txt" + ) def test_resolve_blob_name_with_simple_filename(azure_blob_storage_io): @@ -319,13 +372,18 @@ def test_resolve_blob_name_with_simple_filename(azure_blob_storage_io): def test_resolve_blob_name_normalizes_backslashes(azure_blob_storage_io): - assert azure_blob_storage_io._resolve_blob_name("dir1\\dir2\\file.txt") == "dir1/dir2/file.txt" + assert ( + azure_blob_storage_io._resolve_blob_name("dir1\\dir2\\file.txt") + == "dir1/dir2/file.txt" + ) def test_resolve_blob_name_with_path_object(azure_blob_storage_io): from pathlib import PurePosixPath - result = azure_blob_storage_io._resolve_blob_name(PurePosixPath("dir1/dir2/file.txt")) + result = azure_blob_storage_io._resolve_blob_name( + PurePosixPath("dir1/dir2/file.txt") + ) assert result == "dir1/dir2/file.txt" @@ -333,7 +391,9 @@ async def test_upload_blob_raises_when_client_async_none(): obj = AzureBlobStorageIO.__new__(AzureBlobStorageIO) obj._client_async = None with pytest.raises(RuntimeError, match="Azure container client not initialized"): - await obj._upload_blob_async(file_name="test.txt", data=b"data", content_type="text/plain") + await obj._upload_blob_async( + file_name="test.txt", data=b"data", content_type="text/plain" + ) async def test_read_file_lazy_initializes_client(azure_blob_storage_io): @@ -417,3 +477,41 @@ async def test_is_file_lazy_initializes_client(azure_blob_storage_io): mock_create.assert_called_once() assert result is True + + +def test_supported_content_type_has_expected_values(): + """Test that SupportedContentType enum contains expected media types.""" + # Text types + assert SupportedContentType.PLAIN_TEXT.value == "text/plain" + assert SupportedContentType.HTML.value == "text/html" + assert SupportedContentType.JSON.value == "application/json" + assert SupportedContentType.XML.value == "application/xml" + assert SupportedContentType.CSV.value == "text/csv" + assert SupportedContentType.MARKDOWN.value == "text/markdown" + + # Image types + assert SupportedContentType.PNG.value == "image/png" + assert SupportedContentType.JPEG.value == "image/jpeg" + assert SupportedContentType.GIF.value == "image/gif" + assert SupportedContentType.WEBP.value == "image/webp" + assert SupportedContentType.SVG.value == "image/svg+xml" + assert SupportedContentType.BMP.value == "image/bmp" + + # Audio types + assert SupportedContentType.WAV.value == "audio/wav" + assert SupportedContentType.MP3.value == "audio/mpeg" + assert SupportedContentType.OGG.value == "audio/ogg" + assert SupportedContentType.FLAC.value == "audio/flac" + assert SupportedContentType.M4A.value == "audio/mp4" + + # Video types + assert SupportedContentType.MP4.value == "video/mp4" + assert SupportedContentType.WEBM.value == "video/webm" + assert SupportedContentType.OGG_VIDEO.value == "video/ogg" + assert SupportedContentType.AVI.value == "video/x-msvideo" + + # Document types + assert SupportedContentType.PDF.value == "application/pdf" + assert SupportedContentType.ZIP.value == "application/zip" + assert SupportedContentType.TAR.value == "application/x-tar" + assert SupportedContentType.GZIP.value == "application/gzip" From db28e7a4b81ff617cd419b32e86a208b218a3f0a Mon Sep 17 00:00:00 2001 From: lxcxjxhx Date: Mon, 20 Jul 2026 04:18:03 +0800 Subject: [PATCH 2/3] fix: integrate MIME type selection into file upload flow - Modify AzureBlobStorageIO.write_file_async to accept content_type parameter - Infer content_type from file extension using mimetypes.guess_type - Pass content_type through serializers (save_data_async, save_b64_image_async, save_formatted_audio_async) - Consolidate SupportedContentType enum imports - Add unit tests for content_type inference and propagation Addresses maintainer feedback on PR #2233 --- doc/code/executor/5_workflow.py | 2 +- pyrit/memory/storage/serializers.py | 103 ++++-- pyrit/memory/storage/storage.py | 26 +- .../azure_blob_storage_target.py | 56 +-- tests/unit/memory/storage/test_serializers.py | 336 +++++++++++++++--- tests/unit/memory/storage/test_storage.py | 83 +++++ 6 files changed, 499 insertions(+), 107 deletions(-) diff --git a/doc/code/executor/5_workflow.py b/doc/code/executor/5_workflow.py index 190c4faefc..efb12eae19 100644 --- a/doc/code/executor/5_workflow.py +++ b/doc/code/executor/5_workflow.py @@ -139,7 +139,7 @@ async def processing_callback() -> str: from pyrit.executor.workflow import XPIAWorkflow from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.prompt_target import AzureBlobStorageTarget -from pyrit.prompt_target.azure_blob_storage_target import SupportedContentType +from pyrit.memory.storage.storage import SupportedContentType from pyrit.score import SubStringScorer logging.basicConfig(level=logging.DEBUG) diff --git a/pyrit/memory/storage/serializers.py b/pyrit/memory/storage/serializers.py index f779cc693e..4d46c78c92 100644 --- a/pyrit/memory/storage/serializers.py +++ b/pyrit/memory/storage/serializers.py @@ -74,20 +74,36 @@ def data_serializer_factory( f"The 'category' argument is mandatory and must be one of the following: {get_args(AllowedCategories)}." ) if value is not None: - if data_type in ["text", "reasoning", "function_call", "tool_call", "function_call_output"]: + if data_type in [ + "text", + "reasoning", + "function_call", + "tool_call", + "function_call_output", + ]: return TextDataTypeSerializer(prompt_text=value, data_type=data_type) if data_type == "image_path": - return ImagePathDataTypeSerializer(category=category, prompt_text=value, extension=extension) + return ImagePathDataTypeSerializer( + category=category, prompt_text=value, extension=extension + ) if data_type == "audio_path": - return AudioPathDataTypeSerializer(category=category, prompt_text=value, extension=extension) + return AudioPathDataTypeSerializer( + category=category, prompt_text=value, extension=extension + ) if data_type == "video_path": - return VideoPathDataTypeSerializer(category=category, prompt_text=value, extension=extension) + return VideoPathDataTypeSerializer( + category=category, prompt_text=value, extension=extension + ) if data_type == "binary_path": - return BinaryPathDataTypeSerializer(category=category, prompt_text=value, extension=extension) + return BinaryPathDataTypeSerializer( + category=category, prompt_text=value, extension=extension + ) if data_type == "error": return ErrorDataTypeSerializer(prompt_text=value) if data_type == "url": - return URLDataTypeSerializer(category=category, prompt_text=value, extension=extension) + return URLDataTypeSerializer( + category=category, prompt_text=value, extension=extension + ) raise ValueError(f"Data type {data_type} not supported") if data_type == "image_path": return ImagePathDataTypeSerializer(category=category, extension=extension) @@ -139,7 +155,9 @@ def _get_storage_io(self) -> StorageIO: # Scenarios where a user utilizes an in-memory DuckDB but also needs to interact # with an Azure Storage Account, ex., XPIAWorkflow. if self._memory.results_storage_io is None: - raise RuntimeError("results_storage_io is not configured but Azure storage URL was detected") + raise RuntimeError( + "results_storage_io is not configured but Azure storage URL was detected" + ) return self._memory.results_storage_io return DiskStorageIO() @@ -153,7 +171,9 @@ def data_on_disk(self) -> bool: """ - async def save_data_async(self, data: bytes, output_filename: str | None = None) -> None: + async def save_data_async( + self, data: bytes, output_filename: str | None = None + ) -> None: """ Save data to storage. @@ -167,10 +187,16 @@ async def save_data_async(self, data: bytes, output_filename: str | None = None) file_path = await self.get_data_filename_async(file_name=output_filename) if self._memory.results_storage_io is None: raise RuntimeError("Storage IO not initialized") - await self._memory.results_storage_io.write_file_async(file_path, data) + # Infer content type from file extension + content_type = self.get_mime_type(str(file_path)) + await self._memory.results_storage_io.write_file_async( + file_path, data, content_type=content_type + ) self.value = str(file_path) - async def save_b64_image_async(self, data: str | bytes, output_filename: str | None = None) -> None: + async def save_b64_image_async( + self, data: str | bytes, output_filename: str | None = None + ) -> None: """ Save a base64-encoded image to storage. @@ -185,7 +211,11 @@ async def save_b64_image_async(self, data: str | bytes, output_filename: str | N image_bytes = base64.b64decode(data) if self._memory.results_storage_io is None: raise RuntimeError("Storage IO not initialized") - await self._memory.results_storage_io.write_file_async(file_path, image_bytes) + # Infer content type from file extension + content_type = self.get_mime_type(str(file_path)) + await self._memory.results_storage_io.write_file_async( + file_path, image_bytes, content_type=content_type + ) self.value = str(file_path) async def save_formatted_audio_async( @@ -213,7 +243,9 @@ async def save_formatted_audio_async( # save audio file locally first if in AzureStorageBlob so we can use wave.open to set audio parameters if self._is_azure_storage_url(str(file_path)): - with tempfile.NamedTemporaryFile(suffix=".wav", dir=DB_DATA_PATH, delete=False) as tmp: + with tempfile.NamedTemporaryFile( + suffix=".wav", dir=DB_DATA_PATH, delete=False + ) as tmp: local_temp_path = Path(tmp.name) try: await asyncio.to_thread( @@ -227,8 +259,14 @@ async def save_formatted_audio_async( async with aiofiles.open(local_temp_path, "rb") as f: audio_data = await f.read() if self._memory.results_storage_io is None: - raise RuntimeError("self._memory.results_storage_io is not initialized") - await self._memory.results_storage_io.write_file_async(file_path, audio_data) + raise RuntimeError( + "self._memory.results_storage_io is not initialized" + ) + # Infer content type from file extension + content_type = self.get_mime_type(str(file_path)) + await self._memory.results_storage_io.write_file_async( + file_path, audio_data, content_type=content_type + ) finally: local_temp_path.unlink(missing_ok=True) @@ -259,7 +297,9 @@ async def read_data_async(self) -> bytes: """ if not self.data_on_disk(): - raise TypeError(f"Data for data Type {self.data_type} is not stored on disk") + raise TypeError( + f"Data for data Type {self.data_type} is not stored on disk" + ) if not self.value: raise RuntimeError("Prompt text not set") @@ -309,7 +349,9 @@ async def get_sha256_async(self) -> str: if isinstance(self.value, str): input_bytes = self.value.encode("utf-8") else: - raise ValueError(f"Invalid data type {self.value}, expected str data type.") + raise ValueError( + f"Invalid data type {self.value}, expected str data type." + ) hash_object = hashlib.sha256(input_bytes) return hash_object.hexdigest() @@ -349,13 +391,19 @@ async def get_data_filename_async(self, file_name: str | None = None) -> Path | if self._is_azure_storage_url(results_path): full_data_directory_path = results_path + self.data_sub_directory - self._file_path = full_data_directory_path + f"/{file_name}.{self.file_extension}" + self._file_path = ( + full_data_directory_path + f"/{file_name}.{self.file_extension}" + ) else: full_data_directory_path = results_path + self.data_sub_directory if self._memory.results_storage_io is None: raise RuntimeError("self._memory.results_storage_io is not initialized") - await self._memory.results_storage_io.create_directory_if_not_exists_async(Path(full_data_directory_path)) - self._file_path = Path(full_data_directory_path, f"{file_name}.{self.file_extension}") + await self._memory.results_storage_io.create_directory_if_not_exists_async( + Path(full_data_directory_path) + ) + self._file_path = Path( + full_data_directory_path, f"{file_name}.{self.file_extension}" + ) return self._file_path @@ -401,7 +449,10 @@ def _is_azure_storage_url(self, path: str) -> bool: """ parsed = urlparse(path) - return parsed.scheme in ("http", "https") and "blob.core.windows.net" in parsed.netloc + return ( + parsed.scheme in ("http", "https") + and "blob.core.windows.net" in parsed.netloc + ) class TextDataTypeSerializer(DataTypeSerializer): @@ -458,7 +509,9 @@ def data_on_disk(self) -> bool: class URLDataTypeSerializer(DataTypeSerializer): """Serializer for URL values and URL-backed local file references.""" - def __init__(self, *, category: str, prompt_text: str, extension: str | None = None) -> None: + def __init__( + self, *, category: str, prompt_text: str, extension: str | None = None + ) -> None: """ Initialize a URL serializer. @@ -488,7 +541,13 @@ def data_on_disk(self) -> bool: class ImagePathDataTypeSerializer(DataTypeSerializer): """Serializer for image path values stored on disk.""" - def __init__(self, *, category: str, prompt_text: str | None = None, extension: str | None = None) -> None: + def __init__( + self, + *, + category: str, + prompt_text: str | None = None, + extension: str | None = None, + ) -> None: """ Initialize an image-path serializer. diff --git a/pyrit/memory/storage/storage.py b/pyrit/memory/storage/storage.py index a230710b09..443a966d6b 100644 --- a/pyrit/memory/storage/storage.py +++ b/pyrit/memory/storage/storage.py @@ -73,7 +73,9 @@ async def read_file_async(self, path: Path | str) -> bytes: """ @abstractmethod - async def write_file_async(self, path: Path | str, data: bytes) -> None: + async def write_file_async( + self, path: Path | str, data: bytes, content_type: str | None = None + ) -> None: """ Asynchronously writes data to the given path. """ @@ -117,13 +119,16 @@ async def read_file_async(self, path: Path | str) -> bytes: async with aiofiles.open(path, "rb") as file: return await file.read() - async def write_file_async(self, path: Path | str, data: bytes) -> None: + async def write_file_async( + self, path: Path | str, data: bytes, content_type: str | None = None + ) -> None: """ Asynchronously writes data to a file on the local disk. Args: path (Path): The path to the file. data (bytes): The content to write to the file. + content_type (str | None): Optional content type (unused for local disk storage). """ path = self._convert_to_path(path) @@ -404,7 +409,9 @@ async def read_file_async(self, path: Path | str) -> bytes: finally: await self._close_client_async() - async def write_file_async(self, path: Path | str, data: bytes) -> None: + async def write_file_async( + self, path: Path | str, data: bytes, content_type: str | None = None + ) -> None: """ Write data to Azure Blob Storage at the specified path. @@ -414,13 +421,24 @@ async def write_file_async(self, path: Path | str, data: bytes) -> None: Args: path (Path | str): Full blob URL or relative blob path. data (bytes): The data to write. + content_type (str | None): Optional MIME type for the blob. If not provided, + it will be inferred from the file extension, falling back to the default + blob_content_type configured in the constructor. """ if not self._client_async: self._client_async = await self._create_container_client_async() blob_name = self._resolve_blob_name(path) + + # Determine content type: explicit > inferred from extension > default + if content_type is None: + from mimetypes import guess_type + + inferred_type, _ = guess_type(blob_name) + content_type = inferred_type if inferred_type else self._blob_content_type + try: await self._upload_blob_async( - file_name=blob_name, data=data, content_type=self._blob_content_type + file_name=blob_name, data=data, content_type=content_type ) except Exception as exc: logger.exception(f"Failed to write file at {blob_name}: {exc}") diff --git a/pyrit/prompt_target/azure_blob_storage_target.py b/pyrit/prompt_target/azure_blob_storage_target.py index d0cbe13dec..6311b1dde2 100644 --- a/pyrit/prompt_target/azure_blob_storage_target.py +++ b/pyrit/prompt_target/azure_blob_storage_target.py @@ -3,7 +3,6 @@ import logging import posixpath -from enum import Enum from typing import ClassVar from urllib.parse import urlparse @@ -13,6 +12,7 @@ from azure.storage.blob.aio import ContainerClient as AsyncContainerClient from pyrit.common import default_values +from pyrit.memory.storage.storage import SupportedContentType from pyrit.models import ComponentIdentifier, Message, construct_response_from_request from pyrit.prompt_target.common.prompt_target import AuthMode, PromptTarget from pyrit.prompt_target.common.target_capabilities import TargetCapabilities @@ -22,16 +22,6 @@ logger = logging.getLogger(__name__) -class SupportedContentType(Enum): - """ - All supported content types for uploading blobs to provided storage account container. - See all options here: https://www.iana.org/assignments/media-types/media-types.xhtml. - """ - - PLAIN_TEXT = "text/plain" - HTML = "text/html" - - class AzureBlobStorageTarget(PromptTarget): """ The AzureBlobStorageTarget takes prompts, saves the prompts to a file, and stores them as a blob in a provided @@ -49,7 +39,9 @@ class AzureBlobStorageTarget(PromptTarget): will be capped at the value provided. """ - AZURE_STORAGE_CONTAINER_ENVIRONMENT_VARIABLE: str = "AZURE_STORAGE_ACCOUNT_CONTAINER_URL" + AZURE_STORAGE_CONTAINER_ENVIRONMENT_VARIABLE: str = ( + "AZURE_STORAGE_ACCOUNT_CONTAINER_URL" + ) SAS_TOKEN_ENVIRONMENT_VARIABLE: str = "AZURE_STORAGE_ACCOUNT_SAS_TOKEN" # A SAS token is the "api_key"; with no token the target falls back to @@ -98,7 +90,8 @@ def __init__( self._blob_content_type: str = blob_content_type.value self._container_url: str = default_values.get_required_value( - env_var_name=self.AZURE_STORAGE_CONTAINER_ENVIRONMENT_VARIABLE, passed_value=container_url + env_var_name=self.AZURE_STORAGE_CONTAINER_ENVIRONMENT_VARIABLE, + passed_value=container_url, ) self._sas_token: str | None = sas_token @@ -135,10 +128,13 @@ async def _create_container_client_async(self) -> None: 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 + env_var_name=self.SAS_TOKEN_ENVIRONMENT_VARIABLE, + passed_value=self._sas_token, ) except ValueError: - logger.info("SAS token not provided. Using DefaultAzureCredential for direct Entra ID authentication.") + 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( @@ -165,7 +161,9 @@ async def _close_client_async(self) -> None: if credential: await credential.close() - async def _upload_blob_async(self, file_name: str, data: bytes, content_type: str) -> None: + async def _upload_blob_async( + self, file_name: str, data: bytes, content_type: str + ) -> None: """ (Async) Handles uploading blob to given storage container. @@ -192,7 +190,9 @@ async def _upload_blob_async(self, file_name: str, data: bytes, content_type: st raise RuntimeError("Blob storage client not initialized") blob_client = self._client_async.get_blob_client(blob=blob_path) if await blob_client.exists(): - logger.info(msg=f"Blob {blob_path} already exists. Deleting it before uploading a new version.") + logger.info( + msg=f"Blob {blob_path} already exists. Deleting it before uploading a new version." + ) await blob_client.delete_blob() logger.info(msg=f"Uploading new blob to {blob_path}") await blob_client.upload_blob(data=data, content_settings=content_settings) @@ -233,7 +233,9 @@ def _parse_url(self) -> tuple[str, str]: return container_url, blob_prefix @limit_requests_per_minute - async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + async def _send_prompt_to_target_async( + self, *, normalized_conversation: list[Message] + ) -> list[Message]: """ (Async) Sends prompt to target, which creates a file and uploads it as a blob to the provided storage container. @@ -252,12 +254,16 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me # default file name is .txt, but can be overridden by prompt metadata file_name = f"{request.conversation_id}.txt" if request.prompt_metadata.get("file_name"): - file_name = self._sanitize_file_name(str(request.prompt_metadata["file_name"])) + file_name = self._sanitize_file_name( + str(request.prompt_metadata["file_name"]) + ) data = str.encode(request.converted_value) blob_url = self._container_url + "/" + file_name - await self._upload_blob_async(file_name=file_name, data=data, content_type=self._blob_content_type) + await self._upload_blob_async( + file_name=file_name, data=data, content_type=self._blob_content_type + ) response = construct_response_from_request( request=request, response_text_pieces=[blob_url], response_type="url" @@ -282,6 +288,12 @@ def _sanitize_file_name(file_name: str) -> str: Raises: ValueError: If *file_name* contains path separators or traversal segments. """ - if file_name != posixpath.basename(file_name) or file_name in (".", "..") or "\\" in file_name: - raise ValueError(f"Invalid file_name '{file_name}': must be a bare filename without path separators.") + if ( + file_name != posixpath.basename(file_name) + or file_name in (".", "..") + or "\\" in file_name + ): + raise ValueError( + f"Invalid file_name '{file_name}': must be a bare filename without path separators." + ) return file_name diff --git a/tests/unit/memory/storage/test_serializers.py b/tests/unit/memory/storage/test_serializers.py index b323a4e8b5..a898ab9a36 100644 --- a/tests/unit/memory/storage/test_serializers.py +++ b/tests/unit/memory/storage/test_serializers.py @@ -38,7 +38,9 @@ def test_data_serializer_factory_text_no_data_throws(sqlite_instance): def test_data_serializer_factory_text_with_data(sqlite_instance): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="text", value="test") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="text", value="test" + ) assert isinstance(serializer, DataTypeSerializer) assert isinstance(serializer, TextDataTypeSerializer) assert serializer.data_type == "text" @@ -47,7 +49,9 @@ def test_data_serializer_factory_text_with_data(sqlite_instance): def test_data_serializer_factory_error_with_data(sqlite_instance): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="error", value="test") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="error", value="test" + ) assert isinstance(serializer, DataTypeSerializer) assert isinstance(serializer, ErrorDataTypeSerializer) assert serializer.data_type == "error" @@ -56,25 +60,33 @@ def test_data_serializer_factory_error_with_data(sqlite_instance): async def test_data_serializer_text_read_data_throws(sqlite_instance): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="text", value="test") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="text", value="test" + ) with pytest.raises(TypeError): await serializer.read_data_async() async def test_data_serializer_text_save_data_throws(sqlite_instance): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="text", value="test") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="text", value="test" + ) with pytest.raises(TypeError): await serializer.save_data_async(b"\x00") async def test_data_serializer_error_read_data_throws(sqlite_instance): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="error", value="test") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="error", value="test" + ) with pytest.raises(TypeError): await serializer.read_data_async() async def test_data_serializer_error_save_data_throws(sqlite_instance): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="error", value="test") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="error", value="test" + ) with pytest.raises(TypeError): await serializer.save_data_async(b"\x00") @@ -91,7 +103,9 @@ async def test_data_serializer_factory_missing_category_raises_value_error(): def test_image_path_normalizer_factory(sqlite_instance): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) assert isinstance(serializer, DataTypeSerializer) assert isinstance(serializer, ImagePathDataTypeSerializer) assert serializer.data_type == "image_path" @@ -99,7 +113,9 @@ def test_image_path_normalizer_factory(sqlite_instance): async def test_image_path_save_data(sqlite_instance): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) await serializer.save_data_async(b"\x00") serializer_value = serializer.value assert serializer_value @@ -111,7 +127,9 @@ async def test_image_path_save_data(sqlite_instance): async def test_image_path_read_data(sqlite_instance): data = b"\x00\x11\x22\x33" - normalizer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + normalizer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) await normalizer.save_data_async(data) assert await normalizer.read_data_async() == data read_normalizer = data_serializer_factory( @@ -123,7 +141,9 @@ async def test_image_path_read_data(sqlite_instance): async def test_image_path_read_data_base64(sqlite_instance): data = b"AAAA" - normalizer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + normalizer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) await normalizer.save_data_async(data) base_64_data = await normalizer.read_data_base64_async() assert base_64_data @@ -132,7 +152,9 @@ async def test_image_path_read_data_base64(sqlite_instance): async def test_path_not_exists(sqlite_instance): file_path = "non_existing_file.txt" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path", value=file_path) + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path", value=file_path + ) with pytest.raises(FileNotFoundError): await serializer.read_data_async() @@ -155,7 +177,9 @@ def test_get_mime_type(sqlite_instance): async def test_save_b64_image(sqlite_instance): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) await serializer.save_b64_image_async("\x00") serializer_value = str(serializer.value) assert serializer_value @@ -167,7 +191,9 @@ async def test_save_b64_image(sqlite_instance): async def test_audio_path_save_data(sqlite_instance): """Test saving audio data to disk.""" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="audio_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="audio_path" + ) await serializer.save_data_async(b"audio_data") assert serializer.value.endswith(".mp3") assert os.path.exists(serializer.value) @@ -177,7 +203,9 @@ async def test_audio_path_save_data(sqlite_instance): async def test_audio_path_read_data(sqlite_instance): """Test reading audio data from disk.""" data = b"audio_content" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="audio_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="audio_path" + ) await serializer.save_data_async(data) read_data = await serializer.read_data_async() assert read_data == data @@ -185,7 +213,9 @@ async def test_audio_path_read_data(sqlite_instance): async def test_video_path_save_data(sqlite_instance): """Test saving video data to disk.""" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="video_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="video_path" + ) video_data = b"video_data" await serializer.save_data_async(video_data) assert serializer.value.endswith(".mp4") # Assuming the default extension is '.mp4' @@ -196,7 +226,9 @@ async def test_video_path_save_data(sqlite_instance): async def test_video_path_read_data(sqlite_instance): """Test reading video data from disk.""" video_data = b"video_content" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="video_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="video_path" + ) await serializer.save_data_async(video_data) read_data = await serializer.read_data_async() assert read_data == video_data @@ -206,7 +238,9 @@ async def test_video_path_save_with_custom_extension(sqlite_instance): """Test saving video data with a custom file extension.""" custom_extension = "avi" serializer = data_serializer_factory( - category="prompt-memory-entries", data_type="video_path", extension=custom_extension + category="prompt-memory-entries", + data_type="video_path", + extension=custom_extension, ) video_data = b"video_data" await serializer.save_data_async(video_data) @@ -217,7 +251,9 @@ async def test_video_path_save_with_custom_extension(sqlite_instance): async def test_get_sha256_from_text(sqlite_instance): """Test SHA256 hash calculation for text data.""" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="text", value="test_string") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="text", value="test_string" + ) sha256_hash = await serializer.get_sha256_async() expected_hash = hashlib.sha256(b"test_string").hexdigest() assert sha256_hash == expected_hash @@ -226,7 +262,9 @@ async def test_get_sha256_from_text(sqlite_instance): async def test_get_sha256_from_image_file(sqlite_instance): """Test SHA256 hash calculation for file data.""" data = b"file_content.png" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) await serializer.save_data_async(data) sha256_hash = await serializer.get_sha256_async() expected_hash = hashlib.sha256(data).hexdigest() @@ -238,7 +276,9 @@ def test_is_azure_storage_url(sqlite_instance): valid_url = "https://mystorageaccount.blob.core.windows.net/container/file.txt" invalid_url = "https://example.com/file.txt" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="url", value=valid_url) + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="url", value=valid_url + ) assert serializer._is_azure_storage_url(valid_url) is True assert serializer._is_azure_storage_url(invalid_url) is False @@ -255,9 +295,14 @@ async def test_read_data_local_file_with_dummy_image(sqlite_instance): with open(image_path, "rb") as f: mock_storage_io.read_file_async.return_value = f.read() - with patch("pyrit.memory.storage.serializers.DiskStorageIO", return_value=mock_storage_io): + with patch( + "pyrit.memory.storage.serializers.DiskStorageIO", + return_value=mock_storage_io, + ): serializer = data_serializer_factory( - category="prompt-memory-entries", data_type="image_path", value=image_path + category="prompt-memory-entries", + data_type="image_path", + value=image_path, ) data = await serializer.read_data_async() @@ -276,7 +321,9 @@ async def test_read_data_local_file_with_dummy_image(sqlite_instance): async def test_get_data_filename(sqlite_instance): """Test get_data_filename when a file_name is provided.""" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) provided_filename = "custom_image_name" filename = await serializer.get_data_filename_async(file_name=provided_filename) assert str(filename).endswith(f"{provided_filename}.{serializer.file_extension}") @@ -287,7 +334,9 @@ async def test_get_data_filename(sqlite_instance): def test_binary_path_normalizer_factory(sqlite_instance): """Test factory creates BinaryPathDataTypeSerializer correctly.""" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="binary_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="binary_path" + ) assert isinstance(serializer, DataTypeSerializer) assert isinstance(serializer, BinaryPathDataTypeSerializer) assert serializer.data_type == "binary_path" @@ -297,7 +346,9 @@ def test_binary_path_normalizer_factory(sqlite_instance): def test_binary_path_normalizer_factory_with_value(sqlite_instance): """Test factory creates BinaryPathDataTypeSerializer with value.""" serializer = data_serializer_factory( - category="prompt-memory-entries", data_type="binary_path", value="/path/to/file.bin" + category="prompt-memory-entries", + data_type="binary_path", + value="/path/to/file.bin", ) assert isinstance(serializer, BinaryPathDataTypeSerializer) assert serializer.data_type == "binary_path" @@ -307,7 +358,9 @@ def test_binary_path_normalizer_factory_with_value(sqlite_instance): async def test_binary_path_save_data(sqlite_instance): """Test saving binary data to disk.""" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="binary_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="binary_path" + ) await serializer.save_data_async(b"\x00\x01\x02\x03") serializer_value = serializer.value assert serializer_value @@ -320,12 +373,16 @@ async def test_binary_path_save_data(sqlite_instance): async def test_binary_path_read_data(sqlite_instance): """Test reading binary data from disk.""" data = b"\x00\x11\x22\x33\x44\x55" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="binary_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="binary_path" + ) await serializer.save_data_async(data) assert await serializer.read_data_async() == data # Test reading with a new serializer initialized with the saved path read_serializer = data_serializer_factory( - category="prompt-memory-entries", data_type="binary_path", value=serializer.value + category="prompt-memory-entries", + data_type="binary_path", + value=serializer.value, ) assert await read_serializer.read_data_async() == data @@ -334,7 +391,9 @@ async def test_binary_path_save_with_custom_extension(sqlite_instance): """Test saving binary data with a custom file extension.""" custom_extension = "pdf" serializer = data_serializer_factory( - category="prompt-memory-entries", data_type="binary_path", extension=custom_extension + category="prompt-memory-entries", + data_type="binary_path", + extension=custom_extension, ) binary_data = b"PDF binary content" await serializer.save_data_async(binary_data) @@ -345,40 +404,60 @@ async def test_binary_path_save_with_custom_extension(sqlite_instance): async def test_binary_path_subdirectory(sqlite_instance): """Test that binary data is stored in the correct subdirectory.""" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="binary_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="binary_path" + ) await serializer.save_data_async(b"test data") assert "/binaries/" in serializer.value or "\\binaries\\" in serializer.value def test_get_storage_io_raises_when_results_storage_io_none(): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) serializer.value = "https://account.blob.core.windows.net/container/path/image.png" mock_memory = MagicMock() mock_memory.results_storage_io = None - with patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory): + with patch.object( + type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory + ): with pytest.raises(RuntimeError, match="results_storage_io is not configured"): serializer._get_storage_io() async def test_save_data_raises_when_results_storage_io_none(): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) mock_memory = MagicMock() mock_memory.results_storage_io = None - with patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory): + with patch.object( + type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory + ): with patch.object( - serializer, "get_data_filename_async", new_callable=AsyncMock, return_value="local/path/img.png" + serializer, + "get_data_filename_async", + new_callable=AsyncMock, + return_value="local/path/img.png", ): with pytest.raises(RuntimeError, match="Storage IO not initialized"): await serializer.save_data_async(b"\x89PNG") async def test_save_b64_image_raises_when_results_storage_io_none(): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) mock_memory = MagicMock() mock_memory.results_storage_io = None - with patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory): + with patch.object( + type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory + ): with patch.object( - serializer, "get_data_filename_async", new_callable=AsyncMock, return_value="local/path/img.png" + serializer, + "get_data_filename_async", + new_callable=AsyncMock, + return_value="local/path/img.png", ): import base64 @@ -394,8 +473,15 @@ async def test_save_formatted_audio_raises_when_results_storage_io_none(): mock_memory = MagicMock() mock_memory.results_storage_io = None azure_url = "https://account.blob.core.windows.net/container/audio/test.wav" - with patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory): - with patch.object(serializer, "get_data_filename_async", new_callable=AsyncMock, return_value=azure_url): + with patch.object( + type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory + ): + with patch.object( + serializer, + "get_data_filename_async", + new_callable=AsyncMock, + return_value=azure_url, + ): with patch("wave.open"): with patch("aiofiles.open", new_callable=MagicMock) as mock_aio: mock_file = MagicMock() @@ -403,11 +489,17 @@ async def test_save_formatted_audio_raises_when_results_storage_io_none(): mock_file.__aexit__ = AsyncMock(return_value=False) mock_file.read = AsyncMock(return_value=b"audio_bytes") mock_aio.return_value = mock_file - with pytest.raises(RuntimeError, match="results_storage_io is not initialized"): - await serializer.save_formatted_audio_async(data=b"\x00\x01\x02") + with pytest.raises( + RuntimeError, match="results_storage_io is not initialized" + ): + await serializer.save_formatted_audio_async( + data=b"\x00\x01\x02" + ) -async def test_save_formatted_audio_writes_local_wav_via_to_thread(sqlite_instance, tmp_path): +async def test_save_formatted_audio_writes_local_wav_via_to_thread( + sqlite_instance, tmp_path +): """save_formatted_audio (local-disk path) should produce a readable WAV via _write_wav_sync.""" import wave @@ -415,7 +507,12 @@ async def test_save_formatted_audio_writes_local_wav_via_to_thread(sqlite_instan serializer = factory(category="prompt-memory-entries", data_type="audio_path") output_path = tmp_path / "out.wav" - with patch.object(serializer, "get_data_filename_async", new_callable=AsyncMock, return_value=str(output_path)): + with patch.object( + serializer, + "get_data_filename_async", + new_callable=AsyncMock, + return_value=str(output_path), + ): pcm = b"\x01\x00\x02\x00\x03\x00\x04\x00" await serializer.save_formatted_audio_async( data=pcm, @@ -457,7 +554,9 @@ def test_write_wav_sync_produces_readable_wav(tmp_path): assert wav_file.readframes(wav_file.getnframes()) == pcm -async def test_save_formatted_audio_writes_azure_wav_via_storage_io(sqlite_instance, tmp_path): +async def test_save_formatted_audio_writes_azure_wav_via_storage_io( + sqlite_instance, tmp_path +): """Azure-storage branch should round-trip data through _write_wav_sync to storage_io.write_file.""" import wave @@ -479,8 +578,15 @@ async def _capture_write(file_path, data): azure_url = "https://account.blob.core.windows.net/container/audio/test.wav" pcm = b"\xaa\xbb\xcc\xdd\xee\xff\x00\x11" - with patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory): - with patch.object(serializer, "get_data_filename_async", new_callable=AsyncMock, return_value=azure_url): + with patch.object( + type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory + ): + with patch.object( + serializer, + "get_data_filename_async", + new_callable=AsyncMock, + return_value=azure_url, + ): # Redirect so the temp_audio.wav write lands in tmp_path with patch.object(common_path, "DB_DATA_PATH", str(tmp_path)): from pyrit.memory.storage import serializers as dts_module @@ -507,25 +613,36 @@ async def _capture_write(file_path, data): async def test_get_data_filename_raises_when_results_storage_io_none(): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) serializer._file_path = None mock_memory = MagicMock() mock_memory.results_storage_io = None mock_memory.results_path = "/local/results" - with patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory): + with patch.object( + type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory + ): with pytest.raises(RuntimeError, match="results_storage_io is not initialized"): await serializer.get_data_filename_async() async def test_get_data_filename_uses_db_data_path_when_results_path_falsy(): - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) serializer._file_path = None mock_memory = MagicMock() mock_memory.results_path = None mock_storage_io = AsyncMock() mock_memory.results_storage_io = mock_storage_io with ( - patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory), + patch.object( + type(serializer), + "_memory", + new_callable=PropertyMock, + return_value=mock_memory, + ), patch("pyrit.common.path.DB_DATA_PATH", "/fallback/db_data"), ): result = await serializer.get_data_filename_async(file_name="test_file") @@ -545,8 +662,18 @@ async def test_save_formatted_audio_azure_storage_unlinks_local_temp(tmp_path): azure_url = "https://account.blob.core.windows.net/container/audio/test.wav" with ( - patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory), - patch.object(serializer, "get_data_filename_async", new_callable=AsyncMock, return_value=azure_url), + patch.object( + type(serializer), + "_memory", + new_callable=PropertyMock, + return_value=mock_memory, + ), + patch.object( + serializer, + "get_data_filename_async", + new_callable=AsyncMock, + return_value=azure_url, + ), patch("pyrit.memory.storage.serializers.DB_DATA_PATH", tmp_path), ): await serializer.save_formatted_audio_async(data=b"\x00\x01\x02\x03") @@ -575,12 +702,19 @@ async def test_set_seed_sha256_async_sets_text_hash(sqlite_instance): await set_seed_sha256_async(seed) - assert seed.value_sha256 == "948edbe7ede5aa7423476ae29dcd7d61e7711a071aea0d83698377effa896525" + assert ( + seed.value_sha256 + == "948edbe7ede5aa7423476ae29dcd7d61e7711a071aea0d83698377effa896525" + ) -async def test_save_formatted_audio_async_cleans_up_temp_file_on_azure_upload_failure(tmp_path): +async def test_save_formatted_audio_async_cleans_up_temp_file_on_azure_upload_failure( + tmp_path, +): """Regression test: temp file must be deleted even when Azure upload fails.""" - serializer = data_serializer_factory(category="prompt-memory-entries", data_type="audio_path") + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="audio_path" + ) mock_memory = MagicMock() mock_storage_io = AsyncMock() @@ -592,8 +726,15 @@ async def test_save_formatted_audio_async_cleans_up_temp_file_on_azure_upload_fa # Record existing wav files BEFORE test runs existing_wav_files = set(tmp_path.glob("*.wav")) - with patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory): - with patch.object(serializer, "get_data_filename_async", new_callable=AsyncMock, return_value=azure_url): + with patch.object( + type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory + ): + with patch.object( + serializer, + "get_data_filename_async", + new_callable=AsyncMock, + return_value=azure_url, + ): with patch("pyrit.memory.storage.serializers.DB_DATA_PATH", tmp_path): with pytest.raises(RuntimeError, match="Azure upload failed"): await serializer.save_formatted_audio_async(data=b"\x00\x01\x02") @@ -601,3 +742,82 @@ async def test_save_formatted_audio_async_cleans_up_temp_file_on_azure_upload_fa # Check no NEW wav files leaked after test leaked_files = set(tmp_path.glob("*.wav")) - existing_wav_files assert leaked_files == set(), f"Temp files leaked: {leaked_files}" + + +async def test_save_data_async_passes_content_type_to_storage(sqlite_instance): + """save_data_async should infer and pass content_type to storage layer.""" + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path" + ) + mock_storage_io = AsyncMock() + mock_memory = MagicMock() + mock_memory.results_storage_io = mock_storage_io + + with patch.object( + type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory + ): + await serializer.save_data_async(b"\x89PNG", output_filename="test.png") + + mock_storage_io.write_file_async.assert_awaited_once() + call_kwargs = mock_storage_io.write_file_async.call_args + assert call_kwargs[1].get("content_type") == "image/png" or ( + len(call_kwargs[0]) > 2 and call_kwargs[0][2] == "image/png" + ) + + +async def test_save_b64_image_async_passes_content_type_to_storage(sqlite_instance): + """save_b64_image_async should infer and pass content_type to storage layer.""" + import base64 + + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="image_path", extension="jpg" + ) + mock_storage_io = AsyncMock() + mock_memory = MagicMock() + mock_memory.results_storage_io = mock_storage_io + + with patch.object( + type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory + ): + b64_data = base64.b64encode(b"\xff\xd8\xff\xe0").decode() + await serializer.save_b64_image_async(b64_data, output_filename="photo.jpg") + + mock_storage_io.write_file_async.assert_awaited_once() + call_kwargs = mock_storage_io.write_file_async.call_args + assert call_kwargs[1].get("content_type") == "image/jpeg" or ( + len(call_kwargs[0]) > 2 and call_kwargs[0][2] == "image/jpeg" + ) + + +async def test_save_formatted_audio_async_passes_content_type_to_storage(tmp_path): + """save_formatted_audio_async should infer and pass content_type to Azure storage layer.""" + serializer = data_serializer_factory( + category="prompt-memory-entries", data_type="audio_path" + ) + mock_storage_io = AsyncMock() + mock_memory = MagicMock() + mock_memory.results_storage_io = mock_storage_io + azure_url = "https://account.blob.core.windows.net/container/audio/test.wav" + + with ( + patch.object( + type(serializer), + "_memory", + new_callable=PropertyMock, + return_value=mock_memory, + ), + patch.object( + serializer, + "get_data_filename_async", + new_callable=AsyncMock, + return_value=azure_url, + ), + patch("pyrit.memory.storage.serializers.DB_DATA_PATH", tmp_path), + ): + await serializer.save_formatted_audio_async(data=b"\x00\x01\x02\x03") + + mock_storage_io.write_file_async.assert_awaited_once() + call_kwargs = mock_storage_io.write_file_async.call_args + assert call_kwargs[1].get("content_type") == "audio/wav" or ( + len(call_kwargs[0]) > 2 and call_kwargs[0][2] == "audio/wav" + ) diff --git a/tests/unit/memory/storage/test_storage.py b/tests/unit/memory/storage/test_storage.py index 6b9b7d10e6..d9688fe5a4 100644 --- a/tests/unit/memory/storage/test_storage.py +++ b/tests/unit/memory/storage/test_storage.py @@ -515,3 +515,86 @@ def test_supported_content_type_has_expected_values(): assert SupportedContentType.ZIP.value == "application/zip" assert SupportedContentType.TAR.value == "application/x-tar" assert SupportedContentType.GZIP.value == "application/gzip" + + +async def test_azure_blob_storage_io_write_file_infers_content_type_from_extension(): + """write_file_async should infer content_type from file extension when not provided.""" + container_url = "https://youraccount.blob.core.windows.net/yourcontainer" + azure_blob_storage_io = AzureBlobStorageIO( + container_url=container_url, blob_content_type=SupportedContentType.PLAIN_TEXT + ) + + mock_container_client = AsyncMock() + + with patch.object( + azure_blob_storage_io, "_create_container_client_async", return_value=None + ): + azure_blob_storage_io._client_async = mock_container_client + azure_blob_storage_io._upload_blob_async = AsyncMock() + + data_to_write = b"test" + path = "https://youraccount.blob.core.windows.net/yourcontainer/testfile.html" + + await azure_blob_storage_io.write_file_async(path, data_to_write) + + azure_blob_storage_io._upload_blob_async.assert_awaited_with( + file_name="testfile.html", + data=data_to_write, + content_type="text/html", + ) + + +async def test_azure_blob_storage_io_write_file_with_explicit_content_type(): + """write_file_async should use explicit content_type when provided.""" + container_url = "https://youraccount.blob.core.windows.net/yourcontainer" + azure_blob_storage_io = AzureBlobStorageIO( + container_url=container_url, blob_content_type=SupportedContentType.PLAIN_TEXT + ) + + mock_container_client = AsyncMock() + + with patch.object( + azure_blob_storage_io, "_create_container_client_async", return_value=None + ): + azure_blob_storage_io._client_async = mock_container_client + azure_blob_storage_io._upload_blob_async = AsyncMock() + + data_to_write = b"test data" + path = "https://youraccount.blob.core.windows.net/yourcontainer/testfile.txt" + + await azure_blob_storage_io.write_file_async( + path, data_to_write, content_type="application/json" + ) + + azure_blob_storage_io._upload_blob_async.assert_awaited_with( + file_name="testfile.txt", + data=data_to_write, + content_type="application/json", + ) + + +async def test_azure_blob_storage_io_write_file_infers_image_content_type(): + """write_file_async should infer image content types from extension.""" + container_url = "https://youraccount.blob.core.windows.net/yourcontainer" + azure_blob_storage_io = AzureBlobStorageIO( + container_url=container_url, blob_content_type=SupportedContentType.PLAIN_TEXT + ) + + mock_container_client = AsyncMock() + + with patch.object( + azure_blob_storage_io, "_create_container_client_async", return_value=None + ): + azure_blob_storage_io._client_async = mock_container_client + azure_blob_storage_io._upload_blob_async = AsyncMock() + + data_to_write = b"\x89PNG\r\n\x1a\n" + path = "https://youraccount.blob.core.windows.net/yourcontainer/image.png" + + await azure_blob_storage_io.write_file_async(path, data_to_write) + + azure_blob_storage_io._upload_blob_async.assert_awaited_with( + file_name="image.png", + data=data_to_write, + content_type="image/png", + ) From db7eb8f215cf57bd2a5c08cb1407b8184a2d6da2 Mon Sep 17 00:00:00 2001 From: lxcxjxhx Date: Mon, 20 Jul 2026 23:19:21 +0800 Subject: [PATCH 3/3] fix: update notebook import to use consolidated SupportedContentType The Jupyter notebook still imported SupportedContentType from pyrit.prompt_target.azure_blob_storage_target, which no longer defines it (consolidated into pyrit.memory.storage.storage). Co-Authored-By: Claude --- doc/code/executor/5_workflow.ipynb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/code/executor/5_workflow.ipynb b/doc/code/executor/5_workflow.ipynb index f97d5e8b04..3d40a87f92 100644 --- a/doc/code/executor/5_workflow.ipynb +++ b/doc/code/executor/5_workflow.ipynb @@ -693,7 +693,7 @@ "from pyrit.executor.workflow import XPIAWorkflow\n", "from pyrit.prompt_normalizer import ConverterConfiguration\n", "from pyrit.prompt_target import AzureBlobStorageTarget\n", - "from pyrit.prompt_target.azure_blob_storage_target import SupportedContentType\n", + "from pyrit.memory.storage.storage import SupportedContentType\n", "from pyrit.score import SubStringScorer\n", "\n", "logging.basicConfig(level=logging.DEBUG)\n", @@ -912,4 +912,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file