From 9417cd85abfd72d00f0ee2c84c9c95db045042f2 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 7 Jun 2023 15:25:11 +0100 Subject: [PATCH] Apply suggestions from code review. --- src/primaite/common/config_values_main.py | 1 + src/primaite/environment/observations.py | 25 ++-- src/primaite/environment/primaite_env.py | 13 +-- src/primaite/main.py | 4 + ...n_with_NODE_STATUSES.yaml => laydown.yaml} | 5 +- .../laydown_with_LINK_TRAFFIC_LEVELS.yaml | 110 ------------------ .../laydown_with_NODE_LINK_TABLE.yaml | 77 ------------ .../obs_tests/laydown_without_obs_space.yaml | 74 ------------ .../main_config_LINK_TRAFFIC_LEVELS.yaml | 96 +++++++++++++++ .../main_config_NODE_LINK_TABLE.yaml | 93 +++++++++++++++ .../obs_tests/main_config_NODE_STATUSES.yaml | 93 +++++++++++++++ ...gent.yaml => main_config_without_obs.yaml} | 2 +- tests/conftest.py | 4 + tests/test_observation_space.py | 49 ++++---- 14 files changed, 338 insertions(+), 308 deletions(-) rename tests/config/obs_tests/{laydown_with_NODE_STATUSES.yaml => laydown.yaml} (95%) delete mode 100644 tests/config/obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml delete mode 100644 tests/config/obs_tests/laydown_with_NODE_LINK_TABLE.yaml delete mode 100644 tests/config/obs_tests/laydown_without_obs_space.yaml create mode 100644 tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml create mode 100644 tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml create mode 100644 tests/config/obs_tests/main_config_NODE_STATUSES.yaml rename tests/config/obs_tests/{main_config_no_agent.yaml => main_config_without_obs.yaml} (98%) diff --git a/src/primaite/common/config_values_main.py b/src/primaite/common/config_values_main.py index 3493f9d2..f822b77f 100644 --- a/src/primaite/common/config_values_main.py +++ b/src/primaite/common/config_values_main.py @@ -9,6 +9,7 @@ class ConfigValuesMain(object): """Init.""" # Generic self.agent_identifier = "" # the agent in use + self.observation_config = None # observation space config self.num_episodes = 0 # number of episodes to train over self.num_steps = 0 # number of steps in an episode self.time_delay = 0 # delay between steps (ms) - applies to generic agents only diff --git a/src/primaite/environment/observations.py b/src/primaite/environment/observations.py index a598d6db..9e71ef1b 100644 --- a/src/primaite/environment/observations.py +++ b/src/primaite/environment/observations.py @@ -1,7 +1,7 @@ """Module for handling configurable observation spaces in PrimAITE.""" import logging from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List, Tuple, Union +from typing import TYPE_CHECKING, Dict, Final, List, Tuple, Union import numpy as np from gym import spaces @@ -56,9 +56,9 @@ class NodeLinkTable(AbstractObservationComponent): ``(12, 7)`` """ - _FIXED_PARAMETERS = 4 - _MAX_VAL = 1_000_000 - _DATA_TYPE = np.int64 + _FIXED_PARAMETERS: int = 4 + _MAX_VAL: int = 1_000_000 + _DATA_TYPE: type = np.int64 def __init__(self, env: "Primaite"): super().__init__(env) @@ -159,7 +159,7 @@ class NodeStatuses(AbstractObservationComponent): :type env: Primaite """ - _DATA_TYPE = np.int64 + _DATA_TYPE: type = np.int64 def __init__(self, env: "Primaite"): super().__init__(env) @@ -231,7 +231,7 @@ class LinkTrafficLevels(AbstractObservationComponent): :type quantisation_levels: int, optional """ - _DATA_TYPE = np.int64 + _DATA_TYPE: type = np.int64 def __init__( self, @@ -239,7 +239,14 @@ class LinkTrafficLevels(AbstractObservationComponent): combine_service_traffic: bool = False, quantisation_levels: int = 5, ): - assert quantisation_levels >= 3 + if quantisation_levels < 3: + _msg = ( + f"quantisation_levels must be 3 or more because the lowest and highest levels are " + f"reserved for 0% and 100% link utilisation, got {quantisation_levels} instead. " + f"Resetting to default value (5)" + ) + _LOGGER.warning(_msg) + quantisation_levels = 5 super().__init__(env) @@ -296,7 +303,7 @@ class ObservationsHandler: Each component can also define further parameters to make them more flexible. """ - registry = { + _REGISTRY: Final[Dict[str, type]] = { "NODE_LINK_TABLE": NodeLinkTable, "NODE_STATUSES": NodeStatuses, "LINK_TRAFFIC_LEVELS": LinkTrafficLevels, @@ -384,7 +391,7 @@ class ObservationsHandler: for component_cfg in obs_space_config["components"]: # Figure out which class can instantiate the desired component comp_type = component_cfg["name"] - comp_class = cls.registry[comp_type] + comp_class = cls._REGISTRY[comp_type] # Create the component with options from the YAML options = component_cfg.get("options") or {} diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 0ff58100..7995c4f7 100644 --- a/src/primaite/environment/primaite_env.py +++ b/src/primaite/environment/primaite_env.py @@ -48,10 +48,10 @@ class Primaite(Env): """PRIMmary AI Training Evironment (Primaite) class.""" # Action Space contants - ACTION_SPACE_NODE_PROPERTY_VALUES = 5 - ACTION_SPACE_NODE_ACTION_VALUES = 4 - ACTION_SPACE_ACL_ACTION_VALUES = 3 - ACTION_SPACE_ACL_PERMISSION_VALUES = 2 + ACTION_SPACE_NODE_PROPERTY_VALUES: int = 5 + ACTION_SPACE_NODE_ACTION_VALUES: int = 4 + ACTION_SPACE_ACL_ACTION_VALUES: int = 3 + ACTION_SPACE_ACL_PERMISSION_VALUES: int = 2 def __init__(self, _config_values, _transaction_list): """ @@ -148,6 +148,8 @@ class Primaite(Env): # stores the observation config from the yaml, default is NODE_LINK_TABLE self.obs_config: dict = {"components": [{"name": "NODE_LINK_TABLE"}]} + if self.config_values.observation_config is not None: + self.obs_config = self.config_values.observation_config # Observation Handler manages the user-configurable observation space. # It will be initialised later. @@ -690,9 +692,6 @@ class Primaite(Env): elif item["itemType"] == "ACTIONS": # Get the action information self.get_action_info(item) - elif item["itemType"] == "OBSERVATION_SPACE": - # Get the observation information - self.save_obs_config(item) elif item["itemType"] == "STEPS": # Get the steps information self.get_steps_info(item) diff --git a/src/primaite/main.py b/src/primaite/main.py index c963dd00..5f8aa5e2 100644 --- a/src/primaite/main.py +++ b/src/primaite/main.py @@ -163,6 +163,10 @@ def load_config_values(): try: # Generic config_values.agent_identifier = config_data["agentIdentifier"] + if "observationSpace" in config_data: + config_values.observation_config = config_data["observationSpace"] + else: + config_values.observation_config = None config_values.num_episodes = int(config_data["numEpisodes"]) config_values.time_delay = int(config_data["timeDelay"]) config_values.config_filename_use_case = ( diff --git a/tests/config/obs_tests/laydown_with_NODE_STATUSES.yaml b/tests/config/obs_tests/laydown.yaml similarity index 95% rename from tests/config/obs_tests/laydown_with_NODE_STATUSES.yaml rename to tests/config/obs_tests/laydown.yaml index 56ff3725..d3b131db 100644 --- a/tests/config/obs_tests/laydown_with_NODE_STATUSES.yaml +++ b/tests/config/obs_tests/laydown.yaml @@ -1,8 +1,5 @@ - itemType: ACTIONS type: NODE -- itemType: OBSERVATION_SPACE - components: - - name: NODE_STATUSES - itemType: STEPS steps: 5 - itemType: PORTS @@ -82,7 +79,7 @@ id: '5' startStep: 0 endStep: 5 - load: 20 + load: 999 protocol: TCP port: '80' source: '1' diff --git a/tests/config/obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml b/tests/config/obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml deleted file mode 100644 index e65ea306..00000000 --- a/tests/config/obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml +++ /dev/null @@ -1,110 +0,0 @@ -- itemType: ACTIONS - type: NODE -- itemType: OBSERVATION_SPACE - components: - - 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 -- 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' - -######################################### -# IERS -- itemType: GREEN_IER - id: '5' - startStep: 0 - endStep: 5 - load: 999 - 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 deleted file mode 100644 index 0ceefbfa..00000000 --- a/tests/config/obs_tests/laydown_with_NODE_LINK_TABLE.yaml +++ /dev/null @@ -1,77 +0,0 @@ -- itemType: ACTIONS - type: NODE -- itemType: OBSERVATION_SPACE - components: - - name: NODE_LINK_TABLE -- 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/laydown_without_obs_space.yaml b/tests/config/obs_tests/laydown_without_obs_space.yaml deleted file mode 100644 index 3ef214da..00000000 --- a/tests/config/obs_tests/laydown_without_obs_space.yaml +++ /dev/null @@ -1,74 +0,0 @@ -- 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_LINK_TRAFFIC_LEVELS.yaml b/tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml new file mode 100644 index 00000000..cdb741f3 --- /dev/null +++ b/tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml @@ -0,0 +1,96 @@ +# 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 +observationSpace: + components: + - name: LINK_TRAFFIC_LEVELS + options: + combine_service_traffic: false + quantisation_levels: 8 + +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: 1_000_000_000 + +# 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/config/obs_tests/main_config_NODE_LINK_TABLE.yaml b/tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml new file mode 100644 index 00000000..19d220c8 --- /dev/null +++ b/tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml @@ -0,0 +1,93 @@ +# 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 +observationSpace: + components: + - name: NODE_LINK_TABLE + +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: 1_000_000_000 + +# 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/config/obs_tests/main_config_NODE_STATUSES.yaml b/tests/config/obs_tests/main_config_NODE_STATUSES.yaml new file mode 100644 index 00000000..25520ccc --- /dev/null +++ b/tests/config/obs_tests/main_config_NODE_STATUSES.yaml @@ -0,0 +1,93 @@ +# 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 +observationSpace: + components: + - name: NODE_STATUSES + +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: 1_000_000_000 + +# 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/config/obs_tests/main_config_no_agent.yaml b/tests/config/obs_tests/main_config_without_obs.yaml similarity index 98% rename from tests/config/obs_tests/main_config_no_agent.yaml rename to tests/config/obs_tests/main_config_without_obs.yaml index f632dca9..43ee251f 100644 --- a/tests/config/obs_tests/main_config_no_agent.yaml +++ b/tests/config/obs_tests/main_config_without_obs.yaml @@ -21,7 +21,7 @@ agentLoadFile: C:\[Path]\[agent_saved_filename.zip] # Environment config values # The high value for the observation space -observationSpaceHighValue: 1000000000 +observationSpaceHighValue: 1_000_000_000 # Reward values # Generic diff --git a/tests/conftest.py b/tests/conftest.py index 1e987223..f3728b63 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,6 +19,10 @@ def _get_primaite_env_from_config( def load_config_values(): config_values.agent_identifier = config_data["agentIdentifier"] + if "observationSpace" in config_data: + config_values.observation_config = config_data["observationSpace"] + else: + config_values.observation_config = None config_values.num_episodes = int(config_data["numEpisodes"]) config_values.time_delay = int(config_data["timeDelay"]) config_values.config_filename_use_case = lay_down_config_path diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index ae862c96..dcf98ae1 100644 --- a/tests/test_observation_space.py +++ b/tests/test_observation_space.py @@ -27,9 +27,8 @@ def env(request): @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", + main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml", + lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", ) ) def test_default_obs_space(env: Primaite): @@ -44,9 +43,8 @@ def test_default_obs_space(env: Primaite): @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", + main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml", + lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", ) ) def test_registering_components(env: Primaite): @@ -61,9 +59,9 @@ def test_registering_components(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_LINK_TABLE.yaml", + main_config_path=TEST_CONFIG_ROOT + / "obs_tests/main_config_NODE_LINK_TABLE.yaml", + lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", ) ) class TestNodeLinkTable: @@ -92,17 +90,17 @@ class TestNodeLinkTable: * Node 1: * 1 (id) * 1 (good hardware state) - * 1 (good OS state) + * 3 (compromised OS state) * 1 (good file system state) - * 1 (good service1 state) - * 1 (good service2 state) + * 1 (good TCP state) + * 1 (good UDP state) * Node 2: * 2 (id) * 1 (good hardware state) * 1 (good OS state) * 1 (good file system state) - * 1 (good service1 state) - * 1 (good service2 state) + * 1 (good TCP state) + * 4 (overwhelmed UDP state) * Node 3 (active node): * 3 (id) * 1 (good hardware state) @@ -115,14 +113,14 @@ class TestNodeLinkTable: * 0 (n/a hardware state) * 0 (n/a OS state) * 0 (n/a file system state) - * 0 (no traffic for service1) + * 999 (999 traffic for service1) * 0 (no traffic for service2) * Link 2: * 5 (id) * 0 (good hardware state) * 0 (good OS state) * 0 (good file system state) - * 0 (no traffic service1) + * 999 (999 traffic service1) * 0 (no traffic for service2) """ act = np.asarray([0, 0, 0, 0]) @@ -131,20 +129,19 @@ class TestNodeLinkTable: assert np.array_equal( obs, [ - [1, 1, 1, 1, 1, 1], - [2, 1, 1, 1, 1, 1], + [1, 1, 3, 1, 1, 1], + [2, 1, 1, 1, 1, 4], [3, 1, 1, 1, 0, 0], - [4, 0, 0, 0, 0, 0], - [5, 0, 0, 0, 0, 0], + [4, 0, 0, 0, 999, 0], + [5, 0, 0, 0, 999, 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_NODE_STATUSES.yaml", + main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_NODE_STATUSES.yaml", + lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", ) ) class TestNodeStatuses: @@ -188,9 +185,9 @@ class TestNodeStatuses: @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", + main_config_path=TEST_CONFIG_ROOT + / "obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml", + lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml", ) ) class TestLinkTrafficLevels: