diff --git a/pyrit/converter/url_converter.py b/pyrit/converter/url_converter.py index 6895a4b126..c507b0b458 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,37 @@ 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 "" (encodes + all reserved characters, including forward slashes, ensuring thorough URL + encoding for path traversal and SSRF payload testing). + + Raises: + TypeError: If ``safe`` is not a string. + """ + super().__init__() + if not isinstance(safe, str): + raise TypeError(f"safe must be a str, got {type(safe).__name__}") + self._safe = safe + + @property + def safe(self) -> str: + """The characters preserved during URL encoding.""" + return self._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 +63,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..11ad912089 100644 --- a/tests/unit/converter/test_url_converter.py +++ b/tests/unit/converter/test_url_converter.py @@ -42,3 +42,53 @@ 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_encodes_forward_slash(): + converter = UrlConverter() + result = await converter.convert_async(prompt="../../etc/passwd", input_type="text") + assert isinstance(result, ConverterResult) + assert "/" not in result.output_text + assert result.output_text == "..%2F..%2Fetc%2Fpasswd" + assert result.output_type == "text" + + ssrf_result = await converter.convert_async(prompt="https://attacker.com/steal", input_type="text") + assert "/" not in ssrf_result.output_text + assert ssrf_result.output_text == "https%3A%2F%2Fattacker.com%2Fsteal" + + +async def test_url_converter_custom_safe(): + converter_slash = UrlConverter(safe="/") + result_slash = await converter_slash.convert_async(prompt="../../etc/passwd", input_type="text") + assert result_slash.output_text == "../../etc/passwd" + + converter_url = UrlConverter(safe="/:") + result_url = await converter_url.convert_async(prompt="https://example.com/test", input_type="text") + assert result_url.output_text == "https://example.com/test" + + converter_query = UrlConverter(safe="?=") + result_query = await converter_query.convert_async(prompt="?query=1&foo=bar", input_type="text") + assert result_query.output_text == "?query=1%26foo=bar" + + +def test_url_converter_property_and_identifier(): + default_converter = UrlConverter() + assert default_converter.safe == "" + assert default_converter.get_identifier().params["safe"] == "" + + custom_converter = UrlConverter(safe="/") + assert custom_converter.safe == "/" + assert custom_converter.get_identifier().params["safe"] == "/" + + +def test_url_converter_invalid_safe_type(): + with pytest.raises(TypeError, match="safe must be a str"): + UrlConverter(safe=123) # type: ignore[arg-type] + + +def test_url_converter_reconstruction(): + default_converter = UrlConverter() + assert default_converter.get_reconstruction_parameters() == {"safe": ""} + + custom_converter = UrlConverter(safe="/") + assert custom_converter.get_reconstruction_parameters() == {"safe": "/"}