#2658: added monitored traffic into config + default obs space

tests/assets/configs/basic_switched_network.yaml

@@ -84,6 +84,11 @@ agents:
               num_files: 1
               num_nics: 2
               include_num_access: false
+              monitored_traffic:
+                icmp:
+                    - NONE
+                tcp:
+                    - DNS
               include_nmne: true
               routers:
                 - hostname: router_1

tests/integration_tests/game_layer/observations/test_nic_observations.py

@@ -5,12 +5,18 @@ import pytest
 import yaml
 from gymnasium import spaces
 
+from primaite.game.agent.interface import ProxyAgent
 from primaite.game.agent.observations.nic_observations import NICObservation
 from primaite.game.game import PrimaiteGame
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.host.host_node import NIC
-from primaite.simulator.network.nmne import CAPTURE_NMNE
+from primaite.simulator.network.hardware.nodes.host.server import Server
 from primaite.simulator.sim_container import Simulation
+from primaite.simulator.system.applications.database_client import DatabaseClient
+from primaite.simulator.system.applications.web_browser import WebBrowser
+from primaite.simulator.system.services.dns.dns_client import DNSClient
+from primaite.simulator.system.services.dns.dns_server import DNSServer
+from primaite.simulator.system.services.web_server.web_server import WebServer
 from tests import TEST_ASSETS_ROOT
 
 BASIC_CONFIG = TEST_ASSETS_ROOT / "configs/basic_switched_network.yaml"
@@ -31,6 +37,32 @@ def simulation(example_network) -> Simulation:
     # set simulation network as example network
     sim.network = example_network
 
+    computer: Computer = example_network.get_node_by_hostname("client_1")
+    server: Server = example_network.get_node_by_hostname("server_1")
+
+    web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
+    web_browser.run()
+
+    # Install DNS Client service on computer
+    computer.software_manager.install(DNSClient)
+    dns_client: DNSClient = computer.software_manager.software.get("DNSClient")
+    # set dns server
+    dns_client.dns_server = server.network_interface[1].ip_address
+
+    # Install Web Server service on server
+    server.software_manager.install(WebServer)
+    web_server_service: WebServer = server.software_manager.software.get("WebServer")
+    web_server_service.start()
+
+    # Install DNS Server service on server
+    server.software_manager.install(DNSServer)
+    dns_server: DNSServer = server.software_manager.software.get("DNSServer")
+    # register arcd.com to DNS
+    dns_server.dns_register(
+        domain_name="arcd.com",
+        domain_ip_address=server.network_interfaces[next(iter(server.network_interfaces))].ip_address,
+    )
+
     return sim
 
 
@@ -102,3 +134,71 @@ def test_config_nic_categories(simulation):
             high_nmne_threshold=9,
             include_nmne=True,
         )
+
+
+def test_nic_monitored_traffic(simulation):
+    monitored_traffic = {"icmp": ["NONE"], "tcp": ["DNS"]}
+
+    pc: Computer = simulation.network.get_node_by_hostname("client_1")
+    pc2: Computer = simulation.network.get_node_by_hostname("client_2")
+
+    nic_obs = NICObservation(
+        where=["network", "nodes", pc.hostname, "NICs", 1], include_nmne=True, monitored_traffic=monitored_traffic
+    )
+
+    simulation.pre_timestep(0)  # apply timestep to whole sim
+    simulation.apply_timestep(0)  # apply timestep to whole sim
+    traffic_obs = nic_obs.observe(simulation.describe_state()).get("monitored_traffic")
+
+    assert traffic_obs["icmp"]["inbound"] == 0
+    assert traffic_obs["icmp"]["outbound"] == 0
+
+    # send a ping
+    pc.ping(target_ip_address=pc2.network_interface[1].ip_address)
+    traffic_obs = nic_obs.observe(simulation.describe_state()).get("monitored_traffic")
+
+    assert traffic_obs["icmp"]["inbound"] == 1
+    assert traffic_obs["icmp"]["outbound"] == 1
+
+    simulation.pre_timestep(1)  # apply timestep to whole sim
+    simulation.apply_timestep(1)  # apply timestep to whole sim
+    traffic_obs = nic_obs.observe(simulation.describe_state()).get("monitored_traffic")
+
+    assert traffic_obs["icmp"]["inbound"] == 0
+    assert traffic_obs["icmp"]["outbound"] == 0
+    assert traffic_obs["tcp"][53]["inbound"] == 0
+    assert traffic_obs["tcp"][53]["outbound"] == 0
+
+    # send a database query
+    browser: WebBrowser = pc.software_manager.software.get("WebBrowser")
+    browser.target_url = f"http://arcd.com/"
+    browser.get_webpage()
+
+    traffic_obs = nic_obs.observe(simulation.describe_state()).get("monitored_traffic")
+    assert traffic_obs["icmp"]["inbound"] == 0
+    assert traffic_obs["icmp"]["outbound"] == 0
+    assert traffic_obs["tcp"][53]["inbound"] == 0
+    assert traffic_obs["tcp"][53]["outbound"] == 1  # getting a webpage sent a dns request out
+
+    simulation.pre_timestep(2)  # apply timestep to whole sim
+    simulation.apply_timestep(2)  # apply timestep to whole sim
+    traffic_obs = nic_obs.observe(simulation.describe_state()).get("monitored_traffic")
+
+    assert traffic_obs["icmp"]["inbound"] == 0
+    assert traffic_obs["icmp"]["outbound"] == 0
+    assert traffic_obs["tcp"][53]["inbound"] == 0
+    assert traffic_obs["tcp"][53]["outbound"] == 0
+
+
+def test_nic_monitored_traffic_config():
+    """Test that the config loads the monitored traffic config correctly."""
+    game: PrimaiteGame = load_config(BASIC_CONFIG)
+
+    # should have icmp and DNS
+    defender_agent: ProxyAgent = game.agents.get("defender")
+    cur_obs = defender_agent.observation_manager.current_observation
+
+    assert cur_obs["NODES"]["HOST0"]["NICS"][1]["monitored_traffic"] == {
+        "icmp": {"inbound": 0, "outbound": 0},
+        "tcp": {53: {"inbound": 0, "outbound": 0}},
+    }

tests/integration_tests/game_layer/observations/test_node_observations.py

@@ -32,6 +32,7 @@ def test_host_observation(simulation):
         num_services=1,
         include_num_access=False,
         include_nmne=False,
+        monitored_traffic=None,
         services=[],
         applications=[],
         folders=[],