From 6547789d5dee8bc90209d60bd31885c02dc108c7 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 11 Jul 2023 12:36:22 +0100 Subject: [PATCH] 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.