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

Merge remote-tracking branch 'origin/dev' into feature/1566-configure_episode-steps-learn-eval

Browse Source 
# Conflicts:
#	src/primaite/config/training_config.py
This commit is contained in:
SunilSamra
2023-07-11 11:39:21 +01:00
parent 563ff72fd6 19a9cef130
commit 6b59ce960d
51 changed files with 577 additions and 894 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Access Control List. Models firewall functionality."""

50
src/primaite/acl/access_control_list.py
View File

@@ -9,20 +9,20 @@ class AccessControlList:
"""Access Control List class."""
def __init__(self):
"""Init."""
self.acl: Dict[str, AccessControlList] = {} # A dictionary of ACL Rules
"""Initialise an empty AccessControlList."""
self.acl: Dict[str, ACLRule] = {} # A dictionary of ACL Rules
def check_address_match(self, _rule, _source_ip_address, _dest_ip_address):
"""
Checks for IP address matches.
def check_address_match(self, _rule: ACLRule, _source_ip_address: str, _dest_ip_address: str) -> bool:
"""Checks for IP address matches.
Args:
_rule: The rule being checked
_source_ip_address: the source IP address to compare
_dest_ip_address: the destination IP address to compare
Returns:
True if match; False otherwise.
:param _rule: The rule object to check
:type _rule: ACLRule
:param _source_ip_address: Source IP address to compare
:type _source_ip_address: str
:param _dest_ip_address: Destination IP address to compare
:type _dest_ip_address: str
:return: True if there is a match, otherwise False.
:rtype: bool
"""
if (
(_rule.get_source_ip() == _source_ip_address and _rule.get_dest_ip() == _dest_ip_address)
@@ -34,7 +34,7 @@ class AccessControlList:
else:
return False
def is_blocked(self, _source_ip_address, _dest_ip_address, _protocol, _port):
def is_blocked(self, _source_ip_address: str, _dest_ip_address: str, _protocol: str, _port: str) -> bool:
"""
Checks for rules that block a protocol / port.
@@ -116,3 +116,27 @@ class AccessControlList:
rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
hash_value = hash(rule)
return hash_value
def get_relevant_rules(self, _source_ip_address, _dest_ip_address, _protocol, _port):
"""Get all ACL rules that relate to the given arguments.
:param _source_ip_address: the source IP address to check
:param _dest_ip_address: the destination IP address to check
:param _protocol: the protocol to check
:param _port: the port to check
:return: Dictionary of all ACL rules that relate to the given arguments
:rtype: Dict[str, ACLRule]
"""
relevant_rules = {}
for rule_key, rule_value in self.acl.items():
if self.check_address_match(rule_value, _source_ip_address, _dest_ip_address):
if (
rule_value.get_protocol() == _protocol or rule_value.get_protocol() == "ANY" or _protocol == "ANY"
) and (
str(rule_value.get_port()) == str(_port) or rule_value.get_port() == "ANY" or str(_port) == "ANY"
):
# There's a matching rule.
relevant_rules[rule_key] = rule_value
return relevant_rules

13
src/primaite/acl/acl_rule.py
View File

@@ -7,14 +7,13 @@ class ACLRule:
def __init__(self, _permission, _source_ip, _dest_ip, _protocol, _port):
"""
Init.
Initialise an ACL Rule.
Args:
_permission: The permission (ALLOW or DENY)
_source_ip: The source IP address
_dest_ip: The destination IP address
_protocol: The rule protocol
_port: The rule port
:param _permission: The permission (ALLOW or DENY)
:param _source_ip: The source IP address
:param _dest_ip: The destination IP address
:param _protocol: The rule protocol
:param _port: The rule port
"""
self.permission = _permission
self.source_ip = _source_ip

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

@@ -0,0 +1 @@
"""Common interface between RL agents from different libraries and PrimAITE."""

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

@@ -42,12 +42,21 @@ 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.
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):
"""
Initialise an agent session from config files.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
"""
if not isinstance(training_config_path, Path):
training_config_path = Path(training_config_path)
self._training_config_path: Final[Union[Path, str]] = training_config_path
@@ -55,7 +64,7 @@ class AgentSessionABC(ABC):
if not isinstance(lay_down_config_path, Path):
lay_down_config_path = Path(lay_down_config_path)
self._lay_down_config_path: Final[Union[Path]] = lay_down_config_path
self._lay_down_config_path: Final[Union[Path, str]] = lay_down_config_path
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path)
self.sb3_output_verbose_level = self._training_config.sb3_output_verbose_level
@@ -305,11 +314,20 @@ 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.
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):
"""
Initialise a hardcoded agent session.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
"""
super().__init__(training_config_path, lay_down_config_path)
self._setup()

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

@@ -1,5 +1,9 @@
from typing import Any, Dict, List, Union
import numpy as np
from primaite.acl.access_control_list import AccessControlList
from primaite.acl.acl_rule import ACLRule
from primaite.agents.agent import HardCodedAgentSessionABC
from primaite.agents.utils import (
get_new_action,
@@ -7,13 +11,17 @@ from primaite.agents.utils import (
transform_action_acl_enum,
transform_change_obs_readable,
)
from primaite.common.custom_typing import NodeUnion
from primaite.common.enums import HardCodedAgentView
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.service_node import ServiceNode
from primaite.pol.ier import IER
class HardCodedACLAgent(HardCodedAgentSessionABC):
"""An Agent Session class that implements a deterministic ACL agent."""
def _calculate_action(self, obs):
def _calculate_action(self, obs: np.ndarray) -> int:
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)
@@ -22,12 +30,19 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
# history and reward feedback
return self._calculate_action_full_view(obs)
def get_blocked_green_iers(self, green_iers, acl, nodes):
"""
Get blocked green IERs.
def get_blocked_green_iers(
self, green_iers: Dict[str, IER], acl: AccessControlList, nodes: Dict[str, NodeUnion]
) -> Dict[Any, Any]:
"""Get blocked green IERs.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param green_iers: Green IERs to check for being
:type green_iers: Dict[str, IER]
:param acl: Firewall rules
:type acl: AccessControlList
:param nodes: Nodes in the network
:type nodes: Dict[str,NodeUnion]
:return: Same as `green_iers` input dict, but filtered to only contain the blocked ones.
:rtype: Dict[str, IER]
"""
blocked_green_iers = {}
@@ -45,12 +60,17 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return blocked_green_iers
def get_matching_acl_rules_for_ier(self, ier, acl, nodes):
"""
Get matching ACL rules for an IER.
def get_matching_acl_rules_for_ier(self, ier: IER, acl: AccessControlList, nodes: Dict[str, NodeUnion]):
"""Get list of ACL rules which are relevant to an IER.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param ier: Information Exchange Request to query against the ACL list
:type ier: IER
:param acl: Firewall rules
:type acl: AccessControlList
:param nodes: Nodes in the network
:type nodes: Dict[str,NodeUnion]
:return: _description_
:rtype: _type_
"""
source_node_id = ier.get_source_node_id()
source_node_address = nodes[source_node_id].ip_address
@@ -58,11 +78,12 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
dest_node_address = nodes[dest_node_id].ip_address
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)
return matching_rules
def get_blocking_acl_rules_for_ier(self, ier, acl, nodes):
def get_blocking_acl_rules_for_ier(
self, ier: IER, acl: AccessControlList, nodes: Dict[str, NodeUnion]
) -> Dict[str, Any]:
"""
Get blocking ACL rules for an IER.
@@ -70,8 +91,14 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
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.
:param ier: Information Exchange Request to query against the ACL list
:type ier: IER
:param acl: Firewall rules
:type acl: AccessControlList
:param nodes: Nodes in the network
:type nodes: Dict[str,NodeUnion]
:return: _description_
:rtype: _type_
"""
matching_rules = self.get_matching_acl_rules_for_ier(ier, acl, nodes)
@@ -82,12 +109,19 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return blocked_rules
def get_allow_acl_rules_for_ier(self, ier, acl, nodes):
"""
Get all allowing ACL rules for an IER.
def get_allow_acl_rules_for_ier(
self, ier: IER, acl: AccessControlList, nodes: Dict[str, NodeUnion]
) -> Dict[str, Any]:
"""Get all allowing ACL rules for an IER.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param ier: Information Exchange Request to query against the ACL list
:type ier: IER
:param acl: Firewall rules
:type acl: AccessControlList
:param nodes: Nodes in the network
:type nodes: Dict[str,NodeUnion]
:return: _description_
:rtype: _type_
"""
matching_rules = self.get_matching_acl_rules_for_ier(ier, acl, nodes)
@@ -100,19 +134,32 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
def get_matching_acl_rules(
self,
source_node_id,
dest_node_id,
protocol,
port,
acl,
nodes,
services_list,
):
"""
Get matching ACL rules.
source_node_id: str,
dest_node_id: str,
protocol: str,
port: str,
acl: AccessControlList,
nodes: Dict[str, Union[ServiceNode, ActiveNode]],
services_list: List[str],
) -> Dict[str, ACLRule]:
"""Filter ACL rules to only those which are relevant to the specified nodes.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param source_node_id: Source node
:type source_node_id: str
:param dest_node_id: Destination nodes
:type dest_node_id: str
:param protocol: Network protocol
:type protocol: str
:param port: Network port
:type port: str
:param acl: Access Control list which will be filtered
:type acl: AccessControlList
:param nodes: The environment's node directory.
:type nodes: Dict[str, Union[ServiceNode, ActiveNode]]
:param services_list: List of services registered for the environment.
:type services_list: List[str]
:return: Filtered version of 'acl'
:rtype: Dict[str, ACLRule]
"""
if source_node_id != "ANY":
source_node_address = nodes[str(source_node_id)].ip_address
@@ -132,19 +179,33 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
def get_allow_acl_rules(
self,
source_node_id,
dest_node_id,
protocol,
port,
acl,
nodes,
services_list,
):
"""
Get the ALLOW ACL rules.
source_node_id: int,
dest_node_id: str,
protocol: int,
port: str,
acl: AccessControlList,
nodes: Dict[str, NodeUnion],
services_list: List[str],
) -> Dict[str, ACLRule]:
"""List ALLOW rules relating to specified nodes.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param source_node_id: Source node id
:type source_node_id: int
:param dest_node_id: Destination node
:type dest_node_id: str
:param protocol: Network protocol
:type protocol: int
:param port: Port
:type port: str
:param acl: Firewall ruleset which is applied to the network
:type acl: AccessControlList
:param nodes: The simulation's node store
:type nodes: Dict[str, NodeUnion]
:param services_list: Services list
:type services_list: List[str]
:return: Filtered ACL Rule directory which includes only those rules which affect the specified source and
desination nodes
:rtype: Dict[str, ACLRule]
"""
matching_rules = self.get_matching_acl_rules(
source_node_id,
@@ -165,19 +226,33 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
def get_deny_acl_rules(
self,
source_node_id,
dest_node_id,
protocol,
port,
acl,
nodes,
services_list,
):
"""
Get the DENY ACL rules.
source_node_id: int,
dest_node_id: str,
protocol: int,
port: str,
acl: AccessControlList,
nodes: Dict[str, NodeUnion],
services_list: List[str],
) -> Dict[str, ACLRule]:
"""List DENY rules relating to specified nodes.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param source_node_id: Source node id
:type source_node_id: int
:param dest_node_id: Destination node
:type dest_node_id: str
:param protocol: Network protocol
:type protocol: int
:param port: Port
:type port: str
:param acl: Firewall ruleset which is applied to the network
:type acl: AccessControlList
:param nodes: The simulation's node store
:type nodes: Dict[str, NodeUnion]
:param services_list: Services list
:type services_list: List[str]
:return: Filtered ACL Rule directory which includes only those rules which affect the specified source and
desination nodes
:rtype: Dict[str, ACLRule]
"""
matching_rules = self.get_matching_acl_rules(
source_node_id,
@@ -196,7 +271,7 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return allowed_rules
def _calculate_action_full_view(self, obs):
def _calculate_action_full_view(self, obs: np.ndarray) -> int:
"""
Calculate a good acl-based action for the blue agent to take.
@@ -224,8 +299,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.
:param obs: current observation from the gym environment
:type obs: np.ndarray
:return: Optimal action to take in the environment (chosen from the discrete action space)
:rtype: int
"""
# obs = convert_to_old_obs(obs)
r_obs = transform_change_obs_readable(obs)
@@ -361,8 +438,9 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
action = get_new_action(action, self._env.action_dict)
return action
def _calculate_action_basic_view(self, obs):
"""Calculate a good acl-based action for the blue agent to take.
def _calculate_action_basic_view(self, obs: np.ndarray) -> int:
"""
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.
@@ -379,8 +457,10 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
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.
:param obs: current observation from the gym environment
:type obs: np.ndarray
:return: Optimal action to take in the environment (chosen from the discrete action space)
:rtype: int
"""
action_dict = self._env.action_dict
r_obs = transform_change_obs_readable(obs)

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

@@ -1,3 +1,5 @@
import numpy as np
from primaite.agents.agent import HardCodedAgentSessionABC
from primaite.agents.utils import get_new_action, transform_action_node_enum, transform_change_obs_readable
@@ -5,12 +7,14 @@ from primaite.agents.utils import get_new_action, transform_action_node_enum, tr
class HardCodedNodeAgent(HardCodedAgentSessionABC):
"""An Agent Session class that implements a deterministic Node agent."""
def _calculate_action(self, obs):
def _calculate_action(self, obs: np.ndarray) -> int:
"""
Calculate a good node-based action for the blue agent to take.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param obs: current observation from the gym environment
:type obs: np.ndarray
:return: Optimal action to take in the environment (chosen from the discrete action space)
:rtype: int
"""
action_dict = self._env.action_dict
r_obs = transform_change_obs_readable(obs)

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

@@ -44,6 +44,18 @@ class RLlibAgent(AgentSessionABC):
"""An AgentSession class that implements a Ray RLlib agent."""
def __init__(self, training_config_path, lay_down_config_path):
"""
Initialise the RLLib Agent training session.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:raises ValueError: If the training config contains an unexpected value for agent_framework (should be "RLLIB")
:raises ValueError: If the training config contains an unexpected value for agent_identifies (should be `PPO`
or `A2C`)
"""
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}"

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

