#2350: configurable NMNE category thresholds

tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py

@@ -5,8 +5,9 @@ from typing import Union
 import yaml
 
 from primaite.config.load import data_manipulation_config_path
-from primaite.game.agent.interface import ProxyAgent, RandomAgent
+from primaite.game.agent.interface import ProxyAgent
 from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent
+from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent
 from primaite.game.game import APPLICATION_TYPES_MAPPING, PrimaiteGame, SERVICE_TYPES_MAPPING
 from primaite.simulator.network.container import Network
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
@@ -43,15 +44,15 @@ def test_example_config():
 
     # green agent 1
     assert "client_2_green_user" in game.agents
-    assert isinstance(game.agents["client_2_green_user"], RandomAgent)
+    assert isinstance(game.agents["client_2_green_user"], ProbabilisticAgent)
 
     # green agent 2
     assert "client_1_green_user" in game.agents
-    assert isinstance(game.agents["client_1_green_user"], RandomAgent)
+    assert isinstance(game.agents["client_1_green_user"], ProbabilisticAgent)
 
     # red agent
-    assert "client_1_data_manipulation_red_bot" in game.agents
-    assert isinstance(game.agents["client_1_data_manipulation_red_bot"], DataManipulationAgent)
+    assert "data_manipulation_attacker" in game.agents
+    assert isinstance(game.agents["data_manipulation_attacker"], DataManipulationAgent)
 
     # blue agent
     assert "defender" in game.agents

tests/integration_tests/configuration_file_parsing/test_game_options_config.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+from typing import Union
+
+import yaml
+
+from primaite.config.load import data_manipulation_config_path
+from primaite.game.game import PrimaiteGame
+from tests import TEST_ASSETS_ROOT
+
+BASIC_CONFIG = TEST_ASSETS_ROOT / "configs/basic_switched_network.yaml"
+
+
+def load_config(config_path: Union[str, Path]) -> PrimaiteGame:
+    """Returns a PrimaiteGame object which loads the contents of a given yaml path."""
+    with open(config_path, "r") as f:
+        cfg = yaml.safe_load(f)
+
+    return PrimaiteGame.from_config(cfg)
+
+
+def test_thresholds():
+    """Test that the game options can be parsed correctly."""
+    game = load_config(data_manipulation_config_path())
+
+    assert game.options.thresholds is not None

tests/integration_tests/game_layer/observations/test_observations.py

@@ -1,6 +1,6 @@
 import pytest
 
-from primaite.game.agent.observations.observations import NicObservation
+from primaite.game.agent.observations.nic_observations import NicObservation
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.host.host_node import NIC
 from primaite.simulator.sim_container import Simulation
@@ -33,3 +33,43 @@ def test_nic(simulation):
     nic.disable()
     observation_state = nic_obs.observe(simulation.describe_state())
     assert observation_state.get("nic_status") == 2  # disabled
+
+
+def test_nic_categories(simulation):
+    """Test the NIC observation nmne count categories."""
+    pc: Computer = simulation.network.get_node_by_hostname("client_1")
+
+    nic_obs = NicObservation(where=["network", "nodes", pc.hostname, "NICs", 1])
+
+    assert nic_obs.high_nmne_threshold == 10  # default
+    assert nic_obs.med_nmne_threshold == 5  # default
+    assert nic_obs.low_nmne_threshold == 0  # default
+
+    nic_obs = NicObservation(
+        where=["network", "nodes", pc.hostname, "NICs", 1],
+        low_nmne_threshold=3,
+        med_nmne_threshold=6,
+        high_nmne_threshold=9,
+    )
+
+    assert nic_obs.high_nmne_threshold == 9
+    assert nic_obs.med_nmne_threshold == 6
+    assert nic_obs.low_nmne_threshold == 3
+
+    with pytest.raises(Exception):
+        # should throw an error
+        NicObservation(
+            where=["network", "nodes", pc.hostname, "NICs", 1],
+            low_nmne_threshold=9,
+            med_nmne_threshold=6,
+            high_nmne_threshold=9,
+        )
+
+    with pytest.raises(Exception):
+        # should throw an error
+        NicObservation(
+            where=["network", "nodes", pc.hostname, "NICs", 1],
+            low_nmne_threshold=3,
+            med_nmne_threshold=9,
+            high_nmne_threshold=9,
+        )

tests/integration_tests/network/test_capture_nmne.py

@@ -1,4 +1,4 @@
-from primaite.game.agent.observations.observations import NicObservation
+from primaite.game.agent.observations.nic_observations import NicObservation
 from primaite.simulator.network.hardware.nodes.host.server import Server
 from primaite.simulator.network.nmne import set_nmne_config
 from primaite.simulator.sim_container import Simulation