diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a96b27c..df46fa5d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added the ability for a DatabaseService to terminate a connection. - Added active_connection to DatabaseClientConnection so that if the connection is terminated active_connection is set to False and the object can no longer be used. - Added additional show functions to enable connection inspection. +- 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] diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 908b5148..d8399332 100644 --- a/src/primaite/game/game.py +++ b/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"], diff --git a/src/primaite/simulator/network/airspace.py b/src/primaite/simulator/network/airspace.py index 3c5c048c..8a00a4a4 100644 --- a/src/primaite/simulator/network/airspace.py +++ b/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 diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index e9a938ce..17308c97 100644 --- a/src/primaite/simulator/network/container.py +++ b/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] = {} diff --git a/src/primaite/simulator/network/hardware/nodes/network/router.py b/src/primaite/simulator/network/hardware/nodes/network/router.py index b7563937..53bb4827 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/router.py +++ b/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: diff --git a/src/primaite/simulator/network/hardware/nodes/network/wireless_router.py b/src/primaite/simulator/network/hardware/nodes/network/wireless_router.py index f66ebd27..9e5d4dd4 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/wireless_router.py +++ b/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): - super().__init__(hostname=hostname, num_ports=0, **kwargs) + def __init__(self, hostname: str, airspace: AirSpace, **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"] diff --git a/tests/assets/configs/wireless_wan_network_config.yaml b/tests/assets/configs/wireless_wan_network_config.yaml new file mode 100644 index 00000000..c8f61bad --- /dev/null +++ b/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 diff --git a/tests/integration_tests/network/test_wireless_router.py b/tests/integration_tests/network/test_wireless_router.py index 0e458974..d739bd0b 100644 --- a/tests/integration_tests/network/test_wireless_router.py +++ b/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): - pc_a, pc_b, router_1, router_2 = setup_network +@pytest.fixture(scope="function") +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."