File size: 6,345 Bytes
323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 c90b376 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 c90b376 323c6ec 4e12424 323c6ec 4e12424 323c6ec 4e12424 323c6ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | import numpy as np
import gymnasium as gym
from pettingzoo import ParallelEnv
from config import Config
class CyberPettingZooEnv(ParallelEnv):
metadata = {'render_modes': ['human'], "name": "cyber_network_defense_v2_gnn"}
def __init__(self):
super().__init__()
self.topology = Config.NETWORK_TOPOLOGY
self.total_nodes = Config.TOTAL_HOSTS()
self.features_per_node = 2
self.possible_agents = ["red", "blue"]
self.agents = self.possible_agents[:]
self.observation_spaces = {
agent: gym.spaces.Dict({
'node_features': gym.spaces.Box(low=0.0, high=1.0, shape=(self.total_nodes, self.features_per_node), dtype=np.float32),
'adjacency_matrix': gym.spaces.Box(low=0.0, high=1.0, shape=(self.total_nodes, self.total_nodes), dtype=np.float32)
}) for agent in self.possible_agents
}
self.action_spaces = {
'red': gym.spaces.Discrete(self.total_nodes * 4),
'blue': gym.spaces.Discrete(self.total_nodes * 2)
}
self.node_id_to_specs = {}
self._build_network_topology_map()
self._generate_static_adjacency_matrix()
def _build_network_topology_map(self):
node_id = 0
for subnet_name, specs in self.topology["subnets"].items():
for _ in range(specs["num_hosts"]):
self.node_id_to_specs[node_id] = {
"subnet": subnet_name,
"base_vulnerability": specs["base_vulnerability"]
}
node_id += 1
def _generate_static_adjacency_matrix(self):
"""Programmatically wires lateral communication pathways between subnets."""
A = np.zeros((self.total_nodes, self.total_nodes), dtype=np.float32)
# Define cross-subnet routing permissions
allowed_connections = {
"public_dmz": ["public_dmz", "dns_services"],
"dns_services": ["public_dmz", "dns_services", "corporate", "active_directory"],
"corporate": ["dns_services", "corporate", "active_directory"],
"active_directory": ["dns_services", "corporate", "active_directory", "secure_core"],
"secure_core": ["active_directory", "secure_core"]
}
# Map connections across individual node indexes
for i in range(self.total_nodes):
for j in range(self.total_nodes):
sub_i = self.node_id_to_specs[i]["subnet"]
sub_j = self.node_id_to_specs[j]["subnet"]
if sub_j in allowed_connections[sub_i]:
A[i, j] = 1.0
# Add Self-Loops (A + I)
A_tilde = A + np.eye(self.total_nodes, dtype=np.float32)
# Compute Symmetric Degree Normalization:
row_sum = np.sum(A_tilde, axis=1)
d_inv_sqrt = np.power(row_sum, -0.5, where=row_sum > 0)
d_inv_sqrt[row_sum == 0] = 0.0
D_inv_sqrt = np.diag(d_inv_sqrt)
self.normalized_adj = np.matmul(np.matmul(D_inv_sqrt, A_tilde), D_inv_sqrt).astype(np.float32)
def reset(self, seed=None, options=None):
self.agents = self.possible_agents[:]
self.state_matrix = np.zeros((self.total_nodes, self.features_per_node), dtype=np.float32)
obs_dict = {
'node_features': self.state_matrix.copy(),
'adjacency_matrix': self.normalized_adj.copy()
}
observations = {'red': obs_dict, 'blue': obs_dict}
infos = {'score_difference': 0.0, 'team_progress': 0.0}
return observations, infos
def _is_subnet_compromised(self, subnet_name):
for node_id, specs in self.node_id_to_specs.items():
if specs["subnet"] == subnet_name and self.state_matrix[node_id, 0] == 1.0:
return True
return False
def step(self, actions):
if not actions:
self.agents = []
return {}, {}, {}, {}, {}
red_action = actions.get('red', 0)
blue_action = actions.get('blue', 0)
red_target_node = red_action // 4
red_action_type = red_action % 4
blue_target_node = blue_action // 2
blue_action_type = blue_action % 2
red_target_node = min(max(0, red_target_node), self.total_nodes - 1)
blue_target_node = min(max(0, blue_target_node), self.total_nodes - 1)
red_target_specs = self.node_id_to_specs[red_target_node]
red_target_subnet = red_target_specs["subnet"]
# Defense Execution
if blue_action_type == 1:
self.state_matrix[blue_target_node, 1] = 1.0
# Attack Execution with Infrastructure Co-Dependencies
if red_action_type == 1:
if self.state_matrix[red_target_node, 1] == 0.0:
exploit_probability = red_target_specs["base_vulnerability"]
if red_target_subnet == "secure_core":
if not self._is_subnet_compromised("active_directory"):
exploit_probability = 0.0
elif red_target_subnet == "active_directory":
if self._is_subnet_compromised("dns_services"):
exploit_probability = min(1.0, exploit_probability * 2.0)
if np.random.rand() < exploit_probability:
self.state_matrix[red_target_node, 0] = 1.0
# Telemetry Preparation
total_compromised = np.sum(self.state_matrix[:, 0])
total_patched = np.sum(self.state_matrix[:, 1])
score_diff = float(total_compromised - total_patched)
progress = float(total_compromised / self.total_nodes) if self.total_nodes > 0 else 0.0
obs_dict = {
'node_features': self.state_matrix.copy(),
'adjacency_matrix': self.normalized_adj.copy()
}
observations = {'red': obs_dict, 'blue': obs_dict}
rewards = {'red': 0.0, 'blue': 0.0}
terminations = {'red': False, 'blue': False}
truncations = {'red': False, 'blue': False}
infos = {'score_difference': score_diff, 'team_progress': progress}
return observations, rewards, terminations, truncations, infos |