2623 - finish testing action mask

2024-07-10 13:46:30 +01:00
parent 7201b7b8e0
commit 0c58c3969a
8 changed files with 198 additions and 27 deletions
--- a/tests/integration_tests/game_layer/test_action_mask.py
+++ b/tests/integration_tests/game_layer/test_action_mask.py
@@ -0,0 +1,161 @@
+# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
+from primaite.session.environment import PrimaiteGymEnv
+from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
+from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
+from primaite.simulator.system.services.service import ServiceOperatingState
+from tests.conftest import TEST_ASSETS_ROOT
+
+CFG_PATH = TEST_ASSETS_ROOT / "configs/test_primaite_session.yaml"
+
+
+def test_mask_contents_correct():
+    env = PrimaiteGymEnv(CFG_PATH)
+    game = env.game
+    sim = game.simulation
+    net = sim.network
+    mask = game.action_mask("defender")
+    agent = env.agent
+    node_list = agent.action_manager.node_names
+    action_map = agent.action_manager.action_map
+
+    # CHECK NIC ENABLE/DISABLE ACTIONS
+    for action_num, action in action_map.items():
+        mask = game.action_mask("defender")
+        act_type, act_params = action
+
+        if act_type == "NODE_NIC_ENABLE":
+            node_name = node_list[act_params["node_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            nic_obj = node_obj.network_interface[act_params["nic_id"] + 1]
+            assert nic_obj.enabled
+            assert not mask[action_num]
+            nic_obj.disable()
+            mask = game.action_mask("defender")
+            assert mask[action_num]
+            nic_obj.enable()
+
+        if act_type == "NODE_NIC_DISABLE":
+            node_name = node_list[act_params["node_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            nic_obj = node_obj.network_interface[act_params["nic_id"] + 1]
+            assert nic_obj.enabled
+            assert mask[action_num]
+            nic_obj.disable()
+            mask = game.action_mask("defender")
+            assert not mask[action_num]
+            nic_obj.enable()
+
+        if act_type == "ROUTER_ACL_ADDRULE":
+            assert mask[action_num]
+
+        if act_type == "ROUTER_ACL_REMOVERULE":
+            assert mask[action_num]
+
+        if act_type == "NODE_RESET":
+            node_name = node_list[act_params["node_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            assert node_obj.operating_state is NodeOperatingState.ON
+            assert mask[action_num]
+            node_obj.operating_state = NodeOperatingState.OFF
+            mask = game.action_mask("defender")
+            assert not mask[action_num]
+            node_obj.operating_state = NodeOperatingState.ON
+
+        if act_type == "NODE_SHUTDOWN":
+            node_name = node_list[act_params["node_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            assert node_obj.operating_state is NodeOperatingState.ON
+            assert mask[action_num]
+            node_obj.operating_state = NodeOperatingState.OFF
+            mask = game.action_mask("defender")
+            assert not mask[action_num]
+            node_obj.operating_state = NodeOperatingState.ON
+
+        if act_type == "NODE_OS_SCAN":
+            node_name = node_list[act_params["node_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            assert node_obj.operating_state is NodeOperatingState.ON
+            assert mask[action_num]
+            node_obj.operating_state = NodeOperatingState.OFF
+            mask = game.action_mask("defender")
+            assert not mask[action_num]
+            node_obj.operating_state = NodeOperatingState.ON
+
+        if act_type == "NODE_STARTUP":
+            node_name = node_list[act_params["node_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            assert node_obj.operating_state is NodeOperatingState.ON
+            assert not mask[action_num]
+            node_obj.operating_state = NodeOperatingState.OFF
+            mask = game.action_mask("defender")
+            assert mask[action_num]
+            node_obj.operating_state = NodeOperatingState.ON
+
+        if act_type == "DONOTHING":
+            assert mask[action_num]
+
+        if act_type == "NODE_SERVICE_DISABLE":
+            assert mask[action_num]
+
+        if act_type in ["NODE_SERVICE_SCAN", "NODE_SERVICE_STOP", "NODE_SERVICE_PAUSE"]:
+            node_name = node_list[act_params["node_id"]]
+            service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            service_obj = node_obj.software_manager.software.get(service_name)
+            assert service_obj.operating_state is ServiceOperatingState.RUNNING
+            assert mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.DISABLED
+            mask = game.action_mask("defender")
+            assert not mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.RUNNING
+
+        if act_type == "NODE_SERVICE_RESUME":
+            node_name = node_list[act_params["node_id"]]
+            service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            service_obj = node_obj.software_manager.software.get(service_name)
+            assert service_obj.operating_state is ServiceOperatingState.RUNNING
+            assert not mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.PAUSED
+            mask = game.action_mask("defender")
+            assert mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.RUNNING
+
+        if act_type == "NODE_SERVICE_START":
+            node_name = node_list[act_params["node_id"]]
+            service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            service_obj = node_obj.software_manager.software.get(service_name)
+            assert service_obj.operating_state is ServiceOperatingState.RUNNING
+            assert not mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.STOPPED
+            mask = game.action_mask("defender")
+            assert mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.RUNNING
+
+        if act_type == "NODE_SERVICE_ENABLE":
+            node_name = node_list[act_params["node_id"]]
+            service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]]
+            node_obj = net.get_node_by_hostname(node_name)
+            service_obj = node_obj.software_manager.software.get(service_name)
+            assert service_obj.operating_state is ServiceOperatingState.RUNNING
+            assert not mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.DISABLED
+            mask = game.action_mask("defender")
+            assert mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.RUNNING
+
+        if act_type in ["NODE_FILE_SCAN", "NODE_FILE_CHECKHASH", "NODE_FILE_DELETE"]:
+            node_name = node_list[act_params["node_id"]]
+            folder_name = agent.action_manager.get_folder_name_by_idx(act_params["node_id"], act_params["folder_id"])
+            file_name = agent.action_manager.get_file_name_by_idx(
+                act_params["node_id"], act_params["folder_id"], act_params["file_id"]
+            )
+            node_obj = net.get_node_by_hostname(node_name)
+            file_obj = node_obj.file_system.get_file(folder_name, file_name, include_deleted=True)
+            assert not file_obj.deleted
+            assert mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.DISABLED
+            mask = game.action_mask("defender")
+            assert mask[action_num]
+            service_obj.operating_state = ServiceOperatingState.RUNNING