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

Ran pre-commit hook on all files and performed changes to fix flake8 failures

Browse Source
This commit is contained in:
Chris McCarthy
2023-05-25 11:42:19 +01:00
parent aa8284897a
commit 4eb1658966
42 changed files with 1371 additions and 1170 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/primaite/environment/__init__.py
View File

@@ -1,2 +1 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

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

@@ -1,39 +1,45 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Main environment module containing the PRIMmary AI Training Evironment (Primaite) class
"""
"""Main environment module containing the PRIMmary AI Training Evironment (Primaite) class."""
import numpy as np
import networkx as nx
import copy
import csv
import yaml
import os.path
import logging
from gym import Env, spaces
from matplotlib import pyplot as plt
import os.path
from datetime import datetime
from primaite.common.enums import *
import networkx as nx
import numpy as np
import yaml
from gym import Env, spaces
from matplotlib import pyplot as plt
from primaite.acl.access_control_list import AccessControlList
from primaite.common.enums import (
ACTION_TYPE,
FILE_SYSTEM_STATE,
HARDWARE_STATE,
NODE_POL_INITIATOR,
NODE_POL_TYPE,
PRIORITY,
SOFTWARE_STATE,
TYPE,
)
from primaite.common.service import Service
from primaite.environment.reward import calculate_reward_function
from primaite.links.link import Link
from primaite.pol.ier import IER
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen
from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
from primaite.pol.green_pol import apply_iers, apply_node_pol
from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.passive_node import PassiveNode
from primaite.nodes.service_node import ServiceNode
from primaite.common.service import Service
from primaite.acl.access_control_list import AccessControlList
from primaite.environment.reward import calculate_reward_function
from primaite.pol.green_pol import apply_iers, apply_node_pol
from primaite.pol.ier import IER
from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol
from primaite.transactions.transaction import Transaction
class Primaite(Env):
"""
PRIMmary AI Training Evironment (Primaite) class
"""
"""PRIMmary AI Training Evironment (Primaite) class."""
# Observation / Action Space contants
OBSERVATION_SPACE_FIXED_PARAMETERS = 4
@@ -42,11 +48,11 @@ class Primaite(Env):
ACTION_SPACE_ACL_ACTION_VALUES = 3
ACTION_SPACE_ACL_PERMISSION_VALUES = 2
OBSERVATION_SPACE_HIGH_VALUE = 1000000 # Highest value within an observation space
OBSERVATION_SPACE_HIGH_VALUE = 1000000 # Highest value within an observation space
def __init__(self, _config_values, _transaction_list):
"""
Init
Init.
Args:
_episode_steps: The number of steps for the episode
@@ -54,7 +60,6 @@ class Primaite(Env):
_transaction_list: The list of transactions to populate
_agent_identifier: Identifier for the agent
"""
super(Primaite, self).__init__()
# Take a copy of the config values
@@ -140,10 +145,12 @@ class Primaite(Env):
# Open the config file and build the environment laydown
try:
self.config_file = open("config/" + self.config_values.config_filename_use_case, "r")
self.config_file = open(
"config/" + self.config_values.config_filename_use_case, "r"
)
self.config_data = yaml.safe_load(self.config_file)
self.load_config()
except Exception as e:
except Exception:
logging.error("Could not load the environment configuration")
logging.error("Exception occured", exc_info=True)
@@ -162,17 +169,17 @@ class Primaite(Env):
try:
plt.tight_layout()
nx.draw_networkx(self.network, with_labels=True)
now = datetime.now() # current date and time
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
path = 'outputs/diagrams'
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")
plt.clf()
except Exception as a:
except Exception:
logging.error("Could not save network diagram")
logging.error("Exception occured", exc_info=True)
print("Could not save network diagram")
@@ -194,16 +201,22 @@ class Primaite(Env):
# - service F state | service F loading
# - service G state | service G loading
# Calculate the number of items that need to be included in the observation space
# Calculate the number of items that need to be included in the
# observation space
num_items = self.num_links + self.num_nodes
# Set the number of observation parameters, being # of services plus id, operating state, file system state and O/S state (i.e. 4)
self.num_observation_parameters = self.num_services + self.OBSERVATION_SPACE_FIXED_PARAMETERS
# Set the number of observation parameters, being # of services plus id,
# operating state, file system state and O/S state (i.e. 4)
self.num_observation_parameters = (
self.num_services + self.OBSERVATION_SPACE_FIXED_PARAMETERS
)
# Define the observation shape
self.observation_shape = (num_items, self.num_observation_parameters)
self.observation_space = spaces.Box(low=0,
high=self.config_values.observation_space_high_value,
shape=self.observation_shape,
dtype=np.int64)
self.observation_space = spaces.Box(
low=0,
high=self.config_values.observation_space_high_value,
shape=self.observation_shape,
dtype=np.int64,
)
# This is the observation that is sent back via the rest and step functions
self.env_obs = np.zeros(self.observation_shape, dtype=np.int64)
@@ -216,7 +229,14 @@ class Primaite(Env):
# [0, 4] - what property it's acting on (0 = nothing, state, o/s state, service state, file system state)
# [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore)
# [0, num services] - resolves to service ID (0 = nothing, resolves to service)
self.action_space = spaces.MultiDiscrete([self.num_nodes, self.ACTION_SPACE_NODE_PROPERTY_VALUES, self.ACTION_SPACE_NODE_ACTION_VALUES, self.num_services])
self.action_space = spaces.MultiDiscrete(
[
self.num_nodes,
self.ACTION_SPACE_NODE_PROPERTY_VALUES,
self.ACTION_SPACE_NODE_ACTION_VALUES,
self.num_services,
]
)
else:
logging.info("Action space type ACL selected")
# Terms (for ACL action space):
@@ -226,42 +246,52 @@ class Primaite(Env):
# [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses)
# [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol)
# [0, num ports] - Port (0 = any, then 1 -> x resolving to port)
self.action_space = spaces.MultiDiscrete([self.ACTION_SPACE_ACL_ACTION_VALUES, self.ACTION_SPACE_ACL_PERMISSION_VALUES, self.num_nodes + 1, self.num_nodes + 1, self.num_services + 1, self.num_ports + 1])
self.action_space = spaces.MultiDiscrete(
[
self.ACTION_SPACE_ACL_ACTION_VALUES,
self.ACTION_SPACE_ACL_PERMISSION_VALUES,
self.num_nodes + 1,
self.num_nodes + 1,
self.num_services + 1,
self.num_ports + 1,
]
)
# Set up a csv to store the results of the training
try:
now = datetime.now() # current date and time
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
header = ['Episode', 'Average Reward']
header = ["Episode", "Average Reward"]
# Check whether the output/rerults folder exists (doesn't exist by default install)
path = 'outputs/results/'
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='')
self.csv_file = open(filename, "w", encoding="UTF8", newline="")
self.csv_writer = csv.writer(self.csv_file)
self.csv_writer.writerow(header)
except Exception as e:
logging.error("Could not create csv file to hold average reward per episode")
except Exception:
logging.error(
"Could not create csv file to hold average reward per episode"
)
logging.error("Exception occured", exc_info=True)
def reset(self):
"""
AI Gym Reset function
AI Gym Reset function.
Returns:
Environment observation space (reset)
"""
csv_data = self.episode_count, self.average_reward
self.csv_writer.writerow(csv_data)
self.episode_count += 1
# Don't need to reset links, as they are cleared and recalculated every step
# Clear the ACL
self.init_acl()
@@ -280,7 +310,7 @@ class Primaite(Env):
def step(self, action):
"""
AI Gym Step function
AI Gym Step function.
Args:
action: Action space from agent
@@ -291,7 +321,6 @@ class Primaite(Env):
done: Indicates episode is complete if True
step_info: Additional information relating to this step
"""
if self.step_count == 0:
print("Episode: " + str(self.episode_count) + " running")
@@ -299,14 +328,16 @@ class Primaite(Env):
done = False
self.step_count += 1
#print("Episode step: " + str(self.stepCount))
# print("Episode step: " + str(self.stepCount))
# Need to clear traffic on all links first
for link_key, link_value in self.links.items():
link_value.clear_traffic()
# Create a Transaction (metric) object for this step
transaction = Transaction(datetime.now(), self.agent_identifier, self.episode_count, self.step_count)
transaction = Transaction(
datetime.now(), self.agent_identifier, self.episode_count, self.step_count
)
# Load the initial observation space into the transaction
transaction.set_obs_space_pre(copy.deepcopy(self.env_obs))
# Load the action space into the transaction
@@ -316,50 +347,97 @@ class Primaite(Env):
self.apply_time_based_updates()
# 2. Apply PoL
apply_node_pol(self.nodes, self.node_pol, self.step_count) # Node PoL
apply_iers(self.network, self.nodes, self.links, self.green_iers, self.acl, self.step_count) # Network PoL
apply_node_pol(self.nodes, self.node_pol, self.step_count) # Node PoL
apply_iers(
self.network,
self.nodes,
self.links,
self.green_iers,
self.acl,
self.step_count,
) # Network PoL
# Take snapshots of nodes and links
self.nodes_post_pol = copy.deepcopy(self.nodes)
self.links_post_pol = copy.deepcopy(self.links)
self.links_post_pol = copy.deepcopy(self.links)
# Reference
apply_node_pol(self.nodes_reference, self.node_pol, self.step_count) # Node PoL
apply_iers(self.network_reference, self.nodes_reference, self.links_reference, self.green_iers, self.acl, self.step_count) # Network PoL
apply_node_pol(self.nodes_reference, self.node_pol, self.step_count) # Node PoL
apply_iers(
self.network_reference,
self.nodes_reference,
self.links_reference,
self.green_iers,
self.acl,
self.step_count,
) # Network PoL
# 3. Implement Red Action
apply_red_agent_iers(self.network, self.nodes, self.links, self.red_iers, self.acl, self.step_count)
apply_red_agent_node_pol(self.nodes, self.red_iers, self.red_node_pol, self.step_count)
# 3. Implement Red Action
apply_red_agent_iers(
self.network,
self.nodes,
self.links,
self.red_iers,
self.acl,
self.step_count,
)
apply_red_agent_node_pol(
self.nodes, self.red_iers, self.red_node_pol, self.step_count
)
# Take snapshots of nodes and links
self.nodes_post_red = copy.deepcopy(self.nodes)
self.links_post_red = copy.deepcopy(self.links)
self.links_post_red = copy.deepcopy(self.links)
# 4. Implement Blue Action
self.interpret_action_and_apply(action)
# 5. Reapply normal and Red agent IER PoL, as we need to see what effect the blue agent action has had (if any) on link status
# 5. Reapply normal and Red agent IER PoL, as we need to see what
# effect the blue agent action has had (if any) on link status
# Need to clear traffic on all links first
for link_key, link_value in self.links.items():
link_value.clear_traffic()
apply_iers(self.network, self.nodes, self.links, self.green_iers, self.acl, self.step_count)
apply_red_agent_iers(self.network, self.nodes, self.links, self.red_iers, self.acl, self.step_count)
apply_iers(
self.network,
self.nodes,
self.links,
self.green_iers,
self.acl,
self.step_count,
)
apply_red_agent_iers(
self.network,
self.nodes,
self.links,
self.red_iers,
self.acl,
self.step_count,
)
# Take snapshots of nodes and links
self.nodes_post_blue = copy.deepcopy(self.nodes)
self.links_post_blue = copy.deepcopy(self.links)
self.links_post_blue = copy.deepcopy(self.links)
# 6. Calculate reward signal (for RL)
reward = calculate_reward_function(self.nodes_post_pol, self.nodes_post_blue, self.nodes_reference, self.green_iers, self.red_iers, self.step_count, self.config_values)
#print("Step reward: " + str(reward))
reward = calculate_reward_function(
self.nodes_post_pol,
self.nodes_post_blue,
self.nodes_reference,
self.green_iers,
self.red_iers,
self.step_count,
self.config_values,
)
# print("Step 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":
# For evaluation, need to trigger the done value = True when step count is reached in order to prevent neverending episode
# For evaluation, need to trigger the done value = True when
# step count is reached in order to prevent neverending episode
done = True
print("Average reward: " + str(self.average_reward))
print("Average reward: " + str(self.average_reward))
# Load the reward into the transaction
transaction.set_reward(reward)
# 7. Output Verbose
#self.output_link_status()
# self.output_link_status()
# 8. Update env_obs
self.update_environent_obs()
@@ -373,38 +451,33 @@ class Primaite(Env):
return self.env_obs, reward, done, self.step_info
def __close__(self):
"""
Override close function
"""
"""Override close function."""
self.csv_file.close()
self.config_file.close()
def init_acl(self):
"""
Initialise the Access Control List
"""
"""Initialise the Access Control List."""
self.acl.remove_all_rules()
self.acl.remove_all_rules()
def output_link_status(self):
"""
Output the link status of all links to the console
"""
"""Output the link status of all links to the console."""
for link_key, link_value in self.links.items():
print("Link ID: " + link_value.get_id())
for protocol in link_value.get_protocol_list():
print(" Protocol: " + protocol.get_name().name + ", Load: " + str(protocol.get_load()))
print(
" Protocol: "
+ protocol.get_name().name
+ ", Load: "
+ str(protocol.get_load())
)
def interpret_action_and_apply(self, _action):
"""
Applies agent actions to the nodes and Access Control List
Applies agent actions to the nodes and Access Control List.
Args:
_action: The action space from the agent
"""
# At the moment, actions are only affecting nodes
if self.action_type == ACTION_TYPE.NODE:
self.apply_actions_to_nodes(_action)
@@ -413,12 +486,11 @@ class Primaite(Env):
def apply_actions_to_nodes(self, _action):
"""
Applies agent actions to the nodes
Applies agent actions to the nodes.
Args:
_action: The action space from the agent
"""
node_id = _action[0]
node_property = _action[1]
property_action = _action[2]
@@ -427,7 +499,7 @@ class Primaite(Env):
# Check that the action is requesting a valid node
try:
node = self.nodes[str(node_id)]
except:
except Exception:
return
if node_property == 0:
@@ -472,7 +544,9 @@ class Primaite(Env):
return
elif property_action == 1:
# Patch (valid action if it's good or compromised)
node.set_service_state(self.services_list[service_index], SOFTWARE_STATE.PATCHING)
node.set_service_state(
self.services_list[service_index], SOFTWARE_STATE.PATCHING
)
else:
# Node is not of Service Type
return
@@ -488,7 +562,10 @@ class Primaite(Env):
elif property_action == 2:
# Repair
# You cannot repair a destroyed file system - it needs restoring
if node.get_file_system_state_actual() != FILE_SYSTEM_STATE.DESTROYED:
if (
node.get_file_system_state_actual()
!= FILE_SYSTEM_STATE.DESTROYED
):
node.set_file_system_state(FILE_SYSTEM_STATE.REPAIRING)
elif property_action == 3:
# Restore
@@ -501,12 +578,11 @@ class Primaite(Env):
def apply_actions_to_acl(self, _action):
"""
Applies agent actions to the Access Control List [TO DO]
Applies agent actions to the Access Control List [TO DO].
Args:
_action: The action space from the agent
"""
action_decision = _action[0]
action_permission = _action[1]
action_source_ip = _action[2]
@@ -517,7 +593,7 @@ class Primaite(Env):
if action_decision == 0:
# It's decided to do nothing
return
else:
else:
# It's decided to create a new ACL rule or remove an existing rule
# Permission value
if action_permission == 0:
@@ -556,18 +632,31 @@ class Primaite(Env):
# Now add or remove
if action_decision == 1:
# Add the rule
self.acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
self.acl.add_rule(
acl_rule_permission,
acl_rule_source,
acl_rule_destination,
acl_rule_protocol,
acl_rule_port,
)
elif action_decision == 2:
# Remove the rule
self.acl.remove_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
self.acl.remove_rule(
acl_rule_permission,
acl_rule_source,
acl_rule_destination,
acl_rule_protocol,
acl_rule_port,
)
else:
return
def apply_time_based_updates(self):
"""
Updates anything that needs to count down and then change state (e.g. reset / patching status)
Updates anything that needs to count down and then change state.
e.g. reset / patching status
"""
for node_key, node in self.nodes.items():
if node.get_state() == HARDWARE_STATE.RESETTING:
node.update_resetting_status()
@@ -605,10 +694,7 @@ class Primaite(Env):
pass
def update_environent_obs(self):
"""
# Updates the observation space based on the node and link status
"""
"""Updates the observation space based on the node and link status."""
item_index = 0
# Do nodes first
@@ -617,15 +703,19 @@ class Primaite(Env):
self.env_obs[item_index][1] = node.get_state().value
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
self.env_obs[item_index][2] = node.get_os_state().value
self.env_obs[item_index][3] = node.get_file_system_state_observed().value
self.env_obs[item_index][
3
] = node.get_file_system_state_observed().value
else:
self.env_obs[item_index][2] = 0
self.env_obs[item_index][3] = 0
self.env_obs[item_index][3] = 0
service_index = 4
if isinstance(node, ServiceNode):
if isinstance(node, ServiceNode):
for service in self.services_list:
if node.has_service(service):
self.env_obs[item_index][service_index] = node.get_service_state(service).value
self.env_obs[item_index][
service_index
] = node.get_service_state(service).value
else:
self.env_obs[item_index][service_index] = 0
service_index += 1
@@ -650,17 +740,14 @@ class Primaite(Env):
item_index += 1
def load_config(self):
"""
# Loads config data in order to build the environment configuration
"""
"""Loads config data in order to build the environment configuration."""
for item in self.config_data:
if item["itemType"] == "NODE":
# Create a node
self.create_node(item)
elif item["itemType"] == "LINK":
# Create a link
self.create_link(item)
self.create_link(item)
elif item["itemType"] == "GREEN_IER":
# Create a Green IER
self.create_green_ier(item)
@@ -697,12 +784,11 @@ class Primaite(Env):
def create_node(self, item):
"""
Creates a node from config data
Creates a node from config data.
Args:
item: A config data item
"""
# All nodes have these parameters
node_id = item["id"]
node_name = item["name"]
@@ -712,19 +798,46 @@ class Primaite(Env):
node_hardware_state = HARDWARE_STATE[item["hardwareState"]]
if node_base_type == "PASSIVE":
node = PassiveNode(node_id, node_name, node_type, node_priority, node_hardware_state, self.config_values)
node = PassiveNode(
node_id,
node_name,
node_type,
node_priority,
node_hardware_state,
self.config_values,
)
elif node_base_type == "ACTIVE":
# Active nodes have IP address, operating system state and file system state
node_ip_address = item["ipAddress"]
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]]
node = ActiveNode(node_id, node_name, node_type, node_priority, node_hardware_state, node_ip_address, node_software_state, node_file_system_state, self.config_values)
node = ActiveNode(
node_id,
node_name,
node_type,
node_priority,
node_hardware_state,
node_ip_address,
node_software_state,
node_file_system_state,
self.config_values,
)
elif node_base_type == "SERVICE":
# Service nodes have IP address, operating system state, file system state and list of services
node_ip_address = item["ipAddress"]
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]]
node = ServiceNode(node_id, node_name, node_type, node_priority, node_hardware_state, node_ip_address, node_software_state, node_file_system_state, self.config_values)
node = ServiceNode(
node_id,
node_name,
node_type,
node_priority,
node_hardware_state,
node_ip_address,
node_software_state,
node_file_system_state,
self.config_values,
)
node_services = item["services"]
for service in node_services:
service_protocol = service["name"]
@@ -752,12 +865,11 @@ class Primaite(Env):
def create_link(self, item):
"""
Creates a link from config data
Creates a link from config data.
Args:
item: A config data item
"""
link_id = item["id"]
link_name = item["name"]
link_bandwidth = item["bandwidth"]
@@ -771,7 +883,13 @@ class Primaite(Env):
self.network.add_edge(source_node, dest_node, id=link_name)
# Add link to link dictionary
self.links[link_name] = Link(link_id, link_bandwidth, source_node.get_name(), dest_node.get_name(), self.services_list)
self.links[link_name] = Link(
link_id,
link_bandwidth,
source_node.get_name(),
dest_node.get_name(),
self.services_list,
)
# Reference
source_node_ref = self.nodes_reference[link_source]
@@ -781,16 +899,21 @@ class Primaite(Env):
self.network_reference.add_edge(source_node_ref, dest_node_ref, id=link_name)
# Add link to link dictionary (reference)
self.links_reference[link_name] = Link(link_id, link_bandwidth, source_node_ref.get_name(), dest_node_ref.get_name(), self.services_list)
self.links_reference[link_name] = Link(
link_id,
link_bandwidth,
source_node_ref.get_name(),
dest_node_ref.get_name(),
self.services_list,
)
def create_green_ier(self, item):
"""
Creates a green IER from config data
Creates a green IER from config data.
Args:
item: A config data item
"""
ier_id = item["id"]
ier_start_step = item["startStep"]
ier_end_step = item["endStep"]
@@ -802,16 +925,25 @@ class Primaite(Env):
ier_mission_criticality = item["missionCriticality"]
# Create IER and add to green IER dictionary
self.green_iers[ier_id] = IER(ier_id, ier_start_step, ier_end_step, ier_load, ier_protocol, ier_port, ier_source, ier_destination, ier_mission_criticality)
self.green_iers[ier_id] = IER(
ier_id,
ier_start_step,
ier_end_step,
ier_load,
ier_protocol,
ier_port,
ier_source,
ier_destination,
ier_mission_criticality,
)
def create_red_ier(self, item):
"""
Creates a red IER from config data
Creates a red IER from config data.
Args:
item: A config data item
"""
ier_id = item["id"]
ier_start_step = item["startStep"]
ier_end_step = item["endStep"]
@@ -823,21 +955,30 @@ class Primaite(Env):
ier_mission_criticality = item["missionCriticality"]
# Create IER and add to red IER dictionary
self.red_iers[ier_id] = IER(ier_id, ier_start_step, ier_end_step, ier_load, ier_protocol, ier_port, ier_source, ier_destination, ier_mission_criticality)
self.red_iers[ier_id] = IER(
ier_id,
ier_start_step,
ier_end_step,
ier_load,
ier_protocol,
ier_port,
ier_source,
ier_destination,
ier_mission_criticality,
)
def create_green_pol(self, item):
"""
Creates a green PoL object from config data
Creates a green PoL object from config data.
Args:
item: A config data item
"""
pol_id = item["id"]
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_node = item["nodeId"]
pol_type = NODE_POL_TYPE[item["type"]]
pol_type = NODE_POL_TYPE[item["type"]]
# State depends on whether this is Operating, O/S, file system or Service PoL type
if pol_type == NODE_POL_TYPE.OPERATING:
@@ -850,16 +991,23 @@ class Primaite(Env):
pol_protocol = item["protocol"]
pol_state = SOFTWARE_STATE[item["state"]]
self.node_pol[pol_id] = NodeStateInstructionGreen(pol_id, pol_start_step, pol_end_step, pol_node, pol_type, pol_protocol, pol_state)
self.node_pol[pol_id] = NodeStateInstructionGreen(
pol_id,
pol_start_step,
pol_end_step,
pol_node,
pol_type,
pol_protocol,
pol_state,
)
def create_red_pol(self, item):
"""
Creates a red PoL object from config data
Creates a red PoL object from config data.
Args:
item: A config data item
"""
pol_id = item["id"]
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
@@ -880,32 +1028,48 @@ class Primaite(Env):
pol_source_node_service = item["sourceNodeService"]
pol_source_node_service_state = item["sourceNodeServiceState"]
self.red_node_pol[pol_id] = NodeStateInstructionRed(pol_id, pol_start_step, pol_end_step, pol_target_node_id, pol_initiator, pol_type, pol_protocol, pol_state, pol_source_node_id, pol_source_node_service, pol_source_node_service_state)
self.red_node_pol[pol_id] = NodeStateInstructionRed(
pol_id,
pol_start_step,
pol_end_step,
pol_target_node_id,
pol_initiator,
pol_type,
pol_protocol,
pol_state,
pol_source_node_id,
pol_source_node_service,
pol_source_node_service_state,
)
def create_acl_rule(self, item):
"""
Creates an ACL rule from config data
Creates an ACL rule from config data.
Args:
item: A config data item
"""
acl_rule_permission = item["permission"]
acl_rule_source = item["source"]
acl_rule_destination = item["destination"]
acl_rule_protocol = item["protocol"]
acl_rule_port = item["port"]
self.acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
self.acl.add_rule(
acl_rule_permission,
acl_rule_source,
acl_rule_destination,
acl_rule_protocol,
acl_rule_port,
)
def create_services_list(self, services):
"""
Creates a list of services (enum) from config data
Creates a list of services (enum) from config data.
Args:
item: A config data item representing the services
"""
service_list = services["serviceList"]
for service in service_list:
@@ -917,12 +1081,11 @@ class Primaite(Env):
def create_ports_list(self, ports):
"""
Creates a list of ports from config data
Creates a list of ports from config data.
Args:
item: A config data item representing the ports
"""
ports_list = ports["portsList"]
for port in ports_list:
@@ -934,35 +1097,34 @@ class Primaite(Env):
def get_action_info(self, action_info):
"""
Extracts action_info
Extracts action_info.
Args:
item: A config data item representing action info
"""
self.action_type = ACTION_TYPE[action_info["type"]]
def get_steps_info(self, steps_info):
"""
Extracts steps_info
Extracts steps_info.
Args:
item: A config data item representing steps info
"""
self.episode_steps = int(steps_info["steps"])
logging.info("Training episodes have " + str(self.episode_steps) + " steps")
def reset_environment(self):
"""
# Resets environment using config data config data in order to build the environment configuration
"""
# Resets environment.
Uses config data config data in order to build the environment
configuration.
"""
for item in self.config_data:
if item["itemType"] == "NODE":
# Reset a node's state (normal and reference)
self.reset_node(item)
self.reset_node(item)
elif item["itemType"] == "ACL_RULE":
# Create an ACL rule (these are cleared on reset, so just need to recreate them)
self.create_acl_rule(item)
@@ -970,7 +1132,6 @@ class Primaite(Env):
# Do nothing (bad formatting or not relevant to reset)
pass
# Reset the IER status so they are not running initially
# Green IERs
for ier_key, ier_value in self.green_iers.items():
@@ -981,12 +1142,11 @@ class Primaite(Env):
def reset_node(self, item):
"""
Resets the statuses of a node
Resets the statuses of a node.
Args:
item: A config data item
"""
# All nodes have these parameters
node_id = item["id"]
node_base_type = item["baseType"]
@@ -1027,10 +1187,3 @@ class Primaite(Env):
else:
# Bad formatting
pass

