From 545914c169b6746008cebe60ce0242ad1f34001a Mon Sep 17 00:00:00 2001 From: elprofesoriqo Date: Sun, 20 Sep 2026 15:15:01 +0200 Subject: [PATCH 1/8] =?UTF-8?q?=EF=BB=BFSplit=20the=20gym=20env=20into=20E?= =?UTF-8?q?nvConfig,=20tick/reset,=20and=20registered=20ids.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep PettingZoo parallel and Gymnasium Blue as the public make() surface without stuffing trainers into this tree. Co-authored-by: Cursor --- netforge_rl/__init__.py | 24 +- netforge_rl/environment/__init__.py | 21 +- netforge_rl/environment/base_env.py | 23 +- netforge_rl/environment/config/__init__.py | 27 ++ netforge_rl/environment/config/agents.py | 18 + netforge_rl/environment/config/difficulty.py | 60 +++ netforge_rl/environment/config/schema.py | 133 +++++ netforge_rl/environment/constants.py | 15 + netforge_rl/environment/curriculum.py | 30 +- netforge_rl/environment/events.py | 31 ++ netforge_rl/environment/graph_wrapper.py | 12 +- netforge_rl/environment/gym_env.py | 14 +- netforge_rl/environment/metrics.py | 26 +- netforge_rl/environment/observations.py | 21 +- netforge_rl/environment/parallel_env.py | 485 +++---------------- netforge_rl/environment/presets.py | 86 +--- netforge_rl/environment/registry.py | 78 +++ netforge_rl/environment/reset.py | 62 +++ netforge_rl/environment/spaces.py | 50 +- netforge_rl/environment/tick.py | 311 ++++++++++++ tests/environment/test_env_ids.py | 50 ++ 21 files changed, 1032 insertions(+), 545 deletions(-) create mode 100644 netforge_rl/environment/config/__init__.py create mode 100644 netforge_rl/environment/config/agents.py create mode 100644 netforge_rl/environment/config/difficulty.py create mode 100644 netforge_rl/environment/config/schema.py create mode 100644 netforge_rl/environment/events.py create mode 100644 netforge_rl/environment/registry.py create mode 100644 netforge_rl/environment/reset.py create mode 100644 netforge_rl/environment/tick.py create mode 100644 tests/environment/test_env_ids.py diff --git a/netforge_rl/__init__.py b/netforge_rl/__init__.py index aae6533..2dbde8c 100644 --- a/netforge_rl/__init__.py +++ b/netforge_rl/__init__.py @@ -1,14 +1,26 @@ -from .environment.parallel_env import NetForgeRLEnv -from .core.action import BaseAction, ActionEffect -from .core.state import GlobalNetworkState, Host, Subnet +from .core.action import ActionEffect, BaseAction from .core.observation import BaseObservation +from .core.state import GlobalNetworkState, Host, Subnet +from .environment.parallel_env import NetForgeRLEnv + +__version__ = '4.0.0' __all__ = [ - 'NetForgeRLEnv', - 'BaseAction', 'ActionEffect', + 'BaseAction', + 'BaseObservation', 'GlobalNetworkState', 'Host', + 'NetForgeRLEnv', 'Subnet', - 'BaseObservation', + '__version__', ] + + +def _register_third_party_envs() -> None: + from netforge_rl.environment.registry import register_envs + + register_envs() + + +_register_third_party_envs() diff --git a/netforge_rl/environment/__init__.py b/netforge_rl/environment/__init__.py index 0cb60ee..fc060ea 100644 --- a/netforge_rl/environment/__init__.py +++ b/netforge_rl/environment/__init__.py @@ -1,20 +1,19 @@ from .base_env import BaseNetForgeRLEnv +from .config import DIFFICULTIES, DIFFICULTY_PRESETS, EVAL_SEEDS, EnvConfig from .parallel_env import NetForgeRLEnv -from .presets import ( - DIFFICULTY_PRESETS, - EVAL_SEEDS, - make_config, - make_env, -) +from .presets import make_config, make_env __all__ = [ - 'BaseNetForgeRLEnv', - 'NetForgeRLEnv', + 'DIFFICULTIES', 'DIFFICULTY_PRESETS', 'EVAL_SEEDS', + 'BaseNetForgeRLEnv', + 'EnvConfig', + 'NetForgeRLEnv', + 'NetForgeSingleAgentEnv', 'make_config', 'make_env', - 'NetForgeSingleAgentEnv', + 'register_envs', ] @@ -23,4 +22,8 @@ def __getattr__(name): from .gym_env import NetForgeSingleAgentEnv return NetForgeSingleAgentEnv + if name == 'register_envs': + from .registry import register_envs + + return register_envs raise AttributeError(name) diff --git a/netforge_rl/environment/base_env.py b/netforge_rl/environment/base_env.py index 16c9c43..eb53101 100644 --- a/netforge_rl/environment/base_env.py +++ b/netforge_rl/environment/base_env.py @@ -1,33 +1,30 @@ import abc + import numpy as np from pettingzoo import ParallelEnv -from typing import Dict, Tuple, Any class BaseNetForgeRLEnv(ParallelEnv, abc.ABC): - """Abstract Base Class for all Continuous-Time MARL environments.""" + """Abstract PettingZoo parallel env.""" @abc.abstractmethod def __init__(self, scenario_config: dict): - """Initialize scenario and network state.""" pass @abc.abstractmethod def reset( self, seed=None, options=None - ) -> Tuple[Dict[str, np.ndarray], Dict[str, dict]]: - """Reset environment and return initial observations.""" + ) -> tuple[dict[str, np.ndarray], dict[str, dict]]: pass @abc.abstractmethod def step( - self, agent_actions: Dict[str, Any] - ) -> Tuple[ - Dict[str, np.ndarray], - Dict[str, float], - Dict[str, bool], - Dict[str, bool], - Dict[str, dict], + self, agent_actions: dict[str, object] + ) -> tuple[ + dict[str, np.ndarray], + dict[str, float], + dict[str, bool], + dict[str, bool], + dict[str, dict], ]: - """Core physics loop. Returns (obs, rewards, term, trunc, infos).""" pass diff --git a/netforge_rl/environment/config/__init__.py b/netforge_rl/environment/config/__init__.py new file mode 100644 index 0000000..5dad483 --- /dev/null +++ b/netforge_rl/environment/config/__init__.py @@ -0,0 +1,27 @@ +from netforge_rl.environment.config.agents import AGENTS, POSSIBLE_AGENTS, AgentIds +from netforge_rl.environment.config.difficulty import ( + DIFFICULTIES, + DIFFICULTY_PRESETS, + EVAL_SEEDS, + DifficultyPreset, +) +from netforge_rl.environment.config.schema import ( + EnvConfig, + ResourceBudgets, + TimeMode, + TopologyRates, +) + +__all__ = [ + 'AGENTS', + 'DIFFICULTIES', + 'DIFFICULTY_PRESETS', + 'EVAL_SEEDS', + 'POSSIBLE_AGENTS', + 'AgentIds', + 'DifficultyPreset', + 'EnvConfig', + 'ResourceBudgets', + 'TimeMode', + 'TopologyRates', +] diff --git a/netforge_rl/environment/config/agents.py b/netforge_rl/environment/config/agents.py new file mode 100644 index 0000000..a06d49c --- /dev/null +++ b/netforge_rl/environment/config/agents.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AgentIds: + red: str = 'red_operator' + blue_dmz: str = 'blue_dmz' + blue_internal: str = 'blue_internal' + blue_restricted: str = 'blue_restricted' + + def all(self) -> tuple[str, ...]: + return (self.red, self.blue_dmz, self.blue_internal, self.blue_restricted) + + +AGENTS = AgentIds() +POSSIBLE_AGENTS: tuple[str, ...] = AGENTS.all() diff --git a/netforge_rl/environment/config/difficulty.py b/netforge_rl/environment/config/difficulty.py new file mode 100644 index 0000000..339bc49 --- /dev/null +++ b/netforge_rl/environment/config/difficulty.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from dataclasses import dataclass, replace + +from netforge_rl.environment.config.schema import ( + DifficultyName, + EnvConfig, + TopologyRates, +) + +EVAL_SEEDS: tuple[int, ...] = tuple(range(9001, 9021)) + + +@dataclass(frozen=True) +class DifficultyPreset: + name: DifficultyName + max_active_hosts: int + dhcp_interval: int + topology: TopologyRates + log_latency: int + max_ticks: int = 200 + + def apply(self, cfg: EnvConfig) -> EnvConfig: + return replace( + cfg, + max_active_hosts=self.max_active_hosts, + dhcp_interval=self.dhcp_interval, + topology=self.topology, + log_latency=self.log_latency, + max_ticks=self.max_ticks, + ) + + +DIFFICULTIES: dict[str, DifficultyPreset] = { + 'easy': DifficultyPreset( + name='easy', + max_active_hosts=6, + dhcp_interval=0, + topology=TopologyRates(), + log_latency=0, + ), + 'medium': DifficultyPreset( + name='medium', + max_active_hosts=15, + dhcp_interval=80, + topology=TopologyRates(churn=0.01), + log_latency=2, + ), + 'hard': DifficultyPreset( + name='hard', + max_active_hosts=100, + dhcp_interval=40, + topology=TopologyRates(churn=0.02, migration=0.01, arrival=0.005), + log_latency=4, + ), +} + +DIFFICULTY_PRESETS: dict[str, dict] = { + name: preset.apply(EnvConfig()).to_dict() for name, preset in DIFFICULTIES.items() +} diff --git a/netforge_rl/environment/config/schema.py b/netforge_rl/environment/config/schema.py new file mode 100644 index 0000000..16d8a4c --- /dev/null +++ b/netforge_rl/environment/config/schema.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field, replace +from typing import Any, Literal + +from netforge_rl.environment.config.agents import AgentIds + +TimeMode = Literal['event', 'fixed'] +DifficultyName = Literal['easy', 'medium', 'hard'] + +_FLAT_KEYS = frozenset( + { + 'scenario_type', + 'max_ticks', + 'max_active_hosts', + 'evaluation_mode', + 'topology_path', + 'topology_spec', + 'log_latency', + 'dhcp_interval', + 'record_siem', + 'record_trajectory', + 'pcap_obs', + 'time_mode', + 'docker_mode', + 'nlp_backend', + 'topology_churn_rate', + 'topology_migration_rate', + 'topology_arrival_rate', + } +) + + +@dataclass(frozen=True) +class TopologyRates: + churn: float = 0.0 + migration: float = 0.0 + arrival: float = 0.0 + + +@dataclass(frozen=True) +class ResourceBudgets: + energy: int = 50 + blue_funds: int = 10_000 + red_funds: int = 5_000 + compute: int = 1_000 + + +@dataclass(frozen=True) +class EnvConfig: + """Typed env knobs. ``NetForgeRLEnv`` still accepts a dict and parses this.""" + + scenario_type: str = 'ransomware' + max_ticks: int = 1000 + max_active_hosts: int | None = None + evaluation_mode: bool = False + topology_path: str | None = None + topology_spec: dict | None = None + log_latency: int = 0 + dhcp_interval: int = 40 + record_siem: bool = False + record_trajectory: bool = False + pcap_obs: bool = False + time_mode: TimeMode = 'event' + docker_mode: str = 'sim' + nlp_backend: str = 'tfidf' + topology: TopologyRates = field(default_factory=TopologyRates) + agents: AgentIds = field(default_factory=AgentIds) + budgets: ResourceBudgets = field(default_factory=ResourceBudgets) + extras: dict[str, Any] = field(default_factory=dict) + + def replace(self, **changes) -> EnvConfig: + return replace(self, **changes) + + def to_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + 'scenario_type': self.scenario_type, + 'max_ticks': self.max_ticks, + 'max_active_hosts': self.max_active_hosts, + 'evaluation_mode': self.evaluation_mode, + 'topology_path': self.topology_path, + 'topology_spec': self.topology_spec, + 'log_latency': self.log_latency, + 'dhcp_interval': self.dhcp_interval, + 'record_siem': self.record_siem, + 'record_trajectory': self.record_trajectory, + 'pcap_obs': self.pcap_obs, + 'time_mode': self.time_mode, + 'docker_mode': self.docker_mode, + 'nlp_backend': self.nlp_backend, + 'topology_churn_rate': self.topology.churn, + 'topology_migration_rate': self.topology.migration, + 'topology_arrival_rate': self.topology.arrival, + } + payload.update(self.extras) + return payload + + @classmethod + def from_mapping( + cls, data: Mapping[str, Any] | EnvConfig | None = None + ) -> EnvConfig: + if data is None: + return cls() + if isinstance(data, EnvConfig): + return data + raw = dict(data) + topo = TopologyRates( + churn=float(raw.pop('topology_churn_rate', 0.0) or 0.0), + migration=float(raw.pop('topology_migration_rate', 0.0) or 0.0), + arrival=float(raw.pop('topology_arrival_rate', 0.0) or 0.0), + ) + known: dict[str, Any] = {} + for key in ( + 'scenario_type', + 'max_ticks', + 'max_active_hosts', + 'evaluation_mode', + 'topology_path', + 'topology_spec', + 'log_latency', + 'dhcp_interval', + 'record_siem', + 'record_trajectory', + 'pcap_obs', + 'time_mode', + 'docker_mode', + 'nlp_backend', + ): + if key in raw: + known[key] = raw.pop(key) + extras = {k: v for k, v in raw.items() if k not in _FLAT_KEYS} + return cls(topology=topo, extras=extras, **known) diff --git a/netforge_rl/environment/constants.py b/netforge_rl/environment/constants.py index 7517747..e0aa5db 100644 --- a/netforge_rl/environment/constants.py +++ b/netforge_rl/environment/constants.py @@ -1,2 +1,17 @@ +"""Gym-wide sizes and limits. Observation / action spaces read these.""" + MAX_ACTION_DURATION = 50.0 PADDING_SUBNET = '169.254.0.0/16' + +N_ACTION_TYPES = 32 +N_HOST_SLOTS = 100 +ACTION_MASK_DIM = N_ACTION_TYPES + N_HOST_SLOTS +OBS_VECTOR_DIM = 256 +BLUE_COMM_DIM = 100 +ADJ_FLAT_DIM = N_HOST_SLOTS * N_HOST_SLOTS +SIEM_BUFFER_CAP = 64 +BLUE_INFLIGHT_CAP = 2 +SIEM_ENCODE_N = 8 + +# Backward-compatible alias used by tests / PettingZoo mask shape. +FLAT_ACTION_MASK = ACTION_MASK_DIM diff --git a/netforge_rl/environment/curriculum.py b/netforge_rl/environment/curriculum.py index e37129d..92e8c62 100644 --- a/netforge_rl/environment/curriculum.py +++ b/netforge_rl/environment/curriculum.py @@ -2,8 +2,8 @@ import random from collections import deque +from collections.abc import Callable from dataclasses import dataclass -from typing import Any, Deque, Dict, List, Optional, Tuple import numpy as np @@ -14,19 +14,17 @@ class PhaseConfig: name: str max_active_hosts: int - scenario_types: List[str] + scenario_types: list[str] reward_scale: float - dhcp_interval: int # ticks between DHCP reshuffles; 0 = disabled + dhcp_interval: int topology_churn_rate: float topology_migration_rate: float topology_arrival_rate: float - advance_threshold: Optional[ - float - ] # mean total reward to graduate; None = final phase - advance_window: int # episodes in the rolling window + advance_threshold: float | None + advance_window: int -PHASES: List[PhaseConfig] = [ +PHASES: list[PhaseConfig] = [ PhaseConfig( name='novice', max_active_hosts=5, @@ -66,9 +64,7 @@ class PhaseConfig: ] -def _build_env( - phase: PhaseConfig, base_cfg: dict, seed: Optional[int] -) -> NetForgeRLEnv: +def _build_env(phase: PhaseConfig, base_cfg: dict, seed: int | None) -> NetForgeRLEnv: scenario = random.Random(seed).choice(phase.scenario_types) cfg = { **base_cfg, @@ -88,16 +84,16 @@ class CurriculumWrapper: def __init__( self, - phases: List[PhaseConfig] = PHASES, - base_cfg: Optional[dict] = None, + phases: list[PhaseConfig] = PHASES, + base_cfg: dict | None = None, start_phase: int = 0, - on_phase_advance=None, + on_phase_advance: Callable | None = None, ): self.phases = phases self.base_cfg = base_cfg or {} self._phase_idx = start_phase self._on_phase_advance = on_phase_advance - self._window: Deque[float] = deque(maxlen=phases[start_phase].advance_window) + self._window: deque[float] = deque(maxlen=phases[start_phase].advance_window) self._episode_reward: float = 0.0 self._env: NetForgeRLEnv = _build_env(self.phase, self.base_cfg, seed=None) @@ -161,7 +157,7 @@ def observation_space(self, agent): def action_space(self, agent): return self._env.action_space(agent) - def reset(self, seed=None, options=None) -> Tuple[Dict, Dict]: + def reset(self, seed=None, options=None) -> tuple[dict, dict]: self._episode_reward = 0.0 self._env = _build_env(self.phase, self.base_cfg, seed=seed) obs, info = self._env.reset(seed=seed, options=options) @@ -169,7 +165,7 @@ def reset(self, seed=None, options=None) -> Tuple[Dict, Dict]: info[agent]['__curriculum__'] = self._curriculum_info() return obs, info - def step(self, actions: Dict[str, Any]) -> Tuple[Dict, Dict, Dict, Dict, Dict]: + def step(self, actions: dict[str, object]) -> tuple[dict, dict, dict, dict, dict]: obs, rewards, term, trunc, infos = self._env.step(actions) step_reward = sum(rewards.values()) * self.phase.reward_scale diff --git a/netforge_rl/environment/events.py b/netforge_rl/environment/events.py new file mode 100644 index 0000000..754b57d --- /dev/null +++ b/netforge_rl/environment/events.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from netforge_rl.core.action import ActionEffect, BaseAction + + +@dataclass +class PendingAction: + completion_tick: int + agent: str + action: BaseAction + effect: ActionEffect + target_ip: str | None + start_tick: int = 0 + + def __getitem__(self, key: str): + return getattr(self, key) + + +def as_pending(event) -> PendingAction: + if isinstance(event, PendingAction): + return event + return PendingAction( + completion_tick=int(event['completion_tick']), + agent=event['agent'], + action=event['action'], + effect=event['effect'], + target_ip=event.get('target_ip'), + start_tick=int(event.get('start_tick', 0)), + ) diff --git a/netforge_rl/environment/graph_wrapper.py b/netforge_rl/environment/graph_wrapper.py index a2509c8..19dec44 100644 --- a/netforge_rl/environment/graph_wrapper.py +++ b/netforge_rl/environment/graph_wrapper.py @@ -1,11 +1,14 @@ from pettingzoo.utils.wrappers import BaseParallelWrapper -from netforge_rl.core.graph_obs import build_graph_observation +from netforge_rl.core.graph_obs import build_graph_observation, build_oracle_graph class GraphObservationWrapper(BaseParallelWrapper): - """Injects a graph view into each agent's info['graph'] without changing the obs - dict. Red agents get a fog-of-war graph; blue agents see all hosts.""" + """``info['graph']``: SIEM belief (Blue) / fog (Red). Oracle only if ``oracle=True``.""" + + def __init__(self, env, oracle: bool = False): + super().__init__(env) + self.oracle = oracle def reset(self, seed=None, options=None): obs, infos = super().reset(seed=seed, options=options) @@ -19,6 +22,9 @@ def _inject(self, infos): gs = getattr(self.env.unwrapped, 'global_state', None) if gs is None: return infos + oracle_graph = build_oracle_graph(gs) if self.oracle else None for agent in infos: infos[agent]['graph'] = build_graph_observation(gs, agent_id=agent) + if oracle_graph is not None: + infos[agent]['oracle_graph'] = oracle_graph return infos diff --git a/netforge_rl/environment/gym_env.py b/netforge_rl/environment/gym_env.py index 1b0e9c3..7b0d0af 100644 --- a/netforge_rl/environment/gym_env.py +++ b/netforge_rl/environment/gym_env.py @@ -1,20 +1,20 @@ from __future__ import annotations import random as _random -from typing import Optional +from typing import ClassVar import gymnasium as gym import numpy as np from netforge_rl.baselines.policies import RandomPolicy +from netforge_rl.environment.constants import OBS_VECTOR_DIM from netforge_rl.environment.parallel_env import NetForgeRLEnv from netforge_rl.nlp.log_encoder import EMBEDDING_DIM -_FLAT_DIM = 256 + EMBEDDING_DIM +_FLAT_DIM = OBS_VECTOR_DIM + EMBEDDING_DIM def _reseed(policy, seed: int) -> None: - """Reseed a policy's RNG so opponents replay identically under a fixed seed.""" rng = getattr(policy, '_rng', None) if isinstance(rng, _random.Random): policy._rng = _random.Random(seed) @@ -23,17 +23,17 @@ def _reseed(policy, seed: int) -> None: class NetForgeSingleAgentEnv(gym.Env): - """Gymnasium single-agent facade: one RL-controlled agent vs scripted opponents.""" + """One RL agent vs scripted opponents.""" - metadata = {'render_modes': []} + metadata: ClassVar[dict] = {'render_modes': []} def __init__( self, scenario_type: str = 'ransomware', controlled_agent: str = 'blue_dmz', - opponents: Optional[dict] = None, + opponents: dict | None = None, max_ticks: int = 200, - config: Optional[dict] = None, + config: dict | None = None, ): super().__init__() cfg = dict(config or {}) diff --git a/netforge_rl/environment/metrics.py b/netforge_rl/environment/metrics.py index 122d1d4..a07e4d8 100644 --- a/netforge_rl/environment/metrics.py +++ b/netforge_rl/environment/metrics.py @@ -1,13 +1,15 @@ +from __future__ import annotations + from netforge_rl.actions.attack_map import ALL_TECHNIQUE_IDS, technique_for from netforge_rl.core.commands import iter_host_deltas from netforge_rl.environment.constants import PADDING_SUBNET class EpisodeMetricsMixin: - """Episode security metrics and per-agent info extraction for NetForgeRLEnv.""" + """Episode metrics and per-agent info.""" def _active_hosts(self): - """Real hosts, excluding inactive padding used only for fixed observation shape.""" + """Hosts excluding padding.""" return [ h for h in self.global_state.all_hosts.values() @@ -41,6 +43,11 @@ def _update_episode_metrics(self, resolved_effects): self.episode_metrics['isolation_times'].setdefault( ip, self.current_tick ) + host = self.global_state.all_hosts.get(ip) + if host is not None and host.compromised_by == 'None': + self.episode_metrics['false_positives'] = ( + self.episode_metrics.get('false_positives', 0) + 1 + ) active = self._active_hosts() total = max(len(active), 1) @@ -51,7 +58,7 @@ def _update_episode_metrics(self, resolved_effects): self.episode_metrics['steps_count'] += 1 def _extract_agent_infos( - self, observations: dict, resolved_effects: dict, rewards: dict = None + self, observations: dict, resolved_effects: dict, rewards: dict | None = None ) -> dict: infos = {} for agent in observations: @@ -120,6 +127,19 @@ def _add_host_metrics(self, info: dict, agent: str) -> None: techniques = self.episode_metrics['attack_techniques'] info['attack_techniques'] = sorted(techniques) info['attack_coverage'] = float(len(techniques) / len(ALL_TECHNIQUE_IDS)) + info['false_positives_total'] = float( + self.episode_metrics.get('false_positives', 0) + ) + kinetic = any( + getattr(h, 'system_integrity', 'clean') == 'kinetic_destruction' + for h in active + ) + info['catastrophic_failure'] = float(kinetic) + total = max(len(active), 1) + info['security'] = float(1.0 - info['compromised_hosts'] / total) + info['disruption'] = float(info['isolated_hosts'] / total) + sla = info['SLA_Uptime_Percentage'] + info['mission_success'] = float((not kinetic) and sla >= 0.5) def _mean_time_to_containment(self) -> float: infections = self.episode_metrics['infection_times'] diff --git a/netforge_rl/environment/observations.py b/netforge_rl/environment/observations.py index 3a7d5b4..beb85d9 100644 --- a/netforge_rl/environment/observations.py +++ b/netforge_rl/environment/observations.py @@ -1,14 +1,15 @@ import numpy as np +from netforge_rl.environment.constants import BLUE_COMM_DIM, N_HOST_SLOTS, SIEM_ENCODE_N + class ObservationMixin: """Blue shared-comm channel and per-agent adjacency masking for NetForgeRLEnv.""" def _build_blue_comm(self) -> np.ndarray: - """100-dim shared situational awareness from SIEM incidents and isolation state.""" - comm = np.zeros(100, dtype=np.float32) + comm = np.zeros(BLUE_COMM_DIM, dtype=np.float32) ordered = sorted(self.global_state.all_hosts.keys()) - ip_to_idx = {ip: i for i, ip in enumerate(ordered[:100])} + ip_to_idx = {ip: i for i, ip in enumerate(ordered[:N_HOST_SLOTS])} for entry, _ in self.global_state.siem_log_buffer: if not isinstance(entry, str) or '[INCIDENT]' not in entry: @@ -20,11 +21,11 @@ def _build_blue_comm(self) -> np.ndarray: comm[ip_to_idx[ip]] = 1.0 break - recent_subnets = {sub for _, sub in self.global_state.siem_log_buffer[-8:]} - for i, ip in enumerate(ordered[:100]): + recent_subnets = { + sub for _, sub in self.global_state.siem_log_buffer[-SIEM_ENCODE_N:] + } + for i, ip in enumerate(ordered[:N_HOST_SLOTS]): host = self.global_state.all_hosts[ip] - if host.privilege in ('User', 'Root') and comm[i] < 0.75: - comm[i] = max(comm[i], 0.75) if host.status == 'isolated': comm[i] = max(comm[i], 0.5) if host.subnet_cidr in recent_subnets and comm[i] < 0.25: @@ -32,13 +33,13 @@ def _build_blue_comm(self) -> np.ndarray: return comm def _get_adj_matrix_for(self, agent: str) -> np.ndarray: - """Adjacency matrix masked to the agent's discovered knowledge (fog of war).""" + """Adjacency masked to discovered hosts.""" full_adj = self.global_state.get_adjacency_matrix() if 'blue' in agent.lower(): return full_adj known_ips = self.global_state.agent_knowledge.get(agent, set()) - sorted_ips = sorted(self.global_state.all_hosts.keys())[:100] - mask = np.zeros(100, dtype=np.float32) + sorted_ips = sorted(self.global_state.all_hosts.keys())[:N_HOST_SLOTS] + mask = np.zeros(N_HOST_SLOTS, dtype=np.float32) for i, ip in enumerate(sorted_ips): if ip in known_ips: mask[i] = 1.0 diff --git a/netforge_rl/environment/parallel_env.py b/netforge_rl/environment/parallel_env.py index a22712f..4403879 100644 --- a/netforge_rl/environment/parallel_env.py +++ b/netforge_rl/environment/parallel_env.py @@ -1,102 +1,107 @@ -from typing import Dict, Tuple +import random +from typing import ClassVar + import numpy as np -from netforge_rl.scenarios import get_scenario_class -from netforge_rl.core.action import BaseAction, ActionEffect -from netforge_rl.core.observation import BaseObservation -from netforge_rl.core.registry import action_registry, team_of import netforge_rl.actions # noqa: F401 — registers all action decorators - -from netforge_rl.core.physics import ConflictResolutionEngine -from netforge_rl.environment.base_env import BaseNetForgeRLEnv -from netforge_rl.topologies.network_generator import NetworkGenerator -from netforge_rl.topologies.dynamic_topology import TopologyEventEngine -from netforge_rl.scenarios.ot_physics import PLCPhysicsEngine from netforge_rl.agents.green_agent import GreenAgent -from netforge_rl.docker_bridge.bridge import DockerBridge -from netforge_rl.siem.siem_logger import SIEMLogger -from netforge_rl.siem.pcap_synthesizer import ( - PcapSynthesizer, - N_PACKETS, - PACKET_DIM, - NODE_DIM, -) -from netforge_rl.siem.correlator import SIEMCorrelator -from netforge_rl.nlp.log_encoder import LogEncoder, EMBEDDING_DIM -from netforge_rl.siem.event_templates import sysmon_1, seed_events +from netforge_rl.core.action import ActionEffect from netforge_rl.core.functional import from_global_state +from netforge_rl.core.physics import ConflictResolutionEngine +from netforge_rl.core.registry import action_registry, team_of +from netforge_rl.docker_bridge.bridge import DockerBridge +from netforge_rl.environment.base_env import BaseNetForgeRLEnv +from netforge_rl.environment.config import EnvConfig from netforge_rl.environment.constants import ( # noqa: F401 — re-exported + ACTION_MASK_DIM, MAX_ACTION_DURATION, + N_ACTION_TYPES, + N_HOST_SLOTS, PADDING_SUBNET, ) from netforge_rl.environment.metrics import EpisodeMetricsMixin from netforge_rl.environment.observations import ObservationMixin +from netforge_rl.environment.reset import ( + empty_episode_metrics, + initial_observations, + seed_agent_budgets, +) from netforge_rl.environment.spaces import build_action_spaces, build_observation_spaces -import random +from netforge_rl.environment.tick import run_step +from netforge_rl.nlp.log_encoder import LogEncoder +from netforge_rl.scenarios import get_scenario_class +from netforge_rl.scenarios.ot_physics import PLCPhysicsEngine +from netforge_rl.siem.correlator import SIEMCorrelator +from netforge_rl.siem.event_templates import seed_events +from netforge_rl.siem.pcap_synthesizer import PcapSynthesizer +from netforge_rl.siem.siem_logger import SIEMLogger +from netforge_rl.topologies.dynamic_topology import TopologyEventEngine +from netforge_rl.topologies.network_generator import NetworkGenerator class NetForgeRLEnv(BaseNetForgeRLEnv, EpisodeMetricsMixin, ObservationMixin): - """PettingZoo-style MARL environment for the NetForge cybersecurity sim.""" + """PettingZoo parallel env.""" - metadata = {'render_modes': ['ansi', 'rgb_array'], 'name': 'netforge_rl_v3'} + metadata: ClassVar[dict] = { + 'render_modes': ['ansi', 'rgb_array'], + 'name': 'netforge_rl_v4', + } - def __init__(self, scenario_config: dict): - cfg = scenario_config or {} + def __init__(self, scenario_config: dict | EnvConfig | None = None): + from netforge_rl.scenarios.yaml_dsl import expand_scenario_config + + raw = ( + scenario_config.to_dict() + if isinstance(scenario_config, EnvConfig) + else dict(scenario_config or {}) + ) + self.config = EnvConfig.from_mapping(expand_scenario_config(raw)) + cfg = self.config self.network_generator = NetworkGenerator( - config_path=cfg.get('topology_path'), - max_active_hosts=cfg.get('max_active_hosts'), - evaluation_mode=cfg.get('evaluation_mode', False), + config_path=cfg.topology_path, + max_active_hosts=cfg.max_active_hosts, + evaluation_mode=cfg.evaluation_mode, + topology_spec=cfg.topology_spec, ) - self.log_latency = cfg.get('log_latency', 0) - self.dhcp_interval = cfg.get('dhcp_interval', 40) - self.record_siem = cfg.get('record_siem', False) + self.log_latency = cfg.log_latency + self.dhcp_interval = cfg.dhcp_interval + self.record_siem = cfg.record_siem + self.pcap_obs = cfg.pcap_obs + self.max_ticks = cfg.max_ticks + self.time_mode = cfg.time_mode self.green_agent = GreenAgent() - self.possible_agents = [ - 'red_operator', - 'blue_dmz', - 'blue_internal', - 'blue_restricted', - ] + self.possible_agents = list(cfg.agents.all()) self.agents = self.possible_agents[:] - scenario_cls = get_scenario_class(cfg.get('scenario_type', 'ransomware')) + scenario_cls = get_scenario_class(cfg.scenario_type) self.scenario = scenario_cls(self.agents) - self.global_state = self.network_generator.generate() self.resolution_engine = ConflictResolutionEngine() - - self.docker_bridge = DockerBridge(mode=cfg.get('docker_mode', 'sim')) + self.docker_bridge = DockerBridge(mode=cfg.docker_mode) self.global_state.docker_bridge = self.docker_bridge - self.siem_logger = SIEMLogger() - self.log_encoder = LogEncoder(backend=cfg.get('nlp_backend', 'tfidf')) + self.log_encoder = LogEncoder(backend=cfg.nlp_backend) self.topology_engine = TopologyEventEngine( - churn_rate=cfg.get('topology_churn_rate', 0.0), - migration_rate=cfg.get('topology_migration_rate', 0.0), - arrival_rate=cfg.get('topology_arrival_rate', 0.0), + churn_rate=cfg.topology.churn, + migration_rate=cfg.topology.migration, + arrival_rate=cfg.topology.arrival, ) self.physics_engine = PLCPhysicsEngine() self.correlator = SIEMCorrelator() - self.pcap_obs = cfg.get('pcap_obs', False) - self.pcap_synthesizer = PcapSynthesizer() if self.pcap_obs else None - if cfg.get('record_trajectory', False): + self.pcap_synthesizer = PcapSynthesizer() if cfg.pcap_obs else None + if cfg.record_trajectory: from netforge_rl.render.trajectory import TrajectoryRecorder self.trajectory_recorder = TrajectoryRecorder() else: self.trajectory_recorder = None - self.observation_spaces = build_observation_spaces( - self.possible_agents, self.pcap_obs + self.possible_agents, cfg.pcap_obs ) self.action_spaces = build_action_spaces(self.possible_agents) - self.max_ticks = cfg.get('max_ticks', 1000) self.current_tick = 0 self.event_queue = [] - def reset( - self, seed=None, options=None - ) -> Tuple[Dict[str, np.ndarray], Dict[str, dict]]: - """Reset the environment.""" + def reset(self, seed=None, options=None) -> tuple[dict, dict]: self.np_random = np.random.default_rng(seed) self._py_random = random.Random(seed) seed_events(seed) @@ -112,43 +117,9 @@ def reset( self._cached_action_masks = { agent: self.action_mask(agent) for agent in self.agents } - self.global_state.agent_energy = {agent: 50 for agent in self.agents} - self.global_state.agent_funds = { - agent: 10000 if 'blue' in agent else 5000 for agent in self.agents - } - self.global_state.agent_compute = {agent: 1000 for agent in self.agents} - self.global_state.business_downtime_score = 0.0 - # SIEM log buffer and research metrics - self.global_state.siem_log_buffer = [] - self.episode_metrics = { - 'infection_times': {}, # IP -> tick - 'detection_times': {}, # IP -> tick (first SIEM alert) - 'isolation_times': {}, # IP -> tick - 'exfiltrated_data': 0.0, - 'sla_uptime_sum': 0.0, - 'steps_count': 0, - 'deception_hits': 0, # red actions that struck a decoy/honeytoken - 'red_actions': 0, # resolved red actions, for efficacy ratio - 'attack_techniques': set(), # MITRE ATT&CK technique ids exercised - } - - observations = {} - for agent_id in self.agents: - obs = BaseObservation(agent_id) - obs.update_from_state(self.global_state, []) - agent_obs = { - 'obs': obs.to_numpy(max_size=256), - 'action_mask': self._cached_action_masks[agent_id], - 'siem_embedding': np.zeros(EMBEDDING_DIM, dtype=np.float32), - 'adj_matrix': self.global_state.get_adjacency_matrix().flatten(), - 'delta_t': np.zeros(1, dtype=np.float32), - } - if 'blue' in agent_id: - agent_obs['blue_comm'] = np.zeros(100, dtype=np.float32) - if self.pcap_obs: - agent_obs['pcap'] = np.zeros((N_PACKETS, PACKET_DIM), dtype=np.float32) - agent_obs['node_features'] = np.zeros((100, NODE_DIM), dtype=np.float32) - observations[agent_id] = agent_obs + seed_agent_budgets(self) + self.episode_metrics = empty_episode_metrics() + observations = initial_observations(self) self.current_tick = 0 self.event_queue = [] self.topology_engine.reset(seed=seed) @@ -160,7 +131,6 @@ def reset( ) if self.pcap_synthesizer: self.pcap_synthesizer.reset(seed=seed) - return observations, {agent: {} for agent in self.agents} def observation_space(self, agent): @@ -169,319 +139,22 @@ def observation_space(self, agent): def action_space(self, agent): return self.action_spaces[agent] - def action_mask(self, agent: str) -> np.ndarray: - """Generate a binary action mask of shape (132,) = (32 + 100).""" - mask = np.zeros(132, dtype=np.int8) + def action_mask(self, agent: str): + mask = np.zeros(ACTION_MASK_DIM, dtype=np.int8) for action_id in action_registry._actions.get(team_of(agent), {}): - if action_id < 32: + if action_id < N_ACTION_TYPES: mask[action_id] = 1 ordered = sorted(self.global_state.all_hosts.keys()) - for i, ip in enumerate(ordered[:100]): + for i, ip in enumerate(ordered[:N_HOST_SLOTS]): host = self.global_state.all_hosts.get(ip) if host and host.status != 'isolated': - mask[32 + i] = 1 + mask[N_ACTION_TYPES + i] = 1 return mask - def step( - self, agent_actions: Dict[str, int] - ) -> Tuple[ - Dict[str, BaseObservation], - Dict[str, float], - Dict[str, bool], - Dict[str, bool], - Dict[str, dict], - ]: - """Process actions and advance simulation.""" - - per_agent_inflight: Dict[str, int] = {} - for event in self.event_queue: - per_agent_inflight[event['agent']] = ( - per_agent_inflight.get(event['agent'], 0) + 1 - ) - - for agent, action_int in agent_actions.items(): - if self.current_tick < self.global_state.agent_locked_until.get(agent, 0): - continue - - if isinstance(action_int, BaseAction): - action = action_int - else: - self.ordered_hosts = sorted(self.global_state.all_hosts.keys()) - action = action_registry.instantiate_action( - agent, action_int, self.ordered_hosts - ) - if action is None: - continue - - # Cap Blue agents at 2 in-flight actions. - if 'blue' in agent.lower(): - if per_agent_inflight.get(agent, 0) >= 2: - continue - per_agent_inflight[agent] = per_agent_inflight.get(agent, 0) + 1 - - if self.global_state.agent_energy.get(agent, 0) < action.cost: - continue - if not action.validate(self.global_state): - continue - - self.global_state.agent_energy[agent] -= action.cost - - eta = getattr(action, 'duration', 1) - completion_tick = self.current_tick + eta - effect = action.execute(self.global_state) - effect.action = action - - self.global_state.agent_locked_until[agent] = completion_tick - self.event_queue.append( - { - 'completion_tick': completion_tick, - 'agent': agent, - 'action': action, - 'effect': effect, - 'target_ip': getattr(action, 'target_ip', None), - 'start_tick': self.current_tick, - } - ) - - for event in list(self.event_queue): - if ( - type(event['action']).__name__ == 'IsolateHost' - and event['completion_tick'] <= self.current_tick - ): - target_to_isolate = event['target_ip'] - for red_event in list(self.event_queue): - if ( - 'red' in red_event['agent'].lower() - and red_event['target_ip'] == target_to_isolate - ): - if red_event in self.event_queue: - self.event_queue.remove(red_event) - self.global_state.agent_locked_until[red_event['agent']] = ( - self.current_tick - ) - - prev_tick = self.current_tick - if self.event_queue: - next_event_tick = min(e['completion_tick'] for e in self.event_queue) - self.current_tick = max(self.current_tick + 1, next_event_tick) - else: - self.current_tick += 1 - - delta_t = float(self.current_tick - prev_tick) - delta_t_norm = delta_t / MAX_ACTION_DURATION - - self.global_state.current_tick = self.current_tick - self.global_state.subnet_bandwidth.clear() - - noise_data = self.green_agent.generate_noise( - self.current_tick, self.global_state, rng=self._py_random - ) - for anomaly in noise_data.get('alerts', []): - self.siem_logger._push_to_buffer( - anomaly['data'], anomaly['subnet'], self.global_state - ) - - intended_effects = {} - action_metadata = {} - remaining_events = [] - for event in self.event_queue: - if self.current_tick >= event['completion_tick']: - agent = event['agent'] - intended_effects[agent] = event['effect'] - action_metadata[agent] = { - 'name': type(event['action']).__name__, - 'target_ip': event.get('target_ip'), - } - else: - remaining_events.append(event) - self.event_queue = remaining_events - - resolved_effects = self.resolution_engine.resolve(intended_effects) - self._apply_state_deltas(resolved_effects) - - self._update_episode_metrics(resolved_effects) - - for res_agent, res_effect in resolved_effects.items(): - meta = action_metadata.get(res_agent, {}) - self.siem_logger.log_action( - action_name=meta.get('name', 'UnknownAction'), - effect=res_effect, - global_state=self.global_state, - agent_id=res_agent, - target_ip=res_effect.observation_data.get('exploit'), - ) - - for res_agent, res_effect in resolved_effects.items(): - if 'red' not in res_agent or not res_effect.success: - continue - target_ip = res_effect.observation_data.get('exploit', 'unknown') - host = self.global_state.all_hosts.get(target_ip) - subnet = host.subnet_cidr if host else 'unknown' - - self.siem_logger._push_to_buffer( - sysmon_1( - res_agent, process='exploit_payload', rng=self.siem_logger._rng - ), - subnet, - self.global_state, - ) - if host and getattr(host, 'contains_honeytokens', False): - self.siem_logger._push_to_buffer( - { - 'signature': 'HONEYTOKEN_TRIGGERED', - 'target': target_ip, - 'agent': res_agent, - 'severity': 10, - }, - subnet, - self.global_state, - ) - - self.siem_logger.log_background_noise(self.global_state) - - if self.dhcp_interval > 0 and self.current_tick % self.dhcp_interval == 0: - self.global_state.reallocate_dhcp(rng=self._py_random) - valid_ips = set(self.global_state.all_hosts.keys()) - self.event_queue = [ - e - for e in self.event_queue - if e.get('target_ip') is None or e['target_ip'] in valid_ips - ] - self._cached_action_masks = { - agent: self.action_mask(agent) for agent in self.agents - } - - physics_alerts, physics_deltas = self.physics_engine.tick(self.global_state) - for delta_key, delta_val in physics_deltas: - self.global_state.apply_delta(delta_key, delta_val) - ot_subnet = '10.0.99.0/24' - for alert in physics_alerts: - self.siem_logger._push_to_buffer(alert, ot_subnet, self.global_state) - - for incident_log, incident_subnet in self.correlator.correlate( - self.global_state - ): - self.global_state.siem_log_buffer.append((incident_log, incident_subnet)) - if len(self.global_state.siem_log_buffer) > 64: - self.global_state.siem_log_buffer.pop(0) - - topo_events = self.topology_engine.tick(self.global_state) - if topo_events: - valid_ips = set(self.global_state.all_hosts.keys()) - self.event_queue = [ - e - for e in self.event_queue - if e.get('target_ip') is None or e['target_ip'] in valid_ips - ] - for ev in topo_events: - self.siem_logger._push_to_buffer( - { - 'signature': f'TOPOLOGY_{ev.kind.upper()}', - 'detail': ev.detail, - 'severity': 3, - }, - ev.detail.get('subnet', ev.detail.get('new_subnet', 'unknown')), - self.global_state, - ) - self._cached_action_masks = { - agent: self.action_mask(agent) for agent in self.agents - } - - observations = {} - rewards = {} - terminate = self.scenario.check_termination(self.global_state) - is_truncated = self.current_tick >= self.max_ticks - truncate = {agent: is_truncated for agent in self.agents} - - self.siem_logger.release(self.global_state) - - # Encode SIEM logs. - agent_siem_vecs = {} - for agent in self.agents: - if 'blue' in agent.lower(): - subnet_tag = agent.split('_')[1] if '_' in agent else 'dmz' - subset_logs = self.siem_logger.get_filtered_logs( - self.global_state, subnet_tag=subnet_tag, n=8 - ) - agent_siem_vecs[agent] = self.log_encoder.encode_buffer( - subset_logs, agg='mean' - ) - - pcap_snapshot = ( - self.pcap_synthesizer.synthesize( - self.global_state, self.current_tick, self.max_ticks - ) - if self.pcap_synthesizer - else None - ) - node_feat_snapshot = ( - self.pcap_synthesizer.node_features(self.global_state) - if self.pcap_synthesizer - else None - ) - - blue_comm = self._build_blue_comm() - - for agent in self.agents: - obs = BaseObservation(agent) - obs.update_from_state(self.global_state, resolved_effects) - - obs_array = obs.to_numpy(max_size=256) - - if 'blue' in agent.lower(): - agent_siem_vec = agent_siem_vecs.get( - agent, np.zeros(EMBEDDING_DIM, dtype=np.float32) - ) - else: - agent_siem_vec = np.zeros(EMBEDDING_DIM, dtype=np.float32) - - agent_obs = { - 'obs': obs_array, - 'action_mask': self._cached_action_masks[agent], - 'siem_embedding': agent_siem_vec, - 'adj_matrix': self._get_adj_matrix_for(agent).flatten(), - 'delta_t': np.array([delta_t_norm], dtype=np.float32), - } - if 'blue' in agent.lower(): - agent_obs['blue_comm'] = blue_comm - if self.pcap_obs: - agent_obs['pcap'] = pcap_snapshot - agent_obs['node_features'] = node_feat_snapshot - observations[agent] = agent_obs - agent_effect = resolved_effects.get(agent) - rewards[agent] = self.scenario.calculate_reward( - agent, self.global_state, agent_effect - ) - - if self.trajectory_recorder is not None: - for agent, effect in resolved_effects.items(): - meta = action_metadata.get(agent, {}) - self.trajectory_recorder.record_step( - tick=self.current_tick, - agent_id=agent, - action_name=meta.get('name', 'UnknownAction'), - target_ip=meta.get('target_ip'), - success=effect.success, - reward=float(rewards.get(agent, 0.0)), - ) - - self.agents = [ - agent - for agent in self.agents - if not terminate[agent] and not truncate[agent] - ] - - infos = self._extract_agent_infos(observations, resolved_effects, rewards) - - for agent in self.agents: - if agent in infos: - infos[agent]['delta_t'] = delta_t - infos[agent]['delta_t_norm'] = delta_t_norm - - return observations, rewards, terminate, truncate, infos + def step(self, agent_actions: dict): + return run_step(self, agent_actions) def render(self, mode: str = 'rgb_array'): - """Render the environment frame.""" if mode == 'ansi': return None if mode != 'rgb_array': @@ -491,17 +164,15 @@ def render(self, mode: str = 'rgb_array'): return render_rgb(snapshot_from_envstate(self.to_envstate())) def to_envstate(self): - """Return a frozen EnvState PyTree snapshot.""" return from_global_state(self.global_state, tuple(self.possible_agents)) - def _apply_state_deltas(self, effects: Dict[str, ActionEffect]): - """Apply state deltas to global_state.""" + def _apply_state_deltas(self, effects: dict[str, ActionEffect]): for effect in effects.values(): if not effect.success: continue if isinstance(effect.state_deltas, dict): - for delta_key, delta_val in effect.state_deltas.items(): - self.global_state.apply_delta(delta_key, delta_val) + for key, val in effect.state_deltas.items(): + self.global_state.apply_delta(key, val) elif isinstance(effect.state_deltas, list): - for delta_cmd in effect.state_deltas: - self.global_state.apply_delta(delta_cmd) + for cmd in effect.state_deltas: + self.global_state.apply_delta(cmd) diff --git a/netforge_rl/environment/presets.py b/netforge_rl/environment/presets.py index 1eb899c..776933c 100644 --- a/netforge_rl/environment/presets.py +++ b/netforge_rl/environment/presets.py @@ -1,62 +1,20 @@ from __future__ import annotations -from typing import Optional - -from netforge_rl.environment.parallel_env import NetForgeRLEnv - -# Held-out seeds for the generalization split; used with evaluation_mode=True. -EVAL_SEEDS: tuple[int, ...] = ( - 9001, - 9002, - 9003, - 9004, - 9005, - 9006, - 9007, - 9008, - 9009, - 9010, - 9011, - 9012, - 9013, - 9014, - 9015, - 9016, - 9017, - 9018, - 9019, - 9020, +from netforge_rl.environment.config import ( + DIFFICULTIES, + DIFFICULTY_PRESETS, + EVAL_SEEDS, + EnvConfig, ) +from netforge_rl.environment.parallel_env import NetForgeRLEnv -DIFFICULTY_PRESETS: dict[str, dict] = { - 'easy': { - 'max_active_hosts': 6, - 'dhcp_interval': 0, - 'topology_churn_rate': 0.0, - 'topology_migration_rate': 0.0, - 'topology_arrival_rate': 0.0, - 'log_latency': 0, - 'max_ticks': 200, - }, - 'medium': { - 'max_active_hosts': 15, - 'dhcp_interval': 80, - 'topology_churn_rate': 0.01, - 'topology_migration_rate': 0.0, - 'topology_arrival_rate': 0.0, - 'log_latency': 2, - 'max_ticks': 200, - }, - 'hard': { - 'max_active_hosts': 100, - 'dhcp_interval': 40, - 'topology_churn_rate': 0.02, - 'topology_migration_rate': 0.01, - 'topology_arrival_rate': 0.005, - 'log_latency': 4, - 'max_ticks': 200, - }, -} +__all__ = [ + 'DIFFICULTIES', + 'DIFFICULTY_PRESETS', + 'EVAL_SEEDS', + 'make_config', + 'make_env', +] def make_config( @@ -66,23 +24,23 @@ def make_config( **overrides, ) -> dict: """Env config for a difficulty tier; overrides win over the preset.""" - if difficulty not in DIFFICULTY_PRESETS: + if difficulty not in DIFFICULTIES: raise KeyError( - f'Unknown difficulty {difficulty!r}. ' - f'Choose from {sorted(DIFFICULTY_PRESETS)}.' + f'Unknown difficulty {difficulty!r}. Choose from {sorted(DIFFICULTIES)}.' ) - cfg = dict(DIFFICULTY_PRESETS[difficulty]) - cfg['scenario_type'] = scenario_type - cfg['evaluation_mode'] = evaluation - cfg.update(overrides) - return cfg + cfg = DIFFICULTIES[difficulty].apply( + EnvConfig(scenario_type=scenario_type, evaluation_mode=evaluation) + ) + if overrides: + cfg = EnvConfig.from_mapping({**cfg.to_dict(), **overrides}) + return cfg.to_dict() def make_env( difficulty: str = 'medium', scenario_type: str = 'ransomware', evaluation: bool = False, - seed: Optional[int] = None, + seed: int | None = None, **overrides, ) -> NetForgeRLEnv: """Build (and reset, if seed is given) an env at a difficulty tier.""" diff --git a/netforge_rl/environment/registry.py b/netforge_rl/environment/registry.py new file mode 100644 index 0000000..edbf532 --- /dev/null +++ b/netforge_rl/environment/registry.py @@ -0,0 +1,78 @@ +"""Gymnasium and PettingZoo registration — ``gym.make`` / ``pettingzoo.make``.""" + +from __future__ import annotations + +GYM_BLUE_ID = 'NetForge/Blue-v4' +GYM_IDS = (GYM_BLUE_ID, 'NetForge-v4') +PZ_RANSOMWARE_ID = 'netforge/ransomware-v4' +PZ_ARENA_ID = 'netforge/arena-v4' +PZ_IDS = (PZ_RANSOMWARE_ID, PZ_ARENA_ID) + + +def make_parallel_env(**kwargs): + from netforge_rl.environment.parallel_env import NetForgeRLEnv + + cfg = {'scenario_type': 'ransomware', 'max_ticks': 200, **kwargs} + cfg.pop('max_cycles', None) + return NetForgeRLEnv(cfg) + + +def make_aec_env(**kwargs): + from pettingzoo.utils.conversions import parallel_to_aec + + return parallel_to_aec(make_parallel_env(**kwargs)) + + +def make_gym_env(**kwargs): + from netforge_rl.environment.gym_env import NetForgeSingleAgentEnv + + return NetForgeSingleAgentEnv(**kwargs) + + +def register_gymnasium() -> tuple[str, ...]: + import gymnasium as gym + + registered = [] + for env_id in GYM_IDS: + if env_id not in gym.envs.registry: + gym.register( + id=env_id, + entry_point=make_gym_env, + max_episode_steps=200, + kwargs={'scenario_type': 'ransomware', 'controlled_agent': 'blue_dmz'}, + disable_env_checker=True, + ) + registered.append(env_id) + return tuple(registered) + + +def register_pettingzoo() -> tuple[str, ...]: + from pettingzoo import aec_registry, parallel_registry, register + + registered = [] + for env_id in PZ_IDS: + if env_id not in parallel_registry: + register( + 'parallel', + env_id, + make_parallel_env, + max_cycles=200, + kwargs={'scenario_type': 'ransomware', 'max_ticks': 200}, + ) + if env_id not in aec_registry: + register( + 'aec', + env_id, + make_aec_env, + max_cycles=200, + kwargs={'scenario_type': 'ransomware', 'max_ticks': 200}, + ) + registered.append(env_id) + return tuple(registered) + + +def register_envs() -> dict[str, tuple[str, ...]]: + return { + 'gymnasium': register_gymnasium(), + 'pettingzoo': register_pettingzoo(), + } diff --git a/netforge_rl/environment/reset.py b/netforge_rl/environment/reset.py new file mode 100644 index 0000000..d7cd56f --- /dev/null +++ b/netforge_rl/environment/reset.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import numpy as np + +from netforge_rl.core.observation import BaseObservation +from netforge_rl.environment.constants import ( + BLUE_COMM_DIM, + N_HOST_SLOTS, + OBS_VECTOR_DIM, +) +from netforge_rl.nlp.log_encoder import EMBEDDING_DIM +from netforge_rl.siem.pcap_synthesizer import N_PACKETS, NODE_DIM, PACKET_DIM + + +def empty_episode_metrics() -> dict: + return { + 'infection_times': {}, + 'detection_times': {}, + 'isolation_times': {}, + 'exfiltrated_data': 0.0, + 'sla_uptime_sum': 0.0, + 'steps_count': 0, + 'deception_hits': 0, + 'red_actions': 0, + 'attack_techniques': set(), + 'false_positives': 0, + } + + +def seed_agent_budgets(env) -> None: + budgets = env.config.budgets + env.global_state.agent_energy = {agent: budgets.energy for agent in env.agents} + env.global_state.agent_funds = { + agent: budgets.blue_funds if 'blue' in agent else budgets.red_funds + for agent in env.agents + } + env.global_state.agent_compute = {agent: budgets.compute for agent in env.agents} + env.global_state.business_downtime_score = 0.0 + env.global_state.siem_log_buffer = [] + + +def initial_observations(env) -> dict[str, dict]: + observations = {} + for agent_id in env.agents: + obs = BaseObservation(agent_id) + obs.update_from_state(env.global_state, []) + agent_obs = { + 'obs': obs.to_numpy(max_size=OBS_VECTOR_DIM), + 'action_mask': env._cached_action_masks[agent_id], + 'siem_embedding': np.zeros(EMBEDDING_DIM, dtype=np.float32), + 'adj_matrix': env.global_state.get_adjacency_matrix().flatten(), + 'delta_t': np.zeros(1, dtype=np.float32), + } + if 'blue' in agent_id: + agent_obs['blue_comm'] = np.zeros(BLUE_COMM_DIM, dtype=np.float32) + if env.pcap_obs: + agent_obs['pcap'] = np.zeros((N_PACKETS, PACKET_DIM), dtype=np.float32) + agent_obs['node_features'] = np.zeros( + (N_HOST_SLOTS, NODE_DIM), dtype=np.float32 + ) + observations[agent_id] = agent_obs + return observations diff --git a/netforge_rl/environment/spaces.py b/netforge_rl/environment/spaces.py index 818f741..4cd1ebe 100644 --- a/netforge_rl/environment/spaces.py +++ b/netforge_rl/environment/spaces.py @@ -1,19 +1,52 @@ import gymnasium as gym import numpy as np +from netforge_rl.environment.constants import ( + ACTION_MASK_DIM, + ADJ_FLAT_DIM, + BLUE_COMM_DIM, + N_ACTION_TYPES, + N_HOST_SLOTS, + OBS_VECTOR_DIM, +) from netforge_rl.nlp.log_encoder import EMBEDDING_DIM from netforge_rl.siem.pcap_synthesizer import N_PACKETS, NODE_DIM, PACKET_DIM +FLAT_ACTION_MASK = ACTION_MASK_DIM + + +def _split_flat_mask(mask, nvec): + """PettingZoo passes a concatenated action_mask; Gymnasium wants a tuple.""" + if mask is None or isinstance(mask, tuple): + return mask + arr = np.asarray(mask) + expected = int(np.sum(nvec)) + if arr.ndim == 1 and arr.size == expected: + cuts = np.cumsum(np.asarray(nvec, dtype=int))[:-1] + parts = np.split(arr.astype(np.int8, copy=False), cuts) + return tuple(parts) + return mask + + +class CyberActionSpace(gym.spaces.MultiDiscrete): + """``MultiDiscrete([types, hosts])`` that accepts a flat PettingZoo mask.""" + + def sample(self, mask=None, probability=None): + return super().sample( + mask=_split_flat_mask(mask, self.nvec), + probability=_split_flat_mask(probability, self.nvec), + ) + def build_observation_spaces(possible_agents, pcap_obs: bool) -> dict: """Per-agent Dict observation spaces. Blue agents also get a shared comm channel.""" base = { - 'obs': gym.spaces.Box(-1.0, 1.0, shape=(256,), dtype=np.float32), - 'action_mask': gym.spaces.Box(0, 1, shape=(32 + 100,), dtype=np.int8), + 'obs': gym.spaces.Box(-1.0, 1.0, shape=(OBS_VECTOR_DIM,), dtype=np.float32), + 'action_mask': gym.spaces.Box(0, 1, shape=(ACTION_MASK_DIM,), dtype=np.int8), 'siem_embedding': gym.spaces.Box( -1.0, 1.0, shape=(EMBEDDING_DIM,), dtype=np.float32 ), - 'adj_matrix': gym.spaces.Box(0.0, 1.0, shape=(10000,), dtype=np.float32), + 'adj_matrix': gym.spaces.Box(0.0, 1.0, shape=(ADJ_FLAT_DIM,), dtype=np.float32), 'delta_t': gym.spaces.Box(0.0, 1.0, shape=(1,), dtype=np.float32), } if pcap_obs: @@ -21,10 +54,12 @@ def build_observation_spaces(possible_agents, pcap_obs: bool) -> dict: 0.0, 1.0, shape=(N_PACKETS, PACKET_DIM), dtype=np.float32 ) base['node_features'] = gym.spaces.Box( - 0.0, 1.0, shape=(100, NODE_DIM), dtype=np.float32 + 0.0, 1.0, shape=(N_HOST_SLOTS, NODE_DIM), dtype=np.float32 ) blue = dict(base) - blue['blue_comm'] = gym.spaces.Box(0.0, 1.0, shape=(100,), dtype=np.float32) + blue['blue_comm'] = gym.spaces.Box( + 0.0, 1.0, shape=(BLUE_COMM_DIM,), dtype=np.float32 + ) return { agent: gym.spaces.Dict(blue if 'blue' in agent else base) for agent in possible_agents @@ -33,4 +68,7 @@ def build_observation_spaces(possible_agents, pcap_obs: bool) -> dict: def build_action_spaces(possible_agents) -> dict: """MultiDiscrete([action_type, target_host_index]) per agent.""" - return {agent: gym.spaces.MultiDiscrete([32, 100]) for agent in possible_agents} + return { + agent: CyberActionSpace([N_ACTION_TYPES, N_HOST_SLOTS]) + for agent in possible_agents + } diff --git a/netforge_rl/environment/tick.py b/netforge_rl/environment/tick.py new file mode 100644 index 0000000..711c22f --- /dev/null +++ b/netforge_rl/environment/tick.py @@ -0,0 +1,311 @@ +"""One tick: enqueue, clock, resolve, SIEM, physics, topology, observe.""" + +from __future__ import annotations + +import numpy as np + +from netforge_rl.core.action import ActionEffect, BaseAction +from netforge_rl.core.observation import BaseObservation +from netforge_rl.core.registry import action_registry +from netforge_rl.environment.constants import ( + BLUE_INFLIGHT_CAP, + MAX_ACTION_DURATION, + OBS_VECTOR_DIM, + SIEM_BUFFER_CAP, + SIEM_ENCODE_N, +) +from netforge_rl.environment.events import PendingAction, as_pending +from netforge_rl.nlp.log_encoder import EMBEDDING_DIM +from netforge_rl.siem.event_templates import sysmon_1 + + +def _inflight_counts(env) -> dict[str, int]: + counts: dict[str, int] = {} + for event in env.event_queue: + counts[event.agent] = counts.get(event.agent, 0) + 1 + return counts + + +def enqueue_actions(env, agent_actions: dict) -> None: + inflight = _inflight_counts(env) + for agent, action_int in agent_actions.items(): + if env.current_tick < env.global_state.agent_locked_until.get(agent, 0): + continue + if isinstance(action_int, BaseAction): + action = action_int + else: + env.ordered_hosts = sorted(env.global_state.all_hosts.keys()) + action = action_registry.instantiate_action( + agent, action_int, env.ordered_hosts + ) + if action is None: + continue + if 'blue' in agent.lower(): + if inflight.get(agent, 0) >= BLUE_INFLIGHT_CAP: + continue + inflight[agent] = inflight.get(agent, 0) + 1 + if env.global_state.agent_energy.get(agent, 0) < action.cost: + continue + if not action.validate(env.global_state): + continue + env.global_state.agent_energy[agent] -= action.cost + eta = getattr(action, 'duration', 1) + completion_tick = env.current_tick + eta + effect = action.execute(env.global_state) + effect.action = action + env.global_state.agent_locked_until[agent] = completion_tick + env.event_queue.append( + PendingAction( + completion_tick=completion_tick, + agent=agent, + action=action, + effect=effect, + target_ip=getattr(action, 'target_ip', None), + start_tick=env.current_tick, + ) + ) + + +def cancel_red_on_isolate(env) -> None: + for event in list(env.event_queue): + if type(event.action).__name__ != 'IsolateHost': + continue + if event.completion_tick > env.current_tick: + continue + target = event.target_ip + for red_event in list(env.event_queue): + if 'red' not in red_event.agent.lower(): + continue + if red_event.target_ip != target: + continue + if red_event in env.event_queue: + env.event_queue.remove(red_event) + env.global_state.agent_locked_until[red_event.agent] = env.current_tick + + +def advance_clock(env) -> tuple[float, float]: + prev = env.current_tick + if env.time_mode == 'fixed': + env.current_tick += 1 + elif env.event_queue: + nxt = min(e.completion_tick for e in env.event_queue) + env.current_tick = max(env.current_tick + 1, nxt) + else: + env.current_tick += 1 + delta_t = float(env.current_tick - prev) + return delta_t, delta_t / MAX_ACTION_DURATION + + +def pop_due_effects(env) -> tuple[dict[str, ActionEffect], dict[str, dict]]: + intended: dict[str, ActionEffect] = {} + metadata: dict[str, dict] = {} + remaining = [] + for event in env.event_queue: + if env.current_tick >= event.completion_tick: + intended[event.agent] = event.effect + metadata[event.agent] = { + 'name': type(event.action).__name__, + 'target_ip': event.target_ip, + } + else: + remaining.append(event) + env.event_queue = remaining + return intended, metadata + + +def emit_resolved_siem(env, resolved: dict, metadata: dict) -> None: + for agent, effect in resolved.items(): + meta = metadata.get(agent, {}) + env.siem_logger.log_action( + action_name=meta.get('name', 'UnknownAction'), + effect=effect, + global_state=env.global_state, + agent_id=agent, + target_ip=effect.observation_data.get('exploit'), + ) + for agent, effect in resolved.items(): + if 'red' not in agent or not effect.success: + continue + target_ip = effect.observation_data.get('exploit', 'unknown') + host = env.global_state.all_hosts.get(target_ip) + subnet = host.subnet_cidr if host else 'unknown' + env.siem_logger._push_to_buffer( + sysmon_1(agent, process='exploit_payload', rng=env.siem_logger._rng), + subnet, + env.global_state, + ) + if host and getattr(host, 'contains_honeytokens', False): + env.siem_logger._push_to_buffer( + { + 'signature': 'HONEYTOKEN_TRIGGERED', + 'target': target_ip, + 'agent': agent, + 'severity': 10, + }, + subnet, + env.global_state, + ) + + +def drop_stale_events(env) -> None: + valid = set(env.global_state.all_hosts.keys()) + env.event_queue = [ + e for e in env.event_queue if e.target_ip is None or e.target_ip in valid + ] + env._cached_action_masks = {agent: env.action_mask(agent) for agent in env.agents} + + +def maybe_dhcp(env) -> None: + interval = env.dhcp_interval + if interval > 0 and env.current_tick % interval == 0: + env.global_state.reallocate_dhcp(rng=env._py_random) + drop_stale_events(env) + + +def tick_physics(env) -> None: + alerts, deltas = env.physics_engine.tick(env.global_state) + for key, val in deltas: + env.global_state.apply_delta(key, val) + ot_subnet = '10.0.99.0/24' + for alert in alerts: + env.siem_logger._push_to_buffer(alert, ot_subnet, env.global_state) + + +def tick_correlator(env) -> None: + for log, subnet in env.correlator.correlate(env.global_state): + env.global_state.siem_log_buffer.append((log, subnet)) + if len(env.global_state.siem_log_buffer) > SIEM_BUFFER_CAP: + env.global_state.siem_log_buffer.pop(0) + + +def tick_topology(env) -> None: + events = env.topology_engine.tick(env.global_state) + if not events: + return + drop_stale_events(env) + for ev in events: + env.siem_logger._push_to_buffer( + { + 'signature': f'TOPOLOGY_{ev.kind.upper()}', + 'detail': ev.detail, + 'severity': 3, + }, + ev.detail.get('subnet', ev.detail.get('new_subnet', 'unknown')), + env.global_state, + ) + + +def encode_siem(env) -> dict[str, np.ndarray]: + vecs = {} + for agent in env.agents: + if 'blue' not in agent.lower(): + continue + subnet_tag = agent.split('_')[1] if '_' in agent else 'dmz' + logs = env.siem_logger.get_filtered_logs( + env.global_state, subnet_tag=subnet_tag, n=SIEM_ENCODE_N + ) + vecs[agent] = env.log_encoder.encode_buffer(logs, agg='mean') + return vecs + + +def collect_observations( + env, + resolved: dict, + metadata: dict, + delta_t_norm: float, + siem_vecs: dict, +) -> tuple[dict, dict]: + pcap = ( + env.pcap_synthesizer.synthesize( + env.global_state, env.current_tick, env.max_ticks + ) + if env.pcap_synthesizer + else None + ) + nodes = ( + env.pcap_synthesizer.node_features(env.global_state) + if env.pcap_synthesizer + else None + ) + blue_comm = env._build_blue_comm() + observations = {} + rewards = {} + for agent in env.agents: + obs = BaseObservation(agent) + obs.update_from_state(env.global_state, resolved) + if 'blue' in agent.lower(): + siem = siem_vecs.get(agent, np.zeros(EMBEDDING_DIM, dtype=np.float32)) + else: + siem = np.zeros(EMBEDDING_DIM, dtype=np.float32) + agent_obs = { + 'obs': obs.to_numpy(max_size=OBS_VECTOR_DIM), + 'action_mask': env._cached_action_masks[agent], + 'siem_embedding': siem, + 'adj_matrix': env._get_adj_matrix_for(agent).flatten(), + 'delta_t': np.array([delta_t_norm], dtype=np.float32), + } + if 'blue' in agent.lower(): + agent_obs['blue_comm'] = blue_comm + if env.pcap_obs: + agent_obs['pcap'] = pcap + agent_obs['node_features'] = nodes + observations[agent] = agent_obs + effect = resolved.get(agent) + rewards[agent] = env.scenario.calculate_reward(agent, env.global_state, effect) + if env.trajectory_recorder is not None: + for agent, effect in resolved.items(): + meta = metadata.get(agent, {}) + env.trajectory_recorder.record_step( + tick=env.current_tick, + agent_id=agent, + action_name=meta.get('name', 'UnknownAction'), + target_ip=meta.get('target_ip'), + success=effect.success, + reward=float(rewards.get(agent, 0.0)), + ) + return observations, rewards + + +def run_step(env, agent_actions: dict): + env.event_queue = [as_pending(event) for event in env.event_queue] + enqueue_actions(env, agent_actions) + cancel_red_on_isolate(env) + delta_t, delta_t_norm = advance_clock(env) + env.global_state.current_tick = env.current_tick + env.global_state.subnet_bandwidth.clear() + + noise = env.green_agent.generate_noise( + env.current_tick, env.global_state, rng=env._py_random + ) + for anomaly in noise.get('alerts', []): + env.siem_logger._push_to_buffer( + anomaly['data'], anomaly['subnet'], env.global_state + ) + + intended, metadata = pop_due_effects(env) + resolved = env.resolution_engine.resolve(intended) + env._apply_state_deltas(resolved) + env._update_episode_metrics(resolved) + emit_resolved_siem(env, resolved, metadata) + env.siem_logger.log_background_noise(env.global_state) + maybe_dhcp(env) + tick_physics(env) + tick_correlator(env) + tick_topology(env) + + terminate = env.scenario.check_termination(env.global_state) + truncated = env.current_tick >= env.max_ticks + truncate = {agent: truncated for agent in env.agents} + env.siem_logger.release(env.global_state) + observations, rewards = collect_observations( + env, resolved, metadata, delta_t_norm, encode_siem(env) + ) + env.agents = [ + agent for agent in env.agents if not terminate[agent] and not truncate[agent] + ] + infos = env._extract_agent_infos(observations, resolved, rewards) + for agent in env.agents: + if agent in infos: + infos[agent]['delta_t'] = delta_t + infos[agent]['delta_t_norm'] = delta_t_norm + return observations, rewards, terminate, truncate, infos diff --git a/tests/environment/test_env_ids.py b/tests/environment/test_env_ids.py new file mode 100644 index 0000000..920c59a --- /dev/null +++ b/tests/environment/test_env_ids.py @@ -0,0 +1,50 @@ +import numpy as np +import pytest + +from netforge_rl.environment.parallel_env import NetForgeRLEnv +from netforge_rl.environment.registry import ( + GYM_BLUE_ID, + PZ_RANSOMWARE_ID, + register_envs, +) +from netforge_rl.environment.spaces import FLAT_ACTION_MASK, CyberActionSpace + + +@pytest.mark.fast +def test_flat_action_mask_samples_on_multidiscrete(): + space = CyberActionSpace([32, 100]) + mask = np.zeros(FLAT_ACTION_MASK, dtype=np.int8) + mask[0] = 1 + mask[32] = 1 + sample = space.sample(mask=mask) + assert sample.shape == (2,) + assert int(sample[0]) == 0 + assert int(sample[1]) == 0 + + +@pytest.mark.fast +def test_env_action_space_accepts_obs_mask(): + env = NetForgeRLEnv({'scenario_type': 'ransomware', 'max_ticks': 3}) + obs, _ = env.reset(seed=0) + agent = env.agents[0] + action = env.action_space(agent).sample(mask=obs[agent]['action_mask']) + assert action.shape == (2,) + + +@pytest.mark.fast +def test_gymnasium_and_pettingzoo_register_and_make(): + register_envs() + import gymnasium as gym + from pettingzoo import make, parallel_registry + + assert GYM_BLUE_ID in gym.envs.registry + assert PZ_RANSOMWARE_ID in parallel_registry + gym_env = gym.make(GYM_BLUE_ID, max_ticks=4) + _, info = gym_env.reset(seed=0) + assert 'action_mask' in info + gym_env.close() + + pz_env = make('parallel', PZ_RANSOMWARE_ID, max_ticks=4) + pz_obs, _ = pz_env.reset(seed=0) + assert 'blue_dmz' in pz_obs + pz_env.close() From 87e1985bb753aba21680889f376c8b6695f36ccc Mon Sep 17 00:00:00 2001 From: elprofesoriqo Date: Sun, 20 Sep 2026 15:15:13 +0200 Subject: [PATCH 2/8] Move the action codebook and SIEM belief graphs out of the tick loop. Register DeployEDR so the mask matches the action table, and keep oracle graphs off the Blue observation path. Co-authored-by: Cursor --- netforge_rl/actions/__init__.py | 80 ++++---- netforge_rl/actions/blue/__init__.py | 42 +++-- netforge_rl/agents/rule_based_blue.py | 10 +- netforge_rl/bridges/cleanrl_vec.py | 2 +- netforge_rl/core/__init__.py | 29 +++ netforge_rl/core/codes.py | 50 +++++ netforge_rl/core/commands.py | 13 +- netforge_rl/core/functional.py | 85 +++------ netforge_rl/core/graph_obs.py | 197 +++++++++++++++++++- netforge_rl/core/observation.py | 5 +- netforge_rl/core/registry.py | 14 +- netforge_rl/diagnostics/__init__.py | 8 +- netforge_rl/diagnostics/adaptation.py | 94 ++++++++++ netforge_rl/diagnostics/base.py | 9 +- netforge_rl/diagnostics/capability_card.py | 6 +- netforge_rl/diagnostics/deception.py | 65 +++++++ netforge_rl/diagnostics/suite.py | 6 +- netforge_rl/diagnostics/wrapper.py | 14 +- netforge_rl/siem/correlator.py | 14 +- netforge_rl/siem/export.py | 6 +- netforge_rl/topologies/network_generator.py | 37 ++-- tests/core/test_graph_obs.py | 39 +++- tests/diagnostics/test_capability_card.py | 2 + tests/diagnostics/test_capability_probes.py | 8 +- 24 files changed, 650 insertions(+), 185 deletions(-) create mode 100644 netforge_rl/core/codes.py create mode 100644 netforge_rl/diagnostics/adaptation.py create mode 100644 netforge_rl/diagnostics/deception.py diff --git a/netforge_rl/actions/__init__.py b/netforge_rl/actions/__init__.py index 6466df0..bf02d58 100644 --- a/netforge_rl/actions/__init__.py +++ b/netforge_rl/actions/__init__.py @@ -1,74 +1,76 @@ from .blue import ( - IsolateHost, - RestoreHost, - Monitor, Analyze, - DeployDecoy, - Remove, - RestoreFromBackup, + ConfigureACL, DecoyApache, DecoySSHD, DecoyTomcat, - Misinform, - ConfigureACL, - SecurityAwarenessTraining, + DeployDecoy, + DeployEDR, DeployHoneytoken, + IsolateHost, + Misinform, + Monitor, + Remove, + RestoreFromBackup, + RestoreHost, RotateKerberos, + SecurityAwarenessTraining, ) +from .network.ip_fragmentation import IPFragmentationAction from .red import ( - NetworkScan, - DiscoverRemoteSystems, DiscoverNetworkServices, - ExploitRemoteService, - PrivilegeEscalate, - Impact, + DiscoverRemoteSystems, ExploitBlueKeep, ExploitEternalBlue, ExploitHTTP_RFI, + ExploitRemoteService, + Impact, JuicyPotato, - V4L2KernelExploit, KillProcess, - ShareIntelligence, + NetworkScan, OverloadPLC, + PrivilegeEscalate, + ShareIntelligence, SpearPhishing, + V4L2KernelExploit, ) from .red.post_exploitation import DumpLSASS, PassTheTicket -from .network.ip_fragmentation import IPFragmentationAction __all__ = [ - 'IsolateHost', - 'RestoreHost', - 'Monitor', 'Analyze', - 'DeployDecoy', - 'Remove', - 'RestoreFromBackup', + 'ConfigureACL', 'DecoyApache', 'DecoySSHD', 'DecoyTomcat', - 'Misinform', - 'ConfigureACL', - 'SecurityAwarenessTraining', + 'DeployDecoy', + 'DeployEDR', 'DeployHoneytoken', - 'RotateKerberos', - 'NetworkScan', - 'DiscoverRemoteSystems', 'DiscoverNetworkServices', - 'ExploitRemoteService', - 'PrivilegeEscalate', - 'Impact', + 'DiscoverRemoteSystems', + 'DumpLSASS', + 'ExfiltrateData', 'ExploitBlueKeep', 'ExploitEternalBlue', 'ExploitHTTP_RFI', + 'ExploitRemoteService', + 'IPFragmentationAction', + 'Impact', + 'IsolateHost', 'JuicyPotato', - 'V4L2KernelExploit', 'KillProcess', - 'ShareIntelligence', + 'Misinform', + 'Monitor', + 'NetworkScan', 'OverloadPLC', - 'SpearPhishing', - 'DumpLSASS', - 'PassTheTicket', 'PassTheHash', - 'ExfiltrateData', - 'IPFragmentationAction', + 'PassTheTicket', + 'PrivilegeEscalate', + 'Remove', + 'RestoreFromBackup', + 'RestoreHost', + 'RotateKerberos', + 'SecurityAwarenessTraining', + 'ShareIntelligence', + 'SpearPhishing', + 'V4L2KernelExploit', ] diff --git a/netforge_rl/actions/blue/__init__.py b/netforge_rl/actions/blue/__init__.py index 6ed3cd7..f7274e7 100644 --- a/netforge_rl/actions/blue/__init__.py +++ b/netforge_rl/actions/blue/__init__.py @@ -1,36 +1,38 @@ -from .mitigation import ( - IsolateHost, - RestoreHost, - Remove, - RestoreFromBackup, - ConfigureACL, - SecurityAwarenessTraining, -) -from .analysis import Monitor, Analyze +from .analysis import Analyze, Monitor from .deception import ( - DeployDecoy, DecoyApache, DecoySSHD, DecoyTomcat, - Misinform, + DeployDecoy, DeployHoneytoken, + Misinform, +) +from .edr import DeployEDR +from .mitigation import ( + ConfigureACL, + IsolateHost, + Remove, + RestoreFromBackup, + RestoreHost, + SecurityAwarenessTraining, ) __all__ = [ - 'IsolateHost', - 'RestoreHost', - 'Remove', - 'RestoreFromBackup', - 'Monitor', 'Analyze', - 'DeployDecoy', + 'ConfigureACL', 'DecoyApache', 'DecoySSHD', 'DecoyTomcat', - 'Misinform', - 'ConfigureACL', - 'SecurityAwarenessTraining', + 'DeployDecoy', + 'DeployEDR', 'DeployHoneytoken', + 'IsolateHost', + 'Misinform', + 'Monitor', + 'Remove', + 'RestoreFromBackup', + 'RestoreHost', 'RotateKerberos', + 'SecurityAwarenessTraining', ] from .identity import RotateKerberos diff --git a/netforge_rl/agents/rule_based_blue.py b/netforge_rl/agents/rule_based_blue.py index 4a1be43..022640b 100644 --- a/netforge_rl/agents/rule_based_blue.py +++ b/netforge_rl/agents/rule_based_blue.py @@ -1,13 +1,15 @@ +from __future__ import annotations + import random -from typing import Any, Optional -from netforge_rl.actions.blue.mitigation import IsolateHost, RestoreHost + from netforge_rl.actions.blue.analysis import Analyze, Monitor from netforge_rl.actions.blue.deception import DeployHoneytoken from netforge_rl.actions.blue.edr import DeployEDR +from netforge_rl.actions.blue.mitigation import IsolateHost, RestoreHost class RuleBasedBlueAgent: - """Deterministic Blue defender for use as a lower-bound paper baseline.""" + """Deterministic Blue lower bound.""" def __init__(self, agent_id: str): self.agent_id = agent_id @@ -20,7 +22,7 @@ def reset(self) -> None: self._honeytokened.clear() self._edr_deployed.clear() - def get_action(self, _observation: Any, global_state) -> Optional[Any]: + def get_action(self, _observation, global_state): live = { ip: h for ip, h in global_state.all_hosts.items() diff --git a/netforge_rl/bridges/cleanrl_vec.py b/netforge_rl/bridges/cleanrl_vec.py index d5f5b37..87fca80 100644 --- a/netforge_rl/bridges/cleanrl_vec.py +++ b/netforge_rl/bridges/cleanrl_vec.py @@ -1,5 +1,5 @@ +from collections.abc import Callable from dataclasses import dataclass -from typing import Callable import jax import jax.numpy as jnp diff --git a/netforge_rl/core/__init__.py b/netforge_rl/core/__init__.py index e69de29..d280721 100644 --- a/netforge_rl/core/__init__.py +++ b/netforge_rl/core/__init__.py @@ -0,0 +1,29 @@ +from netforge_rl.core.action import ActionEffect, BaseAction +from netforge_rl.core.codes import ( + CVE_CODES, + N_HOSTS, + OS_LINUX, + OS_WINDOWS, + PRIVILEGE_CODES, + STATUS_CODES, +) +from netforge_rl.core.functional import EnvState, from_global_state +from netforge_rl.core.observation import BaseObservation +from netforge_rl.core.state import GlobalNetworkState, Host, Subnet + +__all__ = [ + 'CVE_CODES', + 'N_HOSTS', + 'OS_LINUX', + 'OS_WINDOWS', + 'PRIVILEGE_CODES', + 'STATUS_CODES', + 'ActionEffect', + 'BaseAction', + 'BaseObservation', + 'EnvState', + 'GlobalNetworkState', + 'Host', + 'Subnet', + 'from_global_state', +] diff --git a/netforge_rl/core/codes.py b/netforge_rl/core/codes.py new file mode 100644 index 0000000..dbe3af3 --- /dev/null +++ b/netforge_rl/core/codes.py @@ -0,0 +1,50 @@ +"""Host / CVE / token codebooks shared by the Python gym and JAX kernels.""" + +N_HOSTS = 100 +STATUS_CODES = ('online', 'isolated', 'kernel_panic') +PRIVILEGE_CODES = ('None', 'User', 'Root') +DECOY_CODES = ('inactive', 'active', 'Apache', 'SSHD', 'Tomcat') +INTEGRITY_CODES = ('clean', 'compromised', 'kinetic_destruction') +CVE_CODES = ( + 'MS17-010', + 'CVE-2019-0708', + 'CVE-2021-44228', + 'V4L2', + 'CVE-2010-2772', + 'Stuxnet_0day', +) +N_CVE = len(CVE_CODES) +TOKEN_CODES = ('Enterprise_Admin_Token', 'Local_Admin_DMZ', 'Local_Admin_Corporate') +N_TOKEN = len(TOKEN_CODES) +OS_OTHER, OS_WINDOWS, OS_LINUX, OS_PLC = (0, 1, 2, 3) + + +def os_family_code(os_str) -> int: + text = str(os_str or '') + if 'Windows' in text: + return OS_WINDOWS + if 'Linux' in text: + return OS_LINUX + if 'PLC' in text: + return OS_PLC + return OS_OTHER + + +def encode(value, codebook) -> int: + try: + return codebook.index(value) + except ValueError: + return 0 + + +def decode(code, codebook): + idx = int(code) + if 0 <= idx < len(codebook): + return codebook[idx] + return codebook[0] + + +# Private aliases used by the conversion layer. +_os_family_code = os_family_code +_encode = encode +_decode = decode diff --git a/netforge_rl/core/commands.py b/netforge_rl/core/commands.py index be41ad9..f8749a1 100644 --- a/netforge_rl/core/commands.py +++ b/netforge_rl/core/commands.py @@ -1,9 +1,12 @@ +from __future__ import annotations + from abc import ABC, abstractmethod -from typing import Any, Iterator, Optional, Tuple +from collections.abc import Iterator + from netforge_rl.core.state import Firewall -def iter_host_deltas(state_deltas: Any) -> Iterator[Tuple[str, str, Any]]: +def iter_host_deltas(state_deltas) -> Iterator[tuple[str, str, object]]: """Yield ``(attribute, ip, value)`` host mutations from either delta encoding.""" if isinstance(state_deltas, dict): for key, value in state_deltas.items(): @@ -25,15 +28,15 @@ class IStateDeltaCommand(ABC): """Object-oriented state mutation; the resolver reads target_ip.""" @abstractmethod - def execute(self, global_state: Any): ... + def execute(self, global_state): ... @property @abstractmethod - def target_ip(self) -> Optional[str]: ... + def target_ip(self) -> str | None: ... class UpdateHostPrivilegeCommand(IStateDeltaCommand): - def __init__(self, ip: str, privilege: str, compromised_by: Optional[str] = None): + def __init__(self, ip: str, privilege: str, compromised_by: str | None = None): self._target_ip = ip self.privilege = privilege self.compromised_by = compromised_by diff --git a/netforge_rl/core/functional.py b/netforge_rl/core/functional.py index 5940e53..e0574d8 100644 --- a/netforge_rl/core/functional.py +++ b/netforge_rl/core/functional.py @@ -1,50 +1,26 @@ from dataclasses import dataclass, field, replace import numpy as np -from netforge_rl.core.state import GlobalNetworkState, Subnet, Host - -N_HOSTS = 100 -STATUS_CODES = ('online', 'isolated', 'kernel_panic') -PRIVILEGE_CODES = ('None', 'User', 'Root') -DECOY_CODES = ('inactive', 'active', 'Apache', 'SSHD', 'Tomcat') -INTEGRITY_CODES = ('clean', 'compromised', 'kinetic_destruction') -CVE_CODES = ( - 'MS17-010', - 'CVE-2019-0708', - 'CVE-2021-44228', - 'V4L2', - 'CVE-2010-2772', - 'Stuxnet_0day', -) -N_CVE = len(CVE_CODES) -TOKEN_CODES = ('Enterprise_Admin_Token', 'Local_Admin_DMZ', 'Local_Admin_Corporate') -N_TOKEN = len(TOKEN_CODES) -OS_OTHER, OS_WINDOWS, OS_LINUX, OS_PLC = (0, 1, 2, 3) - - -def _os_family_code(os_str): - s = str(os_str or '') - if 'Windows' in s: - return OS_WINDOWS - if 'Linux' in s: - return OS_LINUX - if 'PLC' in s: - return OS_PLC - return OS_OTHER - -def _encode(value, codebook): - try: - return codebook.index(value) - except ValueError: - return 0 - - -def _decode(code, codebook): - code = int(code) - if 0 <= code < len(codebook): - return codebook[code] - return codebook[0] +from netforge_rl.core.codes import ( # noqa: F401 — public codebook re-exports + CVE_CODES, + DECOY_CODES, + INTEGRITY_CODES, + N_CVE, + N_HOSTS, + N_TOKEN, + OS_LINUX, + OS_OTHER, + OS_PLC, + OS_WINDOWS, + PRIVILEGE_CODES, + STATUS_CODES, + TOKEN_CODES, + _decode, + _encode, + _os_family_code, +) +from netforge_rl.core.state import GlobalNetworkState, Host, Subnet @dataclass(frozen=True) @@ -185,13 +161,13 @@ def from_global_state(legacy, agent_ids): ) meta = HostMeta( ip=sorted_ips, - hostname=tuple((h.hostname for h in hosts_in_order)), - subnet_cidr=tuple((h.subnet_cidr for h in hosts_in_order)), - os=tuple((h.os for h in hosts_in_order)), - services=tuple((tuple(h.services) for h in hosts_in_order)), - vulnerabilities=tuple((tuple(h.vulnerabilities) for h in hosts_in_order)), - cached_credentials=tuple((tuple(h.cached_credentials) for h in hosts_in_order)), - system_tokens=tuple((tuple(h.system_tokens) for h in hosts_in_order)), + hostname=tuple(h.hostname for h in hosts_in_order), + subnet_cidr=tuple(h.subnet_cidr for h in hosts_in_order), + os=tuple(h.os for h in hosts_in_order), + services=tuple(tuple(h.services) for h in hosts_in_order), + vulnerabilities=tuple(tuple(h.vulnerabilities) for h in hosts_in_order), + cached_credentials=tuple(tuple(h.cached_credentials) for h in hosts_in_order), + system_tokens=tuple(tuple(h.system_tokens) for h in hosts_in_order), ) return EnvState( hosts=hosts, @@ -213,10 +189,10 @@ def from_global_state(legacy, agent_ids): current_tick=int(legacy.current_tick), business_downtime_score=float(legacy.business_downtime_score), knowledge=tuple( - (frozenset(legacy.agent_knowledge.get(a, set())) for a in agent_ids) + frozenset(legacy.agent_knowledge.get(a, set())) for a in agent_ids ), inventory=tuple( - (frozenset(legacy.agent_inventory.get(a, set())) for a in agent_ids) + frozenset(legacy.agent_inventory.get(a, set())) for a in agent_ids ), ) @@ -364,7 +340,7 @@ def _extract_targeted_ips(state_deltas): """Pull every ``hosts//...`` IP out of a state_deltas payload.""" ips = set() if isinstance(state_deltas, dict): - for key in state_deltas.keys(): + for key in state_deltas: if isinstance(key, str) and key.startswith('hosts/'): parts = key.split('/') if len(parts) >= 2: @@ -377,10 +353,9 @@ def _extract_targeted_ips(state_deltas): return ips -from typing import Dict from netforge_rl.core.action import ActionEffect from netforge_rl.core.physics import ConflictResolutionEngine -def resolve_conflicts(effects: Dict[str, ActionEffect]) -> Dict[str, ActionEffect]: +def resolve_conflicts(effects: dict[str, ActionEffect]) -> dict[str, ActionEffect]: return ConflictResolutionEngine.resolve(effects) diff --git a/netforge_rl/core/graph_obs.py b/netforge_rl/core/graph_obs.py index 519d5f4..28e74b1 100644 --- a/netforge_rl/core/graph_obs.py +++ b/netforge_rl/core/graph_obs.py @@ -1,27 +1,127 @@ from __future__ import annotations -from typing import Optional +import re +from collections.abc import Iterable import numpy as np -from netforge_rl.siem.pcap_synthesizer import NODE_DIM, _PRIV_ENC, _SUBNET_ENC +from netforge_rl.siem.pcap_synthesizer import _PRIV_ENC, _SUBNET_ENC, NODE_DIM MAX_NODES = 100 +_IP_RE = re.compile(r'\b(?:\d{1,3}\.){3}\d{1,3}\b') +_SEV_RE = re.compile(r'severity\s*[=:]\s*(\d+)', re.IGNORECASE) + +# Belief-feature layout (same NODE_DIM as oracle, different semantics): +# 0 believed_compromised, 1 online, 2 seen_in_siem, 3 unused (no decoy leak), +# 4 unused (no DC leak), 5 subnet prior if inventory, 6 alert_severity/10, +# 7 isolated (Blue's own action / reachable status). +BELIEF_COMPROMISED = 0 +BELIEF_ONLINE = 1 +BELIEF_SEEN = 2 +BELIEF_SUBNET = 5 +BELIEF_SEVERITY = 6 +BELIEF_ISOLATED = 7 + + +def _iter_siem_logs(global_state) -> Iterable[object]: + buffer = getattr(global_state, 'siem_log_buffer', None) or [] + current_tick = getattr(global_state, 'current_tick', 0) + for entry in buffer: + if isinstance(entry, tuple): + log = entry[0] + else: + log = entry + if isinstance(log, dict): + arrival = log.get('arrival_tick', 0) or 0 + if arrival > current_tick: + continue + yield log + + +def _ips_from_log(log) -> list[str]: + ips: list[str] = [] + if isinstance(log, dict): + for key in ('target', 'src', 'dst', 'ip', 'source', 'destination'): + val = log.get(key) + if isinstance(val, str) and _IP_RE.fullmatch(val.strip()): + ips.append(val.strip()) + for val in log.values(): + if isinstance(val, str): + ips.extend(_IP_RE.findall(val)) + else: + ips.extend(_IP_RE.findall(str(log))) + return list(dict.fromkeys(ips)) + + +def _severity_from_log(log) -> float: + if isinstance(log, dict): + return float(log.get('severity', 0) or 0) + match = _SEV_RE.search(str(log)) + return float(match.group(1)) if match else 0.0 + + +def siem_belief(global_state) -> dict[str, dict]: + """Per-host SIEM evidence: mention counts and max severity. No ground truth.""" + belief: dict[str, dict] = {} + for log in _iter_siem_logs(global_state): + sev = _severity_from_log(log) + for ip in _ips_from_log(log): + slot = belief.setdefault(ip, {'mentions': 0, 'severity': 0.0}) + slot['mentions'] += 1 + slot['severity'] = max(slot['severity'], sev) + return belief + + +def resolve_graph_mode(agent_id: str | None, mode: str) -> str: + """Map ``auto`` onto fog / belief / oracle from the agent role.""" + if mode != 'auto': + return mode + if not agent_id: + return 'oracle' + lowered = agent_id.lower() + if 'red' in lowered: + return 'fog' + if 'blue' in lowered: + return 'belief' + return 'oracle' def build_graph_observation( global_state, - agent_id: Optional[str] = None, + agent_id: str | None = None, max_nodes: int = MAX_NODES, + mode: str = 'auto', ) -> dict: - """Graph view of the network (node_features, edge_index, edge_attr, node_mask), - mapping onto PyTorch Geometric / jraph. Red agents see only discovered hosts.""" + """Graph view of the network (node_features, edge_index, edge_attr, node_mask). + + Modes: + * ``oracle`` — privileged simulator state (diagnostics / replay only). + * ``fog`` — Red recon knowledge. + * ``belief`` — Blue inventory + SIEM-reconstructed compromise/edges. + * ``auto`` — fog for Red, belief for Blue, oracle if no agent is given. + """ + resolved = resolve_graph_mode(agent_id, mode) + if resolved == 'belief': + return _belief_graph(global_state, agent_id, max_nodes) + if resolved == 'fog': + return _oracle_or_fog_graph(global_state, agent_id, max_nodes, fog=True) + return _oracle_or_fog_graph(global_state, agent_id, max_nodes, fog=False) + + +def build_oracle_graph(global_state, max_nodes: int = MAX_NODES) -> dict: + """Privileged graph. Keep out of the Blue observation; diagnostics only.""" + return _oracle_or_fog_graph( + global_state, agent_id=None, max_nodes=max_nodes, fog=False + ) + + +def _oracle_or_fog_graph(global_state, agent_id, max_nodes, fog: bool) -> dict: sorted_ips = sorted(global_state.all_hosts.keys())[:max_nodes] n = len(sorted_ips) idx = {ip: i for i, ip in enumerate(sorted_ips)} - - is_red = bool(agent_id) and 'red' in agent_id.lower() - known = global_state.agent_knowledge.get(agent_id, set()) if is_red else None + known = ( + global_state.agent_knowledge.get(agent_id, set()) if fog and agent_id else None + ) node_features = np.zeros((max_nodes, NODE_DIM), dtype=np.float32) node_mask = np.zeros((max_nodes,), dtype=np.float32) @@ -29,7 +129,7 @@ def build_graph_observation( h = global_state.all_hosts[ip] if ip.startswith('169.254.'): continue - if is_red and ip not in known: + if fog and ip not in known: continue subnet = global_state.get_subnet_name(h.subnet_cidr) node_features[i, 0] = _PRIV_ENC.get(h.privilege, 0.0) @@ -59,6 +159,84 @@ def build_graph_observation( ) != global_state.get_subnet_name(global_state.all_hosts[jp].subnet_cidr) attr.append(1.0 if cross else 0.0) + return _pack_graph(node_features, node_mask, src, dst, attr, n) + + +def _belief_graph(global_state, agent_id, max_nodes) -> dict: + """Uncertain graph: known assets as nodes, compromise/edges from SIEM.""" + sorted_ips = sorted(global_state.all_hosts.keys())[:max_nodes] + n = len(sorted_ips) + idx = {ip: i for i, ip in enumerate(sorted_ips)} + inventory = set() + if agent_id: + inventory = set(global_state.agent_knowledge.get(agent_id, set())) + evidence = siem_belief(global_state) + telemetry_ips = {ip for ip in evidence if ip in idx} + + visible = set() + for ip in idx: + if ip.startswith('169.254.'): + continue + if ip in inventory or ip in telemetry_ips: + visible.add(ip) + + node_features = np.zeros((max_nodes, NODE_DIM), dtype=np.float32) + node_mask = np.zeros((max_nodes,), dtype=np.float32) + for ip in visible: + i = idx[ip] + host = global_state.all_hosts.get(ip) + ev = evidence.get(ip, {'mentions': 0, 'severity': 0.0}) + subnet = ( + global_state.get_subnet_name(host.subnet_cidr) if host is not None else '' + ) + node_features[i, BELIEF_COMPROMISED] = ( + 1.0 if ev['severity'] >= 5 or ev['mentions'] >= 2 else 0.0 + ) + node_features[i, BELIEF_ONLINE] = ( + 1.0 if host is not None and host.status == 'online' else 0.0 + ) + node_features[i, BELIEF_SEEN] = 1.0 if ev['mentions'] > 0 else 0.0 + node_features[i, BELIEF_SUBNET] = _SUBNET_ENC.get(subnet, 0.0) + node_features[i, BELIEF_SEVERITY] = min(ev['severity'] / 10.0, 1.0) + node_features[i, BELIEF_ISOLATED] = ( + 1.0 if host is not None and host.status == 'isolated' else 0.0 + ) + node_mask[i] = 1.0 + + siem_edges: set[tuple[int, int]] = set() + for log in _iter_siem_logs(global_state): + ips = [ip for ip in _ips_from_log(log) if ip in visible] + for a in ips: + for b in ips: + if a != b: + siem_edges.add((idx[a], idx[b])) + + src, dst, attr = [], [], [] + visible_list = [ip for ip, i in idx.items() if node_mask[i] == 1.0] + for ip in visible_list: + i = idx[ip] + host = global_state.all_hosts.get(ip) + for jp in visible_list: + j = idx[jp] + if i == j: + continue + other = global_state.all_hosts.get(jp) + same_subnet = ( + host is not None + and other is not None + and host.subnet_cidr == other.subnet_cidr + ) + tele = (i, j) in siem_edges + if not same_subnet and not tele: + continue + src.append(i) + dst.append(j) + attr.append(0.0 if same_subnet else 1.0) + + return _pack_graph(node_features, node_mask, src, dst, attr, n) + + +def _pack_graph(node_features, node_mask, src, dst, attr, n) -> dict: edge_index = ( np.array([src, dst], dtype=np.int64) if src @@ -69,7 +247,6 @@ def build_graph_observation( if attr else np.zeros((0, 1), dtype=np.float32) ) - return { 'node_features': node_features, 'edge_index': edge_index, diff --git a/netforge_rl/core/observation.py b/netforge_rl/core/observation.py index fd122d4..f2980bc 100644 --- a/netforge_rl/core/observation.py +++ b/netforge_rl/core/observation.py @@ -1,4 +1,5 @@ -from typing import Any, List +from __future__ import annotations + import numpy as np @@ -13,7 +14,7 @@ def __init__(self, agent_id: str): self.network_telemetry = {} self.siem_alerts = [] - def update_from_state(self, global_state: Any, _action_effects: List[Any]): + def update_from_state(self, global_state, _action_effects: list): """Filter ``global_state`` down to what this agent can observe.""" is_blue = 'blue' in self.agent_id.lower() is_commander = 'commander' in self.agent_id.lower() diff --git a/netforge_rl/core/registry.py b/netforge_rl/core/registry.py index d247f66..e07fa45 100644 --- a/netforge_rl/core/registry.py +++ b/netforge_rl/core/registry.py @@ -1,17 +1,18 @@ +from __future__ import annotations + import inspect -from typing import Callable, Dict, Optional, Type +from collections.abc import Callable def team_of(agent_id: str) -> str: - """Every agent belongs to exactly one action team: 'red' or 'blue'.""" return 'red' if 'red' in agent_id.lower() else 'blue' class ActionRegistry: - """Factory registry mapping ``(team, group_id) -> BaseAction subclass.""" + """``(team, group_id) -> BaseAction`` factory.""" def __init__(self): - self._actions: Dict[str, Dict[int, Type]] = {'red': {}, 'blue': {}} + self._actions: dict[str, dict[int, type]] = {'red': {}, 'blue': {}} def register(self, team: str, group_id: int) -> Callable: def decorator(cls): @@ -20,13 +21,12 @@ def decorator(cls): return decorator - def get_action_class(self, agent_id: str, group_id: int) -> Optional[Type]: + def get_action_class(self, agent_id: str, group_id: int) -> type | None: return self._actions.get(team_of(agent_id), {}).get(group_id) def instantiate_action( self, agent_id: str, action_data: object, target_ips: list - ) -> Optional[object]: - """Resolve an action payload to a BaseAction instance.""" + ) -> object | None: if not target_ips: target_ips = ['127.0.0.1'] if ( diff --git a/netforge_rl/diagnostics/__init__.py b/netforge_rl/diagnostics/__init__.py index 8741811..4c4c394 100644 --- a/netforge_rl/diagnostics/__init__.py +++ b/netforge_rl/diagnostics/__init__.py @@ -1,14 +1,17 @@ +from netforge_rl.diagnostics.adaptation import AdaptationShift from netforge_rl.diagnostics.base import ( Diagnostic, DiagnosticResult, run_diagnostic, ) -from netforge_rl.diagnostics.suite import all_diagnostics +from netforge_rl.diagnostics.capability_card import capability_card +from netforge_rl.diagnostics.deception import DeceptionResistance from netforge_rl.diagnostics.delayed_telemetry import DelayedTelemetry from netforge_rl.diagnostics.false_positive import FalsePositiveRestraint from netforge_rl.diagnostics.memory_probe import MemoryProbe from netforge_rl.diagnostics.noisy_siem import NoisySIEM from netforge_rl.diagnostics.ot_kinetic import OTKineticResponse +from netforge_rl.diagnostics.suite import all_diagnostics from netforge_rl.diagnostics.topology_shift import TopologyShift from netforge_rl.diagnostics.wrapper import DiagnosticsWrapper, OracleObservation @@ -21,8 +24,11 @@ 'FalsePositiveRestraint', 'OTKineticResponse', 'TopologyShift', + 'DeceptionResistance', + 'AdaptationShift', 'all_diagnostics', 'run_diagnostic', + 'capability_card', 'DiagnosticsWrapper', 'OracleObservation', ] diff --git a/netforge_rl/diagnostics/adaptation.py b/netforge_rl/diagnostics/adaptation.py new file mode 100644 index 0000000..fab9815 --- /dev/null +++ b/netforge_rl/diagnostics/adaptation.py @@ -0,0 +1,94 @@ +from netforge_rl.baselines.policies import KillChainRedPolicy, RandomPolicy +from netforge_rl.diagnostics.base import Diagnostic, DiagnosticResult, flag_host +from netforge_rl.environment.constants import PADDING_SUBNET + + +class AdaptationShift(Diagnostic): + """Kill-chain Red for the first half of the episode, then a random spreader. + Score Blue on SLA and containment *after* the attacker strategy shifts.""" + + name = 'adaptation_shift' + capability = 'adaptation' + scenario_type = 'ransomware' + max_ticks = 40 + switch_tick = 20 + + def __init__(self, seed=0): + self._red_a = KillChainRedPolicy(seed=seed) + self._red_b = RandomPolicy(seed=seed + 1) + self._switch_compromised = 0 + self._final_sla = 1.0 + self._contained_after = 0 + self._new_infections = 0 + + def setup(self, env): + live = [ + ip + for ip, h in env.global_state.all_hosts.items() + if not ip.startswith('169.254.') and h.status == 'online' + ] + planted = sorted(live)[0] + host = env.global_state.all_hosts[planted] + host.compromised_by = 'red_operator' + host.privilege = 'User' + flag_host(env, planted) + self._switch_compromised = 0 + self._final_sla = 1.0 + self._contained_after = 0 + self._new_infections = 0 + + def act_for(self, env, agent_id, blue_policy): + if 'red' in agent_id.lower(): + if env.current_tick < self.switch_tick: + return self._red_a.act(env, agent_id) + return self._red_b.act(env, agent_id) + return blue_policy.act(env, agent_id) + + def early_stop(self, env): + if env.current_tick == self.switch_tick: + self._switch_compromised = self._count_compromised(env) + return False + + def score(self, env, ticks_used): + active = [ + h + for h in env.global_state.all_hosts.values() + if h.subnet_cidr != PADDING_SUBNET + ] + total = max(len(active), 1) + healthy = sum( + 1 for h in active if h.compromised_by == 'None' and h.status == 'online' + ) + self._final_sla = healthy / total + final_comp = self._count_compromised(env) + self._new_infections = max(0, final_comp - self._switch_compromised) + isolated_comp = sum( + 1 for h in active if h.compromised_by != 'None' and h.status == 'isolated' + ) + self._contained_after = isolated_comp + growth_penalty = min(1.0, self._new_infections / max(total, 1)) + contain = isolated_comp / max(final_comp, 1) + score = max(0.0, 0.5 * self._final_sla + 0.5 * contain - 0.25 * growth_penalty) + score = min(1.0, score) + return DiagnosticResult( + diagnostic=self.name, + capability=self.capability, + policy='', + score=score, + details={ + 'switch_tick': self.switch_tick, + 'switch_compromised': self._switch_compromised, + 'final_compromised': final_comp, + 'new_infections': self._new_infections, + 'final_sla': self._final_sla, + 'ticks_used': ticks_used, + }, + ) + + @staticmethod + def _count_compromised(env) -> int: + return sum( + 1 + for h in env.global_state.all_hosts.values() + if h.subnet_cidr != PADDING_SUBNET and h.compromised_by != 'None' + ) diff --git a/netforge_rl/diagnostics/base.py b/netforge_rl/diagnostics/base.py index c1dfbff..a3a3878 100644 --- a/netforge_rl/diagnostics/base.py +++ b/netforge_rl/diagnostics/base.py @@ -55,10 +55,15 @@ def run_diagnostic(diag, policy, *, seed=0): ticks = 0 while env.agents and ticks < diag.max_ticks: - actions = {a: policy.act(env, a) for a in env.agents} + if hasattr(diag, 'act_for'): + actions = {a: diag.act_for(env, a, policy) for a in env.agents} + else: + actions = {a: policy.act(env, a) for a in env.agents} env.step(actions) ticks += 1 if diag.early_stop(env): break - return diag.score(env, ticks_used=ticks) + result = diag.score(env, ticks_used=ticks) + result.policy = getattr(policy, 'name', '') + return result diff --git a/netforge_rl/diagnostics/capability_card.py b/netforge_rl/diagnostics/capability_card.py index 43bc569..6ed9cdb 100644 --- a/netforge_rl/diagnostics/capability_card.py +++ b/netforge_rl/diagnostics/capability_card.py @@ -1,8 +1,8 @@ from __future__ import annotations import json +from collections.abc import Sequence from pathlib import Path -from typing import Optional, Sequence import numpy as np @@ -13,7 +13,7 @@ def capability_card( policy_factory, seeds: Sequence[int] = (0, 1, 2), - out_dir: Optional[str] = None, + out_dir: str | None = None, name: str = 'policy', ) -> dict: """Score a policy on every diagnostic capability, averaged over seeds.""" @@ -53,7 +53,7 @@ def _render_radar(card: dict, path: Path) -> None: matplotlib.use('Agg') import matplotlib.pyplot as plt - except Exception: + except (ImportError, OSError): return caps = list(card['capabilities']) diff --git a/netforge_rl/diagnostics/deception.py b/netforge_rl/diagnostics/deception.py new file mode 100644 index 0000000..4a608a4 --- /dev/null +++ b/netforge_rl/diagnostics/deception.py @@ -0,0 +1,65 @@ +from netforge_rl.diagnostics.base import Diagnostic, DiagnosticResult, flag_host + + +class DeceptionResistance(Diagnostic): + """One real foothold and one decoy both raise SIEM alerts. Score Blue on + isolating the real host without quarantining the trap (deception resistance).""" + + name = 'deception_resistance' + capability = 'deception' + scenario_type = 'ransomware' + max_ticks = 25 + + def __init__(self): + self._real_ip = None + self._decoy_ip = None + + def setup(self, env): + live = [ + ip + for ip, h in env.global_state.all_hosts.items() + if not ip.startswith('169.254.') and h.status == 'online' + ] + ordered = sorted(live) + self._real_ip = ordered[0] + self._decoy_ip = ordered[min(1, len(ordered) - 1)] + if self._decoy_ip == self._real_ip and len(ordered) > 1: + self._decoy_ip = ordered[1] + real = env.global_state.all_hosts[self._real_ip] + real.compromised_by = 'red_operator' + real.privilege = 'User' + real.decoy = 'inactive' + decoy = env.global_state.all_hosts[self._decoy_ip] + decoy.decoy = 'active' + decoy.contains_honeytokens = True + decoy.compromised_by = 'None' + decoy.privilege = 'None' + flag_host(env, self._real_ip, severity=8) + flag_host(env, self._decoy_ip, severity=8) + + def score(self, env, ticks_used): + real = env.global_state.all_hosts.get(self._real_ip) + decoy = env.global_state.all_hosts.get(self._decoy_ip) + isolated_real = real is not None and real.status == 'isolated' + isolated_decoy = decoy is not None and decoy.status == 'isolated' + if isolated_real and not isolated_decoy: + score = 1.0 + elif isolated_real and isolated_decoy: + score = 0.5 + elif isolated_decoy and not isolated_real: + score = 0.0 + else: + score = 0.25 + return DiagnosticResult( + diagnostic=self.name, + capability=self.capability, + policy='', + score=score, + details={ + 'real_ip': self._real_ip, + 'decoy_ip': self._decoy_ip, + 'isolated_real': isolated_real, + 'isolated_decoy': isolated_decoy, + 'ticks_used': ticks_used, + }, + ) diff --git a/netforge_rl/diagnostics/suite.py b/netforge_rl/diagnostics/suite.py index 277be29..fc07769 100644 --- a/netforge_rl/diagnostics/suite.py +++ b/netforge_rl/diagnostics/suite.py @@ -1,3 +1,5 @@ +from netforge_rl.diagnostics.adaptation import AdaptationShift +from netforge_rl.diagnostics.deception import DeceptionResistance from netforge_rl.diagnostics.delayed_telemetry import DelayedTelemetry from netforge_rl.diagnostics.false_positive import FalsePositiveRestraint from netforge_rl.diagnostics.memory_probe import MemoryProbe @@ -7,7 +9,7 @@ def all_diagnostics(): - """The full capability-probe suite, one instance per capability.""" + """Capability probes, one per card axis.""" return [ MemoryProbe(), NoisySIEM(), @@ -15,4 +17,6 @@ def all_diagnostics(): FalsePositiveRestraint(), OTKineticResponse(), TopologyShift(), + DeceptionResistance(), + AdaptationShift(), ] diff --git a/netforge_rl/diagnostics/wrapper.py b/netforge_rl/diagnostics/wrapper.py index f9cba55..5815219 100644 --- a/netforge_rl/diagnostics/wrapper.py +++ b/netforge_rl/diagnostics/wrapper.py @@ -1,16 +1,17 @@ -from typing import Any, List +from typing import Any + import numpy as np from pettingzoo.utils.wrappers import BaseParallelWrapper +from netforge_rl.core.graph_obs import build_oracle_graph from netforge_rl.core.observation import BaseObservation class OracleObservation(BaseObservation): - """Perfect information view of the network state (No fog-of-war).""" + """Privileged view. Diagnostics only.""" - def update_from_state(self, global_state: Any, _action_effects: List[Any]): + def update_from_state(self, global_state: Any, _action_effects: list[Any]): if global_state: - # Oracle knows ALL ips for ip, host in global_state.all_hosts.items(): self.visible_hosts[ip] = { 'state': 'compromised' @@ -26,13 +27,13 @@ def update_from_state(self, global_state: Any, _action_effects: List[Any]): arrival = log.get('arrival_tick', 0) if isinstance(log, dict) else 0 if arrival <= current_tick: self.siem_alerts.append(log) - self.network_telemetry['global_alert_level'] = 1.0 # Perfect clarity + self.network_telemetry['global_alert_level'] = 1.0 self.network_telemetry['total_isolated_subnets'] = 0 self.network_telemetry['active_alerts'] = len(self.siem_alerts) class DiagnosticsWrapper(BaseParallelWrapper): - """PettingZoo Wrapper to compute the empirical Information Asymmetry gap.""" + """Adds oracle obs, information_asymmetry, oracle_graph.""" def reset(self, seed=None, options=None): obs, infos = super().reset(seed=seed, options=options) @@ -62,5 +63,6 @@ def _inject_oracle(self, obs, infos): infos[agent]['oracle_obs'] = oracle_vec infos[agent]['information_asymmetry'] = distance + infos[agent]['oracle_graph'] = build_oracle_graph(global_state) return obs, infos diff --git a/netforge_rl/siem/correlator.py b/netforge_rl/siem/correlator.py index bcf7303..822b6c7 100644 --- a/netforge_rl/siem/correlator.py +++ b/netforge_rl/siem/correlator.py @@ -1,7 +1,5 @@ -from typing import List from netforge_rl.core.state import GlobalNetworkState - _STAGE_KEYWORDS = { 'RECON': [ 'port_scan', @@ -43,7 +41,7 @@ def __init__(self): def reset(self) -> None: self._last_correlated_tick = -1 - def correlate(self, global_state: 'GlobalNetworkState') -> List[str]: + def correlate(self, global_state: 'GlobalNetworkState') -> list[str]: """Returns incident log strings for this tick. Call once per step.""" if global_state.current_tick == self._last_correlated_tick: return [] @@ -85,10 +83,12 @@ def correlate(self, global_state: 'GlobalNetworkState') -> List[str]: if len(stages) >= 2 or confidence >= 0.4: incidents.append( ( - f'[INCIDENT] stage={stage_name} target={ip} ' - f'confidence={confidence:.2f} ' - f'stages_observed={",".join(sorted(stages))} ' - f'tick={global_state.current_tick}', + ( + f'[INCIDENT] stage={stage_name} target={ip} ' + f'confidence={confidence:.2f} ' + f'stages_observed={",".join(sorted(stages))} ' + f'tick={global_state.current_tick}' + ), subnet, ) ) diff --git a/netforge_rl/siem/export.py b/netforge_rl/siem/export.py index 502e0f8..494b174 100644 --- a/netforge_rl/siem/export.py +++ b/netforge_rl/siem/export.py @@ -2,10 +2,10 @@ import json import re +from collections.abc import Iterable from pathlib import Path -from typing import Iterable, Optional -# Re-maps NetForge's Windows/Sysmon event XML into OCSF-style JSON records. +# Sysmon/Windows XML → OCSF JSON. _EVENTID = re.compile(r'(\d+)') _COMPUTER = re.compile(r'(.*?)') _DATA = re.compile(r'(.*?)', re.DOTALL) @@ -25,7 +25,7 @@ } -def siem_to_ocsf(log_line: str, subnet: str, tick: Optional[int] = None) -> dict: +def siem_to_ocsf(log_line: str, subnet: str, tick: int | None = None) -> dict: """Map one NetForge SIEM log line to an OCSF-style event record.""" tags = [t.strip('[]') for t in re.findall(r'\[[A-Z_]+\]', log_line)] eid_match = _EVENTID.search(log_line) diff --git a/netforge_rl/topologies/network_generator.py b/netforge_rl/topologies/network_generator.py index 5dcae84..ea24fd6 100644 --- a/netforge_rl/topologies/network_generator.py +++ b/netforge_rl/topologies/network_generator.py @@ -14,10 +14,12 @@ def __init__( config_path: Optional[str] = None, max_active_hosts: Optional[int] = None, evaluation_mode: bool = False, + topology_spec: Optional[dict] = None, ): self.config_path = config_path self.max_active_hosts = max_active_hosts self.evaluation_mode = evaluation_mode + self.topology_spec = topology_spec def generate(self, seed: Optional[int] = None) -> GlobalNetworkState: """Generates the architecture.""" @@ -26,7 +28,9 @@ def generate(self, seed: Optional[int] = None) -> GlobalNetworkState: actual_seed = _EVAL_SEED_OFFSET + (seed % _EVAL_SEED_OFFSET) rng = random.Random(actual_seed) if actual_seed is not None else random.Random() - if self.config_path and Path(self.config_path).exists(): + if self.topology_spec: + state = self._load_from_mapping(self.topology_spec, rng) + elif self.config_path and Path(self.config_path).exists(): state = self._load_from_yaml(self.config_path, rng) else: state = self._generate_procedural(rng) @@ -216,7 +220,15 @@ def _configure_procedural_vision(self, state: GlobalNetworkState): state.update_knowledge(blue_id, host.ip) def _load_from_yaml(self, path: str, rng: random.Random) -> GlobalNetworkState: - """Load a fixed network topology from a YAML config file. + """Load a fixed network topology from a YAML config file.""" + import yaml # noqa: PLC0415 (lazy import to keep non-yaml envs lightweight) + + with open(path) as fh: + cfg = yaml.safe_load(fh) + return self._load_from_mapping(cfg, rng) + + def _load_from_mapping(self, cfg: dict, _rng: random.Random) -> GlobalNetworkState: + """Instantiate a topology from a dict (YAML file or inline scenario pack). YAML format example:: @@ -233,20 +245,15 @@ def _load_from_yaml(self, path: str, rng: random.Random) -> GlobalNetworkState: cvss_score: 7.5 decoy: "inactive" """ - import yaml # noqa: PLC0415 (lazy import to keep non-yaml envs lightweight) - - with open(path) as fh: - cfg = yaml.safe_load(fh) + from netforge_rl.core.state import Subnet # noqa: PLC0415 state = GlobalNetworkState() + topo = cfg.get('topology', cfg) - for sn in cfg.get('subnets', []): - from netforge_rl.core.state import Subnet # noqa: PLC0415 - + for sn in topo.get('subnets', []): state.add_subnet(Subnet(cidr=sn['cidr'], name=sn['name'])) - active_hosts = [] - for hcfg in cfg.get('hosts', []): + for hcfg in topo.get('hosts', []): host = Host( ip=hcfg['ip'], hostname=hcfg.get('hostname', hcfg['ip']), @@ -260,8 +267,14 @@ def _load_from_yaml(self, path: str, rng: random.Random) -> GlobalNetworkState: host.human_vulnerability_score = float( hcfg.get('human_vulnerability_score', 0.3) ) + host.contains_honeytokens = bool(hcfg.get('contains_honeytokens', False)) + if hcfg.get('is_domain_controller'): + host.is_domain_controller = True + if 'PLC' in host.os or hcfg.get('ot'): + host.system_integrity = hcfg.get('system_integrity', 'clean') + host.temperature = float(hcfg.get('temperature', 50.0)) + host.pressure = float(hcfg.get('pressure', 100.0)) state.register_host(host) - active_hosts.append(host) padding_needed = 100 - len(state.all_hosts) for p in range(padding_needed): diff --git a/tests/core/test_graph_obs.py b/tests/core/test_graph_obs.py index 11551ce..524b92d 100644 --- a/tests/core/test_graph_obs.py +++ b/tests/core/test_graph_obs.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import numpy as np import pytest -from netforge_rl.core.graph_obs import build_graph_observation +from netforge_rl.core.graph_obs import build_graph_observation, build_oracle_graph from netforge_rl.environment.graph_wrapper import GraphObservationWrapper from netforge_rl.environment.parallel_env import NetForgeRLEnv from netforge_rl.siem.pcap_synthesizer import NODE_DIM @@ -32,9 +34,28 @@ def test_padding_hosts_are_masked(env): @pytest.mark.fast def test_red_fog_of_war_hides_undiscovered_hosts(env): - blue = build_graph_observation(env.global_state, agent_id='blue_dmz') + oracle = build_oracle_graph(env.global_state) red = build_graph_observation(env.global_state, agent_id='red_operator') - assert red['node_mask'].sum() < blue['node_mask'].sum() + assert red['node_mask'].sum() < oracle['node_mask'].sum() + + +@pytest.mark.fast +def test_blue_belief_does_not_leak_compromise(env): + live = [ + ip + for ip, h in env.global_state.all_hosts.items() + if not ip.startswith('169.254.') + ] + ip = sorted(live)[0] + env.global_state.all_hosts[ip].compromised_by = 'red_operator' + env.global_state.all_hosts[ip].privilege = 'Root' + idx = sorted(env.global_state.all_hosts.keys()).index(ip) + belief = build_graph_observation( + env.global_state, agent_id='blue_dmz', mode='belief' + ) + oracle = build_graph_observation(env.global_state, mode='oracle') + assert oracle['node_features'][idx, 2] == 1.0 + assert belief['node_features'][idx, 0] == 0.0 @pytest.mark.fast @@ -46,8 +67,18 @@ def test_edges_only_between_visible_nodes(env): @pytest.mark.integration -def test_wrapper_injects_graph_into_infos(): +def test_wrapper_injects_belief_and_optional_oracle(): env = GraphObservationWrapper(NetForgeRLEnv({'scenario_type': 'ransomware'})) obs, infos = env.reset(seed=0) assert 'graph' in infos['blue_dmz'] + assert 'oracle_graph' not in infos['blue_dmz'] assert infos['blue_dmz']['graph']['node_features'].shape == (100, NODE_DIM) + + env = GraphObservationWrapper( + NetForgeRLEnv({'scenario_type': 'ransomware'}), oracle=True + ) + _, infos = env.reset(seed=0) + assert 'oracle_graph' in infos['blue_dmz'] + assert int(infos['blue_dmz']['oracle_graph']['node_mask'].sum()) >= int( + infos['blue_dmz']['graph']['node_mask'].sum() + ) diff --git a/tests/diagnostics/test_capability_card.py b/tests/diagnostics/test_capability_card.py index 472ea11..d7bf808 100644 --- a/tests/diagnostics/test_capability_card.py +++ b/tests/diagnostics/test_capability_card.py @@ -14,6 +14,8 @@ def test_card_has_all_capabilities(): 'precision', 'safety', 'generalization', + 'deception', + 'adaptation', } assert 0.0 <= card['overall'] <= 1.0 diff --git a/tests/diagnostics/test_capability_probes.py b/tests/diagnostics/test_capability_probes.py index 823062a..7b818c2 100644 --- a/tests/diagnostics/test_capability_probes.py +++ b/tests/diagnostics/test_capability_probes.py @@ -6,9 +6,9 @@ @pytest.mark.fast -def test_suite_has_six_distinct_capabilities(): +def test_suite_has_eight_distinct_capabilities(): diags = all_diagnostics() - assert len(diags) == 6 + assert len(diags) == 8 caps = {d.capability for d in diags} assert caps == { 'memory', @@ -17,11 +17,13 @@ def test_suite_has_six_distinct_capabilities(): 'precision', 'safety', 'generalization', + 'deception', + 'adaptation', } @pytest.mark.integration -@pytest.mark.parametrize('idx', range(6)) +@pytest.mark.parametrize('idx', range(8)) def test_each_probe_runs_and_scores_in_unit_range(idx): diag = all_diagnostics()[idx] result = run_diagnostic(diag, HeuristicBluePolicy(seed=0), seed=0) From e1e9ad3e6e2b342b9fab960e0c5d2ecef32684f3 Mon Sep 17 00:00:00 2001 From: elprofesoriqo Date: Sun, 20 Sep 2026 15:15:25 +0200 Subject: [PATCH 3/8] Add Arena 2026 splits, metrics, and named diagnostic questions. Rank on mission/SLA/FPs/security/CVaR instead of mean reward, and keep Windows CLI catalog text ASCII. Co-authored-by: Cursor --- benchmarks/configs/arena_2026_dev.yaml | 9 ++ benchmarks/results/leaderboard.json | 2 +- netforge_rl/arena/__init__.py | 45 ++++++ netforge_rl/arena/adversarial.py | 51 ++++++ netforge_rl/arena/agent_card.py | 71 ++++++++ netforge_rl/arena/constraints.py | 71 ++++++++ netforge_rl/arena/evaluate.py | 148 +++++++++++++++++ netforge_rl/arena/leaderboard.py | 50 ++++++ netforge_rl/arena/metrics.py | 138 ++++++++++++++++ netforge_rl/arena/ood.py | 66 ++++++++ netforge_rl/arena/questions/__init__.py | 19 +++ netforge_rl/arena/questions/catalog.py | 145 +++++++++++++++++ netforge_rl/arena/questions/policies.py | 9 ++ netforge_rl/arena/questions/run.py | 71 ++++++++ .../arena/questions/runners/__init__.py | 18 +++ netforge_rl/arena/questions/runners/belief.py | 58 +++++++ netforge_rl/arena/questions/runners/clock.py | 51 ++++++ .../arena/questions/runners/population.py | 22 +++ netforge_rl/arena/questions/runners/probes.py | 58 +++++++ netforge_rl/arena/questions/runners/reward.py | 33 ++++ netforge_rl/arena/questions/runners/risk.py | 34 ++++ netforge_rl/arena/questions/types.py | 18 +++ netforge_rl/arena/spec.py | 153 ++++++++++++++++++ netforge_rl/arena/submissions.py | 136 ++++++++++++++++ tests/arena/test_arena.py | 94 +++++++++++ tests/arena/test_infra.py | 96 +++++++++++ tests/arena/test_questions.py | 77 +++++++++ 27 files changed, 1742 insertions(+), 1 deletion(-) create mode 100644 benchmarks/configs/arena_2026_dev.yaml create mode 100644 netforge_rl/arena/__init__.py create mode 100644 netforge_rl/arena/adversarial.py create mode 100644 netforge_rl/arena/agent_card.py create mode 100644 netforge_rl/arena/constraints.py create mode 100644 netforge_rl/arena/evaluate.py create mode 100644 netforge_rl/arena/leaderboard.py create mode 100644 netforge_rl/arena/metrics.py create mode 100644 netforge_rl/arena/ood.py create mode 100644 netforge_rl/arena/questions/__init__.py create mode 100644 netforge_rl/arena/questions/catalog.py create mode 100644 netforge_rl/arena/questions/policies.py create mode 100644 netforge_rl/arena/questions/run.py create mode 100644 netforge_rl/arena/questions/runners/__init__.py create mode 100644 netforge_rl/arena/questions/runners/belief.py create mode 100644 netforge_rl/arena/questions/runners/clock.py create mode 100644 netforge_rl/arena/questions/runners/population.py create mode 100644 netforge_rl/arena/questions/runners/probes.py create mode 100644 netforge_rl/arena/questions/runners/reward.py create mode 100644 netforge_rl/arena/questions/runners/risk.py create mode 100644 netforge_rl/arena/questions/types.py create mode 100644 netforge_rl/arena/spec.py create mode 100644 netforge_rl/arena/submissions.py create mode 100644 tests/arena/test_arena.py create mode 100644 tests/arena/test_infra.py create mode 100644 tests/arena/test_questions.py diff --git a/benchmarks/configs/arena_2026_dev.yaml b/benchmarks/configs/arena_2026_dev.yaml new file mode 100644 index 0000000..d3fc874 --- /dev/null +++ b/benchmarks/configs/arena_2026_dev.yaml @@ -0,0 +1,9 @@ +arena: 2026 +split: dev +difficulty: medium +max_ticks: 200 +seeds: [9001, 9002, 9003, 9004, 9005] +scenarios: [ransomware, apt_espionage, cloud_hybrid, iot_grid, ot_stuxnet] +baselines: [random, heuristic-blue, killchain-red] +metrics: [mission_success, sla_uptime, security, false_positives, catastrophic_failure, cvar05_blue_return] +compute_budget_env_steps: 1000000 diff --git a/benchmarks/results/leaderboard.json b/benchmarks/results/leaderboard.json index aefa205..095c63b 100644 --- a/benchmarks/results/leaderboard.json +++ b/benchmarks/results/leaderboard.json @@ -52,7 +52,7 @@ { "name": "smoke", "team": "blue", - "submitted_at": "2026-07-11T12:39:14Z", + "submitted_at": "2026-09-20T11:00:29Z", "n_episodes": 1, "score": 50.0033, "mean_compromised": 0.0, diff --git a/netforge_rl/arena/__init__.py b/netforge_rl/arena/__init__.py new file mode 100644 index 0000000..39421e3 --- /dev/null +++ b/netforge_rl/arena/__init__.py @@ -0,0 +1,45 @@ +from netforge_rl.arena.adversarial import evaluate_population +from netforge_rl.arena.agent_card import build_agent_card, format_agent_card +from netforge_rl.arena.constraints import ConstrainedEnvWrapper +from netforge_rl.arena.evaluate import evaluate_split, run_episode +from netforge_rl.arena.leaderboard import ( + format_leaderboard, + load_leaderboard, + save_leaderboard, +) +from netforge_rl.arena.metrics import ArenaMetrics, aggregate, cvar +from netforge_rl.arena.ood import evaluate_ood +from netforge_rl.arena.questions import ( + QUESTIONS, + get_question, + list_questions, + run_question, +) +from netforge_rl.arena.spec import ARENA_2026, ARENA_NAME, ArenaSpec, current_arena +from netforge_rl.arena.submissions import evaluate_submission, load_submission + +__all__ = [ + 'ARENA_2026', + 'ARENA_NAME', + 'QUESTIONS', + 'ArenaMetrics', + 'ArenaSpec', + 'ConstrainedEnvWrapper', + 'aggregate', + 'build_agent_card', + 'current_arena', + 'cvar', + 'evaluate_ood', + 'evaluate_population', + 'evaluate_split', + 'evaluate_submission', + 'format_agent_card', + 'format_leaderboard', + 'get_question', + 'list_questions', + 'load_leaderboard', + 'load_submission', + 'run_episode', + 'run_question', + 'save_leaderboard', +] diff --git a/netforge_rl/arena/adversarial.py b/netforge_rl/arena/adversarial.py new file mode 100644 index 0000000..b14e797 --- /dev/null +++ b/netforge_rl/arena/adversarial.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence + +from netforge_rl.arena.evaluate import evaluate_split +from netforge_rl.arena.spec import ARENA_2026 +from netforge_rl.baselines.policies import ( + HeuristicRedPolicy, + KillChainRedPolicy, + RandomPolicy, +) + +DEFAULT_RED_POPULATION: dict[str, Callable] = { + 'random-red': lambda: RandomPolicy(seed=0), + 'heuristic-red': lambda: HeuristicRedPolicy(seed=0), + 'killchain-red': lambda: KillChainRedPolicy(seed=0), +} + + +def evaluate_population( + blue_policy_factory: Callable, + red_population: Mapping[str, Callable] | None = None, + *, + scenarios: Sequence[str] = ('ransomware',), + seeds: Sequence[int] = (0, 1), + max_ticks: int = 60, +) -> dict: + """Evaluate one Blue policy against a population of Red opponents.""" + pool = dict(red_population or DEFAULT_RED_POPULATION) + per_red = {} + for name, factory in pool.items(): + per_red[name] = evaluate_split( + factory, + blue_policy_factory, + split='dev', + spec=ARENA_2026, + scenarios=scenarios, + seeds=seeds, + max_ticks=max_ticks, + ) + sla = [ + v['overall']['mean']['sla_uptime'] + for v in per_red.values() + if v['overall']['n_episodes'] + ] + return { + 'n_opponents': len(pool), + 'per_red': per_red, + 'mean_sla_across_red': round(sum(sla) / max(len(sla), 1), 4), + 'worst_sla_across_red': round(min(sla) if sla else 0.0, 4), + } diff --git a/netforge_rl/arena/agent_card.py b/netforge_rl/arena/agent_card.py new file mode 100644 index 0000000..43a97c2 --- /dev/null +++ b/netforge_rl/arena/agent_card.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from collections.abc import Sequence + +from netforge_rl.arena.evaluate import evaluate_split +from netforge_rl.arena.spec import ARENA_2026 +from netforge_rl.diagnostics.capability_card import capability_card + + +def build_agent_card( + policy_factory, + *, + name: str = 'policy', + seeds: Sequence[int] = (0, 1), + opponent_factory=None, + run_arena: bool = True, + out_dir: str | None = None, +) -> dict: + """Capability probes plus optional Arena-dev metrics.""" + card = capability_card(policy_factory, seeds=seeds, out_dir=out_dir, name=name) + caps = card['capabilities'] + fpr = round(max(0.0, 1.0 - float(caps.get('precision', 1.0))), 4) + headlines = { + 'Safety': caps.get('safety', 0.0), + 'OOD Generalization': caps.get('generalization', 0.0), + 'Memory': caps.get('memory', 0.0), + 'False Positive Rate': fpr, + 'Temporal': caps.get('temporal', 0.0), + 'Deception Resistance': caps.get('deception', 0.0), + 'Adaptation': caps.get('adaptation', 0.0), + 'Attention': caps.get('attention', 0.0), + } + card['headlines'] = headlines + card['false_positive_rate'] = fpr + + if run_arena: + from netforge_rl.baselines.policies import KillChainRedPolicy + + red_factory = opponent_factory or (lambda: KillChainRedPolicy(seed=0)) + arena = evaluate_split( + red_factory, + policy_factory, + split='dev', + scenarios=('ransomware',), + seeds=tuple(seeds)[:2], + max_ticks=40, + spec=ARENA_2026, + ) + card['arena_dev'] = arena['overall'] + mean = arena['overall'].get('mean', {}) + card['headlines']['Mission Success'] = mean.get('mission_success', 0.0) + card['headlines']['CVaR Return'] = ( + arena['overall'].get('cvar05', {}).get('blue_return', 0.0) + ) + if out_dir: + import json + from pathlib import Path + + path = Path(out_dir) / f'{name}_card.json' + path.write_text(json.dumps(card, indent=2)) + return card + + +def format_agent_card(card: dict) -> str: + lines = [f'Agent Card: {card.get("name", "policy")}'] + overall = card.get('overall') + if overall is not None: + lines.append(f' Overall {overall:.2f}') + for label, value in card.get('headlines', {}).items(): + lines.append(f' {label:<22} {float(value):.2f}') + return '\n'.join(lines) diff --git a/netforge_rl/arena/constraints.py b/netforge_rl/arena/constraints.py new file mode 100644 index 0000000..9d11570 --- /dev/null +++ b/netforge_rl/arena/constraints.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from collections.abc import Mapping + +from pettingzoo.utils.wrappers import BaseParallelWrapper + +DEFAULT_CONSTRAINTS: dict[str, dict] = { + 'sla_uptime': {'min': 0.90}, + 'false_positives_total': {'max': 2.0}, + 'disruption': {'max': 0.40}, + 'catastrophic_failure': {'max': 0.0}, +} + + +class ConstrainedEnvWrapper(BaseParallelWrapper): + """Lagrangian costs from Blue info (SLA, FPs, disruption, kinetic).""" + + def __init__( + self, + env, + constraints: Mapping[str, dict] | None = None, + lambda_cost: float = 0.0, + team: str = 'blue', + ): + super().__init__(env) + self.constraints = dict(constraints or DEFAULT_CONSTRAINTS) + self.lambda_cost = float(lambda_cost) + self.team = team + + def reset(self, seed=None, options=None): + obs, infos = super().reset(seed=seed, options=options) + return obs, self._annotate(infos, {}) + + def step(self, actions): + obs, rewards, term, trunc, infos = super().step(actions) + costs = self._costs(infos) + if self.lambda_cost: + rewards = { + agent: float(rewards[agent]) + - ( + self.lambda_cost * costs.get(agent, 0.0) + if self.team in agent + else 0.0 + ) + for agent in rewards + } + return obs, rewards, term, trunc, self._annotate(infos, costs) + + def _costs(self, infos: dict) -> dict[str, float]: + costs = {} + for agent, info in infos.items(): + cost = 0.0 + violations = {} + for key, spec in self.constraints.items(): + value = float(info.get(key, 0.0)) + if 'min' in spec and value < spec['min']: + gap = spec['min'] - value + violations[key] = gap + cost += gap + if 'max' in spec and value > spec['max']: + gap = value - spec['max'] + violations[key] = gap + cost += gap + info['constraint_violations'] = violations + costs[agent] = cost + return costs + + def _annotate(self, infos: dict, costs: dict) -> dict: + for agent, info in infos.items(): + info['constraint_cost'] = float(costs.get(agent, 0.0)) + return infos diff --git a/netforge_rl/arena/evaluate.py b/netforge_rl/arena/evaluate.py new file mode 100644 index 0000000..9cac677 --- /dev/null +++ b/netforge_rl/arena/evaluate.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field + +import numpy as np + +from netforge_rl.arena.metrics import ArenaMetrics, aggregate, metrics_from_env +from netforge_rl.arena.spec import ARENA_2026, ArenaSpec, SplitName +from netforge_rl.environment.graph_wrapper import GraphObservationWrapper +from netforge_rl.scenarios.yaml_dsl import make_env_from_spec, resolve_scenario + + +@dataclass +class EpisodeRecord: + scenario: str + seed: int + split: str + metrics: ArenaMetrics + extras: dict = field(default_factory=dict) + + +def _act(policy, env, agent_id, obs) -> np.ndarray: + fn = getattr(policy, 'act', None) + if fn is None: + return np.zeros(len(env.action_space(agent_id).nvec), dtype=np.int64) + mode = getattr(policy, '_nf_act', None) + if mode == 'obs': + return np.asarray(fn(obs, agent_id), dtype=np.int64) + if mode != 'env': + try: + out = fn(env, agent_id) + policy._nf_act = 'env' + return np.asarray(out, dtype=np.int64) + except TypeError: + policy._nf_act = 'obs' + return np.asarray(fn(obs, agent_id), dtype=np.int64) + return np.asarray(fn(env, agent_id), dtype=np.int64) + + +def run_episode( + red_policy, + blue_policy, + *, + scenario: str = 'ransomware', + seed: int = 0, + max_ticks: int = 200, + evaluation_mode: bool = False, + difficulty: str = 'medium', + graph_obs: bool = True, + env_overrides: dict | None = None, + wrap_env: Callable | None = None, +) -> EpisodeRecord: + spec = resolve_scenario(scenario) + overrides = dict(env_overrides or {}) + overrides.setdefault('max_ticks', max_ticks) + overrides.setdefault('evaluation_mode', evaluation_mode) + env = make_env_from_spec(spec, difficulty=difficulty, seed=None, **overrides) + if graph_obs: + env = GraphObservationWrapper(env, oracle=False) + if wrap_env is not None: + env = wrap_env(env) + obs, _ = env.reset(seed=seed) + for pol in (red_policy, blue_policy): + if hasattr(pol, 'reset'): + pol.reset() + if hasattr(pol, 'bind_env'): + pol.bind_env(env) + + blue_return = red_return = 0.0 + last_info: dict = {} + while env.agents: + actions = {} + for agent_id in env.agents: + policy = red_policy if 'red' in agent_id else blue_policy + actions[agent_id] = _act(policy, env, agent_id, obs.get(agent_id, {})) + obs, rewards, term, trunc, last_info = env.step(actions) + for agent_id, reward in rewards.items(): + if 'red' in agent_id: + red_return += float(reward) + else: + blue_return += float(reward) + if all(term.values()) or all(trunc.values()): + break + + inner = env.unwrapped if hasattr(env, 'unwrapped') else env + metrics = metrics_from_env(inner, last_info, blue_return, red_return) + return EpisodeRecord( + scenario=spec.name, + seed=seed, + split='custom', + metrics=metrics, + ) + + +def evaluate_split( + red_policy_factory: Callable, + blue_policy_factory: Callable, + *, + split: SplitName = 'dev', + spec: ArenaSpec | None = None, + scenarios: Sequence[str] | None = None, + seeds: Sequence[int] | None = None, + max_ticks: int | None = None, + difficulty: str | None = None, + graph_obs: bool = True, + env_overrides: dict | None = None, + wrap_env: Callable | None = None, +) -> dict: + arena = spec or ARENA_2026 + split_info = arena.split(split) + scenarios = tuple(scenarios or arena.scenarios) + seeds = tuple(seeds or split_info.seeds) + max_ticks = max_ticks if max_ticks is not None else arena.max_ticks + difficulty = difficulty or arena.difficulty + + records: list[EpisodeRecord] = [] + per_scenario: dict[str, list[ArenaMetrics]] = {s: [] for s in scenarios} + for scenario in scenarios: + for seed in seeds: + rec = run_episode( + red_policy_factory(), + blue_policy_factory(), + scenario=scenario, + seed=int(seed), + max_ticks=max_ticks, + evaluation_mode=split_info.evaluation_mode, + difficulty=difficulty, + graph_obs=graph_obs, + env_overrides=env_overrides, + wrap_env=wrap_env, + ) + rec.split = split + records.append(rec) + per_scenario[scenario].append(rec.metrics) + + overall = aggregate([r.metrics for r in records]) + return { + 'arena': arena.version, + 'split': split, + 'public': split_info.public, + 'n_episodes': len(records), + 'scenarios': list(scenarios), + 'overall': overall.as_dict(), + 'per_scenario': { + name: aggregate(vals).as_dict() for name, vals in per_scenario.items() + }, + } diff --git a/netforge_rl/arena/leaderboard.py b/netforge_rl/arena/leaderboard.py new file mode 100644 index 0000000..0323dbe --- /dev/null +++ b/netforge_rl/arena/leaderboard.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import json +from pathlib import Path + + +def save_leaderboard(entries: list[dict], path: str | Path) -> Path: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + ranked = sorted( + entries, + key=lambda e: ( + e.get('overall', {}).get('mean', {}).get('mission_success', 0.0), + -e.get('overall', {}).get('mean', {}).get('false_positives', 99.0), + e.get('overall', {}).get('mean', {}).get('security', 0.0), + ), + reverse=True, + ) + for i, entry in enumerate(ranked, 1): + entry['rank'] = i + path.write_text(json.dumps(ranked, indent=2)) + return path + + +def load_leaderboard(path: str | Path) -> list[dict]: + path = Path(path) + if not path.exists(): + return [] + return json.loads(path.read_text()) + + +def format_leaderboard(entries: list[dict]) -> str: + if not entries: + return '(empty leaderboard)' + lines = [ + f'{"rank":<6}{"name":<22}{"split":<8}{"mission":>8}{"SLA":>8}{"FP":>8}{"CVaR":>10}' + ] + for e in entries: + mean = e.get('overall', {}).get('mean', {}) + cvar = e.get('overall', {}).get('cvar05', {}) + lines.append( + f'{e.get("rank", 0):<6}' + f'{e.get("name", "?"):<22}' + f'{e.get("split", "?"):<8}' + f'{mean.get("mission_success", 0):8.2f}' + f'{mean.get("sla_uptime", 0):8.2f}' + f'{mean.get("false_positives", 0):8.2f}' + f'{cvar.get("blue_return", 0):10.2f}' + ) + return '\n'.join(lines) diff --git a/netforge_rl/arena/metrics.py b/netforge_rl/arena/metrics.py new file mode 100644 index 0000000..355ded0 --- /dev/null +++ b/netforge_rl/arena/metrics.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass, field + +import numpy as np + +from netforge_rl.environment.constants import PADDING_SUBNET + + +@dataclass +class ArenaMetrics: + """Per-episode competition metrics. Reward is recorded, not the ranking axis.""" + + mission_success: float + disruption: float + security: float + false_positives: float + exfiltration: float + compromised_hosts: float + isolated_hosts: float + sla_uptime: float + catastrophic_failure: float + blue_return: float + red_return: float + steps: int + detection_rate: float = 0.0 + mttc: float = 0.0 + deception_efficacy: float = 0.0 + + def as_dict(self) -> dict: + return asdict(self) + + +def metrics_from_env( + env, last_info: dict, blue_return: float, red_return: float +) -> ArenaMetrics: + sample = next((v for k, v in last_info.items() if 'blue' in k), {}) + if not sample: + sample = next(iter(last_info.values()), {}) + active = [ + h + for h in env.global_state.all_hosts.values() + if h.subnet_cidr != PADDING_SUBNET + ] + total = max(len(active), 1) + kinetic = any( + getattr(h, 'system_integrity', 'clean') == 'kinetic_destruction' for h in active + ) + sla = float(sample.get('SLA_Uptime_Percentage', 1.0)) + compromised = float(sample.get('compromised_hosts', 0.0)) + isolated = float(sample.get('isolated_hosts', 0.0)) + fps = float(sample.get('false_positives_total', sample.get('false_positives', 0.0))) + return ArenaMetrics( + mission_success=float((not kinetic) and sla >= 0.5), + disruption=isolated / total, + security=max(0.0, 1.0 - compromised / total), + false_positives=fps, + exfiltration=float(sample.get('Total_Exfiltrated_Data', 0.0)), + compromised_hosts=compromised, + isolated_hosts=isolated, + sla_uptime=sla, + catastrophic_failure=float(kinetic), + blue_return=float(blue_return), + red_return=float(red_return), + steps=int(getattr(env, 'current_tick', 0)), + detection_rate=float(sample.get('detection_rate', 0.0)), + mttc=float(sample.get('MTTC', 0.0)), + deception_efficacy=float(sample.get('deception_efficacy', 0.0)), + ) + + +def cvar(values: list[float], alpha: float = 0.05) -> float: + """Lower-tail Conditional Value at Risk (expected value in the worst α).""" + if not values: + return 0.0 + arr = np.sort(np.asarray(values, dtype=float)) + k = max(1, int(np.ceil(alpha * len(arr)))) + return float(arr[:k].mean()) + + +def worst_case(values: list[float]) -> float: + return float(min(values)) if values else 0.0 + + +def tail_failure_rate(flags: list[float], threshold: float = 0.5) -> float: + if not flags: + return 0.0 + arr = np.asarray(flags, dtype=float) + return float(np.mean(arr >= threshold)) + + +@dataclass +class AggregateReport: + n_episodes: int + mean: dict = field(default_factory=dict) + std: dict = field(default_factory=dict) + cvar05: dict = field(default_factory=dict) + worst: dict = field(default_factory=dict) + catastrophic_rate: float = 0.0 + + def as_dict(self) -> dict: + return asdict(self) + + +_SCALAR_KEYS = ( + 'mission_success', + 'disruption', + 'security', + 'false_positives', + 'exfiltration', + 'compromised_hosts', + 'isolated_hosts', + 'sla_uptime', + 'catastrophic_failure', + 'blue_return', + 'red_return', + 'detection_rate', + 'mttc', +) + + +def aggregate(metrics: list[ArenaMetrics], alpha: float = 0.05) -> AggregateReport: + if not metrics: + return AggregateReport(n_episodes=0) + series = {k: [float(getattr(m, k)) for m in metrics] for k in _SCALAR_KEYS} + return AggregateReport( + n_episodes=len(metrics), + mean={k: round(float(np.mean(v)), 4) for k, v in series.items()}, + std={ + k: round(float(np.std(v, ddof=1) if len(v) > 1 else 0.0), 4) + for k, v in series.items() + }, + cvar05={k: round(cvar(v, alpha), 4) for k, v in series.items()}, + worst={k: round(worst_case(v), 4) for k, v in series.items()}, + catastrophic_rate=round( + tail_failure_rate(series['catastrophic_failure'], 0.5), 4 + ), + ) diff --git a/netforge_rl/arena/ood.py b/netforge_rl/arena/ood.py new file mode 100644 index 0000000..08ce458 --- /dev/null +++ b/netforge_rl/arena/ood.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence + +from netforge_rl.arena.evaluate import evaluate_split +from netforge_rl.arena.spec import ARENA_2026 + +OOD_TASKS: tuple[dict, ...] = ( + { + 'name': 'unseen_topology_size', + 'env_overrides': {'max_active_hosts': 40}, + 'scenarios': ('ransomware',), + 'description': 'Train on medium, evaluate on a larger live host count.', + }, + { + 'name': 'unseen_attacker', + 'red': 'killchain-red', + 'scenarios': ('apt_espionage',), + 'description': 'Kill-chain Red if the policy was trained vs heuristic.', + }, + { + 'name': 'unseen_telemetry', + 'env_overrides': {'log_latency': 8}, + 'scenarios': ('ransomware',), + 'description': 'Delayed SIEM relative to the training latency.', + }, + { + 'name': 'unseen_scenario', + 'scenarios': ('ot_stuxnet',), + 'description': 'Held-out scenario family (kinetic OT).', + }, +) + + +def evaluate_ood( + red_policy_factory: Callable, + blue_policy_factory: Callable, + *, + tasks: Sequence[dict] = OOD_TASKS, + seeds: Sequence[int] = (0, 1), + max_ticks: int = 60, + difficulty: str = 'medium', +) -> dict: + """Train-distribution vs unseen size / attacker / telemetry / scenario.""" + results = {} + for task in tasks: + red_factory = red_policy_factory + if task.get('red') == 'killchain-red': + from netforge_rl.baselines.policies import KillChainRedPolicy + + def red_factory(): + return KillChainRedPolicy(seed=0) + + results[task['name']] = evaluate_split( + red_factory, + blue_policy_factory, + split='dev', + spec=ARENA_2026, + scenarios=task.get('scenarios'), + seeds=seeds, + max_ticks=max_ticks, + difficulty=difficulty, + env_overrides=task.get('env_overrides'), + ) + results[task['name']]['description'] = task['description'] + return {'tasks': results} diff --git a/netforge_rl/arena/questions/__init__.py b/netforge_rl/arena/questions/__init__.py new file mode 100644 index 0000000..2821a50 --- /dev/null +++ b/netforge_rl/arena/questions/__init__.py @@ -0,0 +1,19 @@ +from netforge_rl.arena.questions.catalog import QUESTIONS +from netforge_rl.arena.questions.run import ( + format_catalog, + get_question, + list_questions, + run_all_runnable, + run_question, +) +from netforge_rl.arena.questions.types import ResearchQuestion + +__all__ = [ + 'QUESTIONS', + 'ResearchQuestion', + 'format_catalog', + 'get_question', + 'list_questions', + 'run_all_runnable', + 'run_question', +] diff --git a/netforge_rl/arena/questions/catalog.py b/netforge_rl/arena/questions/catalog.py new file mode 100644 index 0000000..cef7621 --- /dev/null +++ b/netforge_rl/arena/questions/catalog.py @@ -0,0 +1,145 @@ +from netforge_rl.arena.questions.types import ResearchQuestion + +QUESTIONS: tuple[ResearchQuestion, ...] = ( + ResearchQuestion( + id='belief-vs-oracle', + title='Belief graph vs privileged oracle', + question=( + 'How much of the true compromise state is missing from Blue SIEM ' + 'belief graph, and does an oracle leak change a telemetry-only policy?' + ), + why='POMDP / privileged-information ablation. Oracle graphs are diagnostic-only.', + protocol='Plant a silent compromise, compare belief vs oracle node masks and labels.', + metrics=('visible_ratio', 'compromise_recall', 'heuristic_sla'), + status='runnable', + note='Train Blue on SIEM belief. Oracle graphs are diagnostics only.', + ), + ResearchQuestion( + id='reward-design', + title='Does the reward isolate the network?', + question=( + 'How does swapping Blue reward (SLA-only, FP-penalized, default) ' + 'change isolation, false positives, and mission success on the same seeds?' + ), + why='A reward that pays for isolation will take the whole net down.', + protocol='RewardDesignWrapper variants, identical seeds, Arena metrics.', + metrics=('isolated_hosts', 'false_positives', 'sla_uptime', 'mission_success'), + status='runnable', + note='Check isolation and FPs, not only return.', + ), + ResearchQuestion( + id='red-population', + title='One Blue vs a Red population', + question=( + 'Does a defender that looks strong vs kill-chain Red still hold SLA ' + 'against random and heuristic attackers?' + ), + why='A defender that only beats one scripted Red overfits.', + protocol='evaluate_population with random / heuristic / kill-chain Red.', + metrics=('mean_sla_across_red', 'worst_sla_across_red'), + status='runnable', + note='Score against the worst Red in the pool, not only the training opponent.', + ), + ResearchQuestion( + id='risk-sensitive', + title='Mean return vs CVaR vs worst case', + question=( + 'Is a high mean Blue return hiding a heavy lower tail or kinetic failures?' + ), + why='Safety-critical defense cannot rank on mean reward.', + protocol='Several seeds; report mean, CVaR-5%, worst-case, catastrophic rate.', + metrics=('mean_blue_return', 'cvar05', 'worst_case', 'catastrophic_rate'), + status='runnable', + note='CVaR and kinetic failures sit next to mean return.', + ), + ResearchQuestion( + id='time-mode', + title='Event-driven vs fixed-step time', + question=( + 'Does the same policy produce different SLA and tick counts under ' + 'durative event time versus a fixed tick clock?' + ), + why='Exploits and isolates take ticks; the clock is part of the gym.', + protocol='Identical seed; env_overrides time_mode=event vs fixed.', + metrics=('steps', 'sla_uptime', 'blue_return'), + status='runnable', + note='Compare event-driven time to a fixed tick clock.', + ), + ResearchQuestion( + id='ood-topology', + title='Train-size vs larger live network', + question='Does a medium-topology defender keep SLA when max_active_hosts grows?', + why='Generalization over topology size, not a new host type.', + protocol='Same policy; default hosts vs max_active_hosts=40.', + metrics=('sla_uptime', 'security', 'compromised_hosts'), + status='runnable', + note='Hold out larger live networks than training.', + ), + ResearchQuestion( + id='adaptation', + title='Mid-episode attacker shift', + question=( + 'After Red switches from kill-chain to random spread, does Blue contain ' + 'new infections or keep fighting the first strategy?' + ), + why='Adaptation is a capability, not a training trick.', + protocol='AdaptationShift diagnostic probe.', + metrics=('score', 'new_infections', 'contained_after'), + status='runnable', + note='Red can change strategy mid-episode.', + ), + ResearchQuestion( + id='deception', + title='Honeypot / decoy resistance', + question='Does Blue waste isolation on decoys, or ignore them and miss real hosts?', + why='Deception is a Blue capability probe, not more Red exploits.', + protocol='DeceptionResistance diagnostic.', + metrics=('score',), + status='runnable', + note='Honeypots and decoys can waste Blue isolation.', + ), + ResearchQuestion( + id='constrained-sla', + title='Defense under SLA / FP constraints', + question=( + 'Can Blue keep security while respecting SLA >= 0.9 and a false-positive cap?' + ), + why='Constrained RL belongs in the gym; Lagrangian/CPO stay outside.', + protocol='ConstrainedEnvWrapper costs on the same seeds as unconstrained.', + metrics=('constraint_cost', 'sla_uptime', 'false_positives', 'security'), + status='runnable', + note='Keep SLA / FP costs on the same seeds as unconstrained.', + ), + ResearchQuestion( + id='gnn-belief', + title='GNN on belief graphs (protocol)', + question='Does GNN-MAPPO trained on SIEM belief transfer to hidden topologies?', + why='The gym ships belief graphs; you bring the GNN trainer.', + protocol='Train on infos[blue][graph]; eval Arena hidden; oracle=cheat ablation.', + metrics=('mission_success', 'ood', 'oracle_gap'), + status='protocol', + note='Do not vendor MAPPO into this tree.', + ), + ResearchQuestion( + id='offline-to-online', + title='Offline SOC logs then online fine-tune (protocol)', + question='Does pretraining on mixed trajectories beat learning from scratch?', + why='Offline RL needs frozen datasets, not a bigger simulator.', + protocol='netforge collect + HDF5/npz; fine-tune; Arena hidden.', + metrics=('mission_success', 'false_positives', 'cvar05'), + status='protocol', + note='Dump mixed trajectories, then fine-tune online.', + ), + ResearchQuestion( + id='league-psro', + title='Self-play league / PSRO (protocol)', + question='Does a Blue that best-responds to a Red league beat a single opponent?', + why='MARL as a game. Payoff matrices ship; PSRO trainers do not.', + protocol='netforge_rl.league.payoff_matrix; historical checkpoints as the pool.', + metrics=('worst_red', 'empirical_game'), + status='protocol', + note='League trainers live outside this repo.', + ), +) + +BY_ID = {q.id: q for q in QUESTIONS} diff --git a/netforge_rl/arena/questions/policies.py b/netforge_rl/arena/questions/policies.py new file mode 100644 index 0000000..8e1b1a6 --- /dev/null +++ b/netforge_rl/arena/questions/policies.py @@ -0,0 +1,9 @@ +from netforge_rl.baselines.policies import HeuristicBluePolicy, KillChainRedPolicy + + +def blue() -> HeuristicBluePolicy: + return HeuristicBluePolicy(seed=0) + + +def killchain() -> KillChainRedPolicy: + return KillChainRedPolicy(seed=0) diff --git a/netforge_rl/arena/questions/run.py b/netforge_rl/arena/questions/run.py new file mode 100644 index 0000000..40dbfa5 --- /dev/null +++ b/netforge_rl/arena/questions/run.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from collections.abc import Sequence + +from netforge_rl.arena.questions.catalog import BY_ID, QUESTIONS +from netforge_rl.arena.questions.runners import RUNNERS +from netforge_rl.arena.questions.types import ResearchQuestion + + +def get_question(qid: str) -> ResearchQuestion: + try: + return BY_ID[qid.strip().lower()] + except KeyError: + raise KeyError( + f'Unknown question {qid!r}. Choose from: {", ".join(BY_ID)}' + ) from None + + +def list_questions(*, runnable_only: bool = False) -> list[dict]: + qs = ( + QUESTIONS + if not runnable_only + else (q for q in QUESTIONS if q.status == 'runnable') + ) + return [q.as_dict() for q in qs] + + +def run_question( + qid: str, + *, + seeds: Sequence[int] = (0,), + max_ticks: int = 20, +) -> dict: + q = get_question(qid) + if q.status != 'runnable': + return { + 'id': q.id, + 'status': 'protocol', + 'question': q.question, + 'protocol': q.protocol, + 'note': q.note, + } + payload = RUNNERS[q.id](seeds=tuple(seeds), max_ticks=int(max_ticks)) + return { + 'id': q.id, + 'title': q.title, + 'question': q.question, + 'status': 'runnable', + 'answer': payload, + } + + +def run_all_runnable( + *, + seeds: Sequence[int] = (0,), + max_ticks: int = 12, + ids: Sequence[str] | None = None, +) -> list[dict]: + chosen = ( + list(ids) + if ids is not None + else [q.id for q in QUESTIONS if q.status == 'runnable'] + ) + return [run_question(qid, seeds=seeds, max_ticks=max_ticks) for qid in chosen] + + +def format_catalog() -> str: + lines = ['id status question'] + for q in QUESTIONS: + lines.append(f'{q.id:<24}{q.status:<11}{q.question}') + return '\n'.join(lines) diff --git a/netforge_rl/arena/questions/runners/__init__.py b/netforge_rl/arena/questions/runners/__init__.py new file mode 100644 index 0000000..4bfed32 --- /dev/null +++ b/netforge_rl/arena/questions/runners/__init__.py @@ -0,0 +1,18 @@ +from netforge_rl.arena.questions.runners.belief import run_belief_vs_oracle +from netforge_rl.arena.questions.runners.clock import run_ood_topology, run_time_mode +from netforge_rl.arena.questions.runners.population import run_red_population +from netforge_rl.arena.questions.runners.probes import run_constrained, run_probe +from netforge_rl.arena.questions.runners.reward import run_reward_design +from netforge_rl.arena.questions.runners.risk import run_risk_sensitive + +RUNNERS = { + 'belief-vs-oracle': run_belief_vs_oracle, + 'reward-design': run_reward_design, + 'red-population': run_red_population, + 'risk-sensitive': run_risk_sensitive, + 'time-mode': run_time_mode, + 'ood-topology': run_ood_topology, + 'adaptation': lambda **kw: run_probe('adaptation', **kw), + 'deception': lambda **kw: run_probe('deception', **kw), + 'constrained-sla': run_constrained, +} diff --git a/netforge_rl/arena/questions/runners/belief.py b/netforge_rl/arena/questions/runners/belief.py new file mode 100644 index 0000000..cf1692c --- /dev/null +++ b/netforge_rl/arena/questions/runners/belief.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + +from netforge_rl.arena.evaluate import run_episode +from netforge_rl.arena.questions.policies import blue, killchain + + +def run_belief_vs_oracle(*, seeds: Sequence[int], max_ticks: int) -> dict: + from netforge_rl.core.graph_obs import build_graph_observation, build_oracle_graph + from netforge_rl.environment.graph_wrapper import GraphObservationWrapper + from netforge_rl.environment.parallel_env import NetForgeRLEnv + + seed = int(seeds[0]) + env = GraphObservationWrapper( + NetForgeRLEnv({'scenario_type': 'ransomware', 'max_ticks': max_ticks}), + oracle=True, + ) + env.reset(seed=seed) + live = [ + ip + for ip, h in env.unwrapped.global_state.all_hosts.items() + if not ip.startswith('169.254.') + ] + planted = min(live) + host = env.unwrapped.global_state.all_hosts[planted] + host.compromised_by = 'red_operator' + host.privilege = 'User' + + gs = env.unwrapped.global_state + belief = build_graph_observation(gs, agent_id='blue_dmz') + oracle = build_oracle_graph(gs) + vis_b = int(belief['node_mask'].sum()) + vis_o = max(int(oracle['node_mask'].sum()), 1) + oracle_comp = set(np.where(oracle['node_features'][:, 2] > 0.5)[0].tolist()) + belief_comp = set(np.where(belief['node_features'][:, 0] > 0.5)[0].tolist()) + recall = len(oracle_comp & belief_comp) / max(len(oracle_comp), 1) + + rec = run_episode( + killchain(), + blue(), + scenario='ransomware', + seed=seed, + max_ticks=max_ticks, + graph_obs=True, + ) + return { + 'planted_host': planted, + 'visible_belief': vis_b, + 'visible_oracle': vis_o, + 'visible_ratio': round(vis_b / vis_o, 4), + 'oracle_compromised_nodes': len(oracle_comp), + 'belief_flagged_nodes': len(belief_comp), + 'compromise_recall': round(recall, 4), + 'heuristic_episode': rec.metrics.as_dict(), + } diff --git a/netforge_rl/arena/questions/runners/clock.py b/netforge_rl/arena/questions/runners/clock.py new file mode 100644 index 0000000..89d373b --- /dev/null +++ b/netforge_rl/arena/questions/runners/clock.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +from collections.abc import Sequence + +from netforge_rl.arena.evaluate import run_episode +from netforge_rl.arena.questions.policies import blue, killchain + + +def run_time_mode(*, seeds: Sequence[int], max_ticks: int) -> dict: + seed = int(seeds[0]) + out = {} + for mode in ('event', 'fixed'): + rec = run_episode( + killchain(), + blue(), + scenario='ransomware', + seed=seed, + max_ticks=max_ticks, + env_overrides={'time_mode': mode}, + ) + out[mode] = { + 'steps': rec.metrics.steps, + 'sla_uptime': rec.metrics.sla_uptime, + 'blue_return': round(rec.metrics.blue_return, 4), + 'security': rec.metrics.security, + } + return {'seed': seed, 'modes': out} + + +def run_ood_topology(*, seeds: Sequence[int], max_ticks: int) -> dict: + seed = int(seeds[0]) + out = {} + for name, overrides in ( + ('medium_default', None), + ('larger_hosts', {'max_active_hosts': 40}), + ): + rec = run_episode( + killchain(), + blue(), + scenario='ransomware', + seed=seed, + max_ticks=max_ticks, + env_overrides=overrides, + ) + out[name] = { + 'sla_uptime': rec.metrics.sla_uptime, + 'security': rec.metrics.security, + 'compromised_hosts': rec.metrics.compromised_hosts, + 'isolated_hosts': rec.metrics.isolated_hosts, + } + return {'seed': seed, 'topologies': out} diff --git a/netforge_rl/arena/questions/runners/population.py b/netforge_rl/arena/questions/runners/population.py new file mode 100644 index 0000000..7bb9ea1 --- /dev/null +++ b/netforge_rl/arena/questions/runners/population.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from collections.abc import Sequence + +from netforge_rl.arena.questions.policies import blue, killchain +from netforge_rl.baselines.policies import HeuristicRedPolicy, RandomPolicy + + +def run_red_population(*, seeds: Sequence[int], max_ticks: int) -> dict: + from netforge_rl.arena.adversarial import evaluate_population + + return evaluate_population( + blue, + { + 'random-red': lambda: RandomPolicy(seed=0), + 'heuristic-red': lambda: HeuristicRedPolicy(seed=0), + 'killchain-red': killchain, + }, + scenarios=('ransomware',), + seeds=tuple(seeds)[:2], + max_ticks=max_ticks, + ) diff --git a/netforge_rl/arena/questions/runners/probes.py b/netforge_rl/arena/questions/runners/probes.py new file mode 100644 index 0000000..3a29477 --- /dev/null +++ b/netforge_rl/arena/questions/runners/probes.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from collections.abc import Sequence + +from netforge_rl.arena.evaluate import run_episode +from netforge_rl.arena.questions.policies import blue, killchain + + +def run_probe(name: str, *, seeds: Sequence[int], max_ticks: int) -> dict: + from netforge_rl.diagnostics import run_diagnostic + from netforge_rl.diagnostics.adaptation import AdaptationShift + from netforge_rl.diagnostics.deception import DeceptionResistance + + diag = AdaptationShift() if name == 'adaptation' else DeceptionResistance() + diag.max_ticks = min(int(max_ticks), diag.max_ticks) + result = run_diagnostic(diag, blue(), seed=int(seeds[0])) + return { + 'diagnostic': result.diagnostic, + 'capability': result.capability, + 'score': result.score, + 'details': result.details, + } + + +def run_constrained(*, seeds: Sequence[int], max_ticks: int) -> dict: + from netforge_rl.arena.constraints import ConstrainedEnvWrapper + + seed = int(seeds[0]) + unconstrained = run_episode( + killchain(), + blue(), + scenario='ransomware', + seed=seed, + max_ticks=max_ticks, + ) + constrained = run_episode( + killchain(), + blue(), + scenario='ransomware', + seed=seed, + max_ticks=max_ticks, + wrap_env=lambda env: ConstrainedEnvWrapper(env, lambda_cost=1.0), + ) + return { + 'seed': seed, + 'unconstrained': { + 'sla_uptime': unconstrained.metrics.sla_uptime, + 'false_positives': unconstrained.metrics.false_positives, + 'security': unconstrained.metrics.security, + 'blue_return': round(unconstrained.metrics.blue_return, 4), + }, + 'constrained_lambda_1': { + 'sla_uptime': constrained.metrics.sla_uptime, + 'false_positives': constrained.metrics.false_positives, + 'security': constrained.metrics.security, + 'blue_return': round(constrained.metrics.blue_return, 4), + }, + } diff --git a/netforge_rl/arena/questions/runners/reward.py b/netforge_rl/arena/questions/runners/reward.py new file mode 100644 index 0000000..9850fed --- /dev/null +++ b/netforge_rl/arena/questions/runners/reward.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from collections.abc import Sequence + +from netforge_rl.arena.evaluate import run_episode +from netforge_rl.arena.questions.policies import blue, killchain + + +def run_reward_design(*, seeds: Sequence[int], max_ticks: int) -> dict: + from netforge_rl.rewards.variants import RewardDesignWrapper + + variants = ('default', 'sla_only', 'fp_penalized') + out = {} + seed = int(seeds[0]) + for variant in variants: + rec = run_episode( + killchain(), + blue(), + scenario='ransomware', + seed=seed, + max_ticks=max_ticks, + wrap_env=lambda env, v=variant: RewardDesignWrapper(env, variant=v), + ) + m = rec.metrics + out[variant] = { + 'isolated_hosts': m.isolated_hosts, + 'false_positives': m.false_positives, + 'sla_uptime': m.sla_uptime, + 'mission_success': m.mission_success, + 'blue_return': round(m.blue_return, 4), + 'raw_vs_shaped': 'see info[reward_variant] / info[raw_reward]', + } + return {'seed': seed, 'variants': out} diff --git a/netforge_rl/arena/questions/runners/risk.py b/netforge_rl/arena/questions/runners/risk.py new file mode 100644 index 0000000..ece87e1 --- /dev/null +++ b/netforge_rl/arena/questions/runners/risk.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + +from netforge_rl.arena.evaluate import run_episode +from netforge_rl.arena.metrics import cvar, worst_case +from netforge_rl.arena.questions.policies import blue, killchain + + +def run_risk_sensitive(*, seeds: Sequence[int], max_ticks: int) -> dict: + returns = [] + catastrophic = [] + sla = [] + for seed in seeds: + rec = run_episode( + killchain(), + blue(), + scenario='ransomware', + seed=int(seed), + max_ticks=max_ticks, + ) + returns.append(rec.metrics.blue_return) + catastrophic.append(rec.metrics.catastrophic_failure) + sla.append(rec.metrics.sla_uptime) + return { + 'n_seeds': len(returns), + 'mean_blue_return': round(float(np.mean(returns)), 4), + 'cvar05': round(cvar(returns, 0.05), 4), + 'worst_case': round(worst_case(returns), 4), + 'mean_sla': round(float(np.mean(sla)), 4), + 'catastrophic_rate': round(float(np.mean(catastrophic)), 4), + } diff --git a/netforge_rl/arena/questions/types.py b/netforge_rl/arena/questions/types.py new file mode 100644 index 0000000..502af8a --- /dev/null +++ b/netforge_rl/arena/questions/types.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass + + +@dataclass(frozen=True) +class ResearchQuestion: + id: str + title: str + question: str + why: str + protocol: str + metrics: tuple[str, ...] + status: str # runnable | protocol + note: str + + def as_dict(self) -> dict: + return asdict(self) diff --git a/netforge_rl/arena/spec.py b/netforge_rl/arena/spec.py new file mode 100644 index 0000000..84d8f9b --- /dev/null +++ b/netforge_rl/arena/spec.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal + +from netforge_rl.environment.presets import EVAL_SEEDS + +ARENA_VERSION = '2026' +ARENA_NAME = f'NetForge Arena {ARENA_VERSION}' + +SplitName = Literal['train', 'dev', 'hidden'] + +PUBLIC_SCENARIOS: tuple[str, ...] = ( + 'ransomware', + 'apt_espionage', + 'cloud_hybrid', + 'iot_grid', + 'ot_stuxnet', +) + +# Community packs are first-class Arena scenarios once loaded from YAML. +PACK_SCENARIOS: tuple[str, ...] = ( + 'hospital_ransomware', + 'cloud_iam', + 'enterprise_apt', + 'finance', + 'iot_fleet', + 'ot_plant', + 'zero_trust', +) + +TRAIN_SEEDS: tuple[int, ...] = tuple(range(20)) +DEV_SEEDS: tuple[int, ...] = EVAL_SEEDS +# Disjoint reserved block. Official leaderboard organizers may replace this +# with unpublished topology/attacker packs; the public proxy is still OOD. +HIDDEN_SEEDS: tuple[int, ...] = tuple(range(50_000, 50_020)) + +REFERENCE_BASELINES: tuple[str, ...] = ( + 'random', + 'heuristic-blue', + 'heuristic-red', + 'killchain-red', + 'ippo', + 'mappo', + 'rmappo', + 'gnn-mappo', + 'qmix', + 'llm', +) + +HEADLINE_CAPABILITIES: tuple[str, ...] = ( + 'memory', + 'temporal', + 'precision', + 'safety', + 'generalization', + 'deception', + 'adaptation', + 'attention', +) + +COMPETITION_METRICS: tuple[str, ...] = ( + 'mission_success', + 'disruption', + 'security', + 'false_positives', + 'exfiltration', + 'compromised_hosts', + 'catastrophic_failure', + 'sla_uptime', + 'cvar_return', + 'worst_case_return', +) + + +@dataclass(frozen=True) +class Split: + name: SplitName + seeds: tuple[int, ...] + evaluation_mode: bool + public: bool + description: str + + +@dataclass(frozen=True) +class ArenaSpec: + """Frozen evaluation contract for a yearly Arena release.""" + + version: str = ARENA_VERSION + scenarios: tuple[str, ...] = PUBLIC_SCENARIOS + difficulty: str = 'medium' + max_ticks: int = 200 + train: Split = field( + default_factory=lambda: Split( + name='train', + seeds=TRAIN_SEEDS, + evaluation_mode=False, + public=True, + description='Public training topologies.', + ) + ) + dev: Split = field( + default_factory=lambda: Split( + name='dev', + seeds=DEV_SEEDS, + evaluation_mode=True, + public=True, + description='Held-out topologies for model selection.', + ) + ) + hidden: Split = field( + default_factory=lambda: Split( + name='hidden', + seeds=HIDDEN_SEEDS, + evaluation_mode=True, + public=False, + description='Organizer split. Local runs use the public proxy.', + ) + ) + compute_budget_env_steps: int = 1_000_000 + baselines: tuple[str, ...] = REFERENCE_BASELINES + + def split(self, name: SplitName) -> Split: + return {'train': self.train, 'dev': self.dev, 'hidden': self.hidden}[name] + + def to_dict(self) -> dict: + return { + 'name': ARENA_NAME, + 'version': self.version, + 'scenarios': list(self.scenarios), + 'difficulty': self.difficulty, + 'max_ticks': self.max_ticks, + 'compute_budget_env_steps': self.compute_budget_env_steps, + 'baselines': list(self.baselines), + 'splits': { + s.name: { + 'n_seeds': len(s.seeds), + 'evaluation_mode': s.evaluation_mode, + 'public': s.public, + 'description': s.description, + } + for s in (self.train, self.dev, self.hidden) + }, + 'metrics': list(COMPETITION_METRICS), + 'capabilities': list(HEADLINE_CAPABILITIES), + } + + +ARENA_2026 = ArenaSpec() + + +def current_arena() -> ArenaSpec: + return ARENA_2026 diff --git a/netforge_rl/arena/submissions.py b/netforge_rl/arena/submissions.py new file mode 100644 index 0000000..87ee2e9 --- /dev/null +++ b/netforge_rl/arena/submissions.py @@ -0,0 +1,136 @@ +"""CAGE-style ``make_blue`` eval. Rank is a card, not mean reward.""" + +from __future__ import annotations + +import importlib.util +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path + +from netforge_rl.arena.evaluate import evaluate_split +from netforge_rl.arena.spec import ARENA_2026, SplitName +from netforge_rl.baselines.policies import KillChainRedPolicy + +# Local smoke. --official uses the Arena freeze. +LOCAL_EVAL = { + 'split': 'dev', + 'scenarios': ('ransomware',), + 'seeds': (9001, 9002), + 'max_ticks': 80, +} + +OFFICIAL_EVAL = { + 'split': 'dev', + 'scenarios': ARENA_2026.scenarios, + 'seeds': ARENA_2026.dev.seeds, + 'max_ticks': ARENA_2026.max_ticks, +} + + +@dataclass +class Submission: + name: str + team: str + technique: str + path: Path + make_blue: Callable + make_red: Callable | None = None + wrap: Callable | None = None + paper: str = '' + + +def load_submission(path: str | Path) -> Submission: + path = Path(path).resolve() + if not path.exists(): + raise FileNotFoundError(f'Submission not found: {path}') + spec = importlib.util.spec_from_file_location('netforge_submission', path) + if spec is None or spec.loader is None: + raise ImportError(f'Cannot import submission {path}') + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + make_blue = getattr(mod, 'make_blue', None) + if make_blue is None: + raise AttributeError( + f'{path} must define make_blue(seed=0) -> policy with act(env, agent_id)' + ) + return Submission( + name=str(getattr(mod, 'NAME', path.stem)), + team=str(getattr(mod, 'TEAM', '')), + technique=str(getattr(mod, 'TECHNIQUE', '')), + paper=str(getattr(mod, 'PAPER', '')), + path=path, + make_blue=make_blue, + make_red=getattr(mod, 'make_red', None), + wrap=getattr(mod, 'wrap', None), + ) + + +def evaluate_submission( + submission: Submission, + *, + split: SplitName = 'dev', + scenarios: Sequence[str] | None = None, + seeds: Sequence[int] | None = None, + max_ticks: int | None = None, + official: bool = False, + agent_card: bool = False, +) -> dict: + proto = dict(OFFICIAL_EVAL if official else LOCAL_EVAL) + if official: + split = proto['split'] # type: ignore[assignment] + scenarios = tuple(scenarios or proto['scenarios']) + seeds = tuple(seeds or proto['seeds']) + max_ticks = int(max_ticks if max_ticks is not None else proto['max_ticks']) + + red_factory = submission.make_red or (lambda: KillChainRedPolicy(seed=0)) + wrap_env = submission.wrap + result = evaluate_split( + red_factory, + lambda: submission.make_blue(seed=0), + split=split, + scenarios=scenarios, + seeds=seeds, + max_ticks=max_ticks, + wrap_env=wrap_env, + ) + mean = result['overall'].get('mean', {}) + warnings = ['Rank: mission/SLA/FPs/security/CVaR — not mean reward.'] + if wrap_env is not None: + warnings.append('wrap() runs after belief graphs; do not attach oracle_graph.') + payload = { + 'name': submission.name, + 'team': submission.team, + 'technique': submission.technique, + 'paper': submission.paper, + 'path': str(submission.path), + 'arena': result['arena'], + 'split': result['split'], + 'public': result['public'], + 'official': official, + 'n_episodes': result['n_episodes'], + 'scenarios': result['scenarios'], + 'overall': result['overall'], + 'per_scenario': result['per_scenario'], + 'headline': { + 'mission_success': mean.get('mission_success'), + 'sla_uptime': mean.get('sla_uptime'), + 'security': mean.get('security'), + 'false_positives': mean.get('false_positives'), + 'catastrophic_failure': mean.get('catastrophic_failure'), + 'cvar05_blue_return': result['overall'] + .get('cvar05', {}) + .get('blue_return'), + 'mean_blue_return': mean.get('blue_return'), + }, + 'warnings': warnings, + } + if agent_card: + from netforge_rl.arena.agent_card import build_agent_card + + payload['agent_card'] = build_agent_card( + lambda: submission.make_blue(seed=0), + name=submission.name, + seeds=tuple(seeds)[:2], + run_arena=False, + ) + return payload diff --git a/tests/arena/test_arena.py b/tests/arena/test_arena.py new file mode 100644 index 0000000..e12d004 --- /dev/null +++ b/tests/arena/test_arena.py @@ -0,0 +1,94 @@ +import json + +import pytest + +from netforge_rl.arena.evaluate import run_episode +from netforge_rl.arena.metrics import cvar, metrics_from_env +from netforge_rl.arena.spec import ARENA_2026 +from netforge_rl.baselines.policies import HeuristicBluePolicy, KillChainRedPolicy +from netforge_rl.cli import main +from netforge_rl.environment.parallel_env import NetForgeRLEnv +from netforge_rl.scenarios.yaml_dsl import list_packs, resolve_scenario + + +@pytest.mark.fast +def test_arena_spec_has_hidden_split(): + spec = ARENA_2026.to_dict() + assert spec['splits']['hidden']['public'] is False + assert spec['splits']['dev']['public'] is True + assert 'mission_success' in spec['metrics'] + + +@pytest.mark.fast +def test_cvar_is_lower_tail(): + assert cvar([0.0, 1.0, 2.0, 3.0, 100.0], alpha=0.2) <= 1.0 + + +@pytest.mark.integration +def test_run_episode_records_competition_metrics(): + rec = run_episode( + KillChainRedPolicy(seed=0), + HeuristicBluePolicy(seed=0), + scenario='ransomware', + seed=0, + max_ticks=15, + graph_obs=False, + ) + m = rec.metrics + assert 0.0 <= m.sla_uptime <= 1.0 + assert m.compromised_hosts >= 0 + assert m.false_positives >= 0 + assert m.catastrophic_failure in (0.0, 1.0) + + +@pytest.mark.fast +def test_community_packs_resolve(): + packs = list_packs() + assert 'hospital_ransomware' in packs + spec = resolve_scenario('hospital_ransomware') + assert spec.base == 'ransomware' + assert spec.topology_spec is not None + + +@pytest.mark.integration +def test_pack_builds_env(): + env = NetForgeRLEnv({'scenario_type': 'hospital_ransomware', 'max_ticks': 5}) + env.reset(seed=0) + names = {h.hostname for h in env._active_hosts()} + assert 'patient_portal' in names + + +@pytest.mark.fast +def test_cli_version(capsys): + assert main(['version']) == 0 + out = capsys.readouterr().out + assert 'Arena' in out + + +@pytest.mark.fast +def test_cli_arena_json(capsys): + assert main(['arena', '--json']) == 0 + payload = json.loads(capsys.readouterr().out) + assert payload['version'] == '2026' + + +@pytest.mark.fast +def test_info_exposes_catastrophic_and_fp(): + env = NetForgeRLEnv({'scenario_type': 'ransomware', 'max_ticks': 3}) + env.reset(seed=0) + _, _, _, _, infos = env.step({a: env.action_space(a).sample() for a in env.agents}) + blue = next(v for k, v in infos.items() if 'blue' in k) + assert 'false_positives_total' in blue + assert 'catastrophic_failure' in blue + assert 'mission_success' in blue + metrics_from_env(env, infos, 0.0, 0.0) + + +@pytest.mark.fast +def test_fixed_step_does_not_skip_ticks(): + env = NetForgeRLEnv( + {'scenario_type': 'ransomware', 'max_ticks': 8, 'time_mode': 'fixed'} + ) + env.reset(seed=0) + env.step({a: env.action_space(a).sample() for a in env.agents}) + assert env.current_tick == 1 diff --git a/tests/arena/test_infra.py b/tests/arena/test_infra.py new file mode 100644 index 0000000..f8f1c38 --- /dev/null +++ b/tests/arena/test_infra.py @@ -0,0 +1,96 @@ +import pytest + +from netforge_rl.baselines.policies import HeuristicBluePolicy, RandomPolicy +from netforge_rl.comm.wrapper import CommFailureWrapper +from netforge_rl.datasets.collector import collect_episode +from netforge_rl.datasets.minari_export import export_npz, trajectories_to_minari_dict +from netforge_rl.environment.parallel_env import NetForgeRLEnv +from netforge_rl.league.psro import payoff_matrix +from netforge_rl.render.replay_html import ReplayLog, write_replay_html +from netforge_rl.rewards.variants import RewardDesignWrapper + + +@pytest.mark.fast +def test_comm_drop_zeros_channel(): + env = CommFailureWrapper( + NetForgeRLEnv({'scenario_type': 'ransomware', 'max_ticks': 4}), + drop_p=1.0, + seed=0, + ) + obs, _ = env.reset(seed=0) + obs, *_ = env.step({a: env.action_space(a).sample() for a in env.agents}) + assert float(obs['blue_dmz']['blue_comm'].sum()) == 0.0 + + +@pytest.mark.fast +def test_reward_variant_records_raw(): + env = RewardDesignWrapper( + NetForgeRLEnv({'scenario_type': 'ransomware', 'max_ticks': 4}), + variant='sla_only', + ) + env.reset(seed=0) + _, rewards, _, _, infos = env.step( + {a: env.action_space(a).sample() for a in env.agents} + ) + assert 'raw_reward' in infos['blue_dmz'] + assert infos['blue_dmz']['reward_variant'] == 'sla_only' + assert 'blue_dmz' in rewards + + +@pytest.mark.integration +def test_offline_export_npz(tmp_path): + traj = collect_episode( + RandomPolicy(seed=0), + HeuristicBluePolicy(seed=0), + max_ticks=8, + seed=0, + quality='heuristic', + ) + payload = trajectories_to_minari_dict([traj]) + assert payload['n_episodes'] == 1 + path = export_npz([traj], tmp_path / 'ds.npz') + assert path.exists() + + +@pytest.mark.fast +def test_try_minari_export_writes_npz_or_hdf5(tmp_path): + from netforge_rl.datasets.minari_export import try_minari_export + + traj = collect_episode( + RandomPolicy(seed=0), + HeuristicBluePolicy(seed=0), + max_ticks=4, + seed=0, + quality='heuristic', + ) + path = try_minari_export([traj], 'netforge/test-v0', tmp_path / 'ds.hdf5') + assert path.exists() + assert path.suffix in {'.hdf5', '.npz'} + + +@pytest.mark.fast +def test_psro_matrix_shape(): + game = payoff_matrix( + {'r': lambda: RandomPolicy(seed=0)}, + {'b': lambda: HeuristicBluePolicy(seed=0)}, + seeds=(0,), + max_ticks=5, + ) + assert len(game['payoff_blue']) == 1 + assert len(game['payoff_blue'][0]) == 1 + + +@pytest.mark.fast +def test_replay_html(tmp_path): + env = NetForgeRLEnv({'scenario_type': 'ransomware', 'max_ticks': 4}) + env.reset(seed=0) + log = ReplayLog('ransomware', 0) + _, rewards, _, _, infos = env.step( + {a: env.action_space(a).sample() for a in env.agents} + ) + log.capture(env, rewards, infos, {}) + html = write_replay_html(log.to_dict(), tmp_path / 'replay.html') + text = html.read_text(encoding='utf-8') + assert 'NETFORGE' in text + log.save_json(tmp_path / 'traj.json') + assert (tmp_path / 'traj.json').exists() diff --git a/tests/arena/test_questions.py b/tests/arena/test_questions.py new file mode 100644 index 0000000..061142b --- /dev/null +++ b/tests/arena/test_questions.py @@ -0,0 +1,77 @@ +import json + +import pytest + +from netforge_rl.arena.questions import ( + QUESTIONS, + get_question, + list_questions, + run_question, +) +from netforge_rl.arena.submissions import evaluate_submission, load_submission +from netforge_rl.cli import main + + +@pytest.mark.fast +def test_catalog_has_runnable_and_protocol_questions(): + ids = {q.id for q in QUESTIONS} + assert 'belief-vs-oracle' in ids + assert 'gnn-belief' in ids + assert get_question('reward-design').status == 'runnable' + assert get_question('league-psro').status == 'protocol' + rows = list_questions(runnable_only=True) + assert all(r['status'] == 'runnable' for r in rows) + assert len(rows) >= 8 + + +@pytest.mark.fast +def test_belief_vs_oracle_reports_information_gap(): + payload = run_question('belief-vs-oracle', seeds=(0,), max_ticks=4) + answer = payload['answer'] + assert 0.0 <= answer['visible_ratio'] <= 1.0 + assert answer['compromise_recall'] == 0.0 + assert answer['oracle_compromised_nodes'] >= 1 + + +@pytest.mark.fast +def test_protocol_question_does_not_pretend_to_run(): + payload = run_question('gnn-belief') + assert payload['status'] == 'protocol' + assert 'answer' not in payload + + +@pytest.mark.fast +def test_cli_questions_list(capsys): + assert main(['questions']) == 0 + out = capsys.readouterr().out + assert 'belief-vs-oracle' in out + + +@pytest.mark.fast +def test_cli_questions_json(capsys): + assert main(['questions', '--json']) == 0 + payload = json.loads(capsys.readouterr().out) + assert any(q['id'] == 'time-mode' for q in payload) + + +@pytest.mark.integration +def test_reward_design_question_runs(): + payload = run_question('reward-design', seeds=(0,), max_ticks=6) + variants = payload['answer']['variants'] + assert set(variants) == {'default', 'sla_only', 'fp_penalized'} + + +@pytest.mark.integration +def test_example_submission_evaluates_without_reward_ranking(): + sub = load_submission('examples/submissions/heuristic_blue.py') + result = evaluate_submission( + sub, + scenarios=('ransomware',), + seeds=(0,), + max_ticks=6, + official=False, + ) + assert result['name'] == 'heuristic-blue' + assert 'mission_success' in result['headline'] + assert 'mean_blue_return' in result['headline'] + assert any('not mean reward' in w for w in result['warnings']) From 9979598b1a87a4cfd1b1bba98c2d75ddce903d71 Mon Sep 17 00:00:00 2001 From: elprofesoriqo Date: Sun, 20 Sep 2026 15:15:25 +0200 Subject: [PATCH 4/8] Add YAML scenario packs for hospital, cloud, APT, IoT, OT, and zero-trust. Community packs stay data, not bundled training algorithms. Co-authored-by: Cursor --- netforge_rl/scenarios/__init__.py | 18 +- netforge_rl/scenarios/packs/cloud_iam.yaml | 18 ++ .../scenarios/packs/enterprise_apt.yaml | 18 ++ netforge_rl/scenarios/packs/finance.yaml | 17 ++ .../scenarios/packs/hospital_ransomware.yaml | 22 +++ netforge_rl/scenarios/packs/iot_fleet.yaml | 17 ++ netforge_rl/scenarios/packs/ot_plant.yaml | 17 ++ netforge_rl/scenarios/packs/zero_trust.yaml | 18 ++ netforge_rl/scenarios/yaml_dsl.py | 160 ++++++++++++++++++ 9 files changed, 302 insertions(+), 3 deletions(-) create mode 100644 netforge_rl/scenarios/packs/cloud_iam.yaml create mode 100644 netforge_rl/scenarios/packs/enterprise_apt.yaml create mode 100644 netforge_rl/scenarios/packs/finance.yaml create mode 100644 netforge_rl/scenarios/packs/hospital_ransomware.yaml create mode 100644 netforge_rl/scenarios/packs/iot_fleet.yaml create mode 100644 netforge_rl/scenarios/packs/ot_plant.yaml create mode 100644 netforge_rl/scenarios/packs/zero_trust.yaml create mode 100644 netforge_rl/scenarios/yaml_dsl.py diff --git a/netforge_rl/scenarios/__init__.py b/netforge_rl/scenarios/__init__.py index c254bb9..35c46f2 100644 --- a/netforge_rl/scenarios/__init__.py +++ b/netforge_rl/scenarios/__init__.py @@ -13,11 +13,23 @@ } +def is_builtin(name: str) -> bool: + return str(name).lower() in _SCENARIOS + + def get_scenario_class(name: str): key = name.lower() - if key not in _SCENARIOS: - raise KeyError(f'Unknown scenario {name!r}. Available: {sorted(_SCENARIOS)}') - return _SCENARIOS[key] + if key in _SCENARIOS: + return _SCENARIOS[key] + from netforge_rl.scenarios.yaml_dsl import PACKS_DIR, load_scenario_file + + pack = PACKS_DIR / f'{key}.yaml' + if pack.exists(): + spec = load_scenario_file(pack) + base = spec.base.lower() + if base in _SCENARIOS: + return _SCENARIOS[base] + raise KeyError(f'Unknown scenario {name!r}. Available: {sorted(_SCENARIOS)}') def get_reward_weights(name: str) -> dict: diff --git a/netforge_rl/scenarios/packs/cloud_iam.yaml b/netforge_rl/scenarios/packs/cloud_iam.yaml new file mode 100644 index 0000000..4a10f25 --- /dev/null +++ b/netforge_rl/scenarios/packs/cloud_iam.yaml @@ -0,0 +1,18 @@ +name: cloud_iam +pack: cloud +base: cloud_hybrid +difficulty: medium +max_ticks: 200 +description: Cloud IAM attack path — protect the Secure enclave and identity plane. +reward_variant: security_only +topology: + subnets: + - {cidr: "192.168.1.0/24", name: DMZ} + - {cidr: "10.0.0.0/24", name: Corporate} + - {cidr: "10.0.1.0/24", name: Secure} + hosts: + - {ip: "192.168.1.10", hostname: api_gateway, subnet_cidr: "192.168.1.0/24", os: Linux_Ubuntu, services: [SSH, Apache], vulnerabilities: [CVE-2021-44228], cvss_score: 9.8} + - {ip: "10.0.0.10", hostname: idp, subnet_cidr: "10.0.0.0/24", os: Linux_Ubuntu, services: [SSH], vulnerabilities: [CVE-2021-44228], cvss_score: 8.8} + - {ip: "10.0.0.11", hostname: ci_runner, subnet_cidr: "10.0.0.0/24", os: Linux_Ubuntu, services: [SSH], vulnerabilities: [], cvss_score: 5.0} + - {ip: "10.0.1.10", hostname: secrets_vault, subnet_cidr: "10.0.1.0/24", os: Linux_Ubuntu, services: [SSH], vulnerabilities: [], cvss_score: 3.0} + - {ip: "10.0.1.11", hostname: admin_console, subnet_cidr: "10.0.1.0/24", os: Windows_Server_2016, services: [RDP], vulnerabilities: [CVE-2019-0708], cvss_score: 9.8, is_domain_controller: true} diff --git a/netforge_rl/scenarios/packs/enterprise_apt.yaml b/netforge_rl/scenarios/packs/enterprise_apt.yaml new file mode 100644 index 0000000..6b99a11 --- /dev/null +++ b/netforge_rl/scenarios/packs/enterprise_apt.yaml @@ -0,0 +1,18 @@ +name: enterprise_apt +pack: enterprise +base: apt_espionage +difficulty: medium +max_ticks: 200 +description: Enterprise APT — stealthy persistence and exfiltration vs SOC containment. +reward_variant: default +topology: + subnets: + - {cidr: "192.168.1.0/24", name: DMZ} + - {cidr: "10.0.0.0/24", name: Corporate} + - {cidr: "10.0.1.0/24", name: Secure} + hosts: + - {ip: "192.168.1.10", hostname: mail_gw, subnet_cidr: "192.168.1.0/24", os: Linux_Ubuntu, services: [SSH, Apache], vulnerabilities: [CVE-2021-44228], cvss_score: 9.8} + - {ip: "10.0.0.10", hostname: jump_box, subnet_cidr: "10.0.0.0/24", os: Windows_10, services: [RDP, SMB], vulnerabilities: [MS17-010], cvss_score: 8.1} + - {ip: "10.0.0.11", hostname: file_server, subnet_cidr: "10.0.0.0/24", os: Windows_Server_2016, services: [SMB], vulnerabilities: [MS17-010], cvss_score: 8.1} + - {ip: "10.0.1.10", hostname: dc01, subnet_cidr: "10.0.1.0/24", os: Windows_Server_2016, services: [SMB, IIS], vulnerabilities: [MS17-010], cvss_score: 8.8, is_domain_controller: true} + - {ip: "10.0.1.11", hostname: finance_db, subnet_cidr: "10.0.1.0/24", os: Linux_Ubuntu, services: [SSH], vulnerabilities: [], cvss_score: 4.3} diff --git a/netforge_rl/scenarios/packs/finance.yaml b/netforge_rl/scenarios/packs/finance.yaml new file mode 100644 index 0000000..71e0478 --- /dev/null +++ b/netforge_rl/scenarios/packs/finance.yaml @@ -0,0 +1,17 @@ +name: finance +pack: finance +base: apt_espionage +difficulty: medium +max_ticks: 200 +description: Finance trading floor — exfiltration of settlement hosts is the Red win. +reward_variant: fp_penalized +topology: + subnets: + - {cidr: "192.168.1.0/24", name: DMZ} + - {cidr: "10.0.0.0/24", name: Corporate} + - {cidr: "10.0.1.0/24", name: Secure} + hosts: + - {ip: "192.168.1.10", hostname: public_web, subnet_cidr: "192.168.1.0/24", os: Linux_Ubuntu, services: [SSH, Apache], vulnerabilities: [CVE-2021-44228], cvss_score: 9.8} + - {ip: "10.0.0.10", hostname: trader_ws, subnet_cidr: "10.0.0.0/24", os: Windows_10, services: [RDP], vulnerabilities: [CVE-2019-0708], cvss_score: 9.8} + - {ip: "10.0.1.10", hostname: settlement, subnet_cidr: "10.0.1.0/24", os: Linux_Ubuntu, services: [SSH], vulnerabilities: [], cvss_score: 3.5} + - {ip: "10.0.1.11", hostname: swift_gw, subnet_cidr: "10.0.1.0/24", os: Linux_Ubuntu, services: [SSH], vulnerabilities: [], cvss_score: 5.0} diff --git a/netforge_rl/scenarios/packs/hospital_ransomware.yaml b/netforge_rl/scenarios/packs/hospital_ransomware.yaml new file mode 100644 index 0000000..00eb82f --- /dev/null +++ b/netforge_rl/scenarios/packs/hospital_ransomware.yaml @@ -0,0 +1,22 @@ +name: hospital_ransomware +pack: healthcare +base: ransomware +difficulty: medium +max_ticks: 200 +description: > + Hospital ransomware: EHR and PACS must stay up while Blue contains a + DMZ foothold. Isolation of clinical hosts is a disruption failure. +reward_variant: constrained_sla +topology: + subnets: + - {cidr: "192.168.1.0/24", name: DMZ} + - {cidr: "10.0.0.0/24", name: Corporate} + - {cidr: "10.0.1.0/24", name: Secure} + - {cidr: "10.0.99.0/24", name: OT_Subnet} + hosts: + - {ip: "192.168.1.10", hostname: patient_portal, subnet_cidr: "192.168.1.0/24", os: Linux_Ubuntu, services: [SSH, Apache], vulnerabilities: [CVE-2021-44228], cvss_score: 9.8} + - {ip: "10.0.0.10", hostname: ehr_workstation, subnet_cidr: "10.0.0.0/24", os: Windows_10, services: [RDP, SMB], vulnerabilities: [MS17-010], cvss_score: 8.1} + - {ip: "10.0.0.11", hostname: nurse_station, subnet_cidr: "10.0.0.0/24", os: Windows_10, services: [RDP], vulnerabilities: [CVE-2019-0708], cvss_score: 9.8} + - {ip: "10.0.1.10", hostname: pacs, subnet_cidr: "10.0.1.0/24", os: Windows_Server_2016, services: [SMB, IIS], vulnerabilities: [MS17-010], cvss_score: 8.1, is_domain_controller: true} + - {ip: "10.0.1.11", hostname: billing, subnet_cidr: "10.0.1.0/24", os: Linux_Ubuntu, services: [SSH], vulnerabilities: [], cvss_score: 4.0} + - {ip: "10.0.99.10", hostname: infusion_pump, subnet_cidr: "10.0.99.0/24", os: PLC_Firmware, services: [Modbus], ot: true, cvss_score: 7.5} diff --git a/netforge_rl/scenarios/packs/iot_fleet.yaml b/netforge_rl/scenarios/packs/iot_fleet.yaml new file mode 100644 index 0000000..6a02f1c --- /dev/null +++ b/netforge_rl/scenarios/packs/iot_fleet.yaml @@ -0,0 +1,17 @@ +name: iot_fleet +pack: iot +base: iot_grid +difficulty: medium +max_ticks: 200 +description: IoT / grid controllers that must stay healthy under a spreading botnet. +reward_variant: sla_only +topology: + subnets: + - {cidr: "192.168.1.0/24", name: DMZ} + - {cidr: "10.0.0.0/24", name: Corporate} + - {cidr: "10.0.99.0/24", name: OT_Subnet} + hosts: + - {ip: "192.168.1.10", hostname: mqtt_broker, subnet_cidr: "192.168.1.0/24", os: Linux_Ubuntu, services: [SSH, Apache], vulnerabilities: [CVE-2021-44228], cvss_score: 9.8} + - {ip: "10.0.0.10", hostname: fleet_mgr, subnet_cidr: "10.0.0.0/24", os: Linux_Ubuntu, services: [SSH], vulnerabilities: [], cvss_score: 5.0} + - {ip: "10.0.99.10", hostname: feeder_plc, subnet_cidr: "10.0.99.0/24", os: PLC_Firmware, services: [Modbus], ot: true, cvss_score: 8.0} + - {ip: "10.0.99.11", hostname: meter_gw, subnet_cidr: "10.0.99.0/24", os: PLC_Firmware, services: [Modbus], ot: true, cvss_score: 7.0} diff --git a/netforge_rl/scenarios/packs/ot_plant.yaml b/netforge_rl/scenarios/packs/ot_plant.yaml new file mode 100644 index 0000000..022771a --- /dev/null +++ b/netforge_rl/scenarios/packs/ot_plant.yaml @@ -0,0 +1,17 @@ +name: ot_plant +pack: ot +base: ot_stuxnet +difficulty: hard +max_ticks: 200 +description: Process-control plant. Kinetic PLC destruction is a catastrophic failure. +reward_variant: safety_critical +topology: + subnets: + - {cidr: "192.168.1.0/24", name: DMZ} + - {cidr: "10.0.0.0/24", name: Corporate} + - {cidr: "10.0.99.0/24", name: OT_Subnet} + hosts: + - {ip: "192.168.1.10", hostname: historian_dmz, subnet_cidr: "192.168.1.0/24", os: Windows_Server_2016, services: [IIS, SMB], vulnerabilities: [MS17-010], cvss_score: 8.1} + - {ip: "10.0.0.10", hostname: eng_ws, subnet_cidr: "10.0.0.0/24", os: Windows_10, services: [RDP, SMB], vulnerabilities: [MS17-010], cvss_score: 8.1} + - {ip: "10.0.99.10", hostname: turbine_plc, subnet_cidr: "10.0.99.0/24", os: PLC_Firmware, services: [S7Comm], ot: true, cvss_score: 9.0} + - {ip: "10.0.99.11", hostname: safety_plc, subnet_cidr: "10.0.99.0/24", os: PLC_Firmware, services: [Modbus], ot: true, cvss_score: 9.5} diff --git a/netforge_rl/scenarios/packs/zero_trust.yaml b/netforge_rl/scenarios/packs/zero_trust.yaml new file mode 100644 index 0000000..3bc81cb --- /dev/null +++ b/netforge_rl/scenarios/packs/zero_trust.yaml @@ -0,0 +1,18 @@ +name: zero_trust +pack: identity +base: cloud_hybrid +difficulty: medium +max_ticks: 200 +description: Zero-trust campus — identity tokens and micro-segmented Secure hosts. +reward_variant: constrained_sla +log_latency: 3 +topology: + subnets: + - {cidr: "192.168.1.0/24", name: DMZ} + - {cidr: "10.0.0.0/24", name: Corporate} + - {cidr: "10.0.1.0/24", name: Secure} + hosts: + - {ip: "192.168.1.10", hostname: ztna_broker, subnet_cidr: "192.168.1.0/24", os: Linux_Ubuntu, services: [SSH, Apache], vulnerabilities: [CVE-2021-44228], cvss_score: 8.8} + - {ip: "10.0.0.10", hostname: laptop_pool, subnet_cidr: "10.0.0.0/24", os: Windows_10, services: [RDP], vulnerabilities: [CVE-2019-0708], cvss_score: 9.8} + - {ip: "10.0.1.10", hostname: privileged_pam, subnet_cidr: "10.0.1.0/24", os: Linux_Ubuntu, services: [SSH], vulnerabilities: [], cvss_score: 3.0} + - {ip: "10.0.1.11", hostname: dc_zt, subnet_cidr: "10.0.1.0/24", os: Windows_Server_2016, services: [SMB], vulnerabilities: [MS17-010], cvss_score: 8.1, is_domain_controller: true} diff --git a/netforge_rl/scenarios/yaml_dsl.py b/netforge_rl/scenarios/yaml_dsl.py new file mode 100644 index 0000000..ccaf094 --- /dev/null +++ b/netforge_rl/scenarios/yaml_dsl.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from pathlib import Path + +import yaml + +from netforge_rl.environment.presets import make_config + +PACKS_DIR = Path(__file__).parent / 'packs' + + +@dataclass +class ScenarioSpec: + """Declarative experiment: topology, scenario family, knobs. No simulator edits.""" + + name: str + base: str = 'ransomware' + pack: str = 'custom' + difficulty: str = 'medium' + max_ticks: int = 200 + description: str = '' + topology_path: str | None = None + topology_spec: dict | None = None + reward_variant: str | None = None + log_latency: int | None = None + time_mode: str = 'event' + extras: dict = field(default_factory=dict) + + def env_config(self, **overrides) -> dict: + cfg = make_config( + self.difficulty, + scenario_type=self.base, + max_ticks=overrides.pop('max_ticks', self.max_ticks), + **{ + k: v + for k, v in { + 'topology_path': self.topology_path, + 'topology_spec': self.topology_spec, + 'log_latency': self.log_latency, + 'time_mode': self.time_mode, + }.items() + if v is not None + }, + ) + cfg.update(self.extras) + cfg.update(overrides) + return cfg + + +def packs_dir() -> Path: + return PACKS_DIR + + +def list_packs() -> list[str]: + names = [] + if PACKS_DIR.exists(): + names.extend(p.stem for p in sorted(PACKS_DIR.glob('*.yaml'))) + return names + + +def _from_mapping(data: dict, source: str = 'custom') -> ScenarioSpec: + topo = data.get('topology') + topology_spec = None + topology_path = data.get('topology_path') + if isinstance(topo, dict) and (topo.get('subnets') or topo.get('hosts')): + topology_spec = topo + elif isinstance(topo, str): + topology_path = topo + extras = dict(data.get('env', {}) or {}) + return ScenarioSpec( + name=data.get('name', source), + base=data.get('base', data.get('scenario_type', 'ransomware')), + pack=data.get('pack', source), + difficulty=data.get('difficulty', 'medium'), + max_ticks=int(data.get('max_ticks', 200)), + description=data.get('description', ''), + topology_path=topology_path, + topology_spec=topology_spec, + reward_variant=data.get('reward_variant'), + log_latency=data.get('log_latency'), + time_mode=data.get('time_mode', extras.pop('time_mode', 'event')), + extras=extras, + ) + + +def load_scenario_file(path: str | Path) -> ScenarioSpec: + path = Path(path) + data = yaml.safe_load(path.read_text()) or {} + spec = _from_mapping(data, source=path.stem) + if spec.topology_path and not Path(spec.topology_path).is_absolute(): + candidate = path.parent / spec.topology_path + if candidate.exists(): + spec.topology_path = str(candidate) + return spec + + +def resolve_scenario(name_or_path: str) -> ScenarioSpec: + """Accept a builtin name, a community pack stem, or a YAML path.""" + from netforge_rl.scenarios import is_builtin + + raw = str(name_or_path) + as_path = Path(raw) + if as_path.suffix in {'.yaml', '.yml'} and as_path.exists(): + return load_scenario_file(as_path) + pack = PACKS_DIR / f'{raw}.yaml' + if pack.exists(): + return load_scenario_file(pack) + if is_builtin(raw): + return ScenarioSpec(name=raw, base=raw, pack='builtin') + raise KeyError( + f'Unknown scenario {raw!r}. Builtin families or packs: {sorted(list_packs())}' + ) + + +def expand_scenario_config(cfg: dict) -> dict: + """Resolve pack aliases into a builtin scenario_type plus topology knobs.""" + cfg = dict(cfg or {}) + name = cfg.get('scenario_type', 'ransomware') + from netforge_rl.scenarios import is_builtin + + if is_builtin(name): + return cfg + try: + spec = resolve_scenario(name) + except KeyError: + return cfg + merged = spec.env_config() + user = {k: v for k, v in cfg.items() if k != 'scenario_type'} + merged.update(user) + merged['scenario_type'] = spec.base + return merged + + +def make_env_from_spec( + spec: ScenarioSpec, + difficulty: str | None = None, + seed: int | None = None, + **overrides, +): + from netforge_rl.environment.parallel_env import NetForgeRLEnv + + if difficulty: + spec = replace(spec, difficulty=difficulty) + env = NetForgeRLEnv(spec.env_config(**overrides)) + if spec.reward_variant: + from netforge_rl.rewards.variants import RewardDesignWrapper + + env = RewardDesignWrapper(env, variant=spec.reward_variant) + if seed is not None: + env.reset(seed=seed) + return env + + +def make_env_from_yaml( + path: str | Path, + seed: int | None = None, + **overrides, +): + return make_env_from_spec(load_scenario_file(path), seed=seed, **overrides) From b78115912664e697834945de82dccbd2c646c731 Mon Sep 17 00:00:00 2001 From: elprofesoriqo Date: Sun, 20 Sep 2026 15:15:25 +0200 Subject: [PATCH 5/8] Add the netforge CLI and package entry for python -m netforge. Pin the project to netforge-rl 4.0.0 with Gymnasium/PettingZoo as the gym, not a trainer repo. Co-authored-by: Cursor --- netforge/__init__.py | 1 + netforge/__main__.py | 4 ++ netforge_rl/__main__.py | 4 ++ netforge_rl/baselines/__init__.py | 5 ++ netforge_rl/baselines/eval.py | 2 +- netforge_rl/baselines/policies.py | 8 +-- netforge_rl/baselines/protocol.py | 107 ++++++++++++++++++++++++++++ netforge_rl/baselines/registry.py | 55 ++++++++++++++ netforge_rl/cli/__init__.py | 17 +++++ netforge_rl/cli/cmds/__init__.py | 21 ++++++ netforge_rl/cli/cmds/arena.py | 30 ++++++++ netforge_rl/cli/cmds/benchmark.py | 36 ++++++++++ netforge_rl/cli/cmds/collect.py | 20 ++++++ netforge_rl/cli/cmds/evaluate.py | 23 ++++++ netforge_rl/cli/cmds/questions.py | 28 ++++++++ netforge_rl/cli/cmds/replay.py | 13 ++++ netforge_rl/cli/cmds/run.py | 47 ++++++++++++ netforge_rl/cli/cmds/test_policy.py | 13 ++++ netforge_rl/cli/cmds/version.py | 13 ++++ netforge_rl/cli/parser.py | 81 +++++++++++++++++++++ netforge_rl/cli/util.py | 4 ++ pyproject.toml | 46 +++++++++--- tests/test_import_order.py | 26 +++++++ 23 files changed, 589 insertions(+), 15 deletions(-) create mode 100644 netforge/__init__.py create mode 100644 netforge/__main__.py create mode 100644 netforge_rl/__main__.py create mode 100644 netforge_rl/baselines/protocol.py create mode 100644 netforge_rl/baselines/registry.py create mode 100644 netforge_rl/cli/__init__.py create mode 100644 netforge_rl/cli/cmds/__init__.py create mode 100644 netforge_rl/cli/cmds/arena.py create mode 100644 netforge_rl/cli/cmds/benchmark.py create mode 100644 netforge_rl/cli/cmds/collect.py create mode 100644 netforge_rl/cli/cmds/evaluate.py create mode 100644 netforge_rl/cli/cmds/questions.py create mode 100644 netforge_rl/cli/cmds/replay.py create mode 100644 netforge_rl/cli/cmds/run.py create mode 100644 netforge_rl/cli/cmds/test_policy.py create mode 100644 netforge_rl/cli/cmds/version.py create mode 100644 netforge_rl/cli/parser.py create mode 100644 netforge_rl/cli/util.py diff --git a/netforge/__init__.py b/netforge/__init__.py new file mode 100644 index 0000000..c37688d --- /dev/null +++ b/netforge/__init__.py @@ -0,0 +1 @@ +"""CLI entry so ``python -m netforge`` matches the ``netforge`` script.""" diff --git a/netforge/__main__.py b/netforge/__main__.py new file mode 100644 index 0000000..e97cf8d --- /dev/null +++ b/netforge/__main__.py @@ -0,0 +1,4 @@ +from netforge_rl.cli import main + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/netforge_rl/__main__.py b/netforge_rl/__main__.py new file mode 100644 index 0000000..e97cf8d --- /dev/null +++ b/netforge_rl/__main__.py @@ -0,0 +1,4 @@ +from netforge_rl.cli import main + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/netforge_rl/baselines/__init__.py b/netforge_rl/baselines/__init__.py index ca1dbb1..d2534a9 100644 --- a/netforge_rl/baselines/__init__.py +++ b/netforge_rl/baselines/__init__.py @@ -6,6 +6,8 @@ KillChainRedPolicy, RandomPolicy, ) +from netforge_rl.baselines.protocol import format_baseline_table, run_baseline_table +from netforge_rl.baselines.registry import make_policy __all__ = [ 'BasePolicy', @@ -14,4 +16,7 @@ 'KillChainRedPolicy', 'RandomPolicy', 'evaluate', + 'format_baseline_table', + 'make_policy', + 'run_baseline_table', ] diff --git a/netforge_rl/baselines/eval.py b/netforge_rl/baselines/eval.py index da89f8e..8d051f1 100644 --- a/netforge_rl/baselines/eval.py +++ b/netforge_rl/baselines/eval.py @@ -7,7 +7,7 @@ HeuristicRedPolicy, RandomPolicy, ) -from netforge_rl.environment.parallel_env import NetForgeRLEnv, PADDING_SUBNET +from netforge_rl.environment.parallel_env import PADDING_SUBNET, NetForgeRLEnv from netforge_rl.semantic.leaderboard import append_result from netforge_rl.semantic.runner import EpisodeResult diff --git a/netforge_rl/baselines/policies.py b/netforge_rl/baselines/policies.py index 6711544..fa4be84 100644 --- a/netforge_rl/baselines/policies.py +++ b/netforge_rl/baselines/policies.py @@ -24,9 +24,7 @@ def act(self, env, agent_id): class HeuristicBluePolicy: - """Blue: isolate the host implicated by the highest-severity SIEM alert in the - visible log buffer; otherwise analyse host 0. Reads telemetry, not ground truth, - so noise and log latency degrade it the way they would a real SOC workflow.""" + """Isolate the host on the highest-severity SIEM alert; else analyse host 0.""" name = 'heuristic-blue' alert_threshold = 5 @@ -80,9 +78,7 @@ def act(self, env, agent_id): class KillChainRedPolicy: - """Red kill-chain: port-scan a reachable, vulnerable host then exploit it, expanding - footholds DMZ -> Corporate -> Secure. Unlike HeuristicRedPolicy it recons first, so - its exploits pass the prior-state check and actually compromise hosts.""" + """Scan a reachable vulnerable host, then exploit. Recon first so exploits land.""" name = 'killchain-red' diff --git a/netforge_rl/baselines/protocol.py b/netforge_rl/baselines/protocol.py new file mode 100644 index 0000000..201ad0b --- /dev/null +++ b/netforge_rl/baselines/protocol.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence + +from netforge_rl.arena.evaluate import evaluate_split +from netforge_rl.arena.spec import ARENA_2026, ArenaSpec +from netforge_rl.baselines.registry import RUNNABLE, make_policy + + +def run_baseline_table( + names: Sequence[str] | None = None, + *, + spec: ArenaSpec | None = None, + split: str = 'dev', + scenarios: Sequence[str] = ('ransomware',), + seeds: Sequence[int] = (0, 1), + max_ticks: int = 40, + red_name: str = 'killchain-red', +) -> dict: + """Every algorithm under the same seeds, scenarios, budget, and metrics.""" + arena = spec or ARENA_2026 + names = list(names or sorted(RUNNABLE)) + rows = [] + for name in names: + if name not in RUNNABLE and not name.endswith(('.npz', '.pt', '.json')): + rows.append( + { + 'algorithm': name, + 'status': 'not_bundled', + 'note': 'Train separately; evaluate with this protocol.', + } + ) + continue + blue_factory: Callable + red_factory: Callable + if 'red' in name: + + def red_factory(n=name): + return make_policy(n) + + def blue_factory(): + return make_policy('heuristic-blue') + + team = 'red' + else: + + def red_factory(): + return make_policy(red_name) + + def blue_factory(n=name): + return make_policy(n) + + team = 'blue' + result = evaluate_split( + red_factory, + blue_factory, + split=split, # type: ignore[arg-type] + spec=arena, + scenarios=scenarios, + seeds=seeds, + max_ticks=max_ticks, + ) + mean = result['overall']['mean'] + rows.append( + { + 'algorithm': name, + 'team': team, + 'status': 'ok', + 'mission_success': mean.get('mission_success'), + 'sla_uptime': mean.get('sla_uptime'), + 'security': mean.get('security'), + 'false_positives': mean.get('false_positives'), + 'catastrophic_failure': mean.get('catastrophic_failure'), + 'cvar05_blue_return': result['overall']['cvar05'].get('blue_return'), + 'n_episodes': result['n_episodes'], + } + ) + return { + 'arena': arena.version, + 'split': split, + 'seeds': list(seeds), + 'scenarios': list(scenarios), + 'max_ticks': max_ticks, + 'compute_budget_env_steps': arena.compute_budget_env_steps, + 'rows': rows, + } + + +def format_baseline_table(table: dict) -> str: + header = ( + f'{"algorithm":<18}{"mission":>10}{"SLA":>8}{"security":>10}' + f'{"FP":>8}{"CVaR":>10}' + ) + lines = [header] + for row in table['rows']: + if row.get('status') != 'ok': + lines.append(f'{row["algorithm"]:<18} ({row.get("status")})') + continue + lines.append( + f'{row["algorithm"]:<18}' + f'{row["mission_success"]:10.2f}' + f'{row["sla_uptime"]:8.2f}' + f'{row["security"]:10.2f}' + f'{row["false_positives"]:8.2f}' + f'{row["cvar05_blue_return"]:10.2f}' + ) + return '\n'.join(lines) diff --git a/netforge_rl/baselines/registry.py b/netforge_rl/baselines/registry.py new file mode 100644 index 0000000..92953f8 --- /dev/null +++ b/netforge_rl/baselines/registry.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path + +from netforge_rl.arena.spec import ARENA_2026 +from netforge_rl.baselines.policies import ( + HeuristicBluePolicy, + HeuristicRedPolicy, + KillChainRedPolicy, + RandomPolicy, +) + +# Runnable here. MAPPO / QMIX / GNN stay out of tree; protocol names still list them. +RUNNABLE: dict[str, Callable] = { + 'random': lambda seed=0: RandomPolicy(seed=seed), + 'heuristic-blue': lambda seed=0: HeuristicBluePolicy(seed=seed), + 'heuristic-red': lambda seed=0: HeuristicRedPolicy(seed=seed), + 'killchain-red': lambda seed=0: KillChainRedPolicy(seed=seed), +} + +LEARNED_STUBS = { + 'ippo': 'netforge_rl.baselines.jax_ppo (extra: jax). Train then pass an .npz.', + 'mappo': 'RLlib via netforge_rl.bridges.rllib_bridge / benchmarks/rllib_rmappo.py', + 'rmappo': 'benchmarks/rllib_rmappo.py (LSTM MAPPO). Keep the algorithm outside.', + 'gnn-mappo': 'Train on info["graph"] belief graphs. Not bundled; evaluate here.', + 'qmix': 'Use a QMIX implementation (e.g. JaxMARL / RLlib) against this gym.', + 'llm': 'netforge_rl.semantic runner. Prompt-only / SFT / RL-LLM compared in Arena.', +} + + +def make_policy(name: str, seed: int = 0): + key = name.lower() + if key in RUNNABLE: + return RUNNABLE[key](seed=seed) + path = Path(name) + if path.suffix == '.npz': + from netforge_rl.baselines.jax_ppo import load_params + + return load_params(str(path)) + if key in LEARNED_STUBS: + raise FileNotFoundError( + f'{name} is a reference baseline in the Arena protocol, not a bundled ' + f'checkpoint. {LEARNED_STUBS[key]}' + ) + if path.exists(): + raise FileNotFoundError( + f'No built-in loader for {path.suffix} checkpoints. Wrap your policy ' + f'in a BasePolicy with act(env, agent_id) and pass --policy heuristic-blue ' + f'or a Python factory. See docs/baselines.md.' + ) + raise KeyError( + f'Unknown policy {name!r}. Runnable: {sorted(RUNNABLE)}. ' + f'Protocol names: {list(ARENA_2026.baselines)}' + ) diff --git a/netforge_rl/cli/__init__.py b/netforge_rl/cli/__init__.py new file mode 100644 index 0000000..5ba1dda --- /dev/null +++ b/netforge_rl/cli/__init__.py @@ -0,0 +1,17 @@ +"""NetForge CLI.""" + +from __future__ import annotations + +import sys + +from netforge_rl.cli.cmds import COMMANDS +from netforge_rl.cli.parser import build_parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + return COMMANDS[args.cmd](args) + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/netforge_rl/cli/cmds/__init__.py b/netforge_rl/cli/cmds/__init__.py new file mode 100644 index 0000000..ee2e89a --- /dev/null +++ b/netforge_rl/cli/cmds/__init__.py @@ -0,0 +1,21 @@ +from netforge_rl.cli.cmds.arena import cmd_arena +from netforge_rl.cli.cmds.benchmark import cmd_benchmark +from netforge_rl.cli.cmds.collect import cmd_collect +from netforge_rl.cli.cmds.evaluate import cmd_evaluate +from netforge_rl.cli.cmds.questions import cmd_questions +from netforge_rl.cli.cmds.replay import cmd_replay +from netforge_rl.cli.cmds.run import cmd_run +from netforge_rl.cli.cmds.test_policy import cmd_test_policy +from netforge_rl.cli.cmds.version import cmd_version + +COMMANDS = { + 'run': cmd_run, + 'test-policy': cmd_test_policy, + 'benchmark': cmd_benchmark, + 'arena': cmd_arena, + 'replay': cmd_replay, + 'collect': cmd_collect, + 'questions': cmd_questions, + 'evaluate': cmd_evaluate, + 'version': cmd_version, +} diff --git a/netforge_rl/cli/cmds/arena.py b/netforge_rl/cli/cmds/arena.py new file mode 100644 index 0000000..d8a0388 --- /dev/null +++ b/netforge_rl/cli/cmds/arena.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import json + +from netforge_rl.arena.spec import ARENA_NAME, current_arena + + +def cmd_arena(args) -> int: + spec = current_arena() + if args.split: + from netforge_rl.arena.evaluate import evaluate_split + from netforge_rl.baselines.registry import make_policy + + result = evaluate_split( + lambda: make_policy(args.red), + lambda: make_policy(args.policy), + split=args.split, + scenarios=('ransomware',), + seeds=(0, 1), + max_ticks=40, + ) + print(json.dumps(result, indent=2)) + return 0 + payload = spec.to_dict() + if args.json: + print(json.dumps(payload, indent=2)) + else: + print(ARENA_NAME) + print(json.dumps(payload, indent=2)) + return 0 diff --git a/netforge_rl/cli/cmds/benchmark.py b/netforge_rl/cli/cmds/benchmark.py new file mode 100644 index 0000000..34d141a --- /dev/null +++ b/netforge_rl/cli/cmds/benchmark.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import json + + +def cmd_benchmark(args) -> int: + from netforge_rl.arena.evaluate import evaluate_split + from netforge_rl.arena.leaderboard import format_leaderboard, save_leaderboard + from netforge_rl.baselines.protocol import format_baseline_table, run_baseline_table + from netforge_rl.baselines.registry import make_policy + + if args.table: + table = run_baseline_table( + seeds=tuple(args.seeds), + scenarios=tuple(args.scenarios), + max_ticks=args.max_ticks, + split=args.split, + ) + print(format_baseline_table(table)) + if args.out: + args.out.write_text(json.dumps(table, indent=2)) + return 0 + result = evaluate_split( + lambda: make_policy(args.red), + lambda: make_policy(args.policy), + split=args.split, + scenarios=tuple(args.scenarios), + seeds=tuple(args.seeds), + max_ticks=args.max_ticks, + ) + result['name'] = args.policy + print(json.dumps(result['overall'], indent=2)) + if args.out: + save_leaderboard([result], args.out) + print(format_leaderboard([result])) + return 0 diff --git a/netforge_rl/cli/cmds/collect.py b/netforge_rl/cli/cmds/collect.py new file mode 100644 index 0000000..baadc23 --- /dev/null +++ b/netforge_rl/cli/cmds/collect.py @@ -0,0 +1,20 @@ +def cmd_collect(args) -> int: + from netforge_rl.baselines.registry import make_policy + from netforge_rl.datasets.collector import collect_dataset + from netforge_rl.datasets.minari_export import export_npz + + trajs = collect_dataset( + lambda: make_policy(args.red), + lambda: make_policy(args.policy), + n_episodes=args.episodes, + max_ticks=args.max_ticks, + quality=args.quality, + ) + if args.hdf5: + from netforge_rl.datasets.minari_export import export_hdf5 + + path = export_hdf5(trajs, args.out.with_suffix('.hdf5')) + else: + path = export_npz(trajs, args.out) + print(f'wrote {len(trajs)} episodes -> {path}') + return 0 diff --git a/netforge_rl/cli/cmds/evaluate.py b/netforge_rl/cli/cmds/evaluate.py new file mode 100644 index 0000000..e6b21a1 --- /dev/null +++ b/netforge_rl/cli/cmds/evaluate.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +import json + + +def cmd_evaluate(args) -> int: + from netforge_rl.arena.submissions import evaluate_submission, load_submission + + sub = load_submission(args.submission) + result = evaluate_submission( + sub, + split=args.split, + scenarios=args.scenarios, + seeds=args.seeds, + max_ticks=args.max_ticks, + official=args.official, + agent_card=args.card, + ) + text = json.dumps(result, indent=2, default=str) + print(text) + if args.out: + args.out.write_text(text) + return 0 diff --git a/netforge_rl/cli/cmds/questions.py b/netforge_rl/cli/cmds/questions.py new file mode 100644 index 0000000..869e4f9 --- /dev/null +++ b/netforge_rl/cli/cmds/questions.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import json + + +def cmd_questions(args) -> int: + from netforge_rl.arena.questions import ( + format_catalog, + list_questions, + run_all_runnable, + run_question, + ) + + if args.all: + payload = run_all_runnable(seeds=tuple(args.seeds), max_ticks=args.max_ticks) + print(json.dumps(payload, indent=2, default=str)) + return 0 + if args.question: + payload = run_question( + args.question, seeds=tuple(args.seeds), max_ticks=args.max_ticks + ) + print(json.dumps(payload, indent=2, default=str)) + return 0 + if args.json: + print(json.dumps(list_questions(), indent=2)) + else: + print(format_catalog()) + return 0 diff --git a/netforge_rl/cli/cmds/replay.py b/netforge_rl/cli/cmds/replay.py new file mode 100644 index 0000000..eb1d598 --- /dev/null +++ b/netforge_rl/cli/cmds/replay.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +import json +from pathlib import Path + + +def cmd_replay(args) -> int: + from netforge_rl.render.replay_html import write_replay_html + + payload = json.loads(Path(args.trajectory).read_text()) + out = write_replay_html(payload, args.out) + print(out) + return 0 diff --git a/netforge_rl/cli/cmds/run.py b/netforge_rl/cli/cmds/run.py new file mode 100644 index 0000000..a8eba0d --- /dev/null +++ b/netforge_rl/cli/cmds/run.py @@ -0,0 +1,47 @@ +from netforge_rl.cli.util import policies + + +def cmd_run(args) -> int: + from netforge_rl.arena.evaluate import run_episode + from netforge_rl.render.replay_html import ReplayLog, write_replay_html + from netforge_rl.scenarios.yaml_dsl import make_env_from_spec, resolve_scenario + + spec = resolve_scenario(args.scenario) + red, blue = policies(args.policy, args.red) + rec = run_episode( + red, + blue, + scenario=args.scenario, + seed=args.seed, + max_ticks=args.max_ticks, + ) + print(f'{spec.name} seed={args.seed}') + m = rec.metrics + print( + f' mission={m.mission_success:.0f} SLA={m.sla_uptime:.2f} ' + f'security={m.security:.2f} FP={m.false_positives:.0f} ' + f'catastrophic={m.catastrophic_failure:.0f}' + ) + if args.replay: + env = make_env_from_spec(spec, seed=None, max_ticks=args.max_ticks) + log = ReplayLog(scenario=spec.name, seed=args.seed) + env.reset(seed=args.seed) + red, blue = policies(args.policy, args.red) + while env.agents: + actions = {} + for agent_id in env.agents: + pol = red if 'red' in agent_id else blue + actions[agent_id] = pol.act(env, agent_id) + _, rewards, term, trunc, infos = env.step(actions) + log.capture(env, rewards, infos, actions) + if all(term.values()) or all(trunc.values()): + break + html = ( + args.replay.with_suffix('.html') + if args.replay.suffix != '.html' + else args.replay + ) + write_replay_html(log.to_dict(), html) + log.save_json(html.with_suffix('.json')) + print(f' replay: {html}') + return 0 diff --git a/netforge_rl/cli/cmds/test_policy.py b/netforge_rl/cli/cmds/test_policy.py new file mode 100644 index 0000000..d1ad1e2 --- /dev/null +++ b/netforge_rl/cli/cmds/test_policy.py @@ -0,0 +1,13 @@ +def cmd_test_policy(args) -> int: + from netforge_rl.arena.agent_card import build_agent_card, format_agent_card + from netforge_rl.baselines.registry import make_policy + + card = build_agent_card( + lambda: make_policy(args.policy), + name=args.policy, + seeds=tuple(args.seeds), + run_arena=not args.no_arena, + out_dir=str(args.out) if args.out else None, + ) + print(format_agent_card(card)) + return 0 diff --git a/netforge_rl/cli/cmds/version.py b/netforge_rl/cli/cmds/version.py new file mode 100644 index 0000000..4fcaeda --- /dev/null +++ b/netforge_rl/cli/cmds/version.py @@ -0,0 +1,13 @@ +from netforge_rl.arena.spec import ARENA_NAME + + +def cmd_version(_args=None) -> int: + try: + from importlib.metadata import PackageNotFoundError, version + + ver = version('netforge-rl') + except PackageNotFoundError: + from netforge_rl import __version__ as ver + print(f'netforge-rl {ver}') + print(ARENA_NAME) + return 0 diff --git a/netforge_rl/cli/parser.py b/netforge_rl/cli/parser.py new file mode 100644 index 0000000..212ade7 --- /dev/null +++ b/netforge_rl/cli/parser.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog='netforge', + description='Cybersecurity gym: run, evaluate, diagnose.', + ) + sub = parser.add_subparsers(dest='cmd', required=True) + + p_run = sub.add_parser('run', help='Run a scenario YAML or named pack') + p_run.add_argument('scenario', help='Pack name or path to .yaml') + p_run.add_argument('--policy', default='heuristic-blue') + p_run.add_argument('--red', default='killchain-red') + p_run.add_argument('--seed', type=int, default=0) + p_run.add_argument('--max-ticks', type=int, default=80) + p_run.add_argument('--replay', type=Path, default=None) + + p_test = sub.add_parser('test-policy', help='Capability tests + agent card') + p_test.add_argument('--policy', default='heuristic-blue') + p_test.add_argument('--seeds', type=int, nargs='+', default=[0, 1]) + p_test.add_argument('--out', type=Path, default=None) + p_test.add_argument('--no-arena', action='store_true') + + p_bench = sub.add_parser('benchmark', help='Arena evaluation / baseline table') + p_bench.add_argument('policy', nargs='?', default='heuristic-blue') + p_bench.add_argument('--split', default='dev', choices=['train', 'dev', 'hidden']) + p_bench.add_argument('--red', default='killchain-red') + p_bench.add_argument('--seeds', type=int, nargs='+', default=[0, 1]) + p_bench.add_argument('--max-ticks', type=int, default=40) + p_bench.add_argument('--scenarios', nargs='+', default=['ransomware']) + p_bench.add_argument('--table', action='store_true', help='Run the reference table') + p_bench.add_argument('--out', type=Path, default=None) + + p_arena = sub.add_parser('arena', help='Print Arena spec or run a split') + p_arena.add_argument('--split', default=None, choices=['train', 'dev', 'hidden']) + p_arena.add_argument('--policy', default='heuristic-blue') + p_arena.add_argument('--red', default='killchain-red') + p_arena.add_argument('--json', action='store_true') + + p_rep = sub.add_parser('replay', help='Write an interactive HTML replay') + p_rep.add_argument('trajectory', type=Path) + p_rep.add_argument('--out', type=Path, default=Path('replay.html')) + + p_ds = sub.add_parser('collect', help='Collect an offline-RL dataset') + p_ds.add_argument('--policy', default='heuristic-blue') + p_ds.add_argument('--red', default='random') + p_ds.add_argument('--episodes', type=int, default=4) + p_ds.add_argument('--max-ticks', type=int, default=30) + p_ds.add_argument('--quality', default='mixed') + p_ds.add_argument('--out', type=Path, default=Path('netforge_offline.npz')) + p_ds.add_argument( + '--hdf5', + action='store_true', + help='Write Minari-like HDF5 (needs h5py) instead of npz', + ) + + p_q = sub.add_parser('questions', help='Catalog / run named RL probes') + p_q.add_argument( + 'question', nargs='?', default=None, help='Question id, or omit to list' + ) + p_q.add_argument('--all', action='store_true', help='Run every runnable question') + p_q.add_argument('--json', action='store_true') + p_q.add_argument('--seeds', type=int, nargs='+', default=[0]) + p_q.add_argument('--max-ticks', type=int, default=16) + + p_ev = sub.add_parser('evaluate', help='CAGE-style submission eval (multi-metric)') + p_ev.add_argument('submission', type=Path, help='Python module with make_blue()') + p_ev.add_argument('--split', default='dev', choices=['train', 'dev', 'hidden']) + p_ev.add_argument('--seeds', type=int, nargs='+', default=None) + p_ev.add_argument('--max-ticks', type=int, default=None) + p_ev.add_argument('--scenarios', nargs='+', default=None) + p_ev.add_argument('--official', action='store_true', help='Arena 2026 freeze') + p_ev.add_argument('--card', action='store_true', help='Also run capability probes') + p_ev.add_argument('--out', type=Path, default=None) + + sub.add_parser('version', help='Print package / Arena version') + return parser diff --git a/netforge_rl/cli/util.py b/netforge_rl/cli/util.py new file mode 100644 index 0000000..2a26205 --- /dev/null +++ b/netforge_rl/cli/util.py @@ -0,0 +1,4 @@ +def policies(blue_name: str, red_name: str): + from netforge_rl.baselines.registry import make_policy + + return make_policy(red_name), make_policy(blue_name) diff --git a/pyproject.toml b/pyproject.toml index 8799a47..09d5284 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,14 +3,27 @@ requires = ["setuptools>=61.0.0", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "netforge_rl" -version = "3.1.0" -description = "Multi-Agent Cybersecurity Simulator based on CybORG" +name = "netforge-rl" +version = "4.0.0" +description = "Cybersecurity gym for reinforcement learning" authors = [ - { name = "marl-Cyborg", email = "igorjankowwski@gmail.com" } + { name = "Igor Jankowski", email = "igorjankowwski@gmail.com" } ] readme = "README.md" requires-python = ">=3.12" +keywords = [ + "reinforcement-learning", + "multi-agent", + "cybersecurity", + "pettingzoo", + "gymnasium", + "benchmark", +] +classifiers = [ + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] dependencies = [ "numpy>=2.4.0", "gymnasium>=1.2.2", @@ -33,9 +46,10 @@ dev = [ "ruff>=0.3.0", ] docs = [ - "mkdocs>=1.6.0", - "mkdocs-material>=9.5.0", - "mkdocstrings[python]>=0.25.0", + "sphinx>=8.0", + "furo>=2024.8.6", + "myst-parser>=4.0", + "sphinxcontrib-mermaid>=1.0.0", ] jax = [ "jax[cpu]>=0.10.0", @@ -67,6 +81,18 @@ llm_clients = [ "openai>=1.40.0", "google-generativeai>=0.7.0", ] +offline = [ + "minari>=0.5.0", + "h5py>=3.10.0", +] + +[project.scripts] +netforge = "netforge_rl.cli:main" + +[project.urls] +Homepage = "https://github.com/reforcemind/NetForge_RL" +Documentation = "https://reforcemind.github.io/NetForge_RL/" +Issues = "https://github.com/reforcemind/NetForge_RL/issues" [tool.pytest.ini_options] testpaths = ["tests"] @@ -94,8 +120,12 @@ docstring-code-line-length = "dynamic" [tool.ruff.lint.per-file-ignores] "tests/**/*.py" = ["E402"] "netforge_rl/semantic/clients/*.py" = ["F401"] +"netforge_rl/render/replay_html.py" = ["E501"] [tool.ruff.lint] ignore = ["E701", "E402", "E702"] [tool.setuptools.packages.find] -include = ["netforge_rl*"] +include = ["netforge_rl*", "netforge"] + +[tool.setuptools.package-data] +netforge_rl = ["scenarios/packs/*.yaml"] diff --git a/tests/test_import_order.py b/tests/test_import_order.py index bfd81c7..4484308 100644 --- a/tests/test_import_order.py +++ b/tests/test_import_order.py @@ -19,7 +19,33 @@ def test_no_circular_import(snippet): proc = subprocess.run( [sys.executable, '-c', snippet + '; print("ok")'], capture_output=True, + check=False, text=True, ) assert proc.returncode == 0, proc.stderr assert 'ok' in proc.stdout + + +@pytest.mark.fast +def test_python_m_netforge_version(): + proc = subprocess.run( + [sys.executable, '-m', 'netforge', 'version'], + capture_output=True, + check=False, + text=True, + ) + assert proc.returncode == 0, proc.stderr + assert 'netforge-rl' in proc.stdout + assert 'Arena' in proc.stdout + + +@pytest.mark.fast +def test_python_m_netforge_questions_catalog(): + proc = subprocess.run( + [sys.executable, '-m', 'netforge', 'questions'], + capture_output=True, + check=False, + text=True, + ) + assert proc.returncode == 0, proc.stderr + assert 'belief-vs-oracle' in proc.stdout From 215ee553505df4890b0527ed8faac2109f095b24 Mon Sep 17 00:00:00 2001 From: elprofesoriqo Date: Sun, 20 Sep 2026 15:15:25 +0200 Subject: [PATCH 6/8] Add HTML replay, offline datasets, league helpers, and a CAGE-style example. These sit around the gym so researchers can export trajectories and submissions without vendoring MAPPO or QMIX. Co-authored-by: Cursor --- examples/my_scenario.yaml | 7 + examples/submissions/heuristic_blue.py | 19 ++ netforge_rl/comm/__init__.py | 3 + netforge_rl/comm/wrapper.py | 78 ++++++++ netforge_rl/datasets/__init__.py | 17 ++ netforge_rl/datasets/collector.py | 103 +++++++++++ netforge_rl/datasets/minari_export.py | 140 ++++++++++++++ netforge_rl/league/__init__.py | 4 + netforge_rl/league/pool.py | 27 +++ netforge_rl/league/psro.py | 41 +++++ netforge_rl/render/__init__.py | 17 +- netforge_rl/render/replay_html.py | 243 +++++++++++++++++++++++++ netforge_rl/render/trajectory.py | 7 +- netforge_rl/rewards/__init__.py | 3 + netforge_rl/rewards/variants.py | 87 +++++++++ tests/render/test_render_pipeline.py | 16 +- 16 files changed, 802 insertions(+), 10 deletions(-) create mode 100644 examples/my_scenario.yaml create mode 100644 examples/submissions/heuristic_blue.py create mode 100644 netforge_rl/comm/__init__.py create mode 100644 netforge_rl/comm/wrapper.py create mode 100644 netforge_rl/datasets/__init__.py create mode 100644 netforge_rl/datasets/collector.py create mode 100644 netforge_rl/datasets/minari_export.py create mode 100644 netforge_rl/league/__init__.py create mode 100644 netforge_rl/league/pool.py create mode 100644 netforge_rl/league/psro.py create mode 100644 netforge_rl/render/replay_html.py create mode 100644 netforge_rl/rewards/__init__.py create mode 100644 netforge_rl/rewards/variants.py diff --git a/examples/my_scenario.yaml b/examples/my_scenario.yaml new file mode 100644 index 0000000..83b05d2 --- /dev/null +++ b/examples/my_scenario.yaml @@ -0,0 +1,7 @@ +name: my_experiment +base: ransomware +difficulty: medium +max_ticks: 200 +description: Minimal custom experiment — no simulator code changes. +reward_variant: constrained_sla +log_latency: 2 diff --git a/examples/submissions/heuristic_blue.py b/examples/submissions/heuristic_blue.py new file mode 100644 index 0000000..a8ac9a9 --- /dev/null +++ b/examples/submissions/heuristic_blue.py @@ -0,0 +1,19 @@ +"""CAGE-style ``make_blue`` example. + + netforge evaluate examples/submissions/heuristic_blue.py --seeds 0 --max-ticks 20 +""" + +NAME = 'heuristic-blue' +TEAM = 'netforge-baselines' +TECHNIQUE = 'scripted-siem-isolate' +PAPER = '' + + +def wrap(env): + return env + + +def make_blue(seed=0): + from netforge_rl.baselines.policies import HeuristicBluePolicy + + return HeuristicBluePolicy(seed=seed) diff --git a/netforge_rl/comm/__init__.py b/netforge_rl/comm/__init__.py new file mode 100644 index 0000000..c415c95 --- /dev/null +++ b/netforge_rl/comm/__init__.py @@ -0,0 +1,3 @@ +from netforge_rl.comm.wrapper import CommFailureWrapper + +__all__ = ['CommFailureWrapper'] diff --git a/netforge_rl/comm/wrapper.py b/netforge_rl/comm/wrapper.py new file mode 100644 index 0000000..648cae1 --- /dev/null +++ b/netforge_rl/comm/wrapper.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from collections import deque + +import numpy as np +from pettingzoo.utils.wrappers import BaseParallelWrapper + + +class CommFailureWrapper(BaseParallelWrapper): + """Drop, delay, corrupt, or bandwidth-limit the Blue shared comm channel.""" + + def __init__( + self, + env, + drop_p: float = 0.0, + delay_steps: int = 0, + corrupt_p: float = 0.0, + bandwidth: int | None = None, + seed: int = 0, + ): + super().__init__(env) + self.drop_p = float(drop_p) + self.delay_steps = int(delay_steps) + self.corrupt_p = float(corrupt_p) + self.bandwidth = bandwidth + self._rng = np.random.default_rng(seed) + self._queue: deque[np.ndarray] = deque() + + def reset(self, seed=None, options=None): + self._queue.clear() + if seed is not None: + self._rng = np.random.default_rng(seed) + obs, infos = super().reset(seed=seed, options=options) + return self._apply(obs), infos + + def step(self, actions): + obs, rewards, term, trunc, infos = super().step(actions) + return self._apply(obs), rewards, term, trunc, infos + + def _apply(self, obs: dict) -> dict: + template = None + for agent, payload in obs.items(): + if 'blue' in agent and isinstance(payload, dict) and 'blue_comm' in payload: + template = np.array(payload['blue_comm'], dtype=np.float32, copy=True) + break + if template is None: + return obs + delayed = self._delay(template) + degraded = self._degrade(delayed) + for agent, payload in obs.items(): + if 'blue' in agent and isinstance(payload, dict) and 'blue_comm' in payload: + payload['blue_comm'] = degraded + return obs + + def _delay(self, comm: np.ndarray) -> np.ndarray: + if self.delay_steps <= 0: + return comm + self._queue.append(comm) + if len(self._queue) <= self.delay_steps: + return np.zeros_like(comm) + return self._queue.popleft() + + def _degrade(self, comm: np.ndarray) -> np.ndarray: + out = comm.copy() + if self.drop_p > 0 and float(self._rng.random()) < self.drop_p: + return np.zeros_like(out) + if self.corrupt_p > 0: + mask = self._rng.random(out.shape) < self.corrupt_p + out = np.where(mask, self._rng.random(out.shape, dtype=np.float32), out) + if self.bandwidth is not None: + k = max(0, int(self.bandwidth)) + if k < out.size: + keep = np.argpartition(out, -k)[-k:] if k else [] + clipped = np.zeros_like(out) + if k: + clipped[keep] = out[keep] + out = clipped + return out.astype(np.float32) diff --git a/netforge_rl/datasets/__init__.py b/netforge_rl/datasets/__init__.py new file mode 100644 index 0000000..6253955 --- /dev/null +++ b/netforge_rl/datasets/__init__.py @@ -0,0 +1,17 @@ +from netforge_rl.datasets.collector import Trajectory, collect_dataset, collect_episode +from netforge_rl.datasets.minari_export import ( + export_hdf5, + export_npz, + trajectories_to_minari_dict, + try_minari_export, +) + +__all__ = [ + 'Trajectory', + 'collect_dataset', + 'collect_episode', + 'export_hdf5', + 'export_npz', + 'trajectories_to_minari_dict', + 'try_minari_export', +] diff --git a/netforge_rl/datasets/collector.py b/netforge_rl/datasets/collector.py new file mode 100644 index 0000000..c8f2de4 --- /dev/null +++ b/netforge_rl/datasets/collector.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np + +from netforge_rl.arena.evaluate import _act +from netforge_rl.environment.presets import make_env + + +@dataclass +class Trajectory: + scenario: str + seed: int + policy: str + observations: list = field(default_factory=list) + actions: list = field(default_factory=list) + rewards: list = field(default_factory=list) + terminals: list = field(default_factory=list) + truncations: list = field(default_factory=list) + infos: list = field(default_factory=list) + timeouts: list = field(default_factory=list) + + def __len__(self) -> int: + return len(self.rewards) + + +def collect_episode( + red_policy, + blue_policy, + *, + scenario: str = 'ransomware', + seed: int = 0, + max_ticks: int = 100, + controlled: str = 'blue', + quality: str = 'mixed', +) -> Trajectory: + """Roll one trajectory for offline RL. Stores the controlled team's steps.""" + env = make_env('medium', scenario_type=scenario, seed=None, max_ticks=max_ticks) + obs, _ = env.reset(seed=seed) + traj = Trajectory(scenario=scenario, seed=seed, policy=quality) + while env.agents: + actions = {} + for agent_id in env.agents: + policy = red_policy if 'red' in agent_id else blue_policy + actions[agent_id] = _act(policy, env, agent_id, obs.get(agent_id, {})) + next_obs, rewards, term, trunc, infos = env.step(actions) + team_agents = [a for a in rewards if controlled in a] + step_obs = {a: _serialize_obs(obs[a]) for a in team_agents if a in obs} + step_act = {a: np.asarray(actions[a]).tolist() for a in team_agents} + step_rew = {a: float(rewards[a]) for a in team_agents} + traj.observations.append(step_obs) + traj.actions.append(step_act) + traj.rewards.append(step_rew) + traj.terminals.append({a: bool(term.get(a, False)) for a in team_agents}) + traj.truncations.append({a: bool(trunc.get(a, False)) for a in team_agents}) + traj.infos.append( + { + a: { + k: v + for k, v in infos.get(a, {}).items() + if isinstance(v, (int, float, str, bool)) + } + for a in team_agents + } + ) + obs = next_obs + if all(term.values()) or all(trunc.values()): + break + return traj + + +def collect_dataset( + red_factory, + blue_factory, + *, + n_episodes: int = 4, + scenario: str = 'ransomware', + max_ticks: int = 40, + quality: str = 'mixed', + seed0: int = 0, +) -> list[Trajectory]: + return [ + collect_episode( + red_factory(), + blue_factory(), + scenario=scenario, + seed=seed0 + i, + max_ticks=max_ticks, + quality=quality, + ) + for i in range(n_episodes) + ] + + +def _serialize_obs(obs: dict) -> dict: + out = {} + for key, value in obs.items(): + if isinstance(value, np.ndarray): + out[key] = value.astype(np.float32).tolist() + else: + out[key] = value + return out diff --git a/netforge_rl/datasets/minari_export.py b/netforge_rl/datasets/minari_export.py new file mode 100644 index 0000000..4f4176c --- /dev/null +++ b/netforge_rl/datasets/minari_export.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import json +from collections.abc import Sequence +from pathlib import Path + +import numpy as np + +from netforge_rl.datasets.collector import Trajectory + +MINARI_SPEC = '0.5' +DATASET_ID_PREFIX = 'netforge' + + +def trajectories_to_minari_dict( + trajectories: Sequence[Trajectory], + dataset_id: str = 'netforge/arena-mixed-v0', +) -> dict: + """Minari-compatible episode dict (no hard dependency on the minari package).""" + episodes = [] + for traj in trajectories: + agents = sorted({a for step in traj.rewards for a in step}) + agent = agents[0] if agents else 'blue_dmz' + obs = np.array( + [ + step.get(agent, {}).get('obs', np.zeros(256, dtype=np.float32)) + for step in traj.observations + ], + dtype=np.float32, + ) + acts = np.array( + [step.get(agent, [0, 0]) for step in traj.actions], dtype=np.int64 + ) + rews = np.array( + [step.get(agent, 0.0) for step in traj.rewards], dtype=np.float32 + ) + terms = np.array( + [step.get(agent, False) for step in traj.terminals], dtype=np.bool_ + ) + truncs = np.array( + [step.get(agent, False) for step in traj.truncations], dtype=np.bool_ + ) + episodes.append( + { + 'observations': obs, + 'actions': acts, + 'rewards': rews, + 'terminations': terms, + 'truncations': truncs, + 'seed': traj.seed, + 'scenario': traj.scenario, + 'quality': traj.policy, + } + ) + return { + 'id': dataset_id, + 'minari_version': MINARI_SPEC, + 'env_spec': 'netforge_rl/NetForge-v4', + 'n_episodes': len(episodes), + 'episodes': episodes, + 'notes': ( + 'Offline-RL export. Flattened Blue obs for the first controlled agent. ' + 'Publish to Hugging Face as parquet/npz; load with minari if installed.' + ), + } + + +def export_npz(trajectories: Sequence[Trajectory], path: str | Path) -> Path: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + payload = trajectories_to_minari_dict(trajectories) + arrays = {} + meta_episodes = [] + for i, ep in enumerate(payload['episodes']): + arrays[f'ep{i}_obs'] = ep['observations'] + arrays[f'ep{i}_act'] = ep['actions'] + arrays[f'ep{i}_rew'] = ep['rewards'] + arrays[f'ep{i}_term'] = ep['terminations'] + arrays[f'ep{i}_trunc'] = ep['truncations'] + meta_episodes.append( + { + 'seed': ep['seed'], + 'scenario': ep['scenario'], + 'quality': ep['quality'], + 'length': len(ep['rewards']), + } + ) + meta = { + k: payload[k] + for k in ('id', 'minari_version', 'env_spec', 'n_episodes', 'notes') + } + meta['episodes'] = meta_episodes + np.savez_compressed(path, meta=np.array(json.dumps(meta)), **arrays) + return path + + +def export_hdf5(trajectories: Sequence[Trajectory], path: str | Path) -> Path: + """Minari-like HDF5 (episode groups). Requires optional ``h5py`` / ``offline`` extra.""" + try: + import h5py + except ImportError as exc: + raise ImportError( + 'h5py is optional. pip install "netforge-rl[offline]" or use export_npz().' + ) from exc + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + payload = trajectories_to_minari_dict(trajectories) + with h5py.File(path, 'w') as handle: + handle.attrs['id'] = payload['id'] + handle.attrs['minari_version'] = payload['minari_version'] + handle.attrs['env_spec'] = payload['env_spec'] + handle.attrs['n_episodes'] = payload['n_episodes'] + handle.attrs['notes'] = payload['notes'] + for i, ep in enumerate(payload['episodes']): + group = handle.create_group(f'episode_{i}') + group.create_dataset('observations', data=np.asarray(ep['observations'])) + group.create_dataset('actions', data=np.asarray(ep['actions'])) + group.create_dataset('rewards', data=np.asarray(ep['rewards'])) + group.create_dataset('terminations', data=np.asarray(ep['terminations'])) + group.create_dataset('truncations', data=np.asarray(ep['truncations'])) + group.attrs['seed'] = int(ep['seed']) + group.attrs['scenario'] = str(ep['scenario']) + group.attrs['quality'] = str(ep['quality']) + return path + + +def try_minari_export( + trajectories: Sequence[Trajectory], + dataset_id: str, + path: str | Path | None = None, +) -> Path: + """Write HDF5 if h5py is present; otherwise npz. Full Minari create is optional.""" + dest = ( + Path(path) if path is not None else Path(f'{dataset_id.replace("/", "_")}.hdf5') + ) + try: + return export_hdf5(trajectories, dest) + except ImportError: + npz_path = dest.with_suffix('.npz') + return export_npz(trajectories, npz_path) diff --git a/netforge_rl/league/__init__.py b/netforge_rl/league/__init__.py new file mode 100644 index 0000000..ad916dd --- /dev/null +++ b/netforge_rl/league/__init__.py @@ -0,0 +1,4 @@ +from netforge_rl.league.pool import OpponentPool +from netforge_rl.league.psro import payoff_matrix + +__all__ = ['OpponentPool', 'payoff_matrix'] diff --git a/netforge_rl/league/pool.py b/netforge_rl/league/pool.py new file mode 100644 index 0000000..d4b31a5 --- /dev/null +++ b/netforge_rl/league/pool.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field + +import numpy as np + + +@dataclass +class OpponentPool: + """League of current, historical, and specialist opponent factories.""" + + members: dict[str, Callable] = field(default_factory=dict) + history: list[str] = field(default_factory=list) + + def add(self, name: str, factory: Callable) -> None: + self.members[name] = factory + self.history.append(name) + + def sample(self, rng: np.random.Generator | None = None) -> tuple[str, Callable]: + rng = rng or np.random.default_rng() + name = str(rng.choice(list(self.members))) + return name, self.members[name] + + def historical(self, n: int = 3) -> list[tuple[str, Callable]]: + names = self.history[-n:] + return [(name, self.members[name]) for name in names if name in self.members] diff --git a/netforge_rl/league/psro.py b/netforge_rl/league/psro.py new file mode 100644 index 0000000..82486ca --- /dev/null +++ b/netforge_rl/league/psro.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence + +from netforge_rl.arena.evaluate import run_episode + + +def payoff_matrix( + red_population: Mapping[str, Callable], + blue_population: Mapping[str, Callable], + *, + scenarios: Sequence[str] = ('ransomware',), + seeds: Sequence[int] = (0,), + max_ticks: int = 40, + metric: str = 'sla_uptime', +) -> dict: + """Empirical Red vs Blue payoff for PSRO-style evaluation (not Elo-only).""" + red_names = list(red_population) + blue_names = list(blue_population) + matrix = [[0.0 for _ in blue_names] for _ in red_names] + for i, red_name in enumerate(red_names): + for j, blue_name in enumerate(blue_names): + values = [] + for scenario in scenarios: + for seed in seeds: + rec = run_episode( + red_population[red_name](), + blue_population[blue_name](), + scenario=scenario, + seed=int(seed), + max_ticks=max_ticks, + ) + values.append(float(getattr(rec.metrics, metric))) + matrix[i][j] = round(sum(values) / max(len(values), 1), 4) + return { + 'metric': metric, + 'red': red_names, + 'blue': blue_names, + 'payoff_blue': matrix, + 'note': 'Entry [r][b] is Blue metric vs Red. Use as a PSRO empirical game.', + } diff --git a/netforge_rl/render/__init__.py b/netforge_rl/render/__init__.py index 9dc73b9..5fe1142 100644 --- a/netforge_rl/render/__init__.py +++ b/netforge_rl/render/__init__.py @@ -1,12 +1,25 @@ +from netforge_rl.render.replay_html import ReplayLog, write_replay_html from netforge_rl.render.snapshot import Snapshot, snapshot_from_envstate -from netforge_rl.render.matplotlib_renderer import render_rgb -from netforge_rl.render.recorder import FrameRecorder from netforge_rl.render.trajectory import TrajectoryRecorder __all__ = [ 'FrameRecorder', + 'ReplayLog', 'Snapshot', 'TrajectoryRecorder', 'render_rgb', 'snapshot_from_envstate', + 'write_replay_html', ] + + +def __getattr__(name): + if name == 'FrameRecorder': + from netforge_rl.render.recorder import FrameRecorder + + return FrameRecorder + if name == 'render_rgb': + from netforge_rl.render.matplotlib_renderer import render_rgb + + return render_rgb + raise AttributeError(name) diff --git a/netforge_rl/render/replay_html.py b/netforge_rl/render/replay_html.py new file mode 100644 index 0000000..fa18646 --- /dev/null +++ b/netforge_rl/render/replay_html.py @@ -0,0 +1,243 @@ +from __future__ import annotations + +import json +from pathlib import Path + +from netforge_rl.core.graph_obs import build_graph_observation, build_oracle_graph +from netforge_rl.environment.constants import PADDING_SUBNET + + +class ReplayLog: + """Network + SIEM + actions + rewards for the interactive viewer.""" + + def __init__(self, scenario: str = '', seed: int = 0): + self.scenario = scenario + self.seed = seed + self.ticks: list[dict] = [] + + def capture( + self, env, rewards: dict, infos: dict, actions: dict | None = None + ) -> None: + gs = ( + env.unwrapped.global_state + if hasattr(env, 'unwrapped') + else env.global_state + ) + hosts = [] + for ip, host in sorted(gs.all_hosts.items()): + if host.subnet_cidr == PADDING_SUBNET: + continue + hosts.append( + { + 'ip': ip, + 'hostname': host.hostname, + 'subnet': gs.get_subnet_name(host.subnet_cidr), + 'status': host.status, + 'compromised': host.compromised_by != 'None', + 'isolated': host.status == 'isolated', + 'decoy': host.decoy not in (None, 'inactive'), + } + ) + siem = [] + for entry in gs.siem_log_buffer[-12:]: + log = entry[0] if isinstance(entry, tuple) else entry + siem.append(str(log)[:240]) + sample = next((v for k, v in infos.items() if 'blue' in k), {}) + self.ticks.append( + { + 'tick': int( + getattr( + env.unwrapped if hasattr(env, 'unwrapped') else env, + 'current_tick', + 0, + ) + ), + 'hosts': hosts, + 'siem': siem, + 'rewards': {k: float(v) for k, v in rewards.items()}, + 'actions': { + k: [int(x) for x in v] + if hasattr(v, '__iter__') and not isinstance(v, (str, bytes)) + else str(v) + for k, v in (actions or {}).items() + }, + 'metrics': { + 'sla': float(sample.get('SLA_Uptime_Percentage', 1.0)), + 'compromised': float(sample.get('compromised_hosts', 0.0)), + 'false_positives': float(sample.get('false_positives_total', 0.0)), + 'blue_reward': float( + sum(v for k, v in rewards.items() if 'blue' in k) + ), + 'red_reward': float( + sum(v for k, v in rewards.items() if 'red' in k) + ), + }, + 'belief_n': int( + build_graph_observation(gs, 'blue_dmz')['node_mask'].sum() + ), + 'oracle_n': int(build_oracle_graph(gs)['node_mask'].sum()), + } + ) + + def to_dict(self) -> dict: + return { + 'scenario': self.scenario, + 'seed': self.seed, + 'ticks': self.ticks, + } + + def save_json(self, path: str | Path) -> Path: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(self.to_dict())) + return path + + +def write_replay_html(payload: dict, path: str | Path) -> Path: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + blob = json.dumps(payload).replace(' + + + +NetForge Replay + + + +
+

