diff --git a/CHANGELOG.md b/CHANGELOG.md index 8426f5dc..d8f30cb0 100644 --- a/CHANGELOG.md +++ b/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] diff --git a/docs/source/configuration/simulation/nodes/images/primaite_example_multi_lan_with_internet_network_dark.png b/docs/source/configuration/simulation/nodes/images/primaite_example_multi_lan_with_internet_network_dark.png index 9c2e6842..a8eab7ee 100644 Binary files a/docs/source/configuration/simulation/nodes/images/primaite_example_multi_lan_with_internet_network_dark.png and b/docs/source/configuration/simulation/nodes/images/primaite_example_multi_lan_with_internet_network_dark.png differ diff --git a/docs/source/configuration/simulation/nodes/images/primaite_example_multi_lan_with_internet_network_light.png b/docs/source/configuration/simulation/nodes/images/primaite_example_multi_lan_with_internet_network_light.png index 92996b89..c3483cf9 100644 Binary files a/docs/source/configuration/simulation/nodes/images/primaite_example_multi_lan_with_internet_network_light.png and b/docs/source/configuration/simulation/nodes/images/primaite_example_multi_lan_with_internet_network_light.png differ diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 336c27df..ab68ea2d 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." diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_wireless_router.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_wireless_router.py deleted file mode 100644 index 494f5a15..00000000 --- a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_wireless_router.py +++ /dev/null @@ -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 - )