Skip to content

Commit ff1fa6b

Browse files
romanlutzCopilot
andauthored
MAINT Simplify Foundry technique mapping (#2336)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 41c8a28e-3109-4b71-b97a-3d8ec8834136
1 parent baf0855 commit ff1fa6b

2 files changed

Lines changed: 360 additions & 157 deletions

File tree

‎pyrit/scenario/scenarios/foundry/red_team_agent.py‎

Lines changed: 106 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,11 @@
1010
"""
1111

1212
import logging
13+
from collections.abc import Callable, Mapping
1314
from dataclasses import dataclass, field
1415
from inspect import signature
15-
from typing import TYPE_CHECKING, Any, TypeVar, cast
16+
from types import MappingProxyType
17+
from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, cast
1618

1719
from pyrit.common import apply_defaults
1820
from pyrit.converter import (
@@ -40,7 +42,13 @@
4042
from pyrit.converter.binary_converter import BinaryConverter
4143
from pyrit.converter.token_smuggling.ascii_smuggler_converter import AsciiSmugglerConverter
4244
from pyrit.datasets import TextJailBreak
43-
from pyrit.executor.attack import CrescendoAttack, PromptSendingAttack, RedTeamingAttack, TreeOfAttacksWithPruningAttack
45+
from pyrit.executor.attack import (
46+
AttackStrategy,
47+
CrescendoAttack,
48+
PromptSendingAttack,
49+
RedTeamingAttack,
50+
TreeOfAttacksWithPruningAttack,
51+
)
4452
from pyrit.executor.attack.core.attack_config import AttackAdversarialConfig, AttackConverterConfig, AttackScoringConfig
4553
from pyrit.models import AttackSeedGroup
4654
from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration
@@ -57,8 +65,6 @@
5765
if TYPE_CHECKING:
5866
from collections.abc import Sequence
5967

60-
from pyrit.executor.attack.core.attack_strategy import AttackStrategy
61-
6268
AttackStrategyT = TypeVar("AttackStrategyT", bound="AttackStrategy[Any, Any]")
6369
logger = logging.getLogger(__name__)
6470

@@ -193,6 +199,40 @@ def default(cls) -> "FoundryTechnique":
193199
return cls.EASY
194200

195201

202+
@dataclass(frozen=True)
203+
class _AttackSpecification:
204+
"""Declarative construction details for a Foundry attack technique."""
205+
206+
attack_type: type[AttackStrategy[Any, Any]]
207+
kwargs: tuple[tuple[str, Any], ...] = ()
208+
209+
210+
@dataclass(frozen=True)
211+
class _ConverterSpecification:
212+
"""Declarative construction details for a Foundry converter technique."""
213+
214+
factory: Callable[..., Converter]
215+
kwargs: tuple[tuple[str, Any], ...] = ()
216+
deferred_kwargs: tuple[tuple[str, Callable[[], Any]], ...] = ()
217+
target_kwarg: str | None = None
218+
219+
def create(self, *, converter_target: PromptTarget) -> Converter:
220+
"""
221+
Create a fresh converter from this specification.
222+
223+
Args:
224+
converter_target (PromptTarget): Target supplied to converters that require one.
225+
226+
Returns:
227+
Converter: The configured converter instance.
228+
"""
229+
kwargs = dict(self.kwargs)
230+
kwargs.update((name, factory()) for name, factory in self.deferred_kwargs)
231+
if self.target_kwarg:
232+
kwargs[self.target_kwarg] = converter_target
233+
return self.factory(**kwargs)
234+
235+
196236
class RedTeamAgent(Scenario):
197237
"""
198238
RedTeamAgent is a preconfigured scenario that automatically generates multiple
@@ -208,6 +248,56 @@ class RedTeamAgent(Scenario):
208248
"""
209249

210250
VERSION: int = 1
251+
_DEFAULT_ATTACK_SPECIFICATION: ClassVar[_AttackSpecification] = _AttackSpecification(PromptSendingAttack)
252+
_ATTACK_SPECIFICATIONS: ClassVar[Mapping[FoundryTechnique, _AttackSpecification]] = MappingProxyType(
253+
{
254+
FoundryTechnique.Crescendo: _AttackSpecification(CrescendoAttack),
255+
FoundryTechnique.MultiTurn: _AttackSpecification(RedTeamingAttack),
256+
FoundryTechnique.Pair: _AttackSpecification(
257+
TreeOfAttacksWithPruningAttack,
258+
kwargs=(("tree_width", 1),),
259+
),
260+
FoundryTechnique.Tap: _AttackSpecification(TreeOfAttacksWithPruningAttack),
261+
}
262+
)
263+
_CONVERTER_SPECIFICATIONS: ClassVar[Mapping[FoundryTechnique, _ConverterSpecification]] = MappingProxyType(
264+
{
265+
FoundryTechnique.AnsiAttack: _ConverterSpecification(AnsiAttackConverter),
266+
FoundryTechnique.AsciiArt: _ConverterSpecification(AsciiArtConverter),
267+
FoundryTechnique.AsciiSmuggler: _ConverterSpecification(AsciiSmugglerConverter),
268+
FoundryTechnique.Atbash: _ConverterSpecification(AtbashConverter),
269+
FoundryTechnique.Base64: _ConverterSpecification(Base64Converter),
270+
FoundryTechnique.Binary: _ConverterSpecification(BinaryConverter),
271+
FoundryTechnique.Caesar: _ConverterSpecification(
272+
CaesarConverter,
273+
kwargs=(("caesar_offset", 3),),
274+
),
275+
FoundryTechnique.CharacterSpace: _ConverterSpecification(CharacterSpaceConverter),
276+
FoundryTechnique.CharSwap: _ConverterSpecification(CharSwapConverter),
277+
FoundryTechnique.Diacritic: _ConverterSpecification(DiacriticConverter),
278+
FoundryTechnique.Flip: _ConverterSpecification(FlipConverter),
279+
FoundryTechnique.Leetspeak: _ConverterSpecification(LeetspeakConverter),
280+
FoundryTechnique.Morse: _ConverterSpecification(MorseConverter),
281+
FoundryTechnique.ROT13: _ConverterSpecification(ROT13Converter),
282+
FoundryTechnique.SuffixAppend: _ConverterSpecification(
283+
SuffixAppendConverter,
284+
kwargs=(("suffix", "!!!"),),
285+
),
286+
FoundryTechnique.StringJoin: _ConverterSpecification(StringJoinConverter),
287+
FoundryTechnique.Tense: _ConverterSpecification(
288+
TenseConverter,
289+
kwargs=(("tense", "past"),),
290+
target_kwarg="converter_target",
291+
),
292+
FoundryTechnique.UnicodeConfusable: _ConverterSpecification(UnicodeConfusableConverter),
293+
FoundryTechnique.UnicodeSubstitution: _ConverterSpecification(UnicodeSubstitutionConverter),
294+
FoundryTechnique.Url: _ConverterSpecification(UrlConverter),
295+
FoundryTechnique.Jailbreak: _ConverterSpecification(
296+
TextJailbreakConverter,
297+
deferred_kwargs=(("jailbreak_template", lambda: TextJailBreak(random_template=True)),),
298+
),
299+
}
300+
)
211301

212302
@apply_defaults
213303
def __init__(
@@ -386,70 +476,22 @@ def _get_attack_from_technique(
386476
Raises:
387477
ValueError: If a converter technique in the composite is not recognized.
388478
"""
389-
attack: AttackStrategy[Any, Any]
390-
391-
attack_type: type[AttackStrategy[Any, Any]] = PromptSendingAttack
392-
attack_kwargs: dict[str, Any] = {}
393-
if composite.attack is not None:
394-
if composite.attack == FoundryTechnique.Crescendo:
395-
attack_type = CrescendoAttack
396-
elif composite.attack == FoundryTechnique.MultiTurn:
397-
attack_type = RedTeamingAttack
398-
elif composite.attack == FoundryTechnique.Pair:
399-
attack_type = TreeOfAttacksWithPruningAttack
400-
attack_kwargs = {"tree_width": 1}
401-
elif composite.attack == FoundryTechnique.Tap:
402-
attack_type = TreeOfAttacksWithPruningAttack
403-
479+
attack_specification = self._ATTACK_SPECIFICATIONS.get(
480+
composite.attack,
481+
self._DEFAULT_ATTACK_SPECIFICATION,
482+
)
404483
converters: list[Converter] = []
405484
for technique in composite.converters:
406-
if technique == FoundryTechnique.AnsiAttack:
407-
converters.append(AnsiAttackConverter())
408-
elif technique == FoundryTechnique.AsciiArt:
409-
converters.append(AsciiArtConverter())
410-
elif technique == FoundryTechnique.AsciiSmuggler:
411-
converters.append(AsciiSmugglerConverter())
412-
elif technique == FoundryTechnique.Atbash:
413-
converters.append(AtbashConverter())
414-
elif technique == FoundryTechnique.Base64:
415-
converters.append(Base64Converter())
416-
elif technique == FoundryTechnique.Binary:
417-
converters.append(BinaryConverter())
418-
elif technique == FoundryTechnique.Caesar:
419-
converters.append(CaesarConverter(caesar_offset=3))
420-
elif technique == FoundryTechnique.CharacterSpace:
421-
converters.append(CharacterSpaceConverter())
422-
elif technique == FoundryTechnique.CharSwap:
423-
converters.append(CharSwapConverter())
424-
elif technique == FoundryTechnique.Diacritic:
425-
converters.append(DiacriticConverter())
426-
elif technique == FoundryTechnique.Flip:
427-
converters.append(FlipConverter())
428-
elif technique == FoundryTechnique.Leetspeak:
429-
converters.append(LeetspeakConverter())
430-
elif technique == FoundryTechnique.Morse:
431-
converters.append(MorseConverter())
432-
elif technique == FoundryTechnique.ROT13:
433-
converters.append(ROT13Converter())
434-
elif technique == FoundryTechnique.SuffixAppend:
435-
converters.append(SuffixAppendConverter(suffix="!!!"))
436-
elif technique == FoundryTechnique.StringJoin:
437-
converters.append(StringJoinConverter())
438-
elif technique == FoundryTechnique.Tense:
439-
converters.append(TenseConverter(tense="past", converter_target=self._adversarial_chat))
440-
elif technique == FoundryTechnique.UnicodeConfusable:
441-
converters.append(UnicodeConfusableConverter())
442-
elif technique == FoundryTechnique.UnicodeSubstitution:
443-
converters.append(UnicodeSubstitutionConverter())
444-
elif technique == FoundryTechnique.Url:
445-
converters.append(UrlConverter())
446-
elif technique == FoundryTechnique.Jailbreak:
447-
jailbreak_template = TextJailBreak(random_template=True)
448-
converters.append(TextJailbreakConverter(jailbreak_template=jailbreak_template))
449-
else:
485+
converter_specification = self._CONVERTER_SPECIFICATIONS.get(technique)
486+
if converter_specification is None:
450487
raise ValueError(f"Unknown technique: {technique}")
488+
converters.append(converter_specification.create(converter_target=self._adversarial_chat))
451489

452-
attack = self._get_attack(attack_type=attack_type, converters=converters, attack_kwargs=attack_kwargs)
490+
attack = self._get_attack(
491+
attack_type=attack_specification.attack_type,
492+
converters=converters,
493+
attack_kwargs=dict(attack_specification.kwargs),
494+
)
453495

454496
return AtomicAttack(
455497
atomic_attack_name=composite.name,

0 commit comments

Comments
 (0)