oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2139 - Fixed unicast and broadcast functionality properly

Browse Source
This commit is contained in:
Chris McCarthy
2024-01-05 22:11:37 +00:00
parent ddf7fbf88b
commit d2d628b676
7 changed files with 341 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
CHANGELOG.md
View File

@@ -38,10 +38,18 @@ SessionManager.
- HTTP Services: `WebBrowser` to simulate a web client and `WebServer`
- Fixed an issue where the services were still able to run even though the node the service is installed on is turned off
- NTP Services: `NTPClient` and `NTPServer`
- **RouterNIC Class**: Introduced a new class `RouterNIC`, extending the standard `NIC` functionality. This class is specifically designed for router operations, optimizing the processing and routing of network traffic.
- **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.
### Changed
- Integrated the RouteTable into the Routers frame processing.
- Frames are now dropped when their TTL reaches 0
- **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.
### Removed
- Removed legacy simulation modules: `acl`, `common`, `environment`, `links`, `nodes`, `pol`

27
src/primaite/simulator/network/hardware/base.py
View File

@@ -274,11 +274,20 @@ class NIC(SimComponent):
def receive_frame(self, frame: Frame) -> bool:
"""
Receive a network frame from the connected link if the NIC is enabled.
Receive a network frame from the connected link, processing it if the NIC is enabled.
The Frame is passed to the Node.
This method decrements the Time To Live (TTL) of the frame, captures it using PCAP (Packet Capture), and checks
if the frame is either a broadcast or destined for this NIC. If the frame is acceptable, it is passed to the
connected node. The method also handles the discarding of frames with TTL expired and logs this event.
:param frame: The network frame being received.
The frame's reception is based on various conditions:
- If the NIC is disabled, the frame is not processed.
- If the TTL of the frame reaches zero after decrement, it is discarded and logged.
- If the frame is a broadcast or its destination MAC/IP address matches this NIC's, it is accepted.
- All other frames are dropped and logged or printed to the console.
:param frame: The network frame being received. This should be an instance of the Frame class.
:return: Returns True if the frame is processed and passed to the node, False otherwise.
"""
if self.enabled:
frame.decrement_ttl()
@@ -288,7 +297,17 @@ class NIC(SimComponent):
frame.set_received_timestamp()
self.pcap.capture(frame)
# If this destination or is broadcast
if frame.ethernet.dst_mac_addr == self.mac_address or frame.ethernet.dst_mac_addr == "ff:ff:ff:ff:ff:ff":
accept_frame = False
# Check if it's a broadcast:
if frame.ethernet.dst_mac_addr == "ff:ff:ff:ff:ff:ff":
if frame.ip.dst_ip_address in {self.ip_address, self.ip_network.broadcast_address}:
accept_frame = True
else:
if frame.ethernet.dst_mac_addr == self.mac_address:
accept_frame = True
if accept_frame:
self._connected_node.receive_frame(frame=frame, from_nic=self)
return True
return False

45
src/primaite/simulator/network/hardware/nodes/router.py
View File

