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

#917 - Integrated the PrimaiteSession into all tests.

Browse Source 
- Ran a full pre-commit hook and thus encountered tons of fixes required
This commit is contained in:
Chris McCarthy
2023-06-30 09:08:13 +01:00
parent 7f912df383
commit 73015802ec
62 changed files with 1880 additions and 802 deletions
Download Patch File Download Diff File Expand all files Collapse all files

190
src/primaite/agents/agent.py
View File

@@ -1,27 +1,31 @@
from __future__ import annotations
import json
import time
from abc import ABC, abstractmethod
from datetime import datetime
from pathlib import Path
from typing import Optional, Final, Dict, Union
from typing import Dict, Final, Optional, Union
from uuid import uuid4
import yaml
import primaite
from primaite import getLogger, SESSIONS_DIR
from primaite.config import lay_down_config
from primaite.config import training_config
from primaite.config import lay_down_config, training_config
from primaite.config.training_config import TrainingConfig
from primaite.data_viz.session_plots import plot_av_reward_per_episode
from primaite.environment.primaite_env import Primaite
_LOGGER = getLogger(__name__)
def _get_session_path(session_timestamp: datetime) -> Path:
def get_session_path(session_timestamp: datetime) -> Path:
"""
Get a temp directory session path the test session will output to.
Get the directory path the session will output to.
This is set in the format of:
~/primaite/sessions/<yyyy-mm-dd>/<yyyy-mm-dd>_<hh-mm-ss>.
:param session_timestamp: This is the datetime that the session started.
:return: The session directory path.
@@ -35,13 +39,15 @@ def _get_session_path(session_timestamp: datetime) -> Path:
class AgentSessionABC(ABC):
"""
An ABC that manages training and/or evaluation of agents in PrimAITE.
This class cannot be directly instantiated and must be inherited from
with all implemented abstract methods implemented.
"""
@abstractmethod
def __init__(
self,
training_config_path,
lay_down_config_path
):
def __init__(self, training_config_path, lay_down_config_path):
if not isinstance(training_config_path, Path):
training_config_path = Path(training_config_path)
self._training_config_path: Final[Union[Path]] = training_config_path
@@ -66,9 +72,8 @@ class AgentSessionABC(ABC):
self._uuid = str(uuid4())
self.session_timestamp: datetime = datetime.now()
"The session timestamp"
self.session_path = _get_session_path(self.session_timestamp)
self.session_path = get_session_path(self.session_timestamp)
"The Session path"
self.checkpoints_path.mkdir(parents=True, exist_ok=True)
@property
def timestamp_str(self) -> str:
@@ -78,17 +83,23 @@ class AgentSessionABC(ABC):
@property
def learning_path(self) -> Path:
"""The learning outputs path."""
return self.session_path / "learning"
path = self.session_path / "learning"
path.mkdir(exist_ok=True, parents=True)
return path
@property
def evaluation_path(self) -> Path:
"""The evaluation outputs path."""
return self.session_path / "evaluation"
path = self.session_path / "evaluation"
path.mkdir(exist_ok=True, parents=True)
return path
@property
def checkpoints_path(self) -> Path:
"""The Session checkpoints path."""
return self.learning_path / "checkpoints"
path = self.learning_path / "checkpoints"
path.mkdir(exist_ok=True, parents=True)
return path
@property
def uuid(self):
@@ -118,14 +129,8 @@ class AgentSessionABC(ABC):
"uuid": self.uuid,
"start_datetime": self.session_timestamp.isoformat(),
"end_datetime": None,
"learning": {
"total_episodes": None,
"total_time_steps": None
},
"evaluation": {
"total_episodes": None,
"total_time_steps": None
},
"learning": {"total_episodes": None, "total_time_steps": None},
"evaluation": {"total_episodes": None, "total_time_steps": None},
"env": {
"training_config": self._training_config.to_dict(
json_serializable=True
@@ -156,11 +161,19 @@ class AgentSessionABC(ABC):
metadata_dict["end_datetime"] = datetime.now().isoformat()
if not self.is_eval:
metadata_dict["learning"]["total_episodes"] = self._env.episode_count # noqa
metadata_dict["learning"]["total_time_steps"] = self._env.total_step_count # noqa
metadata_dict["learning"][
"total_episodes"
] = self._env.episode_count # noqa
metadata_dict["learning"][
"total_time_steps"
] = self._env.total_step_count # noqa
else:
metadata_dict["evaluation"]["total_episodes"] = self._env.episode_count # noqa
metadata_dict["evaluation"]["total_time_steps"] = self._env.total_step_count # noqa
metadata_dict["evaluation"][
"total_episodes"
] = self._env.episode_count # noqa
metadata_dict["evaluation"][
"total_time_steps"
] = self._env.total_step_count # noqa
filepath = self.session_path / "session_metadata.json"
_LOGGER.debug(f"Updating Session Metadata file: {filepath}")
@@ -187,26 +200,47 @@ class AgentSessionABC(ABC):
@abstractmethod
def learn(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs,
):
"""
Train the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param kwargs: Any agent-specific key-word args to be passed.
"""
if self._can_learn:
_LOGGER.info("Finished learning")
_LOGGER.debug("Writing transactions")
self._update_session_metadata_file()
self._can_evaluate = True
self.is_eval = False
self._plot_av_reward_per_episode(learning_session=True)
@abstractmethod
def evaluate(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs,
):
"""
Evaluate the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param kwargs: Any agent-specific key-word args to be passed.
"""
self._env.set_as_eval() # noqa
self.is_eval = True
self._plot_av_reward_per_episode(learning_session=False)
_LOGGER.info("Finished evaluation")
@abstractmethod
@@ -216,6 +250,7 @@ class AgentSessionABC(ABC):
@classmethod
@abstractmethod
def load(cls, path: Union[str, Path]) -> AgentSessionABC:
"""Load an agent from file."""
if not isinstance(path, Path):
path = Path(path)
@@ -246,21 +281,56 @@ class AgentSessionABC(ABC):
else:
# Session path does not exist
msg = f"Failed to load PrimAITE Session, path does not exist: {path}"
msg = (
f"Failed to load PrimAITE Session, path does not exist: {path}"
)
_LOGGER.error(msg)
raise FileNotFoundError(msg)
pass
@abstractmethod
def save(self):
"""Save the agent."""
self._agent.save(self.session_path)
@abstractmethod
def export(self):
"""Export the agent to transportable file format."""
pass
def close(self):
"""Closes the agent."""
self._env.episode_av_reward_writer.close() # noqa
self._env.transaction_writer.close() # noqa
def _plot_av_reward_per_episode(self, learning_session: bool = True):
# self.close()
title = f"PrimAITE Session {self.timestamp_str} "
subtitle = str(self._training_config)
csv_file = f"average_reward_per_episode_{self.timestamp_str}.csv"
image_file = f"average_reward_per_episode_{self.timestamp_str}.png"
if learning_session:
title += "(Learning)"
path = self.learning_path / csv_file
image_path = self.learning_path / image_file
else:
title += "(Evaluation)"
path = self.evaluation_path / csv_file
image_path = self.evaluation_path / image_file
fig = plot_av_reward_per_episode(path, title, subtitle)
fig.write_image(image_path)
_LOGGER.debug(f"Saved average rewards per episode plot to: {path}")
class HardCodedAgentSessionABC(AgentSessionABC):
"""
An Agent Session ABC for evaluation deterministic agents.
This class cannot be directly instantiated and must be inherited from
with all implemented abstract methods implemented.
"""
def __init__(self, training_config_path, lay_down_config_path):
super().__init__(training_config_path, lay_down_config_path)
self._setup()
@@ -270,13 +340,12 @@ class HardCodedAgentSessionABC(AgentSessionABC):
training_config_path=self._training_config_path,
lay_down_config_path=self._lay_down_config_path,
session_path=self.session_path,
timestamp_str=self.timestamp_str
timestamp_str=self.timestamp_str,
)
super()._setup()
self._can_learn = False
self._can_evaluate = True
def _save_checkpoint(self):
pass
@@ -284,11 +353,20 @@ class HardCodedAgentSessionABC(AgentSessionABC):
pass
def learn(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs,
):
"""
Train the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param kwargs: Any agent-specific key-word args to be passed.
"""
_LOGGER.warning("Deterministic agents cannot learn")
@abstractmethod
@@ -296,20 +374,31 @@ class HardCodedAgentSessionABC(AgentSessionABC):
pass
def evaluate(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs,
):
"""
Evaluate the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param kwargs: Any agent-specific key-word args to be passed.
"""
self._env.set_as_eval() # noqa
self.is_eval = True
if not time_steps:
time_steps = self._training_config.num_steps
if not episodes:
episodes = self._training_config.num_episodes
obs = self._env.reset()
for episode in range(episodes):
# Reset env and collect initial observation
obs = self._env.reset()
for step in range(time_steps):
# Calculate action
action = self._calculate_action(obs)
@@ -322,15 +411,18 @@ class HardCodedAgentSessionABC(AgentSessionABC):
# Introduce a delay between steps
time.sleep(self._training_config.time_delay / 1000)
obs = self._env.reset()
self._env.close()
super().evaluate()
@classmethod
def load(cls):
"""Load an agent from file."""
_LOGGER.warning("Deterministic agents cannot be loaded")
def save(self):
"""Save the agent."""
_LOGGER.warning("Deterministic agents cannot be saved")
def export(self):
"""Export the agent to transportable file format."""
_LOGGER.warning("Deterministic agents cannot be exported")

231
src/primaite/agents/hardcoded_acl.py
View File

@@ -11,9 +11,13 @@ from primaite.common.enums import HardCodedAgentView
class HardCodedACLAgent(HardCodedAgentSessionABC):
"""An Agent Session class that implements a deterministic ACL agent."""
def _calculate_action(self, obs):
if self._training_config.hard_coded_agent_view == HardCodedAgentView.BASIC:
if (
self._training_config.hard_coded_agent_view
== HardCodedAgentView.BASIC
):
# Basic view action using only the current observation
return self._calculate_action_basic_view(obs)
else:
@@ -22,6 +26,12 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return self._calculate_action_full_view(obs)
def get_blocked_green_iers(self, green_iers, acl, nodes):
"""
Get blocked green IERs.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
blocked_green_iers = {}
for green_ier_id, green_ier in green_iers.items():
@@ -33,8 +43,9 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
port = green_ier.get_port()
# Can be blocked by an ACL or by default (no allow rule exists)
if acl.is_blocked(source_node_address, dest_node_address, protocol,
port):
if acl.is_blocked(
source_node_address, dest_node_address, protocol, port
):
blocked_green_iers[green_ier_id] = green_ier
return blocked_green_iers
@@ -42,8 +53,10 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
def get_matching_acl_rules_for_ier(self, ier, acl, nodes):
"""
Get matching ACL rules for an IER.
"""
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
source_node_id = ier.get_source_node_id()
source_node_address = nodes[source_node_id].ip_address
dest_node_id = ier.get_dest_node_id()
@@ -51,17 +64,22 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
protocol = ier.get_protocol() # e.g. 'TCP'
port = ier.get_port()
matching_rules = acl.get_relevant_rules(source_node_address,
dest_node_address, protocol,
port)
matching_rules = acl.get_relevant_rules(
source_node_address, dest_node_address, protocol, port
)
return matching_rules
def get_blocking_acl_rules_for_ier(self, ier, acl, nodes):
"""
Get blocking ACL rules for an IER.
Warning: Can return empty dict but IER can still be blocked by default (No ALLOW rule, therefore blocked)
"""
.. warning::
Can return empty dict but IER can still be blocked by default
(No ALLOW rule, therefore blocked).
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
matching_rules = self.get_matching_acl_rules_for_ier(ier, acl, nodes)
blocked_rules = {}
@@ -74,8 +92,10 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
def get_allow_acl_rules_for_ier(self, ier, acl, nodes):
"""
Get all allowing ACL rules for an IER.
"""
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
matching_rules = self.get_matching_acl_rules_for_ier(ier, acl, nodes)
allowed_rules = {}
@@ -85,9 +105,22 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return allowed_rules
def get_matching_acl_rules(self, source_node_id, dest_node_id, protocol,
port, acl,
nodes, services_list):
def get_matching_acl_rules(
self,
source_node_id,
dest_node_id,
protocol,
port,
acl,
nodes,
services_list,
):
"""
Get matching ACL rules.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
if source_node_id != "ANY":
source_node_address = nodes[str(source_node_id)].ip_address
else:
@@ -100,21 +133,39 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
if protocol != "ANY":
protocol = services_list[
protocol - 1] # -1 as dont have to account for ANY in list of services
protocol - 1
] # -1 as dont have to account for ANY in list of services
matching_rules = acl.get_relevant_rules(source_node_address,
dest_node_address, protocol,
port)
matching_rules = acl.get_relevant_rules(
source_node_address, dest_node_address, protocol, port
)
return matching_rules
def get_allow_acl_rules(self, source_node_id, dest_node_id, protocol,
port, acl,
nodes, services_list):
matching_rules = self.get_matching_acl_rules(source_node_id,
dest_node_id,
protocol, port, acl,
nodes,
services_list)
def get_allow_acl_rules(
self,
source_node_id,
dest_node_id,
protocol,
port,
acl,
nodes,
services_list,
):
"""
Get the ALLOW ACL rules.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
matching_rules = self.get_matching_acl_rules(
source_node_id,
dest_node_id,
protocol,
port,
acl,
nodes,
services_list,
)
allowed_rules = {}
for rule_key, rule_value in matching_rules.items():
@@ -123,14 +174,31 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return allowed_rules
def get_deny_acl_rules(self, source_node_id, dest_node_id, protocol, port,
acl,
nodes, services_list):
matching_rules = self.get_matching_acl_rules(source_node_id,
dest_node_id,
protocol, port, acl,
nodes,
services_list)
def get_deny_acl_rules(
self,
source_node_id,
dest_node_id,
protocol,
port,
acl,
nodes,
services_list,
):
"""
Get the DENY ACL rules.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
matching_rules = self.get_matching_acl_rules(
source_node_id,
dest_node_id,
protocol,
port,
acl,
nodes,
services_list,
)
allowed_rules = {}
for rule_key, rule_value in matching_rules.items():
@@ -141,7 +209,7 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
def _calculate_action_full_view(self, obs):
"""
Given an observation and the environment calculate a good acl-based action for the blue agent to take
Calculate a good acl-based action for the blue agent to take.
Knowledge of just the observation space is insufficient for a perfect solution, as we need to know:
@@ -167,8 +235,10 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
nodes once a service becomes overwhelmed. However currently the ACL action space has no way of reversing
an overwhelmed state, so we don't do this.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
#obs = convert_to_old_obs(obs)
# obs = convert_to_old_obs(obs)
r_obs = transform_change_obs_readable(obs)
_, _, _, *s = r_obs
@@ -184,7 +254,6 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
for service_num, service_states in enumerate(s):
for x, service_state in enumerate(service_states):
if service_state == "COMPROMISED":
action_source_id = x + 1 # +1 as 0 is any
action_destination_id = "ANY"
action_protocol = service_num + 1 # +1 as 0 is any
@@ -215,19 +284,23 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
action_permission = "ALLOW"
action_source_ip = rule.get_source_ip()
action_source_id = int(
get_node_of_ip(action_source_ip, self._env.nodes))
get_node_of_ip(action_source_ip, self._env.nodes)
)
action_destination_ip = rule.get_dest_ip()
action_destination_id = int(
get_node_of_ip(action_destination_ip,
self._env.nodes))
get_node_of_ip(
action_destination_ip, self._env.nodes
)
)
action_protocol_name = rule.get_protocol()
action_protocol = (
self._env.services_list.index(
action_protocol_name) + 1
self._env.services_list.index(action_protocol_name)
+ 1
) # convert name e.g. 'TCP' to index
action_port_name = rule.get_port()
action_port = self._env.ports_list.index(
action_port_name) + 1 # convert port name e.g. '80' to index
action_port = (
self._env.ports_list.index(action_port_name) + 1
) # convert port name e.g. '80' to index
found_action = True
break
@@ -258,21 +331,21 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
if not found_action:
# Which Green IERS are blocked
blocked_green_iers = self.get_blocked_green_iers(
self._env.green_iers, self._env.acl,
self._env.nodes)
self._env.green_iers, self._env.acl, self._env.nodes
)
for ier_key, ier in blocked_green_iers.items():
# Which ALLOW rules are allowing this IER (none)
allowing_rules = self.get_allow_acl_rules_for_ier(ier,
self._env.acl,
self._env.nodes)
allowing_rules = self.get_allow_acl_rules_for_ier(
ier, self._env.acl, self._env.nodes
)
# If there are no blocking rules, it may be being blocked by default
# If there is already an allow rule
node_id_to_check = int(ier.get_source_node_id())
service_name_to_check = ier.get_protocol()
service_id_to_check = self._env.services_list.index(
service_name_to_check)
service_name_to_check
)
# Service state of the the source node in the ier
service_state = s[service_id_to_check][node_id_to_check - 1]
@@ -283,11 +356,13 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
action_source_id = int(ier.get_source_node_id())
action_destination_id = int(ier.get_dest_node_id())
action_protocol_name = ier.get_protocol()
action_protocol = self._env.services_list.index(
action_protocol_name) + 1 # convert name e.g. 'TCP' to index
action_protocol = (
self._env.services_list.index(action_protocol_name) + 1
) # convert name e.g. 'TCP' to index
action_port_name = ier.get_port()
action_port = self._env.ports_list.index(
action_port_name) + 1 # convert port name e.g. '80' to index
action_port = (
self._env.ports_list.index(action_port_name) + 1
) # convert port name e.g. '80' to index
found_action = True
break
@@ -311,19 +386,25 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return action
def _calculate_action_basic_view(self, obs):
"""Given an observation calculate a good acl-based action for the blue agent to take
"""Calculate a good acl-based action for the blue agent to take.
Uses ONLY information from the current observation with NO knowledge of previous actions taken and
NO reward feedback.
Uses ONLY information from the current observation with NO knowledge
of previous actions taken and NO reward feedback.
We rely on randomness to select the precise action, as we want to block all traffic originating from
a compromised node, without being able to tell:
We rely on randomness to select the precise action, as we want to
block all traffic originating from a compromised node, without being
able to tell:
1. Which ACL rules already exist
1. Which actions the agent has already tried.
2. Which actions the agent has already tried.
There is a high probability that the correct rule will not be deleted before the state becomes overwhelmed.
There is a high probability that the correct rule will not be deleted
before the state becomes overwhelmed.
Currently a deny rule does not overwrite an allow rule. The allow rules must be deleted.
Currently, a deny rule does not overwrite an allow rule. The allow
rules must be deleted.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
action_dict = self._env.action_dict
r_obs = transform_change_obs_readable(obs)
@@ -333,27 +414,35 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
s = [*s]
number_of_nodes = len(
[i for i in o if i != "NONE"]) # number of nodes (not links)
[i for i in o if i != "NONE"]
) # number of nodes (not links)
for service_num, service_states in enumerate(s):
comprimised_states = [n for n, i in enumerate(service_states) if
i == "COMPROMISED"]
comprimised_states = [
n for n, i in enumerate(service_states) if i == "COMPROMISED"
]
if len(comprimised_states) == 0:
# No states are COMPROMISED, try the next service
continue
compromised_node = np.random.choice(
comprimised_states) + 1 # +1 as 0 would be any
compromised_node = (
np.random.choice(comprimised_states) + 1
) # +1 as 0 would be any
action_decision = "DELETE"
action_permission = "ALLOW"
action_source_ip = compromised_node
# Randomly select a destination ID to block
action_destination_ip = np.random.choice(
list(range(1, number_of_nodes + 1)) + ["ANY"])
action_destination_ip = int(
action_destination_ip) if action_destination_ip != "ANY" else action_destination_ip
list(range(1, number_of_nodes + 1)) + ["ANY"]
)
action_destination_ip = (
int(action_destination_ip)
if action_destination_ip != "ANY"
else action_destination_ip
)
action_protocol = service_num + 1 # +1 as 0 is any
# Randomly select a port to block
# Bad assumption that number of protocols equals number of ports AND no rules exist with an ANY port
# Bad assumption that number of protocols equals number of ports
# AND no rules exist with an ANY port
action_port = np.random.choice(list(range(1, len(s) + 1)))
action = [

64
src/primaite/agents/hardcoded_node.py
View File

@@ -1,16 +1,21 @@
from primaite.agents.agent import HardCodedAgentSessionABC
from primaite.agents.utils import (
get_new_action,
transform_change_obs_readable,
)
from primaite.agents.utils import (
transform_action_node_enum,
transform_change_obs_readable,
)
class HardCodedNodeAgent(HardCodedAgentSessionABC):
"""An Agent Session class that implements a deterministic Node agent."""
def _calculate_action(self, obs):
"""Given an observation calculate a good node-based action for the blue agent to take"""
"""
Calculate a good node-based action for the blue agent to take.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
"""
action_dict = self._env.action_dict
r_obs = transform_change_obs_readable(obs)
_, o, os, *s = r_obs
@@ -18,7 +23,8 @@ class HardCodedNodeAgent(HardCodedAgentSessionABC):
if len(r_obs) == 4: # only 1 service
s = [*s]
# Check in order of most important states (order doesn't currently matter, but it probably should)
# Check in order of most important states (order doesn't currently
# matter, but it probably should)
# First see if any OS states are compromised
for x, os_state in enumerate(os):
if os_state == "COMPROMISED":
@@ -26,8 +32,12 @@ class HardCodedNodeAgent(HardCodedAgentSessionABC):
action_node_property = "OS"
property_action = "PATCHING"
action_service_index = 0 # does nothing isn't relevant for os
action = [action_node_id, action_node_property,
property_action, action_service_index]
action = [
action_node_id,
action_node_property,
property_action,
action_service_index,
]
action = transform_action_node_enum(action)
action = get_new_action(action, action_dict)
# We can only perform 1 action on each step
@@ -44,8 +54,12 @@ class HardCodedNodeAgent(HardCodedAgentSessionABC):
property_action = "PATCHING"
action_service_index = service_num
action = [action_node_id, action_node_property,
property_action, action_service_index]
action = [
action_node_id,
action_node_property,
property_action,
action_service_index,
]
action = transform_action_node_enum(action)
action = get_new_action(action, action_dict)
# We can only perform 1 action on each step
@@ -63,8 +77,12 @@ class HardCodedNodeAgent(HardCodedAgentSessionABC):
property_action = "PATCHING"
action_service_index = service_num
action = [action_node_id, action_node_property,
property_action, action_service_index]
action = [
action_node_id,
action_node_property,
property_action,
action_service_index,
]
action = transform_action_node_enum(action)
action = get_new_action(action, action_dict)
# We can only perform 1 action on each step
@@ -75,10 +93,18 @@ class HardCodedNodeAgent(HardCodedAgentSessionABC):
if os_state == "OFF":
action_node_id = x + 1
action_node_property = "OPERATING"
property_action = "ON" # Why reset it when we can just turn it on
action_service_index = 0 # does nothing isn't relevant for operating state
action = [action_node_id, action_node_property,
property_action, action_service_index]
property_action = (
"ON" # Why reset it when we can just turn it on
)
action_service_index = (
0 # does nothing isn't relevant for operating state
)
action = [
action_node_id,
action_node_property,
property_action,
action_service_index,
]
action = transform_action_node_enum(action, action_dict)
action = get_new_action(action, action_dict)
# We can only perform 1 action on each step
@@ -89,8 +115,12 @@ class HardCodedNodeAgent(HardCodedAgentSessionABC):
action_node_property = "NONE"
property_action = "ON"
action_service_index = 0
action = [action_node_id, action_node_property, property_action,
action_service_index]
action = [
action_node_id,
action_node_property,
property_action,
action_service_index,
]
action = transform_action_node_enum(action)
action = get_new_action(action, action_dict)

115
src/primaite/agents/rllib.py
View File

@@ -1,28 +1,35 @@
from __future__ import annotations
import json
from datetime import datetime
from pathlib import Path
from typing import Optional
from typing import Optional, Union
import tensorflow as tf
from ray.rllib.algorithms import Algorithm
from ray.rllib.algorithms.a2c import A2CConfig
from ray.rllib.algorithms.ppo import PPOConfig
from ray.tune.logger import UnifiedLogger
from ray.tune.registry import register_env
import tensorflow as tf
from primaite import getLogger
from primaite.agents.agent import AgentSessionABC
from primaite.common.enums import AgentFramework, AgentIdentifier, \
DeepLearningFramework
from primaite.common.enums import (
AgentFramework,
AgentIdentifier,
DeepLearningFramework,
)
from primaite.environment.primaite_env import Primaite
_LOGGER = getLogger(__name__)
def _env_creator(env_config):
return Primaite(
training_config_path=env_config["training_config_path"],
lay_down_config_path=env_config["lay_down_config_path"],
session_path=env_config["session_path"],
timestamp_str=env_config["timestamp_str"]
timestamp_str=env_config["timestamp_str"],
)
@@ -37,16 +44,15 @@ def _custom_log_creator(session_path: Path):
class RLlibAgent(AgentSessionABC):
"""An AgentSession class that implements a Ray RLlib agent."""
def __init__(
self,
training_config_path,
lay_down_config_path
):
def __init__(self, training_config_path, lay_down_config_path):
super().__init__(training_config_path, lay_down_config_path)
if not self._training_config.agent_framework == AgentFramework.RLLIB:
msg = (f"Expected RLLIB agent_framework, "
f"got {self._training_config.agent_framework}")
msg = (
f"Expected RLLIB agent_framework, "
f"got {self._training_config.agent_framework}"
)
_LOGGER.error(msg)
raise ValueError(msg)
if self._training_config.agent_identifier == AgentIdentifier.PPO:
@@ -54,8 +60,10 @@ class RLlibAgent(AgentSessionABC):
elif self._training_config.agent_identifier == AgentIdentifier.A2C:
self._agent_config_class = A2CConfig
else:
msg = ("Expected PPO or A2C agent_identifier, "
f"got {self._training_config.agent_identifier.value}")
msg = (
"Expected PPO or A2C agent_identifier, "
f"got {self._training_config.agent_identifier.value}"
)
_LOGGER.error(msg)
raise ValueError(msg)
self._agent_config: PPOConfig
@@ -86,8 +94,12 @@ class RLlibAgent(AgentSessionABC):
metadata_dict = json.load(file)
metadata_dict["end_datetime"] = datetime.now().isoformat()
metadata_dict["total_episodes"] = self._current_result["episodes_total"]
metadata_dict["total_time_steps"] = self._current_result["timesteps_total"]
metadata_dict["total_episodes"] = self._current_result[
"episodes_total"
]
metadata_dict["total_time_steps"] = self._current_result[
"timesteps_total"
]
filepath = self.session_path / "session_metadata.json"
_LOGGER.debug(f"Updating Session Metadata file: {filepath}")
@@ -106,43 +118,48 @@ class RLlibAgent(AgentSessionABC):
training_config_path=self._training_config_path,
lay_down_config_path=self._lay_down_config_path,
session_path=self.session_path,
timestamp_str=self.timestamp_str
)
timestamp_str=self.timestamp_str,
),
)
self._agent_config.training(
train_batch_size=self._training_config.num_steps
)
self._agent_config.framework(
framework="tf"
)
self._agent_config.framework(framework="tf")
self._agent_config.rollouts(
num_rollout_workers=1,
num_envs_per_worker=1,
horizon=self._training_config.num_steps
horizon=self._training_config.num_steps,
)
self._agent: Algorithm = self._agent_config.build(
logger_creator=_custom_log_creator(self.session_path)
)
def _save_checkpoint(self):
checkpoint_n = self._training_config.checkpoint_every_n_episodes
episode_count = self._current_result["episodes_total"]
if checkpoint_n > 0 and episode_count > 0:
if (
(episode_count % checkpoint_n == 0)
or (episode_count == self._training_config.num_episodes)
if (episode_count % checkpoint_n == 0) or (
episode_count == self._training_config.num_episodes
):
self._agent.save(self.checkpoints_path)
self._agent.save(str(self.checkpoints_path))
def learn(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs,
):
"""
Evaluate the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param kwargs: Any agent-specific key-word args to be passed.
"""
# Temporarily override train_batch_size and horizon
if time_steps:
self._agent_config.train_batch_size = time_steps
@@ -150,37 +167,53 @@ class RLlibAgent(AgentSessionABC):
if not episodes:
episodes = self._training_config.num_episodes
_LOGGER.info(f"Beginning learning for {episodes} episodes @"
f" {time_steps} time steps...")
_LOGGER.info(
f"Beginning learning for {episodes} episodes @"
f" {time_steps} time steps..."
)
for i in range(episodes):
self._current_result = self._agent.train()
self._save_checkpoint()
if self._training_config.deep_learning_framework != DeepLearningFramework.TORCH:
if (
self._training_config.deep_learning_framework
!= DeepLearningFramework.TORCH
):
policy = self._agent.get_policy()
tf.compat.v1.summary.FileWriter(
self.session_path / "ray_results",
policy.get_session().graph
self.session_path / "ray_results", policy.get_session().graph
)
super().learn()
self._agent.stop()
def evaluate(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs,
):
"""
Evaluate the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param kwargs: Any agent-specific key-word args to be passed.
"""
raise NotImplementedError
def _get_latest_checkpoint(self):
raise NotImplementedError
@classmethod
def load(cls):
def load(cls, path: Union[str, Path]) -> RLlibAgent:
"""Load an agent from file."""
raise NotImplementedError
def save(self):
"""Save the agent."""
raise NotImplementedError
def export(self):
"""Export the agent to transportable file format."""
raise NotImplementedError

103
src/primaite/agents/sb3.py
View File

@@ -1,27 +1,30 @@
from typing import Optional
from __future__ import annotations
from pathlib import Path
from typing import Optional, Union
import numpy as np
from stable_baselines3 import PPO, A2C
from stable_baselines3 import A2C, PPO
from stable_baselines3.ppo import MlpPolicy as PPOMlp
from primaite import getLogger
from primaite.agents.agent import AgentSessionABC
from primaite.common.enums import AgentIdentifier, AgentFramework
from primaite.common.enums import AgentFramework, AgentIdentifier
from primaite.environment.primaite_env import Primaite
_LOGGER = getLogger(__name__)
class SB3Agent(AgentSessionABC):
def __init__(
self,
training_config_path,
lay_down_config_path
):
"""An AgentSession class that implements a Stable Baselines3 agent."""
def __init__(self, training_config_path, lay_down_config_path):
super().__init__(training_config_path, lay_down_config_path)
if not self._training_config.agent_framework == AgentFramework.SB3:
msg = (f"Expected SB3 agent_framework, "
f"got {self._training_config.agent_framework}")
msg = (
f"Expected SB3 agent_framework, "
f"got {self._training_config.agent_framework}"
)
_LOGGER.error(msg)
raise ValueError(msg)
if self._training_config.agent_identifier == AgentIdentifier.PPO:
@@ -29,8 +32,10 @@ class SB3Agent(AgentSessionABC):
elif self._training_config.agent_identifier == AgentIdentifier.A2C:
self._agent_class = A2C
else:
msg = ("Expected PPO or A2C agent_identifier, "
f"got {self._training_config.agent_identifier.value}")
msg = (
"Expected PPO or A2C agent_identifier, "
f"got {self._training_config.agent_identifier}"
)
_LOGGER.error(msg)
raise ValueError(msg)
@@ -52,25 +57,26 @@ class SB3Agent(AgentSessionABC):
training_config_path=self._training_config_path,
lay_down_config_path=self._lay_down_config_path,
session_path=self.session_path,
timestamp_str=self.timestamp_str
timestamp_str=self.timestamp_str,
)
self._agent = self._agent_class(
PPOMlp,
self._env,
verbose=self.output_verbose_level,
n_steps=self._training_config.num_steps,
tensorboard_log=self._tensorboard_log_path
tensorboard_log=self._tensorboard_log_path,
)
def _save_checkpoint(self):
checkpoint_n = self._training_config.checkpoint_every_n_episodes
episode_count = self._env.episode_count
if checkpoint_n > 0 and episode_count > 0:
if (
(episode_count % checkpoint_n == 0)
or (episode_count == self._training_config.num_episodes)
if (episode_count % checkpoint_n == 0) or (
episode_count == self._training_config.num_episodes
):
checkpoint_path = self.checkpoints_path / f"sb3ppo_{episode_count}.zip"
checkpoint_path = (
self.checkpoints_path / f"sb3ppo_{episode_count}.zip"
)
self._agent.save(checkpoint_path)
_LOGGER.debug(f"Saved agent checkpoint: {checkpoint_path}")
@@ -78,33 +84,54 @@ class SB3Agent(AgentSessionABC):
pass
def learn(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs,
):
"""
Train the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param kwargs: Any agent-specific key-word args to be passed.
"""
if not time_steps:
time_steps = self._training_config.num_steps
if not episodes:
episodes = self._training_config.num_episodes
self.is_eval = False
_LOGGER.info(f"Beginning learning for {episodes} episodes @"
f" {time_steps} time steps...")
_LOGGER.info(
f"Beginning learning for {episodes} episodes @"
f" {time_steps} time steps..."
)
for i in range(episodes):
self._agent.learn(total_timesteps=time_steps)
self._save_checkpoint()
self._env.close()
self.close()
super().learn()
def evaluate(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
deterministic: bool = True,
**kwargs
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
deterministic: bool = True,
**kwargs,
):
"""
Evaluate the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param deterministic: Whether the evaluation is deterministic.
:param kwargs: Any agent-specific key-word args to be passed.
"""
if not time_steps:
time_steps = self._training_config.num_steps
@@ -116,27 +143,31 @@ class SB3Agent(AgentSessionABC):
deterministic_str = "deterministic"
else:
deterministic_str = "non-deterministic"
_LOGGER.info(f"Beginning {deterministic_str} evaluation for "
f"{episodes} episodes @ {time_steps} time steps...")
_LOGGER.info(
f"Beginning {deterministic_str} evaluation for "
f"{episodes} episodes @ {time_steps} time steps..."
)
for episode in range(episodes):
obs = self._env.reset()
for step in range(time_steps):
action, _states = self._agent.predict(
obs,
deterministic=deterministic
obs, deterministic=deterministic
)
if isinstance(action, np.ndarray):
action = np.int64(action)
obs, rewards, done, info = self._env.step(action)
_LOGGER.info(f"Finished evaluation")
super().evaluate()
@classmethod
def load(self):
def load(cls, path: Union[str, Path]) -> SB3Agent:
"""Load an agent from file."""
raise NotImplementedError
def save(self):
"""Save the agent."""
raise NotImplementedError
def export(self):
"""Export the agent to transportable file format."""
raise NotImplementedError

230
src/primaite/agents/utils.py
View File

@@ -4,9 +4,9 @@ from primaite.common.enums import (
HardwareState,
LinkStatus,
NodeHardwareAction,
NodePOLType,
NodeSoftwareAction,
SoftwareState,
NodePOLType
)
@@ -16,14 +16,17 @@ def transform_action_node_readable(action):
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:
elif (
action_node_property == "OS" or action_node_property == "SERVICE"
) and action[2] <= 1:
property_action = NodeSoftwareAction(action[2]).name
else:
property_action = "NONE"
@@ -38,6 +41,9 @@ def transform_action_acl_readable(action):
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"}
@@ -62,6 +68,9 @@ def is_valid_node_action(action):
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)
@@ -77,7 +86,10 @@ def is_valid_node_action(action):
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"]:
if node_property != "OPERATING" and node_action not in [
"NONE",
"PATCHING",
]:
# Software States can only do Nothing or Patch
return False
return True
@@ -92,6 +104,9 @@ def is_valid_acl_action(action):
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)
@@ -118,7 +133,12 @@ def is_valid_acl_action(action):
def is_valid_acl_action_extra(action):
"""Harsher version of valid acl actions, does not allow 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
@@ -136,13 +156,15 @@ def is_valid_acl_action_extra(action):
return True
def transform_change_obs_readable(obs):
"""Transform list of transactions to readable list of each observation property
"""
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]]
@@ -151,7 +173,9 @@ def transform_change_obs_readable(obs):
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]]
service_states = [
SoftwareState(i).name if i <= 4 else i for i in obs[:, service]
]
new_obs.append(service_states)
return new_obs
@@ -159,10 +183,13 @@ def transform_change_obs_readable(obs):
def transform_obs_readable(obs):
"""
example:
np.array([[1,2,1,3],[2,1,1,1]]) -> [[1, 'OFF', 'GOOD', 'COMPROMISED'], [2, 'ON', 'GOOD', 'GOOD']]
"""
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
@@ -172,7 +199,12 @@ def transform_obs_readable(obs):
def convert_to_new_obs(obs, num_nodes=10):
"""Convert original gym Box observation space to new multiDiscrete observation space"""
"""
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
@@ -180,7 +212,9 @@ def convert_to_new_obs(obs, num_nodes=10):
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
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])
@@ -190,13 +224,17 @@ def convert_to_old_obs(obs, num_nodes=10, num_links=10, num_services=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 = 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
@@ -209,14 +247,19 @@ def convert_to_old_obs(obs, num_nodes=10, num_links=10, num_services=1):
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
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)
@@ -236,20 +279,27 @@ def describe_obs_change(obs1, obs2, num_nodes=10, num_links=10, num_services=1):
def _describe_obs_change_helper(obs_change, is_link):
""" "
Helper funcion to describe what has changed
"""
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]
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
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 = [
@@ -263,8 +313,8 @@ def _describe_obs_change_helper(obs_change, is_link):
if not is_link:
desc = f"ID {obs_change[0]}:"
for NodePOLType, state in list(zip(NodePOLTypes, states)):
desc = desc + " " + NodePOLType + " changed to " + state + "."
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]}."
@@ -273,12 +323,14 @@ def _describe_obs_change_helper(obs_change, is_link):
def transform_action_node_enum(action):
"""
Convert a node action from readable string format, to enumerated format
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
@@ -291,24 +343,33 @@ def transform_action_node_enum(action):
action_service_index = action[3]
new_action = [action_node_id, action_node_property, property_action, action_service_index]
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
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:
elif (
action_node_property == "OS" or action_node_property == "SERVICE"
) and action[2] <= 1:
property_action = NodeSoftwareAction(action[2]).name
else:
property_action = "NONE"
@@ -319,9 +380,11 @@ def transform_action_node_readable(action):
def node_action_description(action):
"""
Generate string describing a node-based 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)
@@ -334,7 +397,9 @@ def node_action_description(action):
if property_action == "NONE":
return ""
if node_property == "OPERATING" or node_property == "OS":
description = f"NODE {node_id}, {node_property}, SET TO {property_action}"
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:
@@ -343,34 +408,13 @@ def node_action_description(action):
return description
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]
"""
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 transform_action_acl_enum(action):
"""
Convert a acl action from readable string format, to enumerated format
"""
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}
@@ -388,8 +432,12 @@ def transform_action_acl_enum(action):
def acl_action_description(action):
"""generate string describing a acl-based 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)
@@ -406,11 +454,13 @@ def acl_action_description(action):
def get_node_of_ip(ip, node_dict):
"""
Get the node ID of an IP address
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:
@@ -418,13 +468,16 @@ def get_node_of_ip(ip, node_dict):
def is_valid_node_action(action):
"""Is the node action an actual valid 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)
@@ -438,7 +491,10 @@ def is_valid_node_action(action):
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"]:
if node_property != "OPERATING" and node_action not in [
"NONE",
"PATCHING",
]:
# Software States can only do Nothing or Patch
return False
return True
@@ -446,13 +502,16 @@ def is_valid_node_action(action):
def is_valid_acl_action(action):
"""
Is the ACL action an actual valid 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)
@@ -463,7 +522,11 @@ def is_valid_acl_action(action):
if action_decision == "NONE":
return False
if action_source_id == action_destination_id and action_source_id != "ANY" and action_destination_id != "ANY":
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":
@@ -475,7 +538,12 @@ def is_valid_acl_action(action):
def is_valid_acl_action_extra(action):
"""Harsher version of valid acl actions, does not allow 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
@@ -494,33 +562,17 @@ def is_valid_acl_action_extra(action):
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
"""
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
def get_action_description(action, action_dict):
"""
Get a string describing/explaining what an action is doing in words
"""
action_array = action_dict[action]
if len(action_array) == 4:
# node actions have length 4
action_description = node_action_description(action_array)
elif len(action_array) == 6:
# acl actions have length 6
action_description = acl_action_description(action_array)
else:
# Should never happen
action_description = "Unrecognised action"
return action_description