@@ -19,6 +19,18 @@ class SB3Agent(AgentSessionABC):
"""An AgentSession class that implements a Stable Baselines3 agent."""
def __init__(self, training_config_path, lay_down_config_path):
"""
Initialise the SB3 Agent training session.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:raises ValueError: If the training config contains an unexpected value for agent_framework (should be "SB3")
:raises ValueError: If the training config contains an unexpected value for agent_identifies (should be `PPO`
or `A2C`)
"""
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}"

3
src/primaite/agents/simple.py
View File

@@ -17,8 +17,7 @@ class DummyAgent(HardCodedAgentSessionABC):
"""
A Dummy Agent.
All action spaces setup so dummy action is always 0 regardless of action
type used.
All action spaces setup so dummy action is always 0 regardless of action type used.
"""
def _calculate_action(self, obs):

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

@@ -1,5 +1,8 @@
from typing import Dict, List, Union
import numpy as np
from primaite.common.custom_typing import NodeUnion
from primaite.common.enums import (
HardwareState,
LinkStatus,
@@ -10,15 +13,17 @@ from primaite.common.enums import (
)
def transform_action_node_readable(action):
"""
Convert a node action from enumerated format to readable format.
def transform_action_node_readable(action: List[int]) -> List[Union[int, str]]:
"""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.
:param action: Agent action, formatted as a list of ints, for more information check out
`primaite.environment.primaite_env.Primaite`
:type action: List[int]
:return: The same action list, but with the encodings translated back into meaningful labels
:rtype: List[Union[int,str]]
"""
action_node_property = NodePOLType(action[1]).name
@@ -33,15 +38,18 @@ def transform_action_node_readable(action):
return new_action
def transform_action_acl_readable(action):
def transform_action_acl_readable(action: List[str]) -> List[Union[str, int]]:
"""
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.
:param action: Agent action, formatted as a list of ints, for more information check out
`primaite.environment.primaite_env.Primaite`
:type action: List[int]
:return: The same action list, but with the encodings translated back into meaningful labels
:rtype: List[Union[int,str]]
"""
action_decisions = {0: "NONE", 1: "CREATE", 2: "DELETE"}
action_permissions = {0: "DENY", 1: "ALLOW"}
@@ -58,8 +66,9 @@ def transform_action_acl_readable(action):
return new_action
def is_valid_node_action(action):
"""Is the node action an actual valid action.
def is_valid_node_action(action: List[int]) -> bool:
"""
Is the node action an actual valid action.
Only uses information about the action to determine if the action has an effect
@@ -67,8 +76,11 @@ def is_valid_node_action(action):
- 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.
:param action: Agent action, formatted as a list of ints, for more information check out
`primaite.environment.primaite_env.Primaite`
:type action: List[int]
:return: Whether the action is valid
:rtype: bool
"""
action_r = transform_action_node_readable(action)
@@ -93,7 +105,7 @@ def is_valid_node_action(action):
return True
def is_valid_acl_action(action):
def is_valid_acl_action(action: List[int]) -> bool:
"""
Is the ACL action an actual valid action.
@@ -103,8 +115,11 @@ def is_valid_acl_action(action):
- 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.
:param action: Agent action, formatted as a list of ints, for more information check out
`primaite.environment.primaite_env.Primaite`
:type action: List[int]
:return: Whether the action is valid
:rtype: bool
"""
action_r = transform_action_acl_readable(action)
@@ -126,12 +141,15 @@ def is_valid_acl_action(action):
return True
def is_valid_acl_action_extra(action):
def is_valid_acl_action_extra(action: List[int]) -> bool:
"""
Harsher version of valid acl actions, does not allow action.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param action: Agent action, formatted as a list of ints, for more information check out
`primaite.environment.primaite_env.Primaite`
:type action: List[int]
:return: Whether the action is valid
:rtype: bool
"""
if is_valid_acl_action(action) is False:
return False
@@ -150,22 +168,23 @@ 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.
def transform_change_obs_readable(obs: np.ndarray) -> List[List[Union[str, int]]]:
"""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.
:param obs: Raw observation from the environment.
:type obs: np.ndarray
:return: The same observation, but the encoded integer values are replaced with readable names.
:rtype: List[List[Union[str, int]]]
"""
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]):
for service in range(4, 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)
@@ -173,14 +192,16 @@ def transform_change_obs_readable(obs):
return new_obs
def transform_obs_readable(obs):
"""
Transform observation to readable format.
def transform_obs_readable(obs: np.ndarray) -> List[List[Union[str, int]]]:
"""Transform observation to readable format.
example
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.
:param obs: Raw observation from the environment.
:type obs: np.ndarray
:return: The same observation, but the encoded integer values are replaced with readable names.
:rtype: List[List[Union[str, int]]]
"""
changed_obs = transform_change_obs_readable(obs)
new_obs = list(zip(*changed_obs))
@@ -190,21 +211,23 @@ def transform_obs_readable(obs):
return new_obs
def convert_to_new_obs(obs, num_nodes=10):
"""
Convert original gym Box observation space to new multiDiscrete observation space.
def convert_to_new_obs(obs: np.ndarray, num_nodes: int = 10) -> np.ndarray:
"""Convert original gym Box observation space to new multiDiscrete observation space.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param obs: observation in the 'old' (NodeLinkTable) format
:type obs: np.ndarray
:param num_nodes: number of nodes in the network, defaults to 10
:type num_nodes: int, optional
:return: reformatted observation
:rtype: np.ndarray
"""
# 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.
def convert_to_old_obs(obs: np.ndarray, num_nodes: int = 10, num_links: int = 10, num_services: int = 1) -> np.ndarray:
"""Convert to old observation.
Links filled with 0's as no information is included in new observation space.
@@ -216,8 +239,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.
:param obs: observation in the 'new' (MultiDiscrete) format
:type obs: np.ndarray
:param num_nodes: number of nodes in the network, defaults to 10
:type num_nodes: int, optional
:param num_links: number of links in the network, defaults to 10
:type num_links: int, optional
:param num_services: number of services on the network, defaults to 1
:type num_services: int, optional
:return: 2-d BOX observation space, in the same format as NodeLinkTable
:rtype: np.ndarray
"""
# Convert back to more readable, original format
reshaped_nodes = obs[:-num_links].reshape(num_nodes, num_services + 2)
@@ -239,17 +271,28 @@ 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: np.ndarray, obs2: np.ndarray, num_nodes: int = 10, num_links: int = 10, num_services: int = 1
) -> str:
"""Build a string describing the difference 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.
:param obs1: First observation
:type obs1: np.ndarray
:param obs2: Second observation
:type obs2: np.ndarray
:param num_nodes: How many nodes are in the network laydown, defaults to 10
:type num_nodes: int, optional
:param num_links: How many links are in the network laydown, defaults to 10
:type num_links: int, optional
:param num_services: How many services are configured for this scenario, defaults to 1
:type num_services: int, optional
:return: A multi-line string with a human-readable description of the difference.
:rtype: str
"""
obs1 = convert_to_old_obs(obs1, num_nodes, num_links, num_services)
obs2 = convert_to_old_obs(obs2, num_nodes, num_links, num_services)
@@ -268,7 +311,7 @@ def describe_obs_change(obs1, obs2, num_nodes=10, num_links=10, num_services=1):
return change_string
def _describe_obs_change_helper(obs_change, is_link):
def _describe_obs_change_helper(obs_change: List[int], is_link: bool) -> str:
"""
Helper funcion to describe what has changed.
@@ -277,8 +320,14 @@ def _describe_obs_change_helper(obs_change, is_link):
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.
:param obs_change: List of integers generated within the `describe_obs_change` function. It should correspond to one
row of the observation table, and have `-1` at locations where the observation hasn't changed, and the new
status where it has changed.
:type obs_change: List[int]
:param is_link: Whether the row of the observation space corresponds to a link. False means it represents a node.
:type is_link: bool
:return: A human-readable description of the difference between the two observation rows.
:rtype: str
"""
# 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]
@@ -304,15 +353,15 @@ def _describe_obs_change_helper(obs_change, is_link):
return desc
def transform_action_node_enum(action):
"""
Convert a node action from readable string format, to enumerated format.
def transform_action_node_enum(action: List[Union[str, int]]) -> List[int]:
"""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.
:param action: Action in 'readable' format
:type action: List[Union[str,int]]
:return: Action with verbs encoded as ints
:rtype: List[int]
"""
action_node_id = action[0]
action_node_property = NodePOLType[action[1]].value
@@ -336,63 +385,14 @@ def transform_action_node_enum(action):
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):
def transform_action_acl_enum(action: List[Union[int, str]]) -> np.ndarray:
"""
Convert acl action from readable str format, to enumerated format.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
:param action: ACL-based action expressed as a list of human-readable ints and strings
:type action: List[Union[int,str]]
:return: The same action but encoded to contain only integers.
:rtype: np.ndarray
"""
action_decisions = {"NONE": 0, "CREATE": 1, "DELETE": 2}
action_permissions = {"DENY": 0, "ALLOW": 1}
@@ -410,35 +410,17 @@ def transform_action_acl_enum(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.
def get_node_of_ip(ip: str, node_dict: Dict[str, NodeUnion]) -> str:
"""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.
:param ip: The IP address of the node whose ID is required
:type ip: str
:param node_dict: The environment's node registry dictionary
:type node_dict: Dict[str,NodeUnion]
:return: The key from the registry dict that corresponds to the node with the IP adress provided by `ip`
:rtype: str
"""
for node_key, node_value in node_dict.items():
node_ip = node_value.ip_address
@@ -446,104 +428,18 @@ def get_node_of_ip(ip, node_dict):
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):
def get_new_action(old_action: np.ndarray, action_dict: Dict[int, List]) -> int:
"""
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.
:param old_action: Action expressed as a list of choices, eg. [1,1,1,0]
:type old_action: np.ndarray
:param action_dict: Dictionary for translating the multidiscrete actions into the list-based actions.
:type action_dict: Dict[int,List]
:return: Action key correspoinding to the input `old_action`
:rtype: int
"""
for key, val in action_dict.items():
if list(val) == list(old_action):

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

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Objects which are shared between many PrimAITE modules."""

