#1724 - Added the primaite/simulator/network/transmission sub-package with modules for each layer. They come together to build a minimal but fairly realistic network Frame. A custom PrimaiteHeader has been included to hold primaite specific metadata required in transmission for reward function and RL agent downstream. Added some basic tests that check the proper configuration of Frames with matching headers for protocols. Updated the frame typehints in NIC and Link classes.

This commit is contained in:
Chris McCarthy
2023-08-01 22:25:00 +01:00
parent f41fc241b7
commit 9d17a9b0d3
11 changed files with 577 additions and 9 deletions

View File

@@ -1,6 +1,6 @@
import pytest
from primaite.simulator.network.physical_layer import Link, NIC
from primaite.simulator.network.transmission.physical_layer import Link, NIC
def test_link_fails_with_same_nic():

View File

@@ -0,0 +1,90 @@
import pytest
from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame
from primaite.simulator.network.transmission.network_layer import ICMPHeader, IPPacket, IPProtocol, Precedence
from primaite.simulator.network.transmission.primaite_layer import AgentSource, DataStatus
from primaite.simulator.network.transmission.transport_layer import Port, TCPFlags, TCPHeader, UDPHeader
def test_frame_minimal_instantiation():
"""Tests that the minimum frame (TCP SYN) using default values."""
frame = Frame(
ethernet=EthernetHeader(src_mac_addr="aa:bb:cc:dd:ee:ff", dst_mac_addr="11:22:33:44:55:66"),
ip=IPPacket(src_ip="192.168.0.10", dst_ip="192.168.0.20"),
tcp=TCPHeader(
src_port=8080,
dst_port=80,
),
)
# Check network layer default values
assert frame.ip.protocol == IPProtocol.TCP
assert frame.ip.ttl == 64
assert frame.ip.precedence == Precedence.ROUTINE
# Check transport layer default values
assert frame.tcp.flags == [TCPFlags.SYN]
# Check primaite custom header default values
assert frame.primaite_header.agent_source == AgentSource.GREEN
assert frame.primaite_header.data_status == DataStatus.GOOD
# Check that model can be dumped down to json and returned as size in Bytes
assert frame.size
def test_frame_creation_fails_tcp_without_header():
"""Tests Frame creation fails if the IPProtocol is TCP but there is no TCPHeader."""
with pytest.raises(ValueError):
Frame(
ethernet=EthernetHeader(src_mac_addr="aa:bb:cc:dd:ee:ff", dst_mac_addr="11:22:33:44:55:66"),
ip=IPPacket(src_ip="192.168.0.10", dst_ip="192.168.0.20", protocol=IPProtocol.TCP),
)
def test_frame_creation_fails_udp_without_header():
"""Tests Frame creation fails if the IPProtocol is UDP but there is no UDPHeader."""
with pytest.raises(ValueError):
Frame(
ethernet=EthernetHeader(src_mac_addr="aa:bb:cc:dd:ee:ff", dst_mac_addr="11:22:33:44:55:66"),
ip=IPPacket(src_ip="192.168.0.10", dst_ip="192.168.0.20", protocol=IPProtocol.UDP),
)
def test_frame_creation_fails_tcp_with_udp_header():
"""Tests Frame creation fails if the IPProtocol is TCP but there is a UDPHeader."""
with pytest.raises(ValueError):
Frame(
ethernet=EthernetHeader(src_mac_addr="aa:bb:cc:dd:ee:ff", dst_mac_addr="11:22:33:44:55:66"),
ip=IPPacket(src_ip="192.168.0.10", dst_ip="192.168.0.20", protocol=IPProtocol.TCP),
udp=UDPHeader(src_port=8080, dst_port=80),
)
def test_frame_creation_fails_udp_with_tcp_header():
"""Tests Frame creation fails if the IPProtocol is UDP but there is a TCPHeader."""
with pytest.raises(ValueError):
Frame(
ethernet=EthernetHeader(src_mac_addr="aa:bb:cc:dd:ee:ff", dst_mac_addr="11:22:33:44:55:66"),
ip=IPPacket(src_ip="192.168.0.10", dst_ip="192.168.0.20", protocol=IPProtocol.UDP),
udp=TCPHeader(src_port=8080, dst_port=80),
)
def test_icmp_frame_creation():
"""Tests Frame creation for ICMP."""
frame = Frame(
ethernet=EthernetHeader(src_mac_addr="aa:bb:cc:dd:ee:ff", dst_mac_addr="11:22:33:44:55:66"),
ip=IPPacket(src_ip="192.168.0.10", dst_ip="192.168.0.20", protocol=IPProtocol.ICMP),
icmp=ICMPHeader(),
)
assert frame
def test_icmp_frame_creation_fails_without_icmp_header():
"""Tests Frame creation for ICMP."""
with pytest.raises(ValueError):
Frame(
ethernet=EthernetHeader(src_mac_addr="aa:bb:cc:dd:ee:ff", dst_mac_addr="11:22:33:44:55:66"),
ip=IPPacket(src_ip="192.168.0.10", dst_ip="192.168.0.20", protocol=IPProtocol.ICMP),
)

View File

@@ -0,0 +1,24 @@
import pytest
from primaite.simulator.network.transmission.network_layer import ICMPHeader, ICMPType
def test_icmp_minimal_header_creation():
"""Checks the minimal ICMPHeader (ping 1 request) creation using default values."""
ping = ICMPHeader()
assert ping.icmp_type == ICMPType.ECHO_REQUEST
assert ping.icmp_code == 0
assert ping.identifier
assert ping.sequence == 1
def test_valid_icmp_type_code_pairing():
"""Tests ICMPHeader creation with valid type and code pairing."""
assert ICMPHeader(icmp_type=ICMPType.DESTINATION_UNREACHABLE, icmp_code=6)
def test_invalid_icmp_type_code_pairing():
"""Tests ICMPHeader creation fails with invalid type and code pairing."""
with pytest.raises(ValueError):
assert ICMPHeader(icmp_type=ICMPType.DESTINATION_UNREACHABLE, icmp_code=16)

View File

@@ -3,7 +3,7 @@ from ipaddress import IPv4Address
import pytest
from primaite.simulator.network.physical_layer import generate_mac_address, NIC
from primaite.simulator.network.transmission.physical_layer import generate_mac_address, NIC
def test_mac_address_generation():