@@ -690,6 +690,47 @@ class RouterICMP(ICMP):
self.router.process_frame(frame, from_nic)
class RouterNIC(NIC):
"""
A Router-specific Network Interface Card (NIC) that extends the standard NIC functionality.
This class overrides the standard Node NIC's Layer 3 (L3) broadcast/unicast checks. It is designed
to handle network frames in a manner specific to routers, allowing them to efficiently process
and route network traffic.
"""
def receive_frame(self, frame: Frame) -> bool:
"""
Receive and process a network frame from the connected link, provided the NIC is enabled.
This method is tailored for router behavior. It decrements the frame's Time To Live (TTL), checks for TTL
expiration, and captures the frame using PCAP (Packet Capture). The frame is accepted if it is destined for
this NIC's MAC address or is a broadcast frame.
Key Differences from Standard NIC:
- Does not perform Layer 3 (IP-based) broadcast checks.
- Only checks for Layer 2 (Ethernet) destination MAC address and broadcast frames.
:param frame: The network frame being received. This should be an instance of the Frame class.
:return: Returns True if the frame is processed and passed to the connected node, False otherwise.
"""
if self.enabled:
frame.decrement_ttl()
if frame.ip and frame.ip.ttl < 1:
self._connected_node.sys_log.info("Frame discarded as TTL limit reached")
return False
frame.set_received_timestamp()
self.pcap.capture(frame)
# If this destination or is broadcast
if frame.ethernet.dst_mac_addr == self.mac_address or frame.ethernet.dst_mac_addr == "ff:ff:ff:ff:ff:ff":
self._connected_node.receive_frame(frame=frame, from_nic=self)
return True
return False
def __str__(self) -> str:
return f"{self.mac_address}/{self.ip_address}"
class Router(Node):
"""
A class to represent a network router node.
@@ -700,7 +741,7 @@ class Router(Node):
"""
num_ports: int
ethernet_ports: Dict[int, NIC] = {}
ethernet_ports: Dict[int, RouterNIC] = {}
acl: AccessControlList
route_table: RouteTable
arp: RouterARPCache
@@ -719,7 +760,7 @@ class Router(Node):
kwargs["icmp"] = RouterICMP(sys_log=kwargs.get("sys_log"), arp_cache=kwargs.get("arp"), router=self)
super().__init__(hostname=hostname, num_ports=num_ports, **kwargs)
for i in range(1, self.num_ports + 1):
nic = NIC(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0")
nic = RouterNIC(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0")
self.connect_nic(nic)
self.ethernet_ports[i] = nic

83
src/primaite/simulator/system/core/session_manager.py
View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from ipaddress import IPv4Address
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
@@ -141,41 +141,76 @@ class SessionManager:
def receive_payload_from_software_manager(
self,
payload: Any,
dst_ip_address: Optional[IPv4Address] = None,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dst_port: Optional[Port] = None,
session_id: Optional[str] = None,
is_reattempt: bool = False,
) -> Union[Any, None]:
"""
Receive a payload from the SoftwareManager.
Receive a payload from the SoftwareManager and send it to the appropriate NIC for transmission.
If no session_id, a Session is established. Once established, the payload is sent to ``send_payload_to_nic``.
This method supports both unicast and Layer 3 broadcast transmissions. If `dst_ip_address` is an
IPv4Network, a broadcast is initiated. For unicast, the destination MAC address is resolved via ARP.
A new session is established if `session_id` is not provided, and an existing session is used otherwise.
:param payload: The payload to be sent.
:param session_id: The Session ID the payload is to originate from. Optional. If None, one will be created.
:param dst_ip_address: The destination IP address or network for broadcast. Optional.
:param dst_port: The destination port for the TCP packet. Optional.
:param session_id: The Session ID from which the payload originates. Optional.
:param is_reattempt: Flag to indicate if this is a reattempt after an ARP request. Default is False.
:return: The outcome of sending the frame, or None if sending was unsuccessful.
"""
is_broadcast = False
outbound_nic = None
dst_mac_address = None
# Use session details if session_id is provided
if session_id:
session = self.sessions_by_uuid[session_id]
dst_ip_address = self.sessions_by_uuid[session_id].with_ip_address
dst_port = self.sessions_by_uuid[session_id].dst_port
dst_ip_address = session.with_ip_address
dst_port = session.dst_port
dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dst_ip_address)
# Determine if the payload is for broadcast or unicast
if dst_mac_address:
outbound_nic = self.arp_cache.get_arp_cache_nic(dst_ip_address)
# Handle broadcast transmission
if isinstance(dst_ip_address, IPv4Network):
is_broadcast = True
dst_ip_address = dst_ip_address.broadcast_address
if dst_ip_address:
# Find a suitable NIC for the broadcast
for nic in self.arp_cache.nics.values():
if dst_ip_address in nic.ip_network and nic.enabled:
dst_mac_address = "ff:ff:ff:ff:ff:ff"
outbound_nic = nic
else:
if not is_reattempt:
self.arp_cache.send_arp_request(dst_ip_address)
return self.receive_payload_from_software_manager(
payload=payload,
dst_ip_address=dst_ip_address,
dst_port=dst_port,
session_id=session_id,
is_reattempt=True,
)
else:
return
# Resolve MAC address for unicast transmission
dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dst_ip_address)
# Resolve outbound NIC for unicast transmission
if dst_mac_address:
outbound_nic = self.arp_cache.get_arp_cache_nic(dst_ip_address)
# If MAC address not found, initiate ARP request
else:
if not is_reattempt:
self.arp_cache.send_arp_request(dst_ip_address)
# Reattempt payload transmission after ARP request
return self.receive_payload_from_software_manager(
payload=payload,
dst_ip_address=dst_ip_address,
dst_port=dst_port,
session_id=session_id,
is_reattempt=True,
)
else:
# Return None if reattempt fails
return
# Check if outbound NIC and destination MAC address are resolved
if not outbound_nic or not dst_mac_address:
return False
# Construct the frame for transmission
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=outbound_nic.mac_address, dst_mac_addr=dst_mac_address),
ip=IPPacket(
@@ -189,15 +224,17 @@ class SessionManager:
payload=payload,
)
if not session_id:
# Manage session for unicast transmission
if not (is_broadcast and session_id):
session_key = self._get_session_key(frame, inbound_frame=False)
session = self.sessions_by_key.get(session_key)
if not session:
# Create new session
# Create a new session if it doesn't exist
session = Session.from_session_key(session_key)
self.sessions_by_key[session_key] = session
self.sessions_by_uuid[session.uuid] = session
# Send the frame through the NIC
return outbound_nic.send_frame(frame)
def receive_frame(self, frame: Frame):