6
src/primaite/common/protocol.py
View File

@@ -7,10 +7,10 @@ class Protocol(object):
def __init__(self, _name):
"""
Init.
Initialise a protocol.
Args:
_name: The protocol name
:param _name: The name of the protocol
:type _name: str
"""
self.name = _name
self.load = 0 # bps

2
src/primaite/common/service.py
View File

@@ -9,7 +9,7 @@ class Service(object):
def __init__(self, name: str, port: str, software_state: SoftwareState):
"""
Init.
Initialise a service.
:param name: The service name.
:param port: The service port.

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

@@ -0,0 +1 @@
"""Configuration parameters for running experiments."""

3
src/primaite/config/lay_down_config.py
View File

@@ -26,8 +26,7 @@ def load(file_path: Union[str, Path], legacy_file: bool = False) -> Dict:
Read in a lay down config yaml file.
:param file_path: The config file path.
:param legacy_file: True if the config file is legacy format, otherwise
False.
:param legacy_file: True if the config file is legacy format, otherwise False.
:return: The lay down config as a dict.
:raises ValueError: If the file_path does not exist.
"""

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

@@ -1,3 +1,4 @@
"""Utility to generate plots of sessions metrics after PrimAITE."""
from enum import Enum

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

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Gym/Gymnasium environment for RL agents consisting of a simulated computer network."""

97
src/primaite/environment/observations.py
View File

@@ -25,6 +25,12 @@ class AbstractObservationComponent(ABC):
@abstractmethod
def __init__(self, env: "Primaite"):
"""
Initialise observation component.
:param env: Primaite training environment.
:type env: Primaite
"""
_LOGGER.info(f"Initialising {self} observation component")
self.env: "Primaite" = env
self.space: spaces.Space
@@ -44,11 +50,13 @@ class AbstractObservationComponent(ABC):
class NodeLinkTable(AbstractObservationComponent):
"""Table with nodes and links as rows and hardware/software status as cols.
"""
Table with nodes and links as rows and hardware/software status as cols.
This will create the observation space formatted as a table of integers.
There is one row per node, followed by one row per link.
The number of columns is 4 plus one per service. They are:
* node/link ID
* node hardware status / 0 for links
* node operating system status (if active/service) / 0 for links
@@ -67,6 +75,12 @@ class NodeLinkTable(AbstractObservationComponent):
_DATA_TYPE: type = np.int64
def __init__(self, env: "Primaite"):
"""
Initialise a NodeLinkTable observation space component.
:param env: Training environment.
:type env: Primaite
"""
super().__init__(env)
# 1. Define the shape of your observation space component
@@ -88,7 +102,8 @@ class NodeLinkTable(AbstractObservationComponent):
self.structure = self.generate_structure()
def update(self):
"""Update the observation based on current environment state.
"""
Update the observation based on current environment state.
The structure of the observation space is described in :class:`.NodeLinkTable`
"""
@@ -169,12 +184,14 @@ class NodeLinkTable(AbstractObservationComponent):
class NodeStatuses(AbstractObservationComponent):
"""Flat list of nodes' hardware, OS, file system, and service states.
"""
Flat list of nodes' hardware, OS, file system, and service states.
The MultiDiscrete observation space can be though of as a one-dimensional vector of discrete states, represented by
integers.
Each node has 3 elements plus 1 per service. It will have the following structure:
.. code-block::
[
node1 hardware state,
node1 OS state,
@@ -190,14 +207,17 @@ class NodeStatuses(AbstractObservationComponent):
node2 serviceN state (one for each service),
...
]
:param env: The environment that forms the basis of the observations
:type env: Primaite
"""
_DATA_TYPE: type = np.int64
def __init__(self, env: "Primaite"):
"""
Initialise a NodeStatuses observation component.
:param env: Training environment.
:type env: Primaite
"""
super().__init__(env)
# 1. Define the shape of your observation space component
@@ -218,7 +238,8 @@ class NodeStatuses(AbstractObservationComponent):
self.structure = self.generate_structure()
def update(self):
"""Update the observation based on current environment state.
"""
Update the observation based on current environment state.
The structure of the observation space is described in :class:`.NodeStatuses`
"""
@@ -271,28 +292,22 @@ class NodeStatuses(AbstractObservationComponent):
class LinkTrafficLevels(AbstractObservationComponent):
"""Flat list of traffic levels encoded into banded categories.
"""
Flat list of traffic levels encoded into banded categories.
For each link, total traffic or traffic per service is encoded into a categorical value.
For example, if ``quantisation_levels=5``, the traffic levels represent these values:
0 = No traffic (0% of bandwidth)
1 = No traffic (0%-33% of bandwidth)
2 = No traffic (33%-66% of bandwidth)
3 = No traffic (66%-100% of bandwidth)
4 = No traffic (100% of bandwidth)
* 0 = No traffic (0% of bandwidth)
* 1 = No traffic (0%-33% of bandwidth)
* 2 = No traffic (33%-66% of bandwidth)
* 3 = No traffic (66%-100% of bandwidth)
* 4 = No traffic (100% of bandwidth)
.. note::
The lowest category always corresponds to no traffic and the highest category to the link being at max capacity.
Any amount of traffic between 0% and 100% (exclusive) is divided evenly into the remaining categories.
:param env: The environment that forms the basis of the observations
:type env: Primaite
:param combine_service_traffic: Whether to consider total traffic on the link, or each protocol individually,
defaults to False
:type combine_service_traffic: bool, optional
:param quantisation_levels: How many bands to consider when converting the traffic amount to a categorical value ,
defaults to 5
:type quantisation_levels: int, optional
"""
_DATA_TYPE: type = np.int64
@@ -303,6 +318,18 @@ class LinkTrafficLevels(AbstractObservationComponent):
combine_service_traffic: bool = False,
quantisation_levels: int = 5,
):
"""
Initialise a LinkTrafficLevels observation component.
:param env: The environment that forms the basis of the observations
:type env: Primaite
:param combine_service_traffic: Whether to consider total traffic on the link, or each protocol individually,
defaults to False
:type combine_service_traffic: bool, optional
:param quantisation_levels: How many bands to consider when converting the traffic amount to a categorical
value, defaults to 5
:type quantisation_levels: int, optional
"""
if quantisation_levels < 3:
_msg = (
f"quantisation_levels must be 3 or more because the lowest and highest levels are "
@@ -333,7 +360,8 @@ class LinkTrafficLevels(AbstractObservationComponent):
self.structure = self.generate_structure()
def update(self):
"""Update the observation based on current environment state.
"""
Update the observation based on current environment state.
The structure of the observation space is described in :class:`.LinkTrafficLevels`
"""
@@ -374,10 +402,11 @@ class LinkTrafficLevels(AbstractObservationComponent):
class ObservationsHandler:
"""Component-based observation space handler.
"""
Component-based observation space handler.
This allows users to configure observation spaces by mixing and matching components.
Each component can also define further parameters to make them more flexible.
This allows users to configure observation spaces by mixing and matching components. Each component can also define
further parameters to make them more flexible.
"""
_REGISTRY: Final[Dict[str, type]] = {
@@ -387,6 +416,7 @@ class ObservationsHandler:
}
def __init__(self):
"""Initialise the observation handler."""
self.registered_obs_components: List[AbstractObservationComponent] = []
# internal the observation space (unflattened version of space if flatten=True)
@@ -414,7 +444,8 @@ class ObservationsHandler:
self._flat_observation = spaces.flatten(self._space, self._observation)
def register(self, obs_component: AbstractObservationComponent):
"""Add a component for this handler to track.
"""
Add a component for this handler to track.
:param obs_component: The component to add.
:type obs_component: AbstractObservationComponent
@@ -423,10 +454,11 @@ class ObservationsHandler:
self.update_space()
def deregister(self, obs_component: AbstractObservationComponent):
"""Remove a component from this handler.
"""
Remove a component from this handler.
:param obs_component: Which component to remove. It must exist within this object's
``registered_obs_components`` attribute.
``registered_obs_components`` attribute.
:type obs_component: AbstractObservationComponent
"""
self.registered_obs_components.remove(obs_component)
@@ -466,11 +498,13 @@ class ObservationsHandler:
@classmethod
def from_config(cls, env: "Primaite", obs_space_config: dict):
"""Parse a config dictinary, return a new observation handler populated with new observation component objects.
"""
Parse a config dictinary, return a new observation handler populated with new observation component objects.
The expected format for the config dictionary is:
..code-block::python
.. code-block:: python
config = {
components: [
{
@@ -510,7 +544,8 @@ class ObservationsHandler:
return handler
def describe_structure(self):
"""Create a list of names for the features of the obs space.
"""
Create a list of names for the features of the obs space.
The order of labels follows the flattened version of the space.
"""

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

@@ -73,8 +73,7 @@ class Primaite(Env):
:param training_config_path: The training config filepath.
:param lay_down_config_path: The lay down config filepath.
:param session_path: The directory path the session is writing to.
:param timestamp_str: The session timestamp in the format:
<yyyy-mm-dd>_<hh-mm-ss>.
:param timestamp_str: The session timestamp in the format: <yyyy-mm-dd>_<hh-mm- ss>.
"""
self.session_path: Final[Path] = session_path
self.timestamp_str: Final[str] = timestamp_str
@@ -668,7 +667,8 @@ class Primaite(Env):
pass
def init_observations(self) -> Tuple[spaces.Space, np.ndarray]:
"""Create the environment's observation handler.
"""
Create the environment's observation handler.
:return: The observation space, initial observation (zeroed out array with the correct shape)
:rtype: Tuple[spaces.Space, np.ndarray]
@@ -1048,7 +1048,8 @@ class Primaite(Env):
self.num_ports = len(self.ports_list)
def get_observation_info(self, observation_info):
"""Extracts observation_info.
"""
Extracts observation_info.
:param observation_info: Config item that defines which type of observation space to use
:type observation_info: str
@@ -1065,7 +1066,8 @@ class Primaite(Env):
self.action_type = ActionType[action_info["type"]]
def save_obs_config(self, obs_config: dict):
"""Cache the config for the observation space.
"""
Cache the config for the observation space.
This is necessary as the observation space can't be built while reading the config,
it must be done after all the nodes, links, and services have been initialised.
@@ -1078,10 +1080,9 @@ class Primaite(Env):
def reset_environment(self):
"""
# Resets environment.
Resets environment.
Uses config data config data in order to build the environment
configuration.
Uses config data config data in order to build the environment configuration.
"""
for item in self.lay_down_config:
if item["item_type"] == "NODE":
@@ -1165,7 +1166,6 @@ class Primaite(Env):
5: [1, 3, 1, 0],
...
}
"""
# reserve 0 action to be a nothing action
actions = {0: [1, 0, 0, 0]}
@@ -1221,7 +1221,6 @@ class Primaite(Env):
Create a dictionary mapping each possible discrete action to a more readable mutlidiscrete action.
The dictionary contains actions of both Node and ACL action types.
"""
node_action_dict = self.create_node_action_dict()
acl_action_dict = self.create_acl_action_dict()

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

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Network connections between nodes in the simulation."""

13
src/primaite/links/link.py
View File

@@ -10,14 +10,13 @@ class Link(object):
def __init__(self, _id, _bandwidth, _source_node_name, _dest_node_name, _services):
"""
Init.
Initialise a Link within the simulated network.
Args:
_id: The IER id
_bandwidth: The bandwidth of the link (bps)
_source_node_name: The name of the source node
_dest_node_name: The name of the destination node
_protocols: The protocols to add to the link
:param _id: The IER id
:param _bandwidth: The bandwidth of the link (bps)
:param _source_node_name: The name of the source node
:param _dest_node_name: The name of the destination node
:param _protocols: The protocols to add to the link
"""
self.id = _id
self.bandwidth = _bandwidth

3
src/primaite/main.py
View File

@@ -14,7 +14,8 @@ def run(
training_config_path: Union[str, Path],
lay_down_config_path: Union[str, Path],
):
"""Run the PrimAITE Session.
"""
Run the PrimAITE Session.
:param training_config_path: The training config filepath.
:param lay_down_config_path: The lay down config filepath.

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

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Nodes represent network hosts in the simulation."""

2
src/primaite/nodes/active_node.py
View File

@@ -26,7 +26,7 @@ class ActiveNode(Node):
config_values: TrainingConfig,
):
"""
Init.
Initialise an active node.
:param node_id: The node ID
:param name: The node name

2
src/primaite/nodes/node.py
View File

@@ -19,7 +19,7 @@ class Node:
config_values: TrainingConfig,
):
"""
Init.
Initialise a node.
:param node_id: The node id.
:param name: The name of the node.

17
src/primaite/nodes/node_state_instruction_green.py
View File

@@ -16,16 +16,15 @@ class NodeStateInstructionGreen(object):
_state,
):
"""
Init.
Initialise the Node State Instruction.
Args:
_id: The node state instruction id
_start_step: The start step of the instruction
_end_step: The end step of the instruction
_node_id: The id of the associated node
_node_pol_type: The pattern of life type
_service_name: The service name
_state: The state (node or service)
:param _id: The node state instruction id
:param _start_step: The start step of the instruction
:param _end_step: The end step of the instruction
:param _node_id: The id of the associated node
:param _node_pol_type: The pattern of life type
:param _service_name: The service name
:param _state: The state (node or service)
"""
self.id = _id
self.start_step = _start_step

