From 531e33aecf2127da6ddc100dfcaa321ca947c65d Mon Sep 17 00:00:00 2001 From: lxcxjxhx Date: Sun, 19 Jul 2026 13:21:40 +0800 Subject: [PATCH 1/6] 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/6] 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/6] 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 From 62210746b92b70563f157b1c406e8d929f8b48e4 Mon Sep 17 00:00:00 2001 From: lxcxjxhx Date: Tue, 21 Jul 2026 17:26:27 +0800 Subject: [PATCH 4/6] fix: address maintainer feedback on content type handling - Replace mimetypes.guess_type with deterministic extension mapping to avoid platform-dependent behavior - Add text-only content type validation to AzureBlobStorageTarget to prevent binary uploads - Fix audio serializer to force .wav extension and use audio/wav content type - Maintain backward compatibility for StorageIO.write_file_async signature --- pyrit/memory/storage/serializers.py | 8 ++- pyrit/memory/storage/storage.py | 59 +++++++++++++++++-- .../azure_blob_storage_target.py | 22 +++++++ 3 files changed, 80 insertions(+), 9 deletions(-) diff --git a/pyrit/memory/storage/serializers.py b/pyrit/memory/storage/serializers.py index 4d46c78c92..4bb457e43c 100644 --- a/pyrit/memory/storage/serializers.py +++ b/pyrit/memory/storage/serializers.py @@ -239,6 +239,9 @@ async def save_formatted_audio_async( Raises: RuntimeError: If storage IO is not initialized. """ + # Force .wav extension since we're writing WAV format + if output_filename and not output_filename.lower().endswith(".wav"): + output_filename = Path(output_filename).stem + ".wav" file_path = await self.get_data_filename_async(file_name=output_filename) # save audio file locally first if in AzureStorageBlob so we can use wave.open to set audio parameters @@ -262,10 +265,9 @@ async def save_formatted_audio_async( 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)) + # Always use audio/wav since we're writing WAV format await self._memory.results_storage_io.write_file_async( - file_path, audio_data, content_type=content_type + file_path, audio_data, content_type="audio/wav" ) finally: local_temp_path.unlink(missing_ok=True) diff --git a/pyrit/memory/storage/storage.py b/pyrit/memory/storage/storage.py index 443a966d6b..e326c9c63b 100644 --- a/pyrit/memory/storage/storage.py +++ b/pyrit/memory/storage/storage.py @@ -61,6 +61,49 @@ class SupportedContentType(Enum): GZIP = "application/gzip" +# Deterministic mapping from file extensions to SupportedContentType. +# This avoids platform-dependent behavior from mimetypes.guess_type. +_EXTENSION_TO_CONTENT_TYPE: dict[str, SupportedContentType] = { + # Text types + ".txt": SupportedContentType.PLAIN_TEXT, + ".text": SupportedContentType.PLAIN_TEXT, + ".html": SupportedContentType.HTML, + ".htm": SupportedContentType.HTML, + ".json": SupportedContentType.JSON, + ".xml": SupportedContentType.XML, + ".csv": SupportedContentType.CSV, + ".md": SupportedContentType.MARKDOWN, + ".markdown": SupportedContentType.MARKDOWN, + # Image types + ".png": SupportedContentType.PNG, + ".jpg": SupportedContentType.JPEG, + ".jpeg": SupportedContentType.JPEG, + ".gif": SupportedContentType.GIF, + ".webp": SupportedContentType.WEBP, + ".svg": SupportedContentType.SVG, + ".bmp": SupportedContentType.BMP, + # Audio types + ".wav": SupportedContentType.WAV, + ".wave": SupportedContentType.WAV, + ".mp3": SupportedContentType.MP3, + ".ogg": SupportedContentType.OGG, + ".flac": SupportedContentType.FLAC, + ".m4a": SupportedContentType.M4A, + # Video types + ".mp4": SupportedContentType.MP4, + ".m4v": SupportedContentType.MP4, + ".webm": SupportedContentType.WEBM, + ".ogv": SupportedContentType.OGG_VIDEO, + ".avi": SupportedContentType.AVI, + # Document types + ".pdf": SupportedContentType.PDF, + ".zip": SupportedContentType.ZIP, + ".tar": SupportedContentType.TAR, + ".gz": SupportedContentType.GZIP, + ".gzip": SupportedContentType.GZIP, +} + + class StorageIO(ABC): """ Abstract interface for storage systems (local disk, Azure Storage Account, etc.). @@ -422,8 +465,8 @@ async def write_file_async( 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. + it will be inferred from the file extension using a deterministic mapping, + 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() @@ -431,10 +474,14 @@ async def write_file_async( # 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 + # Use deterministic mapping instead of platform-dependent mimetypes.guess_type + ext = Path(blob_name).suffix.lower() + inferred_content_type = _EXTENSION_TO_CONTENT_TYPE.get(ext) + content_type = ( + inferred_content_type.value + if inferred_content_type + else self._blob_content_type + ) try: await self._upload_blob_async( diff --git a/pyrit/prompt_target/azure_blob_storage_target.py b/pyrit/prompt_target/azure_blob_storage_target.py index 6311b1dde2..b89d8a785b 100644 --- a/pyrit/prompt_target/azure_blob_storage_target.py +++ b/pyrit/prompt_target/azure_blob_storage_target.py @@ -64,6 +64,18 @@ class AzureBlobStorageTarget(PromptTarget): ) ) + # Text-only content types supported by this target (which handles text prompts) + _TEXT_CONTENT_TYPES = frozenset( + { + SupportedContentType.PLAIN_TEXT, + SupportedContentType.HTML, + SupportedContentType.JSON, + SupportedContentType.XML, + SupportedContentType.CSV, + SupportedContentType.MARKDOWN, + } + ) + def __init__( self, *, @@ -82,11 +94,21 @@ def __init__( sas_token (str, Optional): The SAS token for authentication. Defaults to the AZURE_STORAGE_ACCOUNT_SAS_TOKEN environment variable. blob_content_type (SupportedContentType): The content type for blobs. + Must be a text-based type (PLAIN_TEXT, HTML, JSON, XML, CSV, MARKDOWN). Defaults to PLAIN_TEXT. max_requests_per_minute (int, Optional): Maximum number of requests per minute. custom_configuration (TargetConfiguration, Optional): Override the default configuration for this target instance. Defaults to None. + + Raises: + ValueError: If blob_content_type is not a text-based type. """ + if blob_content_type not in self._TEXT_CONTENT_TYPES: + raise ValueError( + f"AzureBlobStorageTarget only supports text-based content types. " + f"Got {blob_content_type.value}. Supported types: " + f"{', '.join(ct.value for ct in self._TEXT_CONTENT_TYPES)}" + ) self._blob_content_type: str = blob_content_type.value self._container_url: str = default_values.get_required_value( From be9ed6671416eb8e27112e523c32b7d269df62e8 Mon Sep 17 00:00:00 2001 From: lxcxjxhx Date: Wed, 22 Jul 2026 19:19:37 +0800 Subject: [PATCH 5/6] fix: address review feedback on MIME type integration - Replace mimetypes.guess_type with deterministic mapping - Preserve backward compatibility of write_file_async - Restrict AzureBlobStorageTarget to textual types - Fix audio serializer WAV/MP3 mismatch - Wire SupportedContentType into upload pipeline --- pyrit/memory/storage/serializers.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pyrit/memory/storage/serializers.py b/pyrit/memory/storage/serializers.py index 4bb457e43c..cca6f3e334 100644 --- a/pyrit/memory/storage/serializers.py +++ b/pyrit/memory/storage/serializers.py @@ -10,7 +10,6 @@ import tempfile import time import wave -from mimetypes import guess_type from pathlib import Path from typing import TYPE_CHECKING, Literal, get_args from urllib.parse import urlparse @@ -18,7 +17,7 @@ import aiofiles from pyrit.common.path import DB_DATA_PATH -from pyrit.memory.storage.storage import DiskStorageIO, StorageIO +from pyrit.memory.storage.storage import DiskStorageIO, StorageIO, _EXTENSION_TO_CONTENT_TYPE if TYPE_CHECKING: from pyrit.memory import MemoryInterface @@ -427,7 +426,7 @@ def get_extension(file_path: str) -> str | None: @staticmethod def get_mime_type(file_path: str) -> str | None: """ - Get the MIME type of the file path. + Get the MIME type of the file path using deterministic mapping. Args: file_path (str): Input file path. @@ -436,8 +435,9 @@ def get_mime_type(file_path: str) -> str | None: str | None: MIME type if detectable; otherwise None. """ - mime_type, _ = guess_type(file_path) - return mime_type + ext = Path(file_path).suffix.lower() + content_type = _EXTENSION_TO_CONTENT_TYPE.get(ext) + return content_type.value if content_type else None def _is_azure_storage_url(self, path: str) -> bool: """ From 33fca0dc2561b6a72b4625aa88185d83aee5606a Mon Sep 17 00:00:00 2001 From: Hyacinth-of-Security <144213008+lxcxjxhx@users.noreply.github.com> Date: Wed, 22 Jul 2026 21:30:34 +0800 Subject: [PATCH 6/6] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- pyrit/prompt_target/azure_blob_storage_target.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyrit/prompt_target/azure_blob_storage_target.py b/pyrit/prompt_target/azure_blob_storage_target.py index b89d8a785b..8f50344913 100644 --- a/pyrit/prompt_target/azure_blob_storage_target.py +++ b/pyrit/prompt_target/azure_blob_storage_target.py @@ -104,10 +104,10 @@ def __init__( ValueError: If blob_content_type is not a text-based type. """ if blob_content_type not in self._TEXT_CONTENT_TYPES: + supported_types = ", ".join(sorted(ct.value for ct in self._TEXT_CONTENT_TYPES)) raise ValueError( - f"AzureBlobStorageTarget only supports text-based content types. " - f"Got {blob_content_type.value}. Supported types: " - f"{', '.join(ct.value for ct in self._TEXT_CONTENT_TYPES)}" + "AzureBlobStorageTarget only supports text-based content types. " + f"Got {blob_content_type.value}. Supported types: {supported_types}" ) self._blob_content_type: str = blob_content_type.value