From 14096b3dd126796505934868a4ae23baf40368d5 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 2 Jun 2023 12:59:01 +0100 Subject: [PATCH] Add tests for observations --- pytest.ini | 2 + src/primaite/environment/observations.py | 13 +- .../laydown_with_LINK_TRAFFIC_LEVELS.yaml | 43 ++++- .../laydown_with_NODE_LINK_TABLE.yaml | 11 +- .../obs_tests/laydown_with_NODE_STATUSES.yaml | 107 +++++++++++ .../obs_tests/laydown_without_obs_space.yaml | 74 ++++++++ .../obs_tests/main_config_no_agent.yaml | 89 +++++++++ tests/test_observation_space.py | 169 +++++++++++++++--- 8 files changed, 476 insertions(+), 32 deletions(-) create mode 100644 tests/config/obs_tests/laydown_with_NODE_STATUSES.yaml create mode 100644 tests/config/obs_tests/laydown_without_obs_space.yaml create mode 100644 tests/config/obs_tests/main_config_no_agent.yaml diff --git a/pytest.ini b/pytest.ini index e618d7a5..b5fae8d0 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,3 +1,5 @@ [pytest] testpaths = tests +markers = + env_config_paths diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index c4402b69..a467a5db 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -165,14 +165,15 @@ class NodeStatuses(AbstractObservationComponent): super().__init__(env) # 1. Define the shape of your observation space component - shape = [ + node_shape = [ len(HardwareState) + 1, len(SoftwareState) + 1, len(FileSystemState) + 1, ] services_shape = [len(SoftwareState) + 1] * self.env.num_services - shape = shape + services_shape + node_shape = node_shape + services_shape + shape = node_shape * self.env.num_nodes # 2. Create Observation space self.space = spaces.MultiDiscrete(shape) @@ -199,7 +200,9 @@ class NodeStatuses(AbstractObservationComponent): for i, service in enumerate(self.env.services_list): if node.has_service(service): service_states[i] = node.get_service_state(service).value - obs.extend([hardware_state, software_state, file_system_state, *service_states]) + obs.extend( + [hardware_state, software_state, file_system_state, *service_states] + ) self.current_observation[:] = obs @@ -303,8 +306,6 @@ class ObservationsHandler: self.registered_obs_components: List[AbstractObservationComponent] = [] self.space: spaces.Space self.current_observation: Union[Tuple[np.ndarray], np.ndarray] - # i can access the registry items like this: - # self.registry.LINK_TRAFFIC_LEVELS def update_obs(self): """Fetch fresh information about the environment.""" @@ -318,6 +319,7 @@ class ObservationsHandler: self.current_observation = current_obs[0] else: self.current_observation = tuple(current_obs) + # TODO: We may need to add ability to flatten the space as not all agents support tuple spaces. def register(self, obs_component: AbstractObservationComponent): """Add a component for this handler to track. @@ -349,6 +351,7 @@ class ObservationsHandler: self.space = component_spaces[0] else: self.space = spaces.Tuple(component_spaces) + # TODO: We may need to add ability to flatten the space as not all agents support tuple spaces. @classmethod def from_config(cls, env: "Primaite", obs_space_config: dict): diff --git a/tests/config/obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml b/tests/config/obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml index d1909125..516bf5cc 100644 --- a/tests/config/obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml +++ b/tests/config/obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml @@ -2,15 +2,20 @@ type: NODE - itemType: OBSERVATION_SPACE components: - - name: NODE_STATUSES + - name: LINK_TRAFFIC_LEVELS + options: + combine_service_traffic: false + quantisation_levels: 8 - itemType: STEPS steps: 5 - itemType: PORTS portsList: - port: '80' + - port: '53' - itemType: SERVICES serviceList: - name: TCP + - name: UDP ######################################## # Nodes @@ -28,6 +33,9 @@ - name: TCP port: '80' state: GOOD + - name: UDP + port: '53' + state: GOOD - itemType: NODE node_id: '2' name: SERVER @@ -42,6 +50,9 @@ - name: TCP port: '80' state: GOOD + - name: UDP + port: '53' + state: GOOD - itemType: NODE node_id: '3' name: SWITCH1 @@ -67,3 +78,33 @@ bandwidth: 1000 source: '3' destination: '2' + +######################################### +# IERS +- itemType: GREEN_IER + id: '5' + startStep: 0 + endStep: 5 + load: 20 + protocol: TCP + port: '80' + source: '1' + destination: '2' + missionCriticality: 5 + +######################################### +# ACL Rules +- itemType: ACL_RULE + id: '6' + permission: ALLOW + source: 192.168.1.1 + destination: 192.168.1.2 + protocol: TCP + port: 80 +- itemType: ACL_RULE + id: '7' + permission: ALLOW + source: 192.168.1.2 + destination: 192.168.1.1 + protocol: TCP + port: 80 diff --git a/tests/config/obs_tests/laydown_with_NODE_LINK_TABLE.yaml b/tests/config/obs_tests/laydown_with_NODE_LINK_TABLE.yaml index 36fb8199..0ceefbfa 100644 --- a/tests/config/obs_tests/laydown_with_NODE_LINK_TABLE.yaml +++ b/tests/config/obs_tests/laydown_with_NODE_LINK_TABLE.yaml @@ -3,17 +3,16 @@ - itemType: OBSERVATION_SPACE components: - name: NODE_LINK_TABLE - options: - - combine_service_traffic: false - - quantisation_levels: 8 - itemType: STEPS steps: 5 - itemType: PORTS portsList: - port: '80' + - port: '53' - itemType: SERVICES serviceList: - name: TCP + - name: UDP ######################################## # Nodes @@ -31,6 +30,9 @@ - name: TCP port: '80' state: GOOD + - name: UDP + port: '53' + state: GOOD - itemType: NODE node_id: '2' name: SERVER @@ -45,6 +47,9 @@ - name: TCP port: '80' state: GOOD + - name: UDP + port: '53' + state: GOOD - itemType: NODE node_id: '3' name: SWITCH1 diff --git a/tests/config/obs_tests/laydown_with_NODE_STATUSES.yaml b/tests/config/obs_tests/laydown_with_NODE_STATUSES.yaml new file mode 100644 index 00000000..56ff3725 --- /dev/null +++ b/tests/config/obs_tests/laydown_with_NODE_STATUSES.yaml @@ -0,0 +1,107 @@ +- itemType: ACTIONS + type: NODE +- itemType: OBSERVATION_SPACE + components: + - name: NODE_STATUSES +- itemType: STEPS + steps: 5 +- itemType: PORTS + portsList: + - port: '80' + - port: '53' +- itemType: SERVICES + serviceList: + - name: TCP + - name: UDP + +######################################## +# Nodes +- itemType: 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: UDP + port: '53' + state: GOOD +- itemType: 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: UDP + port: '53' + state: OVERWHELMED +- itemType: 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 +- itemType: LINK + id: '4' + name: link1 + bandwidth: 1000 + source: '1' + destination: '3' +- itemType: LINK + id: '5' + name: link2 + bandwidth: 1000 + source: '3' + destination: '2' + +######################################### +# IERS +- itemType: GREEN_IER + id: '5' + startStep: 0 + endStep: 5 + load: 20 + protocol: TCP + port: '80' + source: '1' + destination: '2' + missionCriticality: 5 + +######################################### +# ACL Rules +- itemType: ACL_RULE + id: '6' + permission: ALLOW + source: 192.168.1.1 + destination: 192.168.1.2 + protocol: TCP + port: 80 +- itemType: ACL_RULE + id: '7' + permission: ALLOW + source: 192.168.1.2 + destination: 192.168.1.1 + protocol: TCP + port: 80 diff --git a/tests/config/obs_tests/laydown_without_obs_space.yaml b/tests/config/obs_tests/laydown_without_obs_space.yaml new file mode 100644 index 00000000..3ef214da --- /dev/null +++ b/tests/config/obs_tests/laydown_without_obs_space.yaml @@ -0,0 +1,74 @@ +- itemType: ACTIONS + type: NODE +- itemType: STEPS + steps: 5 +- itemType: PORTS + portsList: + - port: '80' + - port: '53' +- itemType: SERVICES + serviceList: + - name: TCP + - name: UDP + +######################################## +# Nodes +- itemType: 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: GOOD + file_system_state: GOOD + services: + - name: TCP + port: '80' + state: GOOD + - name: UDP + port: '53' + state: GOOD +- itemType: 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: UDP + port: '53' + state: GOOD +- itemType: 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 +- itemType: LINK + id: '4' + name: link1 + bandwidth: 1000 + source: '1' + destination: '3' +- itemType: LINK + id: '5' + name: link2 + bandwidth: 1000 + source: '3' + destination: '2' diff --git a/tests/config/obs_tests/main_config_no_agent.yaml b/tests/config/obs_tests/main_config_no_agent.yaml new file mode 100644 index 00000000..f632dca9 --- /dev/null +++ b/tests/config/obs_tests/main_config_no_agent.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: NONE +# Number of episodes to run per session +numEpisodes: 1 +# Time delay between steps (for generic agents) +timeDelay: 1 +# Filename of the scenario / laydown +configFilename: one_node_states_on_off_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 Hardware State +offShouldBeOn: -10 +offShouldBeResetting: -5 +onShouldBeOff: -2 +onShouldBeResetting: -5 +resettingShouldBeOn: -5 +resettingShouldBeOff: -2 +resetting: -3 +# Node Software 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/test_observation_space.py b/tests/test_observation_space.py index a13121b9..314728ae 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -1,45 +1,168 @@ """Test env creation and behaviour with different observation spaces.""" +import numpy as np +import pytest -from primaite.environment.observations import NodeStatuses, ObservationsHandler +from primaite.environment.observations import ( + NodeLinkTable, + NodeStatuses, + ObservationsHandler, +) +from primaite.environment.primaite_env import Primaite from tests import TEST_CONFIG_ROOT from tests.conftest import _get_primaite_env_from_config -def test_creating_env_with_box_obs(): - """Try creating env with box observation space.""" +@pytest.fixture +def env(request): + """Build Primaite environment for integration tests of observation space.""" + marker = request.node.get_closest_marker("env_config_paths") + main_config_path = marker.args[0]["main_config_path"] + lay_down_config_path = marker.args[0]["lay_down_config_path"] env = _get_primaite_env_from_config( - main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml", - lay_down_config_path=TEST_CONFIG_ROOT / "box_obs_space_laydown_config.yaml", + main_config_path=main_config_path, + lay_down_config_path=lay_down_config_path, ) - env.update_environent_obs() - - # we have three nodes and two links, with one service - # therefore the box observation space will have: - # * 5 columns (four fixed and one for the service) - # * 5 rows (3 nodes + 2 links) - assert env.env_obs.shape == (5, 5) + yield env -def test_creating_env_with_multidiscrete_obs(): - """Try creating env with MultiDiscrete observation space.""" - env = _get_primaite_env_from_config( +@pytest.mark.env_config_paths( + dict( main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml", lay_down_config_path=TEST_CONFIG_ROOT - / "multidiscrete_obs_space_laydown_config.yaml", + / "obs_tests/laydown_without_obs_space.yaml", ) +) +def test_default_obs_space(env: Primaite): + """Create environment with no obs space defined in config and check that the default obs space was created.""" env.update_environent_obs() - # we have three nodes and two links, with one service - # the nodes have hardware, OS, FS, and service, the links just have bandwidth, - # therefore we need 3*4 + 2 observations - assert env.env_obs.shape == (3 * 4 + 2,) + components = env.obs_handler.registered_obs_components + + assert len(components) == 1 + assert isinstance(components[0], NodeLinkTable) -def test_component_registration(): - """Test that we can register and deregister a component.""" +@pytest.mark.env_config_paths( + dict( + main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml", + lay_down_config_path=TEST_CONFIG_ROOT + / "obs_tests/laydown_without_obs_space.yaml", + ) +) +def test_registering_components(env: Primaite): + """Test regitering and deregistering a component.""" handler = ObservationsHandler() - component = NodeStatuses() + component = NodeStatuses(env) handler.register(component) assert component in handler.registered_obs_components handler.deregister(component) assert component not in handler.registered_obs_components + + +@pytest.mark.env_config_paths( + dict( + main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_no_agent.yaml", + lay_down_config_path=TEST_CONFIG_ROOT + / "obs_tests/laydown_with_NODE_LINK_TABLE.yaml", + ) +) +class TestNodeLinkTable: + """Test the NodeLinkTable observation component (in isolation).""" + + def test_obs_shape(self, env: Primaite): + """Try creating env with box observation space.""" + env.update_environent_obs() + + # we have three nodes and two links, with two service + # therefore the box observation space will have: + # * 5 rows (3 nodes + 2 links) + # * 6 columns (four fixed and two for the services) + assert env.env_obs.shape == (5, 6) + + # def test_value(self, env: Primaite): + # """""" + # ... + + +@pytest.mark.env_config_paths( + dict( + main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_no_agent.yaml", + lay_down_config_path=TEST_CONFIG_ROOT + / "obs_tests/laydown_with_NODE_STATUSES.yaml", + ) +) +class TestNodeStatuses: + """Test the NodeStatuses observation component (in isolation).""" + + def test_obs_shape(self, env: Primaite): + """Try creating env with NodeStatuses as the only component.""" + assert env.env_obs.shape == (15) + + def test_values(self, env: Primaite): + """Test that the hardware and software states are encoded correctly. + + The laydown has: + * one node with a compromised operating system state + * one node with two services, and the second service is overwhelmed. + * all other states are good or null + Therefore, the expected state is: + * node 1: + * hardware = good (1) + * OS = compromised (3) + * file system = good (1) + * service 1 = good (1) + * service 2 = good (1) + * node 2: + * hardware = good (1) + * OS = good (1) + * file system = good (1) + * service 1 = good (1) + * service 2 = overwhelmed (4) + * node 3 (switch): + * hardware = good (1) + * OS = good (1) + * file system = good (1) + * service 1 = n/a (0) + * service 2 = n/a (0) + """ + act = np.asarray([0, 0, 0, 0]) + obs, _, _, _ = env.step(act) + assert np.array_equal(obs, [1, 3, 1, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 0, 0]) + + +@pytest.mark.env_config_paths( + dict( + main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_no_agent.yaml", + lay_down_config_path=TEST_CONFIG_ROOT + / "obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml", + ) +) +class TestLinkTrafficLevels: + """Test the LinkTrafficLevels 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 links and two services, so the shape should be 2 * 2 + assert env.env_obs.shape == (2 * 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 20 bits of data over both links the whole time (via the first service) + * link bandwidth of 1000, therefore the utilisation is 2% + """ + act = np.asarray([0, 0, 0, 0]) + obs, reward, done, info = env.step(act) + obs, reward, done, info = env.step(act) + + # 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 20 bits of data via link1 and link2 on service 1. + # therefore the first and third elements should be 1 and all others 0 + assert np.array_equal(obs, [1, 0, 1, 0])