22
src/primaite/simulator/system/core/software_manager.py
View File

@@ -1,4 +1,4 @@
from ipaddress import IPv4Address
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
@@ -130,20 +130,28 @@ class SoftwareManager:
def send_payload_to_session_manager(
self,
payload: Any,
dest_ip_address: Optional[IPv4Address] = None,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[Port] = None,
session_id: Optional[str] = None,
) -> bool:
"""
Send a payload to the SessionManager.
Sends a payload to the SessionManager for network transmission.
This method is responsible for initiating the process of sending network payloads. It supports both
unicast and Layer 3 broadcast transmissions. For broadcasts, the destination IP should be specified
as an IPv4Network.
:param payload: The payload to be sent.
:param dest_ip_address: The ip address of the payload destination.
:param dest_port: The port of the payload destination.
:param session_id: The Session ID the payload is to originate from. Optional.
:param dest_ip_address: The IP address or network (for broadcasts) of the payload destination.
:param dest_port: The destination port for the payload. Optional.
:param session_id: The Session ID from which the payload originates. Optional.
:return: True if the payload was successfully sent, False otherwise.
"""
return self.session_manager.receive_payload_from_software_manager(
payload=payload, dst_ip_address=dest_ip_address, dst_port=dest_port, session_id=session_id
payload=payload,
dst_ip_address=dest_ip_address,
dst_port=dest_port,
session_id=session_id,
)
def receive_payload_from_session_manager(self, payload: Any, port: Port, protocol: IPProtocol, session_id: str):

21
src/primaite/simulator/system/software.py
View File

@@ -2,8 +2,8 @@ import copy
from abc import abstractmethod
from datetime import datetime
from enum import Enum
from ipaddress import IPv4Address
from typing import Any, Dict, Optional
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, Union
from primaite.simulator.core import _LOGGER, RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file_system import FileSystem, Folder
@@ -317,19 +317,22 @@ class IOSoftware(Software):
self,
payload: Any,
session_id: Optional[str] = None,
dest_ip_address: Optional[IPv4Address] = None,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[Port] = None,
**kwargs,
) -> bool:
"""
Sends a payload to the SessionManager.
Sends a payload to the SessionManager for network transmission.
This method is responsible for initiating the process of sending network payloads. It supports both
unicast and Layer 3 broadcast transmissions. For broadcasts, the destination IP should be specified
as an IPv4Network. It delegates the actual sending process to the SoftwareManager.
:param payload: The payload to be sent.
:param dest_ip_address: The ip address of the payload destination.
:param dest_port: The port of the payload destination.
:param session_id: The Session ID the payload is to originate from. Optional.
:return: True if successful, False otherwise.
:param dest_ip_address: The IP address or network (for broadcasts) of the payload destination.
:param dest_port: The destination port for the payload. Optional.
:param session_id: The Session ID from which the payload originates. Optional.
:return: True if the payload was successfully sent, False otherwise.
"""
if not self._can_perform_action():
return False