129
src/primaite/environment/reward.py
View File

@@ -1,15 +1,21 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Implements reward function
"""
from primaite.common.enums import *
"""Implements reward function."""
from primaite.common.enums import FILE_SYSTEM_STATE, HARDWARE_STATE, SOFTWARE_STATE
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.service_node import ServiceNode
def calculate_reward_function(initial_nodes, final_nodes, reference_nodes, green_iers, red_iers, step_count, config_values):
def calculate_reward_function(
initial_nodes,
final_nodes,
reference_nodes,
green_iers,
red_iers,
step_count,
config_values,
):
"""
Compares the states of the initial and final nodes/links to get a reward
Compares the states of the initial and final nodes/links to get a reward.
Args:
initial_nodes: The nodes before red and blue agents take effect
@@ -20,29 +26,36 @@ def calculate_reward_function(initial_nodes, final_nodes, reference_nodes, green
step_count: current step
config_values: Config values
"""
reward_value = 0
# For each node, compare operating state, o/s operating state, service states
for node_key, final_node in final_nodes.items():
initial_node = initial_nodes[node_key]
reference_node = reference_nodes[node_key]
# Operating State
reward_value += score_node_operating_state(final_node, initial_node, reference_node, config_values)
reward_value += score_node_operating_state(
final_node, initial_node, reference_node, config_values
)
# Operating System State
if (isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode)):
reward_value += score_node_os_state(final_node, initial_node, reference_node, config_values)
if isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode):
reward_value += score_node_os_state(
final_node, initial_node, reference_node, config_values
)
# Service State
if (isinstance(final_node, ServiceNode)):
reward_value += score_node_service_state(final_node, initial_node, reference_node, config_values)
if isinstance(final_node, ServiceNode):
reward_value += score_node_service_state(
final_node, initial_node, reference_node, config_values
)
# File System State
if isinstance(final_node, ActiveNode):
reward_value += score_node_file_system(final_node, initial_node, reference_node, config_values)
reward_value += score_node_file_system(
final_node, initial_node, reference_node, config_values
)
# Go through each red IER - penalise if it is running
for ier_key, ier_value in red_iers.items():
start_step = ier_value.get_start_step()
@@ -57,14 +70,17 @@ def calculate_reward_function(initial_nodes, final_nodes, reference_nodes, green
stop_step = ier_value.get_end_step()
if step_count >= start_step and step_count <= stop_step:
if not ier_value.get_is_running():
reward_value += config_values.green_ier_blocked * ier_value.get_mission_criticality()
reward_value += (
config_values.green_ier_blocked
* ier_value.get_mission_criticality()
)
return reward_value
def score_node_operating_state(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the operating state of a node
Calculates score relating to the operating state of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -72,8 +88,7 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
reference_node: The node if there had been no red or blue effect
config_values: Config values
"""
score = 0
score = 0
final_node_operating_state = final_node.get_state()
initial_node_operating_state = initial_node.get_state()
reference_node_operating_state = reference_node.get_state()
@@ -81,7 +96,7 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
if final_node_operating_state == reference_node_operating_state:
# All is well - we're no different from the reference situation
score += config_values.all_ok
else:
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
if initial_node_operating_state == HARDWARE_STATE.ON:
@@ -95,7 +110,7 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
if final_node_operating_state == HARDWARE_STATE.ON:
score += config_values.on_should_be_off
elif final_node_operating_state == HARDWARE_STATE.RESETTING:
score += config_values.resetting_should_be_off
score += config_values.resetting_should_be_off
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.RESETTING:
@@ -112,9 +127,10 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
return score
def score_node_os_state(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the operating system state of a node
Calculates score relating to the operating system state of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -122,8 +138,7 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
reference_node: The node if there had been no red or blue effect
config_values: Config values
"""
score = 0
score = 0
final_node_os_state = final_node.get_os_state()
initial_node_os_state = initial_node.get_os_state()
reference_node_os_state = reference_node.get_os_state()
@@ -131,7 +146,7 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
if final_node_os_state == reference_node_os_state:
# All is well - we're no different from the reference situation
score += config_values.all_ok
else:
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
if initial_node_os_state == SOFTWARE_STATE.GOOD:
@@ -145,18 +160,18 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
if final_node_os_state == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_patching
elif final_node_os_state == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_patching
score += config_values.compromised_should_be_patching
elif final_node_os_state == SOFTWARE_STATE.PATCHING:
score += config_values.patching
score += config_values.patching
else:
pass
pass
elif initial_node_os_state == SOFTWARE_STATE.COMPROMISED:
if final_node_os_state == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_compromised
elif final_node_os_state == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_compromised
elif final_node_os_state == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised
score += config_values.compromised
else:
pass
else:
@@ -164,9 +179,10 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
return score
def score_node_service_state(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the service state(s) of a node
Calculates score relating to the service state(s) of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -174,12 +190,11 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
reference_node: The node if there had been no red or blue effect
config_values: Config values
"""
score = 0
score = 0
final_node_services = final_node.get_services()
initial_node_services = initial_node.get_services()
reference_node_services = reference_node.get_services()
for service_key, final_service in final_node_services.items():
reference_service = reference_node_services[service_key]
initial_service = initial_node_services[service_key]
@@ -203,11 +218,11 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
if final_service.get_state() == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_patching
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_patching
score += config_values.compromised_should_be_patching
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_patching
score += config_values.overwhelmed_should_be_patching
elif final_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching
score += config_values.patching
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.COMPROMISED:
@@ -216,9 +231,9 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
elif final_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_compromised
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised
score += config_values.compromised
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_compromised
score += config_values.overwhelmed_should_be_compromised
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
@@ -227,9 +242,9 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
elif final_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_overwhelmed
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_overwhelmed
score += config_values.compromised_should_be_overwhelmed
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed
score += config_values.overwhelmed
else:
pass
else:
@@ -237,17 +252,17 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
return score
def score_node_file_system(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the file system state of a node
Calculates score relating to the file system state of a node.
Args:
final_node: The node after red and blue agents take effect
initial_node: The node before red and blue agents take effect
reference_node: The node if there had been no red or blue effect
"""
score = 0
score = 0
final_node_file_system_state = final_node.get_file_system_state_actual()
initial_node_file_system_state = initial_node.get_file_system_state_actual()
reference_node_file_system_state = reference_node.get_file_system_state_actual()
@@ -259,7 +274,7 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
if final_node_file_system_state == reference_node_file_system_state:
# All is well - we're no different from the reference situation
score += config_values.all_ok
else:
else:
# We're different from the reference situation
# Need to compare initial and final state of node (i.e. after red and blue actions)
if initial_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
@@ -277,15 +292,15 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
if final_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
score += config_values.good_should_be_repairing
elif final_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
score += config_values.restoring_should_be_repairing
score += config_values.restoring_should_be_repairing
elif final_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
score += config_values.corrupt_should_be_repairing
score += config_values.corrupt_should_be_repairing
elif final_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
score += config_values.destroyed_should_be_repairing
score += config_values.destroyed_should_be_repairing
elif final_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
score += config_values.repairing
else:
pass
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
if final_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
score += config_values.good_should_be_restoring
@@ -294,9 +309,9 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
elif final_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
score += config_values.corrupt_should_be_restoring
elif final_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
score += config_values.destroyed_should_be_restoring
score += config_values.destroyed_should_be_restoring
elif final_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
score += config_values.restoring
score += config_values.restoring
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
@@ -307,9 +322,9 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
elif final_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
score += config_values.restoring_should_be_corrupt
elif final_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
score += config_values.destroyed_should_be_corrupt
score += config_values.destroyed_should_be_corrupt
elif final_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
score += config_values.corrupt
score += config_values.corrupt
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
@@ -320,9 +335,9 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
elif final_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
score += config_values.restoring_should_be_destroyed
elif final_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
score += config_values.corrupt_should_be_destroyed
score += config_values.corrupt_should_be_destroyed
elif final_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
score += config_values.destroyed
score += config_values.destroyed
else:
pass
else:
@@ -332,9 +347,9 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
if final_node_scanning_state == reference_node_scanning_state:
# All is well - we're no different from the reference situation
score += config_values.all_ok
else:
else:
# We're different from the reference situation
# We're scanning the file system which incurs a penalty (as it slows down systems)
score += config_values.scanning
return score
return score