oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

901 - merged dev into my branch

Browse Source
This commit is contained in:
SunilSamra
2023-06-13 08:54:33 +01:00
parent ed8b53f5ef dc0349c37b
commit ce5d38a1cd
75 changed files with 3640 additions and 2306 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
src/primaite/environment/observations.py
View File

@@ -355,21 +355,50 @@ class AccessControlList(AbstractObservationComponent):
# 3. Initialise observation with zeroes
self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE)
# Dictionary to map services to numbers for obs space
self.services_dict = {}
def update(self):
"""Update the observation based on current environment state.
The structure of the observation space is described in :class:`.AccessControlList`
"""
obs = []
for acl_rule in self.env.acl:
for acl_rule in self.env.acl.acl:
permission = acl_rule.permission
source_ip = acl_rule.source_ip
dest_ip = acl_rule.dest_ip
protocol = acl_rule.protocol
port = acl_rule.port
position = self.env.acl.index(acl_rule)
position = self.env.acl.acl.index(acl_rule)
if permission == "DENY":
permission_int = 0
else:
permission_int = 1
obs.extend([permission, source_ip, dest_ip, protocol, port, position])
if source_ip == "ANY":
source_ip = 0
if dest_ip == "ANY":
dest_ip = 0
if port == "ANY":
port = 0
if protocol == "ANY":
protocol_int = 0
else:
while True:
if protocol in self.service_dict:
protocol_int = self.services_dict[protocol]
break
else:
self.services_dict[protocol] = len(self.services_dict) + 1
continue
# [0 - DENY, 1 - ALLOW] Permission
# [0 - ANY, x - IP Address/Protocol/Port]
print(permission_int, source_ip, dest_ip, protocol_int, port)
obs.extend(
[permission_int, source_ip, dest_ip, protocol_int, port, position]
)
self.current_observation[:] = obs

225
src/primaite/environment/primaite_env.py
View File

