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

Initial commit of v1.0.0. Updated the .gitignore for the standard Python gitignore. Added Azure DevOps release pipeline for proper artifact release from the start.

Browse Source
This commit is contained in:
Chris McCarthy
2023-03-28 17:33:34 +01:00
parent fdbd5903bc
commit 7800f1f66e
45 changed files with 5086 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
PRIMAITE/environment/__init__.py Normal file
View File

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

989
PRIMAITE/environment/primaite.py Normal file
View File

@@ -0,0 +1,989 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
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
from datetime import datetime
from common.enums import *
from links.link import Link
from pol.ier import IER
from nodes.node_state_instruction import NodeStateInstruction
from pol.green_pol import apply_iers, apply_node_pol
from pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol
from nodes.active_node import ActiveNode
from nodes.passive_node import PassiveNode
from nodes.service_node import ServiceNode
from common.service import Service
from acl.access_control_list import AccessControlList
from environment.reward import calculate_reward_function
from transactions.transaction import Transaction
class PRIMAITE(Env):
"""
PRIMmary AI Training Evironment (PRIMAITE) class
"""
# Observation / Action Space contants
OBSERVATION_SPACE_FIXED_PARAMETERS = 3
ACTION_SPACE_NODE_PROPERTY_VALUES = 4
ACTION_SPACE_NODE_ACTION_VALUES = 4
ACTION_SPACE_ACL_ACTION_VALUES = 3
ACTION_SPACE_ACL_PERMISSION_VALUES = 2
OBSERVATION_SPACE_HIGH_VALUE = 1000000 # Highest value within an observation space
def __init__(self, _config_values, _transaction_list):
"""
Init
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
"""
super(PRIMAITE, self).__init__()
# Take a copy of the config values
self.config_values = _config_values
# Number of steps in an episode
self.episode_steps = 0
# Transaction list
self.transaction_list = _transaction_list
# The agent in use
self.agent_identifier = self.config_values.agent_identifier
# Create a dictionary to hold all the nodes
self.nodes = {}
# Create a dictionary to hold a reference set of nodes
self.nodes_reference = {}
# Create a dictionary to hold all the links
self.links = {}
# Create a dictionary to hold a reference set of links
self.links_reference = {}
# Create a dictionary to hold all the green IERs (this will come from an external source)
self.green_iers = {}
# Create a dictionary to hold all the node PoLs (this will come from an external source)
self.node_pol = {}
# Create a dictionary to hold all the red agent IERs (this will come from an external source)
self.red_iers = {}
# Create a dictionary to hold all the red agent node PoLs (this will come from an external source)
self.red_node_pol = {}
# Create the Access Control List
self.acl = AccessControlList()
# Create a list of services (enums)
self.services_list = []
# Create a list of ports
self.ports_list = []
# Create graph (network)
self.network = nx.MultiGraph()
# Create a graph (network) reference
self.network_reference = nx.MultiGraph()
# Create step count
self.step_count = 0
# Create step info dictionary
self.step_info = {}
# Total reward
self.total_reward = 0
# Average reward
self.average_reward = 0
# Episode count
self.episode_count = 0
# Number of nodes - gets a value by examining the nodes dictionary after it's been populated
self.num_nodes = 0
# Number of links - gets a value by examining the links dictionary after it's been populated
self.num_links = 0
# Number of services - gets a value when config is loaded
self.num_services = 0
# Number of ports - gets a value when config is loaded
self.num_ports = 0
# The action type
self.action_type = 0
# 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_data = yaml.safe_load(self.config_file)
self.load_config()
except Exception as e:
logging.error("Could not load the environment configuration")
logging.error("Exception occured", exc_info=True)
# Store the node objects as node attributes
# (This is so we can access them as objects)
for node in self.network:
self.network.nodes[node]["self"] = node
for node in self.network_reference:
self.network_reference.nodes[node]["self"] = node
self.num_nodes = len(self.nodes)
self.num_links = len(self.links)
# Visualise in PNG
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")
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:
logging.error("Could not save network diagram")
logging.error("Exception occured", exc_info=True)
print("Could not save network diagram")
# Define Observation Space
# x = number of nodes and links (i.e. items)
# y = number of parameters to be sent
# For each item, we send:
# - [For Nodes] | [For Links]
# - node ID | link ID
# - operating state | N/A
# - operating system state | N/A
# - service A state | service A loading
# - service B state | service B loading
# - service C state | service C loading
# - service D state | service D loading
# - service E state | service E loading
# - 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
num_items = self.num_links + self.num_nodes
# Set the number of observation parameters, being # of services plus id, operating sytem system and O/S state (i.e. 3)
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)
# 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)
# Define Action Space - depends on action space type (Node or ACL)
if self.action_type == ACTION_TYPE.NODE:
logging.info("Action space type NODE selected")
# Terms (for node action space):
# [0, num nodes] - node ID (0 = nothing, node ID)
# [0, 3] - what property it's acting on (0 = nothing, state, o/s state, service state)
# [0, 3] - action on property (0 = nothing, On, Off, Reset / Patch)
# [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])
else:
logging.info("Action space type ACL selected")
# Terms (for ACL action space):
# [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule)
# [0, 1] - Permission (0 = DENY, 1 = ALLOW)
# [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses)
# [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])
# 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='')
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")
logging.error("Exception occured", exc_info=True)
def reset(self):
"""
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()
# Reset the node statuses and recreate the ACL from config
# Does this for both live and reference nodes
self.reset_environment()
# Reset counters and totals
self.total_reward = 0
self.step_count = 0
self.average_reward = 0
# Update observations space and return
self.update_environent_obs()
return self.env_obs
def step(self, action):
"""
AI Gym Step function
Args:
action: Action space from agent
Returns:
env_obs: Observation space
reward: Reward value for this step
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")
# TEMP
done = False
self.step_count += 1
#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)
# 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
transaction.set_action_space(copy.deepcopy(action))
# 1. Perform any time-based activities (e.g. a component moving from patching to good)
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
# Take snapshots of nodes and links
self.nodes_post_pol = copy.deepcopy(self.nodes)
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
# 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)
# 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
# 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)
# Take snapshots of nodes and links
self.nodes_post_blue = copy.deepcopy(self.nodes)
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))
self.total_reward += reward
if self.step_count == self.episode_steps:
self.average_reward = self.total_reward / self.step_count
print("Average reward: " + str(self.average_reward))
# Load the reward into the transaction
transaction.set_reward(reward)
# 7. Output Verbose
#self.output_link_status()
# 8. Update env_obs
self.update_environent_obs()
# Load the new observation space into the transaction
transaction.set_obs_space_post(copy.deepcopy(self.env_obs))
# 9. Add the transaction to the list of transactions
self.transaction_list.append(copy.deepcopy(transaction))
# Return
return self.env_obs, reward, done, self.step_info
def __close__(self):
"""
Override close function
"""
self.csv_file.close()
self.config_file.close()
def init_acl(self):
"""
Initialise the Access Control List
"""
self.acl.remove_all_rules()
def output_link_status(self):
"""
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()))
def interpret_action_and_apply(self, _action):
"""
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)
else:
self.apply_actions_to_acl(_action)
def apply_actions_to_nodes(self, _action):
"""
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]
service_index = _action[3]
# Check that the action is requesting a valid node
try:
node = self.nodes[str(node_id)]
except:
return
if node_property == 0:
# This is the do nothing action
return
elif node_property == 1:
# This is an action on the node Operating State
if property_action == 0:
# Do nothing
return
elif property_action == 1:
# Turn on (only applicable if it's OFF, not if it's patching)
if node.get_state() == HARDWARE_STATE.OFF:
node.turn_on()
elif property_action == 2:
# Turn off
node.turn_off()
elif property_action == 3:
# Reset (only applicable if it's ON)
if node.get_state() == HARDWARE_STATE.ON:
node.reset()
else:
return
elif node_property == 2:
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
# This is an action on the node Operating System State
if property_action == 0:
# Do nothing
return
elif property_action == 1:
# Patch (valid action if it's good or compromised)
node.set_os_state(SOFTWARE_STATE.PATCHING)
else:
# Node is not of Active or Service Type
return
elif node_property == 3:
# This is an action on a node Service State
if isinstance(node, ServiceNode):
# This is an action on a node Service State
if property_action == 0:
# Do nothing
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)
else:
# Node is not of Service Type
return
else:
return
def apply_actions_to_acl(self, _action):
"""
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]
action_destination_ip = _action[3]
action_protocol = _action[4]
action_port = _action[5]
if action_decision == 0:
# It's decided to do nothing
return
else:
# It's decided to create a new ACL rule or remove an existing rule
# Permission value
if action_permission == 0:
acl_rule_permission = "DENY"
else:
acl_rule_permission = "ALLOW"
# Source IP value
if action_source_ip == 0:
acl_rule_source = "ANY"
else:
node = list(self.nodes.values())[action_source_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
acl_rule_source = node.get_ip_address()
else:
return
# Destination IP value
if action_destination_ip == 0:
acl_rule_destination = "ANY"
else:
node = list(self.nodes.values())[action_destination_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
acl_rule_destination = node.get_ip_address()
else:
return
# Protocol value
if action_protocol == 0:
acl_rule_protocol = "ANY"
else:
acl_rule_protocol = self.services_list[action_protocol - 1]
# Port value
if action_port == 0:
acl_rule_port = "ANY"
else:
acl_rule_port = self.ports_list[action_port - 1]
# 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)
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)
else:
return
def apply_time_based_updates(self):
"""
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()
else:
pass
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
if node.get_os_state() == SOFTWARE_STATE.PATCHING:
node.update_os_patching_status()
else:
pass
else:
pass
if isinstance(node, ServiceNode):
node.update_services_patching_status()
else:
pass
for node_key, node in self.nodes_reference.items():
if node.get_state() == HARDWARE_STATE.RESETTING:
node.update_resetting_status()
else:
pass
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
if node.get_os_state() == SOFTWARE_STATE.PATCHING:
node.update_os_patching_status()
else:
pass
else:
pass
if isinstance(node, ServiceNode):
node.update_services_patching_status()
else:
pass
def update_environent_obs(self):
"""
# Updates the observation space based on the node and link status
"""
item_index = 0
# Do nodes first
for node_key, node in self.nodes.items():
self.env_obs[item_index][0] = int(node.get_id())
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
else:
self.env_obs[item_index][2] = 0
service_index = 3
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
else:
self.env_obs[item_index][service_index] = 0
service_index += 1
else:
# Not a service node
for service in self.services_list:
self.env_obs[item_index][service_index] = 0
service_index += 1
item_index += 1
# Now do links
for link_key, link in self.links.items():
self.env_obs[item_index][0] = int(link.get_id())
self.env_obs[item_index][1] = 0
self.env_obs[item_index][2] = 0
protocol_list = link.get_protocol_list()
protocol_index = 0
for protocol in protocol_list:
self.env_obs[item_index][protocol_index + 3] = protocol.get_load()
protocol_index += 1
item_index += 1
def load_config(self):
"""
# 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)
elif item["itemType"] == "GREEN_IER":
# Create a Green IER
self.create_green_ier(item)
elif item["itemType"] == "GREEN_POL":
# Create a Green PoL
self.create_green_pol(item)
elif item["itemType"] == "RED_IER":
# Create a Red IER
self.create_red_ier(item)
elif item["itemType"] == "RED_POL":
# Create a Red PoL
self.create_red_pol(item)
elif item["itemType"] == "ACL_RULE":
# Create an ACL rule
self.create_acl_rule(item)
elif item["itemType"] == "SERVICES":
# Create the list of services
self.create_services_list(item)
elif item["itemType"] == "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)
pass
logging.info("Environment configuration loaded")
print("Environment configuration loaded")
def create_node(self, item):
"""
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"]
node_base_type = item["baseType"]
node_type = TYPE[item["nodeType"]]
node_priority = PRIORITY[item["priority"]]
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)
elif node_base_type == "ACTIVE":
# Active nodes have IP address and operating system state
node_ip_address = item["ipAddress"]
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node = ActiveNode(node_id, node_name, node_type, node_priority, node_hardware_state, node_ip_address, node_software_state, self.config_values)
elif node_base_type == "SERVICE":
# Service nodes have IP address, operating system state and list of services
node_ip_address = item["ipAddress"]
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node = ServiceNode(node_id, node_name, node_type, node_priority, node_hardware_state, node_ip_address, node_software_state, self.config_values)
node_services = item["services"]
for service in node_services:
service_protocol = service["name"]
service_port = service["port"]
service_state = SOFTWARE_STATE[service["state"]]
node.add_service(Service(service_protocol, service_port, service_state))
else:
# Bad formatting
pass
# Copy the node for the reference version
node_ref = copy.deepcopy(node)
# Add node to node dictionary
self.nodes[node_id] = node
# Add reference node to reference node dictionary
self.nodes_reference[node_id] = node_ref
# Add node to network
self.network.add_nodes_from([node])
# Add node to network (reference)
self.network_reference.add_nodes_from([node_ref])
def create_link(self, item):
"""
Creates a link from config data
Args:
item: A config data item
"""
link_id = item["id"]
link_name = item["name"]
link_bandwidth = item["bandwidth"]
link_source = item["source"]
link_destination = item["destination"]
source_node = self.nodes[link_source]
dest_node = self.nodes[link_destination]
# Add link to network
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)
# Reference
source_node_ref = self.nodes_reference[link_source]
dest_node_ref = self.nodes_reference[link_destination]
# Add link to network (reference)
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)
def create_green_ier(self, item):
"""
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"]
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"]
# 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)
def create_red_ier(self, item):
"""
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"]
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"]
# 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)
def create_green_pol(self, item):
"""
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["node"]
pol_type = NODE_POL_TYPE[item["type"]]
pol_protocol = item["protocol"]
# State depends on whether this is Operating, O/S or Service PoL type
if pol_type == NODE_POL_TYPE.OPERATING:
pol_state = HARDWARE_STATE[item["state"]]
else:
pol_state = SOFTWARE_STATE[item["state"]]
self.node_pol[pol_id] = NodeStateInstruction(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
Args:
item: A config data item
"""
pol_id = item["id"]
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_node = item["node"]
pol_type = NODE_POL_TYPE[item["type"]]
pol_protocol = item["protocol"]
# State depends on whether this is Operating, O/S or Service PoL type
if pol_type == NODE_POL_TYPE.OPERATING:
pol_state = HARDWARE_STATE[item["state"]]
else:
pol_state = SOFTWARE_STATE[item["state"]]
pol_is_entry_node = item["isEntryNode"]
self.red_node_pol[pol_id] = NodeStateInstruction(pol_id, pol_start_step, pol_end_step, pol_node, pol_type, pol_protocol, pol_state, pol_is_entry_node)
def create_acl_rule(self, item):
"""
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)
def create_services_list(self, services):
"""
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:
service_name = service["name"]
self.services_list.append(service_name)
# Set the number of services
self.num_services = len(self.services_list)
def create_ports_list(self, ports):
"""
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:
port_value = port["port"]
self.ports_list.append(port_value)
# Set the number of ports
self.num_ports = len(self.ports_list)
def get_action_info(self, 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
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
"""
for item in self.config_data:
if item["itemType"] == "NODE":
# Reset a node's state (normal and reference)
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)
else:
# 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():
ier_value.set_is_running(False)
# Red IERs
for ier_key, ier_value in self.red_iers.items():
ier_value.set_is_running(False)
def reset_node(self, item):
"""
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"]
node_hardware_state = HARDWARE_STATE[item["hardwareState"]]
node = self.nodes[node_id]
node_ref = self.nodes_reference[node_id]
# Reset the hardware state (common for all node types)
node.set_state(node_hardware_state)
node_ref.set_state(node_hardware_state)
if node_base_type == "ACTIVE":
# Active nodes have operating system state
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node.set_os_state(node_software_state)
node_ref.set_os_state(node_software_state)
elif node_base_type == "SERVICE":
# Service nodes have operating system state and list of services
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node.set_os_state(node_software_state)
node_ref.set_os_state(node_software_state)
# Update service states
node_services = item["services"]
for service in node_services:
service_protocol = service["name"]
service_state = SOFTWARE_STATE[service["state"]]
# Update node service state
node.set_service_state(service_protocol, service_state)
# Update reference node service state
node_ref.set_service_state(service_protocol, service_state)
else:
# Bad formatting
pass

224
PRIMAITE/environment/reward.py Normal file
View File

@@ -0,0 +1,224 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""
Implements reward function
"""
from common.enums import *
from nodes.active_node import ActiveNode
from nodes.service_node import ServiceNode
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
Args:
initial_nodes: The nodes before red and blue agents take effect
final_nodes: The nodes after red and blue agents take effect
reference_nodes: The nodes if there had been no red or blue effect
green_iers: The green IERs (should be running)
red_iers: Should be stopeed (ideally) by the blue agent
step_count: current step
"""
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)
# 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)
# Service State
if (isinstance(final_node, ServiceNode)):
reward_value += score_node_service_state(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()
stop_step = ier_value.get_end_step()
if step_count >= start_step and step_count <= stop_step:
if ier_value.get_is_running():
reward_value += config_values.red_ier_running
# Go through each green IER - penalise if it's not running (weighted)
for ier_key, ier_value in green_iers.items():
start_step = ier_value.get_start_step()
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()
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
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
final_node_operating_state = final_node.get_state()
initial_node_operating_state = initial_node.get_state()
reference_node_operating_state = reference_node.get_state()
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:
# 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:
if final_node_operating_state == HARDWARE_STATE.OFF:
score += config_values.off_should_be_on
elif final_node_operating_state == HARDWARE_STATE.RESETTING:
score += config_values.resetting_should_be_on
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.OFF:
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
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.RESETTING:
if final_node_operating_state == HARDWARE_STATE.ON:
score += config_values.on_should_be_resetting
elif final_node_operating_state == HARDWARE_STATE.OFF:
score += config_values.off_should_be_resetting
else:
pass
else:
pass
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
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
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()
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:
# 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:
if final_node_os_state == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_good
elif final_node_os_state == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_good
else:
pass
elif initial_node_os_state == SOFTWARE_STATE.PATCHING:
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
else:
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
else:
pass
else:
pass
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
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
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]
if final_service.get_state() == reference_service.get_state():
# All is well - we're no different from the reference situation
score += config_values.all_ok
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_service.get_state() == SOFTWARE_STATE.GOOD:
if final_service.get_state() == SOFTWARE_STATE.PATCHING:
score += config_values.patching_should_be_good
elif final_service.get_state() == SOFTWARE_STATE.COMPROMISED:
score += config_values.compromised_should_be_good
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_good
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.PATCHING:
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
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_patching
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.COMPROMISED:
if final_service.get_state() == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_compromised
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
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed_should_be_compromised
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
if final_service.get_state() == SOFTWARE_STATE.GOOD:
score += config_values.good_should_be_overwhelmed
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
elif final_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
score += config_values.overwhelmed
else:
pass
else:
pass
return score