Merge remote-tracking branch 'origin/dev' into feature/2533-improve-dev-tools-in-primaite-cli

CHANGELOG.md

@@ -23,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Introduced Developer CLI tools to assist with developing/debugging PrimAITE
   - Can be enabled via `primaite dev-mode enable`
   - Activating dev-mode will change the location where the sessions will be output - by default will output where the PrimAITE repository is located
+- Refactored all air-space usage to that a new instance of AirSpace is created for each instance of Network. This 1:1 relationship between network and airspace will allow parallelization.
 
 
 ## [Unreleased]

src/primaite/game/game.py

@@ -14,7 +14,6 @@ from primaite.game.agent.scripted_agents.probabilistic_agent import Probabilisti
 from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
 from primaite.game.agent.scripted_agents.tap001 import TAP001
 from primaite.game.science import graph_has_cycle, topological_sort
-from primaite.simulator.network.airspace import AIR_SPACE
 from primaite.simulator.network.hardware.base import NodeOperatingState
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.host.host_node import NIC
@@ -222,7 +221,6 @@ class PrimaiteGame:
         :return: A PrimaiteGame object.
         :rtype: PrimaiteGame
         """
-        AIR_SPACE.clear()
         game = cls()
         game.options = PrimaiteGameOptions(**cfg["game"])
         game.save_step_metadata = cfg.get("io_settings", {}).get("save_step_metadata") or False
@@ -274,7 +272,7 @@ class PrimaiteGame:
             elif n_type == "firewall":
                 new_node = Firewall.from_config(node_cfg)
             elif n_type == "wireless_router":
-                new_node = WirelessRouter.from_config(node_cfg)
+                new_node = WirelessRouter.from_config(node_cfg, airspace=net.airspace)
             elif n_type == "printer":
                 new_node = Printer(
                     hostname=node_cfg["hostname"],

src/primaite/simulator/network/airspace.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Any, Dict, Final, List, Optional
+from typing import Any, Dict, List, Optional
 
 from prettytable import PrettyTable
 
@@ -14,7 +14,7 @@ from primaite.simulator.system.core.packet_capture import PacketCapture
 
 _LOGGER = getLogger(__name__)
 
-__all__ = ["AIR_SPACE", "AirSpaceFrequency", "WirelessNetworkInterface", "IPWirelessNetworkInterface"]
+__all__ = ["AirSpaceFrequency", "WirelessNetworkInterface", "IPWirelessNetworkInterface"]
 
 
 class AirSpace:
@@ -100,18 +100,6 @@ class AirSpace:
                 wireless_interface.receive_frame(frame)
 
 
-AIR_SPACE: Final[AirSpace] = AirSpace()
-"""
-A singleton instance of the AirSpace class, representing the global wireless airspace.
-
-This instance acts as the central management point for all wireless communications within the simulated network
-environment. By default, there is only one airspace in the simulation, making this variable a singleton that
-manages the registration, removal, and transmission of wireless frames across all wireless network interfaces configured
-in the simulation. It ensures that wireless frames are appropriately transmitted to and received by wireless
-interfaces based on their operational status and frequency band.
-"""
-
-
 class AirSpaceFrequency(Enum):
     """Enumeration representing the operating frequencies for wireless communications."""
 
@@ -149,6 +137,7 @@ class WirelessNetworkInterface(NetworkInterface, ABC):
     and may define additional properties and methods specific to wireless technology.
     """
 
+    airspace: AirSpace
     frequency: AirSpaceFrequency = AirSpaceFrequency.WIFI_2_4
 
     def enable(self):
@@ -171,7 +160,7 @@ class WirelessNetworkInterface(NetworkInterface, ABC):
         self.pcap = PacketCapture(
             hostname=self._connected_node.hostname, port_num=self.port_num, port_name=self.port_name
         )
-        AIR_SPACE.add_wireless_interface(self)
+        self.airspace.add_wireless_interface(self)
 
     def disable(self):
         """Disable the network interface."""
@@ -182,7 +171,7 @@ class WirelessNetworkInterface(NetworkInterface, ABC):
             self._connected_node.sys_log.info(f"Network Interface {self} disabled")
         else:
             _LOGGER.debug(f"Interface {self} disabled")
