From 55f13ae6547d5e05e297536a934781aac32e32e0 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 6 Jun 2023 11:00:41 +0100 Subject: [PATCH] 893 - added new tests to test action space size and node is completing both sets of actions in a single episode and created new main config --- src/primaite/environment/primaite_env.py | 1 + ..._space_fixed_blue_actions_main_config.yaml | 89 +++++++++++++++++++ .../single_action_space_lay_down_config.yaml | 50 ++++++----- tests/conftest.py | 40 +++++++++ tests/test_single_action_space.py | 57 +++++++++--- 5 files changed, 199 insertions(+), 38 deletions(-) create mode 100644 tests/config/single_action_space_fixed_blue_actions_main_config.yaml diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 9b0bbeec..be16590f 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -1273,5 +1273,6 @@ class Primaite(Env): # Combine the Node dict and ACL dict combined_action_dict = {**acl_action_dict, **new_node_action_dict} + print("combined_action_dict entry", combined_action_dict.items()) # print(len(list(combined_action_dict.values()))) return combined_action_dict 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 new file mode 100644 index 00000000..becbc0f3 --- /dev/null +++ b/tests/config/single_action_space_fixed_blue_actions_main_config.yaml @@ -0,0 +1,89 @@ +# Main Config File + +# Generic config values +# Choose one of these (dependent on Agent being trained) +# "STABLE_BASELINES3_PPO" +# "STABLE_BASELINES3_A2C" +# "GENERIC" +agentIdentifier: GENERIC_TEST +# Number of episodes to run per session +numEpisodes: 1 +# Time delay between steps (for generic agents) +timeDelay: 1 +# Filename of the scenario / laydown +configFilename: single_action_space_lay_down_config.yaml +# Type of session to be run (TRAINING or EVALUATION) +sessionType: TRAINING +# Determine whether to load an agent from file +loadAgent: False +# File path and file name of agent if you're loading one in +agentLoadFile: C:\[Path]\[agent_saved_filename.zip] + +# Environment config values +# The high value for the observation space +observationSpaceHighValue: 1000000000 + +# Reward values +# Generic +allOk: 0 +# Node Operating State +offShouldBeOn: -10 +offShouldBeResetting: -5 +onShouldBeOff: -2 +onShouldBeResetting: -5 +resettingShouldBeOn: -5 +resettingShouldBeOff: -2 +resetting: -3 +# Node O/S or Service State +goodShouldBePatching: 2 +goodShouldBeCompromised: 5 +goodShouldBeOverwhelmed: 5 +patchingShouldBeGood: -5 +patchingShouldBeCompromised: 2 +patchingShouldBeOverwhelmed: 2 +patching: -3 +compromisedShouldBeGood: -20 +compromisedShouldBePatching: -20 +compromisedShouldBeOverwhelmed: -20 +compromised: -20 +overwhelmedShouldBeGood: -20 +overwhelmedShouldBePatching: -20 +overwhelmedShouldBeCompromised: -20 +overwhelmed: -20 +# Node File System State +goodShouldBeRepairing: 2 +goodShouldBeRestoring: 2 +goodShouldBeCorrupt: 5 +goodShouldBeDestroyed: 10 +repairingShouldBeGood: -5 +repairingShouldBeRestoring: 2 +repairingShouldBeCorrupt: 2 +repairingShouldBeDestroyed: 0 +repairing: -3 +restoringShouldBeGood: -10 +restoringShouldBeRepairing: -2 +restoringShouldBeCorrupt: 1 +restoringShouldBeDestroyed: 2 +restoring: -6 +corruptShouldBeGood: -10 +corruptShouldBeRepairing: -10 +corruptShouldBeRestoring: -10 +corruptShouldBeDestroyed: 2 +corrupt: -10 +destroyedShouldBeGood: -20 +destroyedShouldBeRepairing: -20 +destroyedShouldBeRestoring: -20 +destroyedShouldBeCorrupt: -20 +destroyed: -20 +scanning: -2 +# IER status +redIerRunning: -5 +greenIerBlocked: -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 diff --git a/tests/config/single_action_space_lay_down_config.yaml b/tests/config/single_action_space_lay_down_config.yaml index 6d44356f..ab3b170e 100644 --- a/tests/config/single_action_space_lay_down_config.yaml +++ b/tests/config/single_action_space_lay_down_config.yaml @@ -15,39 +15,41 @@ node_type: COMPUTER priority: P1 hardware_state: 'ON' + ip_address: 192.168.0.14 + software_state: GOOD + file_system_state: GOOD + services: + - name: ftp + port: '21' + state: COMPROMISED +- itemType: NODE + node_id: '2' + name: server_1 + node_class: SERVICE + node_type: SERVER + priority: P1 + hardware_state: 'ON' ip_address: 192.168.0.1 software_state: GOOD file_system_state: GOOD services: - name: ftp port: '21' - state: GOOD + state: COMPROMISED - itemType: POSITION positions: - node: '1' x_pos: 309 y_pos: 78 -- itemType: RED_POL - id: '1' - startStep: 1 - endStep: 3 - targetNodeId: '1' - initiator: DIRECT - type: FILE - protocol: NA - state: CORRUPT - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA -- itemType: RED_POL - id: '2' - startStep: 3 + - node: '2' + x_pos: 200 + y_pos: 78 +- itemType: RED_IER + id: '3' + startStep: 2 endStep: 15 - targetNodeId: '1' - initiator: DIRECT - type: FILE - protocol: NA - state: GOOD - sourceNodeId: NA - sourceNodeService: NA - sourceNodeServiceState: NA + load: 1000 + protocol: ftp + port: CORRUPT + source: '1' + destination: '2' diff --git a/tests/conftest.py b/tests/conftest.py index 1e987223..dd732e78 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -164,10 +164,13 @@ def _get_primaite_env_from_config( # Load in config data load_config_values() env = Primaite(config_values, []) + # Get the number of steps (which is stored in the child config file) config_values.num_steps = env.episode_steps if env.config_values.agent_identifier == "GENERIC": run_generic(env, config_values) + elif env.config_values.agent_identifier == "GENERIC_TEST": + run_generic_set_actions(env, config_values) return env @@ -197,3 +200,40 @@ def run_generic(env, config_values): # env.reset() # env.close() + + +def run_generic_set_actions(env, config_values): + """Run against a generic agent with specified blue agent actions.""" + # Reset the environment at the start of the episode + # env.reset() + for episode in range(0, config_values.num_episodes): + for step in range(0, config_values.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 + if step == 5: + # [1, 1, 2, 1, 1, 1] + # Creates an ACL rule + # Deny traffic from server_1 to node_1 on port FTP + action = 7 + elif step == 7: + # [1, 1, 2, 0] Node Action + # Sets Node 1 Hardware State to OFF + # Does not resolve any service + action = 16 + print(action, "ran") + # Run the simulation step on the live environment + obs, reward, done, info = env.step(action) + + # Break if done is True + if done: + break + + # Introduce a delay between steps + time.sleep(config_values.time_delay / 1000) + + # Reset the environment at the end of the episode + # env.reset() + + # env.close() diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index 203a6232..fda4c96c 100644 --- a/tests/test_single_action_space.py +++ b/tests/test_single_action_space.py @@ -1,26 +1,55 @@ +from primaite.common.enums import HardwareState from tests import TEST_CONFIG_ROOT from tests.conftest import _get_primaite_env_from_config -def test_single_action_space(): +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( main_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", ) - """ - nv.action_space.n is the total number of actions in the Discrete action space - This is the number of actions the agent has to choose from. + # Retrieve the action space dictionary values from environment + env_action_space_dict = env.action_dict.values() + # Flags to check the conditions of the action space + contains_acl_actions = False + contains_node_actions = False + both_action_spaces = False + # Loop through each element of the list (which is every value from the dictionary) + for dict_item in env_action_space_dict: + # Node action detected + if len(dict_item) == 4: + contains_node_actions = True + # Link action detected + elif len(dict_item) == 6: + contains_acl_actions = True + # If both are there then the ANY action type is working + if contains_node_actions and contains_acl_actions: + both_action_spaces = True + # Check condition should be True + assert both_action_spaces - The total number of actions that an agent can type when a NODE action type is selected is: 6 - The total number of actions that an agent can take when an ACL action type is selected is: 7 - These action spaces are combined and the total number of actions is: 12 - This is due to both actions containing the action to "Do nothing", so it needs to be removed from one of the spaces, - to avoid duplicate actions. - - As a result, 12 is the total number of action spaces. - """ - # e - assert env.action_space.n == 12 +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( + main_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", + ) + # Retrieve hardware state of computer_1 node in laydown config + # Agent turned this off in Step 5 + computer_node_hardware_state = env.nodes["1"].hardware_state + # Retrieve the Access Control List object stored by the environment at the end of the episode + 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() + # 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) + # Therefore these statements below MUST be true + assert computer_node_hardware_state == HardwareState.OFF and num_of_rules == 1