25
src/primaite/nodes/node_state_instruction_red.py
View File

@@ -24,20 +24,19 @@ class NodeStateInstructionRed(object):
_pol_source_node_service_state,
):
"""
Init.
Initialise the Node State Instruction for the red agent.
Args:
_id: The node state instruction id
_start_step: The start step of the instruction
_end_step: The end step of the instruction
_target_node_id: The id of the associated node
-pol_initiator: The way the PoL is applied (DIRECT, IER or SERVICE)
_pol_type: The pattern of life type
-pol_protocol: The pattern of life protocol/service affected
_pol_state: The state (node or service)
_pol_source_node_id: The source node Id (used for initiator type SERVICE)
_pol_source_node_service: The source node service (used for initiator type SERVICE)
_pol_source_node_service_state: The source node service state (used for initiator type SERVICE)
:param _id: The node state instruction id
:param _start_step: The start step of the instruction
:param _end_step: The end step of the instruction
:param _target_node_id: The id of the associated node
:param -pol_initiator: The way the PoL is applied (DIRECT, IER or SERVICE)
:param _pol_type: The pattern of life type
:param pol_protocol: The pattern of life protocol/service affected
:param _pol_state: The state (node or service)
:param _pol_source_node_id: The source node Id (used for initiator type SERVICE)
:param _pol_source_node_service: The source node service (used for initiator type SERVICE)
:param _pol_source_node_service_state: The source node service state (used for initiator type SERVICE)
"""
self.id = _id
self.start_step = _start_step

