Attack-Graph-Learning-Environment

A reinforcement learning framework for modeling adversarial cyber-attack and defense interactions on enterprise network attack graphs. The environment encodes a network topology as a directed graph derived from an attack graph model (e.g., MAL — Meta Attack Language), where edge weights are derived directly from CVSSv3 exploitability and impact metrics.

Overview

The framework implements a Stackelberg game between two RL agents:

• Attacker (DQN): Navigates the attack graph from a reconnaissance entry node to a target asset, maximising cumulative CVSS-weighted reward.
• Defender (GAT-AC / GAT-DQN): Observes Suricata-like alert streams processed through a Graph Attention Network (GAT) and applies structural countermeasures (network filtering, software patching, host isolation).

The environment is fully compatible with the Gymnasium interface (gym.Env), supporting both reset(seed, options) and step(action) with standard observation and action spaces.

Goal nodes are configurable: pass any node ID when constructing the environment, or let it auto-detect the deepest Access metaconcept node.

Environment

GraphEnvironment

The core MDP. Initialised as GraphEnvironment(graph_json, goal_node=None, config_path='config.json'). Pass any node ID for goal_node, or leave it as None to auto-detect the deepest Access metaconcept node. State space is the attack graph G = (V, E) where:

• Each node v ∈ V represents an attack step (Reconnaissance, CVE exploit, Privilege escalation, Host compromise, etc.)
• Each edge (u, v) ∈ E carries a weight derived from the CVSSv3 vector of the associated vulnerability.

Attacker reward for traversing edge (u → v) involving CVE c:

