From 7a02661c66d963d34738a6dd8d36c64db8337435 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Thu, 6 Jul 2023 11:07:21 +0100 Subject: [PATCH] 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]