2
src/primaite/nodes/passive_node.py
View File

@@ -18,7 +18,7 @@ class PassiveNode(Node):
config_values: TrainingConfig,
):
"""
Init.
Initialise a passive node.
:param node_id: The node id.
:param name: The name of the node.

8
src/primaite/nodes/service_node.py
View File

@@ -27,7 +27,7 @@ class ServiceNode(ActiveNode):
config_values: TrainingConfig,
):
"""
Init.
Initialise a Service Node.
:param node_id: The node ID
:param name: The node name
@@ -77,8 +77,7 @@ class ServiceNode(ActiveNode):
Indicates whether a service is in a running state on the node.
:param protocol_name: The service (protocol)
:return: True if service (protocol) is in a running state on the
node, otherwise False.
:return: True if service (protocol) is in a running state on the node, otherwise False.
"""
for service_key, service_value in self.services.items():
if service_key == protocol_name:
@@ -93,8 +92,7 @@ class ServiceNode(ActiveNode):
Indicates whether a service is in an overwhelmed state on the node.
:param protocol_name: The service (protocol)
:return: True if service (protocol) is in an overwhelmed state on the
node, otherwise False.
:return: True if service (protocol) is in an overwhelmed state on the node, otherwise False.
"""
for service_key, service_value in self.services.items():
if service_key == protocol_name:

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