@@ -1,12 +1,11 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Main environment module containing the PRIMmary AI Training Evironment (Primaite) class."""
import copy
import csv
import logging
import os.path
from datetime import datetime
from typing import Dict, Tuple
from pathlib import Path
from typing import Dict, Tuple, Union
import networkx as nx
import numpy as np
@@ -24,11 +23,13 @@ from primaite.common.enums import (
NodePOLInitiator,
NodePOLType,
NodeType,
ObservationType,
Priority,
RulePermissionType,
SoftwareState,
)
from primaite.common.service import Service
from primaite.config import training_config
from primaite.config.training_config import TrainingConfig
from primaite.environment.observations import ObservationsHandler
from primaite.environment.reward import calculate_reward_function
from primaite.links.link import Link
@@ -56,29 +57,41 @@ class Primaite(Env):
ACTION_SPACE_ACL_ACTION_VALUES: int = 3
ACTION_SPACE_ACL_PERMISSION_VALUES: int = 2
def __init__(self, _config_values, _transaction_list):
def __init__(
self,
training_config_path: Union[str, Path],
lay_down_config_path: Union[str, Path],
transaction_list,
session_path: Path,
timestamp_str: str,
):
"""
Init.
The Primaite constructor.
Args:
_episode_steps: The number of steps for the episode
_config_filename: The name of config file
_transaction_list: The list of transactions to populate
_agent_identifier: Identifier for the agent
:param training_config_path: The training config filepath.
:param lay_down_config_path: The lay down config filepath.
:param transaction_list: The list of transactions to populate.
:param session_path: The directory path the session is writing to.
:param timestamp_str: The session timestamp in the format:
<yyyy-mm-dd>_<hh-mm-ss>.
"""
super(Primaite, self).__init__()
self._training_config_path = training_config_path
self._lay_down_config_path = lay_down_config_path
# Take a copy of the config values
self.config_values = _config_values
self.training_config: TrainingConfig = training_config.load(
training_config_path
)
# Number of steps in an episode
self.episode_steps = 0
self.episode_steps = self.training_config.num_steps
super(Primaite, self).__init__()
# Transaction list
self.transaction_list = _transaction_list
self.transaction_list = transaction_list
# The agent in use
self.agent_identifier = self.config_values.agent_identifier
self.agent_identifier = self.training_config.agent_identifier
# Create a dictionary to hold all the nodes
self.nodes: Dict[str, NodeUnion] = {}
@@ -122,6 +135,9 @@ class Primaite(Env):
# Create step count
self.step_count = 0
self.total_step_count: int = 0
"""The total number of time steps completed."""
# Create step info dictionary
self.step_info = {}
@@ -149,29 +165,21 @@ class Primaite(Env):
# The action type
self.action_type = 0
# TODO fix up with TrainingConfig
# stores the observation config from the yaml, default is NODE_LINK_TABLE
self.obs_config: dict = {"components": [{"name": "NODE_LINK_TABLE"}]}
if self.config_values.observation_config is not None:
self.obs_config = self.config_values.observation_config
if self.training_config.observation_space is not None:
self.obs_config = self.training_config.observation_space
# Observation Handler manages the user-configurable observation space.
# It will be initialised later.
self.obs_handler: ObservationsHandler
# Set by main_config
# Adds a DENY ALL or ALLOW ALL to the end of the Access Control List
self.acl_implicit_rule = RulePermissionType.DENY
# Sets a limit to how many ACL
self.max_acl_rules = 0
# Open the config file and build the environment laydown
try:
self.config_file = open(self.config_values.config_filename_use_case, "r")
self.config_data = yaml.safe_load(self.config_file)
self.load_config()
except Exception:
_LOGGER.error("Could not load the environment configuration")
_LOGGER.error("Exception occured", exc_info=True)
with open(self._lay_down_config_path, "r") as file:
# Open the config file and build the environment laydown
self.lay_down_config = yaml.safe_load(file)
self.load_lay_down_config()
# Store the node objects as node attributes
# (This is so we can access them as objects)
@@ -188,15 +196,10 @@ class Primaite(Env):
try:
plt.tight_layout()
nx.draw_networkx(self.network, with_labels=True)
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
datetime.now() # current date and time
path = "outputs/diagrams"
is_dir = os.path.isdir(path)
if not is_dir:
os.makedirs(path)
filename = "outputs/diagrams/network_" + time + ".png"
plt.savefig(filename, format="PNG")
file_path = session_path / f"network_{timestamp_str}.png"
plt.savefig(file_path, format="PNG")
plt.clf()
except Exception:
_LOGGER.error("Could not save network diagram")
@@ -207,7 +210,7 @@ class Primaite(Env):
self.observation_space, self.env_obs = self.init_observations()
# Define Action Space - depends on action space type (Node or ACL)
if self.action_type == ActionType.NODE:
if self.training_config.action_type == ActionType.NODE:
_LOGGER.info("Action space type NODE selected")
# Terms (for node action space):
# [0, num nodes] - node ID (0 = nothing, node ID)
@@ -216,7 +219,7 @@ class Primaite(Env):
# [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa
self.action_dict = self.create_node_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict))
elif self.action_type == ActionType.ACL:
elif self.training_config.action_type == ActionType.ACL:
_LOGGER.info("Action space type ACL selected")
# Terms (for ACL action space):
# [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule)
@@ -227,25 +230,21 @@ class Primaite(Env):
# [0, num ports] - Port (0 = any, then 1 -> x resolving to port)
self.action_dict = self.create_acl_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict))
elif self.action_type == ActionType.ANY:
elif self.training_config.action_type == ActionType.ANY:
_LOGGER.info("Action space type ANY selected - Node + ACL")
self.action_dict = self.create_node_and_acl_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict))
else:
_LOGGER.info("Invalid action type selected")
_LOGGER.info(
f"Invalid action type selected: {self.training_config.action_type}"
)
# Set up a csv to store the results of the training
try:
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
header = ["Episode", "Average Reward"]
# Check whether the output/rerults folder exists (doesn't exist by default install)
path = "outputs/results/"
is_dir = os.path.isdir(path)
if not is_dir:
os.makedirs(path)
filename = "outputs/results/average_reward_per_episode_" + time + ".csv"
self.csv_file = open(filename, "w", encoding="UTF8", newline="")
file_name = f"average_reward_per_episode_{timestamp_str}.csv"
file_path = session_path / file_name
self.csv_file = open(file_path, "w", encoding="UTF8", newline="")
self.csv_writer = csv.writer(self.csv_file)
self.csv_writer.writerow(header)
except Exception:
@@ -304,7 +303,8 @@ class Primaite(Env):
done = False
self.step_count += 1
# print("Episode step: " + str(self.stepCount))
self.total_step_count += 1
# print("Episode step: " + str(self.step_count))
# Need to clear traffic on all links first
for link_key, link_value in self.links.items():
@@ -376,13 +376,13 @@ class Primaite(Env):
self.green_iers,
self.red_iers,
self.step_count,
self.config_values,
self.training_config,
)
print(f" Step {self.step_count} Reward: {str(reward)}")
# print(f" Step {self.step_count} Reward: {str(reward)}")
self.total_reward += reward
if self.step_count == self.episode_steps:
self.average_reward = self.total_reward / self.step_count
if self.config_values.session_type == "EVALUATION":
if self.training_config.session_type == "EVALUATION":
# For evaluation, need to trigger the done value = True when
# step count is reached in order to prevent neverending episode
done = True
@@ -407,7 +407,6 @@ class Primaite(Env):
def __close__(self):
"""Override close function."""
self.csv_file.close()
self.config_file.close()
def init_acl(self):
"""Initialise the Access Control List."""
@@ -433,9 +432,10 @@ class Primaite(Env):
_action: The action space from the agent
"""
# At the moment, actions are only affecting nodes
if self.action_type == ActionType.NODE:
if self.training_config.action_type == ActionType.NODE:
self.apply_actions_to_nodes(_action)
elif self.action_type == ActionType.ACL:
elif self.training_config.action_type == ActionType.ACL:
self.apply_actions_to_acl(_action)
elif (
len(self.action_dict[_action]) == 6
@@ -673,44 +673,39 @@ class Primaite(Env):
self.obs_handler.update_obs()
self.env_obs = self.obs_handler.current_observation
def load_config(self):
def load_lay_down_config(self):
"""Loads config data in order to build the environment configuration."""
for item in self.config_data:
if item["itemType"] == "NODE":
for item in self.lay_down_config:
if item["item_type"] == "NODE":
# Create a node
self.create_node(item)
elif item["itemType"] == "LINK":
elif item["item_type"] == "LINK":
# Create a link
self.create_link(item)
elif item["itemType"] == "GREEN_IER":
elif item["item_type"] == "GREEN_IER":
# Create a Green IER
self.create_green_ier(item)
elif item["itemType"] == "GREEN_POL":
elif item["item_type"] == "GREEN_POL":
# Create a Green PoL
self.create_green_pol(item)
elif item["itemType"] == "RED_IER":
elif item["item_type"] == "RED_IER":
# Create a Red IER
self.create_red_ier(item)
elif item["itemType"] == "RED_POL":
elif item["item_type"] == "RED_POL":
# Create a Red PoL
self.create_red_pol(item)
elif item["itemType"] == "ACL_RULE":
elif item["item_type"] == "ACL_RULE":
# Create an ACL rule
self.create_acl_rule(item)
elif item["itemType"] == "SERVICES":
elif item["item_type"] == "SERVICES":
# Create the list of services
self.create_services_list(item)
elif item["itemType"] == "PORTS":
elif item["item_type"] == "PORTS":
# Create the list of ports
self.create_ports_list(item)
elif item["itemType"] == "ACTIONS":
# Get the action information
self.get_action_info(item)
elif item["itemType"] == "STEPS":
# Get the steps information
self.get_steps_info(item)
else:
# Do nothing (bad formatting)
item_type = item["item_type"]
_LOGGER.error(f"Invalid item_type: {item_type}")
pass
_LOGGER.info("Environment configuration loaded")
@@ -738,7 +733,7 @@ class Primaite(Env):
node_type,
node_priority,
node_hardware_state,
self.config_values,
self.training_config,
)
elif node_class == "ACTIVE":
# Active nodes have IP address, Software State and file system state
@@ -754,7 +749,7 @@ class Primaite(Env):
node_ip_address,
node_software_state,
node_file_system_state,
self.config_values,
self.training_config,
)
elif node_class == "SERVICE":
# Service nodes have IP address, Software State, file system state and list of services
@@ -770,7 +765,7 @@ class Primaite(Env):
node_ip_address,
node_software_state,
node_file_system_state,
self.config_values,
self.training_config,
)
node_services = item["services"]
for service in node_services:
@@ -849,14 +844,14 @@ class Primaite(Env):
item: A config data item
"""
ier_id = item["id"]
ier_start_step = item["startStep"]
ier_end_step = item["endStep"]
ier_start_step = item["start_step"]
ier_end_step = item["end_step"]
ier_load = item["load"]
ier_protocol = item["protocol"]
ier_port = item["port"]
ier_source = item["source"]
ier_destination = item["destination"]
ier_mission_criticality = item["missionCriticality"]
ier_mission_criticality = item["mission_criticality"]
# Create IER and add to green IER dictionary
self.green_iers[ier_id] = IER(
@@ -879,14 +874,14 @@ class Primaite(Env):
item: A config data item
"""
ier_id = item["id"]
ier_start_step = item["startStep"]
ier_end_step = item["endStep"]
ier_start_step = item["start_step"]
ier_end_step = item["end_step"]
ier_load = item["load"]
ier_protocol = item["protocol"]
ier_port = item["port"]
ier_source = item["source"]
ier_destination = item["destination"]
ier_mission_criticality = item["missionCriticality"]
ier_mission_criticality = item["mission_criticality"]
# Create IER and add to red IER dictionary
self.red_iers[ier_id] = IER(
@@ -909,8 +904,8 @@ class Primaite(Env):
item: A config data item
"""
pol_id = item["id"]
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_start_step = item["start_step"]
pol_end_step = item["end_step"]
pol_node = item["nodeId"]
pol_type = NodePOLType[item["type"]]
@@ -943,8 +938,8 @@ class Primaite(Env):
item: A config data item
"""
pol_id = item["id"]
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_start_step = item["start_step"]
pol_end_step = item["end_step"]
pol_target_node_id = item["targetNodeId"]
pol_initiator = NodePOLInitiator[item["initiator"]]
pol_type = NodePOLType[item["type"]]
@@ -1004,7 +999,7 @@ class Primaite(Env):
Args:
item: A config data item representing the services
"""
service_list = services["serviceList"]
service_list = services["service_list"]
for service in service_list:
service_name = service["name"]
@@ -1020,7 +1015,7 @@ class Primaite(Env):
Args:
item: A config data item representing the ports
"""
ports_list = ports["portsList"]
ports_list = ports["ports_list"]
for port in ports_list:
port_value = port["port"]
@@ -1029,6 +1024,14 @@ class Primaite(Env):
# Set the number of ports
self.num_ports = len(self.ports_list)
def get_observation_info(self, observation_info):
"""Extracts observation_info.
:param observation_info: Config item that defines which type of observation space to use
:type observation_info: str
"""
self.observation_type = ObservationType[observation_info["type"]]
def get_action_info(self, action_info):
"""
Extracts action_info.
@@ -1050,16 +1053,6 @@ class Primaite(Env):
"""
self.obs_config = obs_config
def get_steps_info(self, steps_info):
"""
Extracts steps_info.
Args:
item: A config data item representing steps info
"""
self.episode_steps = int(steps_info["steps"])
_LOGGER.info("Training episodes have " + str(self.episode_steps) + " steps")
def reset_environment(self):
"""
# Resets environment.
@@ -1067,11 +1060,11 @@ class Primaite(Env):
Uses config data config data in order to build the environment
configuration.
"""
for item in self.config_data:
if item["itemType"] == "NODE":
for item in self.lay_down_config:
if item["item_type"] == "NODE":
# Reset a node's state (normal and reference)
self.reset_node(item)
elif item["itemType"] == "ACL_RULE":
elif item["item_type"] == "ACL_RULE":
# Create an ACL rule (these are cleared on reset, so just need to recreate them)
self.create_acl_rule(item)
else:
@@ -1173,7 +1166,7 @@ class Primaite(Env):
def create_acl_action_dict(self):
"""Creates a dictionary mapping each possible discrete action to more readable multidiscrete action."""
# reserve 0 action to be a nothing action
actions = {0: [0, 0, 0, 0, 0, 0, 0]}
actions = {0: [0, 0, 0, 0, 0, 0]}
action_key = 1
# 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE
@@ -1185,16 +1178,14 @@ class Primaite(Env):
for dest_ip in range(self.num_nodes + 1):
for protocol in range(self.num_services + 1):
for port in range(self.num_ports + 1):
for position in range(self.max_acl_rules - 1):
action = [
action_decision,
action_permission,
source_ip,
dest_ip,
protocol,
port,
position,
]
action = [
action_decision,
action_permission,
source_ip,
dest_ip,
protocol,
port,
]
# Check to see if its an action we want to include as possible i.e. not a nothing action
if is_valid_acl_action_extra(action):
actions[action_key] = action