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/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..8584b1a --- /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/__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/__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/__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/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/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/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/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/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/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/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/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/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/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/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/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) 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/pyproject.toml b/pyproject.toml index 8799a47..a22807c 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,13 @@ 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] +select = ["E4", "E7", "E9", "F"] 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/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']) 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) 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() diff --git a/tests/render/test_render_pipeline.py b/tests/render/test_render_pipeline.py index d2170f4..daf722d 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,10 @@ 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 + + rec = FrameRecorder(fps=4) with pytest.raises(ValueError): rec.save(tmp_path / 'empty.gif') 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