@@ -1,3 +1,4 @@
"""Contains default jupyter notebooks which demonstrate PrimAITE functionality."""
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
import importlib.util
import os

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

@@ -1 +1,2 @@
"""Pattern of Life- Represents the actions of users on the network."""
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

23
src/primaite/pol/ier.py
View File

@@ -23,19 +23,18 @@ class IER(object):
_running=False,
):
"""
Init.
Initialise an Information Exchange Request.
Args:
_id: The IER id
_start_step: The step when this IER should start
_end_step: The step when this IER should end
_load: The load this IER should put on a link (bps)
_protocol: The protocol of this IER
_port: The port this IER runs on
_source_node_id: The source node ID
_dest_node_id: The destination node ID
_mission_criticality: Criticality of this IER to the mission (0 none, 5 mission critical)
_running: Indicates whether the IER is currently running
:param _id: The IER id
:param _start_step: The step when this IER should start
:param _end_step: The step when this IER should end
:param _load: The load this IER should put on a link (bps)
:param _protocol: The protocol of this IER
:param _port: The port this IER runs on
:param _source_node_id: The source node ID
:param _dest_node_id: The destination node ID
:param _mission_criticality: Criticality of this IER to the mission (0 none, 5 mission critical)
:param _running: Indicates whether the IER is currently running
"""
self.id = _id
self.start_step = _start_step

