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

#917 - Synced with dev and added better logging

Browse Source
This commit is contained in:
Chris McCarthy
2023-06-28 12:01:01 +01:00
parent 498e6a7ac1
commit 7482192046
12 changed files with 170 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -15,7 +15,8 @@ from gym import Env, spaces
from matplotlib import pyplot as plt
from primaite.acl.access_control_list import AccessControlList
from primaite.agents.utils import is_valid_acl_action_extra, is_valid_node_action
from primaite.agents.utils import is_valid_acl_action_extra, \
is_valid_node_action
from primaite.common.custom_typing import NodeUnion
from primaite.common.enums import (
ActionType,
@@ -36,13 +37,15 @@ from primaite.environment.reward import calculate_reward_function
from primaite.links.link import Link
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.node import Node
from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen
from primaite.nodes.node_state_instruction_green import \
NodeStateInstructionGreen
from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
from primaite.nodes.passive_node import PassiveNode
from primaite.nodes.service_node import ServiceNode
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.pol.red_agent_pol import apply_red_agent_iers, \
apply_red_agent_node_pol
from primaite.transactions.transaction import Transaction
from primaite.transactions.transactions_to_file import \
write_transaction_to_file
@@ -61,12 +64,12 @@ class Primaite(Env):
ACTION_SPACE_ACL_PERMISSION_VALUES: int = 2
def __init__(
self,
training_config_path: Union[str, Path],
lay_down_config_path: Union[str, Path],
transaction_list,
session_path: Path,
timestamp_str: str,
self,
training_config_path: Union[str, Path],
lay_down_config_path: Union[str, Path],
transaction_list,
session_path: Path,
timestamp_str: str,
):
"""
The Primaite constructor.
@@ -86,6 +89,7 @@ class Primaite(Env):
self.training_config: TrainingConfig = training_config.load(
training_config_path
)
_LOGGER.info(f"Using: {str(self.training_config)}")
# Number of steps in an episode
self.episode_steps = self.training_config.num_steps
@@ -207,16 +211,14 @@ class Primaite(Env):
plt.savefig(file_path, format="PNG")
plt.clf()
except Exception:
_LOGGER.error("Could not save network diagram")
_LOGGER.error("Exception occured", exc_info=True)
print("Could not save network diagram")
_LOGGER.error("Could not save network diagram", exc_info=True)
# Initiate observation space
self.observation_space, self.env_obs = self.init_observations()
# Define Action Space - depends on action space type (Node or ACL)
if self.training_config.action_type == ActionType.NODE:
_LOGGER.info("Action space type NODE selected")
_LOGGER.debug("Action space type NODE selected")
# Terms (for node action space):
# [0, num nodes] - node ID (0 = nothing, node ID)
# [0, 4] - what property it's acting on (0 = nothing, state, SoftwareState, service state, file system state) # noqa
@@ -225,7 +227,7 @@ class Primaite(Env):
self.action_dict = self.create_node_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict))
elif self.training_config.action_type == ActionType.ACL:
_LOGGER.info("Action space type ACL selected")
_LOGGER.debug("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)
@@ -236,11 +238,11 @@ class Primaite(Env):
self.action_dict = self.create_acl_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict))
elif self.training_config.action_type == ActionType.ANY:
_LOGGER.info("Action space type ANY selected - Node + ACL")
_LOGGER.debug("Action space type ANY selected - Node + ACL")
self.action_dict = self.create_node_and_acl_action_dict()
self.action_space = spaces.Discrete(len(self.action_dict))
else:
_LOGGER.info(
_LOGGER.error(
f"Invalid action type selected: {self.training_config.action_type}"
)
# Set up a csv to store the results of the training
@@ -301,17 +303,14 @@ class Primaite(Env):
done: Indicates episode is complete if True
step_info: Additional information relating to this step
"""
# Introduce a delay between steps
time.sleep(self.training_config.time_delay / 1000)
if self.step_count == 0:
print(f"Episode: {str(self.episode_count)}")
_LOGGER.info(f"Episode: {str(self.episode_count)}")
# TEMP
done = False
self.step_count += 1
self.total_step_count += 1
# print("Episode step: " + str(self.step_count))
# Need to clear traffic on all links first
for link_key, link_value in self.links.items():
@@ -322,7 +321,8 @@ class Primaite(Env):
# Create a Transaction (metric) object for this step
transaction = Transaction(
datetime.now(), self.agent_identifier, self.episode_count, self.step_count
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))
@@ -352,7 +352,8 @@ class Primaite(Env):
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_node_pol(self.nodes_reference, self.node_pol,
self.step_count) # Node PoL
apply_iers(
self.network_reference,
self.nodes_reference,
@@ -389,7 +390,7 @@ class Primaite(Env):
self.step_count,
self.training_config,
)
# print(f" Step {self.step_count} Reward: {str(reward)}")
_LOGGER.debug(f" Step {self.step_count} Reward: {str(reward)}")
self.total_reward += reward
if self.step_count == self.episode_steps:
self.average_reward = self.total_reward / self.step_count
@@ -397,7 +398,7 @@ class Primaite(Env):
# For evaluation, need to trigger the done value = True when
# step count is reached in order to prevent neverending episode
done = True
print(f" Average Reward: {str(self.average_reward)}")
_LOGGER.info(f" Average Reward: {str(self.average_reward)}")
# Load the reward into the transaction
transaction.set_reward(reward)
@@ -428,6 +429,7 @@ class Primaite(Env):
self.timestamp_str
)
self.csv_file.close()
def init_acl(self):
"""Initialise the Access Control List."""
self.acl.remove_all_rules()
@@ -435,9 +437,9 @@ class Primaite(Env):
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())
_LOGGER.debug("Link ID: " + link_value.get_id())
for protocol in link_value.protocol_list:
print(
_LOGGER.debug(
" Protocol: "
+ protocol.get_name().name
+ ", Load: "
@@ -457,11 +459,11 @@ class Primaite(Env):
elif self.training_config.action_type == ActionType.ACL:
self.apply_actions_to_acl(_action)
elif (
len(self.action_dict[_action]) == 6
len(self.action_dict[_action]) == 6
): # ACL actions in multidiscrete form have len 6
self.apply_actions_to_acl(_action)
elif (
len(self.action_dict[_action]) == 4
len(self.action_dict[_action]) == 4
): # Node actions in multdiscrete (array) from have len 4
self.apply_actions_to_nodes(_action)
else:
@@ -529,7 +531,8 @@ class Primaite(Env):
elif property_action == 1:
# Patch (valid action if it's good or compromised)
node.set_service_state(
self.services_list[service_index], SoftwareState.PATCHING
self.services_list[service_index],
SoftwareState.PATCHING
)
else:
# Node is not of Service Type
@@ -589,7 +592,8 @@ class Primaite(Env):
acl_rule_source = "ANY"
else:
node = list(self.nodes.values())[action_source_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
if isinstance(node, ServiceNode) or isinstance(node,
ActiveNode):
acl_rule_source = node.ip_address
else:
return
@@ -598,7 +602,8 @@ class Primaite(Env):
acl_rule_destination = "ANY"
else:
node = list(self.nodes.values())[action_destination_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
if isinstance(node, ServiceNode) or isinstance(node,
ActiveNode):
acl_rule_destination = node.ip_address
else:
return
@@ -683,7 +688,8 @@ class Primaite(Env):
:return: The observation space, initial observation (zeroed out array with the correct shape)
:rtype: Tuple[spaces.Space, np.ndarray]
"""
self.obs_handler = ObservationsHandler.from_config(self, self.obs_config)
self.obs_handler = ObservationsHandler.from_config(self,
self.obs_config)
return self.obs_handler.space, self.obs_handler.current_observation
@@ -727,8 +733,7 @@ class Primaite(Env):
_LOGGER.error(f"Invalid item_type: {item_type}")
pass
_LOGGER.info("Environment configuration loaded")
print("Environment configuration loaded")
_LOGGER.debug("Environment configuration loaded")
def create_node(self, item):
"""
@@ -791,7 +796,8 @@ class Primaite(Env):
service_protocol = service["name"]
service_port = service["port"]
service_state = SoftwareState[service["state"]]
node.add_service(Service(service_protocol, service_port, service_state))
node.add_service(
Service(service_protocol, service_port, service_state))
else:
# Bad formatting
pass
@@ -844,7 +850,8 @@ class Primaite(Env):
dest_node_ref: Node = self.nodes_reference[link_destination]
# Add link to network (reference)
self.network_reference.add_edge(source_node_ref, dest_node_ref, id=link_name)
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(
@@ -1119,7 +1126,8 @@ class Primaite(Env):
# All nodes have these parameters
node_id = item["node_id"]
node_class = item["node_class"]
node_hardware_state: HardwareState = HardwareState[item["hardware_state"]]
node_hardware_state: HardwareState = HardwareState[
item["hardware_state"]]
node: NodeUnion = self.nodes[node_id]
node_ref = self.nodes_reference[node_id]
@@ -1185,7 +1193,8 @@ class Primaite(Env):
# Use MAX to ensure we get them all
for node_action in range(4):
for service_state in range(self.num_services):
action = [node, node_property, node_action, service_state]
action = [node, node_property, node_action,
service_state]
# check to see if it's a nothing action (has no effect)
if is_valid_node_action(action):
actions[action_key] = action