-        AIR_SPACE.remove_wireless_interface(self)
+        self.airspace.remove_wireless_interface(self)
 
     def send_frame(self, frame: Frame) -> bool:
         """
@@ -198,7 +187,7 @@ class WirelessNetworkInterface(NetworkInterface, ABC):
         if self.enabled:
             frame.set_sent_timestamp()
             self.pcap.capture_outbound(frame)
-            AIR_SPACE.transmit(frame, self)
+            self.airspace.transmit(frame, self)
             return True
         # Cannot send Frame as the network interface is not enabled
         return False

src/primaite/simulator/network/container.py

@@ -5,9 +5,11 @@ import matplotlib.pyplot as plt
 import networkx as nx
 from networkx import MultiGraph
 from prettytable import MARKDOWN, PrettyTable
+from pydantic import Field
 
 from primaite import getLogger
 from primaite.simulator.core import RequestManager, RequestType, SimComponent
+from primaite.simulator.network.airspace import AirSpace
 from primaite.simulator.network.hardware.base import Link, Node, WiredNetworkInterface
 from primaite.simulator.network.hardware.nodes.host.server import Printer
 from primaite.simulator.system.applications.application import Application
@@ -28,7 +30,9 @@ class Network(SimComponent):
     """
 
     nodes: Dict[str, Node] = {}
+
     links: Dict[str, Link] = {}
+    airspace: AirSpace = Field(default_factory=lambda: AirSpace())
     _node_id_map: Dict[int, Node] = {}
     _link_id_map: Dict[int, Node] = {}
 

src/primaite/simulator/network/hardware/nodes/network/router.py

@@ -1546,7 +1546,7 @@ class Router(NetworkNode):
         print(table)
 
     @classmethod
-    def from_config(cls, cfg: dict) -> "Router":
+    def from_config(cls, cfg: dict, **kwargs) -> "Router":
         """Create a router based on a config dict.
 
         Schema:

src/primaite/simulator/network/hardware/nodes/network/wireless_router.py

@@ -3,7 +3,7 @@ from typing import Any, Dict, Union
 
 from pydantic import validate_call
 
-from primaite.simulator.network.airspace import AirSpaceFrequency, IPWirelessNetworkInterface
+from primaite.simulator.network.airspace import AirSpace, AirSpaceFrequency, IPWirelessNetworkInterface
 from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
 from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router, RouterInterface
 from primaite.simulator.network.transmission.data_link_layer import Frame
@@ -121,11 +121,14 @@ class WirelessRouter(Router):
 
     network_interfaces: Dict[str, Union[RouterInterface, WirelessAccessPoint]] = {}
     network_interface: Dict[int, Union[RouterInterface, WirelessAccessPoint]] = {}
+    airspace: AirSpace
 
-    def __init__(self, hostname: str, **kwargs):
+    def __init__(self, hostname: str, airspace: AirSpace, **kwargs):
-        super().__init__(hostname=hostname, num_ports=0, **kwargs)
+        super().__init__(hostname=hostname, num_ports=0, airspace=airspace, **kwargs)
 
-        self.connect_nic(WirelessAccessPoint(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0"))
+        self.connect_nic(
+            WirelessAccessPoint(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", airspace=airspace)
+        )
 
         self.connect_nic(RouterInterface(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0"))
 
@@ -215,7 +218,7 @@ class WirelessRouter(Router):
         )
 
     @classmethod
-    def from_config(cls, cfg: Dict) -> "WirelessRouter":
+    def from_config(cls, cfg: Dict, **kwargs) -> "WirelessRouter":
         """Generate the wireless router from config.
 
         Schema:
@@ -245,7 +248,7 @@ class WirelessRouter(Router):
         operating_state = (
             NodeOperatingState.ON if not (p := cfg.get("operating_state")) else NodeOperatingState[p.upper()]
         )
-        router = cls(hostname=cfg["hostname"], operating_state=operating_state)
+        router = cls(hostname=cfg["hostname"], operating_state=operating_state, airspace=kwargs["airspace"])
         if "router_interface" in cfg:
             ip_address = cfg["router_interface"]["ip_address"]
             subnet_mask = cfg["router_interface"]["subnet_mask"]

tests/assets/configs/wireless_wan_network_config.yaml