14
src/primaite/pol/red_agent_pol.py
View File

@@ -296,11 +296,17 @@ def apply_red_agent_node_pol(
pass
def is_red_ier_incoming(node, iers, node_pol_type):
"""
Checks if the RED IER is incoming.
def is_red_ier_incoming(node: NodeUnion, iers: Dict[str, IER], node_pol_type: NodePOLType) -> bool:
"""Checks if the RED IER is incoming.
TODO: Write more descriptive docstring with params and returns.
:param node: Destination node of the IER
:type node: NodeUnion
:param iers: Directory of IERs
:type iers: Dict[str,IER]
:param node_pol_type: Type of Pattern-Of-Life
:type node_pol_type: NodePOLType
:return: Whether the RED IER is incoming.
:rtype: bool
"""
node_id = node.node_id

8
src/primaite/primaite_session.py
View File

@@ -1,3 +1,4 @@
"""Main entry point to PrimAITE. Configure training/evaluation experiments and input/output."""
from __future__ import annotations
from pathlib import Path
@@ -21,8 +22,7 @@ class PrimaiteSession:
"""
The PrimaiteSession class.
Provides a single learning and evaluation entry point for all training
and lay down configurations.
Provides a single learning and evaluation entry point for all training and lay down configurations.
"""
def __init__(
@@ -38,12 +38,12 @@ class PrimaiteSession:
"""
if not isinstance(training_config_path, Path):
training_config_path = Path(training_config_path)
self._training_config_path: Final[Union[Path]] = training_config_path
self._training_config_path: Final[Union[Path, str]] = training_config_path
self._training_config: Final[TrainingConfig] = training_config.load(self._training_config_path)
if not isinstance(lay_down_config_path, Path):
lay_down_config_path = Path(lay_down_config_path)
self._lay_down_config_path: Final[Union[Path]] = lay_down_config_path
self._lay_down_config_path: Final[Union[Path, str]] = lay_down_config_path
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path)
self._agent_session: AgentSessionABC = None # noqa

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

