diff --git a/pyrit/converter/url_converter.py b/pyrit/converter/url_converter.py index 6895a4b126..5446cb2c2e 100644 --- a/pyrit/converter/url_converter.py +++ b/pyrit/converter/url_converter.py @@ -4,7 +4,7 @@ import urllib.parse from pyrit.converter.converter import Converter, ConverterResult -from pyrit.models import PromptDataType +from pyrit.models import ComponentIdentifier, PromptDataType class UrlConverter(Converter): @@ -15,6 +15,26 @@ class UrlConverter(Converter): SUPPORTED_INPUT_TYPES = ("text",) SUPPORTED_OUTPUT_TYPES = ("text",) + def __init__(self, *, safe: str = "/") -> None: + """ + Initialize the converter with safe characters that should not be converted. + + Args: + safe (str): Characters that should not be URL-encoded. Defaults to "/", + matching urllib.parse.quote. Set to "" to encode forward slashes. + """ + super().__init__() + self._safe = safe + + def _build_identifier(self) -> ComponentIdentifier: + """ + Build the converter identifier with encoding parameters. + + Returns: + ComponentIdentifier: The identifier for this converter. + """ + return self._create_identifier(params={"safe": self._safe}) + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: """ Convert the given prompt into a URL-encoded string. @@ -32,4 +52,4 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text if not self.input_supported(input_type): raise ValueError("Input type not supported") - return ConverterResult(output_text=urllib.parse.quote(prompt), output_type="text") + return ConverterResult(output_text=urllib.parse.quote(prompt, safe=self._safe), output_type="text") diff --git a/tests/unit/converter/test_url_converter.py b/tests/unit/converter/test_url_converter.py index 18b0d2cb64..df712eaddb 100644 --- a/tests/unit/converter/test_url_converter.py +++ b/tests/unit/converter/test_url_converter.py @@ -42,3 +42,39 @@ async def test_url_converter_input_not_supported(): converter = UrlConverter() with pytest.raises(ValueError): await converter.convert_async(prompt="hello", input_type="image_path") + + +async def test_url_converter_default_preserves_forward_slash_async() -> None: + converter = UrlConverter() + result = await converter.convert_async(prompt="https://example.com/a b") + assert result.output_text == "https%3A//example.com/a%20b" + + +@pytest.mark.parametrize( + ("safe", "prompt", "expected"), + [ + ("", "../../etc/passwd", "..%2F..%2Fetc%2Fpasswd"), + ("", "https://example.com/test", "https%3A%2F%2Fexample.com%2Ftest"), + ("/:", "https://example.com/test", "https://example.com/test"), + ("?=", "?query=1&foo=bar", "?query=1%26foo=bar"), + ], +) +async def test_url_converter_custom_safe_async(*, safe: str, prompt: str, expected: str) -> None: + converter = UrlConverter(safe=safe) + result = await converter.convert_async(prompt=prompt) + assert result.output_text == expected + + +@pytest.mark.parametrize("safe", ["/", "", "/:"]) +def test_url_converter_identifier_and_reconstruction(safe: str) -> None: + converter = UrlConverter(safe=safe) + assert converter.get_identifier().params["safe"] == safe + params = converter.get_reconstruction_parameters() + assert params == {"safe": safe} + assert UrlConverter(**params).get_identifier() == converter.get_identifier() + + +def test_url_converter_default_reconstruction() -> None: + converter = UrlConverter() + assert converter.get_identifier().params["safe"] == "/" + assert converter.get_reconstruction_parameters() == {"safe": "/"}