@@ -0,0 +1,77 @@
+game:
+  max_episode_length: 256
+  ports:
+    - ARP
+  protocols:
+    - ICMP
+    - TCP
+    - UDP
+
+simulation:
+  network:
+    nodes:
+      - type: computer
+        hostname: pc_a
+        ip_address: 192.168.0.2
+        subnet_mask: 255.255.255.0
+        default_gateway: 192.168.0.1
+        start_up_duration: 0
+
+      - type: computer
+        hostname: pc_b
+        ip_address: 192.168.2.2
+        subnet_mask: 255.255.255.0
+        default_gateway: 192.168.2.1
+        start_up_duration: 0
+
+      - type: wireless_router
+        hostname: router_1
+        start_up_duration: 0
+
+        router_interface:
+          ip_address: 192.168.0.1
+          subnet_mask: 255.255.255.0
+
+        wireless_access_point:
+          ip_address: 192.168.1.1
+          subnet_mask: 255.255.255.0
+          frequency: WIFI_2_4
+        acl:
+          1:
+            action: PERMIT
+        routes:
+          - address: 192.168.2.0 # PC B subnet
+            subnet_mask: 255.255.255.0
+            next_hop_ip_address: 192.168.1.2
+            metric: 0
+
+      - type: wireless_router
+        hostname: router_2
+        start_up_duration: 0
+
+        router_interface:
+          ip_address: 192.168.2.1
+          subnet_mask: 255.255.255.0
+
+        wireless_access_point:
+          ip_address: 192.168.1.2
+          subnet_mask: 255.255.255.0
+          frequency: WIFI_2_4
+        acl:
+          1:
+            action: PERMIT
+        routes:
+          - address: 192.168.0.0 # PC A subnet
+            subnet_mask: 255.255.255.0
+            next_hop_ip_address: 192.168.1.1
+            metric: 0
+    links:
+      - endpoint_a_hostname: pc_a
+        endpoint_a_port: 1
+        endpoint_b_hostname: router_1
+        endpoint_b_port: 2
+
+      - endpoint_a_hostname: pc_b
+        endpoint_a_port: 1
+        endpoint_b_hostname: router_2
+        endpoint_b_port: 2

tests/integration_tests/network/test_wireless_router.py

@@ -1,16 +1,18 @@
 import pytest
+import yaml
 
-from primaite.simulator.network.airspace import AIR_SPACE, AirSpaceFrequency
+from primaite.game.game import PrimaiteGame
 from primaite.simulator.network.container import Network
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.network.router import ACLAction
 from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
 from primaite.simulator.network.transmission.network_layer import IPProtocol
 from primaite.simulator.network.transmission.transport_layer import Port
+from tests import TEST_ASSETS_ROOT
 
 
 @pytest.fixture(scope="function")
-def setup_network():
+def wireless_wan_network():
     network = Network()
 
     # Configure PC A
@@ -25,7 +27,7 @@ def setup_network():
     network.add_node(pc_a)
 
     # Configure Router 1
-    router_1 = WirelessRouter(hostname="router_1", start_up_duration=0)
+    router_1 = WirelessRouter(hostname="router_1", start_up_duration=0, airspace=network.airspace)
     router_1.power_on()
     network.add_node(router_1)
 
@@ -49,7 +51,7 @@ def setup_network():
     network.add_node(pc_b)
 
     # Configure Router 2
-    router_2 = WirelessRouter(hostname="router_2", start_up_duration=0)
+    router_2 = WirelessRouter(hostname="router_2", start_up_duration=0, airspace=network.airspace)
     router_2.power_on()
     network.add_node(router_2)
 
@@ -63,7 +65,7 @@ def setup_network():
     router_1.configure_wireless_access_point("192.168.1.1", "255.255.255.0")
     router_2.configure_wireless_access_point("192.168.1.2", "255.255.255.0")
 
