e11fd2ced4
- Dropped support for overriding the num_episodes and num_steps at the agent level. It's just not needed and will add complexity when overriding and writing output files.
554 lines
17 KiB
Python
554 lines
17 KiB
Python
import numpy as np
|
|
|
|
from primaite.common.enums import (
|
|
HardwareState,
|
|
LinkStatus,
|
|
NodeHardwareAction,
|
|
NodePOLType,
|
|
NodeSoftwareAction,
|
|
SoftwareState,
|
|
)
|
|
|
|
|
|
def transform_action_node_readable(action):
|
|
"""
|
|
Convert a node action from enumerated format to readable format.
|
|
|
|
example:
|
|
[1, 3, 1, 0] -> [1, 'SERVICE', 'PATCHING', 0]
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
action_node_property = NodePOLType(action[1]).name
|
|
|
|
if action_node_property == "OPERATING":
|
|
property_action = NodeHardwareAction(action[2]).name
|
|
elif (action_node_property == "OS" or action_node_property == "SERVICE") and action[2] <= 1:
|
|
property_action = NodeSoftwareAction(action[2]).name
|
|
else:
|
|
property_action = "NONE"
|
|
|
|
new_action = [action[0], action_node_property, property_action, action[3]]
|
|
return new_action
|
|
|
|
|
|
def transform_action_acl_readable(action):
|
|
"""
|
|
Transform an ACL action to a more readable format.
|
|
|
|
example:
|
|
[0, 1, 2, 5, 0, 1] -> ['NONE', 'ALLOW', 2, 5, 'ANY', 1]
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
action_decisions = {0: "NONE", 1: "CREATE", 2: "DELETE"}
|
|
action_permissions = {0: "DENY", 1: "ALLOW"}
|
|
|
|
action_decision = action_decisions[action[0]]
|
|
action_permission = action_permissions[action[1]]
|
|
|
|
# For IPs, Ports and Protocols, 0 means any, otherwise its just an index
|
|
new_action = [action_decision, action_permission] + list(action[2:6])
|
|
for n, val in enumerate(list(action[2:6])):
|
|
if val == 0:
|
|
new_action[n + 2] = "ANY"
|
|
|
|
return new_action
|
|
|
|
|
|
def is_valid_node_action(action):
|
|
"""Is the node action an actual valid action.
|
|
|
|
Only uses information about the action to determine if the action has an effect
|
|
|
|
Does NOT consider:
|
|
- Node ID not valid to perform an operation - e.g. selected node has no service so cannot patch
|
|
- Node already being in that state (turning an ON node ON)
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
action_r = transform_action_node_readable(action)
|
|
|
|
node_property = action_r[1]
|
|
node_action = action_r[2]
|
|
|
|
# print("node property", node_property, "\nnode action", node_action)
|
|
|
|
if node_property == "NONE":
|
|
return False
|
|
if node_action == "NONE":
|
|
return False
|
|
if node_property == "OPERATING" and node_action == "PATCHING":
|
|
# Operating State cannot PATCH
|
|
return False
|
|
if node_property != "OPERATING" and node_action not in [
|
|
"NONE",
|
|
"PATCHING",
|
|
]:
|
|
# Software States can only do Nothing or Patch
|
|
return False
|
|
return True
|
|
|
|
|
|
def is_valid_acl_action(action):
|
|
"""
|
|
Is the ACL action an actual valid action.
|
|
|
|
Only uses information about the action to determine if the action has an effect.
|
|
|
|
Does NOT consider:
|
|
- Trying to create identical rules
|
|
- Trying to create a rule which is a subset of another rule (caused by "ANY")
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
action_r = transform_action_acl_readable(action)
|
|
|
|
action_decision = action_r[0]
|
|
action_permission = action_r[1]
|
|
action_source_id = action_r[2]
|
|
action_destination_id = action_r[3]
|
|
|
|
if action_decision == "NONE":
|
|
return False
|
|
if action_source_id == action_destination_id and action_source_id != "ANY" and action_destination_id != "ANY":
|
|
# ACL rule towards itself
|
|
return False
|
|
if action_permission == "DENY":
|
|
# DENY is unnecessary, we can create and delete allow rules instead
|
|
# No allow rule = blocked/DENY by feault. ALLOW overrides existing DENY.
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def is_valid_acl_action_extra(action):
|
|
"""
|
|
Harsher version of valid acl actions, does not allow action.
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
if is_valid_acl_action(action) is False:
|
|
return False
|
|
|
|
action_r = transform_action_acl_readable(action)
|
|
action_protocol = action_r[4]
|
|
action_port = action_r[5]
|
|
|
|
# Don't allow protocols or ports to be ANY
|
|
# in the future we might want to do the opposite, and only have ANY option for ports and service
|
|
if action_protocol == "ANY":
|
|
return False
|
|
if action_port == "ANY":
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def transform_change_obs_readable(obs):
|
|
"""
|
|
Transform list of transactions to readable list of each observation property.
|
|
|
|
example:
|
|
np.array([[1,2,1,3],[2,1,1,1]]) -> [[1, 2], ['OFF', 'ON'], ['GOOD', 'GOOD'], ['COMPROMISED', 'GOOD']]
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
ids = [i for i in obs[:, 0]]
|
|
operating_states = [HardwareState(i).name for i in obs[:, 1]]
|
|
os_states = [SoftwareState(i).name for i in obs[:, 2]]
|
|
new_obs = [ids, operating_states, os_states]
|
|
|
|
for service in range(3, obs.shape[1]):
|
|
# Links bit/s don't have a service state
|
|
service_states = [SoftwareState(i).name if i <= 4 else i for i in obs[:, service]]
|
|
new_obs.append(service_states)
|
|
|
|
return new_obs
|
|
|
|
|
|
def transform_obs_readable(obs):
|
|
"""
|
|
Transform observation to readable format.
|
|
|
|
np.array([[1,2,1,3],[2,1,1,1]]) -> [[1, 'OFF', 'GOOD', 'COMPROMISED'], [2, 'ON', 'GOOD', 'GOOD']]
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
changed_obs = transform_change_obs_readable(obs)
|
|
new_obs = list(zip(*changed_obs))
|
|
# Convert list of tuples to list of lists
|
|
new_obs = [list(i) for i in new_obs]
|
|
|
|
return new_obs
|
|
|
|
|
|
def convert_to_new_obs(obs, num_nodes=10):
|
|
"""
|
|
Convert original gym Box observation space to new multiDiscrete observation space.
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
# Remove ID columns, remove links and flatten to MultiDiscrete observation space
|
|
new_obs = obs[:num_nodes, 1:].flatten()
|
|
return new_obs
|
|
|
|
|
|
def convert_to_old_obs(obs, num_nodes=10, num_links=10, num_services=1):
|
|
"""
|
|
Convert to old observation.
|
|
|
|
Links filled with 0's as no information is included in new observation space.
|
|
|
|
example:
|
|
obs = array([1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 1, 1, 1])
|
|
|
|
new_obs = array([[ 1, 1, 1, 1],
|
|
[ 2, 1, 1, 1],
|
|
[ 3, 1, 1, 1],
|
|
...
|
|
[20, 0, 0, 0]])
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
# Convert back to more readable, original format
|
|
reshaped_nodes = obs[:-num_links].reshape(num_nodes, num_services + 2)
|
|
|
|
# Add empty links back and add node ID back
|
|
s = np.zeros(
|
|
[reshaped_nodes.shape[0] + num_links, reshaped_nodes.shape[1] + 1],
|
|
dtype=np.int64,
|
|
)
|
|
s[:, 0] = range(1, num_nodes + num_links + 1) # Adding ID back
|
|
s[:num_nodes, 1:] = reshaped_nodes # put values back in
|
|
new_obs = s
|
|
|
|
# Add links back in
|
|
links = obs[-num_links:]
|
|
# Links will be added to the last protocol/service slot but they are not specific to that service
|
|
new_obs[num_nodes:, -1] = links
|
|
|
|
return new_obs
|
|
|
|
|
|
def describe_obs_change(obs1, obs2, num_nodes=10, num_links=10, num_services=1):
|
|
"""
|
|
Return string describing change between two observations.
|
|
|
|
example:
|
|
obs_1 = array([[1, 1, 1, 1, 3], [2, 1, 1, 1, 1]])
|
|
obs_2 = array([[1, 1, 1, 1, 1], [2, 1, 1, 1, 1]])
|
|
output = 'ID 1: SERVICE 2 set to GOOD'
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
obs1 = convert_to_old_obs(obs1, num_nodes, num_links, num_services)
|
|
obs2 = convert_to_old_obs(obs2, num_nodes, num_links, num_services)
|
|
list_of_changes = []
|
|
for n, row in enumerate(obs1 - obs2):
|
|
if row.any() != 0:
|
|
relevant_changes = np.where(row != 0, obs2[n], -1)
|
|
relevant_changes[0] = obs2[n, 0] # ID is always relevant
|
|
is_link = relevant_changes[0] > num_nodes
|
|
desc = _describe_obs_change_helper(relevant_changes, is_link)
|
|
list_of_changes.append(desc)
|
|
|
|
change_string = "\n ".join(list_of_changes)
|
|
if len(list_of_changes) > 0:
|
|
change_string = "\n " + change_string
|
|
return change_string
|
|
|
|
|
|
def _describe_obs_change_helper(obs_change, is_link):
|
|
"""
|
|
Helper funcion to describe what has changed.
|
|
|
|
example:
|
|
[ 1 -1 -1 -1 1] -> "ID 1: Service 1 changed to GOOD"
|
|
|
|
Handles multiple changes e.g. 'ID 1: SERVICE 1 changed to PATCHING. SERVICE 2 set to GOOD.'
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
# Indexes where a change has occured, not including 0th index
|
|
index_changed = [i for i in range(1, len(obs_change)) if obs_change[i] != -1]
|
|
# Node pol types, Indexes >= 3 are service nodes
|
|
NodePOLTypes = [NodePOLType(i).name if i < 3 else NodePOLType(3).name + " " + str(i - 3) for i in index_changed]
|
|
# Account for hardware states, software sattes and links
|
|
states = [
|
|
LinkStatus(obs_change[i]).name
|
|
if is_link
|
|
else HardwareState(obs_change[i]).name
|
|
if i == 1
|
|
else SoftwareState(obs_change[i]).name
|
|
for i in index_changed
|
|
]
|
|
|
|
if not is_link:
|
|
desc = f"ID {obs_change[0]}:"
|
|
for node_pol_type, state in list(zip(NodePOLTypes, states)):
|
|
desc = desc + " " + node_pol_type + " changed to " + state + "."
|
|
else:
|
|
desc = f"ID {obs_change[0]}: Link traffic changed to {states[0]}."
|
|
|
|
return desc
|
|
|
|
|
|
def transform_action_node_enum(action):
|
|
"""
|
|
Convert a node action from readable string format, to enumerated format.
|
|
|
|
example:
|
|
[1, 'SERVICE', 'PATCHING', 0] -> [1, 3, 1, 0]
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
action_node_id = action[0]
|
|
action_node_property = NodePOLType[action[1]].value
|
|
|
|
if action[1] == "OPERATING":
|
|
property_action = NodeHardwareAction[action[2]].value
|
|
elif action[1] == "OS" or action[1] == "SERVICE":
|
|
property_action = NodeSoftwareAction[action[2]].value
|
|
else:
|
|
property_action = 0
|
|
|
|
action_service_index = action[3]
|
|
|
|
new_action = [
|
|
action_node_id,
|
|
action_node_property,
|
|
property_action,
|
|
action_service_index,
|
|
]
|
|
|
|
return new_action
|
|
|
|
|
|
def transform_action_node_readable(action):
|
|
"""
|
|
Convert a node action from enumerated format to readable format.
|
|
|
|
example:
|
|
[1, 3, 1, 0] -> [1, 'SERVICE', 'PATCHING', 0]
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
action_node_property = NodePOLType(action[1]).name
|
|
|
|
if action_node_property == "OPERATING":
|
|
property_action = NodeHardwareAction(action[2]).name
|
|
elif (action_node_property == "OS" or action_node_property == "SERVICE") and action[2] <= 1:
|
|
property_action = NodeSoftwareAction(action[2]).name
|
|
else:
|
|
property_action = "NONE"
|
|
|
|
new_action = [action[0], action_node_property, property_action, action[3]]
|
|
return new_action
|
|
|
|
|
|
def node_action_description(action):
|
|
"""
|
|
Generate string describing a node-based action.
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
if isinstance(action[1], (int, np.int64)):
|
|
# transform action to readable format
|
|
action = transform_action_node_readable(action)
|
|
|
|
node_id = action[0]
|
|
node_property = action[1]
|
|
property_action = action[2]
|
|
service_id = action[3]
|
|
|
|
if property_action == "NONE":
|
|
return ""
|
|
if node_property == "OPERATING" or node_property == "OS":
|
|
description = f"NODE {node_id}, {node_property}, SET TO {property_action}"
|
|
elif node_property == "SERVICE":
|
|
description = f"NODE {node_id} FROM SERVICE {service_id}, SET TO {property_action}"
|
|
else:
|
|
return ""
|
|
|
|
return description
|
|
|
|
|
|
def transform_action_acl_enum(action):
|
|
"""
|
|
Convert acl action from readable str format, to enumerated format.
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
action_decisions = {"NONE": 0, "CREATE": 1, "DELETE": 2}
|
|
action_permissions = {"DENY": 0, "ALLOW": 1}
|
|
|
|
action_decision = action_decisions[action[0]]
|
|
action_permission = action_permissions[action[1]]
|
|
|
|
# For IPs, Ports and Protocols, ANY has value 0, otherwise its just an index
|
|
new_action = [action_decision, action_permission] + list(action[2:6])
|
|
for n, val in enumerate(list(action[2:6])):
|
|
if val == "ANY":
|
|
new_action[n + 2] = 0
|
|
|
|
new_action = np.array(new_action)
|
|
return new_action
|
|
|
|
|
|
def acl_action_description(action):
|
|
"""
|
|
Generate string describing an acl-based action.
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
if isinstance(action[0], (int, np.int64)):
|
|
# transform action to readable format
|
|
action = transform_action_acl_readable(action)
|
|
if action[0] == "NONE":
|
|
description = "NO ACL RULE APPLIED"
|
|
else:
|
|
description = (
|
|
f"{action[0]} RULE: {action[1]} traffic from IP {action[2]} to IP {action[3]},"
|
|
f" for protocol/service index {action[4]} on port index {action[5]}"
|
|
)
|
|
|
|
return description
|
|
|
|
|
|
def get_node_of_ip(ip, node_dict):
|
|
"""
|
|
Get the node ID of an IP address.
|
|
|
|
node_dict: dictionary of nodes where key is ID, and value is the node (can be ontained from env.nodes)
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
for node_key, node_value in node_dict.items():
|
|
node_ip = node_value.ip_address
|
|
if node_ip == ip:
|
|
return node_key
|
|
|
|
|
|
def is_valid_node_action(action):
|
|
"""Is the node action an actual valid action.
|
|
|
|
Only uses information about the action to determine if the action has an effect
|
|
|
|
Does NOT consider:
|
|
- Node ID not valid to perform an operation - e.g. selected node has no service so cannot patch
|
|
- Node already being in that state (turning an ON node ON)
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
action_r = transform_action_node_readable(action)
|
|
|
|
node_property = action_r[1]
|
|
node_action = action_r[2]
|
|
|
|
if node_property == "NONE":
|
|
return False
|
|
if node_action == "NONE":
|
|
return False
|
|
if node_property == "OPERATING" and node_action == "PATCHING":
|
|
# Operating State cannot PATCH
|
|
return False
|
|
if node_property != "OPERATING" and node_action not in [
|
|
"NONE",
|
|
"PATCHING",
|
|
]:
|
|
# Software States can only do Nothing or Patch
|
|
return False
|
|
return True
|
|
|
|
|
|
def is_valid_acl_action(action):
|
|
"""
|
|
Is the ACL action an actual valid action.
|
|
|
|
Only uses information about the action to determine if the action has an effect
|
|
|
|
Does NOT consider:
|
|
- Trying to create identical rules
|
|
- Trying to create a rule which is a subset of another rule (caused by "ANY")
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
action_r = transform_action_acl_readable(action)
|
|
|
|
action_decision = action_r[0]
|
|
action_permission = action_r[1]
|
|
action_source_id = action_r[2]
|
|
action_destination_id = action_r[3]
|
|
|
|
if action_decision == "NONE":
|
|
return False
|
|
if action_source_id == action_destination_id and action_source_id != "ANY" and action_destination_id != "ANY":
|
|
# ACL rule towards itself
|
|
return False
|
|
if action_permission == "DENY":
|
|
# DENY is unnecessary, we can create and delete allow rules instead
|
|
# No allow rule = blocked/DENY by feault. ALLOW overrides existing DENY.
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def is_valid_acl_action_extra(action):
|
|
"""
|
|
Harsher version of valid acl actions, does not allow action.
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
if is_valid_acl_action(action) is False:
|
|
return False
|
|
|
|
action_r = transform_action_acl_readable(action)
|
|
action_protocol = action_r[4]
|
|
action_port = action_r[5]
|
|
|
|
# Don't allow protocols or ports to be ANY
|
|
# in the future we might want to do the opposite, and only have ANY option for ports and service
|
|
if action_protocol == "ANY":
|
|
return False
|
|
if action_port == "ANY":
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def get_new_action(old_action, action_dict):
|
|
"""
|
|
Get new action (e.g. 32) from old action e.g. [1,1,1,0].
|
|
|
|
Old_action can be either node or acl action type
|
|
|
|
TODO: Add params and return in docstring.
|
|
TODO: Typehint params and return.
|
|
"""
|
|
for key, val in action_dict.items():
|
|
if list(val) == list(old_action):
|
|
return key
|
|
# Not all possible actions are included in dict, only valid action are
|
|
# if action is not in the dict, its an invalid action so return 0
|
|
return 0
|