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

@@ -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

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

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'

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()

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