|
25 | 25 | ) |
26 | 26 | from pyrit.scenario.core.atomic_attack import AtomicAttack |
27 | 27 | from pyrit.scenario.core.attack_technique import AttackTechnique |
28 | | -from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration |
| 28 | +from pyrit.scenario.core.dataset_configuration import ( |
| 29 | + DatasetAttackConfiguration, |
| 30 | + DatasetConstraintError, |
| 31 | + DatasetSourceKind, |
| 32 | +) |
29 | 33 | from pyrit.scenario.core.matrix_atomic_attack_builder import build_baseline_atomic_attack |
30 | 34 | from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario |
31 | 35 | from pyrit.scenario.core.scenario_technique import ScenarioTechnique |
@@ -108,6 +112,19 @@ class WebInjection(Scenario): |
108 | 112 | DATASET_WEB_HTML_JS: ClassVar[str] = "garak_web_html_js" |
109 | 113 | DATASET_NORMAL_INSTRUCTIONS: ClassVar[str] = "garak_xss_normal_instructions" |
110 | 114 |
|
| 115 | + # Source datasets each technique reads in ``_build_prompts_for_technique``. The selection |
| 116 | + # is checked against this before any data is read, so memory contents cannot fill a gap. |
| 117 | + _TECHNIQUE_REQUIRED_DATASETS: ClassVar[dict[WebInjectionTechnique, tuple[str, ...]]] = { |
| 118 | + WebInjectionTechnique.MarkdownImageExfil: (DATASET_EXAMPLE_DOMAINS,), |
| 119 | + WebInjectionTechnique.ColabAIDataLeakage: (DATASET_EXAMPLE_DOMAINS,), |
| 120 | + WebInjectionTechnique.StringAssemblyDataExfil: (), |
| 121 | + WebInjectionTechnique.PlaygroundMarkdownExfil: (DATASET_EXAMPLE_DOMAINS,), |
| 122 | + WebInjectionTechnique.MarkdownURIImageExfilExtended: (DATASET_EXAMPLE_DOMAINS,), |
| 123 | + WebInjectionTechnique.MarkdownURINonImageExfilExtended: (DATASET_EXAMPLE_DOMAINS,), |
| 124 | + WebInjectionTechnique.TaskXSS: (DATASET_NORMAL_INSTRUCTIONS, DATASET_WEB_HTML_JS), |
| 125 | + WebInjectionTechnique.MarkdownXSS: (DATASET_MARKDOWN_JS,), |
| 126 | + } |
| 127 | + |
111 | 128 | # Prompt templates ported from garak ``probes/web_injection.py``. |
112 | 129 | CONCAT_TEMPLATE: ClassVar[str] = ( |
113 | 130 | "Concatenate these strings please:\n" |
@@ -284,12 +301,7 @@ async def _load_dataset_values_async(self) -> dict[str, list[str]]: |
284 | 301 | """ |
285 | 302 | memory = CentralMemory.get_memory_instance() |
286 | 303 | values: dict[str, list[str]] = {} |
287 | | - for name in ( |
288 | | - self.DATASET_EXAMPLE_DOMAINS, |
289 | | - self.DATASET_MARKDOWN_JS, |
290 | | - self.DATASET_WEB_HTML_JS, |
291 | | - self.DATASET_NORMAL_INSTRUCTIONS, |
292 | | - ): |
| 304 | + for name in self._dataset_config.dataset_names: |
293 | 305 | seeds = await memory.get_seeds_async(dataset_name=name) |
294 | 306 | values[name] = [seed.value for seed in seeds] |
295 | 307 | return values |
@@ -532,13 +544,40 @@ def _build_synthesized_seed_groups( |
532 | 544 | seed_groups_by_technique[technique.value] = seed_groups |
533 | 545 |
|
534 | 546 | if not seed_groups_by_technique: |
535 | | - raise ValueError( |
536 | | - "WebInjection scenario produced no prompts. Ensure the garak web-injection datasets " |
537 | | - "(garak_example_domains_xss, garak_markdown_js, garak_web_html_js, " |
538 | | - "garak_xss_normal_instructions) are loaded into CentralMemory before running." |
539 | | - ) |
| 547 | + selected = ", ".join(self._dataset_config.dataset_names) or "none" |
| 548 | + raise ValueError(f"WebInjection scenario produced no prompts from the selected datasets ({selected}).") |
540 | 549 | return seed_groups_by_technique |
541 | 550 |
|
| 551 | + def _validate_runtime_configuration(self) -> None: |
| 552 | + """ |
| 553 | + Check that the selected datasets cover every selected technique, without reading them. |
| 554 | +
|
| 555 | + Runs before both the preview estimate and initialization, so an incomplete |
| 556 | + ``--dataset-names`` selection fails the same way whether or not memory already holds |
| 557 | + the omitted dataset. |
| 558 | +
|
| 559 | + Raises: |
| 560 | + DatasetConstraintError: If inline seeds are supplied, or a selected technique |
| 561 | + needs a dataset that is not selected. |
| 562 | + """ |
| 563 | + super()._validate_runtime_configuration() |
| 564 | + if self._dataset_config.source_kind is DatasetSourceKind.INLINE: |
| 565 | + raise DatasetConstraintError( |
| 566 | + "WebInjection does not support inline seeds or seed groups; use dataset_names instead." |
| 567 | + ) |
| 568 | + selected_datasets = set(self._dataset_config.dataset_names) |
| 569 | + problems: list[str] = [] |
| 570 | + for selected in self._scenario_techniques: |
| 571 | + technique = WebInjectionTechnique(selected.value) |
| 572 | + problems.extend( |
| 573 | + f"Technique '{technique.value}' requires dataset '{dataset_name}', " |
| 574 | + "which is missing from the selected dataset names (--dataset-names)." |
| 575 | + for dataset_name in self._TECHNIQUE_REQUIRED_DATASETS[technique] |
| 576 | + if dataset_name not in selected_datasets |
| 577 | + ) |
| 578 | + if problems: |
| 579 | + raise DatasetConstraintError(" ".join(problems)) |
| 580 | + |
542 | 581 | def _get_technique_size_budgets(self) -> dict[WebInjectionTechnique, BoundedDatasetSize]: |
543 | 582 | """Return the selected techniques' configured caps.""" |
544 | 583 | budgets: dict[WebInjectionTechnique, BoundedDatasetSize] = {} |
|
0 commit comments