diff --git a/pyrit/memory/storage/serializers.py b/pyrit/memory/storage/serializers.py index f779cc693e..fbf7d7647e 100644 --- a/pyrit/memory/storage/serializers.py +++ b/pyrit/memory/storage/serializers.py @@ -209,6 +209,7 @@ async def save_formatted_audio_async( Raises: RuntimeError: If storage IO is not initialized. """ + self.file_extension = "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 @@ -346,6 +347,9 @@ async def get_data_filename_async(self, file_name: str | None = None) -> Path | results_path = str(DB_DATA_PATH) file_name = file_name if file_name else str(ticks) + file_suffix = Path(file_name).suffix + if file_suffix: + file_name = file_name[: -len(file_suffix)] if self._is_azure_storage_url(results_path): full_data_directory_path = results_path + self.data_sub_directory diff --git a/pyrit/memory/storage/storage.py b/pyrit/memory/storage/storage.py index 1cd5bac2a8..300b0653ba 100644 --- a/pyrit/memory/storage/storage.py +++ b/pyrit/memory/storage/storage.py @@ -7,7 +7,7 @@ from abc import ABC, abstractmethod from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar from urllib.parse import urlparse import aiofiles @@ -157,6 +157,41 @@ class AzureBlobStorageIO(StorageIO): Implementation of StorageIO for Azure Blob Storage. """ + _EXTENSION_TO_CONTENT_TYPE: ClassVar[dict[str, str]] = { + ".txt": "text/plain", + ".html": "text/html", + ".htm": "text/html", + ".csv": "text/csv", + ".md": "text/markdown", + ".json": "application/json", + ".xml": "application/xml", + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".svg": "image/svg+xml", + ".bmp": "image/bmp", + ".wav": "audio/wav", + ".mp3": "audio/mpeg", + ".ogg": "audio/ogg", + ".flac": "audio/flac", + ".m4a": "audio/mp4", + ".mp4": "video/mp4", + ".webm": "video/webm", + ".ogv": "video/ogg", + ".avi": "video/x-msvideo", + ".pdf": "application/pdf", + ".doc": "application/msword", + ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ".xls": "application/vnd.ms-excel", + ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ".ppt": "application/vnd.ms-powerpoint", + ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation", + ".rtf": "application/rtf", + ".zip": "application/zip", + } + def __init__( self, *, @@ -382,8 +417,9 @@ async def write_file_async(self, path: Path | str, data: bytes) -> None: if not self._client_async: self._client_async = await self._create_container_client_async() blob_name = self._resolve_blob_name(path) + content_type = self._EXTENSION_TO_CONTENT_TYPE.get(Path(blob_name).suffix.lower(), self._blob_content_type) 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=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_serializers.py b/tests/unit/memory/storage/test_serializers.py index b323a4e8b5..6ee4799a82 100644 --- a/tests/unit/memory/storage/test_serializers.py +++ b/tests/unit/memory/storage/test_serializers.py @@ -5,6 +5,7 @@ import os import re import tempfile +from pathlib import Path from typing import get_args from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch @@ -13,10 +14,12 @@ from pyrit.memory.storage import ( AllowedCategories, + AzureBlobStorageIO, BinaryPathDataTypeSerializer, DataTypeSerializer, ErrorDataTypeSerializer, ImagePathDataTypeSerializer, + StorageIO, TextDataTypeSerializer, data_serializer_factory, set_message_piece_sha256_async, @@ -25,6 +28,35 @@ from pyrit.models import MessagePiece, SeedPrompt +class LegacyStorageIO(StorageIO): + """ + Test double representing an existing third-party ``StorageIO`` implementation. + + Its ``write_file_async(path, data)`` method intentionally retains the original + two-argument contract. Tests using this class ensure serializers do not pass a + new content-type keyword argument that would break pre-existing custom storage + backends when Azure Blob Storage adds MIME metadata internally. + """ + + def __init__(self) -> None: + self.writes: list[tuple[Path | str, bytes]] = [] + + async def read_file_async(self, path: Path | str) -> bytes: + return b"" + + async def write_file_async(self, path: Path | str, data: bytes) -> None: + self.writes.append((path, data)) + + async def path_exists_async(self, path: Path | str) -> bool: + return False + + async def is_file_async(self, path: Path | str) -> bool: + return False + + async def create_directory_if_not_exists_async(self, path: Path | str) -> None: + return None + + def test_allowed_categories(): entries = get_args(AllowedCategories) assert len(entries) == 2 @@ -285,6 +317,29 @@ async def test_get_data_filename(sqlite_instance): assert not os.path.exists(filename) # File should not exist yet +async def test_get_data_filename_does_not_duplicate_extension(sqlite_instance): + serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + + filename = await serializer.get_data_filename_async(file_name="photo.png") + + assert Path(filename).name == "photo.png" + + +async def test_save_data_supports_legacy_storage_io_write_signature(): + storage = LegacyStorageIO() + mock_memory = MagicMock() + mock_memory.results_path = "https://account.blob.core.windows.net/container/results" + mock_memory.results_storage_io = storage + serializer = data_serializer_factory(category="prompt-memory-entries", data_type="image_path") + + with patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory): + await serializer.save_data_async(b"\x89PNG", output_filename="photo.png") + + assert storage.writes == [ + ("https://account.blob.core.windows.net/container/results/prompt-memory-entries/images/photo.png", b"\x89PNG") + ] + + 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") @@ -433,6 +488,36 @@ async def test_save_formatted_audio_writes_local_wav_via_to_thread(sqlite_instan assert wav_file.readframes(wav_file.getnframes()) == pcm +async def test_save_formatted_audio_uses_wav_filename_content_and_metadata(tmp_path): + import io + import wave + + storage = AzureBlobStorageIO(container_url="https://account.blob.core.windows.net/container") + mock_container_client = AsyncMock() + storage._client_async = mock_container_client + mock_memory = MagicMock() + mock_memory.results_path = "https://account.blob.core.windows.net/container/results" + mock_memory.results_storage_io = storage + serializer = data_serializer_factory(category="prompt-memory-entries", data_type="audio_path") + pcm = b"\x01\x00\x02\x00\x03\x00\x04\x00" + + with ( + patch.object(type(serializer), "_memory", new_callable=PropertyMock, return_value=mock_memory), + patch("pyrit.memory.storage.serializers.DB_DATA_PATH", tmp_path), + ): + await serializer.save_formatted_audio_async(data=pcm, output_filename="recording.mp3") + + upload_kwargs = mock_container_client.upload_blob.await_args.kwargs + assert upload_kwargs["name"] == "results/prompt-memory-entries/audio/recording.wav" + assert upload_kwargs["content_settings"].content_type == "audio/wav" + assert serializer.value.endswith("/audio/recording.wav") + with wave.open(io.BytesIO(upload_kwargs["data"]), "rb") as wav_file: + assert wav_file.getnchannels() == 1 + assert wav_file.getsampwidth() == 2 + assert wav_file.getframerate() == 16000 + assert wav_file.readframes(wav_file.getnframes()) == pcm + + def test_write_wav_sync_produces_readable_wav(tmp_path): """_write_wav_sync should produce a WAV file readable by wave.open with the same metadata and frames.""" import wave diff --git a/tests/unit/memory/storage/test_storage.py b/tests/unit/memory/storage/test_storage.py index d7c9720885..f8bf53736e 100644 --- a/tests/unit/memory/storage/test_storage.py +++ b/tests/unit/memory/storage/test_storage.py @@ -162,6 +162,43 @@ async def test_azure_blob_storage_io_write_file_with_relative_path(): ) +@pytest.mark.parametrize( + ("path", "expected_content_type"), + [ + ("notes.HTML", "text/html"), + ("photo.JPEG", "image/jpeg"), + ("recording.wav", "audio/wav"), + ("movie.mp4", "video/mp4"), + ("report.pdf", "application/pdf"), + ], +) +async def test_azure_blob_storage_io_write_file_sets_content_type_from_extension(path, expected_content_type): + storage = AzureBlobStorageIO(container_url="https://account.blob.core.windows.net/container") + mock_container_client = AsyncMock() + storage._client_async = mock_container_client + + await storage.write_file_async(path, b"data") + + upload_kwargs = mock_container_client.upload_blob.await_args.kwargs + assert upload_kwargs["name"] == path + assert upload_kwargs["content_settings"].content_type == expected_content_type + + +@pytest.mark.parametrize("path", ["data.unknown", "README"]) +async def test_azure_blob_storage_io_write_file_uses_configured_fallback_content_type(path): + storage = AzureBlobStorageIO( + container_url="https://account.blob.core.windows.net/container", + blob_content_type=SupportedContentType.PLAIN_TEXT, + ) + mock_container_client = AsyncMock() + storage._client_async = mock_container_client + + await storage.write_file_async(path, b"data") + + upload_kwargs = mock_container_client.upload_blob.await_args.kwargs + assert upload_kwargs["content_settings"].content_type == SupportedContentType.PLAIN_TEXT.value + + 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"