FIX: apply font_color to PDFs PDFConverter generates

pyrit/converter/pdf_converter.py

@@ -18,6 +18,11 @@
 from pyrit.memory import DataTypeSerializer, data_serializer_factory
 from pyrit.models import ComponentIdentifier, PromptDataType, SeedPrompt
 
+# Generated PDFs are meant to be read, so they default to black text. Injected text defaults to white so it
+# stays hidden on a typical page.
+_GENERATED_TEXT_COLOR: tuple[int, int, int] = (0, 0, 0)
+_INJECTED_TEXT_COLOR: tuple[int, int, int] = (255, 255, 255)
+
 
 class PDFConverter(Converter):
     """
@@ -43,7 +48,7 @@ def __init__(
         prompt_template: SeedPrompt | None = None,
         font_type: str = "Helvetica",
         font_size: int = 12,
-        font_color: tuple[int, int, int] = (255, 255, 255),
+        font_color: tuple[int, int, int] | None = None,
         page_width: int = 210,
         page_height: int = 297,
         column_width: int = 0,
@@ -58,7 +63,9 @@ def __init__(
             prompt_template (SeedPrompt | None, optional): A ``SeedPrompt`` object representing a template.
             font_type (str): Font type for the PDF. Defaults to "Helvetica".
             font_size (int): Font size for the PDF. Defaults to 12.
-            font_color (tuple): Font color for the PDF in RGB format. Defaults to (255, 255, 255).
+            font_color (tuple[int, int, int] | None): Text color in RGB format (0-255). Generated PDFs use it,
+                or black when it isn't set. Injection items without their own ``font_color`` use it, or white
+                when it isn't set, so injected text stays hidden by default.
             page_width (int): Width of the PDF page in mm. Defaults to 210 (A4 width).
             page_height (int): Height of the PDF page in mm. Defaults to 297 (A4 height).
             column_width (int): Width of each column in the PDF. Defaults to 0 (full page width).
@@ -87,7 +94,9 @@ def __init__(
         self._injection_items = injection_items or []
 
         # Validate font color
-        if not (isinstance(font_color, tuple) and len(font_color) == 3 and all(0 <= c <= 255 for c in font_color)):
+        if font_color is not None and not (
+            isinstance(font_color, tuple) and len(font_color) == 3 and all(0 <= c <= 255 for c in font_color)
+        ):
             raise ValueError(f"Invalid font_color: {font_color}. Must be a tuple of three integers (0-255).")
 
         # If a valid path is provided, load it into memory as BytesIO
@@ -129,6 +138,7 @@ def _build_identifier(self) -> ComponentIdentifier:
                 "page_height": self._page_height,
                 "prompt_template_hash": template_hash,
                 "existing_pdf_path": existing_pdf_path,
+                "font_color": list(self._font_color) if self._font_color is not None else None,
             }
         )
 
@@ -235,6 +245,8 @@ def _generate_pdf(self, content: str) -> bytes:
         reportlab_font = font_map.get(self._font_type, "Helvetica")
 
         c.setFont(reportlab_font, self._font_size)
+        fill_color = self._font_color or _GENERATED_TEXT_COLOR
+        c.setFillColorRGB(*(channel / 255.0 for channel in fill_color))
 
         margin = 10 * mm
         x = margin
@@ -263,6 +275,7 @@ def _generate_pdf(self, content: str) -> bytes:
                 if y < margin + line_height:
                     c.showPage()
                     c.setFont(reportlab_font, self._font_size)
+                    c.setFillColorRGB(*(channel / 255.0 for channel in fill_color))
                     y = page_height_pt - margin
 
                 # Draw the text (ReportLab origin is bottom-left)
@@ -312,7 +325,7 @@ def _modify_existing_pdf(self) -> bytes:
                         text = item.get("text", "")
                         font = item.get("font", self._font_type)
                         font_size = item.get("font_size", self._font_size)
-                        font_color = item.get("font_color", self._font_color)
+                        font_color = item.get("font_color", self._font_color or _INJECTED_TEXT_COLOR)
 
                         # Coordinate validation before calling _inject_text_into_page
                         if not (0 <= x <= page_width and 0 <= y <= page_height):

tests/unit/converter/test_pdf_converter.py

@@ -2,8 +2,10 @@
 # Licensed under the MIT license.
 
 
+import re
 from io import BytesIO
 from pathlib import Path
+from typing import Any
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -574,3 +576,80 @@ async def test_existing_pdf_output_preserves_source_stem(sqlite_instance):
     suffix = "_Jonathon_Sanchez.pdf"
     assert output_name.endswith(suffix)
     assert output_name[: -len(suffix)].isdigit()
+
+
+_FILL_COLOR_OP = re.compile(rb"([\d.]+) ([\d.]+) ([\d.]+) rg")
+
+
+def _page_fill_colors(pdf_path: Path) -> list[set[tuple[int, int, int]]]:
+    """Return the RGB fill colors (0-255) each page's content stream sets."""
+    reader = PdfReader(pdf_path)
+    return [
+        {
+            (round(float(red) * 255), round(float(green) * 255), round(float(blue) * 255))
+            for red, green, blue in _FILL_COLOR_OP.findall(page.get_contents().get_data())
+        }
+        for page in reader.pages
+    ]
+
+
+@pytest.mark.usefixtures("patch_central_database")
+@pytest.mark.parametrize(
+    ("font_color", "expected"),
+    [(None, (0, 0, 0)), ((255, 0, 0), (255, 0, 0)), ((255, 255, 255), (255, 255, 255))],
+)
+async def test_convert_async_generated_pdf_uses_font_color_on_every_page_async(
+    *, font_color: tuple[int, int, int] | None, expected: tuple[int, int, int]
+) -> None:
+    converter = PDFConverter(font_color=font_color)
+
+    result = await converter.convert_async(prompt="\n".join(f"line {index}" for index in range(120)))
+
+    page_colors = _page_fill_colors(Path(result.output_text))
+    assert len(page_colors) > 1
+    assert all(colors == {expected} for colors in page_colors)
+
+
+@pytest.mark.usefixtures("patch_central_database")
+@pytest.mark.parametrize(
+    ("converter_color", "item_color", "expected"),
+    [
+        (None, None, (255, 255, 255)),
+        ((0, 0, 255), None, (0, 0, 255)),
+        ((0, 0, 255), (255, 0, 0), (255, 0, 0)),
+        (None, (255, 0, 0), (255, 0, 0)),
+    ],
+)
+async def test_convert_async_injected_text_color_async(
+    *,
+    mock_pdf_path: Path,
+    converter_color: tuple[int, int, int] | None,
+    item_color: tuple[int, int, int] | None,
+    expected: tuple[int, int, int],
+) -> None:
+    item: dict[str, Any] = {"page": 0, "x": 50, "y": 700, "text": "Injected Text"}
+    if item_color is not None:
+        item["font_color"] = item_color
+    converter = PDFConverter(existing_pdf=mock_pdf_path, injection_items=[item], font_color=converter_color)
+
+    result = await converter.convert_async(prompt="")
+
+    assert expected in _page_fill_colors(Path(result.output_text))[0]
+
+
+@pytest.mark.usefixtures("patch_central_database")
+def test_identifier_distinguishes_font_color() -> None:
+    default = PDFConverter().get_identifier()
+    red = PDFConverter(font_color=(255, 0, 0)).get_identifier()
+    white = PDFConverter(font_color=(255, 255, 255)).get_identifier()
+
+    assert "font_color" not in default.params
+    assert red.params["font_color"] == [255, 0, 0]
+    assert len({default.hash, red.hash, white.hash}) == 3
+
+
+@pytest.mark.usefixtures("patch_central_database")
+@pytest.mark.parametrize("font_color", [(256, 0, 0), (0, 0), [0, 0, 0]])
+def test_init_rejects_invalid_font_color(font_color: Any) -> None:
+    with pytest.raises(ValueError, match="Invalid font_color"):
+        PDFConverter(font_color=font_color)