File size: 2,605 Bytes
ec23564 3b452a0 ec23564 3b452a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | from __future__ import annotations
import json
from pathlib import Path
from typing import Any, Dict, Optional
from .dataset_loader import DatasetLoader
from .injection_engine import InjectionEngine
from .policy_engine import PolicyEngine
from .scenario_types import HiddenEpisodeState
class ScenarioGenerator:
def __init__(
self,
dataset_loader: DatasetLoader,
injection_engine: InjectionEngine,
policy_engine: PolicyEngine,
config_dir: Optional[Path] = None,
) -> None:
self._dataset_loader = dataset_loader
self._injection_engine = injection_engine
self._policy_engine = policy_engine
self._config_dir = config_dir or Path(__file__).resolve().parent / "data" / "config"
self._attacker_profiles = self._load_attacker_profiles()
def generate_episode(self, difficulty: str, seed: Optional[int] = None) -> HiddenEpisodeState:
scenario = self._dataset_loader.get_scenario(difficulty=difficulty, seed=seed)
return self._build_hidden_state(scenario)
def generate_episode_by_id(self, difficulty: str, scenario_id: str) -> HiddenEpisodeState:
scenario = self._dataset_loader.get_scenario_by_id(difficulty=difficulty, scenario_id=scenario_id)
return self._build_hidden_state(scenario)
def _build_hidden_state(self, scenario) -> HiddenEpisodeState:
attack_metadata = self._injection_engine.build_attack_metadata(scenario)
attacker_profile = self._attacker_profiles.get("default_attacker", {})
attack_metadata["attacker_profile"] = attacker_profile
activation = dict(scenario.metadata.get("activation", {}))
attack_metadata["activation"] = activation
attack_active = activation.get("mode", "immediate") == "immediate"
attack_triggered_step = 0 if attack_active else None
attack_trigger_reason = "initial_immediate_attack" if attack_active else None
return HiddenEpisodeState(
scenario=scenario,
policy=self._policy_engine.load_policy_rules(),
attack_metadata=attack_metadata,
attack_active=attack_active,
attack_triggered_step=attack_triggered_step,
attack_trigger_reason=attack_trigger_reason,
last_action_result=f"Scenario {scenario.scenario_id} loaded. {scenario.benign_task}",
)
def _load_attacker_profiles(self) -> Dict[str, Dict[str, Any]]:
path = self._config_dir / "attacker_profiles.json"
with path.open("r", encoding="utf-8-sig") as handle:
return json.load(handle)
|