#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:
@@ -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():
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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():
|
||||
Reference in New Issue
Block a user