R_exploit(c) = AV · AC · PR · UI · NORM_F_ATTEMPT    (exploitability component)
R_impact(c)  = (1-C') · (1-I') · (1-A') · NORM_F    (CIA impact component)

where AV, AC, PR, UI, C', I', A' are the normalised CVSSv3 metric values.

Observation (per-node features, dim = 17):

#	Feature	Description
0	active	Node mask (1 = active, 0 = patched/removed)
1	risk	Unconditional compromise probability (noisy-OR forward propagation)
2	centrality	Betweenness centrality
3	vuln	Binary: 1 if CVE node
4	critical	Binary: 1 if goal node
5	z_score	Normalised z-score of severity stream for node's IP
6	entropy	Shannon entropy of alert severity distribution
7	alerted	Binary: 1 if node has active alert
8	alert_vol_src	Normalised source alert volume
9	alert_vol_dst	Normalised destination alert volume
10	cum_alerts	Cumulative per-node alert count / step
11	alert_recency	1/(steps_since_last_alert + 1)
12	goal_dist	Topological hop distance to goal, normalised
13	max_cvss	Max CVSS on incoming edges, normalised
14	def_filtered	Defender action flag: filtered
15	def_patched	Defender action flag: patched
16	def_restored	Defender action flag: restored

Risk propagation uses iterative noisy-OR forward pass:

P(v) = 1 - ∏_{u ∈ parents(v)} (1 - P(u) · exp(-k · dist(u,v)))

until convergence with tolerance ε = 1e-6.

Defender Actions

ID	Action	Effect
0	Do Nothing	No structural change
1	Mask all outgoing edges from target node
2	Mask target CVE node
3	Restore Connection	Unmask last removed edge

Alert Generation

The defender observes a synthetic IDS alert stream produced by a four-layer generative model (utils/alert_generator.py). Each attacker traversal emits a cluster of correlated, stochastically-detected alerts for the landed node, plus a bundle of benign false positives concentrated on a small set of chronically noisy hosts:

1. Clustering — one exploit emits ceil(LogNormal(μ, σ)) correlated alerts, not a single token.
2. Timing — each exploit names its own temporal model (Hawkes / Poisson / periodic / burst), so port scans look bursty, beacons look periodic, and single-shot RCEs look like one clump.
3. Detection thinning — each alert survives with P_detect = σ(β₀ + β_av·av + β_ac·ac + β_auth·auth + β_stealth·stealth), drawn fresh per traversal so the same node alerts on some visits and is silent on others.
4. False-positive mixture — benign alerts are Zipf-concentrated on a fixed-per-episode noisy-host subset, reproducing the operational SOC regime (Gini ≈ 0.7).

This is what makes the z_score, entropy and alert-volume observation features carry real signal. The whole pipeline is configured from the optional alert_generator block in config.json; the false-positive rate is the first-class knob:

env.set_false_positive_rate(0.30)   # live; re-derives the noisy-host subset

"alert_generator": { "false_positive_rate": 0.30 }

See docs/ALERT_GENERATION.md for the full reference (every knob, the maths of each layer, archetypes, ablation toggles).

Gymnasium Wrappers

environment/gym_env.py provides two standard gymnasium.Env wrappers:

• NetworkAttackEnv — attacker perspective. Discrete(N) action space, one-hot observation of length N.
• NetworkDefenderEnv — defender perspective. MultiDiscrete([4, N]) action space, flat node-feature observation of shape (N × 17,).

Both expose a valid_action_mask() method for action masking with frameworks such as Stable Baselines 3 with MaskablePPO or RLlib.

Agent Architectures

DQN Attacker

A 4-layer MLP (128 hidden units, ReLU activations) with a one-hot state encoding over graph nodes. Trained with Double DQN (target network with soft updates, τ = 0.005), action masking, and ε-greedy exploration with exponential decay.

GAT Actor-Critic Defender

Input: node features x ∈ R^{N×17}, edge_index
  → GATConv(17, 64, heads=4) + BatchNorm → 256-dim
  → GATConv(256, 256, heads=1) + BatchNorm → 256-dim
  → global_mean_pool → R^256
  → shared_fc(256, 256)
  ↗ actor_type(512, 4)       — action type distribution
  ↗ actor_node_score(N, 1)   — node selection score (masked softmax)
  ↗ critic(256, 1)            — V(s) estimate

Training uses advantage-based policy gradient with GAE(λ=0.95):

Â_t = ∑_{k=0}^{∞} (γλ)^k δ_{t+k},  δ_t = r_t + γ·V(s_{t+1}) − V(s_t)
L(θ) = L^{ACTOR} + c₁·L^{CRITIC} − c₂·L^{ENT}

GAT DQN Defender

Same GAT encoder as the Actor-Critic, with a single Q-network head outputting one scalar per (action_type, target_node) pair. Trained with Double DQN, experience replay, and soft target-network updates (τ = 0.005).

Installation

git clone <repo-url>
cd rl-network-defense
pip install -r requirements.txt

PyTorch Geometric requires a matching CUDA build:

pip install torch-geometric
pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-$(python -c "import torch; print(torch.__version__)").html

Set your NVD API key in config.json:

"NVD_KEY": "YOUR_KEY_HERE"

Usage

Train the Attacker

python scripts/train_attacker.py

Train the Actor-Critic Defender

python scripts/train_defender.py --agent ac

Or with a custom goal node:

import json
from environment.graph_env import GraphEnvironment
from agents.defender.ac_defender import AC_Def_Agent

with open("attack_graphs/ag.json") as f:
    graph_data = json.load(f)

env = GraphEnvironment(graph_data, goal_node="42")  # protect node 42
agent = AC_Def_Agent()
agent.train_agent(env, num_episodes=5000,
                  output_path="policy-models/defender/ac_defender.pth",
                  attacker_model_path="policy-models/attacker/dqn_attacker.pth")

Train the DQN Defender

python scripts/train_defender.py --agent dqn

Zero-Shot Evaluation Across Attack Graph Variants

python scripts/evaluate_defenders.py \
    --variants attack_graphs/variants/ \
    --ac-model  policy-models/defender/ac_defender.pth \
    --dqn-model policy-models/defender/dqn_defender.pth

Attack Graph Variants

The attack_graphs/variants/ directory contains 1,000 structurally diverse attack graph variants used for zero-shot generalisation evaluation. Variants are generated by randomising the topology while preserving the core asset structure. The generation_log.json records the generation parameters for each variant. The goal node is auto-detected per variant as the highest-ID Access node.

Authoring Attack Graphs — Topology Builder

The environment can still be instantiated directly from any ready-made attack-graph JSON (attack_graphs/*.json) exactly as before. In addition, the self-contained graph_generator/ package lets you author your own attack graphs from a high-level network topology specification, using the ViolenceLang MAL language. The generated output is byte-compatible with what GraphEnvironment.initialize_from_json() consumes, so it is immediately trainable.

Pipeline

topology spec                violence_generator              RL-ready attack graph
{onlineHosts,         ──►   (ViolenceLang .mar +      ──►   {metadata, assets,
 expectedConnections}       vendored mal-toolbox)            associations, attackers}
                                                             → attack_graphs/<name>.json

The generator is a refactor of the original mal/ViolenceLang/ViolenceGenerator/violence_model_generator.py into an importable module with a CLI. The vendored mal-toolbox (v0.0.21) lives under graph_generator/vendor/, so no mal/ folder and no pip install of mal-toolbox are required — the package is self-contained.

Topology specification format

A topology is a JSON object listing the online hosts and the directed reachability between them:

{
  "onlineHosts": [
    {
      "id": 1,
      "hostname": "web",
      "ip": ["192.168.1.10"],
      "os": "Linux",
      "status": "online",
      "group": "DMZ",
      "vulnerabilities": [
        {"cve": "CVE-2021-44228", "cvss": "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:C/C:H/I:H/A:H"}
      ]
    }
  ],
  "expectedConnections": [
    {"source": 1, "destination": 2}
  ]
}

The legacy per-host layout (hosts with embedded connections) is also accepted and normalised automatically.

CVSS → MAL-asset mapping

For each vulnerability the generator instantiates MAL assets and associations from the CVSSv3 base vector, preserving the original ViolenceLang rules:

CVSS condition	Asset created	Association
AV:L	Local exploit	AvL
AV:N	Network exploit (+ Internet)	AvN, HasInternet
AV:A	Adjacent exploit	AvA
C:H (via Local/Network)	Privileges (Root)	ViHLN / ViHN
I:H/I:L (via Network)	Privileges (User)	ViHN
I:H (via Adjacent)	Privileges (Root)	VcHA
AC:H	Unsuccesfull → Chain	AcHl, Complex, multistage
A:H/A:L	Denial	VaHL
Root privileges gained	Host compromise	Standard
inter-host expectedConnections	Access	Reachability, DoReconOnReachableHost, CanExploit

The attacker entry point is set on the first asset with the Scan attack step, and the goal node is auto-detected by the environment as the highest-ID Access node (or pass goal_node explicitly).

Web app (recommended)

pip install flask python-jsonschema-objects   # one-time; mal-toolbox is vendored
./graph_generator/run.sh                       # → http://127.0.0.1:5000

The browser UI lets you create / load / edit hosts, build CVSS vectors via dropdowns, define connections, validate CVEs against the local database, and generate the attack graph with one click (written to attack_graphs/<name>.json).

Command line

# defaults output to attack_graphs/<topology_stem>.json
python -m graph_generator.violence_generator graph_generator/topologies/sl300_big.json

# explicit output path
python -m graph_generator.violence_generator my_topology.json -o attack_graphs/my.json

Python API

from graph_generator.violence_generator import generate_attack_graph
from graph_generator.rl_bridge import build_gym_env_from_topology, validate_cves_against_db

# topology → RL-ready dict (and optionally a saved file)
graph = generate_attack_graph("my_topology.json", output_path="attack_graphs/my.json")

# topology → live Gymnasium env in one step
env = build_gym_env_from_topology("my_topology.json", perspective="attacker",
                                  save_to="attack_graphs/my.json")

# pre-flight: which CVEs resolve in the remediation database?
report = validate_cves_against_db("my_topology.json")

Note on edge rewards. GraphEnvironment recomputes each edge's CVSS-weighted reward by looking the CVE up in database/vulnerability-remediation-database.db — it does not use the CVSS string from the topology directly. A CVE that is missing from the database, or present with an empty cvss_string, therefore contributes no severity weight to its edges. The web app's Validate CVEs button (and validate_cves_against_db) flags these cases before you generate.

Configuration

All hyperparameters are centralised in config.json. Key sections:

• environment — CVSS normalisation factors, NVD key, database path
• attacker — DQN training hyperparameters
• defender — GAT-PPO training hyperparameters
• replay_training — Alert-replay specific settings
• paths — Model and data file paths

License

See LICENSE.