# Copyright (c) Microsoft Corporation. # Licensed under the MIT license. """ AtomicAttack class for executing single attack configurations against datasets. This module provides the AtomicAttack class that represents an atomic test combining an attack, a dataset, and execution parameters. Multiple AtomicAttacks can be grouped together into larger test scenarios for comprehensive security testing. Eventually it's a good goal to unify attacks as much as we can. But there are times when that may not be possible or make sense. So this class exists to have a common interface for scenarios. """ from __future__ import annotations import logging from typing import TYPE_CHECKING, Any from pyrit.common.utils import to_sha256 from pyrit.executor.attack import AttackExecutor, AttackExecutorResult from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution from pyrit.models import ( AtomicAttackEvaluationIdentifier, AtomicAttackIdentifier, AttackResult, AttackSeedGroup, ScenarioRunPlanGroupKind, config_hash, ) from pyrit.scenario.core.technique_requirements import PreparedSeedGroup, prepare_seed_group if TYPE_CHECKING: from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.score import TrueFalseScorer logger = logging.getLogger(__name__) class AtomicAttack: """ Represents a single atomic attack test combining an attack strategy and dataset. An AtomicAttack is an executable unit that executes a configured attack against all objectives in a dataset. Multiple AtomicAttacks can be grouped together into larger test scenarios for comprehensive security testing and evaluation. The AtomicAttack uses AttackSeedGroups as the single source of truth for objectives, prepended conversations, and next messages. Each AttackSeedGroup must have an objective set. An ``AttackTechnique`` bundles the attack strategy with an optional ``AttackTechniqueSeedGroup``, cleanly separating "how to attack" from "what to attack" (the objective). """ def __init__( self, *, atomic_attack_name: str, display_group: str | None = None, technique_name: str | None = None, attack_technique: AttackTechnique, seed_groups: list[AttackSeedGroup], adversarial_chat: PromptTarget | None = None, objective_scorer: TrueFalseScorer | None = None, memory_labels: dict[str, str] | None = None, group_kind: ScenarioRunPlanGroupKind = ScenarioRunPlanGroupKind.ATTACK, **attack_execute_params: Any, ) -> None: """ Initialize an atomic attack with an attack strategy and seed groups. Args: atomic_attack_name: Unique key for this atomic attack. Used for resume tracking and result persistence — must be unique across all ``AtomicAttack`` instances in a scenario. display_group: Optional label for grouping results in user-facing output (console printer, reports). When ``None``, falls back to ``atomic_attack_name``. technique_name: Optional catalog name for the technique that built this atomic attack. attack_technique: An AttackTechnique bundling the attack strategy and optional technique seeds. seed_groups: List of seed attack groups. Each must be a ``AttackSeedGroup`` (which guarantees exactly one objective). adversarial_chat: Optional chat target for generating adversarial prompts or simulated conversations. objective_scorer: Optional scorer for evaluating simulated conversations. memory_labels: Additional labels to apply to prompts. group_kind: What this group runs, recorded in the scenario's run plan. ``build_baseline_atomic_attack`` passes ``BASELINE`` and Adaptive scenarios pass ``ADAPTIVE``. It describes the group and does not change its identity or what it executes. **attack_execute_params: Additional parameters to pass to the attack execution method. Raises: ValueError: If seed_groups list is empty. TypeError: If any entry of ``seed_groups`` is not a ``AttackSeedGroup``. """ self.atomic_attack_name = atomic_attack_name self.display_group = display_group or atomic_attack_name self._technique_name = technique_name self._group_kind = group_kind self._attack_technique = attack_technique # Validate seed_groups if not seed_groups: raise ValueError("seed_groups list cannot be empty") # Validate that each seed_group is actually a AttackSeedGroup (which Pydantic # already ensured holds the AtomicAttack invariant of "exactly one objective" # at construction time). A plain SeedGroup or AttackTechniqueSeedGroup is not # accepted here even though they share a base class. for sg in seed_groups: if not isinstance(sg, AttackSeedGroup): raise TypeError(f"seed_groups must contain AttackSeedGroup instances; got {type(sg).__name__}.") self._seed_groups = seed_groups self._validate_unique_objective_hashes() attack_technique.requirements.validate_target( target=attack_technique.attack.get_objective_target(), attack_class=attack_technique.attack.__class__, ) self._prepared_seed_groups: dict[str, PreparedSeedGroup] = { group.logical_id: prepare_seed_group( seed_group=group, requirements=attack_technique.requirements.seed_group, seed_technique=attack_technique.seed_technique, ) for group in seed_groups } self._adversarial_chat = adversarial_chat self._objective_scorer = objective_scorer self._memory_labels = memory_labels or {} self._attack_execute_params = attack_execute_params # Set via set_scenario_result_id() by Scenario._execute_scenario_async # before run_async. When set, each persisted AttackResult is linked to # the scenario via the attribution_parent_id foreign key on # AttackResultEntry. self._scenario_result_id: str | None = None logger.info( f"Initialized atomic attack with {len(self._seed_groups)} seed groups, " f"attack type: {type(self._attack_technique.attack).__name__}" ) def set_scenario_result_id(self, scenario_result_id: str | None) -> None: """ Bind this atomic attack to a scenario result for attribution. Called by ``Scenario._execute_scenario_async`` before each ``run_async`` so persisted ``AttackResult`` rows carry the ``attribution_parent_id`` foreign key back to the scenario. Pass ``None`` to clear the binding (e.g. when running an atomic attack outside of a scenario). Args: scenario_result_id (str | None): The scenario result UUID this atomic attack belongs to, or ``None`` to detach. """ self._scenario_result_id = scenario_result_id def _validate_unique_objective_hashes(self) -> None: """ Ensure each seed group in this atomic attack has a unique objective hash. Within a single ``AtomicAttack`` (one ``atomic_attack_name``, one technique), the objective text identifies a unit of work. Duplicates would mean two indistinguishable rows on the write side, which makes resume reconciliation ambiguous — the hash-based resume key treats a set of hashes as already-done, with no way to distinguish which of two duplicate rows is "the one" that is still outstanding. The hash is currently derived from objective text only. A future iteration may hash the full ``SeedGroup`` (minus technique-specific fields) so two seed groups that share an objective string but differ in other inputs can coexist in one atomic attack. Raises: ValueError: If two seed groups share the same ``objective_sha256``. """ seen: dict[str, int] = {} for sg in self._seed_groups: if sg.objective is None: continue sha = to_sha256(sg.objective.value) if sha in seen: raise ValueError( f"AtomicAttack '{self.atomic_attack_name}' has duplicate objective hash " f"{sha[:12]}... across seed_groups; each (objective, technique) pair must be unique." ) seen[sha] = 1 @property def attack_technique(self) -> AttackTechnique: """The attack technique for this atomic attack.""" return self._attack_technique @property def technique_name(self) -> str | None: """Catalog name of the technique that built this attack.""" return self._technique_name @property def group_kind(self) -> ScenarioRunPlanGroupKind: """What this group runs, as recorded in the scenario's run plan.""" return self._group_kind @property def technique_eval_hash(self) -> str: """ Behavioral evaluation hash for this atomic attack's technique configuration. Builds an ``AtomicAttack`` identifier from this attack's technique (without any seed group) and runs it through ``AtomicAttackEvaluationIdentifier`` so target/scorer/seed-identifier noise is stripped per the standard atomic-attack eval rules. The result is stable across resume runs and across different seed groups, which is what makes it usable as the resume disambiguator alongside ``atomic_attack_name``. """ composite = AtomicAttackIdentifier.build( technique_identifier=self._attack_technique.get_identifier(), seed_group=None, ) return AtomicAttackEvaluationIdentifier(composite).eval_hash @property def logical_group_id(self) -> str: """The stable identity of this planned atomic-attack group.""" return config_hash( { "atomic_attack_name": self.atomic_attack_name, "technique_eval_hash": self.technique_eval_hash, } ) @property def objectives(self) -> list[str]: """ The objectives from the seed groups. Returns: list[str]: List of objectives from all seed groups. """ return [sg.objective.value for sg in self._seed_groups if sg.objective is not None] @property def seed_groups(self) -> list[AttackSeedGroup]: """ A copy of the seed groups list for this atomic attack. Returns: list[AttackSeedGroup]: A copy of the seed groups list. """ return list(self._seed_groups) @property def seed_group_adaptations(self) -> dict[str, str]: """Applied adaptations keyed by the retained original group's logical ID.""" return { group.logical_id: prepared.adaptation for group in self._seed_groups if (prepared := self._prepared_seed_groups[group.logical_id]).adaptation is not None } def drop_seed_groups_with_hashes(self, *, hashes: set[str]) -> None: """ Drop seed groups whose ``objective_sha256`` is in ``hashes``. This is the resume filter: within an atomic attack, ``objective_sha256`` is the stable identity (enforced unique by ``__init__``). Content-derived keys are robust to reordering and resampling, so resume produces the right remaining-work set even when ``get_seed_groups()`` is rebuilt from scratch on each ``run_async()``. Args: hashes (set[str]): SHA256 hashes of objective text for seed groups to drop (typically those that have already produced a non-error ``AttackResult``). """ self._seed_groups = [ sg for sg in self._seed_groups if sg.objective is None or to_sha256(sg.objective.value) not in hashes ] def keep_seed_groups_with_hashes(self, *, hashes: set[str]) -> set[str]: """ Keep only seed groups whose ``objective_sha256`` is in ``hashes``. Inverse of ``drop_seed_groups_with_hashes``: used on resume to replay the originally-sampled subset and ignore any seed groups that were added since (or that landed in this run's fresh ``random.sample`` draw and are no longer in the persisted set). Args: hashes (set[str]): SHA256 hashes of objective text for seed groups to keep. Returns: set[str]: The hashes that were actually retained (intersection of ``hashes`` and the current seed_groups' hashes). The caller can union these across atomic attacks to detect persisted hashes that no longer exist in the dataset. """ retained: set[str] = set() new_groups: list[AttackSeedGroup] = [] for sg in self._seed_groups: if sg.objective is None: continue sha = to_sha256(sg.objective.value) if sha in hashes: retained.add(sha) new_groups.append(sg) self._seed_groups = new_groups return retained async def run_async( self, *, executor: AttackExecutor | None = None, return_partial_on_failure: bool = True, **attack_params: Any, ) -> AttackExecutorResult[AttackResult]: """ Execute the atomic attack against all seed groups. This method uses ``AttackExecutor`` to run the configured attack against all seed groups. Concurrency is owned by the executor: pass a shared ``AttackExecutor`` instance to share a single budget across multiple atomic attacks (this is how ``Scenario`` parallelizes them). Source identity and adaptation metadata are passed to the executor before completed or failed results are persisted. When return_partial_on_failure=True (default), this method will return an AttackExecutorResult containing both completed results and incomplete objectives (those that didn't finish execution due to exceptions). This allows scenarios to save progress and retry only the incomplete objectives. Note: "completed" means the execution finished, not that the attack objective was achieved. "incomplete" means execution didn't finish (threw an exception). Args: executor (AttackExecutor | None): Optional ``AttackExecutor`` to run the attack with. When provided, its concurrency budget is used and is shared with anything else holding a reference to it. When ``None``, a fresh ``AttackExecutor(max_concurrency=1)`` is created for this call. return_partial_on_failure (bool): If True, returns partial results even when some objectives don't complete execution. If False, raises an exception on any execution failure. Defaults to True. **attack_params: Execution inputs overriding constructor-supplied defaults for this call. Memory labels merge with existing labels; call-time values win on shared keys. Returns: AttackExecutorResult[AttackResult]: Result containing completed attack results and incomplete objectives (those that didn't finish execution). Raises: ValueError: If inputs replace owned executor arguments, or execution fails completely and return_partial_on_failure=False. """ if executor is None: executor = AttackExecutor(max_concurrency=1) execution_params = {**self._attack_execute_params, **attack_params} memory_labels = {**self._memory_labels, **(execution_params.pop("memory_labels", None) or {})} reserved = execution_params.keys() & { "attack", "seed_groups", "adversarial_chat", "objective_scorer", "return_partial_on_failure", "attribution", "attributions", "atomic_attack_identifier", "atomic_attack_identifiers", "result_metadata", } if reserved: raise ValueError(f"AtomicAttack owns these executor arguments: {sorted(reserved)}") logger.info( f"Starting atomic attack execution with {len(self._seed_groups)} seed groups " f"(executor max_concurrency={getattr(executor, '_max_concurrency', 'unknown')})" ) try: technique = self._attack_technique technique_identifier = technique.get_identifier() technique_eval_hash = self.technique_eval_hash execution_seed_groups = [ self._prepared_seed_groups[group.logical_id].seed_group.model_copy(deep=True) for group in self._seed_groups ] result_identifiers: list[AtomicAttackIdentifier] = [] result_metadata: list[dict[str, Any]] = [] for source_group in self._seed_groups: identifier = AtomicAttackIdentifier.build( technique_identifier=technique_identifier, seed_group=source_group, ) result_identifiers.append( AtomicAttackIdentifier.from_component_identifier(identifier.with_eval_hash(technique_eval_hash)) ) adaptation = self._prepared_seed_groups[source_group.logical_id].adaptation result_metadata.append( { "seed_group_adaptation": { "source_seed_group_id": source_group.logical_id, "technique_eval_hash": technique_eval_hash, "adaptation": adaptation, } } if adaptation is not None else {} ) attributions: list[AttackResultAttribution] | None = None if self._scenario_result_id is not None: attributions = [ AttackResultAttribution( parent_id=self._scenario_result_id, parent_collection=self.atomic_attack_name, parent_eval_hash=technique_eval_hash, seed_group_id=seed_group.logical_id, ) for seed_group in self._seed_groups ] untyped_results = await executor.execute_attack_from_seed_groups_async( attack=technique.attack, seed_groups=execution_seed_groups, adversarial_chat=self._adversarial_chat, objective_scorer=self._objective_scorer, memory_labels=memory_labels, return_partial_on_failure=return_partial_on_failure, attributions=attributions, atomic_attack_identifiers=result_identifiers, result_metadata=result_metadata, **execution_params, ) completed_results: list[AttackResult] = [] for result in untyped_results.completed_results: if not isinstance(result, AttackResult): raise ValueError(f"Attack returned unsupported result type: {type(result).__name__}") completed_results.append(result) results = AttackExecutorResult[AttackResult]( completed_results=completed_results, incomplete_objectives=untyped_results.incomplete_objectives, input_indices=untyped_results.input_indices, incomplete_result_ids=untyped_results.incomplete_result_ids, ) # Log completion status if results.has_incomplete: logger.warning( f"Atomic attack execution completed with {len(results.completed_results)} completed " f"and {len(results.incomplete_objectives)} incomplete objectives" ) else: logger.info( f"Atomic attack execution completed successfully with {len(results.completed_results)} results" ) return results except Exception as e: logger.error(f"Atomic attack '{self.atomic_attack_name}' execution failed: {str(e)}") raise ValueError(f"Failed to execute atomic attack '{self.atomic_attack_name}': {str(e)}") from e