Skip to content

Commit 8cf62e2

Browse files
committed
MAINT: extract technique factory resolution into its own module
resolve_technique_factories, resolve_technique_factories_for_techniques and TechniqueResolutionError move from matrix_atomic_attack_builder to a new pyrit/scenario/core/technique_resolution module, so non-matrix scenario shapes can resolve techniques without importing the matrix builder (issue #2869). consumers and patch targets updated; resolution tests split into their own module with full parity. Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com>
1 parent 7b53310 commit 8cf62e2

7 files changed

Lines changed: 194 additions & 152 deletions

File tree

‎pyrit/scenario/core/matrix_atomic_attack_builder.py‎

Lines changed: 1 addition & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -28,16 +28,7 @@
2828
from pyrit.prompt_normalizer import ConverterConfiguration
2929
from pyrit.scenario.core.atomic_attack import AtomicAttack
3030
from pyrit.scenario.core.attack_technique import AttackTechnique
31-
32-
33-
class TechniqueResolutionError(ValueError):
34-
"""
35-
Raised when a selected scenario technique has no registered factory.
36-
37-
Subclasses ``ValueError`` so existing ``except ValueError`` handlers keep working,
38-
mirroring ``DatasetConstraintError``.
39-
"""
40-
31+
from pyrit.scenario.core.technique_resolution import resolve_technique_factories
4132

4233
if TYPE_CHECKING:
4334
from collections.abc import Callable, Mapping, Sequence
@@ -46,7 +37,6 @@ class TechniqueResolutionError(ValueError):
4637
from pyrit.prompt_target import PromptTarget
4738
from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory
4839
from pyrit.scenario.core.scenario_context import ScenarioContext
49-
from pyrit.scenario.core.scenario_technique import ScenarioTechnique
5040
from pyrit.score import Scorer, TrueFalseScorer
5141

5242
logger = logging.getLogger(__name__)
@@ -142,75 +132,6 @@ def build_baseline_atomic_attack(
142132
)
143133

144134

145-
def resolve_technique_factories(
146-
*,
147-
context: ScenarioContext,
148-
extra_factories: dict[str, AttackTechniqueFactory] | None = None,
149-
) -> dict[str, AttackTechniqueFactory]:
150-
"""
151-
Resolve a run's selected techniques to their registered ``AttackTechniqueFactory`` instances.
152-
153-
Reads the ``AttackTechniqueRegistry`` singleton and keeps only the factories whose name
154-
matches a selected technique, preserving selection order. Raises if any selected
155-
technique has no registered factory so the run cannot silently omit requested work.
156-
157-
Args:
158-
context (ScenarioContext): The resolved runtime inputs for this run.
159-
extra_factories (dict[str, AttackTechniqueFactory] | None): Scenario-local factories
160-
merged on top of the registry before filtering, so a scenario can offer techniques
161-
without registering them globally. Entries override registry factories of the same
162-
name.
163-
164-
Returns:
165-
dict[str, AttackTechniqueFactory]: Mapping of technique name to factory, ordered by
166-
the selected techniques.
167-
168-
Raises:
169-
TechniqueResolutionError: If any selected technique has no registered factory.
170-
"""
171-
return resolve_technique_factories_for_techniques(
172-
scenario_techniques=context.scenario_techniques,
173-
extra_factories=extra_factories,
174-
)
175-
176-
177-
def resolve_technique_factories_for_techniques(
178-
*,
179-
scenario_techniques: Sequence[ScenarioTechnique],
180-
extra_factories: dict[str, AttackTechniqueFactory] | None = None,
181-
) -> dict[str, AttackTechniqueFactory]:
182-
"""
183-
Resolve selected concrete techniques to their canonical factories.
184-
185-
Args:
186-
scenario_techniques (Sequence[ScenarioTechnique]): Concrete techniques to resolve.
187-
extra_factories (dict[str, AttackTechniqueFactory] | None): Scenario-local factories
188-
merged on top of the registry.
189-
190-
Returns:
191-
dict[str, AttackTechniqueFactory]: Selected factories in technique order.
192-
193-
Raises:
194-
TechniqueResolutionError: If any selected technique has no registered factory.
195-
"""
196-
from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry
197-
198-
all_factories = dict(AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise())
199-
if extra_factories:
200-
all_factories.update(extra_factories)
201-
202-
missing = list(dict.fromkeys(t.value for t in scenario_techniques if t.value not in all_factories))
203-
204-
if missing:
205-
raise TechniqueResolutionError(
206-
"The following selected attack techniques have no registered factory: "
207-
f"{', '.join(missing)}. Register the techniques (or pass them via "
208-
"extra_factories) before starting the run."
209-
)
210-
211-
return {technique.value: all_factories[technique.value] for technique in scenario_techniques}
212-
213-
214135
def filter_compatible_seed_groups(
215136
*,
216137
factory: AttackTechniqueFactory,

‎pyrit/scenario/core/scenario.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -710,6 +710,8 @@ def _build_technique_size_components(
710710

711711
from pyrit.scenario.core.matrix_atomic_attack_builder import (
712712
filter_compatible_seed_groups,
713+
)
714+
from pyrit.scenario.core.technique_resolution import (
713715
resolve_technique_factories_for_techniques,
714716
)
715717

@@ -755,6 +757,8 @@ def _get_technique_compatibility_bounds(
755757
"""
756758
from pyrit.scenario.core.matrix_atomic_attack_builder import (
757759
filter_compatible_seed_groups,
760+
)
761+
from pyrit.scenario.core.technique_resolution import (
758762
resolve_technique_factories_for_techniques,
759763
)
760764

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT license.
3+
4+
"""
5+
Technique factory resolution for scenarios.
6+
7+
Scenarios of different shapes (matrix atomic builders, adaptive scenarios,
8+
composite builders, per-objective builders) all need to resolve selected
9+
:class:`~pyrit.scenario.core.scenario_technique.ScenarioTechnique` entries to
10+
their registered ``AttackTechniqueFactory`` instances. Keeping that logic here
11+
lets non-matrix scenarios use it without importing the matrix builder.
12+
"""
13+
14+
from __future__ import annotations
15+
16+
from typing import TYPE_CHECKING
17+
18+
if TYPE_CHECKING:
19+
from collections.abc import Sequence
20+
21+
from pyrit.scenario.core.scenario_context import ScenarioContext
22+
from pyrit.scenario.core.scenario_technique import ScenarioTechnique
23+
24+
25+
class TechniqueResolutionError(ValueError):
26+
"""
27+
Raised when a selected scenario technique has no registered factory.
28+
29+
Subclasses ``ValueError`` so existing ``except ValueError`` handlers keep working,
30+
mirroring ``DatasetConstraintError``.
31+
"""
32+
33+
34+
def resolve_technique_factories(
35+
*,
36+
context: ScenarioContext,
37+
extra_factories: dict[str, AttackTechniqueFactory] | None = None,
38+
) -> dict[str, AttackTechniqueFactory]:
39+
"""
40+
Resolve a run's selected techniques to their registered ``AttackTechniqueFactory`` instances.
41+
42+
Reads the ``AttackTechniqueRegistry`` singleton and keeps only the factories whose name
43+
matches a selected technique, preserving selection order. Raises if any selected
44+
technique has no registered factory so the run cannot silently omit requested work.
45+
46+
Args:
47+
context (ScenarioContext): The resolved runtime inputs for this run.
48+
extra_factories (dict[str, AttackTechniqueFactory] | None): Scenario-local factories
49+
merged on top of the registry before filtering, so a scenario can offer techniques
50+
without registering them globally. Entries override registry factories of the same
51+
name.
52+
53+
Returns:
54+
dict[str, AttackTechniqueFactory]: Mapping of technique name to factory, ordered by
55+
the selected techniques.
56+
57+
Raises:
58+
TechniqueResolutionError: If any selected technique has no registered factory.
59+
"""
60+
return resolve_technique_factories_for_techniques(
61+
scenario_techniques=context.scenario_techniques,
62+
extra_factories=extra_factories,
63+
)
64+
65+
66+
def resolve_technique_factories_for_techniques(
67+
*,
68+
scenario_techniques: Sequence[ScenarioTechnique],
69+
extra_factories: dict[str, AttackTechniqueFactory] | None = None,
70+
) -> dict[str, AttackTechniqueFactory]:
71+
"""
72+
Resolve selected concrete techniques to their canonical factories.
73+
74+
Args:
75+
scenario_techniques (Sequence[ScenarioTechnique]): Concrete techniques to resolve.
76+
extra_factories (dict[str, AttackTechniqueFactory] | None): Scenario-local factories
77+
merged on top of the registry.
78+
79+
Returns:
80+
dict[str, AttackTechniqueFactory]: Selected factories in technique order.
81+
82+
Raises:
83+
TechniqueResolutionError: If any selected technique has no registered factory.
84+
"""
85+
from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry
86+
87+
all_factories = dict(AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise())
88+
if extra_factories:
89+
all_factories.update(extra_factories)
90+
91+
missing = list(dict.fromkeys(t.value for t in scenario_techniques if t.value not in all_factories))
92+
93+
if missing:
94+
raise TechniqueResolutionError(
95+
"The following selected attack techniques have no registered factory: "
96+
f"{', '.join(missing)}. Register the techniques (or pass them via "
97+
"extra_factories) before starting the run."
98+
)
99+
100+
return {technique.value: all_factories[technique.value] for technique in scenario_techniques}

‎pyrit/scenario/scenarios/benchmark/adversarial.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,12 @@
3131
from pyrit.scenario.core.matrix_atomic_attack_builder import (
3232
MatrixAtomicAttackBuilder,
3333
filter_compatible_seed_groups,
34+
)
35+
from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario
36+
from pyrit.scenario.core.technique_resolution import (
3437
resolve_technique_factories,
3538
resolve_technique_factories_for_techniques,
3639
)
37-
from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario
3840

3941
if TYPE_CHECKING:
4042
from pyrit.prompt_target import PromptTarget

‎tests/unit/scenario/core/test_matrix_atomic_attack_builder.py‎

Lines changed: 0 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,8 @@
2626
from pyrit.scenario.core.matrix_atomic_attack_builder import (
2727
MatrixAtomicAttackBuilder,
2828
MatrixCombo,
29-
TechniqueResolutionError,
3029
build_baseline_atomic_attack,
3130
build_matrix_atomic_attacks,
32-
resolve_technique_factories,
3331
)
3432
from pyrit.scenario.core.scenario_context import ScenarioContext
3533
from pyrit.score import TrueFalseScorer
@@ -388,69 +386,6 @@ def _patch_registry(factories: dict):
388386

389387

390388
@pytest.mark.usefixtures("patch_central_database")
391-
class TestResolveTechniqueFactories:
392-
"""``resolve_technique_factories`` filters the registry to the selected techniques."""
393-
394-
def test_keeps_only_selected_in_order(self):
395-
factories = {
396-
"alpha": _mock_factory(name="alpha"),
397-
"beta": _mock_factory(name="beta"),
398-
"gamma": _mock_factory(name="gamma"),
399-
}
400-
context = _context(techniques=[_technique("beta"), _technique("alpha")])
401-
with _patch_registry(factories):
402-
resolved = resolve_technique_factories(context=context)
403-
assert list(resolved.keys()) == ["beta", "alpha"]
404-
405-
def test_raises_when_any_selected_technique_is_missing(self):
406-
factories = {"alpha": _mock_factory(name="alpha")}
407-
context = _context(techniques=[_technique("alpha"), _technique("missing")])
408-
with _patch_registry(factories), pytest.raises(TechniqueResolutionError, match="missing"):
409-
resolve_technique_factories(context=context)
410-
411-
def test_raises_when_all_selected_techniques_missing(self):
412-
"""A nonempty selection resolving to nothing must fail loudly, not run baseline-only."""
413-
factories = {"alpha": _mock_factory(name="alpha")}
414-
context = _context(techniques=[_technique("missing_a"), _technique("missing_b")])
415-
with _patch_registry(factories), pytest.raises(TechniqueResolutionError, match="missing_a"):
416-
resolve_technique_factories(context=context)
417-
418-
def test_empty_selection_resolves_without_error(self):
419-
context = _context(techniques=[])
420-
with _patch_registry({}):
421-
assert resolve_technique_factories(context=context) == {}
422-
423-
def test_error_lists_each_missing_technique_once_in_selection_order(self):
424-
factories = {"alpha": _mock_factory(name="alpha")}
425-
context = _context(
426-
techniques=[
427-
_technique("missing_a"),
428-
_technique("alpha"),
429-
_technique("missing_b"),
430-
_technique("missing_a"),
431-
]
432-
)
433-
with _patch_registry(factories), pytest.raises(TechniqueResolutionError) as exc_info:
434-
resolve_technique_factories(context=context)
435-
message = str(exc_info.value)
436-
assert message.index("missing_a") < message.index("missing_b")
437-
assert message.count("missing_a") == 1
438-
439-
def test_extra_factories_merged_and_override_registry(self):
440-
registry_factories = {"alpha": _mock_factory(name="alpha")}
441-
local_alpha = _mock_factory(name="alpha")
442-
local_only = _mock_factory(name="local")
443-
context = _context(techniques=[_technique("alpha"), _technique("local")])
444-
with _patch_registry(registry_factories):
445-
resolved = resolve_technique_factories(
446-
context=context,
447-
extra_factories={"alpha": local_alpha, "local": local_only},
448-
)
449-
assert list(resolved.keys()) == ["alpha", "local"]
450-
assert resolved["alpha"] is local_alpha # extra overrides the registry factory of the same name
451-
assert resolved["local"] is local_only # local-only factory is selectable without global registration
452-
453-
454389
@pytest.mark.usefixtures("patch_central_database")
455390
class TestBuildMatrixAtomicAttacks:
456391
"""``build_matrix_atomic_attacks`` wires the context into the builder in one call."""
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT license.
3+
4+
"""Tests for scenario technique factory resolution (``technique_resolution``)."""
5+
6+
import pytest
7+
8+
from pyrit.scenario.core.technique_resolution import (
9+
TechniqueResolutionError,
10+
resolve_technique_factories,
11+
)
12+
from tests.unit.scenario.core.test_matrix_atomic_attack_builder import (
13+
_context,
14+
_mock_factory,
15+
_patch_registry,
16+
_technique,
17+
)
18+
19+
20+
class TestResolveTechniqueFactories:
21+
"""``resolve_technique_factories`` filters the registry to the selected techniques."""
22+
23+
def test_keeps_only_selected_in_order(self):
24+
factories = {
25+
"alpha": _mock_factory(name="alpha"),
26+
"beta": _mock_factory(name="beta"),
27+
"gamma": _mock_factory(name="gamma"),
28+
}
29+
context = _context(techniques=[_technique("beta"), _technique("alpha")])
30+
with _patch_registry(factories):
31+
resolved = resolve_technique_factories(context=context)
32+
assert list(resolved.keys()) == ["beta", "alpha"]
33+
34+
def test_raises_when_any_selected_technique_is_missing(self):
35+
factories = {"alpha": _mock_factory(name="alpha")}
36+
context = _context(techniques=[_technique("alpha"), _technique("missing")])
37+
with _patch_registry(factories), pytest.raises(TechniqueResolutionError, match="missing"):
38+
resolve_technique_factories(context=context)
39+
40+
def test_raises_when_all_selected_techniques_missing(self):
41+
"""A nonempty selection resolving to nothing must fail loudly, not run baseline-only."""
42+
factories = {"alpha": _mock_factory(name="alpha")}
43+
context = _context(techniques=[_technique("missing_a"), _technique("missing_b")])
44+
with _patch_registry(factories), pytest.raises(TechniqueResolutionError, match="missing_a"):
45+
resolve_technique_factories(context=context)
46+
47+
def test_empty_selection_resolves_without_error(self):
48+
context = _context(techniques=[])
49+
with _patch_registry({}):
50+
assert resolve_technique_factories(context=context) == {}
51+
52+
def test_error_lists_each_missing_technique_once_in_selection_order(self):
53+
factories = {"alpha": _mock_factory(name="alpha")}
54+
context = _context(
55+
techniques=[
56+
_technique("missing_a"),
57+
_technique("alpha"),
58+
_technique("missing_b"),
59+
_technique("missing_a"),
60+
]
61+
)
62+
with _patch_registry(factories), pytest.raises(TechniqueResolutionError) as exc_info:
63+
resolve_technique_factories(context=context)
64+
message = str(exc_info.value)
65+
assert message.index("missing_a") < message.index("missing_b")
66+
assert message.count("missing_a") == 1
67+
68+
def test_extra_factories_merged_and_override_registry(self):
69+
registry_factories = {"alpha": _mock_factory(name="alpha")}
70+
local_alpha = _mock_factory(name="alpha")
71+
local_only = _mock_factory(name="local")
72+
context = _context(techniques=[_technique("alpha"), _technique("local")])
73+
with _patch_registry(registry_factories):
74+
resolved = resolve_technique_factories(
75+
context=context,
76+
extra_factories={"alpha": local_alpha, "local": local_only},
77+
)
78+
assert list(resolved.keys()) == ["alpha", "local"]
79+
assert resolved["alpha"] is local_alpha # extra overrides the registry factory of the same name
80+
assert resolved["local"] is local_only # local-only factory is selectable without global registration

0 commit comments

Comments
 (0)