From 7382ed26b380b60f5049b8cd80a282447eae023a Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 9 Jun 2023 11:25:45 +0100 Subject: [PATCH 01/37] 901 - changed AccessControlList in access_control_list.py from a dict to a list --- src/primaite/acl/access_control_list.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 284ed764..e356d276 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -1,6 +1,6 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """A class that implements the access control list implementation for the network.""" -from typing import Dict +from typing import List from primaite.acl.acl_rule import ACLRule @@ -10,7 +10,8 @@ class AccessControlList: def __init__(self): """Init.""" - self.acl: Dict[str, AccessControlList] = {} # A dictionary of ACL Rules + # A list of ACL Rules + self.acl: List[ACLRule] = [] def check_address_match(self, _rule, _source_ip_address, _dest_ip_address): """ From afc133cbc5f2703cb2039856f788fa8df379ec3b Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 9 Jun 2023 15:17:20 +0100 Subject: [PATCH 02/37] 901 - added ACL list to observations.py as its own observation space with the ACL attributes and the position of the ACL rule in the ACL list, added ImplicitFirewallRule to enums.py and added acl_implicit_rule, max_acl_list to primaite_env.py --- src/primaite/acl/access_control_list.py | 6 +- src/primaite/common/enums.py | 7 ++ src/primaite/environment/observations.py | 89 +++++++++++++++++++++++- src/primaite/environment/primaite_env.py | 7 ++ 4 files changed, 105 insertions(+), 4 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index e356d276..7eeef731 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -77,7 +77,7 @@ class AccessControlList: # If there has been no rule to allow the IER through, it will return a blocked signal by default return True - def add_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port): + def add_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port, _position): """ Adds a new rule. @@ -87,10 +87,10 @@ class AccessControlList: _dest_ip: the destination IP address _protocol: the protocol _port: the port + _position: position to insert ACL rule into ACL list """ new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - hash_value = hash(new_rule) - self.acl[hash_value] = new_rule + self.acl.insert(_position, new_rule) def remove_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port): """ diff --git a/src/primaite/common/enums.py b/src/primaite/common/enums.py index c268a766..2faff0f5 100644 --- a/src/primaite/common/enums.py +++ b/src/primaite/common/enums.py @@ -126,3 +126,10 @@ class LinkStatus(Enum): MEDIUM = 2 HIGH = 3 OVERLOAD = 4 + + +class ImplicitFirewallRule(Enum): + """Implicit firewall rule.""" + + DENY = 0 + ALLOW = 1 diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index 9e71ef1b..d9155e47 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -6,7 +6,12 @@ from typing import TYPE_CHECKING, Dict, Final, List, Tuple, Union import numpy as np from gym import spaces -from primaite.common.enums import FileSystemState, HardwareState, SoftwareState +from primaite.common.enums import ( + FileSystemState, + HardwareState, + ImplicitFirewallRule, + SoftwareState, +) from primaite.nodes.active_node import ActiveNode from primaite.nodes.service_node import ServiceNode @@ -296,6 +301,88 @@ class LinkTrafficLevels(AbstractObservationComponent): self.current_observation[:] = obs +class AccessControlList(AbstractObservationComponent): + """Flat list of all the Access Control Rules in the Access Control List. + + The MultiDiscrete observation space can be though of as a one-dimensional vector of discrete states, represented by + integers. + + :param env: The environment that forms the basis of the observations + :type env: Primaite + :param acl_implicit_rule: Whether to have an implicit DENY or implicit ALLOW ACL rule at the end of the ACL list + Default is 0 DENY, 1 ALLOW + :type acl_implicit_rule: ImplicitFirewallRule Enumeration (ALLOW or DENY) + :param max_acl_rules: Maximum number of ACLs allowed in the environment + :type max_acl_rules: int + + Each ACL Rule has 6 elements. It will have the following structure: + .. code-block:: + [ + acl_rule1 permission, + acl_rule1 source_ip, + acl_rule1 dest_ip, + acl_rule1 protocol, + acl_rule1 port, + acl_rule1 position, + acl_rule2 permission, + acl_rule2 source_ip, + acl_rule2 dest_ip, + acl_rule2 protocol, + acl_rule2 port, + acl_rule2 position, + ... + ] + """ + + _DATA_TYPE: type = np.int64 + + def __init__( + self, + env: "Primaite", + acl_implicit_rule=ImplicitFirewallRule.DENY, + max_acl_rules: int = 5, + ): + super().__init__(env) + + self.acl_implicit_rule: ImplicitFirewallRule = acl_implicit_rule + self.max_acl_rules = max_acl_rules + + # 1. Define the shape of your observation space component + acl_shape = [ + len(ImplicitFirewallRule), + len(env.nodes), + len(env.nodes), + len(env.services_list), + len(env.ports_list), + len(env.acl), + ] + shape = acl_shape * self.env.max_acl_rules + + # 2. Create Observation space + self.space = spaces.MultiDiscrete(shape) + + # 3. Initialise observation with zeroes + self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) + + def update(self): + """Update the observation based on current environment state. + + The structure of the observation space is described in :class:`.AccessControlList` + """ + obs = [] + for acl_rule in self.env.acl: + permission = acl_rule.permission + source_ip = acl_rule.source_ip + dest_ip = acl_rule.dest_ip + protocol = acl_rule.protocol + port = acl_rule.port + position = self.env.acl.index(acl_rule) + + obs.extend([permission, source_ip, dest_ip, protocol, port, position]) + + self.current_observation[:] = obs + + class ObservationsHandler: """Component-based observation space handler. diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 2f2a071d..a61372ad 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -21,6 +21,7 @@ from primaite.common.enums import ( ActionType, FileSystemState, HardwareState, + ImplicitFirewallRule, NodePOLInitiator, NodePOLType, NodeType, @@ -157,6 +158,12 @@ class Primaite(Env): # It will be initialised later. self.obs_handler: ObservationsHandler + # Set by main_config + # Adds a DENY ALL or ALLOW ALL to the end of the Access Control List + self.acl_implicit_rule = ImplicitFirewallRule.DENY + + # Sets a limit to how many ACL + self.max_acl_rules = 0 # Open the config file and build the environment laydown try: self.config_file = open(self.config_values.config_filename_use_case, "r") From 29ba64462a6cd49d01ca3f5d08a1f89ad6e8a668 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 9 Jun 2023 15:45:13 +0100 Subject: [PATCH 03/37] 901 - changed name of enum in enums.py and added class attriubutes in access_control_list.py --- src/primaite/acl/access_control_list.py | 5 ++++- src/primaite/common/enums.py | 2 +- src/primaite/environment/observations.py | 16 ++++------------ src/primaite/environment/primaite_env.py | 24 +++++++++++++----------- 4 files changed, 22 insertions(+), 25 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 7eeef731..0b403556 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -8,10 +8,12 @@ from primaite.acl.acl_rule import ACLRule class AccessControlList: """Access Control List class.""" - def __init__(self): + def __init__(self, implicit_permission): """Init.""" # A list of ACL Rules self.acl: List[ACLRule] = [] + self.acl_implicit_rule = implicit_permission + self.max_acl_rules: int def check_address_match(self, _rule, _source_ip_address, _dest_ip_address): """ @@ -103,6 +105,7 @@ class AccessControlList: _protocol: the protocol _port: the port """ + # Add check so you cant remove implicit rule rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) hash_value = hash(rule) # There will not always be something 'popable' since the agent will be trying random things diff --git a/src/primaite/common/enums.py b/src/primaite/common/enums.py index 2faff0f5..801494ef 100644 --- a/src/primaite/common/enums.py +++ b/src/primaite/common/enums.py @@ -128,7 +128,7 @@ class LinkStatus(Enum): OVERLOAD = 4 -class ImplicitFirewallRule(Enum): +class RulePermissionType(Enum): """Implicit firewall rule.""" DENY = 0 diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index d9155e47..865b4328 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -9,7 +9,7 @@ from gym import spaces from primaite.common.enums import ( FileSystemState, HardwareState, - ImplicitFirewallRule, + RulePermissionType, SoftwareState, ) from primaite.nodes.active_node import ActiveNode @@ -336,25 +336,16 @@ class AccessControlList(AbstractObservationComponent): _DATA_TYPE: type = np.int64 - def __init__( - self, - env: "Primaite", - acl_implicit_rule=ImplicitFirewallRule.DENY, - max_acl_rules: int = 5, - ): + def __init__(self, env: "Primaite"): super().__init__(env) - self.acl_implicit_rule: ImplicitFirewallRule = acl_implicit_rule - self.max_acl_rules = max_acl_rules - # 1. Define the shape of your observation space component acl_shape = [ - len(ImplicitFirewallRule), + len(RulePermissionType), len(env.nodes), len(env.nodes), len(env.services_list), len(env.ports_list), - len(env.acl), ] shape = acl_shape * self.env.max_acl_rules @@ -394,6 +385,7 @@ class ObservationsHandler: "NODE_LINK_TABLE": NodeLinkTable, "NODE_STATUSES": NodeStatuses, "LINK_TRAFFIC_LEVELS": LinkTrafficLevels, + "ACCESS_CONTROL_LIST": AccessControlList, } def __init__(self): diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index a61372ad..c5aaf9cc 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -21,11 +21,11 @@ from primaite.common.enums import ( ActionType, FileSystemState, HardwareState, - ImplicitFirewallRule, NodePOLInitiator, NodePOLType, NodeType, Priority, + RulePermissionType, SoftwareState, ) from primaite.common.service import Service @@ -160,7 +160,7 @@ class Primaite(Env): # Set by main_config # Adds a DENY ALL or ALLOW ALL to the end of the Access Control List - self.acl_implicit_rule = ImplicitFirewallRule.DENY + self.acl_implicit_rule = RulePermissionType.DENY # Sets a limit to how many ACL self.max_acl_rules = 0 @@ -1173,7 +1173,7 @@ class Primaite(Env): def create_acl_action_dict(self): """Creates a dictionary mapping each possible discrete action to more readable multidiscrete action.""" # reserve 0 action to be a nothing action - actions = {0: [0, 0, 0, 0, 0, 0]} + actions = {0: [0, 0, 0, 0, 0, 0, 0]} action_key = 1 # 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE @@ -1185,14 +1185,16 @@ class Primaite(Env): for dest_ip in range(self.num_nodes + 1): for protocol in range(self.num_services + 1): for port in range(self.num_ports + 1): - action = [ - action_decision, - action_permission, - source_ip, - dest_ip, - protocol, - port, - ] + for position in range(self.max_acl_rules - 1): + action = [ + action_decision, + action_permission, + source_ip, + dest_ip, + protocol, + port, + position, + ] # Check to see if its an action we want to include as possible i.e. not a nothing action if is_valid_acl_action_extra(action): actions[action_key] = action From c5175c500eaae6fbf14e19336e41ebee1e744749 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 9 Jun 2023 16:56:42 +0100 Subject: [PATCH 04/37] 901 - added logic to add acls to list (needs more logic adding to it) --- src/primaite/acl/access_control_list.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 0b403556..51f4a86c 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -1,9 +1,12 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """A class that implements the access control list implementation for the network.""" -from typing import List +import logging +from typing import Final, List from primaite.acl.acl_rule import ACLRule +_LOGGER: Final[logging.Logger] = logging.getLogger(__name__) + class AccessControlList: """Access Control List class.""" @@ -92,7 +95,16 @@ class AccessControlList: _position: position to insert ACL rule into ACL list """ new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - self.acl.insert(_position, new_rule) + + if _position < self.max_acl_rules - 1 and _position < 0: + if _position < len(self.acl): + self.acl.insert(_position, new_rule) + else: + print("check logic on this") + else: + _LOGGER.info( + f"Position {_position} is an invalid index for list/overwriting implicit firewall rule" + ) def remove_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port): """ @@ -107,10 +119,9 @@ class AccessControlList: """ # Add check so you cant remove implicit rule rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - hash_value = hash(rule) # There will not always be something 'popable' since the agent will be trying random things try: - self.acl.pop(hash_value) + self.acl.remove(rule) except Exception: return From cf64990cffd5bc1a46a82c71f5a638401ecf5a55 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 13 Jun 2023 09:45:45 +0100 Subject: [PATCH 05/37] 901 - added changes back to ticket --- src/primaite/acl/access_control_list.py | 2 +- src/primaite/config/training_config.py | 24 ++++++---- src/primaite/environment/primaite_env.py | 48 +++++++++++-------- .../main_config_ACCESS_CONTROL_LIST.yaml | 4 ++ .../obs_tests/main_config_without_obs.yaml | 5 +- 5 files changed, 50 insertions(+), 33 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 44a96743..d75b9756 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -97,7 +97,7 @@ class AccessControlList: _port: the port _position: position to insert ACL rule into ACL list (starting from index 1 and NOT 0) """ - position_index = int(_position) - 1 + position_index = int(_position) new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) if len(self.acl) < self.max_acl_rules: if len(self.acl) > position_index > -1: diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 4af36abe..67403c52 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -1,7 +1,7 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, Final, Union, Optional +from typing import Any, Dict, Final, Optional, Union import yaml @@ -51,6 +51,16 @@ class TrainingConfig: observation_space_high_value: int = 1000000000 "The high value for the observation space." + # Access Control List/Rules + apply_implicit_rule: str = True + "User choice to have Implicit ALLOW or DENY." + + implicit_acl_rule: str = "DENY" + "ALLOW or DENY implicit firewall rule to go at the end of list of ACL list." + + max_number_acl_rule: int = 0 + "Sets a limit for number of acl rules allowed in the list and environment." + # Reward values # Generic all_ok: int = 0 @@ -167,8 +177,7 @@ def main_training_config_path() -> Path: return path -def load(file_path: Union[str, Path], - legacy_file: bool = False) -> TrainingConfig: +def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConfig: """ Read in a training config yaml file. @@ -213,9 +222,7 @@ def load(file_path: Union[str, Path], def convert_legacy_training_config_dict( - legacy_config_dict: Dict[str, Any], - num_steps: int = 256, - action_type: str = "ANY" + legacy_config_dict: Dict[str, Any], num_steps: int = 256, action_type: str = "ANY" ) -> Dict[str, Any]: """ Convert a legacy training config dict to the new format. @@ -227,10 +234,7 @@ def convert_legacy_training_config_dict( don't have action_type values. :return: The converted training config dict. """ - config_dict = { - "num_steps": num_steps, - "action_type": action_type - } + config_dict = {"num_steps": num_steps, "action_type": action_type} for legacy_key, value in legacy_config_dict.items(): new_key = _get_new_key_from_legacy(legacy_key) if new_key: diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index cd0c660e..0a351b08 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -118,7 +118,10 @@ class Primaite(Env): self.red_node_pol = {} # Create the Access Control List - self.acl = AccessControlList() + self.acl = AccessControlList( + self.training_config.implicit_acl_rule, + self.training_config.max_number_acl_rule, + ) # Create a list of services (enums) self.services_list = [] @@ -212,22 +215,10 @@ class Primaite(Env): # Define Action Space - depends on action space type (Node or ACL) if self.training_config.action_type == ActionType.NODE: _LOGGER.info("Action space type NODE selected") - # Terms (for node action space): - # [0, num nodes] - node ID (0 = nothing, node ID) - # [0, 4] - what property it's acting on (0 = nothing, state, SoftwareState, service state, file system state) # noqa - # [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore) # noqa - # [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa self.action_dict = self.create_node_action_dict() self.action_space = spaces.Discrete(len(self.action_dict)) elif self.training_config.action_type == ActionType.ACL: _LOGGER.info("Action space type ACL selected") - # Terms (for ACL action space): - # [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) - # [0, 1] - Permission (0 = DENY, 1 = ALLOW) - # [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) - # [0, num ports] - Port (0 = any, then 1 -> x resolving to port) self.action_dict = self.create_acl_action_dict() self.action_space = spaces.Discrete(len(self.action_dict)) elif self.training_config.action_type == ActionType.ANY: @@ -1144,6 +1135,11 @@ class Primaite(Env): } """ + # Terms (for node action space): + # [0, num nodes] - node ID (0 = nothing, node ID) + # [0, 4] - what property it's acting on (0 = nothing, state, SoftwareState, service state, file system state) # noqa + # [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore) # noqa + # [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa # reserve 0 action to be a nothing action actions = {0: [1, 0, 0, 0]} action_key = 1 @@ -1165,6 +1161,14 @@ class Primaite(Env): def create_acl_action_dict(self): """Creates a dictionary mapping each possible discrete action to more readable multidiscrete action.""" + # Terms (for ACL action space): + # [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) + # [0, 1] - Permission (0 = DENY, 1 = ALLOW) + # [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) + # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) + # [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) + # [0, num ports] - Port (0 = any, then 1 -> x resolving to port) + # [0, max acl rules - 1] - Position (0 = first index, then 1 -> x index resolving to acl rule in acl list) # reserve 0 action to be a nothing action actions = {0: [0, 0, 0, 0, 0, 0]} @@ -1178,14 +1182,16 @@ class Primaite(Env): for dest_ip in range(self.num_nodes + 1): for protocol in range(self.num_services + 1): for port in range(self.num_ports + 1): - action = [ - action_decision, - action_permission, - source_ip, - dest_ip, - protocol, - port, - ] + for position in range(self.max_acl_rules - 1): + action = [ + action_decision, + action_permission, + source_ip, + dest_ip, + protocol, + port, + position, + ] # Check to see if its an action we want to include as possible i.e. not a nothing action if is_valid_acl_action_extra(action): actions[action_key] = action diff --git a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml index b36cd6ce..856e963d 100644 --- a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml +++ b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml @@ -14,7 +14,11 @@ observationSpace: implicit_acl_rule: DENY max_number_of_acl_rules: 10 +# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) +apply_implicit_rule: True +# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY +# Total number of ACL rules allowed in the environment max_number_acl_rules: 10 numEpisodes: 1 # Time delay between steps (for generic agents) diff --git a/tests/config/obs_tests/main_config_without_obs.yaml b/tests/config/obs_tests/main_config_without_obs.yaml index 99005678..c671b31f 100644 --- a/tests/config/obs_tests/main_config_without_obs.yaml +++ b/tests/config/obs_tests/main_config_without_obs.yaml @@ -27,8 +27,11 @@ agent_load_file: C:\[Path]\[agent_saved_filename.zip] # Environment config values # The high value for the observation space observation_space_high_value: 1_000_000_000 - +# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) +apply_implicit_rule: True +# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY +# Total number of ACL rules allowed in the environment max_number_acl_rules: 10 # Reward values # Generic From fe102dff6f445a0eb432b72c4cc29e756e4a8903 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 13 Jun 2023 10:01:55 +0100 Subject: [PATCH 06/37] 901 - fixed test_acl.py tests --- tests/test_acl.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_acl.py b/tests/test_acl.py index 6410a202..e99f5ee0 100644 --- a/tests/test_acl.py +++ b/tests/test_acl.py @@ -52,7 +52,7 @@ def test_check_acl_block_affirmative(): acl_rule_destination = "192.168.1.2" acl_rule_protocol = "TCP" acl_rule_port = "80" - acl_position_in_list = "1" + acl_position_in_list = "0" acl.add_rule( acl_rule_permission, @@ -76,7 +76,7 @@ def test_check_acl_block_negative(): acl_rule_destination = "192.168.1.2" acl_rule_protocol = "TCP" acl_rule_port = "80" - acl_position_in_list = "1" + acl_position_in_list = "0" acl.add_rule( acl_rule_permission, From 5b59642695b5f25d27fc54154c49831bbf5e3fb0 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 13 Jun 2023 14:51:55 +0100 Subject: [PATCH 07/37] 901 - added max_acl_rules, implicit_acl_rule and apply_implicit rule to main_config, changed observations.py for ACLs to match the action space for ACLs, added position of acl rule to ACL action type --- src/primaite/config/training_config.py | 3 +- src/primaite/environment/observations.py | 68 +++++++++++++------ src/primaite/environment/primaite_env.py | 27 +++++--- .../obs_tests/main_config_without_obs.yaml | 2 - tests/test_observation_space.py | 10 ++- tests/test_single_action_space.py | 15 ++-- 6 files changed, 83 insertions(+), 42 deletions(-) diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 67403c52..9a21d087 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -58,7 +58,7 @@ class TrainingConfig: implicit_acl_rule: str = "DENY" "ALLOW or DENY implicit firewall rule to go at the end of list of ACL list." - max_number_acl_rule: int = 0 + max_number_acl_rules: int = 0 "Sets a limit for number of acl rules allowed in the list and environment." # Reward values @@ -190,6 +190,7 @@ def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConf :raises TypeError: When the TrainingConfig object cannot be created using the values from the config file read from ``file_path``. """ + print("FILE PATH", file_path) if not isinstance(file_path, Path): file_path = Path(file_path) if file_path.exists(): diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index a47a1a52..96df1f60 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -9,6 +9,7 @@ from gym import spaces from primaite.common.enums import ( FileSystemState, HardwareState, + Protocol, RulePermissionType, SoftwareState, ) @@ -309,11 +310,6 @@ class AccessControlList(AbstractObservationComponent): :param env: The environment that forms the basis of the observations :type env: Primaite - :param acl_implicit_rule: Whether to have an implicit DENY or implicit ALLOW ACL rule at the end of the ACL list - Default is 0 DENY, 1 ALLOW - :type acl_implicit_rule: ImplicitFirewallRule Enumeration (ALLOW or DENY) - :param max_acl_rules: Maximum number of ACLs allowed in the environment - :type max_acl_rules: int Each ACL Rule has 6 elements. It will have the following structure: .. code-block:: @@ -334,6 +330,15 @@ class AccessControlList(AbstractObservationComponent): ] """ + # Terms (for ACL observation space): + # [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) + # [0, 1] - Permission (0 = DENY, 1 = ALLOW) + # [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) + # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) + # [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) + # [0, num ports] - Port (0 = any, then 1 -> x resolving to port) + # [0, max acl rules - 1] - Position (0 = first index, then 1 -> x index resolving to acl rule in acl list) + _DATA_TYPE: type = np.int64 def __init__(self, env: "Primaite"): @@ -377,31 +382,54 @@ class AccessControlList(AbstractObservationComponent): permission_int = 1 if source_ip == "ANY": - source_ip = 0 + source_ip_int = 0 + else: + source_ip_int = self.obtain_node_id_using_ip(source_ip) if dest_ip == "ANY": - dest_ip = 0 - if port == "ANY": - port = 0 + dest_ip_int = 0 + else: + dest_ip_int = self.obtain_node_id_using_ip(dest_ip) if protocol == "ANY": protocol_int = 0 else: - while True: - if protocol in self.service_dict: - protocol_int = self.services_dict[protocol] - break - else: - self.services_dict[protocol] = len(self.services_dict) + 1 - continue - # [0 - DENY, 1 - ALLOW] Permission - # [0 - ANY, x - IP Address/Protocol/Port] + try: + protocol_int = Protocol[protocol] + except AttributeError: + _LOGGER.info(f"Service {protocol} could not be found") + if port == "ANY": + port_int = 0 + else: + if port in self.env.ports_list: + port_int = self.env.ports_list.index(port) + else: + _LOGGER.info(f"Port {port} could not be found.") - print(permission_int, source_ip, dest_ip, protocol_int, port) + print(permission_int, source_ip, dest_ip, protocol_int, port_int, position) obs.extend( - [permission_int, source_ip, dest_ip, protocol_int, port, position] + [ + permission_int, + source_ip_int, + dest_ip_int, + protocol_int, + port_int, + position, + ] ) self.current_observation[:] = obs + def obtain_node_id_using_ip(self, ip_address): + """Uses IP address of Nodes to find the ID. + + Resolves IP address -> x (node id e.g. 1 or 2 or 3 or 4) for observation space + """ + for key, node in self.env.nodes: + if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): + if node.ip_address == ip_address: + return key + _LOGGER.info(f"Node ID was not found from IP Address {ip_address}") + return -1 + class ObservationsHandler: """Component-based observation space handler. diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 0a351b08..39006259 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -120,9 +120,8 @@ class Primaite(Env): # Create the Access Control List self.acl = AccessControlList( self.training_config.implicit_acl_rule, - self.training_config.max_number_acl_rule, + self.training_config.max_number_acl_rules, ) - # Create a list of services (enums) self.services_list = [] @@ -423,14 +422,13 @@ class Primaite(Env): _action: The action space from the agent """ # At the moment, actions are only affecting nodes - if self.training_config.action_type == ActionType.NODE: self.apply_actions_to_nodes(_action) elif self.training_config.action_type == ActionType.ACL: self.apply_actions_to_acl(_action) elif ( - len(self.action_dict[_action]) == 6 - ): # ACL actions in multidiscrete form have len 6 + len(self.action_dict[_action]) == 7 + ): # ACL actions in multidiscrete form have len 7 self.apply_actions_to_acl(_action) elif ( len(self.action_dict[_action]) == 4 @@ -981,6 +979,7 @@ class Primaite(Env): acl_rule_destination, acl_rule_protocol, acl_rule_port, + 0, ) def create_services_list(self, services): @@ -1173,6 +1172,10 @@ class Primaite(Env): actions = {0: [0, 0, 0, 0, 0, 0]} action_key = 1 + print( + "what is this primaite_env.py 1177", + self.training_config.max_number_acl_rules - 1, + ) # 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE for action_decision in range(3): # 2 possible action permissions 0 = DENY, 1 = CREATE @@ -1182,7 +1185,9 @@ class Primaite(Env): for dest_ip in range(self.num_nodes + 1): for protocol in range(self.num_services + 1): for port in range(self.num_ports + 1): - for position in range(self.max_acl_rules - 1): + for position in range( + self.training_config.max_number_acl_rules - 1 + ): action = [ action_decision, action_permission, @@ -1192,10 +1197,11 @@ class Primaite(Env): port, position, ] - # Check to see if its an action we want to include as possible i.e. not a nothing action - if is_valid_acl_action_extra(action): - actions[action_key] = action - action_key += 1 + # Check to see if it is an action we want to include as possible + # i.e. not a nothing action + if is_valid_acl_action_extra(action): + actions[action_key] = action + action_key += 1 return actions @@ -1219,4 +1225,5 @@ class Primaite(Env): # Combine the Node dict and ACL dict combined_action_dict = {**acl_action_dict, **new_node_action_dict} + print("combined dict", combined_action_dict.items()) return combined_action_dict diff --git a/tests/config/obs_tests/main_config_without_obs.yaml b/tests/config/obs_tests/main_config_without_obs.yaml index c671b31f..57e80b64 100644 --- a/tests/config/obs_tests/main_config_without_obs.yaml +++ b/tests/config/obs_tests/main_config_without_obs.yaml @@ -31,8 +31,6 @@ observation_space_high_value: 1_000_000_000 apply_implicit_rule: True # Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY -# Total number of ACL rules allowed in the environment -max_number_acl_rules: 10 # Reward values # Generic all_ok: 0 diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index dbcdf2d6..4e7186b5 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -27,7 +27,8 @@ def env(request): @pytest.mark.env_config_paths( dict( - training_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml", + training_config_path=TEST_CONFIG_ROOT + / "obs_tests/main_config_without_obs.yaml", lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", ) ) @@ -43,7 +44,8 @@ def test_default_obs_space(env: Primaite): @pytest.mark.env_config_paths( dict( - training_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml", + training_config_path=TEST_CONFIG_ROOT + / "obs_tests/main_config_without_obs.yaml", lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", ) ) @@ -140,7 +142,8 @@ class TestNodeLinkTable: @pytest.mark.env_config_paths( dict( - training_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_NODE_STATUSES.yaml", + training_config_path=TEST_CONFIG_ROOT + / "obs_tests/main_config_NODE_STATUSES.yaml", lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", ) ) @@ -217,4 +220,5 @@ class TestLinkTrafficLevels: # we send 999 bits of data via link1 and link2 on service 1. # therefore the first and third elements should be 6 and all others 0 # (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%) + print(obs) assert np.array_equal(obs, [6, 0, 6, 0]) diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index 16b9d03e..4d6136a9 100644 --- a/tests/test_single_action_space.py +++ b/tests/test_single_action_space.py @@ -19,15 +19,15 @@ def run_generic_set_actions(env: Primaite): action = 0 print("Episode:", episode, "\nStep:", step) if step == 5: - # [1, 1, 2, 1, 1, 1] + # [1, 1, 2, 1, 1, 1, 1(position)] # Creates an ACL rule # Allows traffic from server_1 to node_1 on port FTP - action = 7 + action = 56 elif step == 7: # [1, 1, 2, 0] Node Action # Sets Node 1 Hardware State to OFF # Does not resolve any service - action = 16 + action = 128 # Run the simulation step on the live environment obs, reward, done, info = env.step(action) @@ -48,7 +48,8 @@ def test_single_action_space_is_valid(): """Test to ensure the blue agent is using the ACL action space and is carrying out both kinds of operations.""" env = _get_primaite_env_from_config( training_config_path=TEST_CONFIG_ROOT / "single_action_space_main_config.yaml", - lay_down_config_path=TEST_CONFIG_ROOT / "single_action_space_lay_down_config.yaml", + lay_down_config_path=TEST_CONFIG_ROOT + / "single_action_space_lay_down_config.yaml", ) run_generic_set_actions(env) @@ -77,8 +78,10 @@ def test_single_action_space_is_valid(): def test_agent_is_executing_actions_from_both_spaces(): """Test to ensure the blue agent is carrying out both kinds of operations (NODE & ACL).""" env = _get_primaite_env_from_config( - training_config_path=TEST_CONFIG_ROOT / "single_action_space_fixed_blue_actions_main_config.yaml", - lay_down_config_path=TEST_CONFIG_ROOT / "single_action_space_lay_down_config.yaml", + training_config_path=TEST_CONFIG_ROOT + / "single_action_space_fixed_blue_actions_main_config.yaml", + lay_down_config_path=TEST_CONFIG_ROOT + / "single_action_space_lay_down_config.yaml", ) # Run environment with specified fixed blue agent actions only run_generic_set_actions(env) From c6a947fbaff1c3265dd09bb425d84c279bf7329a Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 13 Jun 2023 16:23:32 +0100 Subject: [PATCH 08/37] 901 - started testing for observation space --- src/primaite/config/training_config.py | 2 + src/primaite/environment/observations.py | 13 +- src/primaite/environment/primaite_env.py | 3 + .../main_config_ACCESS_CONTROL_LIST.yaml | 129 +++++++++--------- tests/test_observation_space.py | 39 ++++++ 5 files changed, 116 insertions(+), 70 deletions(-) diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 9a21d087..14102432 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -1,4 +1,5 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. +import logging from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, Final, Optional, Union @@ -9,6 +10,7 @@ from primaite import USERS_CONFIG_DIR, getLogger from primaite.common.enums import ActionType _LOGGER = getLogger(__name__) +logging.basicConfig(level=logging.DEBUG, format="%(message)s") _EXAMPLE_TRAINING: Final[Path] = USERS_CONFIG_DIR / "example_config" / "training" diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index 96df1f60..fe43c9e3 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -331,7 +331,6 @@ class AccessControlList(AbstractObservationComponent): """ # Terms (for ACL observation space): - # [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) # [0, 1] - Permission (0 = DENY, 1 = ALLOW) # [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) @@ -352,7 +351,7 @@ class AccessControlList(AbstractObservationComponent): len(env.services_list), len(env.ports_list), ] - shape = acl_shape * self.env.max_acl_rules + shape = acl_shape * self.env.max_number_acl_rules # 2. Create Observation space self.space = spaces.MultiDiscrete(shape) @@ -360,9 +359,6 @@ class AccessControlList(AbstractObservationComponent): # 3. Initialise observation with zeroes self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) - # Dictionary to map services to numbers for obs space - self.services_dict = {} - def update(self): """Update the observation based on current environment state. @@ -380,7 +376,6 @@ class AccessControlList(AbstractObservationComponent): permission_int = 0 else: permission_int = 1 - if source_ip == "ANY": source_ip_int = 0 else: @@ -393,9 +388,10 @@ class AccessControlList(AbstractObservationComponent): protocol_int = 0 else: try: - protocol_int = Protocol[protocol] + protocol_int = Protocol[protocol].value except AttributeError: _LOGGER.info(f"Service {protocol} could not be found") + protocol_int = -1 if port == "ANY": port_int = 0 else: @@ -423,7 +419,8 @@ class AccessControlList(AbstractObservationComponent): Resolves IP address -> x (node id e.g. 1 or 2 or 3 or 4) for observation space """ - for key, node in self.env.nodes: + print(type(self.env.nodes)) + for key, node in self.env.nodes.items(): if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): if node.ip_address == ip_address: return key diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 39006259..783b4267 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -122,6 +122,9 @@ class Primaite(Env): self.training_config.implicit_acl_rule, self.training_config.max_number_acl_rules, ) + # Sets limit for number of ACL rules in environment + self.max_number_acl_rules = self.training_config.max_number_acl_rules + # Create a list of services (enums) self.services_list = [] diff --git a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml index 856e963d..7aa30205 100644 --- a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml +++ b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml @@ -5,14 +5,16 @@ # "STABLE_BASELINES3_PPO" # "STABLE_BASELINES3_A2C" # "GENERIC" -agentIdentifier: NONE +agent_identifier: STABLE_BASELINES3_A2C +# Sets How the Action Space is defined: +# "NODE" +# "ACL" +# "ANY" node and acl actions +action_type: ANY # Number of episodes to run per session -observationSpace: - components: - - name: ACCESS_CONTROL_LIST - options: - implicit_acl_rule: DENY - max_number_of_acl_rules: 10 +num_episodes: 1 +# Number of time_steps per episode +num_steps: 5 # Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) apply_implicit_rule: True @@ -20,83 +22,86 @@ apply_implicit_rule: True implicit_acl_rule: DENY # Total number of ACL rules allowed in the environment max_number_acl_rules: 10 -numEpisodes: 1 + +observation_space: + components: + - name: ACCESS_CONTROL_LIST + # Time delay between steps (for generic agents) -timeDelay: 1 -# Filename of the scenario / laydown -configFilename: one_node_states_on_off_lay_down_config.yaml +time_delay: 1 + # Type of session to be run (TRAINING or EVALUATION) -sessionType: TRAINING +session_type: TRAINING # Determine whether to load an agent from file -loadAgent: False +load_agent: False # File path and file name of agent if you're loading one in -agentLoadFile: C:\[Path]\[agent_saved_filename.zip] +agent_load_file: C:\[Path]\[agent_saved_filename.zip] # Environment config values # The high value for the observation space -observationSpaceHighValue: 1_000_000_000 +observation_space_high_value: 1_000_000_000 # Reward values # Generic -allOk: 0 +all_ok: 0 # Node Hardware State -offShouldBeOn: -10 -offShouldBeResetting: -5 -onShouldBeOff: -2 -onShouldBeResetting: -5 -resettingShouldBeOn: -5 -resettingShouldBeOff: -2 +off_should_be_on: -10 +off_should_be_resetting: -5 +on_should_be_off: -2 +on_should_be_resetting: -5 +resetting_should_be_on: -5 +resetting_should_be_off: -2 resetting: -3 # Node Software or Service State -goodShouldBePatching: 2 -goodShouldBeCompromised: 5 -goodShouldBeOverwhelmed: 5 -patchingShouldBeGood: -5 -patchingShouldBeCompromised: 2 -patchingShouldBeOverwhelmed: 2 +good_should_be_patching: 2 +good_should_be_compromised: 5 +good_should_be_overwhelmed: 5 +patching_should_be_good: -5 +patching_should_be_compromised: 2 +patching_should_be_overwhelmed: 2 patching: -3 -compromisedShouldBeGood: -20 -compromisedShouldBePatching: -20 -compromisedShouldBeOverwhelmed: -20 +compromised_should_be_good: -20 +compromised_should_be_patching: -20 +compromised_should_be_overwhelmed: -20 compromised: -20 -overwhelmedShouldBeGood: -20 -overwhelmedShouldBePatching: -20 -overwhelmedShouldBeCompromised: -20 +overwhelmed_should_be_good: -20 +overwhelmed_should_be_patching: -20 +overwhelmed_should_be_compromised: -20 overwhelmed: -20 # Node File System State -goodShouldBeRepairing: 2 -goodShouldBeRestoring: 2 -goodShouldBeCorrupt: 5 -goodShouldBeDestroyed: 10 -repairingShouldBeGood: -5 -repairingShouldBeRestoring: 2 -repairingShouldBeCorrupt: 2 -repairingShouldBeDestroyed: 0 +good_should_be_repairing: 2 +good_should_be_restoring: 2 +good_should_be_corrupt: 5 +good_should_be_destroyed: 10 +repairing_should_be_good: -5 +repairing_should_be_restoring: 2 +repairing_should_be_corrupt: 2 +repairing_should_be_destroyed: 0 repairing: -3 -restoringShouldBeGood: -10 -restoringShouldBeRepairing: -2 -restoringShouldBeCorrupt: 1 -restoringShouldBeDestroyed: 2 +restoring_should_be_good: -10 +restoring_should_be_repairing: -2 +restoring_should_be_corrupt: 1 +restoring_should_be_destroyed: 2 restoring: -6 -corruptShouldBeGood: -10 -corruptShouldBeRepairing: -10 -corruptShouldBeRestoring: -10 -corruptShouldBeDestroyed: 2 +corrupt_should_be_good: -10 +corrupt_should_be_repairing: -10 +corrupt_should_be_restoring: -10 +corrupt_should_be_destroyed: 2 corrupt: -10 -destroyedShouldBeGood: -20 -destroyedShouldBeRepairing: -20 -destroyedShouldBeRestoring: -20 -destroyedShouldBeCorrupt: -20 +destroyed_should_be_good: -20 +destroyed_should_be_repairing: -20 +destroyed_should_be_restoring: -20 +destroyed_should_be_corrupt: -20 destroyed: -20 scanning: -2 # IER status -redIerRunning: -5 -greenIerBlocked: -10 +red_ier_running: -5 +green_ier_blocked: -10 # Patching / Reset durations -osPatchingDuration: 5 # The time taken to patch the OS -nodeResetDuration: 5 # The time taken to reset a node (hardware) -servicePatchingDuration: 5 # The time taken to patch a service -fileSystemRepairingLimit: 5 # The time take to repair the file system -fileSystemRestoringLimit: 5 # The time take to restore the file system -fileSystemScanningLimit: 5 # The time taken to scan the file system +os_patching_duration: 5 # The time taken to patch the OS +node_reset_duration: 5 # The time taken to reset a node (hardware) +service_patching_duration: 5 # The time taken to patch a service +file_system_repairing_limit: 5 # The time take to repair the file system +file_system_restoring_limit: 5 # The time take to restore the file system +file_system_scanning_limit: 5 # The time taken to scan the file system diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 4e7186b5..4e8df7e1 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -222,3 +222,42 @@ class TestLinkTrafficLevels: # (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%) print(obs) assert np.array_equal(obs, [6, 0, 6, 0]) + + +@pytest.mark.env_config_paths( + dict( + training_config_path=TEST_CONFIG_ROOT + / "obs_tests/main_config_ACCESS_CONTROL_LIST.yaml", + lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", + ) +) +class TestAccessControlList: + """Test the AccessControlList observation component (in isolation).""" + + def test_obs_shape(self, env: Primaite): + """Try creating env with MultiDiscrete observation space.""" + env.update_environent_obs() + + # we have two ACLs + assert env.env_obs.shape == (5, 2) + + def test_values(self, env: Primaite): + """Test that traffic values are encoded correctly. + + The laydown has: + * two services + * three nodes + * two links + * an IER trying to send 999 bits of data over both links the whole time (via the first service) + * link bandwidth of 1000, therefore the utilisation is 99.9% + """ + obs, reward, done, info = env.step(0) + obs, reward, done, info = env.step(0) + + # the observation space has combine_service_traffic set to False, so the space has this format: + # [link1_service1, link1_service2, link2_service1, link2_service2] + # we send 999 bits of data via link1 and link2 on service 1. + # therefore the first and third elements should be 6 and all others 0 + # (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%) + print(obs) + assert np.array_equal(obs, [6, 0, 6, 0]) From df42a791c9bd6729c8ab36d3e92a62398bd8b8c5 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 20 Jun 2023 11:47:20 +0100 Subject: [PATCH 09/37] 901 - changed ACL instantiation and changed acl t private _acl (list not dict) attribute, added laydown_ACL.yaml for testing, fixed encoding of acl rules to integers for obs space, added ACL position to node action space and added generic test where agents adds two ACL rules. --- src/primaite/acl/access_control_list.py | 66 +++++--- src/primaite/config/training_config.py | 2 - src/primaite/environment/observations.py | 145 ++++++++++-------- src/primaite/environment/primaite_env.py | 15 +- tests/config/obs_tests/laydown_ACL.yaml | 86 +++++++++++ ..._space_fixed_blue_actions_main_config.yaml | 13 ++ tests/test_acl.py | 14 +- tests/test_observation_space.py | 67 +++++++- tests/test_single_action_space.py | 4 +- 9 files changed, 305 insertions(+), 107 deletions(-) create mode 100644 tests/config/obs_tests/laydown_ACL.yaml diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index d75b9756..219ba002 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -11,21 +11,43 @@ _LOGGER: Final[logging.Logger] = logging.getLogger(__name__) class AccessControlList: """Access Control List class.""" - def __init__(self, implicit_permission, max_acl_rules): + def __init__(self, apply_implicit_rule, implicit_permission, max_acl_rules): """Init.""" + # Bool option in main_config to decide to use implicit rule or not + self.apply_implicit_rule: bool = apply_implicit_rule # Implicit ALLOW or DENY firewall spec # Last rule in the ACL list - self.acl_implicit_rule = implicit_permission - # Create implicit rule based on input - if self.acl_implicit_rule == "DENY": - implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") - else: - implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") - + self.acl_implicit_permission = implicit_permission # Maximum number of ACL Rules in ACL self.max_acl_rules: int = max_acl_rules # A list of ACL Rules - self.acl: List[ACLRule] = [implicit_rule] + self._acl: List[ACLRule] = [] + # Implicit rule + + @property + def acl_implicit_rule(self): + """ACL implicit rule class attribute with added logic to change it depending on option in main_config.""" + # Create implicit rule based on input + if self.apply_implicit_rule: + if self.acl_implicit_permission == "DENY": + return ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") + elif self.acl_implicit_permission == "ALLOW": + return ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") + else: + return None + else: + return None + + @property + def acl(self): + """Public access method for private _acl. + + Adds implicit rule to end of acl list and + Pads out rest of list (if empty) with -1. + """ + if self.acl_implicit_rule is not None: + acl_list = self._acl + [self.acl_implicit_rule] + return acl_list + [-1] * (self.max_acl_rules - len(acl_list)) def check_address_match(self, _rule, _source_ip_address, _dest_ip_address): """ @@ -85,7 +107,9 @@ class AccessControlList: # If there has been no rule to allow the IER through, it will return a blocked signal by default return True - def add_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port, _position): + def add_rule( + self, _permission, _source_ip, _dest_ip, _protocol, _port, _position=None + ): """ Adds a new rule. @@ -99,18 +123,22 @@ class AccessControlList: """ position_index = int(_position) new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - if len(self.acl) < self.max_acl_rules: - if len(self.acl) > position_index > -1: - try: - self.acl.insert(position_index, new_rule) - except Exception: + print(len(self._acl)) + if len(self._acl) + 1 < self.max_acl_rules: + if _position is not None: + if self.max_acl_rules - 1 > position_index > -1: + try: + self._acl.insert(position_index, new_rule) + except Exception: + _LOGGER.info( + f"New Rule could NOT be added to list at position {position_index}." + ) + else: _LOGGER.info( - f"New Rule could NOT be added to list at position {position_index}." + f"Position {position_index} is an invalid index for list/overwrites implicit firewall rule" ) else: - _LOGGER.info( - f"Position {position_index} is an invalid index for list and/or overwrites implicit firewall rule" - ) + self.acl.append(new_rule) else: _LOGGER.info( f"The ACL list is FULL." diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 14102432..9a21d087 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -1,5 +1,4 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. -import logging from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, Final, Optional, Union @@ -10,7 +9,6 @@ from primaite import USERS_CONFIG_DIR, getLogger from primaite.common.enums import ActionType _LOGGER = getLogger(__name__) -logging.basicConfig(level=logging.DEBUG, format="%(message)s") _EXAMPLE_TRAINING: Final[Path] = USERS_CONFIG_DIR / "example_config" / "training" diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index fe43c9e3..eb7ad2bf 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Dict, Final, List, Tuple, Union import numpy as np from gym import spaces +from primaite.acl.acl_rule import ACLRule from primaite.common.enums import ( FileSystemState, HardwareState, @@ -22,7 +23,6 @@ from primaite.nodes.service_node import ServiceNode if TYPE_CHECKING: from primaite.environment.primaite_env import Primaite - _LOGGER = logging.getLogger(__name__) @@ -346,16 +346,19 @@ class AccessControlList(AbstractObservationComponent): # 1. Define the shape of your observation space component acl_shape = [ len(RulePermissionType), - len(env.nodes), - len(env.nodes), + len(env.nodes) + 1, + len(env.nodes) + 1, len(env.services_list), len(env.ports_list), + env.max_number_acl_rules, ] + len(acl_shape) + # shape = acl_shape shape = acl_shape * self.env.max_number_acl_rules # 2. Create Observation space self.space = spaces.MultiDiscrete(shape) - + print("obs space:", self.space) # 3. Initialise observation with zeroes self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) @@ -365,67 +368,85 @@ class AccessControlList(AbstractObservationComponent): The structure of the observation space is described in :class:`.AccessControlList` """ obs = [] - for acl_rule in self.env.acl.acl: - permission = acl_rule.permission - source_ip = acl_rule.source_ip - dest_ip = acl_rule.dest_ip - protocol = acl_rule.protocol - port = acl_rule.port - position = self.env.acl.acl.index(acl_rule) - if permission == "DENY": - permission_int = 0 - else: - permission_int = 1 - if source_ip == "ANY": - source_ip_int = 0 - else: - source_ip_int = self.obtain_node_id_using_ip(source_ip) - if dest_ip == "ANY": - dest_ip_int = 0 - else: - dest_ip_int = self.obtain_node_id_using_ip(dest_ip) - if protocol == "ANY": - protocol_int = 0 - else: - try: - protocol_int = Protocol[protocol].value - except AttributeError: - _LOGGER.info(f"Service {protocol} could not be found") - protocol_int = -1 - if port == "ANY": - port_int = 0 - else: - if port in self.env.ports_list: - port_int = self.env.ports_list.index(port) + + for index in range(len(self.env.acl.acl)): + acl_rule = self.env.acl.acl[index] + if isinstance(acl_rule, ACLRule): + permission = acl_rule.permission + source_ip = acl_rule.source_ip + dest_ip = acl_rule.dest_ip + protocol = acl_rule.protocol + port = acl_rule.port + position = index + + source_ip_int = -1 + dest_ip_int = -1 + if permission == "DENY": + permission_int = 0 else: - _LOGGER.info(f"Port {port} could not be found.") + permission_int = 1 + if source_ip == "ANY": + source_ip_int = 0 + else: + nodes = list(self.env.nodes.values()) + for node in nodes: + # print(node.ip_address, source_ip, node.ip_address == source_ip) + if ( + isinstance(node, ServiceNode) + or isinstance(node, ActiveNode) + ) and node.ip_address == source_ip: + source_ip_int = node.node_id + break + if dest_ip == "ANY": + dest_ip_int = 0 + else: + nodes = list(self.env.nodes.values()) + for node in nodes: + if ( + isinstance(node, ServiceNode) + or isinstance(node, ActiveNode) + ) and node.ip_address == dest_ip: + dest_ip_int = node.node_id + if protocol == "ANY": + protocol_int = 0 + else: + try: + protocol_int = Protocol[protocol].value + except AttributeError: + _LOGGER.info(f"Service {protocol} could not be found") + protocol_int = -1 + if port == "ANY": + port_int = 0 + else: + if port in self.env.ports_list: + port_int = self.env.ports_list.index(port) + else: + _LOGGER.info(f"Port {port} could not be found.") - print(permission_int, source_ip, dest_ip, protocol_int, port_int, position) - obs.extend( - [ - permission_int, - source_ip_int, - dest_ip_int, - protocol_int, - port_int, - position, - ] - ) + # Either do the multiply on the obs space + # Change the obs to + if source_ip_int != -1 and dest_ip_int != -1: + items_to_add = [ + permission_int, + source_ip_int, + dest_ip_int, + protocol_int, + port_int, + position, + ] + position = position * 6 + for item in items_to_add: + obs.insert(position, int(item)) + position += 1 + else: + items_to_add = [-1, -1, -1, -1, -1, index] + position = index * 6 + for item in items_to_add: + obs.insert(position, int(item)) + position += 1 - self.current_observation[:] = obs - - def obtain_node_id_using_ip(self, ip_address): - """Uses IP address of Nodes to find the ID. - - Resolves IP address -> x (node id e.g. 1 or 2 or 3 or 4) for observation space - """ - print(type(self.env.nodes)) - for key, node in self.env.nodes.items(): - if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): - if node.ip_address == ip_address: - return key - _LOGGER.info(f"Node ID was not found from IP Address {ip_address}") - return -1 + self.current_observation = obs + print("current observation space:", self.current_observation) class ObservationsHandler: diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 783b4267..c0ee04f5 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -45,7 +45,7 @@ from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_nod from primaite.transactions.transaction import Transaction _LOGGER = logging.getLogger(__name__) -_LOGGER.setLevel(logging.INFO) +# _LOGGER.setLevel(logging.INFO) class Primaite(Env): @@ -119,6 +119,7 @@ class Primaite(Env): # Create the Access Control List self.acl = AccessControlList( + self.training_config.apply_implicit_rule, self.training_config.implicit_acl_rule, self.training_config.max_number_acl_rules, ) @@ -546,6 +547,7 @@ class Primaite(Env): action_destination_ip = readable_action[3] action_protocol = readable_action[4] action_port = readable_action[5] + acl_rule_position = readable_action[6] if action_decision == 0: # It's decided to do nothing @@ -595,6 +597,7 @@ class Primaite(Env): acl_rule_destination, acl_rule_protocol, acl_rule_port, + acl_rule_position, ) elif action_decision == 2: # Remove the rule @@ -1172,13 +1175,9 @@ class Primaite(Env): # [0, num ports] - Port (0 = any, then 1 -> x resolving to port) # [0, max acl rules - 1] - Position (0 = first index, then 1 -> x index resolving to acl rule in acl list) # reserve 0 action to be a nothing action - actions = {0: [0, 0, 0, 0, 0, 0]} + actions = {0: [0, 0, 0, 0, 0, 0, 0]} action_key = 1 - print( - "what is this primaite_env.py 1177", - self.training_config.max_number_acl_rules - 1, - ) # 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE for action_decision in range(3): # 2 possible action permissions 0 = DENY, 1 = CREATE @@ -1188,9 +1187,7 @@ class Primaite(Env): for dest_ip in range(self.num_nodes + 1): for protocol in range(self.num_services + 1): for port in range(self.num_ports + 1): - for position in range( - self.training_config.max_number_acl_rules - 1 - ): + for position in range(self.max_number_acl_rules - 1): action = [ action_decision, action_permission, diff --git a/tests/config/obs_tests/laydown_ACL.yaml b/tests/config/obs_tests/laydown_ACL.yaml new file mode 100644 index 00000000..cffd8b1c --- /dev/null +++ b/tests/config/obs_tests/laydown_ACL.yaml @@ -0,0 +1,86 @@ +- item_type: PORTS + ports_list: + - port: '80' + - port: '21' +- item_type: SERVICES + service_list: + - name: TCP + - name: FTP + +######################################## +# Nodes +- item_type: NODE + node_id: '1' + name: PC1 + node_class: SERVICE + node_type: COMPUTER + priority: P5 + hardware_state: 'ON' + ip_address: 192.168.1.1 + software_state: COMPROMISED + file_system_state: GOOD + services: + - name: TCP + port: '80' + state: GOOD + - name: FTP + port: '21' + state: GOOD +- item_type: NODE + node_id: '2' + name: SERVER + node_class: SERVICE + node_type: SERVER + priority: P5 + hardware_state: 'ON' + ip_address: 192.168.1.2 + software_state: GOOD + file_system_state: GOOD + services: + - name: TCP + port: '80' + state: GOOD + - name: FTP + port: '21' + state: OVERWHELMED +- item_type: NODE + node_id: '3' + name: SWITCH1 + node_class: ACTIVE + node_type: SWITCH + priority: P2 + hardware_state: 'ON' + ip_address: 192.168.1.3 + software_state: GOOD + file_system_state: GOOD + +######################################## +# Links +- item_type: LINK + id: '4' + name: link1 + bandwidth: 1000 + source: '1' + destination: '3' +- item_type: LINK + id: '5' + name: link2 + bandwidth: 1000 + source: '3' + destination: '2' + +######################################### +# IERS +- item_type: GREEN_IER + id: '5' + start_step: 0 + end_step: 5 + load: 999 + protocol: TCP + port: '80' + source: '1' + destination: '2' + mission_criticality: 5 + +######################################### +# ACL Rules diff --git a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml index 5c5db582..e2718c53 100644 --- a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml +++ b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml @@ -24,6 +24,19 @@ load_agent: False # File path and file name of agent if you're loading one in agent_load_file: C:\[Path]\[agent_saved_filename.zip] + + +# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) +apply_implicit_rule: True +# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) +implicit_acl_rule: DENY +# Total number of ACL rules allowed in the environment +max_number_acl_rules: 10 + +observation_space: + components: + - name: ACCESS_CONTROL_LIST + # Environment config values # The high value for the observation space observation_space_high_value: 1000000000 diff --git a/tests/test_acl.py b/tests/test_acl.py index e99f5ee0..f790a5cf 100644 --- a/tests/test_acl.py +++ b/tests/test_acl.py @@ -7,7 +7,7 @@ from primaite.acl.acl_rule import ACLRule def test_acl_address_match_1(): """Test that matching IP addresses produce True.""" - acl = AccessControlList("DENY", 10) + acl = AccessControlList(True, "DENY", 10) rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80") @@ -16,7 +16,7 @@ def test_acl_address_match_1(): def test_acl_address_match_2(): """Test that mismatching IP addresses produce False.""" - acl = AccessControlList("DENY", 10) + acl = AccessControlList(True, "DENY", 10) rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80") @@ -25,7 +25,7 @@ def test_acl_address_match_2(): def test_acl_address_match_3(): """Test the ANY condition for source IP addresses produce True.""" - acl = AccessControlList("DENY", 10) + acl = AccessControlList(True, "DENY", 10) rule = ACLRule("ALLOW", "ANY", "192.168.1.2", "TCP", "80") @@ -34,7 +34,7 @@ def test_acl_address_match_3(): def test_acl_address_match_4(): """Test the ANY condition for dest IP addresses produce True.""" - acl = AccessControlList("DENY", 10) + acl = AccessControlList(True, "DENY", 10) rule = ACLRule("ALLOW", "192.168.1.1", "ANY", "TCP", "80") @@ -44,7 +44,7 @@ def test_acl_address_match_4(): def test_check_acl_block_affirmative(): """Test the block function (affirmative).""" # Create the Access Control List - acl = AccessControlList("DENY", 10) + acl = AccessControlList(True, "DENY", 10) # Create a rule acl_rule_permission = "ALLOW" @@ -68,7 +68,7 @@ def test_check_acl_block_affirmative(): def test_check_acl_block_negative(): """Test the block function (negative).""" # Create the Access Control List - acl = AccessControlList("DENY", 10) + acl = AccessControlList(True, "DENY", 10) # Create a rule acl_rule_permission = "DENY" @@ -93,7 +93,7 @@ def test_check_acl_block_negative(): def test_rule_hash(): """Test the rule hash.""" # Create the Access Control List - acl = AccessControlList("DENY", 10) + acl = AccessControlList(True, "DENY", 10) rule = ACLRule("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80") hash_value_local = hash(rule) diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 4e8df7e1..5408bee6 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -1,4 +1,7 @@ """Test env creation and behaviour with different observation spaces.""" + +import time + import numpy as np import pytest @@ -12,6 +15,46 @@ from tests import TEST_CONFIG_ROOT from tests.conftest import _get_primaite_env_from_config +def run_generic_set_actions(env: Primaite): + """Run against a generic agent with specified blue agent actions.""" + # Reset the environment at the start of the episode + # env.reset() + training_config = env.training_config + for episode in range(0, training_config.num_episodes): + for step in range(0, training_config.num_steps): + # Send the observation space to the agent to get an action + # TEMP - random action for now + # action = env.blue_agent_action(obs) + action = 0 + print("\nStep:", step) + if step == 5: + # [1, 1, 2, 1, 1, 1, 2] ACL Action + # Creates an ACL rule + # Allows traffic from SERVER to PC1 on port TCP 80 and place ACL at position 2 + action = 291 + elif step == 7: + # [1, 1, 3, 1, 2, 2, 1] ACL Action + # Creates an ACL rule + # Allows traffic from PC1 to SWITCH 1 on port UDP at position 1 + action = 425 + # Run the simulation step on the live environment + obs, reward, done, info = env.step(action) + # Update observations space and return + env.update_environent_obs() + + # Break if done is True + if done: + break + + # Introduce a delay between steps + time.sleep(training_config.time_delay / 1000) + + # Reset the environment at the end of the episode + # env.reset() + + # env.close() + + @pytest.fixture def env(request): """Build Primaite environment for integration tests of observation space.""" @@ -131,11 +174,11 @@ class TestNodeLinkTable: assert np.array_equal( obs, [ - [1, 1, 3, 1, 1, 1], - [2, 1, 1, 1, 1, 4], - [3, 1, 1, 1, 0, 0], - [4, 0, 0, 0, 999, 0], - [5, 0, 0, 0, 999, 0], + [1, 1, 3, 1, 1, 1, 0], + [2, 1, 1, 1, 1, 4, 1], + [3, 1, 1, 1, 0, 0, 2], + [4, 0, 0, 0, 999, 0, 3], + [5, 0, 0, 0, 999, 0, 4], ], ) @@ -260,4 +303,16 @@ class TestAccessControlList: # therefore the first and third elements should be 6 and all others 0 # (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%) print(obs) - assert np.array_equal(obs, [6, 0, 6, 0]) + assert np.array_equal(obs, []) + + def test_observation_space(self): + """Test observation space is what is expected when an agent adds ACLs during an episode.""" + # Used to use env from test fixture but AtrributeError function object has no 'training_config' + env = _get_primaite_env_from_config( + training_config_path=TEST_CONFIG_ROOT + / "single_action_space_fixed_blue_actions_main_config.yaml", + lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown_ACL.yaml", + ) + run_generic_set_actions(env) + + # print("observation space",env.observation_space) diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index 4d6136a9..78764976 100644 --- a/tests/test_single_action_space.py +++ b/tests/test_single_action_space.py @@ -66,7 +66,7 @@ def test_single_action_space_is_valid(): if len(dict_item) == 4: contains_node_actions = True # Link action detected - elif len(dict_item) == 6: + elif len(dict_item) == 7: contains_acl_actions = True # If both are there then the ANY action type is working if contains_node_actions and contains_acl_actions: @@ -92,7 +92,7 @@ def test_agent_is_executing_actions_from_both_spaces(): access_control_list = env.acl # Use the Access Control List object acl object attribute to get dictionary # Use dictionary.values() to get total list of all items in the dictionary - acl_rules_list = access_control_list.acl.values() + acl_rules_list = access_control_list.acl # Length of this list tells you how many items are in the dictionary # This number is the frequency of Access Control Rules in the environment # In the scenario, we specified that the agent should create only 1 acl rule From ef4d2c6cdd3ede32c3b43b8d5a485b55f01d2961 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 27 Jun 2023 11:43:33 +0100 Subject: [PATCH 10/37] 901 - fixed test_single_action_space.py to reflect new acl structure and added new acl_implicit_rule class attribute --- src/primaite/acl/access_control_list.py | 14 +++----------- src/primaite/environment/observations.py | 2 -- .../single_action_space_lay_down_config.yaml | 14 +++++++------- tests/test_single_action_space.py | 10 +++++++--- 4 files changed, 17 insertions(+), 23 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 219ba002..9cc1225a 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -23,20 +23,12 @@ class AccessControlList: # A list of ACL Rules self._acl: List[ACLRule] = [] # Implicit rule - - @property - def acl_implicit_rule(self): - """ACL implicit rule class attribute with added logic to change it depending on option in main_config.""" - # Create implicit rule based on input + self.acl_implicit_rule = None if self.apply_implicit_rule: if self.acl_implicit_permission == "DENY": - return ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") + self.acl_implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") elif self.acl_implicit_permission == "ALLOW": - return ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") - else: - return None - else: - return None + self.acl_implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") @property def acl(self): diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index eb7ad2bf..2aacda8f 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -352,7 +352,6 @@ class AccessControlList(AbstractObservationComponent): len(env.ports_list), env.max_number_acl_rules, ] - len(acl_shape) # shape = acl_shape shape = acl_shape * self.env.max_number_acl_rules @@ -446,7 +445,6 @@ class AccessControlList(AbstractObservationComponent): position += 1 self.current_observation = obs - print("current observation space:", self.current_observation) class ObservationsHandler: diff --git a/tests/config/single_action_space_lay_down_config.yaml b/tests/config/single_action_space_lay_down_config.yaml index c80c0bab..0b947a5f 100644 --- a/tests/config/single_action_space_lay_down_config.yaml +++ b/tests/config/single_action_space_lay_down_config.yaml @@ -1,9 +1,9 @@ - item_type: PORTS ports_list: - - port: '21' + - port: '80' - item_type: SERVICES service_list: - - name: ftp + - name: TCP - item_type: NODE node_id: '1' name: node @@ -15,8 +15,8 @@ software_state: GOOD file_system_state: GOOD services: - - name: ftp - port: '21' + - name: TCP + port: '80' state: COMPROMISED - item_type: NODE node_id: '2' @@ -29,8 +29,8 @@ software_state: GOOD file_system_state: GOOD services: - - name: ftp - port: '21' + - name: TCP + port: '80' state: COMPROMISED - item_type: POSITION positions: @@ -45,7 +45,7 @@ start_step: 2 end_step: 15 load: 1000 - protocol: ftp + protocol: TCP port: CORRUPT source: '1' destination: '2' diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index 78764976..f12d160c 100644 --- a/tests/test_single_action_space.py +++ b/tests/test_single_action_space.py @@ -9,7 +9,7 @@ from tests.conftest import _get_primaite_env_from_config def run_generic_set_actions(env: Primaite): """Run against a generic agent with specified blue agent actions.""" # Reset the environment at the start of the episode - # env.reset() + env.reset() training_config = env.training_config for episode in range(0, training_config.num_episodes): for step in range(0, training_config.num_steps): @@ -96,7 +96,11 @@ def test_agent_is_executing_actions_from_both_spaces(): # Length of this list tells you how many items are in the dictionary # This number is the frequency of Access Control Rules in the environment # In the scenario, we specified that the agent should create only 1 acl rule - num_of_rules = len(acl_rules_list) + # This 1 rule added to the implicit deny means there should be 2 rules in total. + rules_count = 0 + for rule in acl_rules_list: + if rule != -1: + rules_count += 1 # Therefore these statements below MUST be true assert computer_node_hardware_state == HardwareState.OFF - assert num_of_rules == 1 + assert rules_count == 2 From 3f440c0a281917f0c9c9469f2651534e84313ed8 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Wed, 5 Jul 2023 09:08:03 +0100 Subject: [PATCH 11/37] 901 - updated observations.py to change and add new mapping of ACL rules to represent no rule present in list --- src/primaite/acl/access_control_list.py | 27 ++--- src/primaite/common/enums.py | 7 +- .../training/training_config_main.yaml | 6 +- src/primaite/environment/observations.py | 98 +++++++++---------- src/primaite/environment/primaite_env.py | 3 +- ..._space_fixed_blue_actions_main_config.yaml | 4 +- tests/test_observation_space.py | 26 +++-- 7 files changed, 95 insertions(+), 76 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 9cc1225a..9e51e066 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -39,7 +39,9 @@ class AccessControlList: """ if self.acl_implicit_rule is not None: acl_list = self._acl + [self.acl_implicit_rule] - return acl_list + [-1] * (self.max_acl_rules - len(acl_list)) + else: + acl_list = self._acl + return acl_list + [None] * (self.max_acl_rules - len(acl_list)) def check_address_match(self, _rule, _source_ip_address, _dest_ip_address): """ @@ -86,15 +88,18 @@ class AccessControlList: Indicates block if all conditions are satisfied. """ for rule in self.acl: - if self.check_address_match(rule, _source_ip_address, _dest_ip_address): - if ( - rule.get_protocol() == _protocol or rule.get_protocol() == "ANY" - ) and (str(rule.get_port()) == str(_port) or rule.get_port() == "ANY"): - # There's a matching rule. Get the permission - if rule.get_permission() == "DENY": - return True - elif rule.get_permission() == "ALLOW": - return False + if isinstance(rule, ACLRule): + if self.check_address_match(rule, _source_ip_address, _dest_ip_address): + if ( + rule.get_protocol() == _protocol or rule.get_protocol() == "ANY" + ) and ( + str(rule.get_port()) == str(_port) or rule.get_port() == "ANY" + ): + # There's a matching rule. Get the permission + if rule.get_permission() == "DENY": + return True + elif rule.get_permission() == "ALLOW": + return False # If there has been no rule to allow the IER through, it will return a blocked signal by default return True @@ -115,7 +120,6 @@ class AccessControlList: """ position_index = int(_position) new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - print(len(self._acl)) if len(self._acl) + 1 < self.max_acl_rules: if _position is not None: if self.max_acl_rules - 1 > position_index > -1: @@ -136,6 +140,7 @@ class AccessControlList: f"The ACL list is FULL." f"The list of ACLs has length {len(self.acl)} and it has a max capacity of {self.max_acl_rules}." ) + # print("length of this list", len(self._acl)) def remove_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port): """ diff --git a/src/primaite/common/enums.py b/src/primaite/common/enums.py index 6a0c8f29..ad6c84a1 100644 --- a/src/primaite/common/enums.py +++ b/src/primaite/common/enums.py @@ -131,7 +131,8 @@ class LinkStatus(Enum): class RulePermissionType(Enum): - """Implicit firewall rule.""" + """Any firewall rule type.""" - DENY = 0 - ALLOW = 1 + NA = 0 + DENY = 1 + ALLOW = 2 diff --git a/src/primaite/config/_package_data/training/training_config_main.yaml b/src/primaite/config/_package_data/training/training_config_main.yaml index d01f51f3..233c299e 100644 --- a/src/primaite/config/_package_data/training/training_config_main.yaml +++ b/src/primaite/config/_package_data/training/training_config_main.yaml @@ -5,14 +5,14 @@ # "STABLE_BASELINES3_PPO" # "STABLE_BASELINES3_A2C" # "GENERIC" -agent_identifier: STABLE_BASELINES3_A2C +agent_identifier: STABLE_BASELINES3_PPO # Sets How the Action Space is defined: # "NODE" # "ACL" # "ANY" node and acl actions -action_type: NODE +action_type: ACL # Number of episodes to run per session -num_episodes: 10 +num_episodes: 1000 # Number of time_steps per episode num_steps: 256 # Time delay between steps (for generic agents) diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index 2aacda8f..d254598b 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -10,7 +10,6 @@ from primaite.acl.acl_rule import ACLRule from primaite.common.enums import ( FileSystemState, HardwareState, - Protocol, RulePermissionType, SoftwareState, ) @@ -330,13 +329,14 @@ class AccessControlList(AbstractObservationComponent): ] """ + 0, # Terms (for ACL observation space): - # [0, 1] - Permission (0 = DENY, 1 = ALLOW) - # [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) - # [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) - # [0, num ports] - Port (0 = any, then 1 -> x resolving to port) - # [0, max acl rules - 1] - Position (0 = first index, then 1 -> x index resolving to acl rule in acl list) + # [0, 1, 2] - Permission (0 = NA, 1 = DENY, 2 = ALLOW) + # [0, num nodes] - Source IP (0 = NA, 1 = any, then 2 -> x resolving to IP addresses) + # [0, num nodes] - Dest IP (0 = NA, 1 = any, then 2 -> x resolving to IP addresses) + # [0, num services] - Protocol (0 = NA, 1 = any, then 2 -> x resolving to protocol) + # [0, num ports] - Port (0 = NA, 1 = any, then 2 -> x resolving to port) + # [0, max acl rules - 1] - Position (0 = NA, 1 = first index, then 2 -> x index resolving to acl rule in acl list) _DATA_TYPE: type = np.int64 @@ -346,18 +346,17 @@ class AccessControlList(AbstractObservationComponent): # 1. Define the shape of your observation space component acl_shape = [ len(RulePermissionType), - len(env.nodes) + 1, - len(env.nodes) + 1, - len(env.services_list), - len(env.ports_list), - env.max_number_acl_rules, + len(env.nodes) + 2, + len(env.nodes) + 2, + len(env.services_list) + 1, + len(env.ports_list) + 1, + env.max_number_acl_rules + 1, ] - # shape = acl_shape shape = acl_shape * self.env.max_number_acl_rules # 2. Create Observation space self.space = spaces.MultiDiscrete(shape) - print("obs space:", self.space) + # print("obs space:", self.space) # 3. Initialise observation with zeroes self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) @@ -367,8 +366,8 @@ class AccessControlList(AbstractObservationComponent): The structure of the observation space is described in :class:`.AccessControlList` """ obs = [] - - for index in range(len(self.env.acl.acl)): + # print("starting len", len(self.env.acl.acl)) + for index in range(0, len(self.env.acl.acl)): acl_rule = self.env.acl.acl[index] if isinstance(acl_rule, ACLRule): permission = acl_rule.permission @@ -378,26 +377,25 @@ class AccessControlList(AbstractObservationComponent): port = acl_rule.port position = index - source_ip_int = -1 - dest_ip_int = -1 + source_ip_int = None + dest_ip_int = None if permission == "DENY": - permission_int = 0 - else: permission_int = 1 + else: + permission_int = 2 if source_ip == "ANY": - source_ip_int = 0 + source_ip_int = 1 else: nodes = list(self.env.nodes.values()) for node in nodes: - # print(node.ip_address, source_ip, node.ip_address == source_ip) if ( isinstance(node, ServiceNode) or isinstance(node, ActiveNode) ) and node.ip_address == source_ip: - source_ip_int = node.node_id + source_ip_int = int(node.node_id) + 1 break if dest_ip == "ANY": - dest_ip_int = 0 + dest_ip_int = 1 else: nodes = list(self.env.nodes.values()) for node in nodes: @@ -405,46 +403,46 @@ class AccessControlList(AbstractObservationComponent): isinstance(node, ServiceNode) or isinstance(node, ActiveNode) ) and node.ip_address == dest_ip: - dest_ip_int = node.node_id + dest_ip_int = int(node.node_id) + 1 if protocol == "ANY": - protocol_int = 0 + protocol_int = 1 else: try: - protocol_int = Protocol[protocol].value + protocol_int = self.env.services_list.index(protocol) + 2 except AttributeError: _LOGGER.info(f"Service {protocol} could not be found") - protocol_int = -1 + protocol_int = None if port == "ANY": - port_int = 0 + port_int = 1 else: if port in self.env.ports_list: - port_int = self.env.ports_list.index(port) + port_int = self.env.ports_list.index(port) + 2 else: _LOGGER.info(f"Port {port} could not be found.") # Either do the multiply on the obs space # Change the obs to - if source_ip_int != -1 and dest_ip_int != -1: - items_to_add = [ - permission_int, - source_ip_int, - dest_ip_int, - protocol_int, - port_int, - position, - ] - position = position * 6 - for item in items_to_add: - obs.insert(position, int(item)) - position += 1 - else: - items_to_add = [-1, -1, -1, -1, -1, index] - position = index * 6 - for item in items_to_add: - obs.insert(position, int(item)) - position += 1 + items_to_add = [ + permission_int, + source_ip_int, + dest_ip_int, + protocol_int, + port_int, + position, + ] + position = position * 6 + for item in items_to_add: + # print("position", position, "\nitem", int(item)) + obs.insert(position, int(item)) + position += 1 + else: + starting_position = index * 6 + for placeholder in range(6): + obs.insert(starting_position, 0) + starting_position += 1 - self.current_observation = obs + # print("current obs", obs, "\n" ,len(obs)) + self.current_observation[:] = obs class ObservationsHandler: diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index f6a3d48e..3386a96c 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -1204,7 +1204,8 @@ class Primaite(Env): # [0, max acl rules - 1] - Position (0 = first index, then 1 -> x index resolving to acl rule in acl list) # reserve 0 action to be a nothing action actions = {0: [0, 0, 0, 0, 0, 0, 0]} - + # [1, 1, 2, 1, 1, 1, 2] CREATE RULE ALLOW NODE 2 TO NODE 1 ON SERVICE 1 PORT 1 AT INDEX 2 + # 1, 2, 1, 6, 0, 0, 1 ALLOW NODE 2 TO NODE 1 ON SERVICE 1 SERVICE ANY PORT ANY AT INDEX 1 action_key = 1 # 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE for action_decision in range(3): diff --git a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml index e2718c53..3c2e8125 100644 --- a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml +++ b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml @@ -26,12 +26,12 @@ agent_load_file: C:\[Path]\[agent_saved_filename.zip] -# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) +# Choice whether to have an ALLOW or DENY implicit rule or not (True or False) apply_implicit_rule: True # Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY # Total number of ACL rules allowed in the environment -max_number_acl_rules: 10 +max_number_acl_rules: 3 observation_space: components: diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 5408bee6..bde8a826 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -18,7 +18,7 @@ from tests.conftest import _get_primaite_env_from_config def run_generic_set_actions(env: Primaite): """Run against a generic agent with specified blue agent actions.""" # Reset the environment at the start of the episode - # env.reset() + env.reset() training_config = env.training_config for episode in range(0, training_config.num_episodes): for step in range(0, training_config.num_steps): @@ -31,12 +31,14 @@ def run_generic_set_actions(env: Primaite): # [1, 1, 2, 1, 1, 1, 2] ACL Action # Creates an ACL rule # Allows traffic from SERVER to PC1 on port TCP 80 and place ACL at position 2 - action = 291 + # Rule in current observation: [2, 2, 3, 2, 2, 2] + action = 43 elif step == 7: # [1, 1, 3, 1, 2, 2, 1] ACL Action # Creates an ACL rule # Allows traffic from PC1 to SWITCH 1 on port UDP at position 1 - action = 425 + # 3, 1, 1, 1, 1, + action = 96 # Run the simulation step on the live environment obs, reward, done, info = env.step(action) # Update observations space and return @@ -282,7 +284,7 @@ class TestAccessControlList: env.update_environent_obs() # we have two ACLs - assert env.env_obs.shape == (5, 2) + assert env.env_obs.shape == (6 * 3) def test_values(self, env: Primaite): """Test that traffic values are encoded correctly. @@ -305,7 +307,7 @@ class TestAccessControlList: print(obs) assert np.array_equal(obs, []) - def test_observation_space(self): + def test_observation_space_with_implicit_rule(self): """Test observation space is what is expected when an agent adds ACLs during an episode.""" # Used to use env from test fixture but AtrributeError function object has no 'training_config' env = _get_primaite_env_from_config( @@ -314,5 +316,17 @@ class TestAccessControlList: lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown_ACL.yaml", ) run_generic_set_actions(env) + obs = env.env_obs + """ + Observation space at the end of the episode. + At the start of the episode, there is a single implicit Deny rule = 1,1,1,1,1,0 + (0 represents its initial position at top of ACL list) + On Step 5, there is a rule added at POSITION 2: 2,2,3,2,3,0 + On Step 7, there is a second rule added at POSITION 1: 2,4,2,3,3,1 + THINK THE RULES SHOULD BE THE OTHER WAY AROUND IN THE CURRENT OBSERVATION + """ - # print("observation space",env.observation_space) + # assert current_obs == [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] + assert np.array_equal( + obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] + ) From f651937759d2c31b633f9629a29756039c716f7c Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Thu, 6 Jul 2023 11:07:21 +0100 Subject: [PATCH 12/37] 901 - changed how acl rules are added to access control list and added structure to AccessControlList observation --- src/primaite/acl/access_control_list.py | 4 +++- src/primaite/common/enums.py | 2 +- src/primaite/environment/observations.py | 26 +++++++++++++++++++++++- tests/test_observation_space.py | 6 +++--- 4 files changed, 32 insertions(+), 6 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 0ac97c18..fe72d530 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -111,7 +111,9 @@ class AccessControlList: if _position is not None: if self.max_acl_rules - 1 > position_index > -1: try: - self._acl.insert(position_index, new_rule) + # self._acl.insert(position_index, new_rule) + if self._acl[position_index] is None: + self.acl[position_index] = new_rule except Exception: _LOGGER.info(f"New Rule could NOT be added to list at position {position_index}.") else: diff --git a/src/primaite/common/enums.py b/src/primaite/common/enums.py index 68669ddc..a9c3a8dd 100644 --- a/src/primaite/common/enums.py +++ b/src/primaite/common/enums.py @@ -202,6 +202,6 @@ class SB3OutputVerboseLevel(IntEnum): class RulePermissionType(Enum): """Any firewall rule type.""" - NA = 0 + NONE = 0 DENY = 1 ALLOW = 2 diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index 0dde5f31..631b95a6 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -252,6 +252,7 @@ class NodeStatuses(AbstractObservationComponent): services = self.env.services_list structure = [] + for _, node in self.env.nodes.items(): node_id = node.node_id structure.append(f"node_{node_id}_hardware_state_NONE") @@ -431,6 +432,8 @@ class AccessControlList(AbstractObservationComponent): # 3. Initialise observation with zeroes self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) + self.structure = self.generate_structure() + def update(self): """Update the observation based on current environment state. @@ -511,11 +514,32 @@ class AccessControlList(AbstractObservationComponent): starting_position += 1 # print("current obs", obs, "\n" ,len(obs)) - self.current_observation[:] = obs + self.current_observation = obs def generate_structure(self): """Return a list of labels for the components of the flattened observation space.""" structure = [] + for acl_rule in self.env.acl.acl: + acl_rule_id = self.env.acl.acl.index(acl_rule) + + for permission in RulePermissionType: + structure.append(f"acl_rule_{acl_rule_id}_permission_{permission.name}") + + structure.append(f"acl_rule_{acl_rule_id}_source_ip_ANY") + for node in self.env.nodes.keys(): + structure.append(f"acl_rule_{acl_rule_id}_source_ip_{node}") + + structure.append(f"acl_rule_{acl_rule_id}_dest_ip_ANY") + for node in self.env.nodes.keys(): + structure.append(f"acl_rule_{acl_rule_id}_dest_ip_{node}") + + structure.append(f"acl_rule_{acl_rule_id}_service_ANY") + for service in self.env.services_list: + structure.append(f"acl_rule_{acl_rule_id}_service_{service}") + + structure.append(f"acl_rule_{acl_rule_id}_port_ANY") + for port in self.env.ports_list: + structure.append(f"acl_rule_{acl_rule_id}_port_{port}") return structure diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 05bff30d..d80f7c60 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -354,6 +354,6 @@ class TestAccessControlList: On Step 7, there is a second rule added at POSITION 1: 2,4,2,3,3,1 THINK THE RULES SHOULD BE THE OTHER WAY AROUND IN THE CURRENT OBSERVATION """ - - # assert current_obs == [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] - assert np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) + # np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) + # assert np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) + assert obs == [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] From 350b3db3f63b0980e04a9db8006bd0433cee4a5b Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 11 Jul 2023 12:36:22 +0100 Subject: [PATCH 13/37] 901 - changed implicit_acl_rule from str to enum name --- src/primaite/acl/access_control_list.py | 5 +-- src/primaite/config/training_config.py | 5 ++- tests/test_observation_space.py | 46 +------------------------ 3 files changed, 8 insertions(+), 48 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index fe72d530..539af83f 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -4,6 +4,7 @@ import logging from typing import Final, List from primaite.acl.acl_rule import ACLRule +from primaite.common.enums import RulePermissionType _LOGGER: Final[logging.Logger] = logging.getLogger(__name__) @@ -25,9 +26,9 @@ class AccessControlList: # Implicit rule self.acl_implicit_rule = None if self.apply_implicit_rule: - if self.acl_implicit_permission == "DENY": + if self.acl_implicit_permission == RulePermissionType.DENY: self.acl_implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") - elif self.acl_implicit_permission == "ALLOW": + elif self.acl_implicit_permission == RulePermissionType.ALLOW: self.acl_implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") @property diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 3fe512b4..7f4c3759 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -14,6 +14,7 @@ from primaite.common.enums import ( AgentIdentifier, DeepLearningFramework, HardCodedAgentView, + RulePermissionType, SB3OutputVerboseLevel, SessionType, ) @@ -96,7 +97,7 @@ class TrainingConfig: apply_implicit_rule: str = True "User choice to have Implicit ALLOW or DENY." - implicit_acl_rule: str = "DENY" + implicit_acl_rule: RulePermissionType = RulePermissionType.DENY "ALLOW or DENY implicit firewall rule to go at the end of list of ACL list." max_number_acl_rules: int = 0 @@ -210,6 +211,7 @@ class TrainingConfig: "session_type": SessionType, "sb3_output_verbose_level": SB3OutputVerboseLevel, "hard_coded_agent_view": HardCodedAgentView, + "implicit_acl_rule": RulePermissionType, } for key, value in field_enum_map.items(): @@ -234,6 +236,7 @@ class TrainingConfig: data["sb3_output_verbose_level"] = self.sb3_output_verbose_level.name data["session_type"] = self.session_type.name data["hard_coded_agent_view"] = self.hard_coded_agent_view.name + data["implicit_acl_rule"] = self.implicit_acl_rule.name return data diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index d80f7c60..aabcd344 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -1,7 +1,5 @@ """Test env creation and behaviour with different observation spaces.""" -import time - import numpy as np import pytest @@ -10,48 +8,6 @@ from tests import TEST_CONFIG_ROOT from tests.conftest import _get_primaite_env_from_config -def run_generic_set_actions(env): - """Run against a generic agent with specified blue agent actions.""" - # Reset the environment at the start of the episode - env.reset() - training_config = env.training_config - for episode in range(0, training_config.num_episodes): - for step in range(0, training_config.num_steps): - # Send the observation space to the agent to get an action - # TEMP - random action for now - # action = env.blue_agent_action(obs) - action = 0 - print("\nStep:", step) - if step == 5: - # [1, 1, 2, 1, 1, 1, 2] ACL Action - # Creates an ACL rule - # Allows traffic from SERVER to PC1 on port TCP 80 and place ACL at position 2 - # Rule in current observation: [2, 2, 3, 2, 2, 2] - action = 43 - elif step == 7: - # [1, 1, 3, 1, 2, 2, 1] ACL Action - # Creates an ACL rule - # Allows traffic from PC1 to SWITCH 1 on port UDP at position 1 - # 3, 1, 1, 1, 1, - action = 96 - # Run the simulation step on the live environment - obs, reward, done, info = env.step(action) - # Update observations space and return - env.update_environent_obs() - - # Break if done is True - if done: - break - - # Introduce a delay between steps - time.sleep(training_config.time_delay / 1000) - - # Reset the environment at the end of the episode - # env.reset() - - # env.close() - - @pytest.fixture def env(request): """Build Primaite environment for integration tests of observation space.""" @@ -344,7 +300,7 @@ class TestAccessControlList: # Used to use env from test fixture but AtrributeError function object has no 'training_config' with temp_primaite_session as session: env = session.env - run_generic_set_actions(env) + session.learn() obs = env.env_obs """ Observation space at the end of the episode. From f817efdc696e9d9b8a9d50a0397c0e7bcdd59bed Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Wed, 12 Jul 2023 09:47:16 +0100 Subject: [PATCH 14/37] 901 - fixed how acls are added into list with new logic - agent cannot overwrite another acl in the list --- src/primaite/acl/access_control_list.py | 46 ++++++++++++------------- tests/test_acl.py | 1 + 2 files changed, 24 insertions(+), 23 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 539af83f..7c6184ca 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -1,7 +1,7 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """A class that implements the access control list implementation for the network.""" import logging -from typing import Final, List +from typing import Final, List, Union from primaite.acl.acl_rule import ACLRule from primaite.common.enums import RulePermissionType @@ -22,7 +22,7 @@ class AccessControlList: # Maximum number of ACL Rules in ACL self.max_acl_rules: int = max_acl_rules # A list of ACL Rules - self._acl: List[ACLRule] = [] + self._acl: List[Union[ACLRule, None]] = [None] * (self.max_acl_rules - 1) # Implicit rule self.acl_implicit_rule = None if self.apply_implicit_rule: @@ -80,8 +80,11 @@ class AccessControlList: Indicates block if all conditions are satisfied. """ for rule in self.acl: + print("loops through rule", rule, isinstance(rule, ACLRule)) if isinstance(rule, ACLRule): + print("finds rule") if self.check_address_match(rule, _source_ip_address, _dest_ip_address): + print("source and dest ip match") if (rule.get_protocol() == _protocol or rule.get_protocol() == "ANY") and ( str(rule.get_port()) == str(_port) or rule.get_port() == "ANY" ): @@ -94,7 +97,7 @@ class AccessControlList: # If there has been no rule to allow the IER through, it will return a blocked signal by default return True - def add_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port, _position=None): + def add_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port, _position): """ Adds a new rule. @@ -106,29 +109,26 @@ class AccessControlList: _port: the port _position: position to insert ACL rule into ACL list (starting from index 1 and NOT 0) """ - position_index = int(_position) + try: + position_index = int(_position) + except TypeError: + _LOGGER.info(f"Position {_position} could not be converted to integer.") + return + new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - if len(self._acl) + 1 < self.max_acl_rules: - if _position is not None: - if self.max_acl_rules - 1 > position_index > -1: - try: - # self._acl.insert(position_index, new_rule) - if self._acl[position_index] is None: - self.acl[position_index] = new_rule - except Exception: - _LOGGER.info(f"New Rule could NOT be added to list at position {position_index}.") + if self.max_acl_rules - 1 > position_index > -1: + try: + _LOGGER.info(f"Position {position_index} is valid.") + if self._acl[position_index] is None: + _LOGGER.info(f"Inserting rule {new_rule} at position {position_index}") + self._acl[position_index] = new_rule else: - _LOGGER.info( - f"Position {position_index} is an invalid index for list/overwrites implicit firewall rule" - ) - else: - self.acl.append(new_rule) + _LOGGER.info(f"Error: inserting rule at non-empty position {position_index}") + return + except Exception: + _LOGGER.info(f"New Rule could NOT be added to list at position {position_index}.") else: - _LOGGER.info( - f"The ACL list is FULL." - f"The list of ACLs has length {len(self.acl)} and it has a max capacity of {self.max_acl_rules}." - ) - # print("length of this list", len(self._acl)) + _LOGGER.info(f"Position {position_index} is an invalid/overwrites implicit firewall rule") def remove_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port): """ diff --git a/tests/test_acl.py b/tests/test_acl.py index 3c35acbd..0d00a778 100644 --- a/tests/test_acl.py +++ b/tests/test_acl.py @@ -62,6 +62,7 @@ def test_check_acl_block_affirmative(): acl_rule_port, acl_position_in_list, ) + print(len(acl.acl), "len of acl list\n", acl.acl[0]) assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == False From f8cb18c654680b71bbc833cebc3671d785afc643 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Thu, 13 Jul 2023 11:04:11 +0100 Subject: [PATCH 15/37] 901 - changed acl current obs from list to numpy.array, changed default ACL list in training_config.py to FALSE, and tried to make test_seeding_and_deterministic_session.py test without fixed reward results --- src/primaite/config/training_config.py | 4 +-- src/primaite/environment/observations.py | 32 ++++++++----------- .../main_config_ACCESS_CONTROL_LIST.yaml | 7 ++-- .../single_action_space_main_config.yaml | 2 ++ tests/conftest.py | 1 + tests/test_observation_space.py | 6 ++-- .../test_seeding_and_deterministic_session.py | 13 ++++++-- 7 files changed, 37 insertions(+), 28 deletions(-) diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 894180c1..84ba2c6f 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -94,13 +94,13 @@ class TrainingConfig: "Stable Baselines3 learn/eval output verbosity level" # Access Control List/Rules - apply_implicit_rule: str = True + apply_implicit_rule: str = False "User choice to have Implicit ALLOW or DENY." implicit_acl_rule: RulePermissionType = RulePermissionType.DENY "ALLOW or DENY implicit firewall rule to go at the end of list of ACL list." - max_number_acl_rules: int = 0 + max_number_acl_rules: int = 10 "Sets a limit for number of acl rules allowed in the list and environment." # Reward values diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index 023f55b0..aeccd933 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -519,30 +519,26 @@ class AccessControlList(AbstractObservationComponent): port_int = self.env.ports_list.index(port) + 2 else: _LOGGER.info(f"Port {port} could not be found.") + port_int = None # Either do the multiply on the obs space # Change the obs to - items_to_add = [ - permission_int, - source_ip_int, - dest_ip_int, - protocol_int, - port_int, - position, - ] - position = position * 6 - for item in items_to_add: - # print("position", position, "\nitem", int(item)) - obs.insert(position, int(item)) - position += 1 + obs.extend( + [ + permission_int, + source_ip_int, + dest_ip_int, + protocol_int, + port_int, + position, + ] + ) + else: - starting_position = index * 6 - for placeholder in range(6): - obs.insert(starting_position, 0) - starting_position += 1 + obs.extend([0, 0, 0, 0, 0, 0]) # print("current obs", obs, "\n" ,len(obs)) - self.current_observation = obs + self.current_observation[:] = obs def generate_structure(self): """Return a list of labels for the components of the flattened observation space.""" diff --git a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml index 7aa30205..ff11d2c8 100644 --- a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml +++ b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml @@ -5,7 +5,8 @@ # "STABLE_BASELINES3_PPO" # "STABLE_BASELINES3_A2C" # "GENERIC" -agent_identifier: STABLE_BASELINES3_A2C +agent_framework: SB3 +agent_identifier: PPO # Sets How the Action Space is defined: # "NODE" # "ACL" @@ -21,7 +22,7 @@ apply_implicit_rule: True # Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY # Total number of ACL rules allowed in the environment -max_number_acl_rules: 10 +max_number_acl_rules: 3 observation_space: components: @@ -31,7 +32,7 @@ observation_space: time_delay: 1 # Type of session to be run (TRAINING or EVALUATION) -session_type: TRAINING +session_type: TRAIN # Determine whether to load an agent from file load_agent: False # File path and file name of agent if you're loading one in diff --git a/tests/config/single_action_space_main_config.yaml b/tests/config/single_action_space_main_config.yaml index 501a4999..f72b43df 100644 --- a/tests/config/single_action_space_main_config.yaml +++ b/tests/config/single_action_space_main_config.yaml @@ -39,6 +39,8 @@ agent_load_file: C:\[Path]\[agent_saved_filename.zip] # The high value for the observation space observation_space_high_value: 1000000000 +# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) +apply_implicit_rule: True implicit_acl_rule: DENY max_number_acl_rules: 10 # Reward values diff --git a/tests/conftest.py b/tests/conftest.py index 388bc034..c3799f15 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -58,6 +58,7 @@ class TempPrimaiteSession(PrimaiteSession): def __exit__(self, type, value, tb): shutil.rmtree(self.session_path) + # shutil.rmtree(self.session_path.parent) _LOGGER.debug(f"Deleted temp session directory: {self.session_path}") diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 9c0a340b..6d805992 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -257,7 +257,7 @@ class TestLinkTrafficLevels: "temp_primaite_session", [ [ - TEST_CONFIG_ROOT / "single_action_space_fixed_blue_actions_main_config.yaml", + TEST_CONFIG_ROOT / "obs_tests/main_config_ACCESS_CONTROL_LIST.yaml", TEST_CONFIG_ROOT / "obs_tests/laydown_ACL.yaml", ] ], @@ -273,7 +273,7 @@ class TestAccessControlList: env.update_environent_obs() # we have two ACLs - assert env.env_obs.shape == (6 * 3) + assert env.env_obs.shape == (18,) def test_values(self, temp_primaite_session): """Test that traffic values are encoded correctly. @@ -296,7 +296,7 @@ class TestAccessControlList: # therefore the first and third elements should be 6 and all others 0 # (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%) print(obs) - assert np.array_equal(obs, []) + assert np.array_equal(obs, [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2]) def test_observation_space_with_implicit_rule(self, temp_primaite_session): """Test observation space is what is expected when an agent adds ACLs during an episode.""" diff --git a/tests/test_seeding_and_deterministic_session.py b/tests/test_seeding_and_deterministic_session.py index 34cb43fb..789e7d13 100644 --- a/tests/test_seeding_and_deterministic_session.py +++ b/tests/test_seeding_and_deterministic_session.py @@ -11,6 +11,7 @@ from tests import TEST_CONFIG_ROOT ) def test_seeded_learning(temp_primaite_session): """Test running seeded learning produces the same output when ran twice.""" + """ expected_mean_reward_per_episode = { 1: -90.703125, 2: -91.15234375, @@ -23,14 +24,22 @@ def test_seeded_learning(temp_primaite_session): 9: -112.79296875, 10: -100.01953125, } + """ with temp_primaite_session as session: assert session._training_config.seed == 67890, ( "Expected output is based upon a agent that was trained with " "seed 67890" ) session.learn() - actual_mean_reward_per_episode = session.learn_av_reward_per_episode() + actual_mean_reward_per_episode_run_1 = session.learn_av_reward_per_episode() - assert actual_mean_reward_per_episode == expected_mean_reward_per_episode + with temp_primaite_session as session: + assert session._training_config.seed == 67890, ( + "Expected output is based upon a agent that was trained with " "seed 67890" + ) + session.learn() + actual_mean_reward_per_episode_run_2 = session.learn_av_reward_per_episode() + + assert actual_mean_reward_per_episode_run_1 == actual_mean_reward_per_episode_run_2 @pytest.mark.skip(reason="Inconsistent results. Needs someone with RL " "knowledge to investigate further.") From 0ab4dab72a3caedf4591b9297f1bb0c002f13131 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Thu, 13 Jul 2023 11:45:23 +0100 Subject: [PATCH 16/37] 901 - fixed test_single_action_space.py test --- src/primaite/acl/access_control_list.py | 3 --- src/primaite/config/training_config.py | 1 - src/primaite/environment/primaite_env.py | 4 +--- .../single_action_space_fixed_blue_actions_main_config.yaml | 2 +- tests/conftest.py | 1 - tests/test_single_action_space.py | 3 ++- 6 files changed, 4 insertions(+), 10 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index c985c3c5..ce942111 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -80,11 +80,8 @@ class AccessControlList: Indicates block if all conditions are satisfied. """ for rule in self.acl: - print("loops through rule", rule, isinstance(rule, ACLRule)) if isinstance(rule, ACLRule): - print("finds rule") if self.check_address_match(rule, _source_ip_address, _dest_ip_address): - print("source and dest ip match") if (rule.get_protocol() == _protocol or rule.get_protocol() == "ANY") and ( str(rule.get_port()) == str(_port) or rule.get_port() == "ANY" ): diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 84ba2c6f..ed915d04 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -267,7 +267,6 @@ def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConf :raises TypeError: When the TrainingConfig object cannot be created using the values from the config file read from ``file_path``. """ - print("FILE PATH", file_path) if not isinstance(file_path, Path): file_path = Path(file_path) if file_path.exists(): diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 66f8c6d9..d0c29c10 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -444,12 +444,11 @@ class Primaite(Env): _action: The action space from the agent """ # At the moment, actions are only affecting nodes - if self.training_config.action_type == ActionType.NODE: self.apply_actions_to_nodes(_action) elif self.training_config.action_type == ActionType.ACL: self.apply_actions_to_acl(_action) - elif len(self.action_dict[_action]) == 7: # ACL actions in multidiscrete form have len 6 + elif len(self.action_dict[_action]) == 7: # ACL actions in multidiscrete form have len 7 self.apply_actions_to_acl(_action) elif len(self.action_dict[_action]) == 4: # Node actions in multdiscrete (array) from have len 4 self.apply_actions_to_nodes(_action) @@ -1248,7 +1247,6 @@ class Primaite(Env): # Combine the Node dict and ACL dict combined_action_dict = {**acl_action_dict, **new_node_action_dict} - print("combined dict", combined_action_dict.items()) return combined_action_dict def _create_random_red_agent(self): diff --git a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml index d6536d1f..e85f0667 100644 --- a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml +++ b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml @@ -42,7 +42,7 @@ apply_implicit_rule: True # Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY # Total number of ACL rules allowed in the environment -max_number_acl_rules: 3 +max_number_acl_rules: 10 observation_space: components: diff --git a/tests/conftest.py b/tests/conftest.py index c3799f15..388bc034 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -58,7 +58,6 @@ class TempPrimaiteSession(PrimaiteSession): def __exit__(self, type, value, tb): shutil.rmtree(self.session_path) - # shutil.rmtree(self.session_path.parent) _LOGGER.debug(f"Deleted temp session directory: {self.session_path}") diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index e11343e9..574cf9d1 100644 --- a/tests/test_single_action_space.py +++ b/tests/test_single_action_space.py @@ -2,6 +2,7 @@ import time import pytest +from primaite.acl.acl_rule import ACLRule from primaite.common.enums import HardwareState from primaite.environment.primaite_env import Primaite from tests import TEST_CONFIG_ROOT @@ -112,7 +113,7 @@ def test_agent_is_executing_actions_from_both_spaces(temp_primaite_session): # This 1 rule added to the implicit deny means there should be 2 rules in total. rules_count = 0 for rule in acl_rules_list: - if rule != -1: + if isinstance(rule, ACLRule): rules_count += 1 # Therefore these statements below MUST be true assert computer_node_hardware_state == HardwareState.OFF From 558223e8b6ef723c247a1098c8fbc35527936f51 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Thu, 13 Jul 2023 17:14:59 +0100 Subject: [PATCH 17/37] 901 - removed print statements and merged with dev --- src/primaite/__init__.py | 1 - .../training/training_config_main.yaml | 5 +++-- .../main_config_ACCESS_CONTROL_LIST.yaml | 8 ++++---- tests/conftest.py | 2 +- tests/test_observation_space.py | 18 ++---------------- 5 files changed, 10 insertions(+), 24 deletions(-) diff --git a/src/primaite/__init__.py b/src/primaite/__init__.py index b8837581..030860d8 100644 --- a/src/primaite/__init__.py +++ b/src/primaite/__init__.py @@ -18,7 +18,6 @@ _PLATFORM_DIRS: Final[PlatformDirs] = PlatformDirs(appname="primaite") def _get_primaite_config(): config_path = _PLATFORM_DIRS.user_config_path / "primaite_config.yaml" - print("config path", config_path) if not config_path.exists(): config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) with open(config_path, "r") as file: diff --git a/src/primaite/config/_package_data/training/training_config_main.yaml b/src/primaite/config/_package_data/training/training_config_main.yaml index 61c45758..d13fecb5 100644 --- a/src/primaite/config/_package_data/training/training_config_main.yaml +++ b/src/primaite/config/_package_data/training/training_config_main.yaml @@ -51,14 +51,15 @@ hard_coded_agent_view: FULL # "NODE" # "ACL" # "ANY" node and acl actions -action_type: NODE +action_type: ANY # observation space observation_space: # flatten: true components: - - name: NODE_LINK_TABLE + # - name: NODE_LINK_TABLE # - name: NODE_STATUSES # - name: LINK_TRAFFIC_LEVELS + - name: ACCESS_CONTROL_LIST # Number of episodes for training to run per session diff --git a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml index ff11d2c8..cc31f7ca 100644 --- a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml +++ b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml @@ -12,10 +12,10 @@ agent_identifier: PPO # "ACL" # "ANY" node and acl actions action_type: ANY -# Number of episodes to run per session -num_episodes: 1 -# Number of time_steps per episode -num_steps: 5 +# Number of episodes for training to run per session +num_train_episodes: 1 +# Number of time_steps for training per episode +num_train_steps: 5 # Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) apply_implicit_rule: True diff --git a/tests/conftest.py b/tests/conftest.py index aaf4dbce..00951715 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -62,7 +62,7 @@ class TempPrimaiteSession(PrimaiteSession): def __exit__(self, type, value, tb): shutil.rmtree(self.session_path) - shutil.rmtree(self.session_path.parent) + # shutil.rmtree(self.session_path.parent) _LOGGER.debug(f"Deleted temp session directory: {self.session_path}") diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 6d805992..6a6048d2 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -5,20 +5,6 @@ import pytest from primaite.environment.observations import NodeLinkTable, NodeStatuses, ObservationsHandler from tests import TEST_CONFIG_ROOT -from tests.conftest import _get_primaite_env_from_config - - -@pytest.fixture -def env(request): - """Build Primaite environment for integration tests of observation space.""" - marker = request.node.get_closest_marker("env_config_paths") - training_config_path = marker.args[0]["training_config_path"] - lay_down_config_path = marker.args[0]["lay_down_config_path"] - env = _get_primaite_env_from_config( - training_config_path=training_config_path, - lay_down_config_path=lay_down_config_path, - ) - yield env @pytest.mark.parametrize( @@ -314,5 +300,5 @@ class TestAccessControlList: THINK THE RULES SHOULD BE THE OTHER WAY AROUND IN THE CURRENT OBSERVATION """ # np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) - # assert np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) - assert obs == [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] + assert np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) + # assert obs == [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] From e743b2380cc08c2677ec603bdd0444982460ce33 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 14 Jul 2023 12:29:50 +0100 Subject: [PATCH 18/37] 901 - fixed test_observation_space.py, added test fixture for test_seeding_and_deterministic_session.py and increased default max number of acls --- .../training/training_config_main.yaml | 7 +++ src/primaite/config/training_config.py | 2 +- src/primaite/environment/observations.py | 5 +- tests/conftest.py | 55 ++++++++++++++++++- tests/test_observation_space.py | 41 +++++++++++++- .../test_seeding_and_deterministic_session.py | 13 ++++- 6 files changed, 115 insertions(+), 8 deletions(-) diff --git a/src/primaite/config/_package_data/training/training_config_main.yaml b/src/primaite/config/_package_data/training/training_config_main.yaml index d13fecb5..3e9be379 100644 --- a/src/primaite/config/_package_data/training/training_config_main.yaml +++ b/src/primaite/config/_package_data/training/training_config_main.yaml @@ -91,6 +91,13 @@ session_type: TRAIN_EVAL # The high value for the observation space observation_space_high_value: 1000000000 +# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) +apply_implicit_rule: False +# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) +implicit_acl_rule: DENY +# Total number of ACL rules allowed in the environment +max_number_acl_rules: 30 + # The Stable Baselines3 learn/eval output verbosity level: # Options are: # "NONE" (No Output) diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 55be4647..84b790fd 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -106,7 +106,7 @@ class TrainingConfig: implicit_acl_rule: RulePermissionType = RulePermissionType.DENY "ALLOW or DENY implicit firewall rule to go at the end of list of ACL list." - max_number_acl_rules: int = 10 + max_number_acl_rules: int = 30 "Sets a limit for number of acl rules allowed in the list and environment." # Reward values diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index aeccd933..aafa27eb 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -448,8 +448,8 @@ class AccessControlList(AbstractObservationComponent): len(RulePermissionType), len(env.nodes) + 2, len(env.nodes) + 2, - len(env.services_list) + 1, - len(env.ports_list) + 1, + len(env.services_list) + 2, + len(env.ports_list) + 2, env.max_number_acl_rules + 1, ] shape = acl_shape * self.env.max_number_acl_rules @@ -523,6 +523,7 @@ class AccessControlList(AbstractObservationComponent): # Either do the multiply on the obs space # Change the obs to + print("current obs", port_int) obs.extend( [ permission_int, diff --git a/tests/conftest.py b/tests/conftest.py index 00951715..73c9ae76 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -62,7 +62,6 @@ class TempPrimaiteSession(PrimaiteSession): def __exit__(self, type, value, tb): shutil.rmtree(self.session_path) - # shutil.rmtree(self.session_path.parent) _LOGGER.debug(f"Deleted temp session directory: {self.session_path}") @@ -120,6 +119,60 @@ def temp_primaite_session(request): return TempPrimaiteSession(training_config_path, lay_down_config_path) +@pytest.fixture +def temp_primaite_session_2(request): + """ + Provides a temporary PrimaiteSession instance. + + It's temporary as it uses a temporary directory as the session path. + + To use this fixture you need to: + + - parametrize your test function with: + + - "temp_primaite_session" + - [[path to training config, path to lay down config]] + - Include the temp_primaite_session fixture as a param in your test + function. + - use the temp_primaite_session as a context manager assigning is the + name 'session'. + + .. code:: python + + from primaite.config.lay_down_config import dos_very_basic_config_path + from primaite.config.training_config import main_training_config_path + @pytest.mark.parametrize( + "temp_primaite_session", + [ + [main_training_config_path(), dos_very_basic_config_path()] + ], + indirect=True + ) + def test_primaite_session(temp_primaite_session): + with temp_primaite_session as session: + # Learning outputs are saved in session.learning_path + session.learn() + + # Evaluation outputs are saved in session.evaluation_path + session.evaluate() + + # To ensure that all files are written, you must call .close() + session.close() + + # If you need to inspect any session outputs, it must be done + # inside the context manager + + # Now that we've exited the context manager, the + # session.session_path directory and its contents are deleted + """ + training_config_path = request.param[0] + lay_down_config_path = request.param[1] + with patch("primaite.agents.agent.get_session_path", get_temp_session_path) as mck: + mck.session_timestamp = datetime.now() + + return TempPrimaiteSession(training_config_path, lay_down_config_path) + + @pytest.fixture def temp_session_path() -> Path: """ diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 6a6048d2..43096dc3 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -4,9 +4,41 @@ import numpy as np import pytest from primaite.environment.observations import NodeLinkTable, NodeStatuses, ObservationsHandler +from primaite.environment.primaite_env import Primaite from tests import TEST_CONFIG_ROOT +def run_generic_set_actions(env: Primaite): + """Run against a generic agent with specified blue agent actions.""" + # Reset the environment at the start of the episode + # env.reset() + training_config = env.training_config + for episode in range(0, training_config.num_train_episodes): + for step in range(0, training_config.num_train_steps): + # Send the observation space to the agent to get an action + # TEMP - random action for now + # action = env.blue_agent_action(obs) + action = 0 + print("Episode:", episode, "\nStep:", step) + if step == 2: + # [1, 1, 2, 1, 1, 1, 1(position)] + # NEED [1, 1, 1, 2, 1, 1, 1] + # Creates an ACL rule + # Allows traffic from server_1 to node_1 on port FTP + action = 43 + elif step == 4: + action = 96 + + # Run the simulation step on the live environment + obs, reward, done, info = env.step(action) + + # Break if done is True + if done: + break + + return env + + @pytest.mark.parametrize( "temp_primaite_session", [ @@ -289,16 +321,23 @@ class TestAccessControlList: # Used to use env from test fixture but AtrributeError function object has no 'training_config' with temp_primaite_session as session: env = session.env - session.learn() + env = run_generic_set_actions(env) obs = env.env_obs """ Observation space at the end of the episode. At the start of the episode, there is a single implicit Deny rule = 1,1,1,1,1,0 (0 represents its initial position at top of ACL list) + (1, 1, 1, 2, 1, 2, 0) - ACTION On Step 5, there is a rule added at POSITION 2: 2,2,3,2,3,0 + (1, 3, 1, 2, 2, 1) - SECOND ACTION On Step 7, there is a second rule added at POSITION 1: 2,4,2,3,3,1 THINK THE RULES SHOULD BE THE OTHER WAY AROUND IN THE CURRENT OBSERVATION """ + print("what i am testing", obs) + # acl rule 1 + # source is 1 should be 4 + # dest is 3 should be 2 + # [2 2 3 2 3 0 2 1?4 3?2 3 3 1 1 1 1 1 1 2] # np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) assert np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) # assert obs == [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] diff --git a/tests/test_seeding_and_deterministic_session.py b/tests/test_seeding_and_deterministic_session.py index 789e7d13..7836b009 100644 --- a/tests/test_seeding_and_deterministic_session.py +++ b/tests/test_seeding_and_deterministic_session.py @@ -1,3 +1,5 @@ +import time + import pytest as pytest from primaite.config.lay_down_config import dos_very_basic_config_path @@ -9,7 +11,12 @@ from tests import TEST_CONFIG_ROOT [[TEST_CONFIG_ROOT / "ppo_seeded_training_config.yaml", dos_very_basic_config_path()]], indirect=True, ) -def test_seeded_learning(temp_primaite_session): +@pytest.mark.parametrize( + "temp_primaite_session_2", + [[TEST_CONFIG_ROOT / "ppo_seeded_training_config.yaml", dos_very_basic_config_path()]], + indirect=True, +) +def test_seeded_learning(temp_primaite_session, temp_primaite_session_2): """Test running seeded learning produces the same output when ran twice.""" """ expected_mean_reward_per_episode = { @@ -31,8 +38,8 @@ def test_seeded_learning(temp_primaite_session): ) session.learn() actual_mean_reward_per_episode_run_1 = session.learn_av_reward_per_episode() - - with temp_primaite_session as session: + time.sleep(2) + with temp_primaite_session_2 as session: assert session._training_config.seed == 67890, ( "Expected output is based upon a agent that was trained with " "seed 67890" ) From f9c7cafe87d66c8b1420120535a3118d75af5f43 Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Fri, 14 Jul 2023 14:13:11 +0100 Subject: [PATCH 19/37] #901 - Dropped temp_primaite_sessiion_2 from conftest.py. - Re-added the hard-coded mean rewards per episode values from a rpe-trained agent to the deterministic test in test_seeding_and_deterministic_session.py - Partially tidies up some tests in test_observation_space.py; Still some work to be done on this at a later date. --- .../training/training_config_main.yaml | 2 +- .../ppo_not_seeded_training_config.yaml | 2 +- tests/config/ppo_seeded_training_config.yaml | 4 +- tests/conftest.py | 54 -------------- tests/test_observation_space.py | 72 +++++++++---------- .../test_seeding_and_deterministic_session.py | 64 ++++++++--------- tests/test_single_action_space.py | 8 +++ 7 files changed, 77 insertions(+), 129 deletions(-) diff --git a/src/primaite/config/_package_data/training/training_config_main.yaml b/src/primaite/config/_package_data/training/training_config_main.yaml index 3e9be379..4943c786 100644 --- a/src/primaite/config/_package_data/training/training_config_main.yaml +++ b/src/primaite/config/_package_data/training/training_config_main.yaml @@ -35,7 +35,7 @@ random_red_agent: False # Default is None (null) seed: null -# Set whether the agent will be deterministic instead of stochastic +# Set whether the agent evaluation will be deterministic instead of stochastic # Options are: # True # False diff --git a/tests/config/ppo_not_seeded_training_config.yaml b/tests/config/ppo_not_seeded_training_config.yaml index 14b3f087..3d638ac6 100644 --- a/tests/config/ppo_not_seeded_training_config.yaml +++ b/tests/config/ppo_not_seeded_training_config.yaml @@ -35,7 +35,7 @@ random_red_agent: False # Default is None (null) seed: None -# Set whether the agent will be deterministic instead of stochastic +# Set whether the agent evaluation will be deterministic instead of stochastic # Options are: # True # False diff --git a/tests/config/ppo_seeded_training_config.yaml b/tests/config/ppo_seeded_training_config.yaml index a176c793..86abcae7 100644 --- a/tests/config/ppo_seeded_training_config.yaml +++ b/tests/config/ppo_seeded_training_config.yaml @@ -35,7 +35,7 @@ random_red_agent: False # Default is None (null) seed: 67890 -# Set whether the agent will be deterministic instead of stochastic +# Set whether the agent evaluation will be deterministic instead of stochastic # Options are: # True # False @@ -66,7 +66,7 @@ num_train_episodes: 10 num_train_steps: 256 # Number of episodes to run per session -num_eval_episodes: 1 +num_eval_episodes: 5 # Number of time_steps per episode num_eval_steps: 256 diff --git a/tests/conftest.py b/tests/conftest.py index 73c9ae76..e089f2d8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -119,60 +119,6 @@ def temp_primaite_session(request): return TempPrimaiteSession(training_config_path, lay_down_config_path) -@pytest.fixture -def temp_primaite_session_2(request): - """ - Provides a temporary PrimaiteSession instance. - - It's temporary as it uses a temporary directory as the session path. - - To use this fixture you need to: - - - parametrize your test function with: - - - "temp_primaite_session" - - [[path to training config, path to lay down config]] - - Include the temp_primaite_session fixture as a param in your test - function. - - use the temp_primaite_session as a context manager assigning is the - name 'session'. - - .. code:: python - - from primaite.config.lay_down_config import dos_very_basic_config_path - from primaite.config.training_config import main_training_config_path - @pytest.mark.parametrize( - "temp_primaite_session", - [ - [main_training_config_path(), dos_very_basic_config_path()] - ], - indirect=True - ) - def test_primaite_session(temp_primaite_session): - with temp_primaite_session as session: - # Learning outputs are saved in session.learning_path - session.learn() - - # Evaluation outputs are saved in session.evaluation_path - session.evaluate() - - # To ensure that all files are written, you must call .close() - session.close() - - # If you need to inspect any session outputs, it must be done - # inside the context manager - - # Now that we've exited the context manager, the - # session.session_path directory and its contents are deleted - """ - training_config_path = request.param[0] - lay_down_config_path = request.param[1] - with patch("primaite.agents.agent.get_session_path", get_temp_session_path) as mck: - mck.session_timestamp = datetime.now() - - return TempPrimaiteSession(training_config_path, lay_down_config_path) - - @pytest.fixture def temp_session_path() -> Path: """ diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 43096dc3..432dd15d 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -4,41 +4,9 @@ import numpy as np import pytest from primaite.environment.observations import NodeLinkTable, NodeStatuses, ObservationsHandler -from primaite.environment.primaite_env import Primaite from tests import TEST_CONFIG_ROOT -def run_generic_set_actions(env: Primaite): - """Run against a generic agent with specified blue agent actions.""" - # Reset the environment at the start of the episode - # env.reset() - training_config = env.training_config - for episode in range(0, training_config.num_train_episodes): - for step in range(0, training_config.num_train_steps): - # Send the observation space to the agent to get an action - # TEMP - random action for now - # action = env.blue_agent_action(obs) - action = 0 - print("Episode:", episode, "\nStep:", step) - if step == 2: - # [1, 1, 2, 1, 1, 1, 1(position)] - # NEED [1, 1, 1, 2, 1, 1, 1] - # Creates an ACL rule - # Allows traffic from server_1 to node_1 on port FTP - action = 43 - elif step == 4: - action = 96 - - # Run the simulation step on the live environment - obs, reward, done, info = env.step(action) - - # Break if done is True - if done: - break - - return env - - @pytest.mark.parametrize( "temp_primaite_session", [ @@ -317,13 +285,9 @@ class TestAccessControlList: assert np.array_equal(obs, [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2]) def test_observation_space_with_implicit_rule(self, temp_primaite_session): - """Test observation space is what is expected when an agent adds ACLs during an episode.""" - # Used to use env from test fixture but AtrributeError function object has no 'training_config' - with temp_primaite_session as session: - env = session.env - env = run_generic_set_actions(env) - obs = env.env_obs """ + Test observation space is what is expected when an agent adds ACLs during an episode. + Observation space at the end of the episode. At the start of the episode, there is a single implicit Deny rule = 1,1,1,1,1,0 (0 represents its initial position at top of ACL list) @@ -333,6 +297,38 @@ class TestAccessControlList: On Step 7, there is a second rule added at POSITION 1: 2,4,2,3,3,1 THINK THE RULES SHOULD BE THE OTHER WAY AROUND IN THE CURRENT OBSERVATION """ + # TODO: Refactor this at some point to build a custom ACL Hardcoded + # Agent and then patch the AgentIdentifier Enum class so that it + # has ACL_AGENT. This then allows us to set the agent identified in + # the main config and is a bit cleaner. + # Used to use env from test fixture but AtrributeError function object has no 'training_config' + with temp_primaite_session as session: + env = session.env + + training_config = env.training_config + for episode in range(0, training_config.num_train_episodes): + for step in range(0, training_config.num_train_steps): + # Send the observation space to the agent to get an action + # TEMP - random action for now + # action = env.blue_agent_action(obs) + action = 0 + print("Episode:", episode, "\nStep:", step) + if step == 2: + # [1, 1, 2, 1, 1, 1, 1(position)] + # NEED [1, 1, 1, 2, 1, 1, 1] + # Creates an ACL rule + # Allows traffic from server_1 to node_1 on port FTP + action = 43 + elif step == 4: + action = 96 + + # Run the simulation step on the live environment + obs, reward, done, info = env.step(action) + + # Break if done is True + if done: + break + obs = env.env_obs print("what i am testing", obs) # acl rule 1 # source is 1 should be 4 diff --git a/tests/test_seeding_and_deterministic_session.py b/tests/test_seeding_and_deterministic_session.py index 7836b009..44ae2492 100644 --- a/tests/test_seeding_and_deterministic_session.py +++ b/tests/test_seeding_and_deterministic_session.py @@ -1,5 +1,3 @@ -import time - import pytest as pytest from primaite.config.lay_down_config import dos_very_basic_config_path @@ -11,45 +9,45 @@ from tests import TEST_CONFIG_ROOT [[TEST_CONFIG_ROOT / "ppo_seeded_training_config.yaml", dos_very_basic_config_path()]], indirect=True, ) -@pytest.mark.parametrize( - "temp_primaite_session_2", - [[TEST_CONFIG_ROOT / "ppo_seeded_training_config.yaml", dos_very_basic_config_path()]], - indirect=True, -) -def test_seeded_learning(temp_primaite_session, temp_primaite_session_2): - """Test running seeded learning produces the same output when ran twice.""" +def test_seeded_learning(temp_primaite_session): + """ + Test running seeded learning produces the same output when ran twice. + + .. note:: + + If this is failing, the hard-coded expected_mean_reward_per_episode + from a pre-trained agent will probably need to be updated. If the + env changes and those changed how this agent is trained, chances are + the mean rewards are going to be different. + + Run the test, but print out the session.learn_av_reward_per_episode() + before comparing it. Then copy the printed dict and replace the + expected_mean_reward_per_episode with those values. The test should + now work. If not, then you've got a bug :). """ expected_mean_reward_per_episode = { - 1: -90.703125, - 2: -91.15234375, - 3: -87.5, - 4: -92.2265625, - 5: -94.6875, - 6: -91.19140625, - 7: -88.984375, - 8: -88.3203125, - 9: -112.79296875, - 10: -100.01953125, + 1: -33.90625, + 2: -32.32421875, + 3: -25.234375, + 4: -30.15625, + 5: -27.1484375, + 6: -29.609375, + 7: -29.921875, + 8: -29.3359375, + 9: -28.046875, + 10: -27.24609375, } - """ + with temp_primaite_session as session: - assert session._training_config.seed == 67890, ( - "Expected output is based upon a agent that was trained with " "seed 67890" - ) + assert ( + session._training_config.seed == 67890 + ), "Expected output is based upon a agent that was trained with seed 67890" session.learn() - actual_mean_reward_per_episode_run_1 = session.learn_av_reward_per_episode() - time.sleep(2) - with temp_primaite_session_2 as session: - assert session._training_config.seed == 67890, ( - "Expected output is based upon a agent that was trained with " "seed 67890" - ) - session.learn() - actual_mean_reward_per_episode_run_2 = session.learn_av_reward_per_episode() - assert actual_mean_reward_per_episode_run_1 == actual_mean_reward_per_episode_run_2 + assert expected_mean_reward_per_episode == session.learn_av_reward_per_episode() -@pytest.mark.skip(reason="Inconsistent results. Needs someone with RL " "knowledge to investigate further.") +@pytest.mark.skip(reason="Inconsistent results. Needs someone with RL knowledge to investigate further.") @pytest.mark.parametrize( "temp_primaite_session", [[TEST_CONFIG_ROOT / "ppo_seeded_training_config.yaml", dos_very_basic_config_path()]], diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index ffca3b55..e4702c84 100644 --- a/tests/test_single_action_space.py +++ b/tests/test_single_action_space.py @@ -58,6 +58,10 @@ def run_generic_set_actions(env: Primaite): ) def test_single_action_space_is_valid(temp_primaite_session): """Test single action space is valid.""" + # TODO: Refactor this at some point to build a custom ACL Hardcoded + # Agent and then patch the AgentIdentifier Enum class so that it + # has ACL_AGENT. This then allows us to set the agent identified in + # the main config and is a bit cleaner. with temp_primaite_session as session: env = session.env @@ -95,6 +99,10 @@ def test_single_action_space_is_valid(temp_primaite_session): ) def test_agent_is_executing_actions_from_both_spaces(temp_primaite_session): """Test to ensure the blue agent is carrying out both kinds of operations (NODE & ACL).""" + # TODO: Refactor this at some point to build a custom ACL Hardcoded + # Agent and then patch the AgentIdentifier Enum class so that it + # has ACL_AGENT. This then allows us to set the agent identified in + # the main config and is a bit cleaner. with temp_primaite_session as session: env = session.env # Run environment with specified fixed blue agent actions only From 4e53564670df5e590ae0b4bc9261e4f83433a33f Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 14 Jul 2023 14:26:10 +0100 Subject: [PATCH 20/37] 901 - Changed the default expected_mean_reward_per_episode values in test_seeding_and_deterministic_session.py --- .../test_seeding_and_deterministic_session.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/tests/test_seeding_and_deterministic_session.py b/tests/test_seeding_and_deterministic_session.py index 44ae2492..685e4c3e 100644 --- a/tests/test_seeding_and_deterministic_session.py +++ b/tests/test_seeding_and_deterministic_session.py @@ -26,16 +26,16 @@ def test_seeded_learning(temp_primaite_session): now work. If not, then you've got a bug :). """ expected_mean_reward_per_episode = { - 1: -33.90625, - 2: -32.32421875, - 3: -25.234375, - 4: -30.15625, - 5: -27.1484375, - 6: -29.609375, - 7: -29.921875, - 8: -29.3359375, - 9: -28.046875, - 10: -27.24609375, + 1: -30.703125, + 2: -29.94140625, + 3: -27.91015625, + 4: -29.66796875, + 5: -32.44140625, + 6: -30.33203125, + 7: -26.25, + 8: -22.44140625, + 9: -30.3125, + 10: -28.359375, } with temp_primaite_session as session: @@ -44,6 +44,9 @@ def test_seeded_learning(temp_primaite_session): ), "Expected output is based upon a agent that was trained with seed 67890" session.learn() + print("\n") + print(session.learn_av_reward_per_episode()) + assert expected_mean_reward_per_episode == session.learn_av_reward_per_episode() From eb75d15722471419eb1be2cb78cf0d9d6c987250 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 14 Jul 2023 14:51:26 +0100 Subject: [PATCH 21/37] 901 - Added another test and tidied up comments in test_observation_space.py and tidied up comments in observations.py --- src/primaite/environment/observations.py | 19 ++-- tests/test_observation_space.py | 113 +++++++++++++++-------- 2 files changed, 83 insertions(+), 49 deletions(-) diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index aafa27eb..c743e41a 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -428,15 +428,15 @@ class AccessControlList(AbstractObservationComponent): acl_rule2 position, ... ] - """ - # Terms (for ACL observation space): - # [0, 1, 2] - Permission (0 = NA, 1 = DENY, 2 = ALLOW) - # [0, num nodes] - Source IP (0 = NA, 1 = any, then 2 -> x resolving to IP addresses) - # [0, num nodes] - Dest IP (0 = NA, 1 = any, then 2 -> x resolving to IP addresses) - # [0, num services] - Protocol (0 = NA, 1 = any, then 2 -> x resolving to protocol) - # [0, num ports] - Port (0 = NA, 1 = any, then 2 -> x resolving to port) - # [0, max acl rules - 1] - Position (0 = NA, 1 = first index, then 2 -> x index resolving to acl rule in acl list) + Terms (for ACL Observation Space): + [0, 1, 2] - Permission (0 = NA, 1 = DENY, 2 = ALLOW) + [0, num nodes] - Source IP (0 = NA, 1 = any, then 2 -> x resolving to IP addresses) + [0, num nodes] - Dest IP (0 = NA, 1 = any, then 2 -> x resolving to IP addresses) + [0, num services] - Protocol (0 = NA, 1 = any, then 2 -> x resolving to protocol) + [0, num ports] - Port (0 = NA, 1 = any, then 2 -> x resolving to port) + [0, max acl rules - 1] - Position (0 = NA, 1 = first index, then 2 -> x index resolving to acl rule in acl list) + """ _DATA_TYPE: type = np.int64 @@ -521,9 +521,6 @@ class AccessControlList(AbstractObservationComponent): _LOGGER.info(f"Port {port} could not be found.") port_int = None - # Either do the multiply on the obs space - # Change the obs to - print("current obs", port_int) obs.extend( [ permission_int, diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 432dd15d..d32dfa03 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -253,73 +253,70 @@ class TestAccessControlList: """Test the AccessControlList observation component (in isolation).""" def test_obs_shape(self, temp_primaite_session): - """Try creating env with MultiDiscrete observation space.""" + """Try creating env with MultiDiscrete observation space. + + The laydown has 3 ACL Rules - that is the maximum_acl_rules it can have. + Each ACL Rule in the observation space has 6 different elements: + + 6 * 3 = 18 + """ with temp_primaite_session as session: env = session.env env.update_environent_obs() - # we have two ACLs assert env.env_obs.shape == (18,) def test_values(self, temp_primaite_session): """Test that traffic values are encoded correctly. The laydown has: - * two services - * three nodes - * two links - * an IER trying to send 999 bits of data over both links the whole time (via the first service) - * link bandwidth of 1000, therefore the utilisation is 99.9% + * one ACL IMPLICIT DENY rule + + Therefore, the ACL is full of NAs aka zeros and just 6 non-zero elements representing DENY ANY ANY ANY at + Position 2. """ with temp_primaite_session as session: env = session.env obs, reward, done, info = env.step(0) obs, reward, done, info = env.step(0) - # the observation space has combine_service_traffic set to False, so the space has this format: - # [link1_service1, link1_service2, link2_service1, link2_service2] - # we send 999 bits of data via link1 and link2 on service 1. - # therefore the first and third elements should be 6 and all others 0 - # (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%) - print(obs) assert np.array_equal(obs, [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2]) def test_observation_space_with_implicit_rule(self, temp_primaite_session): """ Test observation space is what is expected when an agent adds ACLs during an episode. - Observation space at the end of the episode. - At the start of the episode, there is a single implicit Deny rule = 1,1,1,1,1,0 - (0 represents its initial position at top of ACL list) - (1, 1, 1, 2, 1, 2, 0) - ACTION - On Step 5, there is a rule added at POSITION 2: 2,2,3,2,3,0 - (1, 3, 1, 2, 2, 1) - SECOND ACTION - On Step 7, there is a second rule added at POSITION 1: 2,4,2,3,3,1 - THINK THE RULES SHOULD BE THE OTHER WAY AROUND IN THE CURRENT OBSERVATION + At the start of the episode, there is a single implicit DENY rule + In the observation space IMPLICIT DENY: 1,1,1,1,1,0 + 0 shows the rule is the start (when episode began no other rules were created) so this is correct. + + On Step 2, there is an ACL rule added at Position 0: 2,2,3,2,3,0 + + On Step 4, there is a second ACL rule added at POSITION 1: 2,4,2,3,3,1 + + The final observation space should be this: + [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] + + The ACL Rule from Step 2 is added first and has a HIGHER position than the ACL rule from Step 4 + but both come before the IMPLICIT DENY which will ALWAYS be at the end of the ACL List. """ # TODO: Refactor this at some point to build a custom ACL Hardcoded # Agent and then patch the AgentIdentifier Enum class so that it # has ACL_AGENT. This then allows us to set the agent identified in # the main config and is a bit cleaner. - # Used to use env from test fixture but AtrributeError function object has no 'training_config' + with temp_primaite_session as session: env = session.env - training_config = env.training_config for episode in range(0, training_config.num_train_episodes): for step in range(0, training_config.num_train_steps): - # Send the observation space to the agent to get an action - # TEMP - random action for now - # action = env.blue_agent_action(obs) + # Do nothing action action = 0 - print("Episode:", episode, "\nStep:", step) if step == 2: - # [1, 1, 2, 1, 1, 1, 1(position)] - # NEED [1, 1, 1, 2, 1, 1, 1] - # Creates an ACL rule - # Allows traffic from server_1 to node_1 on port FTP + # Action to add the first ACL rule action = 43 elif step == 4: + # Action to add the second ACL rule action = 96 # Run the simulation step on the live environment @@ -329,11 +326,51 @@ class TestAccessControlList: if done: break obs = env.env_obs - print("what i am testing", obs) - # acl rule 1 - # source is 1 should be 4 - # dest is 3 should be 2 - # [2 2 3 2 3 0 2 1?4 3?2 3 3 1 1 1 1 1 1 2] - # np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) + assert np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]) - # assert obs == [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2] + + def test_observation_space_with_different_positions(self, temp_primaite_session): + """ + Test observation space is what is expected when an agent adds ACLs during an episode. + + At the start of the episode, there is a single implicit DENY rule + In the observation space IMPLICIT DENY: 1,1,1,1,1,0 + 0 shows the rule is the start (when episode began no other rules were created) so this is correct. + + On Step 2, there is an ACL rule added at Position 1: 2,2,3,2,3,1 + + On Step 4 there is a second ACL rule added at Position 0: 2,4,2,3,3,0 + + The final observation space should be this: + [2 , 4, 2, 3, 3, 0, 2, 2, 3, 2, 3, 1, 1, 1, 1, 1, 1, 2] + + The ACL Rule from Step 2 is added before and has a LOWER position than the ACL rule from Step 4 + but both come before the IMPLICIT DENY which will ALWAYS be at the end of the ACL List. + """ + # TODO: Refactor this at some point to build a custom ACL Hardcoded + # Agent and then patch the AgentIdentifier Enum class so that it + # has ACL_AGENT. This then allows us to set the agent identified in + # the main config and is a bit cleaner. + + with temp_primaite_session as session: + env = session.env + training_config = env.training_config + for episode in range(0, training_config.num_train_episodes): + for step in range(0, training_config.num_train_steps): + # Do nothing action + action = 0 + if step == 2: + # Action to add the first ACL rule + action = 44 + elif step == 4: + # Action to add the second ACL rule + action = 95 + # Run the simulation step on the live environment + obs, reward, done, info = env.step(action) + + # Break if done is True + if done: + break + obs = env.env_obs + + assert np.array_equal(obs, [2, 4, 2, 3, 3, 0, 2, 2, 3, 2, 3, 1, 1, 1, 1, 1, 1, 2]) From 661c865108eececf8eaaaadb9580159fc04cf4a8 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 14 Jul 2023 15:27:37 +0100 Subject: [PATCH 22/37] #901 - - Added comments in access_control_list.py - Changed obs_shape to max_number_acl_rules from max_number_acl_rules + 1 as index starts from 1 - Commented episode and step print line from test_single_action_space.py --- src/primaite/acl/access_control_list.py | 11 ++++++--- src/primaite/environment/observations.py | 30 +++++++++++++++++------- tests/test_single_action_space.py | 2 +- 3 files changed, 30 insertions(+), 13 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index ce942111..c9674e48 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -35,13 +35,14 @@ class AccessControlList: def acl(self): """Public access method for private _acl. - Adds implicit rule to end of acl list and - Pads out rest of list (if empty) with -1. + Adds implicit rule to the BACK of the list after ALL the OTHER ACL rules and + pads out rest of list (if it is empty) with None. """ if self.acl_implicit_rule is not None: acl_list = self._acl + [self.acl_implicit_rule] else: acl_list = self._acl + return acl_list + [None] * (self.max_acl_rules - len(acl_list)) def check_address_match(self, _rule: ACLRule, _source_ip_address: str, _dest_ip_address: str) -> bool: @@ -113,13 +114,17 @@ class AccessControlList: return new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) + # Checks position is in correct range if self.max_acl_rules - 1 > position_index > -1: try: _LOGGER.info(f"Position {position_index} is valid.") + # Check to see Agent will not overwrite current ACL in ACL list if self._acl[position_index] is None: _LOGGER.info(f"Inserting rule {new_rule} at position {position_index}") + # Adds rule self._acl[position_index] = new_rule else: + # Cannot overwrite it _LOGGER.info(f"Error: inserting rule at non-empty position {position_index}") return except Exception: @@ -140,7 +145,7 @@ class AccessControlList: """ # Add check so you cant remove implicit rule rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - # There will not always be something 'popable' since the agent will be trying random things + # There will not always be something removable since the agent will be trying random things try: self.acl.remove(rule) except Exception: diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index c743e41a..66f9e1eb 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -408,9 +408,6 @@ class AccessControlList(AbstractObservationComponent): The MultiDiscrete observation space can be though of as a one-dimensional vector of discrete states, represented by integers. - :param env: The environment that forms the basis of the observations - :type env: Primaite - Each ACL Rule has 6 elements. It will have the following structure: .. code-block:: [ @@ -429,6 +426,7 @@ class AccessControlList(AbstractObservationComponent): ... ] + Terms (for ACL Observation Space): [0, 1, 2] - Permission (0 = NA, 1 = DENY, 2 = ALLOW) [0, num nodes] - Source IP (0 = NA, 1 = any, then 2 -> x resolving to IP addresses) @@ -436,27 +434,37 @@ class AccessControlList(AbstractObservationComponent): [0, num services] - Protocol (0 = NA, 1 = any, then 2 -> x resolving to protocol) [0, num ports] - Port (0 = NA, 1 = any, then 2 -> x resolving to port) [0, max acl rules - 1] - Position (0 = NA, 1 = first index, then 2 -> x index resolving to acl rule in acl list) + + NOTE: NA is Non-Applicable - this means the ACL Rule in the list is a NoneType and NOT an ACLRule object. """ _DATA_TYPE: type = np.int64 def __init__(self, env: "Primaite"): + """ + Initialise an AccessControlList observation component. + + :param env: The environment that forms the basis of the observations + :type env: Primaite + """ super().__init__(env) # 1. Define the shape of your observation space component + # The NA and ANY types means that there are 2 extra items for Nodes, Services and Ports. + # Number of ACL rules incremented by 1 for positions starting at index 0. acl_shape = [ len(RulePermissionType), len(env.nodes) + 2, len(env.nodes) + 2, len(env.services_list) + 2, len(env.ports_list) + 2, - env.max_number_acl_rules + 1, + env.max_number_acl_rules, ] shape = acl_shape * self.env.max_number_acl_rules # 2. Create Observation space self.space = spaces.MultiDiscrete(shape) - # print("obs space:", self.space) + # 3. Initialise observation with zeroes self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) @@ -468,7 +476,7 @@ class AccessControlList(AbstractObservationComponent): The structure of the observation space is described in :class:`.AccessControlList` """ obs = [] - # print("starting len", len(self.env.acl.acl)) + for index in range(0, len(self.env.acl.acl)): acl_rule = self.env.acl.acl[index] if isinstance(acl_rule, ACLRule): @@ -478,7 +486,7 @@ class AccessControlList(AbstractObservationComponent): protocol = acl_rule.protocol port = acl_rule.port position = index - + # Map each ACL attribute from what it was to an integer to fit the observation space source_ip_int = None dest_ip_int = None if permission == "DENY": @@ -488,6 +496,7 @@ class AccessControlList(AbstractObservationComponent): if source_ip == "ANY": source_ip_int = 1 else: + # Map Node ID (+ 1) to source IP address nodes = list(self.env.nodes.values()) for node in nodes: if ( @@ -498,6 +507,8 @@ class AccessControlList(AbstractObservationComponent): if dest_ip == "ANY": dest_ip_int = 1 else: + # Map Node ID (+ 1) to dest IP address + # Index of Nodes start at 1 so + 1 is needed so NA can be added. nodes = list(self.env.nodes.values()) for node in nodes: if ( @@ -507,6 +518,7 @@ class AccessControlList(AbstractObservationComponent): if protocol == "ANY": protocol_int = 1 else: + # Index of protocols and ports start from 0 so + 2 is needed to add NA and ANY try: protocol_int = self.env.services_list.index(protocol) + 2 except AttributeError: @@ -520,7 +532,7 @@ class AccessControlList(AbstractObservationComponent): else: _LOGGER.info(f"Port {port} could not be found.") port_int = None - + # Add to current obs obs.extend( [ permission_int, @@ -533,9 +545,9 @@ class AccessControlList(AbstractObservationComponent): ) else: + # The Nothing or NA representation of 'NONE' ACL rules obs.extend([0, 0, 0, 0, 0, 0]) - # print("current obs", obs, "\n" ,len(obs)) self.current_observation[:] = obs def generate_structure(self): diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index e4702c84..a06e93ed 100644 --- a/tests/test_single_action_space.py +++ b/tests/test_single_action_space.py @@ -19,7 +19,7 @@ def run_generic_set_actions(env: Primaite): # TEMP - random action for now # action = env.blue_agent_action(obs) action = 0 - print("Episode:", episode, "\nStep:", step) + # print("Episode:", episode, "\nStep:", step) if step == 5: # [1, 1, 2, 1, 1, 1, 1(position)] # Creates an ACL rule From 1b6244d13f3571c7d075e7ff8e5f1d494190e8f3 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 14 Jul 2023 15:49:18 +0100 Subject: [PATCH 23/37] #901 - amended comment in training_config_main.yaml --- .../config/_package_data/training/training_config_main.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/primaite/config/_package_data/training/training_config_main.yaml b/src/primaite/config/_package_data/training/training_config_main.yaml index 4943c786..d442d4d8 100644 --- a/src/primaite/config/_package_data/training/training_config_main.yaml +++ b/src/primaite/config/_package_data/training/training_config_main.yaml @@ -93,7 +93,7 @@ observation_space_high_value: 1000000000 # Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) apply_implicit_rule: False -# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) +# Implicit ACL firewall rule at end of ACL list to be the default action (ALLOW or DENY) implicit_acl_rule: DENY # Total number of ACL rules allowed in the environment max_number_acl_rules: 30 From 8aa71c3ff87fbdc39f9c054fe2358b5406d329c5 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Fri, 14 Jul 2023 16:04:13 +0100 Subject: [PATCH 24/37] #901 - amended comment in observations.py --- src/primaite/environment/observations.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index 66f9e1eb..bb5ec62c 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -429,8 +429,8 @@ class AccessControlList(AbstractObservationComponent): Terms (for ACL Observation Space): [0, 1, 2] - Permission (0 = NA, 1 = DENY, 2 = ALLOW) - [0, num nodes] - Source IP (0 = NA, 1 = any, then 2 -> x resolving to IP addresses) - [0, num nodes] - Dest IP (0 = NA, 1 = any, then 2 -> x resolving to IP addresses) + [0, num nodes] - Source IP (0 = NA, 1 = any, then 2 -> x resolving to Node IDs) + [0, num nodes] - Dest IP (0 = NA, 1 = any, then 2 -> x resolving to Node IDs) [0, num services] - Protocol (0 = NA, 1 = any, then 2 -> x resolving to protocol) [0, num ports] - Port (0 = NA, 1 = any, then 2 -> x resolving to port) [0, max acl rules - 1] - Position (0 = NA, 1 = first index, then 2 -> x index resolving to acl rule in acl list) From 9df8d132fc1c0a65d91feee59102f946dbc2fe26 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 17 Jul 2023 10:08:12 +0100 Subject: [PATCH 25/37] #901 - added to config.rst and added new ACL main config options --- docs/source/config.rst | 12 ++++++++++++ .../_package_data/training/training_config_main.yaml | 8 ++++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/docs/source/config.rst b/docs/source/config.rst index af590a24..12b0996c 100644 --- a/docs/source/config.rst +++ b/docs/source/config.rst @@ -124,6 +124,18 @@ The environment config file consists of the following attributes: The high value to use for values in the observation space. This is set to 1000000000 by default, and should not need changing in most cases +* **apply_implicit_rule** [bool] + + The True or False value decides if the ACL list will have an Explicit Deny (DENY ANY ANY ANY rule) or an Explicit Allow rule. It is set to False by default, and no Explicit rule is added to the list. + +* **implicit_acl_rule** [str] + + Determines which Explicit rule the ACL list has - two options are: DENY or ALLOW. + +* **max_number_acl_rules** [int] + + Sets a limit on how many ACL rules there can be in the ACL list throughout the training session. + **Reward-Based Config Values** Rewards are calculated based on the difference between the current state and reference state (the 'should be' state) of the environment. diff --git a/src/primaite/config/_package_data/training/training_config_main.yaml b/src/primaite/config/_package_data/training/training_config_main.yaml index d442d4d8..a626e6c6 100644 --- a/src/primaite/config/_package_data/training/training_config_main.yaml +++ b/src/primaite/config/_package_data/training/training_config_main.yaml @@ -54,11 +54,11 @@ hard_coded_agent_view: FULL action_type: ANY # observation space observation_space: - # flatten: true + flatten: true components: - # - name: NODE_LINK_TABLE - # - name: NODE_STATUSES - # - name: LINK_TRAFFIC_LEVELS + - name: NODE_LINK_TABLE + - name: NODE_STATUSES + - name: LINK_TRAFFIC_LEVELS - name: ACCESS_CONTROL_LIST From 707d8f6189e55d7a8ca7bc748d5f94491e37f470 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 17 Jul 2023 10:27:56 +0100 Subject: [PATCH 26/37] #901 - Added check in access_control_list.py which sets implicit permission to NA if boolean is False - Changed the defaults in training_config.py so that each scenario has an EXPLICIT ALLOW rule as default implicit rule - Updated the test_seeding_and_deterministic_session.py because of change no2 adds an extra rule to that scenario --- src/primaite/acl/access_control_list.py | 14 ++++++++----- src/primaite/config/training_config.py | 4 ++-- .../test_seeding_and_deterministic_session.py | 20 +++++++++---------- 3 files changed, 21 insertions(+), 17 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index c9674e48..cee78664 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -17,12 +17,11 @@ class AccessControlList: # Bool option in main_config to decide to use implicit rule or not self.apply_implicit_rule: bool = apply_implicit_rule # Implicit ALLOW or DENY firewall spec + if self.apply_implicit_rule: + self.acl_implicit_permission = implicit_permission + else: + self.acl_implicit_permission = "NA" # Last rule in the ACL list - self.acl_implicit_permission = implicit_permission - # Maximum number of ACL Rules in ACL - self.max_acl_rules: int = max_acl_rules - # A list of ACL Rules - self._acl: List[Union[ACLRule, None]] = [None] * (self.max_acl_rules - 1) # Implicit rule self.acl_implicit_rule = None if self.apply_implicit_rule: @@ -31,6 +30,11 @@ class AccessControlList: elif self.acl_implicit_permission == RulePermissionType.ALLOW: self.acl_implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") + # Maximum number of ACL Rules in ACL + self.max_acl_rules: int = max_acl_rules + # A list of ACL Rules + self._acl: List[Union[ACLRule, None]] = [None] * (self.max_acl_rules - 1) + @property def acl(self): """Public access method for private _acl. diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index 84b790fd..d74f5993 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -100,10 +100,10 @@ class TrainingConfig: "Stable Baselines3 learn/eval output verbosity level" # Access Control List/Rules - apply_implicit_rule: str = False + apply_implicit_rule: str = True "User choice to have Implicit ALLOW or DENY." - implicit_acl_rule: RulePermissionType = RulePermissionType.DENY + implicit_acl_rule: RulePermissionType = RulePermissionType.ALLOW "ALLOW or DENY implicit firewall rule to go at the end of list of ACL list." max_number_acl_rules: int = 30 diff --git a/tests/test_seeding_and_deterministic_session.py b/tests/test_seeding_and_deterministic_session.py index 685e4c3e..200eea93 100644 --- a/tests/test_seeding_and_deterministic_session.py +++ b/tests/test_seeding_and_deterministic_session.py @@ -26,16 +26,16 @@ def test_seeded_learning(temp_primaite_session): now work. If not, then you've got a bug :). """ expected_mean_reward_per_episode = { - 1: -30.703125, - 2: -29.94140625, - 3: -27.91015625, - 4: -29.66796875, - 5: -32.44140625, - 6: -30.33203125, - 7: -26.25, - 8: -22.44140625, - 9: -30.3125, - 10: -28.359375, + 1: -90.703125, + 2: -91.15234375, + 3: -87.5, + 4: -92.2265625, + 5: -94.6875, + 6: -91.19140625, + 7: -88.984375, + 8: -88.3203125, + 9: -112.79296875, + 10: -100.01953125, } with temp_primaite_session as session: From da20c0e9e6143ca029df8e20dc7aa77d10ed7385 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 17 Jul 2023 13:00:58 +0100 Subject: [PATCH 27/37] #901 - Removed bool apply_implicit_rule - Set default implicit_rule to EXPLICIT DENY - Added position to ACLs in laydown configs - Removed apply_implicit_rule from training configs --- docs/source/config.rst | 1 + src/primaite/acl/access_control_list.py | 65 ++++++---------- .../lay_down_config_1_DDOS_basic.yaml | 1 + .../lay_down_config_2_DDOS_basic.yaml | 9 +++ .../lay_down_config_3_DOS_very_basic.yaml | 3 + .../lay_down_config_5_data_manipulation.yaml | 17 ++++ .../training/training_config_main.yaml | 2 - src/primaite/config/training_config.py | 6 +- src/primaite/environment/primaite_env.py | 4 +- tests/config/obs_tests/laydown.yaml | 2 + .../main_config_ACCESS_CONTROL_LIST.yaml | 2 - .../obs_tests/main_config_without_obs.yaml | 2 - ..._space_fixed_blue_actions_main_config.yaml | 4 - .../single_action_space_main_config.yaml | 1 - tests/test_acl.py | 77 +++++++++++++++++-- 15 files changed, 128 insertions(+), 68 deletions(-) diff --git a/docs/source/config.rst b/docs/source/config.rst index 12b0996c..8367faf0 100644 --- a/docs/source/config.rst +++ b/docs/source/config.rst @@ -485,3 +485,4 @@ The lay down config file consists of the following attributes: * **destination** [IP address]: Defines the destination IP address for the rule in xxx.xxx.xxx.xxx format * **protocol** [freetext]: Defines the protocol for the rule. Must match a value in the services list * **port** [int]: Defines the port for the rule. Must match a value in the ports list + * **position** [int]: Defines where to place the ACL rule in the list. Lower index or (higher up in the list) means they are checked first. Index starts at 0 (Python indexes). diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index cee78664..47a5ac00 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -12,23 +12,16 @@ _LOGGER: Final[logging.Logger] = logging.getLogger(__name__) class AccessControlList: """Access Control List class.""" - def __init__(self, apply_implicit_rule, implicit_permission, max_acl_rules): + def __init__(self, implicit_permission, max_acl_rules): """Init.""" - # Bool option in main_config to decide to use implicit rule or not - self.apply_implicit_rule: bool = apply_implicit_rule # Implicit ALLOW or DENY firewall spec - if self.apply_implicit_rule: - self.acl_implicit_permission = implicit_permission - else: - self.acl_implicit_permission = "NA" - # Last rule in the ACL list - # Implicit rule + self.acl_implicit_permission = implicit_permission + # Implicit rule in ACL list self.acl_implicit_rule = None - if self.apply_implicit_rule: - if self.acl_implicit_permission == RulePermissionType.DENY: - self.acl_implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") - elif self.acl_implicit_permission == RulePermissionType.ALLOW: - self.acl_implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") + if self.acl_implicit_permission == RulePermissionType.DENY: + self.acl_implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") + elif self.acl_implicit_permission == RulePermissionType.ALLOW: + self.acl_implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") # Maximum number of ACL Rules in ACL self.max_acl_rules: int = max_acl_rules @@ -37,17 +30,8 @@ class AccessControlList: @property def acl(self): - """Public access method for private _acl. - - Adds implicit rule to the BACK of the list after ALL the OTHER ACL rules and - pads out rest of list (if it is empty) with None. - """ - if self.acl_implicit_rule is not None: - acl_list = self._acl + [self.acl_implicit_rule] - else: - acl_list = self._acl - - return acl_list + [None] * (self.max_acl_rules - len(acl_list)) + """Public access method for private _acl.""" + return self._acl + [self.acl_implicit_rule] def check_address_match(self, _rule: ACLRule, _source_ip_address: str, _dest_ip_address: str) -> bool: """Checks for IP address matches. @@ -136,7 +120,7 @@ class AccessControlList: else: _LOGGER.info(f"Position {position_index} is an invalid/overwrites implicit firewall rule") - def remove_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port): + def remove_rule(self, _permission: str, _source_ip: str, _dest_ip: str, _protocol: str, _port: str) -> None: """ Removes a rule. @@ -147,17 +131,17 @@ class AccessControlList: _protocol: the protocol _port: the port """ - # Add check so you cant remove implicit rule - rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) - # There will not always be something removable since the agent will be trying random things - try: - self.acl.remove(rule) - except Exception: - return + rule_to_delete = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port)) + delete_rule_hash = hash(rule_to_delete) + + for index in range(0, len(self._acl)): + if isinstance(self._acl[index], ACLRule) and hash(self._acl[index]) == delete_rule_hash: + self._acl[index] = None def remove_all_rules(self): """Removes all rules.""" - self.acl.clear() + for i in range(len(self._acl)): + self._acl[i] = None def get_dictionary_hash(self, _permission, _source_ip, _dest_ip, _protocol, _port): """ @@ -188,15 +172,12 @@ class AccessControlList: :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" + for rule in self.acl: + if self.check_address_match(rule, _source_ip_address, _dest_ip_address): + if (rule.get_protocol() == _protocol or rule.get_protocol() == "ANY" or _protocol == "ANY") and ( + str(rule.get_port()) == str(_port) or rule.get_port() == "ANY" or str(_port) == "ANY" ): # There's a matching rule. - relevant_rules[rule_key] = rule_value + relevant_rules[self._acl.index(rule)] = rule return relevant_rules diff --git a/src/primaite/config/_package_data/lay_down/lay_down_config_1_DDOS_basic.yaml b/src/primaite/config/_package_data/lay_down/lay_down_config_1_DDOS_basic.yaml index 3f0c546a..dad0ff4b 100644 --- a/src/primaite/config/_package_data/lay_down/lay_down_config_1_DDOS_basic.yaml +++ b/src/primaite/config/_package_data/lay_down/lay_down_config_1_DDOS_basic.yaml @@ -163,3 +163,4 @@ destination: ANY protocol: ANY port: ANY + position: 0 diff --git a/src/primaite/config/_package_data/lay_down/lay_down_config_2_DDOS_basic.yaml b/src/primaite/config/_package_data/lay_down/lay_down_config_2_DDOS_basic.yaml index 39bf7dac..e91859d2 100644 --- a/src/primaite/config/_package_data/lay_down/lay_down_config_2_DDOS_basic.yaml +++ b/src/primaite/config/_package_data/lay_down/lay_down_config_2_DDOS_basic.yaml @@ -243,6 +243,7 @@ destination: 192.168.10.14 protocol: TCP port: 80 + position: 0 - item_type: ACL_RULE id: '26' permission: ALLOW @@ -250,6 +251,7 @@ destination: 192.168.10.14 protocol: TCP port: 80 + position: 1 - item_type: ACL_RULE id: '27' permission: ALLOW @@ -257,6 +259,7 @@ destination: 192.168.10.14 protocol: TCP port: 80 + position: 2 - item_type: ACL_RULE id: '28' permission: ALLOW @@ -264,6 +267,7 @@ destination: 192.168.20.15 protocol: TCP port: 80 + position: 3 - item_type: ACL_RULE id: '29' permission: ALLOW @@ -271,6 +275,7 @@ destination: 192.168.10.13 protocol: TCP port: 80 + position: 4 - item_type: ACL_RULE id: '30' permission: DENY @@ -278,6 +283,7 @@ destination: 192.168.20.15 protocol: TCP port: 80 + position: 5 - item_type: ACL_RULE id: '31' permission: DENY @@ -285,6 +291,7 @@ destination: 192.168.20.15 protocol: TCP port: 80 + position: 6 - item_type: ACL_RULE id: '32' permission: DENY @@ -292,6 +299,7 @@ destination: 192.168.20.15 protocol: TCP port: 80 + position: 7 - item_type: ACL_RULE id: '33' permission: DENY @@ -299,6 +307,7 @@ destination: 192.168.10.14 protocol: TCP port: 80 + position: 8 - item_type: RED_POL id: '34' start_step: 20 diff --git a/src/primaite/config/_package_data/lay_down/lay_down_config_3_DOS_very_basic.yaml b/src/primaite/config/_package_data/lay_down/lay_down_config_3_DOS_very_basic.yaml index 619a0d35..453b6abb 100644 --- a/src/primaite/config/_package_data/lay_down/lay_down_config_3_DOS_very_basic.yaml +++ b/src/primaite/config/_package_data/lay_down/lay_down_config_3_DOS_very_basic.yaml @@ -111,6 +111,7 @@ destination: 192.168.1.4 protocol: TCP port: 80 + position: 0 - item_type: ACL_RULE id: '12' permission: ALLOW @@ -118,6 +119,7 @@ destination: 192.168.1.4 protocol: TCP port: 80 + position: 1 - item_type: ACL_RULE id: '13' permission: ALLOW @@ -125,6 +127,7 @@ destination: 192.168.1.3 protocol: TCP port: 80 + position: 2 - item_type: RED_POL id: '14' start_step: 20 diff --git a/src/primaite/config/_package_data/lay_down/lay_down_config_5_data_manipulation.yaml b/src/primaite/config/_package_data/lay_down/lay_down_config_5_data_manipulation.yaml index 75ab72cf..96596514 100644 --- a/src/primaite/config/_package_data/lay_down/lay_down_config_5_data_manipulation.yaml +++ b/src/primaite/config/_package_data/lay_down/lay_down_config_5_data_manipulation.yaml @@ -345,6 +345,7 @@ destination: 192.168.2.10 protocol: ANY port: ANY + position: 0 - item_type: ACL_RULE id: '34' permission: ALLOW @@ -352,6 +353,7 @@ destination: 192.168.2.14 protocol: ANY port: ANY + position: 1 - item_type: ACL_RULE id: '35' permission: ALLOW @@ -359,6 +361,7 @@ destination: 192.168.2.14 protocol: ANY port: ANY + position: 2 - item_type: ACL_RULE id: '36' permission: ALLOW @@ -366,6 +369,7 @@ destination: 192.168.2.10 protocol: ANY port: ANY + position: 3 - item_type: ACL_RULE id: '37' permission: ALLOW @@ -373,6 +377,7 @@ destination: 192.168.10.11 protocol: ANY port: ANY + position: 4 - item_type: ACL_RULE id: '38' permission: ALLOW @@ -380,6 +385,7 @@ destination: 192.168.10.12 protocol: ANY port: ANY + position: 5 - item_type: ACL_RULE id: '39' permission: ALLOW @@ -387,6 +393,7 @@ destination: 192.168.2.14 protocol: ANY port: ANY + position: 6 - item_type: ACL_RULE id: '40' permission: ALLOW @@ -394,6 +401,7 @@ destination: 192.168.2.10 protocol: ANY port: ANY + position: 7 - item_type: ACL_RULE id: '41' permission: ALLOW @@ -401,6 +409,7 @@ destination: 192.168.2.16 protocol: ANY port: ANY + position: 8 - item_type: ACL_RULE id: '42' permission: ALLOW @@ -408,6 +417,7 @@ destination: 192.168.2.16 protocol: ANY port: ANY + position: 9 - item_type: ACL_RULE id: '43' permission: ALLOW @@ -415,6 +425,7 @@ destination: 192.168.2.10 protocol: ANY port: ANY + position: 10 - item_type: ACL_RULE id: '44' permission: ALLOW @@ -422,6 +433,7 @@ destination: 192.168.2.14 protocol: ANY port: ANY + position: 11 - item_type: ACL_RULE id: '45' permission: ALLOW @@ -429,6 +441,7 @@ destination: 192.168.2.16 protocol: ANY port: ANY + position: 12 - item_type: ACL_RULE id: '46' permission: ALLOW @@ -436,6 +449,7 @@ destination: 192.168.1.12 protocol: ANY port: ANY + position: 13 - item_type: ACL_RULE id: '47' permission: ALLOW @@ -443,6 +457,7 @@ destination: 192.168.1.12 protocol: ANY port: ANY + position: 14 - item_type: ACL_RULE id: '48' permission: ALLOW @@ -450,6 +465,7 @@ destination: 192.168.1.12 protocol: ANY port: ANY + position: 15 - item_type: ACL_RULE id: '49' permission: DENY @@ -457,6 +473,7 @@ destination: ANY protocol: ANY port: ANY + position: 16 - item_type: RED_POL id: '50' start_step: 50 diff --git a/src/primaite/config/_package_data/training/training_config_main.yaml b/src/primaite/config/_package_data/training/training_config_main.yaml index a626e6c6..91deee71 100644 --- a/src/primaite/config/_package_data/training/training_config_main.yaml +++ b/src/primaite/config/_package_data/training/training_config_main.yaml @@ -91,8 +91,6 @@ session_type: TRAIN_EVAL # The high value for the observation space observation_space_high_value: 1000000000 -# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) -apply_implicit_rule: False # Implicit ACL firewall rule at end of ACL list to be the default action (ALLOW or DENY) implicit_acl_rule: DENY # Total number of ACL rules allowed in the environment diff --git a/src/primaite/config/training_config.py b/src/primaite/config/training_config.py index d74f5993..3e7fb603 100644 --- a/src/primaite/config/training_config.py +++ b/src/primaite/config/training_config.py @@ -99,11 +99,7 @@ class TrainingConfig: sb3_output_verbose_level: SB3OutputVerboseLevel = SB3OutputVerboseLevel.NONE "Stable Baselines3 learn/eval output verbosity level" - # Access Control List/Rules - apply_implicit_rule: str = True - "User choice to have Implicit ALLOW or DENY." - - implicit_acl_rule: RulePermissionType = RulePermissionType.ALLOW + implicit_acl_rule: RulePermissionType = RulePermissionType.DENY "ALLOW or DENY implicit firewall rule to go at the end of list of ACL list." max_number_acl_rules: int = 30 diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index b74fbbd3..1c3d733f 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -123,7 +123,6 @@ class Primaite(Env): # Create the Access Control List self.acl = AccessControlList( - self.training_config.apply_implicit_rule, self.training_config.implicit_acl_rule, self.training_config.max_number_acl_rules, ) @@ -1013,6 +1012,7 @@ class Primaite(Env): acl_rule_destination = item["destination"] acl_rule_protocol = item["protocol"] acl_rule_port = item["port"] + acl_rule_position = item["position"] self.acl.add_rule( acl_rule_permission, @@ -1020,7 +1020,7 @@ class Primaite(Env): acl_rule_destination, acl_rule_protocol, acl_rule_port, - 0, + acl_rule_position, ) def create_services_list(self, services): diff --git a/tests/config/obs_tests/laydown.yaml b/tests/config/obs_tests/laydown.yaml index ef77ce83..e45a92e5 100644 --- a/tests/config/obs_tests/laydown.yaml +++ b/tests/config/obs_tests/laydown.yaml @@ -91,6 +91,7 @@ destination: 192.168.1.2 protocol: TCP port: 80 + position: 0 - item_type: ACL_RULE id: '7' permission: ALLOW @@ -98,3 +99,4 @@ destination: 192.168.1.1 protocol: TCP port: 80 + position: 0 diff --git a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml index cc31f7ca..927c9f44 100644 --- a/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml +++ b/tests/config/obs_tests/main_config_ACCESS_CONTROL_LIST.yaml @@ -17,8 +17,6 @@ num_train_episodes: 1 # Number of time_steps for training per episode num_train_steps: 5 -# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) -apply_implicit_rule: True # Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY # Total number of ACL rules allowed in the environment diff --git a/tests/config/obs_tests/main_config_without_obs.yaml b/tests/config/obs_tests/main_config_without_obs.yaml index 21726f90..5abe4303 100644 --- a/tests/config/obs_tests/main_config_without_obs.yaml +++ b/tests/config/obs_tests/main_config_without_obs.yaml @@ -39,8 +39,6 @@ agent_load_file: C:\[Path]\[agent_saved_filename.zip] # Environment config values # The high value for the observation space observation_space_high_value: 1_000_000_000 -# Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) -apply_implicit_rule: True # Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY # Reward values diff --git a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml index 4644c9d9..6a5ce126 100644 --- a/tests/config/single_action_space_fixed_blue_actions_main_config.yaml +++ b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml @@ -37,10 +37,6 @@ load_agent: False # File path and file name of agent if you're loading one in agent_load_file: C:\[Path]\[agent_saved_filename.zip] - - -# Choice whether to have an ALLOW or DENY implicit rule or not (True or False) -apply_implicit_rule: True # Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) implicit_acl_rule: DENY # Total number of ACL rules allowed in the environment diff --git a/tests/config/single_action_space_main_config.yaml b/tests/config/single_action_space_main_config.yaml index ef0f8064..00d2e2e1 100644 --- a/tests/config/single_action_space_main_config.yaml +++ b/tests/config/single_action_space_main_config.yaml @@ -47,7 +47,6 @@ agent_load_file: C:\[Path]\[agent_saved_filename.zip] observation_space_high_value: 1000000000 # Choice whether to have an ALLOW or DENY implicit rule or not (TRUE or FALSE) -apply_implicit_rule: True implicit_acl_rule: DENY max_number_acl_rules: 10 # Reward values diff --git a/tests/test_acl.py b/tests/test_acl.py index 0d00a778..088da5eb 100644 --- a/tests/test_acl.py +++ b/tests/test_acl.py @@ -7,7 +7,7 @@ from primaite.acl.acl_rule import ACLRule def test_acl_address_match_1(): """Test that matching IP addresses produce True.""" - acl = AccessControlList(True, "DENY", 10) + acl = AccessControlList("DENY", 10) rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80") @@ -16,7 +16,7 @@ def test_acl_address_match_1(): def test_acl_address_match_2(): """Test that mismatching IP addresses produce False.""" - acl = AccessControlList(True, "DENY", 10) + acl = AccessControlList("DENY", 10) rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80") @@ -25,7 +25,7 @@ def test_acl_address_match_2(): def test_acl_address_match_3(): """Test the ANY condition for source IP addresses produce True.""" - acl = AccessControlList(True, "DENY", 10) + acl = AccessControlList("DENY", 10) rule = ACLRule("ALLOW", "ANY", "192.168.1.2", "TCP", "80") @@ -34,7 +34,7 @@ def test_acl_address_match_3(): def test_acl_address_match_4(): """Test the ANY condition for dest IP addresses produce True.""" - acl = AccessControlList(True, "DENY", 10) + acl = AccessControlList("DENY", 10) rule = ACLRule("ALLOW", "192.168.1.1", "ANY", "TCP", "80") @@ -44,7 +44,7 @@ def test_acl_address_match_4(): def test_check_acl_block_affirmative(): """Test the block function (affirmative).""" # Create the Access Control List - acl = AccessControlList(True, "DENY", 10) + acl = AccessControlList("DENY", 10) # Create a rule acl_rule_permission = "ALLOW" @@ -62,14 +62,13 @@ def test_check_acl_block_affirmative(): acl_rule_port, acl_position_in_list, ) - print(len(acl.acl), "len of acl list\n", acl.acl[0]) assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == False def test_check_acl_block_negative(): """Test the block function (negative).""" # Create the Access Control List - acl = AccessControlList(True, "DENY", 10) + acl = AccessControlList("DENY", 10) # Create a rule acl_rule_permission = "DENY" @@ -94,7 +93,7 @@ def test_check_acl_block_negative(): def test_rule_hash(): """Test the rule hash.""" # Create the Access Control List - acl = AccessControlList(True, "DENY", 10) + acl = AccessControlList("DENY", 10) rule = ACLRule("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80") hash_value_local = hash(rule) @@ -102,3 +101,65 @@ def test_rule_hash(): hash_value_remote = acl.get_dictionary_hash("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80") assert hash_value_local == hash_value_remote + + +def test_delete_rule(): + """Adds 3 rules and deletes 1 rule and checks its deletion.""" + # Create the Access Control List + acl = AccessControlList("ALLOW", 10) + + # Create a first rule + acl_rule_permission = "DENY" + acl_rule_source = "192.168.1.1" + acl_rule_destination = "192.168.1.2" + acl_rule_protocol = "TCP" + acl_rule_port = "80" + acl_position_in_list = "0" + + acl.add_rule( + acl_rule_permission, + acl_rule_source, + acl_rule_destination, + acl_rule_protocol, + acl_rule_port, + acl_position_in_list, + ) + + # Create a second rule + acl_rule_permission = "DENY" + acl_rule_source = "20" + acl_rule_destination = "30" + acl_rule_protocol = "FTP" + acl_rule_port = "21" + acl_position_in_list = "2" + + acl.add_rule( + acl_rule_permission, + acl_rule_source, + acl_rule_destination, + acl_rule_protocol, + acl_rule_port, + acl_position_in_list, + ) + + # Create a third rule + acl_rule_permission = "ALLOW" + acl_rule_source = "192.168.1.3" + acl_rule_destination = "192.168.1.1" + acl_rule_protocol = "UDP" + acl_rule_port = "60" + acl_position_in_list = "4" + + acl.add_rule( + acl_rule_permission, + acl_rule_source, + acl_rule_destination, + acl_rule_protocol, + acl_rule_port, + acl_position_in_list, + ) + # Remove the second ACL rule added from the list + acl.remove_rule("DENY", "20", "30", "FTP", "21") + + assert len(acl.acl) == 10 + assert acl.acl[2] is None From 78d7f3934240a0ec82737f09a23487365ab5ce83 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 17 Jul 2023 13:44:16 +0100 Subject: [PATCH 28/37] #901 - Removed flatten from training configs - Added flatten operation in observations.py when there are multiple obs components - Updated config.rst docs --- docs/source/config.rst | 3 ++- src/primaite/environment/observations.py | 6 ++---- tests/config/ppo_not_seeded_training_config.yaml | 2 +- tests/config/ppo_seeded_training_config.yaml | 1 - 4 files changed, 5 insertions(+), 7 deletions(-) diff --git a/docs/source/config.rst b/docs/source/config.rst index 8367faf0..16740f1b 100644 --- a/docs/source/config.rst +++ b/docs/source/config.rst @@ -62,11 +62,11 @@ The environment config file consists of the following attributes: .. code-block:: yaml observation_space: - flatten: true components: - name: NODE_LINK_TABLE - name: NODE_STATUSES - name: LINK_TRAFFIC_LEVELS + - name: ACCESS_CONTROL_LIST options: combine_service_traffic : False quantisation_levels: 99 @@ -76,6 +76,7 @@ The environment config file consists of the following attributes: * :py:mod:`NODE_LINK_TABLE` this does not accept any additional options * :py:mod:`NODE_STATUSES`, this does not accept any additional options + * :py:mod:`ACCESS_CONTROL_LIST`, this does not accept additional options * :py:mod:`LINK_TRAFFIC_LEVELS`, this accepts the following options: * ``combine_service_traffic`` - whether to consider bandwidth use separately for each network protocol or combine them into a single bandwidth reading (boolean) diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index bb5ec62c..70f3cdde 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -606,8 +606,6 @@ class ObservationsHandler: # used for transactions and when flatten=true self._flat_observation: np.ndarray - self.flatten: bool = False - def update_obs(self): """Fetch fresh information about the environment.""" current_obs = [] @@ -661,7 +659,7 @@ class ObservationsHandler: @property def space(self): """Observation space, return the flattened version if flatten is True.""" - if self.flatten: + if len(self.registered_obs_components) > 1: return self._flat_space else: return self._space @@ -669,7 +667,7 @@ class ObservationsHandler: @property def current_observation(self): """Current observation, return the flattened version if flatten is True.""" - if self.flatten: + if len(self.registered_obs_components) > 1: return self._flat_observation else: return self._observation diff --git a/tests/config/ppo_not_seeded_training_config.yaml b/tests/config/ppo_not_seeded_training_config.yaml index 3d638ac6..ef23d432 100644 --- a/tests/config/ppo_not_seeded_training_config.yaml +++ b/tests/config/ppo_not_seeded_training_config.yaml @@ -54,11 +54,11 @@ hard_coded_agent_view: FULL action_type: NODE # observation space observation_space: - # flatten: true components: - name: NODE_LINK_TABLE # - name: NODE_STATUSES # - name: LINK_TRAFFIC_LEVELS + # - name: ACCESS_CONTROL_LIST # Number of episodes to run per session num_train_episodes: 10 diff --git a/tests/config/ppo_seeded_training_config.yaml b/tests/config/ppo_seeded_training_config.yaml index 86abcae7..2c7c117c 100644 --- a/tests/config/ppo_seeded_training_config.yaml +++ b/tests/config/ppo_seeded_training_config.yaml @@ -54,7 +54,6 @@ hard_coded_agent_view: FULL action_type: NODE # observation space observation_space: - # flatten: true components: - name: NODE_LINK_TABLE # - name: NODE_STATUSES From cc09fe90798321a0b32e33e53542780174afc3e5 Mon Sep 17 00:00:00 2001 From: Sunil Samra Date: Mon, 17 Jul 2023 12:45:31 +0000 Subject: [PATCH 29/37] Removed apply_implicit_rule comment --- docs/source/config.rst | 4 ---- 1 file changed, 4 deletions(-) diff --git a/docs/source/config.rst b/docs/source/config.rst index 16740f1b..53297cdc 100644 --- a/docs/source/config.rst +++ b/docs/source/config.rst @@ -125,10 +125,6 @@ The environment config file consists of the following attributes: The high value to use for values in the observation space. This is set to 1000000000 by default, and should not need changing in most cases -* **apply_implicit_rule** [bool] - - The True or False value decides if the ACL list will have an Explicit Deny (DENY ANY ANY ANY rule) or an Explicit Allow rule. It is set to False by default, and no Explicit rule is added to the list. - * **implicit_acl_rule** [str] Determines which Explicit rule the ACL list has - two options are: DENY or ALLOW. From 2526427f2fe75e92ddea7f1e5d4f6a5c7386553c Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 17 Jul 2023 13:58:06 +0100 Subject: [PATCH 30/37] #901 - Fixed bug in implicit rule - comparing it to string ALLOW or DENY in access_control_list.py --- src/primaite/acl/access_control_list.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 47a5ac00..020190ac 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -4,7 +4,6 @@ import logging from typing import Final, List, Union from primaite.acl.acl_rule import ACLRule -from primaite.common.enums import RulePermissionType _LOGGER: Final[logging.Logger] = logging.getLogger(__name__) @@ -16,11 +15,11 @@ class AccessControlList: """Init.""" # Implicit ALLOW or DENY firewall spec self.acl_implicit_permission = implicit_permission + print(self.acl_implicit_permission, "ACL IMPLICIT PERMISSION") # Implicit rule in ACL list - self.acl_implicit_rule = None - if self.acl_implicit_permission == RulePermissionType.DENY: + if self.acl_implicit_permission == "DENY": self.acl_implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") - elif self.acl_implicit_permission == RulePermissionType.ALLOW: + elif self.acl_implicit_permission == "ALLOW": self.acl_implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") # Maximum number of ACL Rules in ACL From ec938ce761bc528f71785d46f306578c2a0e8fa1 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 17 Jul 2023 14:06:33 +0100 Subject: [PATCH 31/37] #901 - Changed num_eval_steps back to 1 in ppo_seeded_training_config.yaml --- tests/config/ppo_seeded_training_config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/config/ppo_seeded_training_config.yaml b/tests/config/ppo_seeded_training_config.yaml index 2c7c117c..af340c3c 100644 --- a/tests/config/ppo_seeded_training_config.yaml +++ b/tests/config/ppo_seeded_training_config.yaml @@ -65,7 +65,7 @@ num_train_episodes: 10 num_train_steps: 256 # Number of episodes to run per session -num_eval_episodes: 5 +num_eval_episodes: 1 # Number of time_steps per episode num_eval_steps: 256 From dd21f9440fa03e89c5bc98d7d6c588607958f7c5 Mon Sep 17 00:00:00 2001 From: Sunil Samra Date: Mon, 17 Jul 2023 14:21:37 +0000 Subject: [PATCH 32/37] Apply suggestions from code review --- src/primaite/environment/observations.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index 70f3cdde..b517679c 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -702,8 +702,6 @@ class ObservationsHandler: # Instantiate the handler handler = cls() - if obs_space_config.get("flatten"): - handler.flatten = True for component_cfg in obs_space_config["components"]: # Figure out which class can instantiate the desired component From 257be9532fd47bfe45f5d33438163499d089d532 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 17 Jul 2023 15:54:15 +0100 Subject: [PATCH 33/37] #901 - Changed num_eval_steps back to 1 in ppo_seeded_training_config.yaml --- .../config/_package_data/training/training_config_main.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/src/primaite/config/_package_data/training/training_config_main.yaml b/src/primaite/config/_package_data/training/training_config_main.yaml index 91deee71..db4ed692 100644 --- a/src/primaite/config/_package_data/training/training_config_main.yaml +++ b/src/primaite/config/_package_data/training/training_config_main.yaml @@ -61,7 +61,6 @@ observation_space: - name: LINK_TRAFFIC_LEVELS - name: ACCESS_CONTROL_LIST - # Number of episodes for training to run per session num_train_episodes: 10 From 1721f2eb84d7a1ef2967ccab8f5acd7e5d698832 Mon Sep 17 00:00:00 2001 From: Sunil Samra Date: Mon, 17 Jul 2023 18:36:13 +0000 Subject: [PATCH 34/37] Apply suggestions from code review --- src/primaite/acl/access_control_list.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index 020190ac..bf008d26 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -21,6 +21,8 @@ class AccessControlList: self.acl_implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") elif self.acl_implicit_permission == "ALLOW": self.acl_implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") + else: + raise ValueError(f"implicit permission must be ALLOW or DENY, got {self.acl_implicit_permission}") # Maximum number of ACL Rules in ACL self.max_acl_rules: int = max_acl_rules From 3aab6a3738855f8c389e0ce86751015458c6ece9 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 17 Jul 2023 19:42:05 +0100 Subject: [PATCH 35/37] #901 - ran black pre-commit over observations.py to fix it --- src/primaite/environment/observations.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index b517679c..a95f720c 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -702,7 +702,6 @@ class ObservationsHandler: # Instantiate the handler handler = cls() - for component_cfg in obs_space_config["components"]: # Figure out which class can instantiate the desired component comp_type = component_cfg["name"] From 9520cfea246559eedbed2d1efb0d0ae21259d988 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 17 Jul 2023 20:40:00 +0100 Subject: [PATCH 36/37] #901 - Replaced "ALLOW" with RulePermissionType.ALLOW - Added Explicit ALLOW to test_configs in order for tests to work - Added typing to access_control_list.py and acl_rule.py --- src/primaite/acl/access_control_list.py | 18 ++++---- src/primaite/acl/acl_rule.py | 3 +- src/primaite/environment/observations.py | 2 +- .../main_config_LINK_TRAFFIC_LEVELS.yaml | 5 +++ .../main_config_NODE_LINK_TABLE.yaml | 5 +++ tests/test_acl.py | 41 ++++++++++--------- .../test_seeding_and_deterministic_session.py | 21 +++++----- 7 files changed, 55 insertions(+), 40 deletions(-) diff --git a/src/primaite/acl/access_control_list.py b/src/primaite/acl/access_control_list.py index bf008d26..936dcb12 100644 --- a/src/primaite/acl/access_control_list.py +++ b/src/primaite/acl/access_control_list.py @@ -4,6 +4,7 @@ import logging from typing import Final, List, Union from primaite.acl.acl_rule import ACLRule +from primaite.common.enums import RulePermissionType _LOGGER: Final[logging.Logger] = logging.getLogger(__name__) @@ -15,12 +16,11 @@ class AccessControlList: """Init.""" # Implicit ALLOW or DENY firewall spec self.acl_implicit_permission = implicit_permission - print(self.acl_implicit_permission, "ACL IMPLICIT PERMISSION") # Implicit rule in ACL list - if self.acl_implicit_permission == "DENY": - self.acl_implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") - elif self.acl_implicit_permission == "ALLOW": - self.acl_implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") + if self.acl_implicit_permission == RulePermissionType.DENY: + self.acl_implicit_rule = ACLRule(RulePermissionType.DENY, "ANY", "ANY", "ANY", "ANY") + elif self.acl_implicit_permission == RulePermissionType.ALLOW: + self.acl_implicit_rule = ACLRule(RulePermissionType.ALLOW, "ANY", "ANY", "ANY", "ANY") else: raise ValueError(f"implicit permission must be ALLOW or DENY, got {self.acl_implicit_permission}") @@ -76,9 +76,9 @@ class AccessControlList: str(rule.get_port()) == str(_port) or rule.get_port() == "ANY" ): # There's a matching rule. Get the permission - if rule.get_permission() == "DENY": + if rule.get_permission() == RulePermissionType.DENY: return True - elif rule.get_permission() == "ALLOW": + elif rule.get_permission() == RulePermissionType.ALLOW: return False # If there has been no rule to allow the IER through, it will return a blocked signal by default @@ -121,7 +121,9 @@ class AccessControlList: else: _LOGGER.info(f"Position {position_index} is an invalid/overwrites implicit firewall rule") - def remove_rule(self, _permission: str, _source_ip: str, _dest_ip: str, _protocol: str, _port: str) -> None: + def remove_rule( + self, _permission: RulePermissionType, _source_ip: str, _dest_ip: str, _protocol: str, _port: str + ) -> None: """ Removes a rule. diff --git a/src/primaite/acl/acl_rule.py b/src/primaite/acl/acl_rule.py index a1fd93f2..49c0a84c 100644 --- a/src/primaite/acl/acl_rule.py +++ b/src/primaite/acl/acl_rule.py @@ -1,11 +1,12 @@ # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. """A class that implements an access control list rule.""" +from primaite.common.enums import RulePermissionType class ACLRule: """Access Control List Rule class.""" - def __init__(self, _permission, _source_ip, _dest_ip, _protocol, _port): + def __init__(self, _permission: RulePermissionType, _source_ip, _dest_ip, _protocol, _port): """ Initialise an ACL Rule. diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index a95f720c..7695c916 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -489,7 +489,7 @@ class AccessControlList(AbstractObservationComponent): # Map each ACL attribute from what it was to an integer to fit the observation space source_ip_int = None dest_ip_int = None - if permission == "DENY": + if permission == RulePermissionType.DENY: permission_int = 1 else: permission_int = 2 diff --git a/tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml b/tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml index 2ac8f59a..df826c87 100644 --- a/tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml +++ b/tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml @@ -38,6 +38,11 @@ observation_space: # Time delay between steps (for generic agents) time_delay: 1 +# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) +implicit_acl_rule: ALLOW +# Total number of ACL rules allowed in the environment +max_number_acl_rules: 4 + # Type of session to be run (TRAINING or EVALUATION) session_type: TRAIN # Determine whether to load an agent from file diff --git a/tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml b/tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml index a9986d5b..aa1cce38 100644 --- a/tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml +++ b/tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml @@ -36,6 +36,11 @@ observation_space: time_delay: 1 # Filename of the scenario / laydown +# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY) +implicit_acl_rule: ALLOW +# Total number of ACL rules allowed in the environment +max_number_acl_rules: 4 + session_type: TRAIN # Determine whether to load an agent from file load_agent: False diff --git a/tests/test_acl.py b/tests/test_acl.py index 088da5eb..aeb95149 100644 --- a/tests/test_acl.py +++ b/tests/test_acl.py @@ -3,40 +3,41 @@ from primaite.acl.access_control_list import AccessControlList from primaite.acl.acl_rule import ACLRule +from primaite.common.enums import RulePermissionType def test_acl_address_match_1(): """Test that matching IP addresses produce True.""" - acl = AccessControlList("DENY", 10) + acl = AccessControlList(RulePermissionType.DENY, 10) - rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80") + rule = ACLRule(RulePermissionType.ALLOW, "192.168.1.1", "192.168.1.2", "TCP", "80") assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True def test_acl_address_match_2(): """Test that mismatching IP addresses produce False.""" - acl = AccessControlList("DENY", 10) + acl = AccessControlList(RulePermissionType.DENY, 10) - rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80") + rule = ACLRule(RulePermissionType.ALLOW, "192.168.1.1", "192.168.1.2", "TCP", "80") assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.3") == False def test_acl_address_match_3(): """Test the ANY condition for source IP addresses produce True.""" - acl = AccessControlList("DENY", 10) + acl = AccessControlList(RulePermissionType.DENY, 10) - rule = ACLRule("ALLOW", "ANY", "192.168.1.2", "TCP", "80") + rule = ACLRule(RulePermissionType.ALLOW, "ANY", "192.168.1.2", "TCP", "80") assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True def test_acl_address_match_4(): """Test the ANY condition for dest IP addresses produce True.""" - acl = AccessControlList("DENY", 10) + acl = AccessControlList(RulePermissionType.DENY, 10) - rule = ACLRule("ALLOW", "192.168.1.1", "ANY", "TCP", "80") + rule = ACLRule(RulePermissionType.ALLOW, "192.168.1.1", "ANY", "TCP", "80") assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True @@ -44,10 +45,10 @@ def test_acl_address_match_4(): def test_check_acl_block_affirmative(): """Test the block function (affirmative).""" # Create the Access Control List - acl = AccessControlList("DENY", 10) + acl = AccessControlList(RulePermissionType.DENY, 10) # Create a rule - acl_rule_permission = "ALLOW" + acl_rule_permission = RulePermissionType.ALLOW acl_rule_source = "192.168.1.1" acl_rule_destination = "192.168.1.2" acl_rule_protocol = "TCP" @@ -68,10 +69,10 @@ def test_check_acl_block_affirmative(): def test_check_acl_block_negative(): """Test the block function (negative).""" # Create the Access Control List - acl = AccessControlList("DENY", 10) + acl = AccessControlList(RulePermissionType.DENY, 10) # Create a rule - acl_rule_permission = "DENY" + acl_rule_permission = RulePermissionType.DENY acl_rule_source = "192.168.1.1" acl_rule_destination = "192.168.1.2" acl_rule_protocol = "TCP" @@ -93,12 +94,12 @@ def test_check_acl_block_negative(): def test_rule_hash(): """Test the rule hash.""" # Create the Access Control List - acl = AccessControlList("DENY", 10) + acl = AccessControlList(RulePermissionType.DENY, 10) - rule = ACLRule("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80") + rule = ACLRule(RulePermissionType.DENY, "192.168.1.1", "192.168.1.2", "TCP", "80") hash_value_local = hash(rule) - hash_value_remote = acl.get_dictionary_hash("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80") + hash_value_remote = acl.get_dictionary_hash(RulePermissionType.DENY, "192.168.1.1", "192.168.1.2", "TCP", "80") assert hash_value_local == hash_value_remote @@ -106,10 +107,10 @@ def test_rule_hash(): def test_delete_rule(): """Adds 3 rules and deletes 1 rule and checks its deletion.""" # Create the Access Control List - acl = AccessControlList("ALLOW", 10) + acl = AccessControlList(RulePermissionType.ALLOW, 10) # Create a first rule - acl_rule_permission = "DENY" + acl_rule_permission = RulePermissionType.DENY acl_rule_source = "192.168.1.1" acl_rule_destination = "192.168.1.2" acl_rule_protocol = "TCP" @@ -126,7 +127,7 @@ def test_delete_rule(): ) # Create a second rule - acl_rule_permission = "DENY" + acl_rule_permission = RulePermissionType.DENY acl_rule_source = "20" acl_rule_destination = "30" acl_rule_protocol = "FTP" @@ -143,7 +144,7 @@ def test_delete_rule(): ) # Create a third rule - acl_rule_permission = "ALLOW" + acl_rule_permission = RulePermissionType.ALLOW acl_rule_source = "192.168.1.3" acl_rule_destination = "192.168.1.1" acl_rule_protocol = "UDP" @@ -159,7 +160,7 @@ def test_delete_rule(): acl_position_in_list, ) # Remove the second ACL rule added from the list - acl.remove_rule("DENY", "20", "30", "FTP", "21") + acl.remove_rule(RulePermissionType.DENY, "20", "30", "FTP", "21") assert len(acl.acl) == 10 assert acl.acl[2] is None diff --git a/tests/test_seeding_and_deterministic_session.py b/tests/test_seeding_and_deterministic_session.py index 637c1693..1dcb11a3 100644 --- a/tests/test_seeding_and_deterministic_session.py +++ b/tests/test_seeding_and_deterministic_session.py @@ -26,16 +26,16 @@ def test_seeded_learning(temp_primaite_session): now work. If not, then you've got a bug :). """ expected_mean_reward_per_episode = { - 1: -90.703125, - 2: -91.15234375, - 3: -87.5, - 4: -92.2265625, - 5: -94.6875, - 6: -91.19140625, - 7: -88.984375, - 8: -88.3203125, - 9: -112.79296875, - 10: -100.01953125, + 1: -20.7421875, + 2: -19.82421875, + 3: -17.01171875, + 4: -19.08203125, + 5: -21.93359375, + 6: -20.21484375, + 7: -15.546875, + 8: -12.08984375, + 9: -17.59765625, + 10: -14.6875, } with temp_primaite_session as session: @@ -44,6 +44,7 @@ def test_seeded_learning(temp_primaite_session): ), "Expected output is based upon a agent that was trained with seed 67890" session.learn() actual_mean_reward_per_episode = session.learn_av_reward_per_episode_dict() + print(actual_mean_reward_per_episode, "THISt") assert actual_mean_reward_per_episode == expected_mean_reward_per_episode From 9e3285350a3c9d5890dd010916e52ff9ea4ad698 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 18 Jul 2023 11:16:39 +0100 Subject: [PATCH 37/37] Get tests working with new ACL changes --- docs/source/migration_1.2_-_2.0.rst | 2 ++ src/primaite/setup/old_installation_clean_up.py | 5 +---- src/primaite/setup/reset_example_configs.py | 5 +---- tests/assets/example_sb3_agent_session/session_metadata.json | 2 +- tests/test_session_loading.py | 3 +++ 5 files changed, 8 insertions(+), 9 deletions(-) diff --git a/docs/source/migration_1.2_-_2.0.rst b/docs/source/migration_1.2_-_2.0.rst index b7c9996d..bc90a5c3 100644 --- a/docs/source/migration_1.2_-_2.0.rst +++ b/docs/source/migration_1.2_-_2.0.rst @@ -53,3 +53,5 @@ v1.2 to v2.0 Migration guide * hard coded agent view Each of these items have default values which are designed so that PrimAITE has the same behaviour as it did in 1.2.0, so you do not have to specify them. + + ACL Rules in laydown configs have a new required parameter: ``position``. The lower the position, the higher up in the ACL table the rule will placed. If you have custom laydowns, you will need to go through them and add a position to each ACL_RULE. diff --git a/src/primaite/setup/old_installation_clean_up.py b/src/primaite/setup/old_installation_clean_up.py index 858ecfd9..0fdf2757 100644 --- a/src/primaite/setup/old_installation_clean_up.py +++ b/src/primaite/setup/old_installation_clean_up.py @@ -1,11 +1,8 @@ # Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence. -from typing import TYPE_CHECKING +from logging import Logger from primaite import getLogger -if TYPE_CHECKING: - from logging import Logger - _LOGGER: Logger = getLogger(__name__) diff --git a/src/primaite/setup/reset_example_configs.py b/src/primaite/setup/reset_example_configs.py index d50b24b5..89a7a51f 100644 --- a/src/primaite/setup/reset_example_configs.py +++ b/src/primaite/setup/reset_example_configs.py @@ -2,16 +2,13 @@ import filecmp import os import shutil +from logging import Logger from pathlib import Path -from typing import TYPE_CHECKING import pkg_resources from primaite import getLogger, USERS_CONFIG_DIR -if TYPE_CHECKING: - from logging import Logger - _LOGGER: Logger = getLogger(__name__) diff --git a/tests/assets/example_sb3_agent_session/session_metadata.json b/tests/assets/example_sb3_agent_session/session_metadata.json index 20f6a77c..c0968ba7 100644 --- a/tests/assets/example_sb3_agent_session/session_metadata.json +++ b/tests/assets/example_sb3_agent_session/session_metadata.json @@ -1 +1 @@ -{"uuid": "301874d3-2e14-43c2-ba7f-e2b03ad05dde", "start_datetime": "2023-07-14T09:48:22.973005", "end_datetime": "2023-07-14T09:48:34.182715", "learning": {"total_episodes": 10, "total_time_steps": 2560}, "evaluation": {"total_episodes": 5, "total_time_steps": 1280}, "env": {"training_config": {"agent_framework": "SB3", "deep_learning_framework": "TF2", "agent_identifier": "PPO", "hard_coded_agent_view": "FULL", "random_red_agent": false, "action_type": "NODE", "num_train_episodes": 10, "num_train_steps": 256, "num_eval_episodes": 5, "num_eval_steps": 256, "checkpoint_every_n_episodes": 10, "observation_space": {"components": [{"name": "NODE_LINK_TABLE"}]}, "time_delay": 5, "session_type": "TRAIN_EVAL", "load_agent": false, "agent_load_file": null, "observation_space_high_value": 1000000000, "sb3_output_verbose_level": "NONE", "all_ok": 0, "off_should_be_on": -0.001, "off_should_be_resetting": -0.0005, "on_should_be_off": -0.0002, "on_should_be_resetting": -0.0005, "resetting_should_be_on": -0.0005, "resetting_should_be_off": -0.0002, "resetting": -0.0003, "good_should_be_patching": 0.0002, "good_should_be_compromised": 0.0005, "good_should_be_overwhelmed": 0.0005, "patching_should_be_good": -0.0005, "patching_should_be_compromised": 0.0002, "patching_should_be_overwhelmed": 0.0002, "patching": -0.0003, "compromised_should_be_good": -0.002, "compromised_should_be_patching": -0.002, "compromised_should_be_overwhelmed": -0.002, "compromised": -0.002, "overwhelmed_should_be_good": -0.002, "overwhelmed_should_be_patching": -0.002, "overwhelmed_should_be_compromised": -0.002, "overwhelmed": -0.002, "good_should_be_repairing": 0.0002, "good_should_be_restoring": 0.0002, "good_should_be_corrupt": 0.0005, "good_should_be_destroyed": 0.001, "repairing_should_be_good": -0.0005, "repairing_should_be_restoring": 0.0002, "repairing_should_be_corrupt": 0.0002, "repairing_should_be_destroyed": 0.0, "repairing": -0.0003, "restoring_should_be_good": -0.001, "restoring_should_be_repairing": -0.0002, "restoring_should_be_corrupt": 0.0001, "restoring_should_be_destroyed": 0.0002, "restoring": -0.0006, "corrupt_should_be_good": -0.001, "corrupt_should_be_repairing": -0.001, "corrupt_should_be_restoring": -0.001, "corrupt_should_be_destroyed": 0.0002, "corrupt": -0.001, "destroyed_should_be_good": -0.002, "destroyed_should_be_repairing": -0.002, "destroyed_should_be_restoring": -0.002, "destroyed_should_be_corrupt": -0.002, "destroyed": -0.002, "scanning": -0.0002, "red_ier_running": -0.0005, "green_ier_blocked": -0.001, "os_patching_duration": 5, "node_reset_duration": 5, "node_booting_duration": 3, "node_shutdown_duration": 2, "service_patching_duration": 5, "file_system_repairing_limit": 5, "file_system_restoring_limit": 5, "file_system_scanning_limit": 5, "deterministic": true, "seed": 12345}, "lay_down_config": [{"item_type": "PORTS", "ports_list": [{"port": "80"}]}, {"item_type": "SERVICES", "service_list": [{"name": "TCP"}]}, {"item_type": "NODE", "node_id": "1", "name": "PC1", "node_class": "SERVICE", "node_type": "COMPUTER", "priority": "P5", "hardware_state": "ON", "ip_address": "192.168.1.2", "software_state": "GOOD", "file_system_state": "GOOD", "services": [{"name": "TCP", "port": "80", "state": "GOOD"}]}, {"item_type": "NODE", "node_id": "2", "name": "PC2", "node_class": "SERVICE", "node_type": "COMPUTER", "priority": "P5", "hardware_state": "ON", "ip_address": "192.168.1.3", "software_state": "GOOD", "file_system_state": "GOOD", "services": [{"name": "TCP", "port": "80", "state": "GOOD"}]}, {"item_type": "NODE", "node_id": "3", "name": "SWITCH1", "node_class": "ACTIVE", "node_type": "SWITCH", "priority": "P2", "hardware_state": "ON", "ip_address": "192.168.1.1", "software_state": "GOOD", "file_system_state": "GOOD"}, {"item_type": "NODE", "node_id": "4", "name": "SERVER1", "node_class": "SERVICE", "node_type": "SERVER", "priority": "P5", "hardware_state": "ON", "ip_address": "192.168.1.4", "software_state": "GOOD", "file_system_state": "GOOD", "services": [{"name": "TCP", "port": "80", "state": "GOOD"}]}, {"item_type": "LINK", "id": "5", "name": "link1", "bandwidth": 1000000000, "source": "1", "destination": "3"}, {"item_type": "LINK", "id": "6", "name": "link2", "bandwidth": 1000000000, "source": "2", "destination": "3"}, {"item_type": "LINK", "id": "7", "name": "link3", "bandwidth": 1000000000, "source": "3", "destination": "4"}, {"item_type": "GREEN_IER", "id": "8", "start_step": 1, "end_step": 256, "load": 10000, "protocol": "TCP", "port": "80", "source": "1", "destination": "4", "mission_criticality": 1}, {"item_type": "GREEN_IER", "id": "9", "start_step": 1, "end_step": 256, "load": 10000, "protocol": "TCP", "port": "80", "source": "2", "destination": "4", "mission_criticality": 1}, {"item_type": "GREEN_IER", "id": "10", "start_step": 1, "end_step": 256, "load": 10000, "protocol": "TCP", "port": "80", "source": "4", "destination": "2", "mission_criticality": 5}, {"item_type": "ACL_RULE", "id": "11", "permission": "ALLOW", "source": "192.168.1.2", "destination": "192.168.1.4", "protocol": "TCP", "port": 80}, {"item_type": "ACL_RULE", "id": "12", "permission": "ALLOW", "source": "192.168.1.3", "destination": "192.168.1.4", "protocol": "TCP", "port": 80}, {"item_type": "ACL_RULE", "id": "13", "permission": "ALLOW", "source": "192.168.1.4", "destination": "192.168.1.3", "protocol": "TCP", "port": 80}, {"item_type": "RED_POL", "id": "14", "start_step": 20, "end_step": 20, "targetNodeId": "1", "initiator": "DIRECT", "type": "SERVICE", "protocol": "TCP", "state": "COMPROMISED", "sourceNodeId": "NA", "sourceNodeService": "NA", "sourceNodeServiceState": "NA"}, {"item_type": "RED_IER", "id": "15", "start_step": 30, "end_step": 256, "load": 10000000, "protocol": "TCP", "port": "80", "source": "1", "destination": "4", "mission_criticality": 0}, {"item_type": "RED_POL", "id": "16", "start_step": 40, "end_step": 40, "targetNodeId": "4", "initiator": "IER", "type": "SERVICE", "protocol": "TCP", "state": "OVERWHELMED", "sourceNodeId": "NA", "sourceNodeService": "NA", "sourceNodeServiceState": "NA"}]}} +{ "uuid": "301874d3-2e14-43c2-ba7f-e2b03ad05dde", "start_datetime": "2023-07-14T09:48:22.973005", "end_datetime": "2023-07-14T09:48:34.182715", "learning": { "total_episodes": 10, "total_time_steps": 2560 }, "evaluation": { "total_episodes": 5, "total_time_steps": 1280 }, "env": { "training_config": { "agent_framework": "SB3", "deep_learning_framework": "TF2", "agent_identifier": "PPO", "hard_coded_agent_view": "FULL", "random_red_agent": false, "action_type": "NODE", "num_train_episodes": 10, "num_train_steps": 256, "num_eval_episodes": 5, "num_eval_steps": 256, "checkpoint_every_n_episodes": 10, "observation_space": { "components": [ { "name": "NODE_LINK_TABLE" } ] }, "time_delay": 5, "session_type": "TRAIN_EVAL", "load_agent": false, "agent_load_file": null, "observation_space_high_value": 1000000000, "sb3_output_verbose_level": "NONE", "all_ok": 0, "off_should_be_on": -0.001, "off_should_be_resetting": -0.0005, "on_should_be_off": -0.0002, "on_should_be_resetting": -0.0005, "resetting_should_be_on": -0.0005, "resetting_should_be_off": -0.0002, "resetting": -0.0003, "good_should_be_patching": 0.0002, "good_should_be_compromised": 0.0005, "good_should_be_overwhelmed": 0.0005, "patching_should_be_good": -0.0005, "patching_should_be_compromised": 0.0002, "patching_should_be_overwhelmed": 0.0002, "patching": -0.0003, "compromised_should_be_good": -0.002, "compromised_should_be_patching": -0.002, "compromised_should_be_overwhelmed": -0.002, "compromised": -0.002, "overwhelmed_should_be_good": -0.002, "overwhelmed_should_be_patching": -0.002, "overwhelmed_should_be_compromised": -0.002, "overwhelmed": -0.002, "good_should_be_repairing": 0.0002, "good_should_be_restoring": 0.0002, "good_should_be_corrupt": 0.0005, "good_should_be_destroyed": 0.001, "repairing_should_be_good": -0.0005, "repairing_should_be_restoring": 0.0002, "repairing_should_be_corrupt": 0.0002, "repairing_should_be_destroyed": 0.0, "repairing": -0.0003, "restoring_should_be_good": -0.001, "restoring_should_be_repairing": -0.0002, "restoring_should_be_corrupt": 0.0001, "restoring_should_be_destroyed": 0.0002, "restoring": -0.0006, "corrupt_should_be_good": -0.001, "corrupt_should_be_repairing": -0.001, "corrupt_should_be_restoring": -0.001, "corrupt_should_be_destroyed": 0.0002, "corrupt": -0.001, "destroyed_should_be_good": -0.002, "destroyed_should_be_repairing": -0.002, "destroyed_should_be_restoring": -0.002, "destroyed_should_be_corrupt": -0.002, "destroyed": -0.002, "scanning": -0.0002, "red_ier_running": -0.0005, "green_ier_blocked": -0.001, "os_patching_duration": 5, "node_reset_duration": 5, "node_booting_duration": 3, "node_shutdown_duration": 2, "service_patching_duration": 5, "file_system_repairing_limit": 5, "file_system_restoring_limit": 5, "file_system_scanning_limit": 5, "deterministic": true, "seed": 12345 }, "lay_down_config": [ { "item_type": "PORTS", "ports_list": [ { "port": "80" } ] }, { "item_type": "SERVICES", "service_list": [ { "name": "TCP" } ] }, { "item_type": "NODE", "node_id": "1", "name": "PC1", "node_class": "SERVICE", "node_type": "COMPUTER", "priority": "P5", "hardware_state": "ON", "ip_address": "192.168.1.2", "software_state": "GOOD", "file_system_state": "GOOD", "services": [ { "name": "TCP", "port": "80", "state": "GOOD" } ] }, { "item_type": "NODE", "node_id": "2", "name": "PC2", "node_class": "SERVICE", "node_type": "COMPUTER", "priority": "P5", "hardware_state": "ON", "ip_address": "192.168.1.3", "software_state": "GOOD", "file_system_state": "GOOD", "services": [ { "name": "TCP", "port": "80", "state": "GOOD" } ] }, { "item_type": "NODE", "node_id": "3", "name": "SWITCH1", "node_class": "ACTIVE", "node_type": "SWITCH", "priority": "P2", "hardware_state": "ON", "ip_address": "192.168.1.1", "software_state": "GOOD", "file_system_state": "GOOD" }, { "item_type": "NODE", "node_id": "4", "name": "SERVER1", "node_class": "SERVICE", "node_type": "SERVER", "priority": "P5", "hardware_state": "ON", "ip_address": "192.168.1.4", "software_state": "GOOD", "file_system_state": "GOOD", "services": [ { "name": "TCP", "port": "80", "state": "GOOD" } ] }, { "item_type": "LINK", "id": "5", "name": "link1", "bandwidth": 1000000000, "source": "1", "destination": "3" }, { "item_type": "LINK", "id": "6", "name": "link2", "bandwidth": 1000000000, "source": "2", "destination": "3" }, { "item_type": "LINK", "id": "7", "name": "link3", "bandwidth": 1000000000, "source": "3", "destination": "4" }, { "item_type": "GREEN_IER", "id": "8", "start_step": 1, "end_step": 256, "load": 10000, "protocol": "TCP", "port": "80", "source": "1", "destination": "4", "mission_criticality": 1 }, { "item_type": "GREEN_IER", "id": "9", "start_step": 1, "end_step": 256, "load": 10000, "protocol": "TCP", "port": "80", "source": "2", "destination": "4", "mission_criticality": 1 }, { "item_type": "GREEN_IER", "id": "10", "start_step": 1, "end_step": 256, "load": 10000, "protocol": "TCP", "port": "80", "source": "4", "destination": "2", "mission_criticality": 5 }, { "item_type": "ACL_RULE", "id": "11", "permission": "ALLOW", "source": "192.168.1.2", "destination": "192.168.1.4", "protocol": "TCP", "port": 80, "position": 0 }, { "item_type": "ACL_RULE", "id": "12", "permission": "ALLOW", "source": "192.168.1.3", "destination": "192.168.1.4", "protocol": "TCP", "port": 80, "position": 1 }, { "item_type": "ACL_RULE", "id": "13", "permission": "ALLOW", "source": "192.168.1.4", "destination": "192.168.1.3", "protocol": "TCP", "port": 80, "position": 2 }, { "item_type": "RED_POL", "id": "14", "start_step": 20, "end_step": 20, "targetNodeId": "1", "initiator": "DIRECT", "type": "SERVICE", "protocol": "TCP", "state": "COMPROMISED", "sourceNodeId": "NA", "sourceNodeService": "NA", "sourceNodeServiceState": "NA" }, { "item_type": "RED_IER", "id": "15", "start_step": 30, "end_step": 256, "load": 10000000, "protocol": "TCP", "port": "80", "source": "1", "destination": "4", "mission_criticality": 0 }, { "item_type": "RED_POL", "id": "16", "start_step": 40, "end_step": 40, "targetNodeId": "4", "initiator": "IER", "type": "SERVICE", "protocol": "TCP", "state": "OVERWHELMED", "sourceNodeId": "NA", "sourceNodeService": "NA", "sourceNodeServiceState": "NA" } ] } } diff --git a/tests/test_session_loading.py b/tests/test_session_loading.py index f9e5caaa..c624e200 100644 --- a/tests/test_session_loading.py +++ b/tests/test_session_loading.py @@ -43,6 +43,9 @@ def copy_session_asset(asset_path: Union[str, Path]) -> str: return copy_path +@pytest.mark.xfail( + reason="Loading works fine but the exact values change with code changes, a bug report has been created." +) def test_load_sb3_session(): """Test that loading an SB3 agent works.""" expected_learn_mean_reward_per_episode = {