From 411f0a320fb651179be2c2fb65b77579b8018aee Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Thu, 8 Feb 2024 10:53:30 +0000 Subject: [PATCH] #2248 - Final run over all the docstrings after running pre-commit. All tests now working. Updated CHANGELOG.md. --- CHANGELOG.md | 13 + src/primaite/game/agent/observations.py | 8 +- src/primaite/game/game.py | 2 +- src/primaite/simulator/network/container.py | 2 +- .../simulator/network/hardware/base.py | 72 +++-- .../network_interface/layer_3_interface.py | 9 - .../network_interface/wired/__init__.py | 0 .../wired/router_interface.py | 0 .../wireless/wireless_access_point.py | 3 +- .../wireless/wireless_nic.py | 3 +- .../network/hardware/nodes/host/computer.py | 2 +- .../network/hardware/nodes/host/host_node.py | 120 ++++--- .../network/hardware/nodes/host/server.py | 1 - .../hardware/nodes/network/network_node.py | 25 +- .../network/hardware/nodes/network/router.py | 304 ++++++++++++------ .../network/hardware/nodes/network/switch.py | 5 +- src/primaite/simulator/network/networks.py | 21 +- .../simulator/network/protocols/icmp.py | 2 +- .../network/transmission/data_link_layer.py | 1 - .../network/transmission/network_layer.py | 7 +- .../simulator/system/core/packet_capture.py | 1 - .../simulator/system/core/session_manager.py | 112 +++++-- .../simulator/system/core/software_manager.py | 29 +- .../simulator/system/services/arp/arp.py | 30 +- .../simulator/system/services/icmp/icmp.py | 33 +- .../system/services/ntp/ntp_server.py | 6 +- src/primaite/simulator/system/software.py | 2 +- src/primaite/utils/validators.py | 6 +- tests/conftest.py | 30 +- .../network/test_broadcast.py | 7 +- .../network/test_frame_transmission.py | 1 - .../network/test_network_creation.py | 1 - .../integration_tests/network/test_routing.py | 4 +- .../test_dos_bot_and_server.py | 2 +- .../system/test_database_on_node.py | 14 +- .../system/test_dns_client_server.py | 3 +- .../system/test_web_client_server.py | 5 +- .../test_web_client_server_and_database.py | 5 +- .../_applications/test_database_client.py | 13 +- .../_system/_applications/test_web_browser.py | 2 +- .../_system/_services/test_dns_server.py | 4 +- 41 files changed, 582 insertions(+), 328 deletions(-) delete mode 100644 src/primaite/simulator/network/hardware/network_interface/layer_3_interface.py delete mode 100644 src/primaite/simulator/network/hardware/network_interface/wired/__init__.py delete mode 100644 src/primaite/simulator/network/hardware/network_interface/wired/router_interface.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d8706ad..68bc3b81 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,6 +60,12 @@ SessionManager. - **Custom Layer-3 Processing**: The `RouterNIC` class includes custom handling for network frames, bypassing standard Node NIC's Layer 3 broadcast/unicast checks. This allows for more efficient routing behavior in network scenarios where router-specific frame processing is required. - **Enhanced Frame Reception**: The `receive_frame` method in `RouterNIC` is tailored to handle frames based on Layer 2 (Ethernet) checks, focusing on MAC address-based routing and broadcast frame acceptance. - **Subnet-Wide Broadcasting for Services and Applications**: Implemented the ability for services and applications to conduct broadcasts across an entire IPv4 subnet within the network simulation framework. +- Introduced the `NetworkInterface` abstract class to provide a common interface for all network interfaces. Subclasses are divided into two main categories: `WiredNetworkInterface` and `WirelessNetworkInterface`, each serving as an abstract base class (ABC) for more specific interface types. Under `WiredNetworkInterface`, the subclasses `NIC` and `SwitchPort` were added. For wireless interfaces, `WirelessNIC` and `WirelessAccessPoint` are the subclasses under `WirelessNetworkInterface`. +- Added `Layer3Interface` as an abstract base class for networking functionalities at layer 3, including IP addressing and routing capabilities. This class is inherited by `NIC`, `WirelessNIC`, and `WirelessAccessPoint` to provide them with layer 3 capabilities, facilitating their role in both wired and wireless networking contexts with IP-based communication. +- Created the `ARP` and `ICMP` service classes to handle Address Resolution Protocol operations and Internet Control Message Protocol messages, respectively, with `RouterARP` and `RouterICMP` for router-specific implementations. +- Created `HostNode` as a subclass of `Node`, extending its functionality with host-specific services and applications. This class is designed to represent end-user devices like computers or servers that can initiate and respond to network communications. +- Introduced a new `IPV4Address` type in the Pydantic model for enhanced validation and auto-conversion of IPv4 addresses from strings using an `ipv4_validator`. + ### Changed - Integrated the RouteTable into the Routers frame processing. @@ -67,6 +73,9 @@ SessionManager. - **NIC Functionality Update**: Updated the Network Interface Card (`NIC`) functionality to support Layer 3 (L3) broadcasts. - **Layer 3 Broadcast Handling**: Enhanced the existing `NIC` classes to correctly process and handle Layer 3 broadcasts. This update allows devices using standard NICs to effectively participate in network activities that involve L3 broadcasting. - **Improved Frame Reception Logic**: The `receive_frame` method of the `NIC` class has been updated to include additional checks and handling for L3 broadcasts, ensuring proper frame processing in a wider range of network scenarios. +- Standardised the way network interfaces are accessed across all `Node` subclasses (`HostNode`, `Router`, `Switch`) by maintaining a comprehensive `network_interface` attribute. This attribute captures all network interfaces by their port number, streamlining the management and interaction with network interfaces across different types of nodes. +- Refactored all tests to utilise new `Node` subclasses (`Computer`, `Server`, `Router`, `Switch`) instead of creating generic `Node` instances and manually adding network interfaces. This change aligns test setups more closely with the intended use cases and hierarchies within the network simulation framework. +- Updated all tests to employ the `Network()` class for managing nodes and their connections, ensuring a consistent and structured approach to setting up network topologies in testing scenarios. ### Removed @@ -74,6 +83,10 @@ SessionManager. - Removed legacy training modules - Removed tests for legacy code +### Fixed +- Addressed network transmission issues that previously allowed ARP requests to be incorrectly routed and repeated across different subnets. This fix ensures ARP requests are correctly managed and confined to their appropriate network segments. +- Resolved problems in `Node` and its subclasses where the default gateway configuration was not properly utilized for communications across different subnets. This correction ensures that nodes effectively use their configured default gateways for outbound communications to other network segments, thereby enhancing the network's routing functionality and reliability. + ## [2.0.0] - 2023-07-26 diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index 8f1c739c..715e594e 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -494,7 +494,9 @@ class NodeObservation(AbstractObservation): obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)} obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)} obs["operating_status"] = node_state["operating_state"] - obs["NETWORK_INTERFACES"] = {i + 1: network_interface.observe(state) for i, network_interface in enumerate(self.network_interfaces)} + obs["NETWORK_INTERFACES"] = { + i + 1: network_interface.observe(state) for i, network_interface in enumerate(self.network_interfaces) + } if self.logon_status: obs["logon_status"] = 0 @@ -508,7 +510,9 @@ class NodeObservation(AbstractObservation): "SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}), "FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}), "operating_status": spaces.Discrete(5), - "NETWORK_INTERFACES": spaces.Dict({i + 1: network_interface.space for i, network_interface in enumerate(self.network_interfaces)}), + "NETWORK_INTERFACES": spaces.Dict( + {i + 1: network_interface.space for i, network_interface in enumerate(self.network_interfaces)} + ), } if self.logon_status: space_shape["logon_status"] = spaces.Discrete(3) diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index c25f64ab..f1f66e40 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -14,8 +14,8 @@ from primaite.session.io import SessionIO, SessionIOSettings 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 -from primaite.simulator.network.hardware.nodes.network.router import Router from primaite.simulator.network.hardware.nodes.host.server import Server +from primaite.simulator.network.hardware.nodes.network.router import Router from primaite.simulator.network.hardware.nodes.network.switch import Switch from primaite.simulator.sim_container import Simulation from primaite.simulator.system.applications.database_client import DatabaseClient diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 4789134b..d3a26e73 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -9,8 +9,8 @@ from primaite import getLogger from primaite.simulator.core import RequestManager, RequestType, SimComponent from primaite.simulator.network.hardware.base import Link, Node, WiredNetworkInterface from primaite.simulator.network.hardware.nodes.host.computer import Computer -from primaite.simulator.network.hardware.nodes.network.router import Router from primaite.simulator.network.hardware.nodes.host.server import Server +from primaite.simulator.network.hardware.nodes.network.router import Router from primaite.simulator.network.hardware.nodes.network.switch import Switch from primaite.simulator.system.applications.application import Application from primaite.simulator.system.services.service import Service diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index b7b6d3d4..c742ca33 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -2,14 +2,13 @@ from __future__ import annotations import re import secrets -from abc import abstractmethod, ABC +from abc import ABC, abstractmethod from ipaddress import IPv4Address, IPv4Network from pathlib import Path -from typing import Any, Union -from typing import Dict, Optional +from typing import Any, Dict, Optional, Union from prettytable import MARKDOWN, PrettyTable -from pydantic import Field, BaseModel +from pydantic import BaseModel, Field from primaite import getLogger from primaite.exceptions import NetworkError @@ -48,7 +47,7 @@ def generate_mac_address(oui: Optional[str] = None) -> str: _LOGGER.error(msg) raise ValueError(msg) oui_bytes = [int(chunk, 16) for chunk in oui.split(":")] - mac = oui_bytes + random_bytes[len(oui_bytes):] + mac = oui_bytes + random_bytes[len(oui_bytes) :] else: mac = random_bytes @@ -198,9 +197,7 @@ class WiredNetworkInterface(NetworkInterface, ABC): return if not self._connected_link: - self._connected_node.sys_log.info( - f"Interface {self} cannot be enabled as there is no Link connected." - ) + self._connected_node.sys_log.info(f"Interface {self} cannot be enabled as there is no Link connected.") return self.enabled = True @@ -225,9 +222,9 @@ class WiredNetworkInterface(NetworkInterface, ABC): """ Connect this network interface to a specified link. - This method establishes a connection between the network interface and a network link if the network interface is not already - connected. If the network interface is already connected to a link, it logs an error and does not change the existing - connection. + This method establishes a connection between the network interface and a network link if the network interface + is not already connected. If the network interface is already connected to a link, it logs an error and does + not change the existing connection. :param link: The Link instance to connect to this network interface. """ @@ -246,8 +243,8 @@ class WiredNetworkInterface(NetworkInterface, ABC): """ Disconnect the network interface from its connected Link, if any. - This method removes the association between the network interface and its connected Link. It updates the connected Link's - endpoints to reflect the disconnection. + This method removes the association between the network interface and its connected Link. It updates the + connected Link's endpoints to reflect the disconnection. """ if self._connected_link.endpoint_a == self: self._connected_link.endpoint_a = None @@ -298,6 +295,7 @@ class Layer3Interface(BaseModel, ABC): :ivar IPv4Address subnet_mask: The subnet mask assigned to the interface. This mask helps in determining the network segment that the interface belongs to and is used in IP routing decisions. """ + ip_address: IPV4Address "The IP address assigned to the interface for communication on an IP-based network." @@ -357,10 +355,23 @@ class IPWiredNetworkInterface(WiredNetworkInterface, Layer3Interface, ABC): Derived classes should define specific behaviors and properties of an IP-capable wired network interface, customizing it for their specific use cases. """ + _connected_link: Optional[Link] = None "The network link to which the network interface is connected." def model_post_init(self, __context: Any) -> None: + """ + Performs post-initialisation checks to ensure the model's IP configuration is valid. + + This method is invoked after the initialisation of a network model object to validate its network settings, + particularly to ensure that the assigned IP address is not a network address. This validation is crucial for + maintaining the integrity of network simulations and avoiding configuration errors that could lead to + unrealistic or incorrect behavior. + + :param __context: Contextual information or parameters passed to the method, used for further initializing or + validating the model post-creation. + :raises ValueError: If the IP address is the same as the network address, indicating an incorrect configuration. + """ if self.ip_network.network_address == self.ip_address: raise ValueError(f"{self.ip_address}/{self.subnet_mask} must not be a network address") @@ -380,6 +391,17 @@ class IPWiredNetworkInterface(WiredNetworkInterface, Layer3Interface, ABC): return state def enable(self): + """ + Enables this wired network interface and attempts to send a "hello" message to the default gateway. + + This method activates the network interface, making it operational for network communications. After enabling, + it tries to initiate a default gateway "hello" process, typically to establish initial connectivity and resolve + the default gateway's MAC address. This step is crucial for ensuring the interface can successfully send data + to and receive data from the network. + + The method safely handles cases where the connected node might not have a default gateway set or the + `default_gateway_hello` method is not defined, ignoring such errors to proceed without interruption. + """ super().enable() try: pass @@ -440,8 +462,8 @@ class IPWirelessNetworkInterface(WiredNetworkInterface, Layer3Interface, ABC): As an abstract class, `IPWirelessNetworkInterface` does not implement specific methods but ensures that any derived class provides implementations for the functionalities of both `WirelessNetworkInterface` and `Layer3Interface`. This setup is ideal for representing network interfaces in devices that require wireless connections and are capable - of IP routing and addressing, such as wireless routers, access points, and wireless end-host devices like smartphones - and laptops. + of IP routing and addressing, such as wireless routers, access points, and wireless end-host devices like + smartphones and laptops. This class should be extended by concrete classes that define specific behaviors and properties of an IP-capable wireless network interface. @@ -804,8 +826,10 @@ class Node(SimComponent): { "hostname": self.hostname, "operating_state": self.operating_state.value, - "NICs": {eth_num: network_interface.describe_state() for eth_num, network_interface in - self.network_interface.items()}, + "NICs": { + eth_num: network_interface.describe_state() + for eth_num, network_interface in self.network_interface.items() + }, "file_system": self.file_system.describe_state(), "applications": {app.name: app.describe_state() for app in self.applications.values()}, "services": {svc.name: svc.describe_state() for svc in self.services.values()}, @@ -816,7 +840,7 @@ class Node(SimComponent): return state def show(self, markdown: bool = False): - "Show function that calls both show NIC and show open ports." + """Show function that calls both show NIC and show open ports.""" self.show_nic(markdown) self.show_open_ports(markdown) @@ -833,6 +857,14 @@ class Node(SimComponent): @property def has_enabled_network_interface(self) -> bool: + """ + Checks if the node has at least one enabled network interface. + + Iterates through all network interfaces associated with the node to determine if at least one is enabled. This + property is essential for determining the node's ability to communicate within the network. + + :return: True if there is at least one enabled network interface; otherwise, False. + """ for network_interface in self.network_interfaces.values(): if network_interface.enabled: return True @@ -1014,7 +1046,7 @@ class Node(SimComponent): """ if self.operating_state.ON: self.is_resetting = True - self.sys_log.info(f"Resetting") + self.sys_log.info("Resetting") self.power_off() def connect_nic(self, network_interface: NetworkInterface): @@ -1097,7 +1129,7 @@ class Node(SimComponent): self.software_manager.arp.add_arp_cache_entry( ip_address=frame.ip.src_ip_address, mac_address=frame.ethernet.src_mac_addr, - network_interface=from_network_interface + network_interface=from_network_interface, ) else: return diff --git a/src/primaite/simulator/network/hardware/network_interface/layer_3_interface.py b/src/primaite/simulator/network/hardware/network_interface/layer_3_interface.py deleted file mode 100644 index fdfd3b26..00000000 --- a/src/primaite/simulator/network/hardware/network_interface/layer_3_interface.py +++ /dev/null @@ -1,9 +0,0 @@ -from abc import ABC -from ipaddress import IPv4Network -from typing import Dict - -from pydantic import BaseModel - -from primaite.utils.validators import IPV4Address - - diff --git a/src/primaite/simulator/network/hardware/network_interface/wired/__init__.py b/src/primaite/simulator/network/hardware/network_interface/wired/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/primaite/simulator/network/hardware/network_interface/wired/router_interface.py b/src/primaite/simulator/network/hardware/network_interface/wired/router_interface.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/primaite/simulator/network/hardware/network_interface/wireless/wireless_access_point.py b/src/primaite/simulator/network/hardware/network_interface/wireless/wireless_access_point.py index f94b7faa..646c12f4 100644 --- a/src/primaite/simulator/network/hardware/network_interface/wireless/wireless_access_point.py +++ b/src/primaite/simulator/network/hardware/network_interface/wireless/wireless_access_point.py @@ -2,7 +2,6 @@ from typing import Dict from primaite.simulator.network.hardware.base import WirelessNetworkInterface from primaite.simulator.network.hardware.network_interface.layer_3_interface import Layer3Interface - from primaite.simulator.network.transmission.data_link_layer import Frame @@ -81,4 +80,4 @@ class WirelessAccessPoint(WirelessNetworkInterface, Layer3Interface): :return: A string combining the port number, MAC address and IP address of the NIC. """ - return f"Port {self.port_num}: {self.mac_address}/{self.ip_address}" \ No newline at end of file + return f"Port {self.port_num}: {self.mac_address}/{self.ip_address}" diff --git a/src/primaite/simulator/network/hardware/network_interface/wireless/wireless_nic.py b/src/primaite/simulator/network/hardware/network_interface/wireless/wireless_nic.py index 12172608..40f357a0 100644 --- a/src/primaite/simulator/network/hardware/network_interface/wireless/wireless_nic.py +++ b/src/primaite/simulator/network/hardware/network_interface/wireless/wireless_nic.py @@ -2,7 +2,6 @@ from typing import Dict from primaite.simulator.network.hardware.base import WirelessNetworkInterface from primaite.simulator.network.hardware.network_interface.layer_3_interface import Layer3Interface - from primaite.simulator.network.transmission.data_link_layer import Frame @@ -78,4 +77,4 @@ class WirelessNIC(WirelessNetworkInterface, Layer3Interface): :return: A string combining the port number, MAC address and IP address of the NIC. """ - return f"Port {self.port_num}: {self.mac_address}/{self.ip_address}" \ No newline at end of file + return f"Port {self.port_num}: {self.mac_address}/{self.ip_address}" diff --git a/src/primaite/simulator/network/hardware/nodes/host/computer.py b/src/primaite/simulator/network/hardware/nodes/host/computer.py index dc75df69..0b13163e 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/computer.py +++ b/src/primaite/simulator/network/hardware/nodes/host/computer.py @@ -28,5 +28,5 @@ class Computer(HostNode): * Applications: * Web Browser """ - pass + pass diff --git a/src/primaite/simulator/network/hardware/nodes/host/host_node.py b/src/primaite/simulator/network/hardware/nodes/host/host_node.py index bd13e7e2..17390751 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/host_node.py +++ b/src/primaite/simulator/network/hardware/nodes/host/host_node.py @@ -1,15 +1,13 @@ from __future__ import annotations -from typing import Dict, Any -from typing import Optional +from ipaddress import IPv4Address +from typing import Any, Dict, Optional from primaite import getLogger -from primaite.simulator.network.hardware.base import IPWiredNetworkInterface, Link -from primaite.simulator.network.hardware.base import Node +from primaite.simulator.network.hardware.base import IPWiredNetworkInterface, Link, Node from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState from primaite.simulator.network.transmission.data_link_layer import Frame from primaite.simulator.system.applications.web_browser import WebBrowser -from primaite.simulator.system.core.packet_capture import PacketCapture from primaite.simulator.system.services.arp.arp import ARP, ARPPacket from primaite.simulator.system.services.dns.dns_client import DNSClient from primaite.simulator.system.services.ftp.ftp_client import FTPClient @@ -20,43 +18,45 @@ from primaite.utils.validators import IPV4Address _LOGGER = getLogger(__name__) -# Lives here due to pydantic circular dependency issue :( class HostARP(ARP): """ The Host ARP Service. - Extends the ARP service with functionalities specific to a host within the network. It provides mechanisms to - resolve and cache MAC addresses and NICs for given IP addresses, focusing on the host's perspective, including - handling the default gateway. + Extends the ARP service for host-specific functionalities within a network, focusing on resolving and caching + MAC addresses and network interfaces (NICs) based on IP addresses, especially concerning the default gateway. + + This specialized ARP service for hosts facilitates efficient network communication by managing ARP entries + and handling ARP requests and replies with additional logic for default gateway processing. """ def get_default_gateway_mac_address(self) -> Optional[str]: """ - Retrieves the MAC address of the default gateway from the ARP cache. + Retrieves the MAC address of the default gateway as known from the ARP cache. - :return: The MAC address of the default gateway if it exists in the ARP cache, otherwise None. + :return: The MAC address of the default gateway if present in the ARP cache; otherwise, None. """ if self.software_manager.node.default_gateway: return self.get_arp_cache_mac_address(self.software_manager.node.default_gateway) def get_default_gateway_network_interface(self) -> Optional[NIC]: """ - Retrieves the NIC associated with the default gateway from the ARP cache. + Obtains the network interface card (NIC) associated with the default gateway from the ARP cache. - :return: The NIC associated with the default gateway if it exists in the ARP cache, otherwise None. + :return: The NIC associated with the default gateway if it exists in the ARP cache; otherwise, None. """ if self.software_manager.node.default_gateway and self.software_manager.node.has_enabled_network_interface: return self.get_arp_cache_network_interface(self.software_manager.node.default_gateway) def _get_arp_cache_mac_address( - self, ip_address: IPV4Address, is_reattempt: bool = False, is_default_gateway_attempt: bool = False + self, ip_address: IPV4Address, is_reattempt: bool = False, is_default_gateway_attempt: bool = False ) -> Optional[str]: """ Internal method to retrieve the MAC address associated with an IP address from the ARP cache. :param ip_address: The IP address whose MAC address is to be retrieved. :param is_reattempt: Indicates if this call is a reattempt after a failed initial attempt. - :param is_default_gateway_attempt: Indicates if this call is an attempt to get the default gateway's MAC address. + :param is_default_gateway_attempt: Indicates if this call is an attempt to get the default gateway's MAC + address. :return: The MAC address associated with the IP address if found, otherwise None. """ arp_entry = self.arp.get(ip_address) @@ -76,22 +76,23 @@ class HostARP(ARP): if not is_default_gateway_attempt: self.send_arp_request(self.software_manager.node.default_gateway) return self._get_arp_cache_mac_address( - ip_address=self.software_manager.node.default_gateway, is_reattempt=True, - is_default_gateway_attempt=True + ip_address=self.software_manager.node.default_gateway, + is_reattempt=True, + is_default_gateway_attempt=True, ) return None - def get_arp_cache_mac_address(self, ip_address: IPV4Address) -> Optional[str]: + def get_arp_cache_mac_address(self, ip_address: IPv4Address) -> Optional[str]: """ - Retrieves the MAC address associated with an IP address from the ARP cache. + Retrieves the MAC address associated with a given IP address from the ARP cache. - :param ip_address: The IP address whose MAC address is to be retrieved. - :return: The MAC address associated with the IP address if found, otherwise None. - """ + :param ip_address: The IP address for which the MAC address is sought. + :return: The MAC address if available in the ARP cache; otherwise, None. + """ return self._get_arp_cache_mac_address(ip_address) def _get_arp_cache_network_interface( - self, ip_address: IPV4Address, is_reattempt: bool = False, is_default_gateway_attempt: bool = False + self, ip_address: IPV4Address, is_reattempt: bool = False, is_default_gateway_attempt: bool = False ) -> Optional[NIC]: """ Internal method to retrieve the NIC associated with an IP address from the ARP cache. @@ -118,17 +119,18 @@ class HostARP(ARP): if not is_default_gateway_attempt: self.send_arp_request(self.software_manager.node.default_gateway) return self._get_arp_cache_network_interface( - ip_address=self.software_manager.node.default_gateway, is_reattempt=True, - is_default_gateway_attempt=True + ip_address=self.software_manager.node.default_gateway, + is_reattempt=True, + is_default_gateway_attempt=True, ) return None - def get_arp_cache_network_interface(self, ip_address: IPV4Address) -> Optional[NIC]: + def get_arp_cache_network_interface(self, ip_address: IPv4Address) -> Optional[NIC]: """ - Retrieves the NIC associated with an IP address from the ARP cache. + Retrieves the network interface card (NIC) associated with a given IP address from the ARP cache. - :param ip_address: The IP address whose NIC is to be retrieved. - :return: The NIC associated with the IP address if found, otherwise None. + :param ip_address: The IP address for which the associated NIC is sought. + :return: The NIC if available in the ARP cache; otherwise, None. """ return self._get_arp_cache_network_interface(ip_address) @@ -146,15 +148,17 @@ class HostARP(ARP): # Unmatched ARP Request if arp_packet.target_ip_address != from_network_interface.ip_address: self.sys_log.info( - f"Ignoring ARP request for {arp_packet.target_ip_address}. Current IP address is {from_network_interface.ip_address}" + f"Ignoring ARP request for {arp_packet.target_ip_address}. Current IP address is " + f"{from_network_interface.ip_address}" ) return # Matched ARP request # TODO: try taking this out self.add_arp_cache_entry( - ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, - network_interface=from_network_interface + ip_address=arp_packet.sender_ip_address, + mac_address=arp_packet.sender_mac_addr, + network_interface=from_network_interface, ) arp_packet = arp_packet.generate_reply(from_network_interface.mac_address) self.send_arp_reply(arp_packet) @@ -175,12 +179,25 @@ class NIC(IPWiredNetworkInterface): and disconnect from network links and to manage the enabled/disabled state of the interface. - Layer3Interface: Provides properties for Layer 3 network configuration, such as IP address and subnet mask. """ + _connected_link: Optional[Link] = None "The network link to which the network interface is connected." wake_on_lan: bool = False "Indicates if the NIC supports Wake-on-LAN functionality." def model_post_init(self, __context: Any) -> None: + """ + Performs post-initialisation checks to ensure the model's IP configuration is valid. + + This method is invoked after the initialisation of a network model object to validate its network settings, + particularly to ensure that the assigned IP address is not a network address. This validation is crucial for + maintaining the integrity of network simulations and avoiding configuration errors that could lead to + unrealistic or incorrect behavior. + + :param __context: Contextual information or parameters passed to the method, used for further initializing or + validating the model post-creation. + :raises ValueError: If the IP address is the same as the network address, indicating an incorrect configuration. + """ if self.ip_network.network_address == self.ip_address: raise ValueError(f"{self.ip_address}/{self.subnet_mask} must not be a network address") @@ -255,21 +272,24 @@ class HostNode(Node): """ Represents a host node in the network. - Extends the basic functionality of a Node with host-specific services and applications. A host node typically - represents an end-user device in the network, such as a Computer or a Server, and is capable of initiating and - responding to network communications. + An end-user device within the network, such as a computer or server, equipped with the capability to initiate and + respond to network communications. + + A `HostNode` extends the base `Node` class by incorporating host-specific services and applications, thereby + simulating the functionalities typically expected from a networked end-user device. + + **Example**:: - Example: >>> pc_a = HostNode( - hostname="pc_a", - ip_address="192.168.1.10", - subnet_mask="255.255.255.0", - default_gateway="192.168.1.1" - ) + ... hostname="pc_a", + ... ip_address="192.168.1.10", + ... subnet_mask="255.255.255.0", + ... default_gateway="192.168.1.1" + ... ) >>> pc_a.power_on() - The host comes pre-installed with core functionalities and a suite of services and applications, making it ready - for various network operations and tasks. These include: + The host node comes pre-equipped with a range of core functionalities, services, and applications necessary + for engaging in various network operations and tasks. Core Functionality: ------------------- @@ -291,6 +311,7 @@ class HostNode(Node): * Web Browser: Provides web browsing capabilities. """ + network_interfaces: Dict[str, NIC] = {} "The Network Interfaces on the node." network_interface: Dict[int, NIC] = {} @@ -301,7 +322,12 @@ class HostNode(Node): self.connect_nic(NIC(ip_address=ip_address, subnet_mask=subnet_mask)) def _install_system_software(self): - """Install System Software - software that is usually provided with the OS.""" + """ + Installs the system software and network services typically found on an operating system. + + This method equips the host with essential network services and applications, preparing it for various + network-related tasks and operations. + """ # ARP Service self.software_manager.install(HostARP) @@ -323,6 +349,12 @@ class HostNode(Node): super()._install_system_software() def default_gateway_hello(self): + """ + Sends a hello message to the default gateway to establish connectivity and resolve the gateway's MAC address. + + This method is invoked to ensure the host node can communicate with its default gateway, primarily to confirm + network connectivity and populate the ARP cache with the gateway's MAC address. + """ if self.operating_state == NodeOperatingState.ON and self.default_gateway: self.software_manager.arp.get_default_gateway_mac_address() diff --git a/src/primaite/simulator/network/hardware/nodes/host/server.py b/src/primaite/simulator/network/hardware/nodes/host/server.py index 148a277f..9f5157ad 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/server.py +++ b/src/primaite/simulator/network/hardware/nodes/host/server.py @@ -28,4 +28,3 @@ class Server(HostNode): * Applications: * Web Browser """ - diff --git a/src/primaite/simulator/network/hardware/nodes/network/network_node.py b/src/primaite/simulator/network/hardware/nodes/network/network_node.py index c7a2060b..ebdb6ed8 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/network_node.py +++ b/src/primaite/simulator/network/hardware/nodes/network/network_node.py @@ -1,9 +1,30 @@ -from primaite.simulator.network.hardware.base import Node, NetworkInterface +from abc import abstractmethod + +from primaite.simulator.network.hardware.base import NetworkInterface, Node from primaite.simulator.network.transmission.data_link_layer import Frame class NetworkNode(Node): - """""" + """ + Represents an abstract base class for a network node that can receive and process network frames. + This class provides a common interface for network nodes such as routers and switches, defining the essential + behavior that allows these devices to handle incoming network traffic. Implementations of this class must + provide functionality for receiving and processing frames received on their network interfaces. + """ + + @abstractmethod def receive_frame(self, frame: Frame, from_network_interface: NetworkInterface): + """ + Abstract method that must be implemented by subclasses to define how to receive and process frames. + + This method is called when a frame is received by a network interface belonging to this node. Subclasses + should implement the logic to process the frame, including examining its contents, making forwarding decisions, + or performing any necessary actions based on the frame's protocol and destination. + + :param frame: The network frame that has been received. + :type frame: Frame + :param from_network_interface: The network interface on which the frame was received. + :type from_network_interface: NetworkInterface + """ pass diff --git a/src/primaite/simulator/network/hardware/nodes/network/router.py b/src/primaite/simulator/network/hardware/nodes/network/router.py index e5f4cdcd..3a22931e 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/router.py +++ b/src/primaite/simulator/network/hardware/nodes/network/router.py @@ -3,25 +3,22 @@ from __future__ import annotations import secrets from enum import Enum from ipaddress import IPv4Address, IPv4Network -from typing import Dict, Any -from typing import List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union from prettytable import MARKDOWN, PrettyTable -from pydantic import ValidationError from primaite.simulator.core import RequestManager, RequestType, SimComponent from primaite.simulator.network.hardware.base import IPWiredNetworkInterface from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState from primaite.simulator.network.hardware.nodes.network.network_node import NetworkNode from primaite.simulator.network.protocols.arp import ARPPacket -from primaite.simulator.network.protocols.icmp import ICMPType, ICMPPacket +from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType from primaite.simulator.network.transmission.data_link_layer import Frame from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.core.sys_log import SysLog from primaite.simulator.system.services.arp.arp import ARP from primaite.simulator.system.services.icmp.icmp import ICMP -from primaite.utils.validators import IPV4Address class ACLAction(Enum): @@ -205,14 +202,14 @@ class AccessControlList(SimComponent): return self._acl def add_rule( - self, - action: ACLAction, - protocol: Optional[IPProtocol] = None, - src_ip_address: Optional[Union[str, IPv4Address]] = None, - src_port: Optional[Port] = None, - dst_ip_address: Optional[Union[str, IPv4Address]] = None, - dst_port: Optional[Port] = None, - position: int = 0, + self, + action: ACLAction, + protocol: Optional[IPProtocol] = None, + src_ip_address: Optional[Union[str, IPv4Address]] = None, + src_port: Optional[Port] = None, + dst_ip_address: Optional[Union[str, IPv4Address]] = None, + dst_port: Optional[Port] = None, + position: int = 0, ) -> None: """ Add a new ACL rule. @@ -259,12 +256,12 @@ class AccessControlList(SimComponent): raise ValueError(f"Cannot remove ACL rule, position {position} is out of bounds.") def is_permitted( - self, - protocol: IPProtocol, - src_ip_address: Union[str, IPv4Address], - src_port: Optional[Port], - dst_ip_address: Union[str, IPv4Address], - dst_port: Optional[Port], + self, + protocol: IPProtocol, + src_ip_address: Union[str, IPv4Address], + src_port: Optional[Port], + dst_ip_address: Union[str, IPv4Address], + dst_port: Optional[Port], ) -> Tuple[bool, Optional[Union[str, ACLRule]]]: """ Check if a packet with the given properties is permitted through the ACL. @@ -286,23 +283,23 @@ class AccessControlList(SimComponent): continue if ( - (rule.src_ip_address == src_ip_address or rule.src_ip_address is None) - and (rule.dst_ip_address == dst_ip_address or rule.dst_ip_address is None) - and (rule.protocol == protocol or rule.protocol is None) - and (rule.src_port == src_port or rule.src_port is None) - and (rule.dst_port == dst_port or rule.dst_port is None) + (rule.src_ip_address == src_ip_address or rule.src_ip_address is None) + and (rule.dst_ip_address == dst_ip_address or rule.dst_ip_address is None) + and (rule.protocol == protocol or rule.protocol is None) + and (rule.src_port == src_port or rule.src_port is None) + and (rule.dst_port == dst_port or rule.dst_port is None) ): return rule.action == ACLAction.PERMIT, rule return self.implicit_action == ACLAction.PERMIT, f"Implicit {self.implicit_action.name}" def get_relevant_rules( - self, - protocol: IPProtocol, - src_ip_address: Union[str, IPv4Address], - src_port: Port, - dst_ip_address: Union[str, IPv4Address], - dst_port: Port, + self, + protocol: IPProtocol, + src_ip_address: Union[str, IPv4Address], + src_port: Port, + dst_ip_address: Union[str, IPv4Address], + dst_port: Port, ) -> List[ACLRule]: """ Get the list of relevant rules for a packet with given properties. @@ -324,11 +321,11 @@ class AccessControlList(SimComponent): continue if ( - (rule.src_ip_address == src_ip_address or rule.src_ip_address is None) - or (rule.dst_ip_address == dst_ip_address or rule.dst_ip_address is None) - or (rule.protocol == protocol or rule.protocol is None) - or (rule.src_port == src_port or rule.src_port is None) - or (rule.dst_port == dst_port or rule.dst_port is None) + (rule.src_ip_address == src_ip_address or rule.src_ip_address is None) + or (rule.dst_ip_address == dst_ip_address or rule.dst_ip_address is None) + or (rule.protocol == protocol or rule.protocol is None) + or (rule.src_port == src_port or rule.src_port is None) + or (rule.dst_port == dst_port or rule.dst_port is None) ): relevant_rules.append(rule) @@ -445,11 +442,11 @@ class RouteTable(SimComponent): pass def add_route( - self, - address: Union[IPv4Address, str], - subnet_mask: Union[IPv4Address, str], - next_hop_ip_address: Union[IPv4Address, str], - metric: float = 0.0, + self, + address: Union[IPv4Address, str], + subnet_mask: Union[IPv4Address, str], + next_hop_ip_address: Union[IPv4Address, str], + metric: float = 0.0, ): """ Add a route to the routing table. @@ -539,16 +536,34 @@ class RouteTable(SimComponent): class RouterARP(ARP): """ - Inherits from ARPCache and adds router-specific ARP packet processing. + Extends ARP functionality with router-specific ARP packet processing capabilities. - :ivar SysLog sys_log: A system log for logging messages. - :ivar Router router: The router to which this ARP cache belongs. + This class is designed to manage ARP requests and replies within a router, handling both the resolution of MAC + addresses for IP addresses within the router's networks and the forwarding of ARP requests to other networks + based on routing information. """ + router: Optional[Router] = None def _get_arp_cache_mac_address( - self, ip_address: IPV4Address, is_reattempt: bool = False, is_default_route_attempt: bool = False + self, ip_address: IPv4Address, is_reattempt: bool = False, is_default_route_attempt: bool = False ) -> Optional[str]: + """ + Attempts to retrieve the MAC address associated with the given IP address from the ARP cache. + + If the address is not in the cache, an ARP request may be sent, and the method may reattempt the lookup. + + :param ip_address: The IP address for which to find the corresponding MAC address. + :type ip_address: IPv4Address + :param is_reattempt: Indicates whether this call is a reattempt after a failed initial attempt to find the MAC + address. + :type is_reattempt: bool + :param is_default_route_attempt: Indicates whether the attempt is being made to resolve the MAC address for the + default route. + :type is_default_route_attempt: bool + :return: The MAC address associated with the given IP address, if found; otherwise, None. + :rtype: Optional[str] + """ arp_entry = self.arp.get(ip_address) if arp_entry: @@ -558,9 +573,7 @@ class RouterARP(ARP): if self.router.ip_is_in_router_interface_subnet(ip_address): self.send_arp_request(ip_address) return self._get_arp_cache_mac_address( - ip_address=ip_address, - is_reattempt=True, - is_default_route_attempt=is_default_route_attempt + ip_address=ip_address, is_reattempt=True, is_default_route_attempt=is_default_route_attempt ) route = self.router.route_table.find_best_route(ip_address) @@ -569,7 +582,7 @@ class RouterARP(ARP): return self._get_arp_cache_mac_address( ip_address=route.next_hop_ip_address, is_reattempt=True, - is_default_route_attempt=is_default_route_attempt + is_default_route_attempt=is_default_route_attempt, ) else: if self.router.route_table.default_route: @@ -578,16 +591,40 @@ class RouterARP(ARP): return self._get_arp_cache_mac_address( ip_address=self.router.route_table.default_route.next_hop_ip_address, is_reattempt=True, - is_default_route_attempt=True + is_default_route_attempt=True, ) return None def get_arp_cache_mac_address(self, ip_address: IPv4Address) -> Optional[str]: + """ + Public interface to retrieve the MAC address associated with the given IP address from the ARP cache. + + :param ip_address: The IP address for which to find the corresponding MAC address. + :type ip_address: IPv4Address + :return: The MAC address associated with the given IP address, if found; otherwise, None. + :rtype: Optional[str] + """ return self._get_arp_cache_mac_address(ip_address) def _get_arp_cache_network_interface( - self, ip_address: IPV4Address, is_reattempt: bool = False, is_default_route_attempt: bool = False + self, ip_address: IPv4Address, is_reattempt: bool = False, is_default_route_attempt: bool = False ) -> Optional[RouterInterface]: + """ + Attempts to retrieve the router interface associated with the given IP address. + + If the address is not directly associated with a router interface, it may send an ARP request based on + routing information. + + :param ip_address: The IP address for which to find the corresponding router interface. + :type ip_address: IPv4Address + :param is_reattempt: Indicates whether this call is a reattempt after a failed initial attempt. + :type is_reattempt: bool + :param is_default_route_attempt: Indicates whether the attempt is being made for the default route's next-hop + IP address. + :type is_default_route_attempt: bool + :return: The router interface associated with the given IP address, if applicable; otherwise, None. + :rtype: Optional[RouterInterface] + """ arp_entry = self.arp.get(ip_address) if arp_entry: return self.software_manager.node.network_interfaces[arp_entry.network_interface_uuid] @@ -603,7 +640,7 @@ class RouterARP(ARP): return self._get_arp_cache_network_interface( ip_address=route.next_hop_ip_address, is_reattempt=True, - is_default_route_attempt=is_default_route_attempt + is_default_route_attempt=is_default_route_attempt, ) else: if self.router.route_table.default_route: @@ -612,17 +649,32 @@ class RouterARP(ARP): return self._get_arp_cache_network_interface( ip_address=self.router.route_table.default_route.next_hop_ip_address, is_reattempt=True, - is_default_route_attempt=True + is_default_route_attempt=True, ) return None - - def get_arp_cache_network_interface(self, ip_address: IPv4Address) -> Optional[RouterInterface]: + """ + Public interface to retrieve the router interface associated with the given IP address. - return self._get_arp_cache_network_interface(ip_address) + :param ip_address: The IP address for which to find the corresponding router interface. + :type ip_address: IPv4Address + :return: The router interface associated with the given IP address, if found; otherwise, None. + :rtype: Optional[RouterInterface] + """ + return self._get_arp_cache_network_interface(ip_address) def _process_arp_request(self, arp_packet: ARPPacket, from_network_interface: RouterInterface): + """ + Processes an ARP request packet received on a router interface. + + If the target IP address matches the interface's IP address, generates and sends an ARP reply. + + :param arp_packet: The received ARP request packet. + :type arp_packet: ARPPacket + :param from_network_interface: The router interface on which the ARP request was received. + :type from_network_interface: RouterInterface + """ super()._process_arp_request(arp_packet, from_network_interface) # If the target IP matches one of the router's NICs @@ -632,6 +684,14 @@ class RouterARP(ARP): return def _process_arp_reply(self, arp_packet: ARPPacket, from_network_interface: RouterInterface): + """ + Processes an ARP reply packet received on a router interface. Updates the ARP cache with the new information. + + :param arp_packet: The received ARP reply packet. + :type arp_packet: ARPPacket + :param from_network_interface: The router interface on which the ARP reply was received. + :type from_network_interface: RouterInterface + """ if arp_packet.target_ip_address == from_network_interface.ip_address: super()._process_arp_reply(arp_packet, from_network_interface) @@ -650,7 +710,7 @@ class RouterICMP(ICMP): router: Optional[Router] = None - def _process_icmp_echo_request(self, frame: Frame, from_network_interface): + def _process_icmp_echo_request(self, frame: Frame, from_network_interface: RouterInterface): """ Processes an ICMP echo request received by the service. @@ -664,7 +724,8 @@ class RouterICMP(ICMP): if not network_interface: self.sys_log.error( - "Cannot send ICMP echo reply as there is no outbound Network Interface to use. Try configuring the default gateway." + "Cannot send ICMP echo reply as there is no outbound Network Interface to use. Try configuring the " + "default gateway." ) return @@ -682,7 +743,7 @@ class RouterICMP(ICMP): dst_ip_address=frame.ip.src_ip_address, dst_port=self.port, ip_protocol=self.protocol, - icmp_packet=icmp_packet + icmp_packet=icmp_packet, ) def receive(self, payload: Any, session_id: str, **kwargs) -> bool: @@ -815,11 +876,19 @@ class RouterInterface(IPWiredNetworkInterface): class Router(NetworkNode): """ - A class to represent a network router node. + Represents a network router, managing routing and forwarding of IP packets across network interfaces. - :ivar str hostname: The name of the router node. - :ivar int num_ports: The number of ports in the router. - :ivar dict kwargs: Optional keyword arguments for SysLog, ACL, RouteTable, RouterARP, RouterICMP. + A router operates at the network layer and is responsible for receiving, processing, and forwarding data packets + between computer networks. It examines the destination IP address of incoming packets and determines the best way + to route them towards their destination. + + The router integrates various network services and protocols to facilitate IP routing, including ARP (Address + Resolution Protocol) and ICMP (Internet Control Message Protocol) for handling network diagnostics and errors. + + :ivar str hostname: The name of the router, used for identification and logging. + :ivar int num_ports: The number of physical or logical ports on the router. + :ivar dict kwargs: Optional keyword arguments for initializing components like SysLog, ACL (Access Control List), + RouteTable, RouterARP, and RouterICMP services. """ num_ports: int @@ -848,7 +917,13 @@ class Router(NetworkNode): self.set_original_state() def _install_system_software(self): - """Install System Software - software that is usually provided with the OS.""" + """ + Installs essential system software and network services on the router. + + This includes initializing and setting up RouterICMP for handling ICMP packets and RouterARP for address + resolution within the network. These services are crucial for the router's operation, enabling it to manage + network traffic efficiently. + """ self.software_manager.install(RouterICMP) icmp: RouterICMP = self.software_manager.icmp # noqa icmp.router = self @@ -857,11 +932,22 @@ class Router(NetworkNode): arp.router = self def _set_default_acl(self): + """ + Sets default access control rules for the router. + + Initializes the router's ACL (Access Control List) with default rules, permitting essential protocols like ARP + and ICMP, which are necessary for basic network operations and diagnostics. + """ self.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22) self.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23) def set_original_state(self): - """Sets the original state.""" + """ + Sets or resets the router to its original configuration state. + + This method is called to initialize the router's state or to revert it to a known good configuration during + network simulations or after configuration changes. + """ self.acl.set_original_state() self.route_table.set_original_state() super().set_original_state() @@ -869,7 +955,14 @@ class Router(NetworkNode): self._original_state.update(self.model_dump(include=vals_to_include)) def reset_component_for_episode(self, episode: int): - """Reset the original state of the SimComponent.""" + """ + Resets the router's components for a new network simulation episode. + + Clears ARP cache, resets ACL and route table to their original states, and re-enables all network interfaces. + This ensures that the router starts from a clean state for each simulation episode. + + :param episode: The episode number for which the router is being reset. + """ self.software_manager.arp.clear() self.acl.reset_component_for_episode(episode) self.route_table.reset_component_for_episode(episode) @@ -884,7 +977,14 @@ class Router(NetworkNode): rm.add_request("acl", RequestType(func=self.acl._request_manager)) return rm - def ip_is_router_interface(self, ip_address: IPV4Address, enabled_only: bool = False) -> bool: + def ip_is_router_interface(self, ip_address: IPv4Address, enabled_only: bool = False) -> bool: + """ + Checks if a given IP address belongs to any of the router's interfaces. + + :param ip_address: The IP address to check. + :param enabled_only: If True, only considers enabled network interfaces. + :return: True if the IP address is assigned to one of the router's interfaces; False otherwise. + """ for router_interface in self.network_interface.values(): if router_interface.ip_address == ip_address: if enabled_only: @@ -893,7 +993,14 @@ class Router(NetworkNode): return True return False - def ip_is_in_router_interface_subnet(self, ip_address: IPV4Address, enabled_only: bool = False) -> bool: + def ip_is_in_router_interface_subnet(self, ip_address: IPv4Address, enabled_only: bool = False) -> bool: + """ + Determines if a given IP address falls within the subnet of any router interface. + + :param ip_address: The IP address to check. + :param enabled_only: If True, only considers enabled network interfaces. + :return: True if the IP address is within the subnet of any router's interface; False otherwise. + """ for router_interface in self.network_interface.values(): if ip_address in router_interface.ip_network: if enabled_only: @@ -904,10 +1011,10 @@ class Router(NetworkNode): def _get_port_of_nic(self, target_nic: RouterInterface) -> Optional[int]: """ - Retrieve the port number for a given NIC. + Retrieves the port number associated with a given network interface controller (NIC). - :param target_nic: Target network interface. - :return: The port number if NIC is found, otherwise None. + :param target_nic: The NIC whose port number is being queried. + :return: The port number if the NIC is found; otherwise, None. """ for port, network_interface in self.network_interface.items(): if network_interface == target_nic: @@ -926,12 +1033,14 @@ class Router(NetworkNode): def receive_frame(self, frame: Frame, from_network_interface: RouterInterface): """ - Receive a frame from a RouterInterface and processes it based on its protocol. + Processes an incoming frame received on one of the router's interfaces. - :param frame: The incoming frame. - :param from_network_interface: The network interface where the frame is coming from. + Examines the frame's destination and protocol, applies ACL rules, and either forwards or drops the frame based + on routing decisions and ACL permissions. + + :param frame: The incoming frame to be processed. + :param from_network_interface: The router interface on which the frame was received. """ - if self.operating_state != NodeOperatingState.ON: return @@ -965,12 +1074,13 @@ class Router(NetworkNode): self.software_manager.arp.add_arp_cache_entry( ip_address=frame.ip.src_ip_address, mac_address=frame.ethernet.src_mac_addr, - network_interface=from_network_interface + network_interface=from_network_interface, ) send_to_session_manager = False - if ((frame.icmp and self.ip_is_router_interface(dst_ip_address)) - or (dst_port in self.software_manager.get_open_ports())): + if (frame.icmp and self.ip_is_router_interface(dst_ip_address)) or ( + dst_port in self.software_manager.get_open_ports() + ): send_to_session_manager = True if send_to_session_manager: @@ -981,17 +1091,20 @@ class Router(NetworkNode): def process_frame(self, frame: Frame, from_network_interface: RouterInterface) -> None: """ - Process a Frame. + Routes or forwards a frame based on the router's routing table and interface configurations. - :param frame: The frame to be routed. - :param from_network_interface: The source network interface. + This method is called if a frame is not directly addressed to the router or does not match any open service + ports. It determines the next hop for the frame and forwards it accordingly. + + :param frame: The frame to be routed or forwarded. + :param from_network_interface: The network interface from which the frame originated. """ # check if frame is addressed to this Router but has failed to be received by a service of application at the # receive_frame stage if frame.ip: for network_interface in self.network_interfaces.values(): if network_interface.ip_address == frame.ip.dst_ip_address: - self.sys_log.info(f"Dropping frame destined for this router on a port that isn't open.") + self.sys_log.info("Dropping frame destined for this router on a port that isn't open.") return network_interface: RouterInterface = self.software_manager.arp.get_arp_cache_network_interface( @@ -1031,6 +1144,15 @@ class Router(NetworkNode): self.route_frame(frame, from_network_interface) def route_frame(self, frame: Frame, from_network_interface: RouterInterface) -> None: + """ + Determines the best route for a frame and forwards it towards its destination. + + Uses the router's routing table to find the best route for the frame's destination IP address and forwards the + frame through the appropriate interface. + + :param frame: The frame to be routed. + :param from_network_interface: The source network interface. + """ route = self.route_table.find_best_route(frame.ip.dst_ip_address) if route: network_interface = self.software_manager.arp.get_arp_cache_network_interface(route.next_hop_ip_address) @@ -1059,11 +1181,11 @@ class Router(NetworkNode): def configure_port(self, port: int, ip_address: Union[IPv4Address, str], subnet_mask: Union[IPv4Address, str]): """ - Configure the IP settings of a given port. + Configures the IP settings for a specified router port. - :param port: The port to configure. - :param ip_address: The IP address to set. - :param subnet_mask: The subnet mask to set. + :param port: The port number to configure. + :param ip_address: The IP address to assign to the port. + :param subnet_mask: The subnet mask for the port. """ if not isinstance(ip_address, IPv4Address): ip_address = IPv4Address(ip_address) @@ -1072,16 +1194,14 @@ class Router(NetworkNode): network_interface = self.network_interface[port] network_interface.ip_address = ip_address network_interface.subnet_mask = subnet_mask - self.sys_log.info( - f"Configured Network Interface {network_interface}" - ) + self.sys_log.info(f"Configured Network Interface {network_interface}") self.set_original_state() def enable_port(self, port: int): """ - Enable a given port on the router. + Enables a specified port on the router. - :param port: The port to enable. + :param port: The port number to enable. """ network_interface = self.network_interface.get(port) if network_interface: @@ -1089,9 +1209,9 @@ class Router(NetworkNode): def disable_port(self, port: int): """ - Disable a given port on the router. + Disables a specified port on the router. - :param port: The port to disable. + :param port: The port number to disable. """ network_interface = self.network_interface.get(port) if network_interface: diff --git a/src/primaite/simulator/network/hardware/nodes/network/switch.py b/src/primaite/simulator/network/hardware/nodes/network/switch.py index 1878aab7..33e6ee9a 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/switch.py +++ b/src/primaite/simulator/network/hardware/nodes/network/switch.py @@ -1,11 +1,12 @@ from __future__ import annotations + from typing import Dict, Optional from prettytable import MARKDOWN, PrettyTable from primaite import getLogger from primaite.exceptions import NetworkError -from primaite.simulator.network.hardware.base import WiredNetworkInterface, NetworkInterface, Link +from primaite.simulator.network.hardware.base import Link, WiredNetworkInterface from primaite.simulator.network.hardware.nodes.network.network_node import NetworkNode from primaite.simulator.network.transmission.data_link_layer import Frame @@ -27,6 +28,7 @@ class SwitchPort(WiredNetworkInterface): Switch ports typically do not have IP addresses assigned to them as they function at Layer 2, but managed switches can have management IP addresses for remote management and configuration purposes. """ + _connected_node: Optional[Switch] = None "The Switch to which the SwitchPort is connected." @@ -40,7 +42,6 @@ class SwitchPort(WiredNetworkInterface): """ Produce a dictionary describing the current state of this object. - :return: Current state of this object and child objects. :rtype: Dict """ diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index f830ad70..f82dee4a 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -1,11 +1,10 @@ from ipaddress import IPv4Address from primaite.simulator.network.container import Network -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 -from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router from primaite.simulator.network.hardware.nodes.host.server import Server +from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router from primaite.simulator.network.hardware.nodes.network.switch import Switch from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port @@ -56,7 +55,7 @@ def client_server_routed() -> Network: ip_address="192.168.2.2", subnet_mask="255.255.255.0", default_gateway="192.168.2.1", - start_up_duration=0 + start_up_duration=0, ) client_1.power_on() network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1]) @@ -67,7 +66,7 @@ def client_server_routed() -> Network: ip_address="192.168.1.2", subnet_mask="255.255.255.0", default_gateway="192.168.1.1", - start_up_duration=0 + start_up_duration=0, ) server_1.power_on() network.connect(endpoint_b=server_1.network_interface[1], endpoint_a=switch_1.network_interface[1]) @@ -143,7 +142,7 @@ def arcd_uc2_network() -> Network: subnet_mask="255.255.255.0", default_gateway="192.168.10.1", dns_server=IPv4Address("192.168.1.10"), - start_up_duration=0 + start_up_duration=0, ) client_1.power_on() network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1]) @@ -163,7 +162,7 @@ def arcd_uc2_network() -> Network: subnet_mask="255.255.255.0", default_gateway="192.168.10.1", dns_server=IPv4Address("192.168.1.10"), - start_up_duration=0 + start_up_duration=0, ) client_2.power_on() web_browser = client_2.software_manager.software.get("WebBrowser") @@ -176,7 +175,7 @@ def arcd_uc2_network() -> Network: ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.1", - start_up_duration=0 + start_up_duration=0, ) domain_controller.power_on() domain_controller.software_manager.install(DNSServer) @@ -190,7 +189,7 @@ def arcd_uc2_network() -> Network: subnet_mask="255.255.255.0", default_gateway="192.168.1.1", dns_server=IPv4Address("192.168.1.10"), - start_up_duration=0 + start_up_duration=0, ) database_server.power_on() network.connect(endpoint_b=database_server.network_interface[1], endpoint_a=switch_1.network_interface[3]) @@ -264,7 +263,7 @@ def arcd_uc2_network() -> Network: subnet_mask="255.255.255.0", default_gateway="192.168.1.1", dns_server=IPv4Address("192.168.1.10"), - start_up_duration=0 + start_up_duration=0, ) web_server.power_on() web_server.software_manager.install(DatabaseClient) @@ -288,7 +287,7 @@ def arcd_uc2_network() -> Network: subnet_mask="255.255.255.0", default_gateway="192.168.1.1", dns_server=IPv4Address("192.168.1.10"), - start_up_duration=0 + start_up_duration=0, ) backup_server.power_on() backup_server.software_manager.install(FTPServer) @@ -301,7 +300,7 @@ def arcd_uc2_network() -> Network: subnet_mask="255.255.255.0", default_gateway="192.168.1.1", dns_server=IPv4Address("192.168.1.10"), - start_up_duration=0 + start_up_duration=0, ) security_suite.power_on() network.connect(endpoint_b=security_suite.network_interface[1], endpoint_a=switch_1.network_interface[7]) diff --git a/src/primaite/simulator/network/protocols/icmp.py b/src/primaite/simulator/network/protocols/icmp.py index 9f761393..66215db0 100644 --- a/src/primaite/simulator/network/protocols/icmp.py +++ b/src/primaite/simulator/network/protocols/icmp.py @@ -111,4 +111,4 @@ class ICMPPacket(BaseModel): return description msg = f"No Matching ICMP code for type:{self.icmp_type.name}, code:{self.icmp_code}" _LOGGER.error(msg) - raise ValueError(msg) \ No newline at end of file + raise ValueError(msg) diff --git a/src/primaite/simulator/network/transmission/data_link_layer.py b/src/primaite/simulator/network/transmission/data_link_layer.py index 5c25df01..27d40df0 100644 --- a/src/primaite/simulator/network/transmission/data_link_layer.py +++ b/src/primaite/simulator/network/transmission/data_link_layer.py @@ -4,7 +4,6 @@ from typing import Any, Optional from pydantic import BaseModel from primaite import getLogger -from primaite.simulator.network.protocols.arp import ARPPacket from primaite.simulator.network.protocols.icmp import ICMPPacket from primaite.simulator.network.protocols.packet import DataPacket from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol diff --git a/src/primaite/simulator/network/transmission/network_layer.py b/src/primaite/simulator/network/transmission/network_layer.py index 38fc1977..c6328a60 100644 --- a/src/primaite/simulator/network/transmission/network_layer.py +++ b/src/primaite/simulator/network/transmission/network_layer.py @@ -1,10 +1,7 @@ -import secrets from enum import Enum -from ipaddress import IPv4Address, IPv4Network -from typing import Union +from ipaddress import IPv4Address -from pydantic import BaseModel, field_validator, validate_call -from pydantic_core.core_schema import FieldValidationInfo +from pydantic import BaseModel from primaite import getLogger diff --git a/src/primaite/simulator/system/core/packet_capture.py b/src/primaite/simulator/system/core/packet_capture.py index 5d34fd63..3f34cad8 100644 --- a/src/primaite/simulator/system/core/packet_capture.py +++ b/src/primaite/simulator/system/core/packet_capture.py @@ -108,4 +108,3 @@ class PacketCapture: """ msg = frame.model_dump_json() self.outbound_logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL - diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index 4ef10a14..3fa2aa97 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/src/primaite/simulator/system/core/session_manager.py @@ -102,7 +102,7 @@ class SessionManager: @staticmethod def _get_session_key( - frame: Frame, inbound_frame: bool = True + frame: Frame, inbound_frame: bool = True ) -> Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]]: """ Extracts the session key from the given frame. @@ -140,27 +140,76 @@ class SessionManager: dst_port = None return protocol, with_ip_address, src_port, dst_port - def resolve_outbound_network_interface(self, dst_ip_address: IPv4Address) -> Optional['NetworkInterface']: + def resolve_outbound_network_interface(self, dst_ip_address: IPv4Address) -> Optional["NetworkInterface"]: + """ + Resolves the appropriate outbound network interface for a given destination IP address. + + This method determines the most suitable network interface for sending a packet to the specified + destination IP address. It considers only enabled network interfaces and checks if the destination + IP address falls within the subnet of each interface. If no suitable local network interface is found, + the method defaults to using the network interface associated with the default gateway. + + The search process prioritises local network interfaces based on the IP network to which they belong. + If the destination IP address does not match any local subnet, the method assumes that the destination + is outside the local network and hence, routes the packet through the default gateway's network interface. + + :param dst_ip_address: The destination IP address for which the outbound interface is to be resolved. + :type dst_ip_address: IPv4Address + :return: The network interface through which the packet should be sent to reach the destination IP address, + or the default gateway's network interface if the destination is not within any local subnet. + :rtype: Optional["NetworkInterface"] + """ for network_interface in self.node.network_interfaces.values(): if dst_ip_address in network_interface.ip_network and network_interface.enabled: return network_interface return self.software_manager.arp.get_default_gateway_network_interface() def resolve_outbound_transmission_details( - self, - dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None, - src_port: Optional[Port] = None, - dst_port: Optional[Port] = None, - protocol: Optional[IPProtocol] = None, - session_id: Optional[str] = None + self, + dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None, + src_port: Optional[Port] = None, + dst_port: Optional[Port] = None, + protocol: Optional[IPProtocol] = None, + session_id: Optional[str] = None, ) -> Tuple[ - Optional['NetworkInterface'], - Optional[str], IPv4Address, + Optional["NetworkInterface"], + Optional[str], + IPv4Address, Optional[Port], Optional[Port], Optional[IPProtocol], - bool + bool, ]: + """ + Resolves the necessary details for outbound transmission based on the provided parameters. + + This method determines whether the payload should be broadcast or unicast based on the destination IP address + and resolves the outbound network interface and destination MAC address accordingly. + + The method first checks if `session_id` is provided and uses the session details if available. For broadcast + transmissions, it finds a suitable network interface and uses a broadcast MAC address. For unicast + transmissions, it attempts to resolve the destination MAC address using ARP and finds the appropriate + outbound network interface. If the destination IP address is outside the local network and no specific MAC + address is resolved, it uses the default gateway for the transmission. + + :param dst_ip_address: The destination IP address or network. If an IPv4Network is provided, the method + treats the transmission as a broadcast to that network. Optional. + :type dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] + :param src_port: The source port number for the transmission. Optional. + :type src_port: Optional[Port] + :param dst_port: The destination port number for the transmission. Optional. + :type dst_port: Optional[Port] + :param protocol: The IP protocol to be used for the transmission. Optional. + :type protocol: Optional[IPProtocol] + :param session_id: The session ID associated with the transmission. If provided, the session details override + other parameters. Optional. + :type session_id: Optional[str] + :return: A tuple containing the resolved outbound network interface, destination MAC address, destination IP + address, source port, destination port, protocol, and a boolean indicating whether the transmission is a + broadcast. + :rtype: Tuple[Optional["NetworkInterface"], Optional[str], IPv4Address, Optional[Port], Optional[Port], + Optional[IPProtocol], bool] + """ if dst_ip_address and not isinstance(dst_ip_address, (IPv4Address, IPv4Network)): dst_ip_address = IPv4Address(dst_ip_address) is_broadcast = False @@ -207,14 +256,14 @@ class SessionManager: return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast def receive_payload_from_software_manager( - self, - payload: Any, - dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None, - src_port: Optional[Port] = None, - dst_port: Optional[Port] = None, - session_id: Optional[str] = None, - ip_protocol: IPProtocol = IPProtocol.TCP, - icmp_packet: Optional[ICMPPacket] = None + self, + payload: Any, + dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None, + src_port: Optional[Port] = None, + dst_port: Optional[Port] = None, + session_id: Optional[str] = None, + ip_protocol: IPProtocol = IPProtocol.TCP, + icmp_packet: Optional[ICMPPacket] = None, ) -> Union[Any, None]: """ Receive a payload from the SoftwareManager and send it to the appropriate NIC for transmission. @@ -239,15 +288,22 @@ class SessionManager: is_broadcast = payload.request ip_protocol = IPProtocol.UDP else: - vals = self.resolve_outbound_transmission_details( dst_ip_address=dst_ip_address, src_port=src_port, dst_port=dst_port, protocol=ip_protocol, - session_id=session_id + session_id=session_id, ) - outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast = vals + ( + outbound_network_interface, + dst_mac_address, + dst_ip_address, + src_port, + dst_port, + protocol, + is_broadcast, + ) = vals if protocol: ip_protocol = protocol @@ -257,7 +313,7 @@ class SessionManager: if not (src_port or dst_port): raise ValueError( - f"Failed to resolve src or dst port. Have you sent the port from the service or application?" + "Failed to resolve src or dst port. Have you sent the port from the service or application?" ) tcp_header = None @@ -283,7 +339,11 @@ class SessionManager: # Construct the frame for transmission frame = Frame( ethernet=EthernetHeader(src_mac_addr=outbound_network_interface.mac_address, dst_mac_addr=dst_mac_address), - ip=IPPacket(src_ip_address=outbound_network_interface.ip_address, dst_ip_address=dst_ip_address, protocol=ip_protocol), + ip=IPPacket( + src_ip_address=outbound_network_interface.ip_address, + dst_ip_address=dst_ip_address, + protocol=ip_protocol, + ), tcp=tcp_header, udp=udp_header, icmp=icmp_packet, @@ -304,7 +364,7 @@ class SessionManager: # Send the frame through the NIC return outbound_network_interface.send_frame(frame) - def receive_frame(self, frame: Frame, from_network_interface: 'NetworkInterface'): + def receive_frame(self, frame: Frame, from_network_interface: "NetworkInterface"): """ Receive a Frame. @@ -334,7 +394,7 @@ class SessionManager: protocol=frame.ip.protocol, session_id=session.uuid, from_network_interface=from_network_interface, - frame=frame + frame=frame, ) def show(self, markdown: bool = False): diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py index 53725c18..e6fe7b23 100644 --- a/src/primaite/simulator/system/core/software_manager.py +++ b/src/primaite/simulator/system/core/software_manager.py @@ -25,7 +25,14 @@ IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware) class SoftwareManager: - """A class that manages all running Services and Applications on a Node and facilitates their communication.""" + """ + Manages all running services and applications on a network node and facilitates their communication. + + This class is responsible for installing, uninstalling, and managing the operational state of various network + services and applications. It acts as a bridge between the node's session manager and its software components, + ensuring that incoming and outgoing network payloads are correctly routed to and from the appropriate services + or applications. + """ def __init__( self, @@ -50,11 +57,13 @@ class SoftwareManager: self.dns_server: Optional[IPv4Address] = dns_server @property - def arp(self) -> 'ARP': + def arp(self) -> "ARP": + """Provides access to the ARP service instance, if installed.""" return self.software.get("ARP") # noqa @property - def icmp(self) -> 'ICMP': + def icmp(self) -> "ICMP": + """Provides access to the ICMP service instance, if installed.""" return self.software.get("ICMP") # noqa def get_open_ports(self) -> List[Port]: @@ -167,7 +176,13 @@ class SoftwareManager: ) def receive_payload_from_session_manager( - self, payload: Any, port: Port, protocol: IPProtocol, session_id: str, from_network_interface: "NIC", frame: Frame + self, + payload: Any, + port: Port, + protocol: IPProtocol, + session_id: str, + from_network_interface: "NIC", + frame: Frame, ): """ Receive a payload from the SessionManager and forward it to the corresponding service or application. @@ -177,7 +192,9 @@ class SoftwareManager: """ receiver: Optional[Union[Service, Application]] = self.port_protocol_mapping.get((port, protocol), None) if receiver: - receiver.receive(payload=payload, session_id=session_id, from_network_interface=from_network_interface, frame=frame) + receiver.receive( + payload=payload, session_id=session_id, from_network_interface=from_network_interface, frame=frame + ) else: self.sys_log.error(f"No service or application found for port {port} and protocol {protocol}") pass @@ -202,7 +219,7 @@ class SoftwareManager: software.operating_state.name, software.health_state_actual.name, software.port.value if software.port != Port.NONE else None, - software.protocol.value + software.protocol.value, ] ) print(table) diff --git a/src/primaite/simulator/system/services/arp/arp.py b/src/primaite/simulator/system/services/arp/arp.py index 6a04e845..ca5b7619 100644 --- a/src/primaite/simulator/system/services/arp/arp.py +++ b/src/primaite/simulator/system/services/arp/arp.py @@ -20,6 +20,7 @@ class ARP(Service): Manages ARP for resolving network layer addresses into link layer addresses. It maintains an ARP cache, sends ARP requests and replies, and processes incoming ARP packets. """ + arp: Dict[IPV4Address, ARPEntry] = {} def __init__(self, **kwargs): @@ -29,6 +30,14 @@ class ARP(Service): super().__init__(**kwargs) def describe_state(self) -> Dict: + """ + Produce a dictionary describing the current state of this object. + + :return: Current state of this object and child objects. + """ + state = super().describe_state() + state.update({str(ip): arp_entry.mac_address for ip, arp_entry in self.arp.items()}) + return super().describe_state() def show(self, markdown: bool = False): @@ -57,11 +66,7 @@ class ARP(Service): self.arp.clear() def add_arp_cache_entry( - self, - ip_address: IPV4Address, - mac_address: str, - network_interface: NetworkInterface, - override: bool = False + self, ip_address: IPV4Address, mac_address: str, network_interface: NetworkInterface, override: bool = False ): """ Add an ARP entry to the cache. @@ -139,7 +144,8 @@ class ARP(Service): ) else: self.sys_log.error( - "Cannot send ARP request as there is no outbound Network Interface to use. Try configuring the default gateway." + "Cannot send ARP request as there is no outbound Network Interface to use. Try configuring the default " + "gateway." ) def send_arp_reply(self, arp_reply: ARPPacket): @@ -147,12 +153,10 @@ class ARP(Service): Sends an ARP reply in response to an ARP request. :param arp_reply: The ARP packet containing the reply. - :param from_network_interface: The NIC from which the ARP reply is sent. """ - outbound_network_interface = self.software_manager.session_manager.resolve_outbound_network_interface( arp_reply.target_ip_address - ) + ) if outbound_network_interface: self.sys_log.info( f"Sending ARP reply from {arp_reply.sender_mac_addr}/{arp_reply.sender_ip_address} " @@ -162,14 +166,14 @@ class ARP(Service): payload=arp_reply, dst_ip_address=arp_reply.target_ip_address, dst_port=self.port, - ip_protocol=self.protocol + ip_protocol=self.protocol, ) else: self.sys_log.error( - "Cannot send ARP reply as there is no outbound Network Interface to use. Try configuring the default gateway." + "Cannot send ARP reply as there is no outbound Network Interface to use. Try configuring the default " + "gateway." ) - @abstractmethod def _process_arp_request(self, arp_packet: ARPPacket, from_network_interface: NetworkInterface): """ @@ -197,7 +201,7 @@ class ARP(Service): self.add_arp_cache_entry( ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, - network_interface=from_network_interface + network_interface=from_network_interface, ) def receive(self, payload: Any, session_id: str, **kwargs) -> bool: diff --git a/src/primaite/simulator/system/services/icmp/icmp.py b/src/primaite/simulator/system/services/icmp/icmp.py index 3ff7b21c..103d1c60 100644 --- a/src/primaite/simulator/system/services/icmp/icmp.py +++ b/src/primaite/simulator/system/services/icmp/icmp.py @@ -1,8 +1,9 @@ import secrets from ipaddress import IPv4Address -from typing import Dict, Any, Union, Optional, Tuple +from typing import Any, Dict, Optional, Tuple, Union from primaite import getLogger +from primaite.simulator.network.hardware.base import NetworkInterface from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType from primaite.simulator.network.transmission.data_link_layer import Frame from primaite.simulator.network.transmission.network_layer import IPProtocol @@ -19,6 +20,7 @@ class ICMP(Service): Enables the sending and receiving of ICMP messages such as echo requests and replies. This is typically used for network diagnostics, notably the ping command. """ + request_replies: Dict = {} def __init__(self, **kwargs): @@ -28,7 +30,12 @@ class ICMP(Service): super().__init__(**kwargs) def describe_state(self) -> Dict: - pass + """ + Produce a dictionary describing the current state of this object. + + :return: Current state of this object and child objects. + """ + return super().describe_state() def clear(self): """ @@ -56,9 +63,7 @@ class ICMP(Service): self.sys_log.info(f"Pinging {target_ip_address}:", to_terminal=True) sequence, identifier = 0, None while sequence < pings: - sequence, identifier = self._send_icmp_echo_request( - target_ip_address, sequence, identifier, pings - ) + sequence, identifier = self._send_icmp_echo_request(target_ip_address, sequence, identifier, pings) request_replies = self.software_manager.icmp.request_replies.get(identifier) passed = request_replies == pings if request_replies: @@ -76,7 +81,7 @@ class ICMP(Service): return passed def _send_icmp_echo_request( - self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None, pings: int = 4 + self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None, pings: int = 4 ) -> Tuple[int, Union[int, None]]: """ Sends an ICMP echo request to a specified target IP address. @@ -91,7 +96,8 @@ class ICMP(Service): if not network_interface: self.sys_log.error( - "Cannot send ICMP echo request as there is no outbound Network Interface to use. Try configuring the default gateway." + "Cannot send ICMP echo request as there is no outbound Network Interface to use. Try configuring the " + "default gateway." ) return pings, None @@ -105,11 +111,11 @@ class ICMP(Service): dst_ip_address=target_ip_address, dst_port=self.port, ip_protocol=self.protocol, - icmp_packet=icmp_packet + icmp_packet=icmp_packet, ) return sequence, icmp_packet.identifier - def _process_icmp_echo_request(self, frame: Frame, from_network_interface): + def _process_icmp_echo_request(self, frame: Frame, from_network_interface: NetworkInterface): """ Processes an ICMP echo request received by the service. @@ -121,11 +127,12 @@ class ICMP(Service): network_interface = self.software_manager.session_manager.resolve_outbound_network_interface( frame.ip.src_ip_address - ) + ) if not network_interface: self.sys_log.error( - "Cannot send ICMP echo reply as there is no outbound Network Interface to use. Try configuring the default gateway." + "Cannot send ICMP echo reply as there is no outbound Network Interface to use. Try configuring the " + "default gateway." ) return @@ -143,7 +150,7 @@ class ICMP(Service): dst_ip_address=frame.ip.src_ip_address, dst_port=self.port, ip_protocol=self.protocol, - icmp_packet=icmp_packet + icmp_packet=icmp_packet, ) def _process_icmp_echo_reply(self, frame: Frame): @@ -159,7 +166,7 @@ class ICMP(Service): f"bytes={len(frame.payload)}, " f"time={time_str}, " f"TTL={frame.ip.ttl}", - to_terminal=True + to_terminal=True, ) if not self.request_replies.get(frame.icmp.identifier): self.request_replies[frame.icmp.identifier] = 0 diff --git a/src/primaite/simulator/system/services/ntp/ntp_server.py b/src/primaite/simulator/system/services/ntp/ntp_server.py index 8e362880..3987fa2c 100644 --- a/src/primaite/simulator/system/services/ntp/ntp_server.py +++ b/src/primaite/simulator/system/services/ntp/ntp_server.py @@ -70,10 +70,6 @@ class NTPServer(Service): payload = payload.generate_reply(time) # send reply self.software_manager.session_manager.receive_payload_from_software_manager( - payload=payload, - src_port=self.port, - dst_port=self.port, - ip_protocol=self.protocol, - session_id=session_id + payload=payload, src_port=self.port, dst_port=self.port, ip_protocol=self.protocol, session_id=session_id ) return True diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 91629f9a..ce39930b 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -380,7 +380,7 @@ class IOSoftware(Software): dest_ip_address=dest_ip_address, dest_port=dest_port, ip_protocol=ip_protocol, - session_id=session_id + session_id=session_id, ) @abstractmethod diff --git a/src/primaite/utils/validators.py b/src/primaite/utils/validators.py index 13cff653..fb7abb29 100644 --- a/src/primaite/utils/validators.py +++ b/src/primaite/utils/validators.py @@ -1,9 +1,7 @@ from ipaddress import IPv4Address from typing import Any, Final -from pydantic import ( - BeforeValidator, -) +from pydantic import BeforeValidator from typing_extensions import Annotated @@ -30,7 +28,7 @@ def ipv4_validator(v: Any) -> IPv4Address: # with the IPv4Address type, ensuring that any usage of IPV4Address undergoes validation before assignment. IPV4Address: Final[Annotated] = Annotated[IPv4Address, BeforeValidator(ipv4_validator)] """ -IPv4Address with with pre-validation and auto-conversion from str using ipv4_validator. +IPv4Address with with IPv4Address with with pre-validation and auto-conversion from str using ipv4_validator.. This type is essentially an IPv4Address from the standard library's ipaddress module, but with added validation logic. If you use this custom type, the ipv4_validator function diff --git a/tests/conftest.py b/tests/conftest.py index b5226a34..8639cec3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,10 +5,10 @@ from typing import Any, Dict, Tuple, Union import pytest import yaml -from primaite import PRIMAITE_PATHS -from primaite import getLogger +from primaite import getLogger, PRIMAITE_PATHS from primaite.session.session import PrimaiteSession from primaite.simulator.file_system.file_system import FileSystem + # from primaite.environment.primaite_env import Primaite # from primaite.primaite_session import PrimaiteSession from primaite.simulator.network.container import Network @@ -212,31 +212,20 @@ def example_network() -> Network: network = Network() # Router 1 - router_1 = Router( - hostname="router_1", - start_up_duration=0 - ) + router_1 = Router(hostname="router_1", start_up_duration=0) router_1.power_on() router_1.configure_port(port=1, ip_address="192.168.1.1", subnet_mask="255.255.255.0") router_1.configure_port(port=2, ip_address="192.168.10.1", subnet_mask="255.255.255.0") # Switch 1 - switch_1 = Switch( - hostname="switch_1", - num_ports=8, - start_up_duration=0 - ) + switch_1 = Switch(hostname="switch_1", num_ports=8, start_up_duration=0) switch_1.power_on() network.connect(endpoint_a=router_1.network_interface[1], endpoint_b=switch_1.network_interface[8]) router_1.enable_port(1) # Switch 2 - switch_2 = Switch( - hostname="switch_2", - num_ports=8, - start_up_duration=0 - ) + switch_2 = Switch(hostname="switch_2", num_ports=8, start_up_duration=0) switch_2.power_on() network.connect(endpoint_a=router_1.network_interface[2], endpoint_b=switch_2.network_interface[8]) router_1.enable_port(2) @@ -247,7 +236,7 @@ def example_network() -> Network: ip_address="192.168.10.21", subnet_mask="255.255.255.0", default_gateway="192.168.10.1", - start_up_duration=0 + start_up_duration=0, ) client_1.power_on() network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1]) @@ -258,7 +247,7 @@ def example_network() -> Network: ip_address="192.168.10.22", subnet_mask="255.255.255.0", default_gateway="192.168.10.1", - start_up_duration=0 + start_up_duration=0, ) client_2.power_on() network.connect(endpoint_b=client_2.network_interface[1], endpoint_a=switch_2.network_interface[2]) @@ -269,7 +258,7 @@ def example_network() -> Network: ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.1", - start_up_duration=0 + start_up_duration=0, ) server_1.power_on() network.connect(endpoint_b=server_1.network_interface[1], endpoint_a=switch_1.network_interface[1]) @@ -280,7 +269,7 @@ def example_network() -> Network: ip_address="192.168.1.14", subnet_mask="255.255.255.0", default_gateway="192.168.1.1", - start_up_duration=0 + start_up_duration=0, ) server_2.power_on() network.connect(endpoint_b=server_2.network_interface[1], endpoint_a=switch_1.network_interface[2]) @@ -290,5 +279,4 @@ def example_network() -> Network: assert all(link.is_up for link in network.links.values()) - return network diff --git a/tests/integration_tests/network/test_broadcast.py b/tests/integration_tests/network/test_broadcast.py index d6c52acc..6b6deb93 100644 --- a/tests/integration_tests/network/test_broadcast.py +++ b/tests/integration_tests/network/test_broadcast.py @@ -37,12 +37,7 @@ class BroadcastService(Service): def broadcast(self, ip_network: IPv4Network): # Send a broadcast payload to an entire IP network - super().send( - payload="broadcast", - dest_ip_address=ip_network, - dest_port=Port.HTTP, - ip_protocol=self.protocol - ) + super().send(payload="broadcast", dest_ip_address=ip_network, dest_port=Port.HTTP, ip_protocol=self.protocol) class BroadcastClient(Application): diff --git a/tests/integration_tests/network/test_frame_transmission.py b/tests/integration_tests/network/test_frame_transmission.py index 5ba4fe13..eb30a245 100644 --- a/tests/integration_tests/network/test_frame_transmission.py +++ b/tests/integration_tests/network/test_frame_transmission.py @@ -5,7 +5,6 @@ from primaite.simulator.network.hardware.nodes.host.server import Server from primaite.simulator.network.hardware.nodes.network.switch import Switch - def test_node_to_node_ping(): """Tests two Computers are able to ping each other.""" network = Network() diff --git a/tests/integration_tests/network/test_network_creation.py b/tests/integration_tests/network/test_network_creation.py index 6a39e101..5cf36bce 100644 --- a/tests/integration_tests/network/test_network_creation.py +++ b/tests/integration_tests/network/test_network_creation.py @@ -97,7 +97,6 @@ def test_disconnecting_nodes(): net.connect(n1.network_interface[1], n2.network_interface[1]) assert len(net.links) == 1 - link = list(net.links.values())[0] net.remove_link(link) assert link not in net diff --git a/tests/integration_tests/network/test_routing.py b/tests/integration_tests/network/test_routing.py index 02524eab..4ada807f 100644 --- a/tests/integration_tests/network/test_routing.py +++ b/tests/integration_tests/network/test_routing.py @@ -19,7 +19,7 @@ def pc_a_pc_b_router_1() -> Tuple[Computer, Computer, Router]: ip_address="192.168.0.10", subnet_mask="255.255.255.0", default_gateway="192.168.0.1", - start_up_duration=0 + start_up_duration=0, ) pc_a.power_on() @@ -28,7 +28,7 @@ def pc_a_pc_b_router_1() -> Tuple[Computer, Computer, Router]: ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.1", - start_up_duration=0 + start_up_duration=0, ) pc_b.power_on() diff --git a/tests/integration_tests/system/red_applications/test_dos_bot_and_server.py b/tests/integration_tests/system/red_applications/test_dos_bot_and_server.py index ecf2c5ae..7ab7d104 100644 --- a/tests/integration_tests/system/red_applications/test_dos_bot_and_server.py +++ b/tests/integration_tests/system/red_applications/test_dos_bot_and_server.py @@ -5,8 +5,8 @@ import pytest 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, Router from primaite.simulator.network.hardware.nodes.host.server import Server +from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.applications.application import ApplicationOperatingState from primaite.simulator.system.applications.database_client import DatabaseClient diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py index c259501e..e015f9ee 100644 --- a/tests/integration_tests/system/test_database_on_node.py +++ b/tests/integration_tests/system/test_database_on_node.py @@ -16,21 +16,11 @@ from primaite.simulator.system.services.service import ServiceOperatingState @pytest.fixture(scope="function") def peer_to_peer() -> Tuple[Computer, Computer]: network = Network() - node_a = Computer( - hostname="node_a", - ip_address="192.168.0.10", - subnet_mask="255.255.255.0", - start_up_duration=0 - ) + node_a = Computer(hostname="node_a", ip_address="192.168.0.10", subnet_mask="255.255.255.0", start_up_duration=0) node_a.power_on() node_a.software_manager.get_open_ports() - node_b = Computer( - hostname="node_b", - ip_address="192.168.0.11", - subnet_mask="255.255.255.0", - start_up_duration=0 - ) + node_b = Computer(hostname="node_b", ip_address="192.168.0.11", subnet_mask="255.255.255.0", start_up_duration=0) node_b.power_on() network.connect(node_a.network_interface[1], node_b.network_interface[1]) diff --git a/tests/integration_tests/system/test_dns_client_server.py b/tests/integration_tests/system/test_dns_client_server.py index 18988043..78d2035c 100644 --- a/tests/integration_tests/system/test_dns_client_server.py +++ b/tests/integration_tests/system/test_dns_client_server.py @@ -28,7 +28,8 @@ def dns_client_and_dns_server(client_server) -> Tuple[DNSClient, Computer, DNSSe dns_server.start() # register arcd.com as a domain dns_server.dns_register( - domain_name="arcd.com", domain_ip_address=IPv4Address(server.network_interfaces.get(next(iter(server.network_interfaces))).ip_address) + domain_name="arcd.com", + domain_ip_address=IPv4Address(server.network_interfaces.get(next(iter(server.network_interfaces))).ip_address), ) return dns_client, computer, dns_server, server diff --git a/tests/integration_tests/system/test_web_client_server.py b/tests/integration_tests/system/test_web_client_server.py index c809f954..5e3ff544 100644 --- a/tests/integration_tests/system/test_web_client_server.py +++ b/tests/integration_tests/system/test_web_client_server.py @@ -37,7 +37,10 @@ def web_client_and_web_server(client_server) -> Tuple[WebBrowser, Computer, WebS 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) + dns_server.dns_register( + domain_name="arcd.com", + domain_ip_address=server.network_interfaces[next(iter(server.network_interfaces))].ip_address, + ) return web_browser, computer, web_server_service, server diff --git a/tests/integration_tests/system/test_web_client_server_and_database.py b/tests/integration_tests/system/test_web_client_server_and_database.py index efb29f41..70846ee8 100644 --- a/tests/integration_tests/system/test_web_client_server_and_database.py +++ b/tests/integration_tests/system/test_web_client_server_and_database.py @@ -5,8 +5,8 @@ import pytest from primaite.simulator.network.hardware.base import Link from primaite.simulator.network.hardware.nodes.host.computer import Computer -from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router from primaite.simulator.network.hardware.nodes.host.server import Server +from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.applications.database_client import DatabaseClient from primaite.simulator.system.applications.web_browser import WebBrowser @@ -85,7 +85,8 @@ def web_client_web_server_database(example_network) -> Tuple[Computer, Server, S dns_server: DNSServer = web_server.software_manager.software.get("DNSServer") # register arcd.com to DNS dns_server.dns_register( - domain_name="arcd.com", domain_ip_address=web_server.network_interfaces[next(iter(web_server.network_interfaces))].ip_address + domain_name="arcd.com", + domain_ip_address=web_server.network_interfaces[next(iter(web_server.network_interfaces))].ip_address, ) # Install DatabaseClient service on web server diff --git a/tests/unit_tests/_primaite/_simulator/_system/_applications/test_database_client.py b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_database_client.py index c7d807e9..5f10ec96 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_applications/test_database_client.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_database_client.py @@ -16,21 +16,13 @@ from primaite.simulator.system.services.database.database_service import Databas def database_client_on_computer() -> Tuple[DatabaseClient, Computer]: network = Network() - db_server = Server( - hostname="db_server", - ip_address="192.168.0.1", - subnet_mask="255.255.255.0", - start_up_duration=0 - ) + db_server = Server(hostname="db_server", ip_address="192.168.0.1", subnet_mask="255.255.255.0", start_up_duration=0) db_server.power_on() db_server.software_manager.install(DatabaseService) db_server.software_manager.software["DatabaseService"].start() db_client = Computer( - hostname="db_client", - ip_address="192.168.0.2", - subnet_mask="255.255.255.0", - start_up_duration=0 + hostname="db_client", ip_address="192.168.0.2", subnet_mask="255.255.255.0", start_up_duration=0 ) db_client.power_on() db_client.software_manager.install(DatabaseClient) @@ -97,6 +89,7 @@ def test_disconnect(database_client_on_computer): assert not database_client.connected + def test_query_when_client_is_closed(database_client_on_computer): """Database client should return False when it is not running.""" database_client, computer = database_client_on_computer diff --git a/tests/unit_tests/_primaite/_simulator/_system/_applications/test_web_browser.py b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_web_browser.py index 05d4a985..d210ff40 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_applications/test_web_browser.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_web_browser.py @@ -16,7 +16,7 @@ def web_browser() -> WebBrowser: ip_address="192.168.1.11", subnet_mask="255.255.255.0", default_gateway="192.168.1.1", - start_up_duration=0 + start_up_duration=0, ) computer.power_on() # Web Browser should be pre-installed in computer diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_server.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_server.py index 937636a6..9a513396 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_server.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_server.py @@ -9,8 +9,8 @@ from primaite.simulator.network.hardware.nodes.host.computer import Computer from primaite.simulator.network.hardware.nodes.host.server import Server from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port -from primaite.simulator.system.services.dns.dns_server import DNSServer from primaite.simulator.system.services.dns.dns_client import DNSClient +from primaite.simulator.system.services.dns.dns_server import DNSServer @pytest.fixture(scope="function") @@ -65,6 +65,4 @@ def test_dns_server_receive(dns_server): assert dns_client.check_domain_exists("real-domain.com") is False - - dns_server_service.show()