@@ -1 +1,2 @@
"""Utilities to prepare the user's data folders."""
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

3
src/primaite/setup/reset_demo_notebooks.py
View File

@@ -15,8 +15,7 @@ def run(overwrite_existing: bool = True):
"""
Resets the demo jupyter notebooks in the users app notebooks directory.
:param overwrite_existing: A bool to toggle replacing existing edited
notebooks on or off.
:param overwrite_existing: A bool to toggle replacing existing edited notebooks on or off.
"""
notebooks_package_data_root = pkg_resources.resource_filename("primaite", "notebooks/_package_data")
for subdir, dirs, files in os.walk(notebooks_package_data_root):

3
src/primaite/setup/reset_example_configs.py
View File

@@ -14,8 +14,7 @@ def run(overwrite_existing=True):
"""
Resets the example config files in the users app config directory.
:param overwrite_existing: A bool to toggle replacing existing edited
config on or off.
:param overwrite_existing: A bool to toggle replacing existing edited config on or off.
"""
configs_package_data_root = pkg_resources.resource_filename("primaite", "config/_package_data")

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

@@ -1 +1,2 @@
"""Record data of the system's state and agent's observations and actions."""
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

3
src/primaite/transactions/transaction.py
View File

@@ -86,8 +86,7 @@ def _turn_obs_space_to_array(obs_space, obs_assets, obs_features) -> List[str]:
Turns observation space into a string array so it can be saved to csv.
:param obs_space: The observation space
:param obs_assets: The number of assets (i.e. nodes or links) in the
observation space
:param obs_assets: The number of assets (i.e. nodes or links) in the observation space
:param obs_features: The number of features associated with the asset
:return: The observation space as an array of strings
"""

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

@@ -0,0 +1 @@
"""Utilities for PrimAITE."""

3
src/primaite/utils/session_output_reader.py
View File

@@ -10,8 +10,7 @@ def av_rewards_dict(av_rewards_csv_file: Union[str, Path]) -> Dict[int, float]:
"""
Read an average rewards per episode csv file and return as a dict.
The dictionary keys are the episode number, and the values are the mean
reward that episode.
The dictionary keys are the episode number, and the values are the mean reward that episode.
:param av_rewards_csv_file: The average rewards per episode csv file path.
:return: The average rewards per episode cdv as a dict.

15
src/primaite/utils/session_output_writer.py
View File

@@ -29,6 +29,18 @@ class SessionOutputWriter:
transaction_writer: bool = False,
learning_session: bool = True,
):
"""
Initialise the Session Output Writer.
:param env: PrimAITE gym environment.
:type env: Primaite
:param transaction_writer: If `true`, this will output a full account of every transaction taken by the agent.
If `false` it will output the average reward per episode, defaults to False
:type transaction_writer: bool, optional
:param learning_session: Set to `true` to indicate that the current session is a training session. This
determines the name of the folder which contains the final output csv. Defaults to True
:type learning_session: bool, optional
"""
self._env = env
self.transaction_writer = transaction_writer
self.learning_session = learning_session
@@ -68,8 +80,7 @@ class SessionOutputWriter:
"""
Write a row of session data.
:param data: The row of data to write. Can be a Tuple or an instance
of Transaction.
:param data: The row of data to write. Can be a Tuple or an instance of Transaction.
"""
if isinstance(data, Transaction):
header, data = data.as_csv_data()