180
tests/integration_tests/network/test_broadcast.py Normal file
View File

@@ -0,0 +1,180 @@
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, List, Tuple
import pytest
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.services.service import Service
class BroadcastService(Service):
"""A service for sending broadcast and unicast messages over a network."""
def __init__(self, **kwargs):
# Set default service properties for broadcasting
kwargs["name"] = "BroadcastService"
kwargs["port"] = Port.HTTP
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
def describe_state(self) -> Dict:
# Implement state description for the service
pass
def unicast(self, ip_address: IPv4Address):
# Send a unicast payload to a specific IP address
super().send(
payload="unicast",
dest_ip_address=ip_address,
dest_port=Port.HTTP,
)
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,
)
class BroadcastClient(Application):
"""A client application to receive broadcast and unicast messages."""
payloads_received: List = []
def __init__(self, **kwargs):
# Set default client properties
kwargs["name"] = "BroadcastClient"
kwargs["port"] = Port.HTTP
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
def describe_state(self) -> Dict:
# Implement state description for the application
pass
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
# Append received payloads to the list and print a message
self.payloads_received.append(payload)
print(f"Payload: {payload} received on node {self.sys_log.hostname}")
@pytest.fixture(scope="function")
def broadcast_network() -> Network:
network = Network()
client_1 = Computer(
hostname="client_1",
ip_address="192.168.1.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
client_1.power_on()
client_1.software_manager.install(BroadcastClient)
application_1 = client_1.software_manager.software["BroadcastClient"]
application_1.run()
client_2 = Computer(
hostname="client_2",
ip_address="192.168.1.3",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
client_2.power_on()
client_2.software_manager.install(BroadcastClient)
application_2 = client_2.software_manager.software["BroadcastClient"]
application_2.run()
server_1 = Server(
hostname="server_1",
ip_address="192.168.1.1",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
)
server_1.power_on()
server_1.software_manager.install(BroadcastService)
service: BroadcastService = server_1.software_manager.software["BroadcastService"]
service.start()
switch_1 = Switch(hostname="switch_1", num_ports=6, start_up_duration=0)
switch_1.power_on()
network.connect(endpoint_a=client_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[1])
network.connect(endpoint_a=client_2.ethernet_port[1], endpoint_b=switch_1.switch_ports[2])
network.connect(endpoint_a=server_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[3])
return network
@pytest.fixture(scope="function")
def broadcast_service_and_clients(broadcast_network) -> Tuple[BroadcastService, BroadcastClient, BroadcastClient]:
client_1: BroadcastClient = broadcast_network.get_node_by_hostname("client_1").software_manager.software[
"BroadcastClient"
]
client_2: BroadcastClient = broadcast_network.get_node_by_hostname("client_2").software_manager.software[
"BroadcastClient"
]
service: BroadcastService = broadcast_network.get_node_by_hostname("server_1").software_manager.software[
"BroadcastService"
]
return service, client_1, client_2
def test_broadcast_correct_subnet(broadcast_service_and_clients):
service, client_1, client_2 = broadcast_service_and_clients
assert not client_1.payloads_received
assert not client_2.payloads_received
service.broadcast(IPv4Network("192.168.1.0/24"))
assert client_1.payloads_received == ["broadcast"]
assert client_2.payloads_received == ["broadcast"]
def test_broadcast_incorrect_subnet(broadcast_service_and_clients):
service, client_1, client_2 = broadcast_service_and_clients
assert not client_1.payloads_received
assert not client_2.payloads_received
service.broadcast(IPv4Network("192.168.2.0/24"))
assert not client_1.payloads_received
assert not client_2.payloads_received
def test_unicast_correct_address(broadcast_service_and_clients):
service, client_1, client_2 = broadcast_service_and_clients
assert not client_1.payloads_received
assert not client_2.payloads_received
service.unicast(IPv4Address("192.168.1.2"))
assert client_1.payloads_received == ["unicast"]
assert not client_2.payloads_received
def test_unicast_incorrect_address(broadcast_service_and_clients):
service, client_1, client_2 = broadcast_service_and_clients
assert not client_1.payloads_received
assert not client_2.payloads_received
service.unicast(IPv4Address("192.168.2.2"))
assert not client_1.payloads_received
assert not client_2.payloads_received