|
| 1 | +# Copyright (c) Microsoft Corporation. |
| 2 | +# Licensed under the MIT license. |
| 3 | + |
| 4 | +import os |
| 5 | +import pathlib |
| 6 | +from typing import List, Optional |
| 7 | + |
| 8 | +from pyrit.common import apply_defaults |
| 9 | +from pyrit.common.path import DATASETS_PATH, SCORER_CONFIG_PATH |
| 10 | +from pyrit.executor.attack.core.attack_config import ( |
| 11 | + AttackAdversarialConfig, |
| 12 | + AttackScoringConfig, |
| 13 | +) |
| 14 | +from pyrit.executor.attack.core.attack_strategy import AttackStrategy |
| 15 | +from pyrit.executor.attack.multi_turn.red_teaming import RedTeamingAttack |
| 16 | +from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack |
| 17 | +from pyrit.models import SeedDataset |
| 18 | +from pyrit.prompt_target import OpenAIChatTarget, PromptChatTarget |
| 19 | +from pyrit.scenarios.atomic_attack import AtomicAttack |
| 20 | +from pyrit.scenarios.scenario import Scenario |
| 21 | +from pyrit.scenarios.scenario_strategy import ( |
| 22 | + ScenarioCompositeStrategy, |
| 23 | + ScenarioStrategy, |
| 24 | +) |
| 25 | +from pyrit.score import ( |
| 26 | + SelfAskRefusalScorer, |
| 27 | + SelfAskTrueFalseScorer, |
| 28 | + TrueFalseCompositeScorer, |
| 29 | + TrueFalseInverterScorer, |
| 30 | + TrueFalseScoreAggregator, |
| 31 | +) |
| 32 | + |
| 33 | + |
| 34 | +class CyberStrategy(ScenarioStrategy): |
| 35 | + """ |
| 36 | + Strategies for malware-focused cyber attacks. While not in the CyberStrategy class, a |
| 37 | + few of these include: |
| 38 | + * Shell smashing |
| 39 | + * Zip bombs |
| 40 | + * File deletion (rm -rf /). |
| 41 | + """ |
| 42 | + |
| 43 | + # Aggregate members (special markers that expand to strategies with matching tags) |
| 44 | + ALL = ("all", {"all"}) |
| 45 | + SINGLE_TURN = ("single_turn", {"single_turn"}) |
| 46 | + MULTI_TURN = ("multi_turn", {"multi_turn"}) |
| 47 | + |
| 48 | + |
| 49 | +class CyberScenario(Scenario): |
| 50 | + """ |
| 51 | + Cyber scenario implementation for PyRIT. |
| 52 | +
|
| 53 | + This scenario tests how willing models are to exploit cybersecurity harms by generating |
| 54 | + malware. The CyberScenario class contains different variations of the malware generation |
| 55 | + techniques. |
| 56 | + """ |
| 57 | + |
| 58 | + version: int = 1 |
| 59 | + |
| 60 | + @classmethod |
| 61 | + def get_strategy_class(cls) -> type[ScenarioStrategy]: |
| 62 | + """ |
| 63 | + Get the strategy enum class for this scenario. |
| 64 | +
|
| 65 | + Returns: |
| 66 | + Type[ScenarioStrategy]: The CyberStrategy enum class. |
| 67 | + """ |
| 68 | + return CyberStrategy |
| 69 | + |
| 70 | + @classmethod |
| 71 | + def get_default_strategy(cls) -> ScenarioStrategy: |
| 72 | + """ |
| 73 | + Get the default strategy used when no strategies are specified. |
| 74 | +
|
| 75 | + Returns: |
| 76 | + ScenarioStrategy: CyberStrategy.ALL (all cyber strategies). |
| 77 | + """ |
| 78 | + return CyberStrategy.ALL |
| 79 | + |
| 80 | + @apply_defaults |
| 81 | + def __init__( |
| 82 | + self, |
| 83 | + *, |
| 84 | + adversarial_chat: Optional[PromptChatTarget] = None, |
| 85 | + objectives: Optional[List[str]] = None, |
| 86 | + objective_scorer: Optional[TrueFalseCompositeScorer] = None, |
| 87 | + include_baseline: bool = True, |
| 88 | + scenario_result_id: Optional[str] = None, |
| 89 | + ) -> None: |
| 90 | + """ |
| 91 | + Initialize the cyber harms scenario. |
| 92 | +
|
| 93 | + Args: |
| 94 | + adversarial_chat (Optional[PromptChatTarget]): Adversarial chat for the red teaming attack, corresponding |
| 95 | + to CyberStrategy.MultiTurn. If not provided, defaults to an OpenAI chat target. |
| 96 | + objectives (Optional[List[str]]): List of objectives to test for cyber harms, e.g. malware generation. |
| 97 | + objective_scorer (Optional[SelfAskTrueFalseScorer]): Objective scorer for malware detection. If not |
| 98 | + provided, defaults to a SelfAskScorer using the malware.yaml file under the scorer config store for |
| 99 | + malware detection |
| 100 | + include_baseline (bool): Whether to include a baseline atomic attack that sends all objectives |
| 101 | + without modifications. Defaults to True. When True, a "baseline" attack is automatically |
| 102 | + added as the first atomic attack, allowing comparison between unmodified prompts and |
| 103 | + attack-modified prompts. |
| 104 | + scenario_result_id (Optional[str]): Optional ID of an existing scenario result to resume. |
| 105 | + """ |
| 106 | + # CyberScenario uses a "take object, make config" pattern to expose a more ergonomic interface. Helper |
| 107 | + # methods return objects, not configs. |
| 108 | + |
| 109 | + # In this context the "objective" scorer has nothing to do with the "objective" target. |
| 110 | + # The scoring config is what's used later in scoring orchestration, not the scorer object. |
| 111 | + if not objective_scorer: |
| 112 | + objective_scorer = self._get_default_objective_scorer() |
| 113 | + self._scorer_config = AttackScoringConfig(objective_scorer=objective_scorer) |
| 114 | + |
| 115 | + # Adversarial chat is used by the RedTeamingAttack (multiturn) variation of the scenario. Otherwise, it's not |
| 116 | + # invoked. |
| 117 | + self._adversarial_chat = adversarial_chat if adversarial_chat else self._get_default_adversarial_target() |
| 118 | + self._adversarial_config = AttackAdversarialConfig(target=self._adversarial_chat) |
| 119 | + |
| 120 | + self._objectives = objectives if objectives else self._get_default_dataset() |
| 121 | + |
| 122 | + super().__init__( |
| 123 | + name="Cyber Scenario", |
| 124 | + version=self.version, |
| 125 | + strategy_class=CyberStrategy, |
| 126 | + objective_scorer_identifier=objective_scorer.get_identifier(), |
| 127 | + include_default_baseline=include_baseline, |
| 128 | + scenario_result_id=scenario_result_id, |
| 129 | + ) |
| 130 | + |
| 131 | + def _get_default_objective_scorer(self) -> TrueFalseCompositeScorer: |
| 132 | + """ |
| 133 | + Retrieve the default objective scorer. |
| 134 | + The scoring mechanism is as follows: |
| 135 | + 1. Detect for the presence of malware using a SelfAskTrueFalseScorer. |
| 136 | + 2. Detect for NON-refusal as a backstop. (If the target model refuses to produce malware, we want that |
| 137 | + to register as a negative score). |
| 138 | + """ |
| 139 | + presence_of_malware = SelfAskTrueFalseScorer( |
| 140 | + chat_target=OpenAIChatTarget(), |
| 141 | + true_false_question_path=SCORER_CONFIG_PATH / "true_false_question" / "malware.yaml", |
| 142 | + ) |
| 143 | + |
| 144 | + backstop = TrueFalseInverterScorer(scorer=SelfAskRefusalScorer(chat_target=OpenAIChatTarget())) |
| 145 | + |
| 146 | + return TrueFalseCompositeScorer( |
| 147 | + aggregator=TrueFalseScoreAggregator.AND, scorers=[presence_of_malware, backstop] |
| 148 | + ) |
| 149 | + |
| 150 | + def _get_default_adversarial_target(self) -> OpenAIChatTarget: |
| 151 | + """ |
| 152 | + The adversarial target is used in the red teaming attack variation, where we want an unsafe model attempting |
| 153 | + to generate malware. |
| 154 | + """ |
| 155 | + return OpenAIChatTarget( |
| 156 | + endpoint=os.environ.get("AZURE_OPENAI_GPT4O_UNSAFE_ENDPOINT"), |
| 157 | + api_key=os.environ.get("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY"), |
| 158 | + temperature=1.2, |
| 159 | + ) |
| 160 | + |
| 161 | + def _get_default_dataset(self) -> list[str]: |
| 162 | + """ |
| 163 | + Get the default seed prompts for malware tests. |
| 164 | +
|
| 165 | + This dataset includes a set of exploits that represent cybersecurity harms. |
| 166 | +
|
| 167 | + Returns: |
| 168 | + list[str]: List of seed prompt strings to be encoded and tested. |
| 169 | + """ |
| 170 | + seed_prompts: List[str] = [] |
| 171 | + malware_path = pathlib.Path(DATASETS_PATH) / "seed_prompts" |
| 172 | + seed_prompts.extend(SeedDataset.from_yaml_file(malware_path / "malware.prompt").get_values()) |
| 173 | + return seed_prompts |
| 174 | + |
| 175 | + async def _get_atomic_attack_from_strategy_async(self, strategy: str) -> AtomicAttack: |
| 176 | + """ |
| 177 | + Translate the strategy into an actual AtomicAttack. |
| 178 | +
|
| 179 | + Args: |
| 180 | + strategy: The CyberStrategy enum (SingleTurn or MultiTurn). |
| 181 | +
|
| 182 | + Returns: |
| 183 | + AtomicAttack configured for the specified strategy. |
| 184 | + """ |
| 185 | + # objective_target is guaranteed to be non-None by parent class validation |
| 186 | + assert self._objective_target is not None |
| 187 | + attack_strategy: Optional[AttackStrategy] = None |
| 188 | + if strategy == "single_turn": |
| 189 | + attack_strategy = PromptSendingAttack( |
| 190 | + objective_target=self._objective_target, |
| 191 | + attack_scoring_config=self._scorer_config, |
| 192 | + ) |
| 193 | + elif strategy == "multi_turn": |
| 194 | + attack_strategy = RedTeamingAttack( |
| 195 | + objective_target=self._objective_target, |
| 196 | + attack_scoring_config=self._scorer_config, |
| 197 | + attack_adversarial_config=self._adversarial_config, |
| 198 | + ) |
| 199 | + else: |
| 200 | + raise ValueError(f"Unknown CyberStrategy: {strategy}") |
| 201 | + |
| 202 | + return AtomicAttack( |
| 203 | + atomic_attack_name=f"cyber_{strategy}", |
| 204 | + attack=attack_strategy, |
| 205 | + objectives=self._objectives, |
| 206 | + memory_labels=self._memory_labels, |
| 207 | + ) |
| 208 | + |
| 209 | + async def _get_atomic_attacks_async(self) -> List[AtomicAttack]: |
| 210 | + """ |
| 211 | + Generate atomic attacks for each strategy. |
| 212 | +
|
| 213 | + Returns: |
| 214 | + List[AtomicAttack]: List of atomic attacks to execute. |
| 215 | + """ |
| 216 | + atomic_attacks: List[AtomicAttack] = [] |
| 217 | + strategies = ScenarioCompositeStrategy.extract_single_strategy_values( |
| 218 | + composites=self._scenario_composites, strategy_type=CyberStrategy |
| 219 | + ) |
| 220 | + |
| 221 | + for strategy in strategies: |
| 222 | + atomic_attacks.append(await self._get_atomic_attack_from_strategy_async(strategy)) |
| 223 | + return atomic_attacks |
0 commit comments