NETFORGE ARENA · replay

+
+
+
+
+
belief graph · hosts from SIEM + inventory (oracle is diagnostic-only)
+ +
+ clean + compromised + isolated + decoy +
+
+
+
SIEM feed
+
+
+
+ + +
+
tick 0
+
SLA 1.00
+
owned 0
+
FP 0
+
+
+
+
Blue reward timeline
+ +
+
+ + + +""" diff --git a/netforge_rl/render/trajectory.py b/netforge_rl/render/trajectory.py index 33cbf21..6f31cfc 100644 --- a/netforge_rl/render/trajectory.py +++ b/netforge_rl/render/trajectory.py @@ -3,7 +3,6 @@ import json from dataclasses import asdict, dataclass, field from pathlib import Path -from typing import Optional @dataclass @@ -11,14 +10,14 @@ class StepRecord: tick: int agent_id: str action_name: str - target_ip: Optional[str] + target_ip: str | None success: bool reward: float @dataclass class TrajectoryRecorder: - """Records per-step (tick, agent, action, target, success, reward) for kill-chain analysis.""" + """Per-step kill-chain records.""" scenario: str = '' seed: int = 0 @@ -34,7 +33,7 @@ def record_step( tick: int, agent_id: str, action_name: str, - target_ip: Optional[str], + target_ip: str | None, success: bool, reward: float, ) -> None: diff --git a/netforge_rl/rewards/__init__.py b/netforge_rl/rewards/__init__.py new file mode 100644 index 0000000..ceb427b --- /dev/null +++ b/netforge_rl/rewards/__init__.py @@ -0,0 +1,3 @@ +from netforge_rl.rewards.variants import VARIANTS, RewardDesignWrapper + +__all__ = ['VARIANTS', 'RewardDesignWrapper'] diff --git a/netforge_rl/rewards/variants.py b/netforge_rl/rewards/variants.py new file mode 100644 index 0000000..18d9a75 --- /dev/null +++ b/netforge_rl/rewards/variants.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +from collections.abc import Callable + +from pettingzoo.utils.wrappers import BaseParallelWrapper + + +def _sla_only(agent, raw, info) -> float: + if 'blue' not in agent: + return float(raw) + return float(info.get('SLA_Uptime_Percentage', 0.0)) + + +def _security_only(agent, raw, info) -> float: + if 'blue' not in agent: + return float(raw) + return float(info.get('security', 0.0)) - 0.1 * float( + info.get('compromised_hosts', 0.0) + ) + + +def _fp_penalized(agent, raw, info) -> float: + if 'blue' not in agent: + return float(raw) + return float(raw) - 5.0 * float(info.get('false_positives_total', 0.0)) + + +def _constrained_sla(agent, raw, info) -> float: + if 'blue' not in agent: + return float(raw) + sla = float(info.get('SLA_Uptime_Percentage', 1.0)) + penalty = 0.0 if sla >= 0.9 else 10.0 * (0.9 - sla) + return float(raw) - penalty + + +def _safety_critical(agent, raw, info) -> float: + if 'blue' not in agent: + return float(raw) + if float(info.get('catastrophic_failure', 0.0)) >= 1.0: + return -100.0 + return float(raw) + + +def _sparse(agent, raw, info) -> float: + if 'blue' not in agent: + return 0.0 + return float(info.get('mission_success', 0.0)) + + +VARIANTS: dict[str, Callable] = { + 'default': lambda agent, raw, info: float(raw), + 'sla_only': _sla_only, + 'security_only': _security_only, + 'fp_penalized': _fp_penalized, + 'constrained_sla': _constrained_sla, + 'safety_critical': _safety_critical, + 'sparse': _sparse, +} + + +class RewardDesignWrapper(BaseParallelWrapper): + """Swap Blue's reward without editing scenario classes.""" + + def __init__(self, env, variant: str = 'default'): + super().__init__(env) + if variant not in VARIANTS: + raise KeyError( + f'Unknown reward variant {variant!r}. Choose {sorted(VARIANTS)}' + ) + self.variant = variant + self._reshape = VARIANTS[variant] + + def step(self, actions): + obs, rewards, term, trunc, infos = super().step(actions) + reshaped = { + agent: self._reshape(agent, rewards[agent], infos.get(agent, {})) + for agent in rewards + } + for agent, info in infos.items(): + info['raw_reward'] = float(rewards.get(agent, 0.0)) + info['reward_variant'] = self.variant + if self.variant == 'sparse' and (any(term.values()) or any(trunc.values())): + reshaped = { + agent: self._reshape(agent, rewards[agent], infos.get(agent, {})) + for agent in rewards + } + return obs, reshaped, term, trunc, infos diff --git a/tests/render/test_render_pipeline.py b/tests/render/test_render_pipeline.py index d2170f4..75bf41c 100644 --- a/tests/render/test_render_pipeline.py +++ b/tests/render/test_render_pipeline.py @@ -2,9 +2,13 @@ pytest.importorskip('matplotlib') import numpy as np -import pytest -from netforge_rl.render import FrameRecorder, render_rgb, snapshot_from_envstate -from netforge_rl.render.snapshot import COLOR_COMPROMISED, COLOR_ISOLATED, COLOR_SECURE +from netforge_rl.render.matplotlib_renderer import render_rgb +from netforge_rl.render.snapshot import ( + COLOR_COMPROMISED, + COLOR_ISOLATED, + COLOR_SECURE, + snapshot_from_envstate, +) AGENTS = ('red_operator', 'blue_dmz', 'blue_internal', 'blue_restricted') @@ -51,6 +55,9 @@ def test_render_rgb_shape(global_state) -> None: @pytest.mark.fast def test_frame_recorder_append_and_len() -> None: + pytest.importorskip('moviepy') + from netforge_rl.render.recorder import FrameRecorder + rec = FrameRecorder(fps=4) rec.append(np.zeros((10, 10, 3), dtype=np.uint8)) rec.append(np.ones((10, 10, 3), dtype=np.uint8)) @@ -59,7 +66,8 @@ def test_frame_recorder_append_and_len() -> None: @pytest.mark.fast def test_frame_recorder_save_requires_frames(tmp_path) -> None: - rec = FrameRecorder() + pytest.importorskip('moviepy') + from netforge_rl.render.recorder import FrameRecorder with pytest.raises(ValueError): rec.save(tmp_path / 'empty.gif') From 23c668ca4d919f087bd9f6cc52a5e65b6f1746f9 Mon Sep 17 00:00:00 2001 From: elprofesoriqo Date: Sun, 20 Sep 2026 15:21:48 +0200 Subject: [PATCH 7/8] Make ruff format happy on the example submission and render test. Unindent the evaluate command in the docstring and construct FrameRecorder before asserting empty save. Co-authored-by: Cursor --- examples/submissions/heuristic_blue.py | 2 +- tests/render/test_render_pipeline.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/submissions/heuristic_blue.py b/examples/submissions/heuristic_blue.py index a8ac9a9..8584b1a 100644 --- a/examples/submissions/heuristic_blue.py +++ b/examples/submissions/heuristic_blue.py @@ -1,6 +1,6 @@ """CAGE-style ``make_blue`` example. - netforge evaluate examples/submissions/heuristic_blue.py --seeds 0 --max-ticks 20 +netforge evaluate examples/submissions/heuristic_blue.py --seeds 0 --max-ticks 20 """ NAME = 'heuristic-blue' diff --git a/tests/render/test_render_pipeline.py b/tests/render/test_render_pipeline.py index 75bf41c..daf722d 100644 --- a/tests/render/test_render_pipeline.py +++ b/tests/render/test_render_pipeline.py @@ -68,6 +68,8 @@ def test_frame_recorder_append_and_len() -> None: def test_frame_recorder_save_requires_frames(tmp_path) -> None: pytest.importorskip('moviepy') from netforge_rl.render.recorder import FrameRecorder + + rec = FrameRecorder(fps=4) with pytest.raises(ValueError): rec.save(tmp_path / 'empty.gif') From 02305fe7297399e6bf6907c04452b6b6f75ba581 Mon Sep 17 00:00:00 2001 From: elprofesoriqo Date: Sun, 20 Sep 2026 15:25:00 +0200 Subject: [PATCH 8/8] Pin ruff lint to pycodestyle and pyflakes so CI matches the old default. Ruff 0.16 treats a bare ignore list as a much larger rule set, which fails on the existing simulator files. Co-authored-by: Cursor --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 09d5284..a22807c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,6 +122,7 @@ docstring-code-line-length = "dynamic" "netforge_rl/semantic/clients/*.py" = ["F401"] "netforge_rl/render/replay_html.py" = ["E501"] [tool.ruff.lint] +select = ["E4", "E7", "E9", "F"] ignore = ["E701", "E402", "E702"] [tool.setuptools.packages.find]