-    AIR_SPACE.show()
+    network.airspace.show()
 
     router_1.route_table.add_route(
         address="192.168.2.0", subnet_mask="255.255.255.0", next_hop_ip_address="192.168.1.2"
@@ -77,11 +79,35 @@ def setup_network():
     return pc_a, pc_b, router_1, router_2
 
 
-def test_cross_router_connectivity(setup_network):
+@pytest.fixture(scope="function")
-    pc_a, pc_b, router_1, router_2 = setup_network
+def wireless_wan_network_from_config_yaml():
+    config_path = TEST_ASSETS_ROOT / "configs" / "wireless_wan_network_config.yaml"
+
+    with open(config_path, "r") as f:
+        config_dict = yaml.safe_load(f)
+    network = PrimaiteGame.from_config(cfg=config_dict).simulation.network
+
+    network.airspace.show()
+
+    return network
+
+
+def test_cross_wireless_wan_connectivity(wireless_wan_network):
+    pc_a, pc_b, router_1, router_2 = wireless_wan_network
     # Ensure that PCs can ping across routers before any frequency change
     assert pc_a.ping(pc_a.default_gateway), "PC A should ping its default gateway successfully."
     assert pc_b.ping(pc_b.default_gateway), "PC B should ping its default gateway successfully."
 
     assert pc_a.ping(pc_b.network_interface[1].ip_address), "PC A should ping PC B across routers successfully."
     assert pc_b.ping(pc_a.network_interface[1].ip_address), "PC B should ping PC A across routers successfully."
+
+
+def test_cross_wireless_wan_connectivity_from_yaml(wireless_wan_network_from_config_yaml):
+    pc_a = wireless_wan_network_from_config_yaml.get_node_by_hostname("pc_a")
+    pc_b = wireless_wan_network_from_config_yaml.get_node_by_hostname("pc_b")
+
+    assert pc_a.ping(pc_a.default_gateway), "PC A should ping its default gateway successfully."
+    assert pc_b.ping(pc_b.default_gateway), "PC B should ping its default gateway successfully."
+
+    assert pc_a.ping(pc_b.network_interface[1].ip_address), "PC A should ping PC B across routers successfully."
+    assert pc_b.ping(pc_a.network_interface[1].ip_address), "PC B should ping PC A across routers successfully."

tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_wireless_router.py

@@ -1,97 +0,0 @@
-from ipaddress import IPv4Address
-
-from primaite.simulator.network.hardware.nodes.network.router import ACLAction
-from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
-from primaite.simulator.network.transmission.network_layer import IPProtocol
-from primaite.simulator.network.transmission.transport_layer import Port
-
-
-def test_wireless_router_from_config():
-    cfg = {
-        "ref": "router_2",
-        "type": "wireless_router",
-        "hostname": "router_2",
-        "router_interface": {
-            "ip_address": "192.168.1.1",
-            "subnet_mask": "255.255.255.0",
-        },
-        "wireless_access_point": {
-            "ip_address": "192.170.1.1",
-            "subnet_mask": "255.255.255.0",
-            "frequency": "WIFI_2_4",
-        },
-        "acl": {
-            0: {
-                "action": "PERMIT",
-                "src_port": "POSTGRES_SERVER",
-                "dst_port": "POSTGRES_SERVER",
-            },
-            1: {
-                "action": "PERMIT",
-                "protocol": "ICMP",
-            },
-            2: {
-                "action": "PERMIT",
-                "src_ip": "100.100.100.1",
-                "dst_ip": "100.100.101.1",
-            },
-            3: {
-                "action": "PERMIT",
-                "src_ip": "100.100.102.0",
-                "dst_ip": "100.100.103.0",
-                "src_wildcard_mask": "0.0.0.255",
-                "dst_wildcard_mask": "0.0.0.255",
-            },
-            20: {
-                "action": "DENY",
-            },
-        },
-    }
-
-    rt = WirelessRouter.from_config(cfg=cfg)
-
-    r0 = rt.acl.acl[0]
-    assert r0.action == ACLAction.PERMIT
-    assert r0.src_port == r0.dst_port == Port.POSTGRES_SERVER
-    assert r0.src_ip_address == r0.dst_ip_address == r0.dst_wildcard_mask == r0.src_wildcard_mask == r0.protocol == None
-
-    r1 = rt.acl.acl[1]
-    assert r1.action == ACLAction.PERMIT
-    assert r1.protocol == IPProtocol.ICMP
-    assert (
-        r1.src_ip_address
-        == r1.dst_ip_address
-        == r1.dst_wildcard_mask
-        == r1.src_wildcard_mask
-        == r1.src_port
-        == r1.dst_port
-        == None
-    )
-
-    r2 = rt.acl.acl[2]
-    assert r2.action == ACLAction.PERMIT
-    assert r2.src_ip_address == IPv4Address("100.100.100.1")
-    assert r2.dst_ip_address == IPv4Address("100.100.101.1")
-    assert r2.src_wildcard_mask == r2.dst_wildcard_mask == None
-    assert r2.src_port == r2.dst_port == r2.protocol == None
-
-    r3 = rt.acl.acl[3]
-    assert r3.action == ACLAction.PERMIT
-    assert r3.src_ip_address == IPv4Address("100.100.102.0")
-    assert r3.dst_ip_address == IPv4Address("100.100.103.0")
-    assert r3.src_wildcard_mask == IPv4Address("0.0.0.255")
-    assert r3.dst_wildcard_mask == IPv4Address("0.0.0.255")
-    assert r3.src_port == r3.dst_port == r3.protocol == None
-
-    r20 = rt.acl.acl[20]
-    assert r20.action == ACLAction.DENY
-    assert (
-        r20.src_ip_address
-        == r20.dst_ip_address
-        == r20.src_wildcard_mask
-        == r20.dst_wildcard_mask
-        == r20.src_port
-        == r20.dst_port
-        == r20.protocol
-        == None
-    )