From c6f71600fc39c2b58bee837b6ed99db0fbcfec18 Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Fri, 25 Aug 2023 09:07:32 +0100 Subject: [PATCH 01/10] #1800 - Fixed the ping functionality so that it actually checks for replies and returns True if the right number of replies have been received. - Added the foundations of a Router class along with ACLRule and RouteTableEntry classes. --- .../simulator/network/hardware/base.py | 25 ++++-- .../network/hardware/nodes/router.py | 86 +++++++++++++++++++ .../network/test_frame_transmission.py | 4 +- .../integration_tests/network/test_routing.py | 27 ++++++ 4 files changed, 134 insertions(+), 8 deletions(-) create mode 100644 src/primaite/simulator/network/hardware/nodes/router.py create mode 100644 tests/integration_tests/network/test_routing.py diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 28e7693a..c64b9b67 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -173,6 +173,9 @@ class NIC(SimComponent): if self.connected_node.operating_state != NodeOperatingState.ON: self.connected_node.sys_log.error(f"NIC {self} cannot be enabled as the endpoint is not turned on") return + if not self.connected_link: + _LOGGER.error(f"NIC {self} cannot be enabled as it is not connected to a Link") + return self.enabled = True self.connected_node.sys_log.info(f"NIC {self} enabled") @@ -210,6 +213,7 @@ class NIC(SimComponent): # TODO: Inform the Node that a link has been connected self.connected_link = link + self.enable() _LOGGER.info(f"NIC {self} connected to Link {link}") def disconnect_link(self): @@ -266,8 +270,10 @@ class NIC(SimComponent): frame.decrement_ttl() frame.set_received_timestamp() self.pcap.capture(frame) - self.connected_node.receive_frame(frame=frame, from_nic=self) - return True + # 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: @@ -688,7 +694,6 @@ class ARPCache: frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_packet) from_nic.send_frame(frame) - class ICMP: """ The ICMP (Internet Control Message Protocol) class. @@ -705,6 +710,8 @@ class ICMP: """ self.sys_log: SysLog = sys_log self.arp: ARPCache = arp_cache + self.request_replies = {} + def process_icmp(self, frame: Frame): """ @@ -733,6 +740,9 @@ class ICMP: src_nic.send_frame(frame) elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY: self.sys_log.info(f"Received echo reply from {frame.ip.src_ip}") + if not self.request_replies.get(frame.icmp.identifier): + self.request_replies[frame.icmp.identifier] = 0 + self.request_replies[frame.icmp.identifier] += 1 def ping( self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None @@ -875,7 +885,7 @@ class Node(SimComponent): return state def show(self): - """Prints a table of the NICs on the Node..""" + """Prints a table of the NICs on the Node.""" from prettytable import PrettyTable table = PrettyTable(["MAC Address", "Address", "Default Gateway", "Speed", "Status"]) @@ -898,7 +908,8 @@ class Node(SimComponent): self.operating_state = NodeOperatingState.ON self.sys_log.info("Turned on") for nic in self.nics.values(): - nic.enable() + if nic.connected_link: + nic.enable() def power_off(self): """Power off the Node, disabling its NICs if it is in the ON state.""" @@ -961,7 +972,9 @@ class Node(SimComponent): sequence, identifier = 0, None while sequence < pings: sequence, identifier = self.icmp.ping(target_ip_address, sequence, identifier) - return True + passed = self.icmp.request_replies[identifier] == pings + self.icmp.request_replies.pop(identifier) + return passed self.sys_log.info("Ping failed as the node is turned off") return False diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py new file mode 100644 index 00000000..c5620b88 --- /dev/null +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -0,0 +1,86 @@ +from enum import Enum +from ipaddress import IPv4Address +from typing import Dict, List, Union + +from primaite.simulator.core import SimComponent +from primaite.simulator.network.hardware.base import Node, NIC +from prettytable import PrettyTable + +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port + + +class ACLAction(Enum): + DENY = 0 + PERMIT = 1 + + +class ACLRule(SimComponent): + action: ACLAction + protocol: IPProtocol + src_ip: IPv4Address + src_wildcard: IPv4Address = IPv4Address("0.0.0.0") + src_port: Port + dst_ip: IPv4Address + dst_port: Port + + +class RouteTableEntry(SimComponent): + pass + + +class Router(Node): + num_ports: int + ethernet_ports: Dict[int, NIC] = {} + acl: List = [] + route_table: Dict = {} + + def __init__(self, hostname: str, num_ports: int = 5, **kwargs): + 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") + self.connect_nic(nic) + self.ethernet_ports[i] = nic + + def describe_state(self) -> Dict: + pass + + def configure_port( + self, + port: int, + ip_address: Union[IPv4Address, str], + subnet_mask: str + ): + if not isinstance(ip_address, IPv4Address): + ip_address = IPv4Address(ip_address) + nic = self.ethernet_ports[port] + nic.ip_address = ip_address + nic.subnet_mask = subnet_mask + self.sys_log.info(f"Configured port {port} with {ip_address=} {subnet_mask=}") + + def enable_port(self, port: int): + nic = self.ethernet_ports.get(port) + if nic: + nic.enable() + + def disable_port(self, port: int): + nic = self.ethernet_ports.get(port) + if nic: + nic.disable() + + def show(self): + """Prints a table of the NICs on the Node.""" + table = PrettyTable(["Port", "MAC Address", "Address", "Speed", "Status"]) + + for port, nic in self.ethernet_ports.items(): + table.add_row( + [ + port, + nic.mac_address, + f"{nic.ip_address}/{nic.ip_network.prefixlen}", + nic.speed, + "Enabled" if nic.enabled else "Disabled", + ] + ) + print(table) diff --git a/tests/integration_tests/network/test_frame_transmission.py b/tests/integration_tests/network/test_frame_transmission.py index 3840c302..d3d6541a 100644 --- a/tests/integration_tests/network/test_frame_transmission.py +++ b/tests/integration_tests/network/test_frame_transmission.py @@ -45,7 +45,7 @@ def test_multi_nic(): node_a.ping("192.168.0.11") - node_c.ping("10.0.0.12") + assert node_c.ping("10.0.0.12") def test_switched_network(): @@ -83,4 +83,4 @@ def test_switched_network(): link_nic_d_switch_2 = Link(endpoint_a=nic_d, endpoint_b=switch_2.switch_ports[2]) link_switch_1_switch_2 = Link(endpoint_a=switch_1.switch_ports[6], endpoint_b=switch_2.switch_ports[6]) - pc_a.ping("192.168.0.13") + assert pc_a.ping("192.168.0.13") diff --git a/tests/integration_tests/network/test_routing.py b/tests/integration_tests/network/test_routing.py new file mode 100644 index 00000000..cca48c0d --- /dev/null +++ b/tests/integration_tests/network/test_routing.py @@ -0,0 +1,27 @@ +from primaite.simulator.network.hardware.base import Node, NIC, Link +from primaite.simulator.network.hardware.nodes.router import Router + + +def test_ping_fails_with_no_route(): + """Tests a larges network of Nodes and Switches with one node pinging another.""" + pc_a = Node(hostname="pc_a") + nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1") + pc_a.connect_nic(nic_a) + pc_a.power_on() + + pc_b = Node(hostname="pc_b") + nic_b = NIC(ip_address="192.168.1.10", subnet_mask="255.255.255.0", gateway="192.168.1.1") + pc_b.connect_nic(nic_b) + pc_b.power_on() + + router_1 = Router(hostname="router_1") + router_1.configure_port(1, "192.168.0.1", "255.255.255.0") + router_1.configure_port(2, "192.168.1.1", "255.255.255.0") + + router_1.power_on() + router_1.show() + + link_nic_a_router_1 = Link(endpoint_a=nic_a, endpoint_b=router_1.ethernet_ports[1]) + link_nic_b_router_1 = Link(endpoint_a=nic_b, endpoint_b=router_1.ethernet_ports[2]) + router_1.power_on() + #assert pc_a.ping("192.168.1.10") \ No newline at end of file From 1bf51c7741f1c45dd6846d9bb3ee611aadd02bf1 Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Wed, 30 Aug 2023 21:38:55 +0100 Subject: [PATCH 02/10] #1800 - Added ACL and routing classes. - Added .show() methods to new router classes to enable inspection of the components as you would a real router. - Removed gateway from the NIC and added default_gateway to Node so that Node has a single default gateway. - Added some routing tests to check that ping can be performed when router between subnets. --- .../simulator/network/hardware/base.py | 146 +++-- .../network/hardware/nodes/router.py | 505 +++++++++++++++++- .../network/test_frame_transmission.py | 21 +- .../network/test_link_connection.py | 9 +- .../network/test_nic_link_connection.py | 3 +- .../integration_tests/network/test_routing.py | 56 +- .../_network/_hardware/nodes/__init__.py | 0 .../_network/_hardware/nodes/test_router.py | 104 ++++ .../_simulator/_network/_hardware/test_nic.py | 14 - 9 files changed, 739 insertions(+), 119 deletions(-) create mode 100644 tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/__init__.py create mode 100644 tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_router.py diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index c64b9b67..921ebbcd 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -77,12 +77,10 @@ class NIC(SimComponent): ip_address: IPv4Address "The IP address assigned to the NIC for communication on an IP-based network." - subnet_mask: str + subnet_mask: IPv4Address "The subnet mask assigned to the NIC." - gateway: IPv4Address - "The default gateway IP address for forwarding network traffic to other networks. Randomly generated upon creation." mac_address: str - "The MAC address of the NIC. Defaults to a randomly set MAC address." + "The MAC address of the NIC. Defaults to a randomly set MAC address. Randomly generated upon creation." speed: int = 100 "The speed of the NIC in Mbps. Default is 100 Mbps." mtu: int = 1500 @@ -111,16 +109,10 @@ class NIC(SimComponent): """ if not isinstance(kwargs["ip_address"], IPv4Address): kwargs["ip_address"] = IPv4Address(kwargs["ip_address"]) - if not isinstance(kwargs["gateway"], IPv4Address): - kwargs["gateway"] = IPv4Address(kwargs["gateway"]) if "mac_address" not in kwargs: kwargs["mac_address"] = generate_mac_address() super().__init__(**kwargs) - if self.ip_address == self.gateway: - msg = f"NIC ip address {self.ip_address} cannot be the same as the gateway {self.gateway}" - _LOGGER.error(msg) - raise ValueError(msg) if self.ip_network.network_address == self.ip_address: msg = ( f"Failed to set IP address {self.ip_address} and subnet mask {self.subnet_mask} as it is a " @@ -274,6 +266,9 @@ class NIC(SimComponent): 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 + else: + self.connected_node.sys_log.info("Dropping frame not for me") + print(frame) return False def __str__(self) -> str: @@ -567,7 +562,21 @@ class ARPCache: self.arp: Dict[IPv4Address, ARPEntry] = {} self.nics: Dict[str, "NIC"] = {} - def _add_arp_cache_entry(self, ip_address: IPv4Address, mac_address: str, nic: NIC): + def show(self): + """Prints a table of ARC Cache.""" + table = PrettyTable(["IP Address", "MAC Address", "Via"]) + table.title = f"{self.sys_log.hostname} ARP Cache" + for ip, arp in self.arp.items(): + table.add_row( + [ + str(ip), + arp.mac_address, + self.nics[arp.nic_uuid].mac_address, + ] + ) + print(table) + + def add_arp_cache_entry(self, ip_address: IPv4Address, mac_address: str, nic: NIC, override: bool = False): """ Add an ARP entry to the cache. @@ -575,9 +584,14 @@ class ARPCache: :param mac_address: The MAC address associated with the IP address. :param nic: The NIC through which the NIC with the IP address is reachable. """ - self.sys_log.info(f"Adding ARP cache entry for {mac_address}/{ip_address} via NIC {nic}") - arp_entry = ARPEntry(mac_address=mac_address, nic_uuid=nic.uuid) - self.arp[ip_address] = arp_entry + for _nic in self.nics.values(): + if _nic.ip_address == ip_address: + return + if override or not self.arp.get(ip_address): + self.sys_log.info(f"Adding ARP cache entry for {mac_address}/{ip_address} via NIC {nic}") + arp_entry = ARPEntry(mac_address=mac_address, nic_uuid=nic.uuid) + + self.arp[ip_address] = arp_entry def _remove_arp_cache_entry(self, ip_address: IPv4Address): """ @@ -607,6 +621,7 @@ class ARPCache: :return: The NIC associated with the IP address, or None if not found. """ arp_entry = self.arp.get(ip_address) + if arp_entry: return self.nics[arp_entry.nic_uuid] @@ -641,6 +656,29 @@ class ARPCache: frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_packet) nic.send_frame(frame) + def send_arp_reply(self, arp_reply: ARPPacket, from_nic: NIC): + """ + Send an ARP reply back through the NIC it came from. + + :param arp_reply: The ARP reply to send. + :param from_nic: The NIC to send the ARP reply from. + """ + self.sys_log.info( + f"Sending ARP reply from {arp_reply.sender_mac_addr}/{arp_reply.sender_ip} " + f"to {arp_reply.target_ip}/{arp_reply.target_mac_addr} " + ) + tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) + + ip_packet = IPPacket( + src_ip=arp_reply.sender_ip, + dst_ip=arp_reply.target_ip, + ) + + ethernet_header = EthernetHeader(src_mac_addr=arp_reply.sender_mac_addr, dst_mac_addr=arp_reply.target_mac_addr) + + frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_reply) + from_nic.send_frame(frame) + def process_arp_packet(self, from_nic: NIC, arp_packet: ARPPacket): """ Process a received ARP packet, handling both ARP requests and responses. @@ -656,7 +694,7 @@ class ARPCache: self.sys_log.info( f"Received ARP response for {arp_packet.sender_ip} from {arp_packet.sender_mac_addr} via NIC {from_nic}" ) - self._add_arp_cache_entry( + self.add_arp_cache_entry( ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic ) return @@ -673,26 +711,13 @@ class ARPCache: return # Matched ARP request - self._add_arp_cache_entry(ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic) + self.add_arp_cache_entry(ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic) arp_packet = arp_packet.generate_reply(from_nic.mac_address) - self.sys_log.info( - f"Sending ARP reply from {arp_packet.sender_mac_addr}/{arp_packet.sender_ip} " - f"to {arp_packet.target_ip}/{arp_packet.target_mac_addr} " - ) + self.send_arp_reply(arp_packet, from_nic) - tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) + def __contains__(self, item) -> bool: + return item in self.arp - # Network Layer - ip_packet = IPPacket( - src_ip=arp_packet.sender_ip, - dst_ip=arp_packet.target_ip, - ) - # Data Link Layer - ethernet_header = EthernetHeader( - src_mac_addr=arp_packet.sender_mac_addr, dst_mac_addr=arp_packet.target_mac_addr - ) - frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_packet) - from_nic.send_frame(frame) class ICMP: """ @@ -712,8 +737,7 @@ class ICMP: self.arp: ARPCache = arp_cache self.request_replies = {} - - def process_icmp(self, frame: Frame): + def process_icmp(self, frame: Frame, from_nic: NIC, is_reattempt: bool = False): """ Process an ICMP packet, including handling echo requests and replies. @@ -722,7 +746,15 @@ class ICMP: if frame.icmp.icmp_type == ICMPType.ECHO_REQUEST: self.sys_log.info(f"Received echo request from {frame.ip.src_ip}") target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip) + src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip) + if not src_nic: + print(self.sys_log.hostname) + print(frame.ip.src_ip) + self.arp.show() + self.arp.send_arp_request(frame.ip.src_ip) + self.process_icmp(frame=frame, from_nic=from_nic, is_reattempt=True) + return tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) # Network Layer @@ -737,6 +769,7 @@ class ICMP: ) frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet) self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") + src_nic.send_frame(frame) elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY: self.sys_log.info(f"Received echo reply from {frame.ip.src_ip}") @@ -745,7 +778,7 @@ class ICMP: self.request_replies[frame.icmp.identifier] += 1 def ping( - self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None + self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None, pings: int = 4 ) -> Tuple[int, Union[int, None]]: """ Send an ICMP echo request (ping) to a target IP address and manage the sequence and identifier. @@ -757,13 +790,21 @@ class ICMP: was not found in the ARP cache. """ nic = self.arp.get_arp_cache_nic(target_ip_address) - # TODO: Eventually this ARP request needs to be done elsewhere. It's not the resonsibility of the + # TODO: Eventually this ARP request needs to be done elsewhere. It's not the responsibility of the # ping function to handle ARP lookups + + # Already tried once and cannot get ARP entry, stop trying + if sequence == -1: + if not nic: + return 4, None + else: + sequence = 0 + # No existing ARP entry if not nic: self.sys_log.info(f"No entry in ARP cache for {target_ip_address}") self.arp.send_arp_request(target_ip_address) - return 0, None + return -1, None # ARP entry exists sequence += 1 @@ -812,6 +853,8 @@ class Node(SimComponent): hostname: str "The node hostname on the network." + default_gateway: Optional[IPv4Address] = None + "The default gateway IP address for forwarding network traffic to other networks." operating_state: NodeOperatingState = NodeOperatingState.OFF "The hardware state of the node." nics: Dict[str, NIC] = {} @@ -843,9 +886,12 @@ class Node(SimComponent): This method initializes the ARP cache, ICMP handler, session manager, and software manager if they are not provided. """ + if kwargs.get("default_gateway"): + if not isinstance(kwargs["default_gateway"], IPv4Address): + kwargs["default_gateway"] = IPv4Address(kwargs["default_gateway"]) if not kwargs.get("sys_log"): kwargs["sys_log"] = SysLog(kwargs["hostname"]) - if not kwargs.get("arp_cache"): + if not kwargs.get("arp"): kwargs["arp"] = ARPCache(sys_log=kwargs.get("sys_log")) if not kwargs.get("icmp"): kwargs["icmp"] = ICMP(sys_log=kwargs.get("sys_log"), arp_cache=kwargs.get("arp")) @@ -886,10 +932,8 @@ class Node(SimComponent): def show(self): """Prints a table of the NICs on the Node.""" - from prettytable import PrettyTable - table = PrettyTable(["MAC Address", "Address", "Default Gateway", "Speed", "Status"]) - + table.title = f"{self.hostname} Network Interface Cards" for nic in self.nics.values(): table.add_row( [ @@ -967,13 +1011,18 @@ class Node(SimComponent): """ if not isinstance(target_ip_address, IPv4Address): target_ip_address = IPv4Address(target_ip_address) + if target_ip_address.is_loopback: + self.sys_log.info("Pinging loopback address") + return any(nic.enabled for nic in self.nics.values()) if self.operating_state == NodeOperatingState.ON: self.sys_log.info(f"Attempting to ping {target_ip_address}") sequence, identifier = 0, None while sequence < pings: - sequence, identifier = self.icmp.ping(target_ip_address, sequence, identifier) - passed = self.icmp.request_replies[identifier] == pings - self.icmp.request_replies.pop(identifier) + sequence, identifier = self.icmp.ping(target_ip_address, sequence, identifier, pings) + request_replies = self.icmp.request_replies.get(identifier) + passed = request_replies == pings + if request_replies: + self.icmp.request_replies.pop(identifier) return passed self.sys_log.info("Ping failed as the node is turned off") return False @@ -997,13 +1046,18 @@ class Node(SimComponent): :param frame: The Frame being received. :param from_nic: The NIC that received the frame. """ + if frame.ip: + if frame.ip.src_ip in self.arp: + self.arp.add_arp_cache_entry( + ip_address=frame.ip.src_ip, mac_address=frame.ethernet.src_mac_addr, nic=from_nic + ) if frame.ip.protocol == IPProtocol.TCP: if frame.tcp.src_port == Port.ARP: self.arp.process_arp_packet(from_nic=from_nic, arp_packet=frame.arp) elif frame.ip.protocol == IPProtocol.UDP: pass elif frame.ip.protocol == IPProtocol.ICMP: - self.icmp.process_icmp(frame=frame) + self.icmp.process_icmp(frame=frame, from_nic=from_nic) class Switch(Node): @@ -1027,7 +1081,7 @@ class Switch(Node): def show(self): """Prints a table of the SwitchPorts on the Switch.""" table = PrettyTable(["Port", "MAC Address", "Speed", "Status"]) - + table.title = f"{self.hostname} Switch Ports" for port_num, port in self.switch_ports.items(): table.add_row([port_num, port.mac_address, port.speed, "Enabled" if port.enabled else "Disabled"]) print(table) diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index c5620b88..528e4a73 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -1,63 +1,514 @@ -from enum import Enum -from ipaddress import IPv4Address -from typing import Dict, List, Union +from __future__ import annotations + +from enum import Enum +from ipaddress import IPv4Address, IPv4Network +from typing import Dict, List, Optional, Tuple, Union -from primaite.simulator.core import SimComponent -from primaite.simulator.network.hardware.base import Node, NIC from prettytable import PrettyTable -from primaite.simulator.network.transmission.network_layer import IPProtocol -from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.core import SimComponent +from primaite.simulator.network.hardware.base import ARPCache, ICMP, NIC, Node +from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame +from primaite.simulator.network.transmission.network_layer import ICMPPacket, ICMPType, IPPacket, IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader +from primaite.simulator.system.core.sys_log import SysLog class ACLAction(Enum): + """Enum for defining the ACL action types.""" + DENY = 0 PERMIT = 1 class ACLRule(SimComponent): - action: ACLAction - protocol: IPProtocol - src_ip: IPv4Address - src_wildcard: IPv4Address = IPv4Address("0.0.0.0") - src_port: Port - dst_ip: IPv4Address - dst_port: Port + def describe_state(self) -> Dict: + pass + + action: ACLAction = ACLAction.DENY + protocol: Optional[IPProtocol] = None + src_ip: Optional[IPv4Address] = None + src_port: Optional[Port] = None + dst_ip: Optional[IPv4Address] = None + dst_port: Optional[Port] = None + + def __str__(self) -> str: + rule_strings = [] + for key, value in self.model_dump(exclude={"uuid", "action_manager"}).items(): + if value is None: + value = "ANY" + if isinstance(value, Enum): + rule_strings.append(f"{key}={value.name}") + else: + rule_strings.append(f"{key}={value}") + return ", ".join(rule_strings) -class RouteTableEntry(SimComponent): - pass +class AccessControlList(SimComponent): + sys_log: SysLog + implicit_action: ACLAction + implicit_rule: ACLRule + max_acl_rules: int = 25 + _acl: List[Optional[ACLRule]] = [None] * 24 + + def __init__(self, **kwargs) -> None: + if not kwargs.get("implicit_action"): + kwargs["implicit_action"] = ACLAction.DENY + if not kwargs.get("max_acl_rules"): + kwargs["max_acl_rules"] = 25 + kwargs["implicit_rule"] = ACLRule(action=kwargs["implicit_action"]) + kwargs["_acl"] = [None] * (kwargs["max_acl_rules"] - 1) + + super().__init__(**kwargs) + + def describe_state(self) -> Dict: + pass + + @property + def acl(self) -> List[Optional[ACLRule]]: + return self._acl + + def add_rule( + self, + action: ACLAction, + protocol: Optional[IPProtocol] = None, + src_ip: Optional[Union[str, IPv4Address]] = None, + src_port: Optional[Port] = None, + dst_ip: Optional[Union[str, IPv4Address]] = None, + dst_port: Optional[Port] = None, + position: int = 0, + ) -> None: + if isinstance(src_ip, str): + src_ip = IPv4Address(src_ip) + if isinstance(dst_ip, str): + dst_ip = IPv4Address(dst_ip) + if 0 <= position < self.max_acl_rules: + self._acl[position] = ACLRule( + action=action, src_ip=src_ip, dst_ip=dst_ip, protocol=protocol, src_port=src_port, dst_port=dst_port + ) + else: + raise ValueError(f"Position {position} is out of bounds.") + + def remove_rule(self, position: int) -> None: + if 0 <= position < self.max_acl_rules: + self._acl[position] = None + else: + raise ValueError(f"Position {position} is out of bounds.") + + def is_permitted( + self, + protocol: IPProtocol, + src_ip: Union[str, IPv4Address], + src_port: Optional[Port], + dst_ip: Union[str, IPv4Address], + dst_port: Optional[Port], + ) -> Tuple[bool, Optional[Union[str, ACLRule]]]: + if not isinstance(src_ip, IPv4Address): + src_ip = IPv4Address(src_ip) + if not isinstance(dst_ip, IPv4Address): + dst_ip = IPv4Address(dst_ip) + for rule in self._acl: + if not rule: + continue + + if ( + (rule.src_ip == src_ip or rule.src_ip is None) + and (rule.dst_ip == dst_ip or rule.dst_ip 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: Union[str, IPv4Address], + src_port: Port, + dst_ip: Union[str, IPv4Address], + dst_port: Port, + ) -> List[ACLRule]: + if not isinstance(src_ip, IPv4Address): + src_ip = IPv4Address(src_ip) + if not isinstance(dst_ip, IPv4Address): + dst_ip = IPv4Address(dst_ip) + relevant_rules = [] + for rule in self._acl: + if rule is None: + continue + + if ( + (rule.src_ip == src_ip or rule.src_ip is None) + or (rule.dst_ip == dst_ip or rule.dst_ip 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) + + return relevant_rules + + def show(self): + """Prints a table of the routes in the RouteTable.""" + """ + action: ACLAction + protocol: Optional[IPProtocol] + src_ip: Optional[IPv4Address] + src_port: Optional[Port] + dst_ip: Optional[IPv4Address] + dst_port: Optional[Port] + """ + table = PrettyTable(["Index", "Action", "Protocol", "Src IP", "Src Port", "Dst IP", "Dst Port"]) + table.title = f"{self.sys_log.hostname} Access Control List" + for index, rule in enumerate(self.acl + [self.implicit_rule]): + if rule: + table.add_row( + [ + index, + rule.action.name if rule.action else "ANY", + rule.protocol.name if rule.protocol else "ANY", + rule.src_ip if rule.src_ip else "ANY", + f"{rule.src_port.value} ({rule.src_port.name})" if rule.src_port else "ANY", + rule.dst_ip if rule.dst_ip else "ANY", + f"{rule.dst_port.value} ({rule.dst_port.name})" if rule.dst_port else "ANY", + ] + ) + print(table) + + +class RouteEntry(SimComponent): + """ + Represents a single entry in a routing table. + + Attributes: + address (IPv4Address): The destination IP address or network address. + subnet_mask (IPv4Address): The subnet mask for the network. + next_hop (IPv4Address): The next hop IP address to which packets should be forwarded. + metric (int): The cost metric for this route. Default is 0.0. + + Example: + >>> entry = RouteEntry( + ... IPv4Address("192.168.1.0"), + ... IPv4Address("255.255.255.0"), + ... IPv4Address("192.168.2.1"), + ... metric=5 + ... ) + """ + + address: IPv4Address + "The destination IP address or network address." + subnet_mask: IPv4Address + "The subnet mask for the network." + next_hop: IPv4Address + "The next hop IP address to which packets should be forwarded." + metric: float = 0.0 + "The cost metric for this route. Default is 0.0." + + def __init__(self, **kwargs): + for key in {"address", "subnet_mask", "next_hop"}: + if not isinstance(kwargs[key], IPv4Address): + kwargs[key] = IPv4Address(kwargs[key]) + super().__init__(**kwargs) + + def describe_state(self) -> Dict: + pass + + +class RouteTable(SimComponent): + """ + Represents a routing table holding multiple route entries. + + Attributes: + routes (List[RouteEntry]): A list of RouteEntry objects. + + Methods: + add_route: Add a route to the routing table. + find_best_route: Find the best route for a given destination IP. + + Example: + >>> rt = RouteTable() + >>> rt.add_route( + ... RouteEntry( + ... IPv4Address("192.168.1.0"), + ... IPv4Address("255.255.255.0"), + ... IPv4Address("192.168.2.1"), + ... metric=5 + ... ) + ... ) + >>> best_route = rt.find_best_route(IPv4Address("192.168.1.34")) + """ + + routes: List[RouteEntry] = [] + sys_log: SysLog + + def describe_state(self) -> Dict: + pass + + def add_route( + self, + address: Union[IPv4Address, str], + subnet_mask: Union[IPv4Address, str], + next_hop: Union[IPv4Address, str], + metric: float = 0.0, + ): + """Add a route to the routing table. + + :param route: A RouteEntry object representing the route. + """ + for key in {address, subnet_mask, next_hop}: + if not isinstance(key, IPv4Address): + key = IPv4Address(key) + route = RouteEntry(address=address, subnet_mask=subnet_mask, next_hop=next_hop, metric=metric) + self.routes.append(route) + + def find_best_route(self, destination_ip: Union[str, IPv4Address]) -> Optional[RouteEntry]: + """ + Find the best route for a given destination IP. + + :param destination_ip: The destination IPv4Address to find the route for. + :return: The best matching RouteEntry, or None if no route matches. + + The algorithm uses Longest Prefix Match and considers metrics to find the best route. + """ + if not isinstance(destination_ip, IPv4Address): + destination_ip = IPv4Address(destination_ip) + best_route = None + longest_prefix = -1 + lowest_metric = float("inf") # Initialise at infinity as any other number we compare to it will be smaller + + for route in self.routes: + route_network = IPv4Network(f"{route.address}/{route.subnet_mask}", strict=False) + prefix_len = route_network.prefixlen + + if destination_ip in route_network: + if prefix_len > longest_prefix or (prefix_len == longest_prefix and route.metric < lowest_metric): + best_route = route + longest_prefix = prefix_len + lowest_metric = route.metric + + return best_route + + def show(self): + """Prints a table of the routes in the RouteTable.""" + table = PrettyTable(["Index", "Address", "Next Hop", "Metric"]) + table.title = f"{self.sys_log.hostname} Route Table" + for index, route in enumerate(self.routes): + network = IPv4Network(f"{route.address}/{route.subnet_mask}") + table.add_row([index, f"{route.address}/{network.prefixlen}", route.next_hop, route.metric]) + print(table) + + +class RouterARPCache(ARPCache): + def __init__(self, sys_log: SysLog, router: Router): + super().__init__(sys_log) + self.router: Router = router + + def process_arp_packet(self, from_nic: NIC, frame: Frame): + """ + Overridden method to process a received ARP packet in a router-specific way. + + :param from_nic: The NIC that received the ARP packet. + :param frame: The original arp frame. + """ + arp_packet = frame.arp + + # ARP Reply + if not arp_packet.request: + for nic in self.router.nics.values(): + if arp_packet.target_ip == nic.ip_address: + # reply to the Router specifically + self.sys_log.info( + f"Received ARP response for {arp_packet.sender_ip} from {arp_packet.sender_mac_addr} via NIC {from_nic}" + ) + self.add_arp_cache_entry( + ip_address=arp_packet.sender_ip, + mac_address=arp_packet.sender_mac_addr, + nic=from_nic, + ) + return + + # Reply for a connected requested + nic = self.get_arp_cache_nic(arp_packet.target_ip) + if nic: + self.sys_log.info(f"Forwarding arp reply for {arp_packet.target_ip}, from {arp_packet.sender_ip}") + arp_packet.sender_mac_addr = nic.mac_address + frame.decrement_ttl() + nic.send_frame(frame) + + # ARP Request + self.sys_log.info( + f"Received ARP request for {arp_packet.target_ip} from " + f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip} " + ) + # Matched ARP request + self.add_arp_cache_entry(ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic) + arp_packet = arp_packet.generate_reply(from_nic.mac_address) + self.send_arp_reply(arp_packet, from_nic) + + # If the target IP matches one of the router's NICs + for nic in self.nics.values(): + if nic.enabled and nic.ip_address == arp_packet.target_ip: + arp_reply = arp_packet.generate_reply(from_nic.mac_address) + self.send_arp_reply(arp_reply, from_nic) + return + + +class RouterICMP(ICMP): + router: Router + + def __init__(self, sys_log: SysLog, arp_cache: ARPCache, router: Router): + super().__init__(sys_log, arp_cache) + self.router = router + + def process_icmp(self, frame: Frame, from_nic: NIC, is_reattempt: bool = False): + if frame.icmp.icmp_type == ICMPType.ECHO_REQUEST: + # determine if request is for router interface or whether it needs to be routed + + for nic in self.router.nics.values(): + if nic.ip_address == frame.ip.dst_ip and nic.enabled: + # reply to the request + self.sys_log.info(f"Received echo request from {frame.ip.src_ip}") + target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip) + src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip) + tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) + + # Network Layer + ip_packet = IPPacket(src_ip=nic.ip_address, dst_ip=frame.ip.src_ip, protocol=IPProtocol.ICMP) + # Data Link Layer + ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) + icmp_reply_packet = ICMPPacket( + icmp_type=ICMPType.ECHO_REPLY, + icmp_code=0, + identifier=frame.icmp.identifier, + sequence=frame.icmp.sequence + 1, + ) + frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet) + self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") + + src_nic.send_frame(frame) + return + + # Route the frame + self.router.route_frame(frame, from_nic) + elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY: + self.sys_log.info(f"Received echo reply from {frame.ip.src_ip}") + if not self.request_replies.get(frame.icmp.identifier): + self.request_replies[frame.icmp.identifier] = 0 + self.request_replies[frame.icmp.identifier] += 1 class Router(Node): num_ports: int ethernet_ports: Dict[int, NIC] = {} - acl: List = [] - route_table: Dict = {} + acl: AccessControlList + route_table: RouteTable + arp: RouterARPCache + icmp: RouterICMP def __init__(self, hostname: str, num_ports: int = 5, **kwargs): + if not kwargs.get("sys_log"): + kwargs["sys_log"] = SysLog(hostname) + if not kwargs.get("acl"): + kwargs["acl"] = AccessControlList(sys_log=kwargs["sys_log"], implicit_action=ACLAction.DENY) + if not kwargs.get("route_table"): + kwargs["route_table"] = RouteTable(sys_log=kwargs["sys_log"]) + if not kwargs.get("arp"): + kwargs["arp"] = RouterARPCache(sys_log=kwargs.get("sys_log"), router=self) + if not kwargs.get("icmp"): + 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") self.connect_nic(nic) self.ethernet_ports[i] = nic + self.arp.nics = self.nics + self.icmp.arp = self.arp + + def _get_port_of_nic(self, target_nic: NIC) -> Optional[int]: + for port, nic in self.ethernet_ports.items(): + if nic == target_nic: + return port + def describe_state(self) -> Dict: pass - def configure_port( - self, - port: int, - ip_address: Union[IPv4Address, str], - subnet_mask: str - ): + def route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None: + if not re_attempt: + # Check if src ip is on network of one of the NICs + nic = self.arp.get_arp_cache_nic(frame.ip.dst_ip) + target_mac = self.arp.get_arp_cache_mac_address(frame.ip.dst_ip) + if not nic: + self.arp.send_arp_request(frame.ip.dst_ip) + return self.route_frame(frame=frame, from_nic=from_nic, re_attempt=True) + for nic in self.nics.values(): + if nic.enabled and frame.ip.dst_ip in nic.ip_network: + from_port = self._get_port_of_nic(from_nic) + to_port = self._get_port_of_nic(nic) + self.sys_log.info(f"Routing frame to internally from port {from_port} to port {to_port}") + frame.decrement_ttl() + frame.ethernet.src_mac_addr = nic.mac_address + frame.ethernet.dst_mac_addr = target_mac + nic.send_frame(frame) + return + else: + self.sys_log.info(f"Destination {frame.ip.dst_ip} is unreachable") + + def receive_frame(self, frame: Frame, from_nic: NIC): + """ + Receive a Frame from the connected NIC and process it. + + Depending on the protocol, the frame is passed to the appropriate handler such as ARP or ICMP, or up to the + SessionManager if no code manager exists. + + :param frame: The Frame being received. + :param from_nic: The NIC that received the frame. + """ + route_frame = False + protocol = frame.ip.protocol + src_ip = frame.ip.src_ip + dst_ip = frame.ip.dst_ip + src_port = None + dst_port = None + if frame.ip.protocol == IPProtocol.TCP: + src_port = frame.tcp.src_port + dst_port = frame.tcp.dst_port + elif frame.ip.protocol == IPProtocol.UDP: + src_port = frame.udp.src_port + dst_port = frame.udp.dst_port + + # Check if it's permitted + permitted, rule = self.acl.is_permitted( + protocol=protocol, src_ip=src_ip, src_port=src_port, dst_ip=dst_ip, dst_port=dst_port + ) + if not permitted: + at_port = self._get_port_of_nic(from_nic) + self.sys_log.info(f"Frame blocked at port {at_port} by rule {rule}") + return + if not self.arp.get_arp_cache_nic(src_ip): + self.arp.add_arp_cache_entry(src_ip, frame.ethernet.src_mac_addr, from_nic) + if frame.ip.protocol == IPProtocol.ICMP: + self.icmp.process_icmp(frame=frame, from_nic=from_nic) + else: + if src_port == Port.ARP: + self.arp.process_arp_packet(from_nic=from_nic, frame=frame) + else: + # All other traffic + route_frame = True + if route_frame: + self.route_frame(frame, from_nic) + + def configure_port(self, port: int, ip_address: Union[IPv4Address, str], subnet_mask: Union[IPv4Address, str]): if not isinstance(ip_address, IPv4Address): ip_address = IPv4Address(ip_address) + if not isinstance(subnet_mask, IPv4Address): + subnet_mask = IPv4Address(subnet_mask) nic = self.ethernet_ports[port] nic.ip_address = ip_address nic.subnet_mask = subnet_mask - self.sys_log.info(f"Configured port {port} with {ip_address=} {subnet_mask=}") + self.sys_log.info(f"Configured port {port} with ip_address={ip_address}/{nic.ip_network.prefixlen}") def enable_port(self, port: int): nic = self.ethernet_ports.get(port) @@ -72,7 +523,7 @@ class Router(Node): def show(self): """Prints a table of the NICs on the Node.""" table = PrettyTable(["Port", "MAC Address", "Address", "Speed", "Status"]) - + table.title = f"{self.hostname} Ethernet Interfaces" for port, nic in self.ethernet_ports.items(): table.add_row( [ diff --git a/tests/integration_tests/network/test_frame_transmission.py b/tests/integration_tests/network/test_frame_transmission.py index d3d6541a..34b76060 100644 --- a/tests/integration_tests/network/test_frame_transmission.py +++ b/tests/integration_tests/network/test_frame_transmission.py @@ -3,14 +3,13 @@ from primaite.simulator.network.hardware.base import Link, NIC, Node, Switch def test_node_to_node_ping(): """Tests two Nodes are able to ping each other.""" - # TODO Add actual checks. Manual check performed for now. node_a = Node(hostname="node_a") - nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1") + nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0") node_a.connect_nic(nic_a) node_a.power_on() node_b = Node(hostname="node_b") - nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0", gateway="192.168.0.1") + nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0") node_b.connect_nic(nic_b) node_b.power_on() @@ -23,19 +22,19 @@ def test_multi_nic(): """Tests that Nodes with multiple NICs can ping each other and the data go across the correct links.""" # TODO Add actual checks. Manual check performed for now. node_a = Node(hostname="node_a") - nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1") + nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0") node_a.connect_nic(nic_a) node_a.power_on() node_b = Node(hostname="node_b") - nic_b1 = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0", gateway="192.168.0.1") - nic_b2 = NIC(ip_address="10.0.0.12", subnet_mask="255.0.0.0", gateway="10.0.0.1") + nic_b1 = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0") + nic_b2 = NIC(ip_address="10.0.0.12", subnet_mask="255.0.0.0") node_b.connect_nic(nic_b1) node_b.connect_nic(nic_b2) node_b.power_on() node_c = Node(hostname="node_c") - nic_c = NIC(ip_address="10.0.0.13", subnet_mask="255.0.0.0", gateway="10.0.0.1") + nic_c = NIC(ip_address="10.0.0.13", subnet_mask="255.0.0.0") node_c.connect_nic(nic_c) node_c.power_on() @@ -52,22 +51,22 @@ def test_switched_network(): """Tests a larges network of Nodes and Switches with one node pinging another.""" # TODO Add actual checks. Manual check performed for now. pc_a = Node(hostname="pc_a") - nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1") + nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0") pc_a.connect_nic(nic_a) pc_a.power_on() pc_b = Node(hostname="pc_b") - nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0", gateway="192.168.0.1") + nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0") pc_b.connect_nic(nic_b) pc_b.power_on() pc_c = Node(hostname="pc_c") - nic_c = NIC(ip_address="192.168.0.12", subnet_mask="255.255.255.0", gateway="192.168.0.1") + nic_c = NIC(ip_address="192.168.0.12", subnet_mask="255.255.255.0") pc_c.connect_nic(nic_c) pc_c.power_on() pc_d = Node(hostname="pc_d") - nic_d = NIC(ip_address="192.168.0.13", subnet_mask="255.255.255.0", gateway="192.168.0.1") + nic_d = NIC(ip_address="192.168.0.13", subnet_mask="255.255.255.0") pc_d.connect_nic(nic_d) pc_d.power_on() diff --git a/tests/integration_tests/network/test_link_connection.py b/tests/integration_tests/network/test_link_connection.py index e08e40b9..ef65f078 100644 --- a/tests/integration_tests/network/test_link_connection.py +++ b/tests/integration_tests/network/test_link_connection.py @@ -4,18 +4,17 @@ from primaite.simulator.network.hardware.base import Link, NIC, Node def test_link_up(): """Tests Nodes, NICs, and Links can all be connected and be in an enabled/up state.""" node_a = Node(hostname="node_a") - nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1") + nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0") node_a.connect_nic(nic_a) node_a.power_on() - assert nic_a.enabled node_b = Node(hostname="node_b") - nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0", gateway="192.168.0.1") + nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0") node_b.connect_nic(nic_b) node_b.power_on() - assert nic_b.enabled - link = Link(endpoint_a=nic_a, endpoint_b=nic_b) + assert nic_a.enabled + assert nic_b.enabled assert link.is_up diff --git a/tests/integration_tests/network/test_nic_link_connection.py b/tests/integration_tests/network/test_nic_link_connection.py index 52a0c735..f051d026 100644 --- a/tests/integration_tests/network/test_nic_link_connection.py +++ b/tests/integration_tests/network/test_nic_link_connection.py @@ -8,7 +8,6 @@ def test_link_fails_with_same_nic(): with pytest.raises(ValueError): nic_a = NIC( ip_address="192.168.1.2", - subnet_mask="255.255.255.0", - gateway="192.168.0.1", + subnet_mask="255.255.255.0" ) Link(endpoint_a=nic_a, endpoint_b=nic_a) diff --git a/tests/integration_tests/network/test_routing.py b/tests/integration_tests/network/test_routing.py index cca48c0d..cb420e22 100644 --- a/tests/integration_tests/network/test_routing.py +++ b/tests/integration_tests/network/test_routing.py @@ -1,27 +1,55 @@ -from primaite.simulator.network.hardware.base import Node, NIC, Link -from primaite.simulator.network.hardware.nodes.router import Router +from typing import Tuple + +import pytest + +from primaite.simulator.network.hardware.base import Link, NIC, Node +from primaite.simulator.network.hardware.nodes.router import ACLAction, Router +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port -def test_ping_fails_with_no_route(): - """Tests a larges network of Nodes and Switches with one node pinging another.""" - pc_a = Node(hostname="pc_a") - nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1") +@pytest.fixture(scope="function") +def pc_a_pc_b_router_1() -> Tuple[Node, Node, Router]: + pc_a = Node(hostname="pc_a", default_gateway="192.168.0.1") + nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0") pc_a.connect_nic(nic_a) pc_a.power_on() - pc_b = Node(hostname="pc_b") - nic_b = NIC(ip_address="192.168.1.10", subnet_mask="255.255.255.0", gateway="192.168.1.1") + pc_b = Node(hostname="pc_b", default_gateway="192.168.1.1") + nic_b = NIC(ip_address="192.168.1.10", subnet_mask="255.255.255.0") pc_b.connect_nic(nic_b) pc_b.power_on() router_1 = Router(hostname="router_1") + router_1.power_on() + router_1.configure_port(1, "192.168.0.1", "255.255.255.0") router_1.configure_port(2, "192.168.1.1", "255.255.255.0") - router_1.power_on() - router_1.show() + Link(endpoint_a=nic_a, endpoint_b=router_1.ethernet_ports[1]) + Link(endpoint_a=nic_b, endpoint_b=router_1.ethernet_ports[2]) + router_1.enable_port(1) + router_1.enable_port(2) - link_nic_a_router_1 = Link(endpoint_a=nic_a, endpoint_b=router_1.ethernet_ports[1]) - link_nic_b_router_1 = Link(endpoint_a=nic_b, endpoint_b=router_1.ethernet_ports[2]) - router_1.power_on() - #assert pc_a.ping("192.168.1.10") \ No newline at end of file + router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22) + + router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23) + return pc_a, pc_b, router_1 + + +def test_ping_default_gateway(pc_a_pc_b_router_1): + pc_a, pc_b, router_1 = pc_a_pc_b_router_1 + + assert pc_a.ping(pc_a.default_gateway) + + +def test_ping_other_router_port(pc_a_pc_b_router_1): + pc_a, pc_b, router_1 = pc_a_pc_b_router_1 + + assert pc_a.ping(pc_b.default_gateway) + + +def test_host_on_other_subnet(pc_a_pc_b_router_1): + pc_a, pc_b, router_1 = pc_a_pc_b_router_1 + + assert pc_a.ping("192.168.1.10") diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/__init__.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_router.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_router.py new file mode 100644 index 00000000..48d0fc06 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_router.py @@ -0,0 +1,104 @@ +from ipaddress import IPv4Address + +from primaite.simulator.network.hardware.nodes.router import AccessControlList, ACLAction, ACLRule +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port + + +def test_add_rule(): + acl = AccessControlList() + acl.add_rule( + action=ACLAction.PERMIT, + protocol=IPProtocol.TCP, + src_ip=IPv4Address("192.168.1.1"), + src_port=Port(8080), + dst_ip=IPv4Address("192.168.1.2"), + dst_port=Port(80), + position=1, + ) + assert acl.acl[1].action == ACLAction.PERMIT + assert acl.acl[1].protocol == IPProtocol.TCP + assert acl.acl[1].src_ip == IPv4Address("192.168.1.1") + assert acl.acl[1].src_port == Port(8080) + assert acl.acl[1].dst_ip == IPv4Address("192.168.1.2") + assert acl.acl[1].dst_port == Port(80) + + +def test_remove_rule(): + acl = AccessControlList() + acl.add_rule( + action=ACLAction.PERMIT, + protocol=IPProtocol.TCP, + src_ip=IPv4Address("192.168.1.1"), + src_port=Port(8080), + dst_ip=IPv4Address("192.168.1.2"), + dst_port=Port(80), + position=1, + ) + acl.remove_rule(1) + assert not acl.acl[1] + + +def test_rules(): + acl = AccessControlList() + acl.add_rule( + action=ACLAction.PERMIT, + protocol=IPProtocol.TCP, + src_ip=IPv4Address("192.168.1.1"), + src_port=Port(8080), + dst_ip=IPv4Address("192.168.1.2"), + dst_port=Port(80), + position=1, + ) + acl.add_rule( + action=ACLAction.DENY, + protocol=IPProtocol.TCP, + src_ip=IPv4Address("192.168.1.3"), + src_port=Port(8080), + dst_ip=IPv4Address("192.168.1.4"), + dst_port=Port(80), + position=2, + ) + assert acl.is_permitted( + protocol=IPProtocol.TCP, + src_ip=IPv4Address("192.168.1.1"), + src_port=Port(8080), + dst_ip=IPv4Address("192.168.1.2"), + dst_port=Port(80), + ) + assert not acl.is_permitted( + protocol=IPProtocol.TCP, + src_ip=IPv4Address("192.168.1.3"), + src_port=Port(8080), + dst_ip=IPv4Address("192.168.1.4"), + dst_port=Port(80), + ) + + +def test_default_rule(): + acl = AccessControlList() + acl.add_rule( + action=ACLAction.PERMIT, + protocol=IPProtocol.TCP, + src_ip=IPv4Address("192.168.1.1"), + src_port=Port(8080), + dst_ip=IPv4Address("192.168.1.2"), + dst_port=Port(80), + position=1, + ) + acl.add_rule( + action=ACLAction.DENY, + protocol=IPProtocol.TCP, + src_ip=IPv4Address("192.168.1.3"), + src_port=Port(8080), + dst_ip=IPv4Address("192.168.1.4"), + dst_port=Port(80), + position=2, + ) + assert not acl.is_permitted( + protocol=IPProtocol.UDP, + src_ip=IPv4Address("192.168.1.5"), + src_port=Port(8080), + dst_ip=IPv4Address("192.168.1.12"), + dst_port=Port(80), + ) diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py index dc508508..11873128 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py @@ -32,10 +32,8 @@ def test_nic_ip_address_type_conversion(): nic = NIC( ip_address="192.168.1.2", subnet_mask="255.255.255.0", - gateway="192.168.0.1", ) assert isinstance(nic.ip_address, IPv4Address) - assert isinstance(nic.gateway, IPv4Address) def test_nic_deserialize(): @@ -43,7 +41,6 @@ def test_nic_deserialize(): nic = NIC( ip_address="192.168.1.2", subnet_mask="255.255.255.0", - gateway="192.168.0.1", ) nic_json = nic.model_dump_json() @@ -51,21 +48,10 @@ def test_nic_deserialize(): assert nic == deserialized_nic -def test_nic_ip_address_as_gateway_fails(): - """Tests NIC creation fails if ip address is the same as the gateway.""" - with pytest.raises(ValueError): - NIC( - ip_address="192.168.0.1", - subnet_mask="255.255.255.0", - gateway="192.168.0.1", - ) - - def test_nic_ip_address_as_network_address_fails(): """Tests NIC creation fails if ip address and subnet mask are a network address.""" with pytest.raises(ValueError): NIC( ip_address="192.168.0.0", subnet_mask="255.255.255.0", - gateway="192.168.0.1", ) From e73d7f49d68e5e6b0e481db098d1e0aa49c044fd Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Thu, 31 Aug 2023 11:03:38 +0100 Subject: [PATCH 03/10] #1800 - Fixed routing and processing of ICMP packets in the Router class --- .../simulator/network/hardware/base.py | 9 +- .../network/hardware/nodes/router.py | 171 ++++++++++-------- 2 files changed, 98 insertions(+), 82 deletions(-) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 921ebbcd..4803150d 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -744,14 +744,12 @@ class ICMP: :param frame: The Frame containing the ICMP packet to process. """ if frame.icmp.icmp_type == ICMPType.ECHO_REQUEST: - self.sys_log.info(f"Received echo request from {frame.ip.src_ip}") + if not is_reattempt: + self.sys_log.info(f"Received echo request from {frame.ip.src_ip}") target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip) src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip) if not src_nic: - print(self.sys_log.hostname) - print(frame.ip.src_ip) - self.arp.show() self.arp.send_arp_request(frame.ip.src_ip) self.process_icmp(frame=frame, from_nic=from_nic, is_reattempt=True) return @@ -932,14 +930,13 @@ class Node(SimComponent): def show(self): """Prints a table of the NICs on the Node.""" - table = PrettyTable(["MAC Address", "Address", "Default Gateway", "Speed", "Status"]) + table = PrettyTable(["MAC Address", "Address", "Speed", "Status"]) table.title = f"{self.hostname} Network Interface Cards" for nic in self.nics.values(): table.add_row( [ nic.mac_address, f"{nic.ip_address}/{nic.ip_network.prefixlen}", - nic.gateway, nic.speed, "Enabled" if nic.enabled else "Disabled", ] diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 528e4a73..7db92938 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -69,14 +69,14 @@ class AccessControlList(SimComponent): return self._acl def add_rule( - self, - action: ACLAction, - protocol: Optional[IPProtocol] = None, - src_ip: Optional[Union[str, IPv4Address]] = None, - src_port: Optional[Port] = None, - dst_ip: Optional[Union[str, IPv4Address]] = None, - dst_port: Optional[Port] = None, - position: int = 0, + self, + action: ACLAction, + protocol: Optional[IPProtocol] = None, + src_ip: Optional[Union[str, IPv4Address]] = None, + src_port: Optional[Port] = None, + dst_ip: Optional[Union[str, IPv4Address]] = None, + dst_port: Optional[Port] = None, + position: int = 0, ) -> None: if isinstance(src_ip, str): src_ip = IPv4Address(src_ip) @@ -96,12 +96,12 @@ class AccessControlList(SimComponent): raise ValueError(f"Position {position} is out of bounds.") def is_permitted( - self, - protocol: IPProtocol, - src_ip: Union[str, IPv4Address], - src_port: Optional[Port], - dst_ip: Union[str, IPv4Address], - dst_port: Optional[Port], + self, + protocol: IPProtocol, + src_ip: Union[str, IPv4Address], + src_port: Optional[Port], + dst_ip: Union[str, IPv4Address], + dst_port: Optional[Port], ) -> Tuple[bool, Optional[Union[str, ACLRule]]]: if not isinstance(src_ip, IPv4Address): src_ip = IPv4Address(src_ip) @@ -112,23 +112,23 @@ class AccessControlList(SimComponent): continue if ( - (rule.src_ip == src_ip or rule.src_ip is None) - and (rule.dst_ip == dst_ip or rule.dst_ip 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 == src_ip or rule.src_ip is None) + and (rule.dst_ip == dst_ip or rule.dst_ip 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: Union[str, IPv4Address], - src_port: Port, - dst_ip: Union[str, IPv4Address], - dst_port: Port, + self, + protocol: IPProtocol, + src_ip: Union[str, IPv4Address], + src_port: Port, + dst_ip: Union[str, IPv4Address], + dst_port: Port, ) -> List[ACLRule]: if not isinstance(src_ip, IPv4Address): src_ip = IPv4Address(src_ip) @@ -140,11 +140,11 @@ class AccessControlList(SimComponent): continue if ( - (rule.src_ip == src_ip or rule.src_ip is None) - or (rule.dst_ip == dst_ip or rule.dst_ip 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 == src_ip or rule.src_ip is None) + or (rule.dst_ip == dst_ip or rule.dst_ip 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) @@ -247,11 +247,11 @@ class RouteTable(SimComponent): pass def add_route( - self, - address: Union[IPv4Address, str], - subnet_mask: Union[IPv4Address, str], - next_hop: Union[IPv4Address, str], - metric: float = 0.0, + self, + address: Union[IPv4Address, str], + subnet_mask: Union[IPv4Address, str], + next_hop: Union[IPv4Address, str], + metric: float = 0.0, ): """Add a route to the routing table. @@ -367,36 +367,46 @@ class RouterICMP(ICMP): # determine if request is for router interface or whether it needs to be routed for nic in self.router.nics.values(): - if nic.ip_address == frame.ip.dst_ip and nic.enabled: - # reply to the request - self.sys_log.info(f"Received echo request from {frame.ip.src_ip}") - target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip) - src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip) - tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) + if nic.ip_address == frame.ip.dst_ip: + if nic.enabled: + # reply to the request + if not is_reattempt: + self.sys_log.info(f"Received echo request from {frame.ip.src_ip}") + target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip) + src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip) + tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) - # Network Layer - ip_packet = IPPacket(src_ip=nic.ip_address, dst_ip=frame.ip.src_ip, protocol=IPProtocol.ICMP) - # Data Link Layer - ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) - icmp_reply_packet = ICMPPacket( - icmp_type=ICMPType.ECHO_REPLY, - icmp_code=0, - identifier=frame.icmp.identifier, - sequence=frame.icmp.sequence + 1, - ) - frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet) - self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") + # Network Layer + ip_packet = IPPacket(src_ip=nic.ip_address, dst_ip=frame.ip.src_ip, protocol=IPProtocol.ICMP) + # Data Link Layer + ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) + icmp_reply_packet = ICMPPacket( + icmp_type=ICMPType.ECHO_REPLY, + icmp_code=0, + identifier=frame.icmp.identifier, + sequence=frame.icmp.sequence + 1, + ) + frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet) + self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") - src_nic.send_frame(frame) + src_nic.send_frame(frame) return # Route the frame self.router.route_frame(frame, from_nic) + elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY: - self.sys_log.info(f"Received echo reply from {frame.ip.src_ip}") - if not self.request_replies.get(frame.icmp.identifier): - self.request_replies[frame.icmp.identifier] = 0 - self.request_replies[frame.icmp.identifier] += 1 + for nic in self.router.nics.values(): + if nic.ip_address == frame.ip.dst_ip: + if nic.enabled: + self.sys_log.info(f"Received echo reply from {frame.ip.src_ip}") + if not self.request_replies.get(frame.icmp.identifier): + self.request_replies[frame.icmp.identifier] = 0 + self.request_replies[frame.icmp.identifier] += 1 + + return + # Route the frame + self.router.route_frame(frame, from_nic) class Router(Node): @@ -436,25 +446,34 @@ class Router(Node): pass def route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None: - if not re_attempt: - # Check if src ip is on network of one of the NICs - nic = self.arp.get_arp_cache_nic(frame.ip.dst_ip) - target_mac = self.arp.get_arp_cache_mac_address(frame.ip.dst_ip) - if not nic: - self.arp.send_arp_request(frame.ip.dst_ip) - return self.route_frame(frame=frame, from_nic=from_nic, re_attempt=True) - for nic in self.nics.values(): - if nic.enabled and frame.ip.dst_ip in nic.ip_network: - from_port = self._get_port_of_nic(from_nic) - to_port = self._get_port_of_nic(nic) - self.sys_log.info(f"Routing frame to internally from port {from_port} to port {to_port}") - frame.decrement_ttl() - frame.ethernet.src_mac_addr = nic.mac_address - frame.ethernet.dst_mac_addr = target_mac - nic.send_frame(frame) - return - else: + # Check if src ip is on network of one of the NICs + nic = self.arp.get_arp_cache_nic(frame.ip.dst_ip) + target_mac = self.arp.get_arp_cache_mac_address(frame.ip.dst_ip) + + if re_attempt and not nic: self.sys_log.info(f"Destination {frame.ip.dst_ip} is unreachable") + return + + if not nic: + self.arp.send_arp_request(frame.ip.dst_ip) + return self.route_frame(frame=frame, from_nic=from_nic, re_attempt=True) + + if not nic.enabled: + # TODO: Add sys_log here + return + + if frame.ip.dst_ip in nic.ip_network: + from_port = self._get_port_of_nic(from_nic) + to_port = self._get_port_of_nic(nic) + self.sys_log.info(f"Routing frame to internally from port {from_port} to port {to_port}") + frame.decrement_ttl() + frame.ethernet.src_mac_addr = nic.mac_address + frame.ethernet.dst_mac_addr = target_mac + nic.send_frame(frame) + return + else: + pass + # TODO: Deal with routing from route tables def receive_frame(self, frame: Frame, from_nic: NIC): """ From 89ad22acebbd39cc6bf7a30808b104f46bd1152b Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Thu, 31 Aug 2023 13:35:56 +0100 Subject: [PATCH 04/10] #1800 - Synced with dev. - Added the UC2 network. - Added a Computer class. --- src/primaite/simulator/network/container.py | 29 +++- .../network/hardware/nodes/computer.py | 44 +++++ src/primaite/simulator/network/networks.py | 154 ++++++++++++++++++ 3 files changed, 221 insertions(+), 6 deletions(-) create mode 100644 src/primaite/simulator/network/hardware/nodes/computer.py create mode 100644 src/primaite/simulator/network/networks.py diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 85676034..ac502d84 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Union +from typing import Any, Dict, Union, Optional from primaite import getLogger from primaite.simulator.core import Action, ActionManager, AllowAllValidator, SimComponent @@ -58,6 +58,19 @@ class Network(SimComponent): node.parent = self _LOGGER.info(f"Added node {node.uuid} to Network {self.uuid}") + def get_node_by_hostname(self, hostname: str) -> Optional[Node]: + """ + Get a Node from the Network by its hostname. + + .. note:: Assumes hostnames on the network are unique. + + :param hostname: The Node hostname. + :return: The Node if it exists in the network. + """ + for node in self.nodes.values(): + if node.hostname == hostname: + return node + def remove_node(self, node: Node) -> None: """ Remove a node from the network. @@ -72,7 +85,8 @@ class Network(SimComponent): node.parent = None _LOGGER.info(f"Removed node {node.uuid} from network {self.uuid}") - def connect(self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs) -> None: + def connect(self, endpoint_a: Union[Node, NIC, SwitchPort], endpoint_b: Union[Node, NIC, SwitchPort], **kwargs) -> \ + None: """Connect two nodes on the network by creating a link between an NIC/SwitchPort of each one. :param endpoint_a: The endpoint to which to connect the link on the first node @@ -81,16 +95,19 @@ class Network(SimComponent): :type endpoint_b: Union[NIC, SwitchPort] :raises RuntimeError: _description_ """ - node_a = endpoint_a.parent - node_b = endpoint_b.parent + node_a: Node = endpoint_a.parent if not isinstance(endpoint_a, Node) else endpoint_a + node_b: Node = endpoint_b.parent if not isinstance(endpoint_b, Node) else endpoint_b if node_a not in self: self.add_node(node_a) if node_b not in self: self.add_node(node_b) if node_a is node_b: - _LOGGER.warn(f"Cannot link endpoint {endpoint_a} to {endpoint_b} because they belong to the same node.") + _LOGGER.warning(f"Cannot link endpoint {endpoint_a} to {endpoint_b} because they belong to the same node.") return - + if isinstance(endpoint_a, Node) and len(endpoint_a.nics) == 1: + endpoint_a = list(endpoint_a.nics.values())[0] + if isinstance(endpoint_b, Node) and len(endpoint_b.nics) == 1: + endpoint_b = list(endpoint_b.nics.values())[0] link = Link(endpoint_a=endpoint_a, endpoint_b=endpoint_b, **kwargs) self.links[link.uuid] = link link.parent = self diff --git a/src/primaite/simulator/network/hardware/nodes/computer.py b/src/primaite/simulator/network/hardware/nodes/computer.py new file mode 100644 index 00000000..8dfb7540 --- /dev/null +++ b/src/primaite/simulator/network/hardware/nodes/computer.py @@ -0,0 +1,44 @@ +from ipaddress import IPv4Address + +from primaite.simulator.network.hardware.base import Node, NIC + + +class Computer(Node): + """ + A basic computer class. + + Example: + >>> pc_a = Computer( + 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() + + Instances of computer come 'pre-packaged' with the following: + + * Core Functionality: + * ARP. + * ICMP. + * Packet Capture. + * Sys Log. + * Services: + * DNS Client. + * FTP Client. + * LDAP Client. + * NTP Client. + * Applications: + * Email Client. + * Web Browser. + * Processes: + * Placeholder. + """ + + def __init__(self, **kwargs): + for key in {"ip_address", "subnet_mask", "default_gateway"}: + if key in kwargs: + if not isinstance(kwargs[key], IPv4Address): + kwargs[key] = IPv4Address(kwargs[key]) + super().__init__(**kwargs) + self.connect_nic(NIC(ip_address=kwargs["ip_address"], subnet_mask=kwargs["subnet_mask"])) diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py new file mode 100644 index 00000000..0eccefa4 --- /dev/null +++ b/src/primaite/simulator/network/networks.py @@ -0,0 +1,154 @@ +from primaite.simulator.network.container import Network +from primaite.simulator.network.hardware.base import Switch, NIC +from primaite.simulator.network.hardware.nodes.computer import Computer +from primaite.simulator.network.hardware.nodes.router import Router, ACLAction +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port + + +def arcd_uc2_network() -> Network: + """ + Models the ARCD Use Case 2 Network. + + +------------+ + | domain_ | + +------------+ controller | + | | | + | +------------+ + | + | + +------------+ | +------------+ + | | | | | + | client_1 +---------+ | +---------+ web_server | + | | | | | | | + +------------+ | | | +------------+ + +--+---------+ +------------+ +------+--+--+ + | | | | | | + | switch_2 +------+ router_1 +------+ switch_1 | + | | | | | | + +--+------+--+ +------------+ +--+---+--+--+ + +------------+ | | | | | +------------+ + | | | | | | | | database | + | client_2 +---------+ | | | +---------+ _server | + | | | | | | | + +------------+ | | | +------------+ + | +------------+ | | + | | security | | | + +---------+ _suite +---------+ | +------------+ + | | | | backup_ | + +------------+ +------------+ server | + | | + +------------+ + + Example: + >>> network = arcd_uc2_network() + >>> network.get_node_by_hostname("client_1").ping("192.168.1.10") + + """ + network = Network() + + # Router 1 + router_1 = Router(hostname="router_1", num_ports=5) + 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) + switch_1.power_on() + network.connect(endpoint_a=router_1.ethernet_ports[1], endpoint_b=switch_1.switch_ports[8]) + router_1.enable_port(1) + + # Switch 2 + switch_2 = Switch(hostname="switch_2", num_ports=8) + switch_2.power_on() + network.connect(endpoint_a=router_1.ethernet_ports[2], endpoint_b=switch_2.switch_ports[8]) + router_1.enable_port(2) + + # Client 1 + client_1 = Computer( + hostname="client_1", + ip_address="192.168.10.21", + subnet_mask="255.255.255.0", + default_gateway="192.168.10.1" + ) + client_1.power_on() + network.connect(endpoint_a=client_1, endpoint_b=switch_2.switch_ports[1]) + + # Client 2 + client_2 = Computer( + hostname="client_2", + ip_address="192.168.10.22", + subnet_mask="255.255.255.0", + default_gateway="192.168.10.1" + ) + client_2.power_on() + network.connect(endpoint_a=client_2, endpoint_b=switch_2.switch_ports[2]) + + # Domain Controller + domain_controller = Computer( + hostname="domain_controller", + ip_address="192.168.1.10", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1" + ) + domain_controller.power_on() + network.connect(endpoint_a=domain_controller, endpoint_b=switch_1.switch_ports[1]) + + # Web Server + web_server = Computer( + hostname="web_server", + ip_address="192.168.1.12", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1" + ) + web_server.power_on() + network.connect(endpoint_a=web_server, endpoint_b=switch_1.switch_ports[2]) + + # Database Server + database_server = Computer( + hostname="database_server", + ip_address="192.168.1.14", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1" + ) + database_server.power_on() + network.connect(endpoint_a=database_server, endpoint_b=switch_1.switch_ports[3]) + + # Backup Server + backup_server = Computer( + hostname="backup_server", + ip_address="192.168.1.16", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1" + ) + backup_server.power_on() + network.connect(endpoint_a=backup_server, endpoint_b=switch_1.switch_ports[4]) + + # Security Suite + security_suite = Computer( + hostname="security_suite", + ip_address="192.168.1.110", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1" + ) + security_suite.power_on() + network.connect(endpoint_a=security_suite, endpoint_b=switch_1.switch_ports[7]) + security_suite_external_nic = NIC(ip_address="192.168.10.110", subnet_mask="255.255.255.0") + security_suite.connect_nic(security_suite_external_nic) + network.connect(endpoint_a=security_suite_external_nic, endpoint_b=switch_2.switch_ports[7]) + + router_1.acl.add_rule( + action=ACLAction.PERMIT, + src_port=Port.ARP, + dst_port=Port.ARP, + position=22 + ) + + router_1.acl.add_rule( + action=ACLAction.PERMIT, + protocol=IPProtocol.ICMP, + position=23 + ) + + return network From 5111affeebbc8a75becbfa2c31b34eed4a8a9ebc Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Fri, 1 Sep 2023 16:58:21 +0100 Subject: [PATCH 05/10] #1800 - Added more docstrings and rst docs. - Extended the .show functionality to enable markdown format too. --- docs/source/simulation.rst | 3 + .../simulation_components/network/network.rst | 114 +++++++++ .../simulation_components/network/router.rst | 73 ++++++ .../simulation_components/network/switch.rst | 8 + src/primaite/simulator/network/container.py | 175 +++++++++++-- .../simulator/network/hardware/base.py | 62 +++-- .../network/hardware/nodes/computer.py | 24 +- .../network/hardware/nodes/router.py | 240 +++++++++++++++--- .../network/hardware/nodes/server.py | 37 +++ src/primaite/simulator/network/networks.py | 103 ++++++-- .../network/transmission/data_link_layer.py | 5 + src/primaite/simulator/system/core/sys_log.py | 17 +- 12 files changed, 753 insertions(+), 108 deletions(-) create mode 100644 docs/source/simulation_components/network/network.rst create mode 100644 docs/source/simulation_components/network/router.rst create mode 100644 docs/source/simulation_components/network/switch.rst create mode 100644 src/primaite/simulator/network/hardware/nodes/server.py diff --git a/docs/source/simulation.rst b/docs/source/simulation.rst index a2784628..7e9fe77f 100644 --- a/docs/source/simulation.rst +++ b/docs/source/simulation.rst @@ -18,3 +18,6 @@ Contents simulation_structure simulation_components/network/base_hardware simulation_components/network/transport_to_data_link_layer + simulation_components/network/router + simulation_components/network/switch + simulation_components/network/network diff --git a/docs/source/simulation_components/network/network.rst b/docs/source/simulation_components/network/network.rst new file mode 100644 index 00000000..e5614980 --- /dev/null +++ b/docs/source/simulation_components/network/network.rst @@ -0,0 +1,114 @@ +.. only:: comment + + © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK + +.. _about: + +Network +======= + +The ``Network`` class serves as the backbone of the simulation. It offers a framework to manage various network +components such as routers, switches, servers, and clients. This document provides a detailed explanation of how to +effectively use the ``Network`` class. + +Example Usage +------------- + +Below demonstrates how to use the Router node to connect Nodes, and block traffic using ACLs. For this demonstration, +we'll use the following Network that has a client, server, two switches, and a router. + +.. code-block:: text + + +------------+ +------------+ +------------+ +------------+ +------------+ + | | | | | | | | | | + | client_1 +------+ switch_2 +------+ router_1 +------+ switch_1 +------+ server_1 | + | | | | | | | | | | + +------------+ +------------+ +------------+ +------------+ +------------+ + +1. Relevant imports + +.. code-block:: python + + from primaite.simulator.network.container import Network + from primaite.simulator.network.hardware.base import Switch, NIC + from primaite.simulator.network.hardware.nodes.computer import Computer + from primaite.simulator.network.hardware.nodes.router import Router, ACLAction + from primaite.simulator.network.hardware.nodes.server import Server + from primaite.simulator.network.transmission.network_layer import IPProtocol + from primaite.simulator.network.transmission.transport_layer import Port + +2. Create the Network + +.. code-block:: python + + network = Network() + +3. Create and configure the Router + +.. code-block:: python + + router_1 = Router(hostname="router_1", num_ports=3) + 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.2.1", subnet_mask="255.255.255.0") + +4. Create and configure the two Switches + +.. code-block:: python + + switch_1 = Switch(hostname="switch_1", num_ports=6) + switch_1.power_on() + switch_2 = Switch(hostname="switch_2", num_ports=6) + switch_2.power_on() + +5. Connect the Switches to the Router + +.. code-block:: python + + network.connect(endpoint_a=router_1.ethernet_ports[1], endpoint_b=switch_1.switch_ports[6]) + router_1.enable_port(1) + network.connect(endpoint_a=router_1.ethernet_ports[2], endpoint_b=switch_2.switch_ports[6]) + router_1.enable_port(2) + +6. Create the Client and Server nodes. + +.. code-block:: python + + client_1 = Computer( + hostname="client_1", + ip_address="192.168.2.2", + subnet_mask="255.255.255.0", + default_gateway="192.168.2.1" + ) + client_1.power_on() + server_1 = Server( + hostname="server_1", + ip_address="192.168.1.2", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1" + ) + server_1.power_on() + +7. Connect the Client and Server to the relevant Switch + +.. code-block:: python + + network.connect(endpoint_a=switch_2.switch_ports[1], endpoint_b=client_1.ethernet_port[1]) + network.connect(endpoint_a=switch_1.switch_ports[1], endpoint_b=server_1.ethernet_port[1]) + +8. Add ACL rules on the Router to allow ARP and ICMP traffic. + +.. code-block:: python + + router_1.acl.add_rule( + action=ACLAction.PERMIT, + src_port=Port.ARP, + dst_port=Port.ARP, + position=22 + ) + + router_1.acl.add_rule( + action=ACLAction.PERMIT, + protocol=IPProtocol.ICMP, + position=23 + ) diff --git a/docs/source/simulation_components/network/router.rst b/docs/source/simulation_components/network/router.rst new file mode 100644 index 00000000..aaa589cc --- /dev/null +++ b/docs/source/simulation_components/network/router.rst @@ -0,0 +1,73 @@ +.. only:: comment + + © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK + +.. _about: + +Router Module +============= + +The router module contains classes for simulating the functions of a network router. + +Router +------ + +The Router class represents a multi-port network router that can receive, process, and route network packets between its ports and other Nodes + +The router maintains internal state including: + +- RouteTable - Routing table to lookup where to forward packets. +- AccessControlList - Access control rules to block or allow packets. +- ARP cache - MAC address lookups for connected devices. +- ICMP handler - Handles ICMP requests to router interfaces. + +The router receives incoming frames on enabled ports. It processes the frame headers and applies the following logic: + +1. Checks the AccessControlList if the packet is permitted. If blocked, it is dropped. +2. For permitted packets, routes the frame based on: + - ARP cache lookups for destination MAC address. + - ICMP echo requests handled directly. + - RouteTable lookup to forward packet out other ports. +3. Updates ARP cache as it learns new information about the Network. + + + +RouteTable +---------- + +The RouteTable holds RouteEntry objects representing routes. It finds the best route for a destination IP using a metric and the longest prefix match algorithm. + +Routes can be added and looked up based on destination IP address. The RouteTable is used by the Router when forwarding packets between other Nodes. + +AccessControlList +----------------- + +The AccessControlList defines Access Control Rules to block or allow packets. Packets are checked against the rules to determine if they are permitted to be forwarded. + +Rules can be added to deny or permit traffic based on IP, port, and protocol. The ACL is checked by the Router when packets are received. + +Packet Processing +----------------- + +-The Router supports the following protocols and packet types: + +ARP +^^^ + +- Handles both ARP requests and responses. +- Updates ARP cache. +- Proxies ARP replies for connected networks. +- Routes ARP requests. + +ICMP +^^^^ + +- Responds to ICMP echo requests to Router interfaces. +- Routes other ICMP messages based on routes. + +TCP/UDP +^^^^^^^ + +- Forwards packets based on routes like IP. +- Applies ACL rules based on protocol, source/destination IP address, and source/destination port numbers. +- Decrements TTL and drops expired TTL packets. diff --git a/docs/source/simulation_components/network/switch.rst b/docs/source/simulation_components/network/switch.rst new file mode 100644 index 00000000..4b3b24bc --- /dev/null +++ b/docs/source/simulation_components/network/switch.rst @@ -0,0 +1,8 @@ +.. only:: comment + + © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK + +.. _about: + +Switch +====== \ No newline at end of file diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index ac502d84..ccb9ce77 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -1,20 +1,41 @@ -from typing import Any, Dict, Union, Optional +from typing import Any, Dict, Union, Optional, List + +import matplotlib.pyplot as plt +import networkx as nx +from networkx import MultiGraph +from prettytable import PrettyTable, MARKDOWN from primaite import getLogger from primaite.simulator.core import Action, ActionManager, AllowAllValidator, SimComponent -from primaite.simulator.network.hardware.base import Link, NIC, Node, SwitchPort +from primaite.simulator.network.hardware.base import Link, NIC, Node, SwitchPort, Switch +from primaite.simulator.network.hardware.nodes.computer import Computer +from primaite.simulator.network.hardware.nodes.router import Router +from primaite.simulator.network.hardware.nodes.server import Server _LOGGER = getLogger(__name__) class Network(SimComponent): - """Top level container object representing the physical network.""" + """ + Top level container object representing the physical network. + + This class manages nodes, links, and other network components. It also + offers methods for rendering the network topology and gathering states. + + :ivar Dict[str, Node] nodes: Dictionary mapping node UUIDs to Node instances. + :ivar Dict[str, Link] links: Dictionary mapping link UUIDs to Link instances. + """ nodes: Dict[str, Node] = {} links: Dict[str, Link] = {} def __init__(self, **kwargs): - """Initialise the network.""" + """" + Initialise the network. + + Constructs the network and sets up its initial state including + the action manager and an empty MultiGraph for topology representation. + """ super().__init__(**kwargs) self.action_manager = ActionManager() @@ -25,15 +46,112 @@ class Network(SimComponent): validator=AllowAllValidator(), ), ) + self._nx_graph = MultiGraph() + + @property + def routers(self) -> List[Router]: + """The Routers in the Network.""" + return [node for node in self.nodes.values() if isinstance(node, Router)] + + @property + def switches(self) -> List[Switch]: + """The Switches in the Network.""" + return [node for node in self.nodes.values() if isinstance(node, Switch)] + + @property + def computers(self) -> List[Computer]: + """The Computers in the Network.""" + return [node for node in self.nodes.values() if isinstance(node, Computer) and not isinstance(node, Server)] + + @property + def servers(self) -> List[Server]: + """The Servers in the Network.""" + return [node for node in self.nodes.values() if isinstance(node, Server)] + + def show(self, nodes: bool = True, ip_addresses: bool = True, links: bool = True, markdown: bool = False): + """ + Print tables describing the Network. + + Generate and print PrettyTable instances that show details about nodes, + IP addresses, and links in the network. Output can be in Markdown format. + + :param nodes: Include node details in the output. Defaults to True. + :param ip_addresses: Include IP address details in the output. Defaults to True. + :param links: Include link details in the output. Defaults to True. + :param markdown: Use Markdown style in table output. Defaults to False. + """ + nodes_type_map = { + "Router": self.routers, + "Switch": self.switches, + "Server": self.servers, + "Computer": self.computers + } + if nodes: + table = PrettyTable(["Node", "Type", "Operating State"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" + table.title = f"Nodes" + for node_type, nodes in nodes_type_map.items(): + for node in nodes: + table.add_row([node.hostname, node_type, node.operating_state.name]) + print(table) + + if ip_addresses: + table = PrettyTable(["Node", "Port", "IP Address", "Subnet Mask", "Default Gateway"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" + table.title = f"IP Addresses" + for nodes in nodes_type_map.values(): + for node in nodes: + for i, port in node.ethernet_port.items(): + table.add_row([node.hostname, i, port.ip_address, port.subnet_mask, node.default_gateway]) + print(table) + + if links: + table = PrettyTable(["Endpoint A", "Endpoint B", "is Up", "Bandwidth (MBits)", "Current Load"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" + table.title = f"Links" + links = list(self.links.values()) + for nodes in nodes_type_map.values(): + for node in nodes: + for link in links[::-1]: + if node in [link.endpoint_a.parent, link.endpoint_b.parent]: + table.add_row( + [ + link.endpoint_a.parent.hostname, + link.endpoint_b.parent.hostname, + link.is_up, + link.bandwidth, + link.current_load_percent + ] + ) + links.remove(link) + print(table) + + def clear_links(self): + """Clear all the links in the network by resetting their component state for the episode.""" + for link in self.links.values(): + link.reset_component_for_episode() + + def draw(self, seed: int = 123): + """ + Draw the Network using NetworkX and matplotlib.pyplot. + + :param seed: An integer seed for reproducible layouts. Default is 123. + """ + pos = nx.spring_layout(self._nx_graph, seed=seed) + nx.draw(self._nx_graph, pos, with_labels=True) + plt.show() def describe_state(self) -> Dict: """ - Produce a dictionary describing the current state of this object. + Produce a dictionary describing the current state of the Network. - Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation. - - :return: Current state of this object and child objects. - :rtype: Dict + :return: A dictionary capturing the current state of the Network and its child objects. """ state = super().describe_state() state.update( @@ -48,14 +166,16 @@ class Network(SimComponent): """ Add an existing node to the network. - :param node: Node instance that the network should keep track of. - :type node: Node + .. note:: If the node is already present in the network, a warning is logged. + + :param node: Node instance that should be kept track of by the network. """ if node in self: _LOGGER.warning(f"Can't add node {node.uuid}. It is already in the network.") return self.nodes[node.uuid] = node node.parent = self + self._nx_graph.add_node(node.hostname) _LOGGER.info(f"Added node {node.uuid} to Network {self.uuid}") def get_node_by_hostname(self, hostname: str) -> Optional[Node]: @@ -75,6 +195,8 @@ class Network(SimComponent): """ Remove a node from the network. + .. note:: If the node is not found in the network, a warning is logged. + :param node: Node instance that is currently part of the network that should be removed. :type node: Node """ @@ -85,18 +207,22 @@ class Network(SimComponent): node.parent = None _LOGGER.info(f"Removed node {node.uuid} from network {self.uuid}") - def connect(self, endpoint_a: Union[Node, NIC, SwitchPort], endpoint_b: Union[Node, NIC, SwitchPort], **kwargs) -> \ - None: - """Connect two nodes on the network by creating a link between an NIC/SwitchPort of each one. - - :param endpoint_a: The endpoint to which to connect the link on the first node - :type endpoint_a: Union[NIC, SwitchPort] - :param endpoint_b: The endpoint to which to connct the link on the second node - :type endpoint_b: Union[NIC, SwitchPort] - :raises RuntimeError: _description_ + def connect( + self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs + ) -> None: """ - node_a: Node = endpoint_a.parent if not isinstance(endpoint_a, Node) else endpoint_a - node_b: Node = endpoint_b.parent if not isinstance(endpoint_b, Node) else endpoint_b + Connect two endpoints on the network by creating a link between their NICs/SwitchPorts. + + .. note:: If the nodes owning the endpoints are not already in the network, they are automatically added. + + :param endpoint_a: The first endpoint to connect. + :type endpoint_a: Union[NIC, SwitchPort] + :param endpoint_b: The second endpoint to connect. + :type endpoint_b: Union[NIC, SwitchPort] + :raises RuntimeError: If any validation or runtime checks fail. + """ + node_a: Node = endpoint_a.parent + node_b: Node = endpoint_b.parent if node_a not in self: self.add_node(node_a) if node_b not in self: @@ -104,12 +230,9 @@ class Network(SimComponent): if node_a is node_b: _LOGGER.warning(f"Cannot link endpoint {endpoint_a} to {endpoint_b} because they belong to the same node.") return - if isinstance(endpoint_a, Node) and len(endpoint_a.nics) == 1: - endpoint_a = list(endpoint_a.nics.values())[0] - if isinstance(endpoint_b, Node) and len(endpoint_b.nics) == 1: - endpoint_b = list(endpoint_b.nics.values())[0] link = Link(endpoint_a=endpoint_a, endpoint_b=endpoint_b, **kwargs) self.links[link.uuid] = link + self._nx_graph.add_edge(endpoint_a.parent.hostname, endpoint_b.parent.hostname) link.parent = self _LOGGER.info(f"Added link {link.uuid} to connect {endpoint_a} and {endpoint_b}") diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 9834d439..674020ee 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1,12 +1,13 @@ from __future__ import annotations +import random import re import secrets from enum import Enum from ipaddress import IPv4Address, IPv4Network from typing import Dict, List, Optional, Tuple, Union -from prettytable import PrettyTable +from prettytable import PrettyTable, MARKDOWN from primaite import getLogger from primaite.exceptions import NetworkError @@ -256,7 +257,6 @@ class NIC(SimComponent): The Frame is passed to the Node. :param frame: The network frame being received. - :type frame: :class:`~primaite.simulator.network.osi_layers.Frame` """ if self.enabled: frame.decrement_ttl() @@ -266,9 +266,6 @@ class NIC(SimComponent): 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 - else: - self.connected_node.sys_log.info("Dropping frame not for me") - print(frame) return False def __str__(self) -> str: @@ -562,9 +559,12 @@ class ARPCache: self.arp: Dict[IPv4Address, ARPEntry] = {} self.nics: Dict[str, "NIC"] = {} - def show(self): + def show(self, markdown: bool = False): """Prints a table of ARC Cache.""" table = PrettyTable(["IP Address", "MAC Address", "Via"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" table.title = f"{self.sys_log.hostname} ARP Cache" for ip, arp in self.arp.items(): table.add_row( @@ -765,12 +765,22 @@ class ICMP: identifier=frame.icmp.identifier, sequence=frame.icmp.sequence + 1, ) - frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet) + payload = secrets.token_urlsafe(int(32/1.3)) # Standard ICMP 32 bytes size + frame = Frame( + ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet, payload=payload + ) self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") src_nic.send_frame(frame) elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY: - self.sys_log.info(f"Received echo reply from {frame.ip.src_ip}") + time = frame.transmission_duration() + time_str = f"{time}ms" if time > 0 else "<1ms" + self.sys_log.info( + f"Reply from {frame.ip.src_ip}: " + f"bytes={len(frame.payload)}, " + f"time={time_str}, " + f"TTL={frame.ip.ttl}" + ) if not self.request_replies.get(frame.icmp.identifier): self.request_replies[frame.icmp.identifier] = 0 self.request_replies[frame.icmp.identifier] += 1 @@ -819,8 +829,8 @@ class ICMP: # Data Link Layer ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) icmp_packet = ICMPPacket(identifier=identifier, sequence=sequence) - frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_packet) - self.sys_log.info(f"Sending echo request to {target_ip_address}") + payload = secrets.token_urlsafe(int(32/1.3)) # Standard ICMP 32 bytes size + frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_packet, payload=payload) nic.send_frame(frame) return sequence, icmp_packet.identifier @@ -857,6 +867,8 @@ class Node(SimComponent): "The hardware state of the node." nics: Dict[str, NIC] = {} "The NICs on the node." + ethernet_port: Dict[int, NIC] = {} + "The NICs on the node by port id." accounts: Dict[str, Account] = {} "All accounts on the node." @@ -928,13 +940,17 @@ class Node(SimComponent): ) return state - def show(self): + def show(self, markdown: bool = False): """Prints a table of the NICs on the Node.""" - table = PrettyTable(["MAC Address", "Address", "Speed", "Status"]) + table = PrettyTable(["Port", "MAC Address", "Address", "Speed", "Status"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" table.title = f"{self.hostname} Network Interface Cards" - for nic in self.nics.values(): + for port, nic in self.ethernet_port.items(): table.add_row( [ + port, nic.mac_address, f"{nic.ip_address}/{nic.ip_network.prefixlen}", nic.speed, @@ -969,6 +985,7 @@ class Node(SimComponent): """ if nic.uuid not in self.nics: self.nics[nic.uuid] = nic + self.ethernet_port[len(self.nics)] = nic nic.connected_node = self nic.parent = self self.sys_log.info(f"Connected NIC {nic}") @@ -990,6 +1007,10 @@ class Node(SimComponent): if isinstance(nic, str): nic = self.nics.get(nic) if nic or nic.uuid in self.nics: + for port, _nic in self.ethernet_port.items(): + if nic == _nic: + self.ethernet_port.pop(port) + break self.nics.pop(nic.uuid) nic.parent = None nic.disable() @@ -1014,7 +1035,7 @@ class Node(SimComponent): self.sys_log.info("Pinging loopback address") return any(nic.enabled for nic in self.nics.values()) if self.operating_state == NodeOperatingState.ON: - self.sys_log.info(f"Attempting to ping {target_ip_address}") + self.sys_log.info(f"Pinging {target_ip_address}:") sequence, identifier = 0, None while sequence < pings: sequence, identifier = self.icmp.ping(target_ip_address, sequence, identifier, pings) @@ -1022,8 +1043,14 @@ class Node(SimComponent): passed = request_replies == pings if request_replies: self.icmp.request_replies.pop(identifier) + else: + request_replies = 0 + self.sys_log.info( + f"Ping statistics for {target_ip_address}: " + f"Packets: Sent = {pings}, " + f"Received = {request_replies}, " + f"Lost = {pings-request_replies} ({(pings-request_replies)/pings*100}% loss)") return passed - self.sys_log.info("Ping failed as the node is turned off") return False def send_frame(self, frame: Frame): @@ -1078,9 +1105,12 @@ class Switch(Node): port.parent = self port.port_num = port_num - def show(self): + def show(self, markdown: bool = False): """Prints a table of the SwitchPorts on the Switch.""" table = PrettyTable(["Port", "MAC Address", "Speed", "Status"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" table.title = f"{self.hostname} Switch Ports" for port_num, port in self.switch_ports.items(): table.add_row([port_num, port.mac_address, port.speed, "Enabled" if port.enabled else "Disabled"]) diff --git a/src/primaite/simulator/network/hardware/nodes/computer.py b/src/primaite/simulator/network/hardware/nodes/computer.py index 8dfb7540..110ad385 100644 --- a/src/primaite/simulator/network/hardware/nodes/computer.py +++ b/src/primaite/simulator/network/hardware/nodes/computer.py @@ -5,7 +5,7 @@ from primaite.simulator.network.hardware.base import Node, NIC class Computer(Node): """ - A basic computer class. + A basic Computer class. Example: >>> pc_a = Computer( @@ -19,20 +19,20 @@ class Computer(Node): Instances of computer come 'pre-packaged' with the following: * Core Functionality: - * ARP. - * ICMP. - * Packet Capture. - * Sys Log. + * ARP + * ICMP + * Packet Capture + * Sys Log * Services: - * DNS Client. - * FTP Client. - * LDAP Client. - * NTP Client. + * DNS Client + * FTP Client + * LDAP Client + * NTP Client * Applications: - * Email Client. - * Web Browser. + * Email Client + * Web Browser * Processes: - * Placeholder. + * Placeholder """ def __init__(self, **kwargs): diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 7db92938..b507143b 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -1,10 +1,11 @@ from __future__ import annotations +import secrets from enum import Enum from ipaddress import IPv4Address, IPv4Network from typing import Dict, List, Optional, Tuple, Union -from prettytable import PrettyTable +from prettytable import PrettyTable, MARKDOWN from primaite.simulator.core import SimComponent from primaite.simulator.network.hardware.base import ARPCache, ICMP, NIC, Node @@ -22,8 +23,16 @@ class ACLAction(Enum): class ACLRule(SimComponent): - def describe_state(self) -> Dict: - pass + """ + Represents an Access Control List (ACL) rule. + + :ivar ACLAction action: Action to be performed (Permit/Deny). Default is DENY. + :ivar Optional[IPProtocol] protocol: Network protocol. Default is None. + :ivar Optional[IPv4Address] src_ip: Source IP address. Default is None. + :ivar Optional[Port] src_port: Source port number. Default is None. + :ivar Optional[IPv4Address] dst_ip: Destination IP address. Default is None. + :ivar Optional[Port] dst_port: Destination port number. Default is None. + """ action: ACLAction = ACLAction.DENY protocol: Optional[IPProtocol] = None @@ -43,8 +52,25 @@ class ACLRule(SimComponent): rule_strings.append(f"{key}={value}") return ", ".join(rule_strings) + def describe_state(self) -> Dict: + """ + Describes the current state of the ACLRule. + + :return: A dictionary representing the current state. + """ + pass + class AccessControlList(SimComponent): + """ + Manages a list of ACLRules to filter network traffic. + + :ivar SysLog sys_log: System logging instance. + :ivar ACLAction implicit_action: Default action for rules. + :ivar ACLRule implicit_rule: Implicit ACL rule, created based on implicit_action. + :ivar int max_acl_rules: Maximum number of ACL rules that can be added. Default is 25. + :ivar List[Optional[ACLRule]] _acl: A list containing the ACL rules. + """ sys_log: SysLog implicit_action: ACLAction implicit_rule: ACLRule @@ -62,10 +88,20 @@ class AccessControlList(SimComponent): super().__init__(**kwargs) def describe_state(self) -> Dict: + """ + Describes the current state of the AccessControlList. + + :return: A dictionary representing the current state. + """ pass @property def acl(self) -> List[Optional[ACLRule]]: + """ + Get the list of ACL rules. + + :return: The list of ACL rules. + """ return self._acl def add_rule( @@ -78,6 +114,18 @@ class AccessControlList(SimComponent): dst_port: Optional[Port] = None, position: int = 0, ) -> None: + """ + Add a new ACL rule. + + :param ACLAction action: Action to be performed (Permit/Deny). + :param Optional[IPProtocol] protocol: Network protocol. + :param Optional[Union[str, IPv4Address]] src_ip: Source IP address. + :param Optional[Port] src_port: Source port number. + :param Optional[Union[str, IPv4Address]] dst_ip: Destination IP address. + :param Optional[Port] dst_port: Destination port number. + :param int position: Position in the ACL list to insert the rule. + :raises ValueError: When the position is out of bounds. + """ if isinstance(src_ip, str): src_ip = IPv4Address(src_ip) if isinstance(dst_ip, str): @@ -90,6 +138,12 @@ class AccessControlList(SimComponent): raise ValueError(f"Position {position} is out of bounds.") def remove_rule(self, position: int) -> None: + """ + Remove an ACL rule from a specific position. + + :param int position: The position of the rule to be removed. + :raises ValueError: When the position is out of bounds. + """ if 0 <= position < self.max_acl_rules: self._acl[position] = None else: @@ -103,6 +157,17 @@ class AccessControlList(SimComponent): dst_ip: 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. + + :param protocol: The protocol of the packet. + :param src_ip: Source IP address of the packet. Accepts string and IPv4Address. + :param src_port: Source port of the packet. Optional. + :param dst_ip: Destination IP address of the packet. Accepts string and IPv4Address. + :param dst_port: Destination port of the packet. Optional. + :return: A tuple with a boolean indicating if the packet is permitted and an optional rule or implicit action + string. + """ if not isinstance(src_ip, IPv4Address): src_ip = IPv4Address(src_ip) if not isinstance(dst_ip, IPv4Address): @@ -130,6 +195,16 @@ class AccessControlList(SimComponent): dst_ip: Union[str, IPv4Address], dst_port: Port, ) -> List[ACLRule]: + """ + Get the list of relevant rules for a packet with given properties. + + :param protocol: The protocol of the packet. + :param src_ip: Source IP address of the packet. Accepts string and IPv4Address. + :param src_port: Source port of the packet. + :param dst_ip: Destination IP address of the packet. Accepts string and IPv4Address. + :param dst_port: Destination port of the packet. + :return: A list of relevant ACLRules. + """ if not isinstance(src_ip, IPv4Address): src_ip = IPv4Address(src_ip) if not isinstance(dst_ip, IPv4Address): @@ -150,17 +225,16 @@ class AccessControlList(SimComponent): return relevant_rules - def show(self): - """Prints a table of the routes in the RouteTable.""" + def show(self, markdown: bool = False): + """ + Display the current ACL rules as a table. + + :param markdown: Whether to display the table in Markdown format. Defaults to False. """ - action: ACLAction - protocol: Optional[IPProtocol] - src_ip: Optional[IPv4Address] - src_port: Optional[Port] - dst_ip: Optional[IPv4Address] - dst_port: Optional[Port] - """ table = PrettyTable(["Index", "Action", "Protocol", "Src IP", "Src Port", "Dst IP", "Dst Port"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" table.title = f"{self.sys_log.hostname} Access Control List" for index, rule in enumerate(self.acl + [self.implicit_rule]): if rule: @@ -213,6 +287,11 @@ class RouteEntry(SimComponent): super().__init__(**kwargs) def describe_state(self) -> Dict: + """ + Describes the current state of the RouteEntry. + + :return: A dictionary representing the current state. + """ pass @@ -220,12 +299,7 @@ class RouteTable(SimComponent): """ Represents a routing table holding multiple route entries. - Attributes: - routes (List[RouteEntry]): A list of RouteEntry objects. - - Methods: - add_route: Add a route to the routing table. - find_best_route: Find the best route for a given destination IP. + :ivar List[RouteEntry] routes: A list of RouteEntry objects. Example: >>> rt = RouteTable() @@ -244,6 +318,11 @@ class RouteTable(SimComponent): sys_log: SysLog def describe_state(self) -> Dict: + """ + Describes the current state of the RouteTable. + + :return: A dictionary representing the current state. + """ pass def add_route( @@ -253,9 +332,13 @@ class RouteTable(SimComponent): next_hop: Union[IPv4Address, str], metric: float = 0.0, ): - """Add a route to the routing table. + """ + Add a route to the routing table. - :param route: A RouteEntry object representing the route. + :param address: The destination address of the route. + :param subnet_mask: The subnet mask of the route. + :param next_hop: The next hop IP for the route. + :param metric: The metric of the route, default is 0.0. """ for key in {address, subnet_mask, next_hop}: if not isinstance(key, IPv4Address): @@ -267,10 +350,10 @@ class RouteTable(SimComponent): """ Find the best route for a given destination IP. - :param destination_ip: The destination IPv4Address to find the route for. - :return: The best matching RouteEntry, or None if no route matches. + This method uses the Longest Prefix Match algorithm and considers metrics to find the best route. - The algorithm uses Longest Prefix Match and considers metrics to find the best route. + :param destination_ip: The destination IP to find the route for. + :return: The best matching RouteEntry, or None if no route matches. """ if not isinstance(destination_ip, IPv4Address): destination_ip = IPv4Address(destination_ip) @@ -290,9 +373,16 @@ class RouteTable(SimComponent): return best_route - def show(self): - """Prints a table of the routes in the RouteTable.""" + def show(self, markdown: bool = False): + """ + Display the current routing table as a table. + + :param markdown: Whether to display the table in Markdown format. Defaults to False. + """ table = PrettyTable(["Index", "Address", "Next Hop", "Metric"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" table.title = f"{self.sys_log.hostname} Route Table" for index, route in enumerate(self.routes): network = IPv4Network(f"{route.address}/{route.subnet_mask}") @@ -301,6 +391,12 @@ class RouteTable(SimComponent): class RouterARPCache(ARPCache): + """ + Inherits from ARPCache and adds router-specific ARP packet processing. + + :ivar SysLog sys_log: A system log for logging messages. + :ivar Router router: The router to which this ARP cache belongs. + """ def __init__(self, sys_log: SysLog, router: Router): super().__init__(sys_log) self.router: Router = router @@ -310,7 +406,7 @@ class RouterARPCache(ARPCache): Overridden method to process a received ARP packet in a router-specific way. :param from_nic: The NIC that received the ARP packet. - :param frame: The original arp frame. + :param frame: The original ARP frame. """ arp_packet = frame.arp @@ -356,6 +452,16 @@ class RouterARPCache(ARPCache): class RouterICMP(ICMP): + """ + A class to represent a router's Internet Control Message Protocol (ICMP) handler. + + :param sys_log: System log for logging network events and errors. + :type sys_log: SysLog + :param arp_cache: The ARP cache for resolving MAC addresses. + :type arp_cache: ARPCache + :param router: The router to which this ICMP handler belongs. + :type router: Router + """ router: Router def __init__(self, sys_log: SysLog, arp_cache: ARPCache, router: Router): @@ -363,6 +469,13 @@ class RouterICMP(ICMP): self.router = router def process_icmp(self, frame: Frame, from_nic: NIC, is_reattempt: bool = False): + """ + Process incoming ICMP frames based on ICMP type. + + :param frame: The incoming frame to process. + :param from_nic: The network interface where the frame is coming from. + :param is_reattempt: Flag to indicate if the process is a reattempt. + """ if frame.icmp.icmp_type == ICMPType.ECHO_REQUEST: # determine if request is for router interface or whether it needs to be routed @@ -386,7 +499,10 @@ class RouterICMP(ICMP): identifier=frame.icmp.identifier, sequence=frame.icmp.sequence + 1, ) - frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet) + payload = secrets.token_urlsafe(int(32/1.3)) # Standard ICMP 32 bytes size + frame = Frame( + ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet, payload=payload + ) self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") src_nic.send_frame(frame) @@ -399,7 +515,14 @@ class RouterICMP(ICMP): for nic in self.router.nics.values(): if nic.ip_address == frame.ip.dst_ip: if nic.enabled: - self.sys_log.info(f"Received echo reply from {frame.ip.src_ip}") + time = frame.transmission_duration() + time_str = f"{time}ms" if time > 0 else "<1ms" + self.sys_log.info( + f"Reply from {frame.ip.src_ip}: " + f"bytes={len(frame.payload)}, " + f"time={time_str}, " + f"TTL={frame.ip.ttl}" + ) if not self.request_replies.get(frame.icmp.identifier): self.request_replies[frame.icmp.identifier] = 0 self.request_replies[frame.icmp.identifier] += 1 @@ -410,6 +533,13 @@ class RouterICMP(ICMP): class Router(Node): + """ + A class to represent a network router node. + + :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, RouterARPCache, RouterICMP. + """ num_ports: int ethernet_ports: Dict[int, NIC] = {} acl: AccessControlList @@ -438,14 +568,32 @@ class Router(Node): self.icmp.arp = self.arp def _get_port_of_nic(self, target_nic: NIC) -> Optional[int]: + """ + Retrieve the port number for a given NIC. + + :param target_nic: Target network interface. + :return: The port number if NIC is found, otherwise None. + """ for port, nic in self.ethernet_ports.items(): if nic == target_nic: return port def describe_state(self) -> Dict: + """ + Describes the current state of the Router. + + :return: A dictionary representing the current state. + """ pass def route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None: + """ + Route a given frame from a source NIC to its destination. + + :param frame: The frame to be routed. + :param from_nic: The source network interface. + :param re_attempt: Flag to indicate if the routing is a reattempt. + """ # Check if src ip is on network of one of the NICs nic = self.arp.get_arp_cache_nic(frame.ip.dst_ip) target_mac = self.arp.get_arp_cache_mac_address(frame.ip.dst_ip) @@ -477,13 +625,10 @@ class Router(Node): def receive_frame(self, frame: Frame, from_nic: NIC): """ - Receive a Frame from the connected NIC and process it. + Receive a frame from a NIC and processes it based on its protocol. - Depending on the protocol, the frame is passed to the appropriate handler such as ARP or ICMP, or up to the - SessionManager if no code manager exists. - - :param frame: The Frame being received. - :param from_nic: The NIC that received the frame. + :param frame: The incoming frame. + :param from_nic: The network interface where the frame is coming from. """ route_frame = False protocol = frame.ip.protocol @@ -520,6 +665,13 @@ class Router(Node): self.route_frame(frame, from_nic) def configure_port(self, port: int, ip_address: Union[IPv4Address, str], subnet_mask: Union[IPv4Address, str]): + """ + Configure the IP settings of a given port. + + :param port: The port to configure. + :param ip_address: The IP address to set. + :param subnet_mask: The subnet mask to set. + """ if not isinstance(ip_address, IPv4Address): ip_address = IPv4Address(ip_address) if not isinstance(subnet_mask, IPv4Address): @@ -530,18 +682,36 @@ class Router(Node): self.sys_log.info(f"Configured port {port} with ip_address={ip_address}/{nic.ip_network.prefixlen}") def enable_port(self, port: int): + """ + Enable a given port on the router. + + :param port: The port to enable. + """ nic = self.ethernet_ports.get(port) if nic: nic.enable() def disable_port(self, port: int): + """ + Disable a given port on the router. + + :param port: The port to disable. + """ nic = self.ethernet_ports.get(port) if nic: nic.disable() - def show(self): + def show(self, markdown: bool = False): + """ + Prints the state of the Ethernet interfaces on the Router. + + :param markdown: Flag to indicate if the output should be in markdown format. + """ """Prints a table of the NICs on the Node.""" table = PrettyTable(["Port", "MAC Address", "Address", "Speed", "Status"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" table.title = f"{self.hostname} Ethernet Interfaces" for port, nic in self.ethernet_ports.items(): table.add_row( diff --git a/src/primaite/simulator/network/hardware/nodes/server.py b/src/primaite/simulator/network/hardware/nodes/server.py new file mode 100644 index 00000000..a3e6f2d7 --- /dev/null +++ b/src/primaite/simulator/network/hardware/nodes/server.py @@ -0,0 +1,37 @@ +from ipaddress import IPv4Address + +from primaite.simulator.network.hardware.base import Node, NIC +from primaite.simulator.network.hardware.nodes.computer import Computer + + +class Server(Computer): + """ + A basic Server class. + + Example: + >>> server_a = Server( + hostname="server_a", + ip_address="192.168.1.10", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1" + ) + >>> server_a.power_on() + + Instances of Server come 'pre-packaged' with the following: + + * Core Functionality: + * ARP + * ICMP + * Packet Capture + * Sys Log + * Services: + * DNS Client + * FTP Client + * LDAP Client + * NTP Client + * Applications: + * Email Client + * Web Browser + * Processes: + * Placeholder + """ diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index 0eccefa4..28e58ca4 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -2,10 +2,80 @@ from primaite.simulator.network.container import Network from primaite.simulator.network.hardware.base import Switch, NIC from primaite.simulator.network.hardware.nodes.computer import Computer from primaite.simulator.network.hardware.nodes.router import Router, ACLAction +from primaite.simulator.network.hardware.nodes.server import Server from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port +def client_server_routed() -> Network: + """ + A basic Client/Server Network routed between subnets. + + +------------+ +------------+ +------------+ +------------+ +------------+ + | | | | | | | | | | + | client_1 +------+ switch_2 +------+ router_1 +------+ switch_1 +------+ server_1 | + | | | | | | | | | | + +------------+ +------------+ +------------+ +------------+ +------------+ + + IP Table: + + """ + network = Network() + + # Router 1 + router_1 = Router(hostname="router_1", num_ports=3) + 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.2.1", subnet_mask="255.255.255.0") + + # Switch 1 + switch_1 = Switch(hostname="switch_1", num_ports=6) + switch_1.power_on() + network.connect(endpoint_a=router_1.ethernet_ports[1], endpoint_b=switch_1.switch_ports[6]) + router_1.enable_port(1) + + # Switch 2 + switch_2 = Switch(hostname="switch_2", num_ports=6) + switch_2.power_on() + network.connect(endpoint_a=router_1.ethernet_ports[2], endpoint_b=switch_2.switch_ports[6]) + router_1.enable_port(2) + + # Client 1 + client_1 = Computer( + hostname="client_1", + ip_address="192.168.2.2", + subnet_mask="255.255.255.0", + default_gateway="192.168.2.1" + ) + client_1.power_on() + network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1]) + + # Server 1 + server_1 = Server( + hostname="server_1", + ip_address="192.168.1.2", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1" + ) + server_1.power_on() + network.connect(endpoint_b=server_1.ethernet_port[1], endpoint_a=switch_1.switch_ports[1]) + + router_1.acl.add_rule( + action=ACLAction.PERMIT, + src_port=Port.ARP, + dst_port=Port.ARP, + position=22 + ) + + router_1.acl.add_rule( + action=ACLAction.PERMIT, + protocol=IPProtocol.ICMP, + position=23 + ) + + return network + + def arcd_uc2_network() -> Network: """ Models the ARCD Use Case 2 Network. @@ -40,9 +110,7 @@ def arcd_uc2_network() -> Network: | | +------------+ - Example: - >>> network = arcd_uc2_network() - >>> network.get_node_by_hostname("client_1").ping("192.168.1.10") + """ network = Network() @@ -73,7 +141,7 @@ def arcd_uc2_network() -> Network: default_gateway="192.168.10.1" ) client_1.power_on() - network.connect(endpoint_a=client_1, endpoint_b=switch_2.switch_ports[1]) + network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1]) # Client 2 client_2 = Computer( @@ -83,60 +151,59 @@ def arcd_uc2_network() -> Network: default_gateway="192.168.10.1" ) client_2.power_on() - network.connect(endpoint_a=client_2, endpoint_b=switch_2.switch_ports[2]) + network.connect(endpoint_b=client_2.ethernet_port[1], endpoint_a=switch_2.switch_ports[2]) # Domain Controller - domain_controller = Computer( + domain_controller = Server( hostname="domain_controller", ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" ) domain_controller.power_on() - network.connect(endpoint_a=domain_controller, endpoint_b=switch_1.switch_ports[1]) + network.connect(endpoint_b=domain_controller.ethernet_port[1], endpoint_a=switch_1.switch_ports[1]) # Web Server - web_server = Computer( + web_server = Server( hostname="web_server", ip_address="192.168.1.12", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" ) web_server.power_on() - network.connect(endpoint_a=web_server, endpoint_b=switch_1.switch_ports[2]) + network.connect(endpoint_b=web_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[2]) # Database Server - database_server = Computer( + database_server = Server( hostname="database_server", ip_address="192.168.1.14", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" ) database_server.power_on() - network.connect(endpoint_a=database_server, endpoint_b=switch_1.switch_ports[3]) + network.connect(endpoint_b=database_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[3]) # Backup Server - backup_server = Computer( + backup_server = Server( hostname="backup_server", ip_address="192.168.1.16", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" ) backup_server.power_on() - network.connect(endpoint_a=backup_server, endpoint_b=switch_1.switch_ports[4]) + network.connect(endpoint_b=backup_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[4]) # Security Suite - security_suite = Computer( + security_suite = Server( hostname="security_suite", ip_address="192.168.1.110", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" ) security_suite.power_on() - network.connect(endpoint_a=security_suite, endpoint_b=switch_1.switch_ports[7]) - security_suite_external_nic = NIC(ip_address="192.168.10.110", subnet_mask="255.255.255.0") - security_suite.connect_nic(security_suite_external_nic) - network.connect(endpoint_a=security_suite_external_nic, endpoint_b=switch_2.switch_ports[7]) + network.connect(endpoint_b=security_suite.ethernet_port[1], endpoint_a=switch_1.switch_ports[7]) + security_suite.connect_nic(NIC(ip_address="192.168.10.110", subnet_mask="255.255.255.0")) + network.connect(endpoint_b=security_suite.ethernet_port[2], endpoint_a=switch_2.switch_ports[7]) router_1.acl.add_rule( action=ACLAction.PERMIT, diff --git a/src/primaite/simulator/network/transmission/data_link_layer.py b/src/primaite/simulator/network/transmission/data_link_layer.py index 1b7ccf7d..ddd9fad3 100644 --- a/src/primaite/simulator/network/transmission/data_link_layer.py +++ b/src/primaite/simulator/network/transmission/data_link_layer.py @@ -124,6 +124,11 @@ class Frame(BaseModel): if not self.received_timestamp: self.received_timestamp = datetime.now() + def transmission_duration(self) -> int: + """The transmission duration in milliseconds.""" + delta = self.received_timestamp - self.sent_timestamp + return int(delta.microseconds / 1000) + @property def size(self) -> float: # noqa - Keep it as MBits as this is how they're expressed """The size of the Frame in Bytes.""" diff --git a/src/primaite/simulator/system/core/sys_log.py b/src/primaite/simulator/system/core/sys_log.py index 4b858c2e..5a7bbbfe 100644 --- a/src/primaite/simulator/system/core/sys_log.py +++ b/src/primaite/simulator/system/core/sys_log.py @@ -1,6 +1,8 @@ import logging from pathlib import Path +from prettytable import PrettyTable, MARKDOWN + from primaite.simulator import TEMP_SIM_OUTPUT @@ -43,7 +45,7 @@ class SysLog: file_handler = logging.FileHandler(filename=log_path) file_handler.setLevel(logging.DEBUG) - log_format = "%(asctime)s %(levelname)s: %(message)s" + log_format = "%(asctime)s::%(levelname)s::%(message)s" file_handler.setFormatter(logging.Formatter(log_format)) self.logger = logging.getLogger(f"{self.hostname}_sys_log") @@ -52,6 +54,19 @@ class SysLog: self.logger.addFilter(_NotJSONFilter()) + def show(self, last_n: int = 10, markdown: bool = False): + table = PrettyTable(["Timestamp", "Level", "Message"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" + table.title = f"{self.hostname} Sys Log" + if self._get_log_path().exists(): + with open(self._get_log_path()) as file: + lines = file.readlines() + for line in lines[-last_n:]: + table.add_row(line.strip().split("::")) + print(table) + def _get_log_path(self) -> Path: """ Constructs the path for the log file based on the hostname. From 05959e5408cd6c32b83a5cb407164b7e89ef72ca Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Mon, 4 Sep 2023 12:14:24 +0100 Subject: [PATCH 06/10] #1800 - Moved the Switch code to a dedicated switch.py module. - Added more switch tests. - Updated ACL tests to use router acl. - Updated more docs. - Moved the Jupyter notebooks to _package_data and fixed up the setup to move all notebooks to ~/primaite/notebooks/example_notebooks. --- CHANGELOG.md | 10 +- MANIFEST.in | 1 + .../simulation_components/network/network.rst | 3 +- .../simulation_components/network/switch.rst | 8 - src/primaite/notebooks/__init__.py | 34 - src/primaite/setup/reset_demo_notebooks.py | 45 +- .../create-simulation_demo.ipynb} | 0 .../network_simulator_demo.ipynb | 688 ++++++++++++++++++ src/primaite/simulator/network/container.py | 23 +- .../simulator/network/hardware/base.py | 115 +-- .../network/hardware/nodes/computer.py | 2 +- .../network/hardware/nodes/router.py | 101 +-- .../network/hardware/nodes/server.py | 3 - .../network/hardware/nodes/switch.py | 121 +++ src/primaite/simulator/network/networks.py | 67 +- src/primaite/simulator/system/core/sys_log.py | 10 +- .../network/test_frame_transmission.py | 41 +- .../network/test_nic_link_connection.py | 5 +- .../network/test_switched_network.py | 25 + .../nodes/{test_router.py => test_acl.py} | 23 +- 20 files changed, 992 insertions(+), 333 deletions(-) delete mode 100644 docs/source/simulation_components/network/switch.rst delete mode 100644 src/primaite/notebooks/__init__.py rename src/primaite/{notebooks/create-simulation.ipynb => simulator/_package_data/create-simulation_demo.ipynb} (100%) create mode 100644 src/primaite/simulator/_package_data/network_simulator_demo.ipynb create mode 100644 src/primaite/simulator/network/hardware/nodes/switch.py create mode 100644 tests/integration_tests/network/test_switched_network.py rename tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/{test_router.py => test_acl.py} (84%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b495c09..2f2918aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,12 +7,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] + + ### Added -- Network Hardware - Added base hardware module with NIC, SwitchPort, Node, Switch, and Link. Nodes and Switches have +- Network Hardware - Added base hardware module with NIC, SwitchPort, Node, and Link. Nodes have fundamental services like ARP, ICMP, and PCAP running them by default. - Network Transmission - Modelled OSI Model layers 1 through to 5 with various classes for creating network frames and transmitting them from a Service/Application, down through the layers, over the wire, and back up through the layers to a Service/Application another machine. +- Introduced `Router` and `Switch` classes to manage networking routes more effectively. + - Added `ACLRule` and `RouteTableEntry` classes as part of the `Router`. +- New `.show()` methods in all network component classes to inspect the state in either plain text or markdown formats. +- Added `Computer` and `Server` class to better differentiate types of network nodes. +- Integrated a new Use Case 2 network into the system. +- New unit tests to verify routing between different subnets using `.ping()`. - system - Added the core structure of Application, Services, and Components. Also added a SoftwareManager and SessionManager. - Permission System - each action can define criteria that will be used to permit or deny agent actions. diff --git a/MANIFEST.in b/MANIFEST.in index 51ae4ddf..2ac7b306 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,2 +1,3 @@ include src/primaite/setup/_package_data/primaite_config.yaml include src/primaite/config/_package_data/*.yaml +include src/primaite/simulator/_package_data/*.ipynb diff --git a/docs/source/simulation_components/network/network.rst b/docs/source/simulation_components/network/network.rst index e5614980..f4d64b16 100644 --- a/docs/source/simulation_components/network/network.rst +++ b/docs/source/simulation_components/network/network.rst @@ -30,10 +30,11 @@ we'll use the following Network that has a client, server, two switches, and a r .. code-block:: python from primaite.simulator.network.container import Network - from primaite.simulator.network.hardware.base import Switch, NIC + from primaite.simulator.network.hardware.base import NIC from primaite.simulator.network.hardware.nodes.computer import Computer from primaite.simulator.network.hardware.nodes.router import Router, ACLAction 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 diff --git a/docs/source/simulation_components/network/switch.rst b/docs/source/simulation_components/network/switch.rst deleted file mode 100644 index 4b3b24bc..00000000 --- a/docs/source/simulation_components/network/switch.rst +++ /dev/null @@ -1,8 +0,0 @@ -.. only:: comment - - © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK - -.. _about: - -Switch -====== \ No newline at end of file diff --git a/src/primaite/notebooks/__init__.py b/src/primaite/notebooks/__init__.py deleted file mode 100644 index bc1dcfcd..00000000 --- a/src/primaite/notebooks/__init__.py +++ /dev/null @@ -1,34 +0,0 @@ -# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -"""Contains default jupyter notebooks which demonstrate PrimAITE functionality.""" - -import importlib.util -import os -import subprocess -import sys -from logging import Logger - -from primaite import getLogger, PRIMAITE_PATHS - -_LOGGER: Logger = getLogger(__name__) - - -def start_jupyter_session() -> None: - """ - Starts a new Jupyter notebook session in the app notebooks directory. - - Currently only works on Windows OS. - - .. todo:: Figure out how to get this working for Linux and MacOS too. - """ - if importlib.util.find_spec("jupyter") is not None: - jupyter_cmd = "python3 -m jupyter lab" - if sys.platform == "win32": - jupyter_cmd = "jupyter lab" - - working_dir = os.getcwd() - os.chdir(PRIMAITE_PATHS.user_notebooks_path) - subprocess.Popen(jupyter_cmd) - os.chdir(working_dir) - else: - # Jupyter is not installed - _LOGGER.error("Cannot start jupyter lab as it is not installed") diff --git a/src/primaite/setup/reset_demo_notebooks.py b/src/primaite/setup/reset_demo_notebooks.py index 1f96c90f..a4ee4c4d 100644 --- a/src/primaite/setup/reset_demo_notebooks.py +++ b/src/primaite/setup/reset_demo_notebooks.py @@ -1,35 +1,46 @@ # © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK import filecmp -import os import shutil from logging import Logger from pathlib import Path -import pkg_resources - from primaite import getLogger, PRIMAITE_PATHS _LOGGER: Logger = getLogger(__name__) +def should_copy_file(src: Path, dest: Path, overwrite_existing: bool) -> bool: + """ + Determine if the file should be copied. + + :param src: The source file Path. + :param dest: The destination file Path. + :param overwrite_existing: A bool to toggle replacing existing edited files on or off. + :return: True if file should be copied, otherwise False. + """ + if not dest.is_file(): + return True + + if overwrite_existing and not filecmp.cmp(src, dest): + return True + + return False + + def run(overwrite_existing: bool = True) -> None: """ - Resets the demo jupyter notebooks in the users app notebooks directory. + Resets the demo Jupyter notebooks in the user's app notebooks directory. :param overwrite_existing: A bool to toggle replacing existing edited notebooks on or off. """ - notebooks_package_data_root = pkg_resources.resource_filename("primaite", "notebooks/_package_data") - for subdir, dirs, files in os.walk(notebooks_package_data_root): - for file in files: - fp = os.path.join(subdir, file) - path_split = os.path.relpath(fp, notebooks_package_data_root).split(os.sep) - target_fp = PRIMAITE_PATHS.user_notebooks_path / Path(*path_split) - target_fp.parent.mkdir(exist_ok=True, parents=True) - copy_file = not target_fp.is_file() + primaite_root = Path(__file__).parent.parent + example_notebooks_user_dir = PRIMAITE_PATHS.user_notebooks_path / "example_notebooks" + example_notebooks_user_dir.mkdir(exist_ok=True, parents=True) - if overwrite_existing and not copy_file: - copy_file = (not filecmp.cmp(fp, target_fp)) and (".ipynb_checkpoints" not in str(target_fp)) + for src_fp in primaite_root.glob("**/*.ipynb"): + dst_fp = example_notebooks_user_dir / src_fp.name - if copy_file: - shutil.copy2(fp, target_fp) - _LOGGER.info(f"Reset example notebook: {target_fp}") + if should_copy_file(src_fp, dst_fp, overwrite_existing): + print(dst_fp) + shutil.copy2(src_fp, dst_fp) + _LOGGER.info(f"Reset example notebook: {dst_fp}") diff --git a/src/primaite/notebooks/create-simulation.ipynb b/src/primaite/simulator/_package_data/create-simulation_demo.ipynb similarity index 100% rename from src/primaite/notebooks/create-simulation.ipynb rename to src/primaite/simulator/_package_data/create-simulation_demo.ipynb diff --git a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb new file mode 100644 index 00000000..252f31fa --- /dev/null +++ b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb @@ -0,0 +1,688 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "03b2013a-b7d1-47ee-b08c-8dab83833720", + "metadata": {}, + "source": [ + "# PrimAITE Router Simulation Demo\n", + "\n", + "This demo uses the ARCD Use Case 2 Network (seen below) to demonstrate the capabilities of the Network simulator in PrimAITE." + ] + }, + { + "cell_type": "raw", + "id": "c8bb5698-e746-4e90-9c2f-efe962acdfa0", + "metadata": {}, + "source": [ + " +------------+\n", + " | domain_ |\n", + " +------------+ controller |\n", + " | | |\n", + " | +------------+\n", + " |\n", + " |\n", + "+------------+ | +------------+\n", + "| | | | |\n", + "| client_1 +---------+ | +---------+ web_server |\n", + "| | | | | | |\n", + "+------------+ | | | +------------+\n", + " +--+---------+ +------------+ +------+--+--+\n", + " | | | | | |\n", + " | switch_2 +------+ router_1 +------+ switch_1 |\n", + " | | | | | |\n", + " +--+------+--+ +------------+ +--+---+--+--+\n", + "+------------+ | | | | | +------------+\n", + "| | | | | | | | database |\n", + "| client_2 +---------+ | | | +---------+ _server |\n", + "| | | | | | |\n", + "+------------+ | | | +------------+\n", + " | +------------+ | |\n", + " | | security | | |\n", + " +---------+ _suite +---------+ | +------------+\n", + " | | | | backup_ |\n", + " +------------+ +------------+ server |\n", + " | |\n", + " +------------+" + ] + }, + { + "cell_type": "markdown", + "id": "415d487c-6457-497d-85d6-99439b3541e7", + "metadata": {}, + "source": [ + "## The Network\n", + "First let's create our network. The network comes 'pre-packaged' with PrimAITE in the `primaite.simulator.network.networks` module.\n", + "\n", + "> ℹ️ You'll see a bunch of logs associated with parts of the Network that aern't an 'electronic' device on the Network and thus don't have a stsrem to log to. Soon these logs are going to be pushed to a Network Logger so we're not clogging up the PrimAITE application logs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "de57ac8c-5b28-4847-a759-2ceaf5593329", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from primaite.simulator.network.networks import arcd_uc2_network" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1e2e4df-67c0-4584-ab27-47e2c7c7fcd2", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network = arcd_uc2_network()" + ] + }, + { + "cell_type": "markdown", + "id": "fb052c56-e9ca-4093-9115-d0c440b5ff53", + "metadata": {}, + "source": [ + "Most of the Network components have a `.show()` function that prints a table of information about that object. We can view the Nodes and Links on the Network by calling `network.show()`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc199741-ef2e-47f5-b2f0-e20049ccf40f", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.show()" + ] + }, + { + "cell_type": "markdown", + "id": "76d2b7e9-280b-4741-a8b3-a84bed219fac", + "metadata": { + "tags": [] + }, + "source": [ + "## Nodes\n", + "\n", + "Now let's inspect some of the nodes. We can directly access a node on the Network by calling .`get_node_by_hostname`. Like Network, a Node, along with some core services like ARP, have a `.show()` method." + ] + }, + { + "cell_type": "markdown", + "id": "84113002-843e-4cab-b899-667b50f25f6b", + "metadata": {}, + "source": [ + "### Router Nodes\n", + "\n", + "First we'll inspect the Router node and some of it's core services." + ] + }, + { + "cell_type": "markdown", + "id": "bf63a178-eee5-4669-bf64-13aea7ecf6cb", + "metadata": {}, + "source": [ + "Calling `router.show()` displays the Ethernet interfaces on the Router. If you need a table in markdown format, pass `markdown=True`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e76d1854-961e-438c-b40f-77fd9c3abe38", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"router_1\").show()" + ] + }, + { + "cell_type": "markdown", + "id": "e000540c-687c-4254-870c-1d814603bdbf", + "metadata": {}, + "source": [ + "Calling `router.arp.show()` displays the Router ARP Cache." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "92de8b42-92d7-4934-9c12-50bf724c9eb2", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"router_1\").arp.show()" + ] + }, + { + "cell_type": "markdown", + "id": "a9ff7ee8-9482-44de-9039-b684866bdc82", + "metadata": {}, + "source": [ + "Calling `router.acl.show()` displays the Access Control List." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5922282a-d22b-4e55-9176-f3f3654c849f", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"router_1\").acl.show()" + ] + }, + { + "cell_type": "markdown", + "id": "71c87884-f793-4c9f-b004-5b0df86cf585", + "metadata": {}, + "source": [ + "Calling `router.router_table.show()` displays the static routes the Router provides." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "327203be-f475-4727-82a1-e992d3b70ed8", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"router_1\").route_table.show()" + ] + }, + { + "cell_type": "markdown", + "id": "eef561a8-3d39-4c8b-bbc8-e8b10b8ed25f", + "metadata": {}, + "source": [ + "Calling `router.sys_log.show()` displays the Router system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3d0aa004-b10c-445f-aaab-340e0e716c74", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"router_1\").sys_log.show(last_n=10)" + ] + }, + { + "cell_type": "markdown", + "id": "25630c90-c54e-4b5d-8bf4-ad1b0722e126", + "metadata": {}, + "source": [ + "### Switch Nodes\n", + "\n", + "Next we'll inspect the Switch node and some of it's core services." + ] + }, + { + "cell_type": "markdown", + "id": "4879394d-2981-40de-a229-e19b09a34e6e", + "metadata": {}, + "source": [ + "Calling `switch.show()` displays the Switch orts on the Switch." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e7fd439b-5442-4e9d-9e7d-86dacb77f458", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"switch_1\").show()" + ] + }, + { + "cell_type": "markdown", + "id": "beb8dbd6-7250-4ac9-9fa2-d2a9c0e5fd19", + "metadata": { + "tags": [] + }, + "source": [ + "Calling `switch.arp.show()` displays the Switch ARP Cache." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d06e1310-4a77-4315-a59f-cb1b49ca2352", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"switch_1\").arp.show()" + ] + }, + { + "cell_type": "markdown", + "id": "fda75ac3-8123-4234-8f36-86547891d8df", + "metadata": {}, + "source": [ + "Calling `switch.sys_log.show()` displays the Switch system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a0d984b7-a7c1-4bbd-aa5a-9d3caecb08dc", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"switch_1\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "2f1d99ad-db4f-4baf-8a35-e1d95f269586", + "metadata": {}, + "source": [ + "### Computer/Server Nodes\n", + "\n", + "Finally, we'll inspect a Computer or Server Node and some of its core services." + ] + }, + { + "cell_type": "markdown", + "id": "c9e2251a-1b47-46e5-840f-7fec3e39c5aa", + "metadata": { + "tags": [] + }, + "source": [ + "Calling `computer.show()` displays the NICs on the Computer/Server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "656c37f6-b145-42af-9714-8d2886d0eff8", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"security_suite\").show()" + ] + }, + { + "cell_type": "markdown", + "id": "f1097a49-a3da-4d79-a06d-ae8af452918f", + "metadata": {}, + "source": [ + "Calling `computer.arp.show()` displays the Computer/Server ARP Cache." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66b267d6-2308-486a-b9aa-cb8d3bcf0753", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"security_suite\").arp.show()" + ] + }, + { + "cell_type": "markdown", + "id": "0d1fcad8-5b1a-4d8b-a49f-aa54a95fcaf0", + "metadata": {}, + "source": [ + "Calling `switch.sys_log.show()` displays the Computer/Server system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1b5debe8-ef1b-445d-8fa9-6a45568f21f3", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"security_suite\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "fcfa1773-798c-4ada-9318-c3ad928217da", + "metadata": {}, + "source": [ + "## Basic Network Comms Check\n", + "\n", + "We can perform a good old ping to check that Nodes are able to communicate with each other." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "495b7de4-b6ce-41a6-9114-f74752ab4491", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.show(nodes=False, links=False)" + ] + }, + { + "cell_type": "markdown", + "id": "3e13922a-217f-4f4e-99b6-57a07613cade", + "metadata": {}, + "source": [ + "We'll first ping client_1's default gateway." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a38abb71-994e-49e8-8f51-e9a550e95b99", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.10.1\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8388e1e9-30e3-4534-8e5a-c6e9144149d2", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_1\").sys_log.show(15)" + ] + }, + { + "cell_type": "markdown", + "id": "02c76d5c-d954-49db-912d-cb9c52f46375", + "metadata": {}, + "source": [ + "Next, we'll ping the interface of the 192.168.1.0/24 Network on the Router (port 1)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ff8e976a-c16b-470c-8923-325713a30d6c", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.1\")" + ] + }, + { + "cell_type": "markdown", + "id": "80280404-a5ab-452f-8a02-771a0d7496b1", + "metadata": {}, + "source": [ + "And finally, we'll ping the web server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c4163f8d-6a72-410c-9f5c-4f881b7de45e", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")" + ] + }, + { + "cell_type": "markdown", + "id": "1194c045-ba77-4427-be30-ed7b5b224850", + "metadata": {}, + "source": [ + "To confirm that the ping was received and processed by the web_server, we can view the sys log" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e79a523a-5780-45b6-8798-c434e0e522bd", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"web_server\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "5928f6dd-1006-45e3-99f3-8f311a875faa", + "metadata": {}, + "source": [ + "## Advanced Network Usage\n", + "\n", + "We can now use the Network to perform some more advaced things." + ] + }, + { + "cell_type": "markdown", + "id": "5e023ef3-7d18-4006-96ee-042a06a481fc", + "metadata": {}, + "source": [ + "Let's attempt to prevent client_2 from being able to ping the web server. First, we'll confirm that it can ping the server first..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "603cf913-e261-49da-a7dd-85e1bb6dec56", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" + ] + }, + { + "cell_type": "markdown", + "id": "5cf962a4-20e6-44ae-9748-7fc5267ae111", + "metadata": {}, + "source": [ + "If we look at the client_2 sys log we can see that the four ICMP echo requests were sent and four ICMP each replies were received:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e047de00-3de4-4823-b26a-2c8d64c7a663", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_2\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "bdc4741d-6e3e-4aec-a69c-c2e9653bd02c", + "metadata": {}, + "source": [ + "Now we'll add an ACL to block ICMP from 192.168.10.22" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6db355ae-b99a-441b-a2c4-4ffe78f46bff", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", + "from primaite.simulator.network.transmission.transport_layer import Port\n", + "from primaite.simulator.network.hardware.nodes.router import ACLAction\n", + "network.get_node_by_hostname(\"router_1\").acl.add_rule(\n", + " action=ACLAction.DENY,\n", + " protocol=IPProtocol.ICMP,\n", + " src_ip=\"192.168.10.22\",\n", + " position=1\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a345e000-8842-4827-af96-adc0fbe390fb", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"router_1\").acl.show()" + ] + }, + { + "cell_type": "markdown", + "id": "3a5bfd9f-04cb-493e-a86c-cd268563a262", + "metadata": {}, + "source": [ + "Now we attempt (and fail) to ping the web server" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a4f4ff31-590f-40fb-b13d-efaa8c2720b6", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" + ] + }, + { + "cell_type": "markdown", + "id": "83e56497-097b-45cb-964e-b15c72547b38", + "metadata": {}, + "source": [ + "We can check that the ping was actually sent by client_2 by viewing the sys log" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f62b8a4e-fd3b-4059-b108-3d4a0b18f2a0", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_2\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "c7040311-a879-4620-86a0-55d0774156e5", + "metadata": {}, + "source": [ + "We can check the router sys log to see why the traffic was blocked" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7e53d776-99da-4d2c-a2a7-bd7ce27bff4c", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"router_1\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "aba0bc7d-da57-477b-b34a-3688b5aab2c6", + "metadata": {}, + "source": [ + "Now a final check to ensure that client_1 can still ping the web_server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d542734b-7582-4af7-8254-bda3de50d091", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d78e9fe3-02c6-4792-944f-5622e26e0412", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_1\").sys_log.show()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index ccb9ce77..239c98a7 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -1,16 +1,17 @@ -from typing import Any, Dict, Union, Optional, List +from typing import Any, Dict, List, Optional, Union import matplotlib.pyplot as plt import networkx as nx from networkx import MultiGraph -from prettytable import PrettyTable, MARKDOWN +from prettytable import MARKDOWN, PrettyTable from primaite import getLogger from primaite.simulator.core import Action, ActionManager, AllowAllValidator, SimComponent -from primaite.simulator.network.hardware.base import Link, NIC, Node, SwitchPort, Switch +from primaite.simulator.network.hardware.base import Link, NIC, Node, SwitchPort from primaite.simulator.network.hardware.nodes.computer import Computer from primaite.simulator.network.hardware.nodes.router import Router from primaite.simulator.network.hardware.nodes.server import Server +from primaite.simulator.network.hardware.nodes.switch import Switch _LOGGER = getLogger(__name__) @@ -30,7 +31,7 @@ class Network(SimComponent): links: Dict[str, Link] = {} def __init__(self, **kwargs): - """" + """ Initialise the network. Constructs the network and sets up its initial state including @@ -84,14 +85,14 @@ class Network(SimComponent): "Router": self.routers, "Switch": self.switches, "Server": self.servers, - "Computer": self.computers + "Computer": self.computers, } if nodes: table = PrettyTable(["Node", "Type", "Operating State"]) if markdown: table.set_style(MARKDOWN) table.align = "l" - table.title = f"Nodes" + table.title = "Nodes" for node_type, nodes in nodes_type_map.items(): for node in nodes: table.add_row([node.hostname, node_type, node.operating_state.name]) @@ -102,7 +103,7 @@ class Network(SimComponent): if markdown: table.set_style(MARKDOWN) table.align = "l" - table.title = f"IP Addresses" + table.title = "IP Addresses" for nodes in nodes_type_map.values(): for node in nodes: for i, port in node.ethernet_port.items(): @@ -114,7 +115,7 @@ class Network(SimComponent): if markdown: table.set_style(MARKDOWN) table.align = "l" - table.title = f"Links" + table.title = "Links" links = list(self.links.values()) for nodes in nodes_type_map.values(): for node in nodes: @@ -126,7 +127,7 @@ class Network(SimComponent): link.endpoint_b.parent.hostname, link.is_up, link.bandwidth, - link.current_load_percent + link.current_load_percent, ] ) links.remove(link) @@ -207,9 +208,7 @@ class Network(SimComponent): node.parent = None _LOGGER.info(f"Removed node {node.uuid} from network {self.uuid}") - def connect( - self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs - ) -> None: + def connect(self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs) -> None: """ Connect two endpoints on the network by creating a link between their NICs/SwitchPorts. diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 674020ee..1193f3ef 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1,13 +1,12 @@ from __future__ import annotations -import random import re import secrets from enum import Enum from ipaddress import IPv4Address, IPv4Network -from typing import Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union -from prettytable import PrettyTable, MARKDOWN +from prettytable import MARKDOWN, PrettyTable from primaite import getLogger from primaite.exceptions import NetworkError @@ -289,7 +288,7 @@ class SwitchPort(SimComponent): "The speed of the SwitchPort in Mbps. Default is 100 Mbps." mtu: int = 1500 "The Maximum Transmission Unit (MTU) of the SwitchPort in Bytes. Default is 1500 B" - connected_node: Optional[Switch] = None + connected_node: Optional[Node] = None "The Node to which the SwitchPort is connected." connected_link: Optional[Link] = None "The Link to which the SwitchPort is connected." @@ -715,7 +714,7 @@ class ARPCache: arp_packet = arp_packet.generate_reply(from_nic.mac_address) self.send_arp_reply(arp_packet, from_nic) - def __contains__(self, item) -> bool: + def __contains__(self, item: Any) -> bool: return item in self.arp @@ -765,7 +764,7 @@ class ICMP: identifier=frame.icmp.identifier, sequence=frame.icmp.sequence + 1, ) - payload = secrets.token_urlsafe(int(32/1.3)) # Standard ICMP 32 bytes size + payload = secrets.token_urlsafe(int(32 / 1.3)) # Standard ICMP 32 bytes size frame = Frame( ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet, payload=payload ) @@ -829,7 +828,7 @@ class ICMP: # Data Link Layer ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) icmp_packet = ICMPPacket(identifier=identifier, sequence=sequence) - payload = secrets.token_urlsafe(int(32/1.3)) # Standard ICMP 32 bytes size + payload = secrets.token_urlsafe(int(32 / 1.3)) # Standard ICMP 32 bytes size frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_packet, payload=payload) nic.send_frame(frame) return sequence, icmp_packet.identifier @@ -1049,7 +1048,8 @@ class Node(SimComponent): f"Ping statistics for {target_ip_address}: " f"Packets: Sent = {pings}, " f"Received = {request_replies}, " - f"Lost = {pings-request_replies} ({(pings-request_replies)/pings*100}% loss)") + f"Lost = {pings-request_replies} ({(pings-request_replies)/pings*100}% loss)" + ) return passed return False @@ -1084,102 +1084,3 @@ class Node(SimComponent): pass elif frame.ip.protocol == IPProtocol.ICMP: self.icmp.process_icmp(frame=frame, from_nic=from_nic) - - -class Switch(Node): - """A class representing a Layer 2 network switch.""" - - num_ports: int = 24 - "The number of ports on the switch." - switch_ports: Dict[int, SwitchPort] = {} - "The SwitchPorts on the switch." - mac_address_table: Dict[str, SwitchPort] = {} - "A MAC address table mapping destination MAC addresses to corresponding SwitchPorts." - - def __init__(self, **kwargs): - super().__init__(**kwargs) - if not self.switch_ports: - self.switch_ports = {i: SwitchPort() for i in range(1, self.num_ports + 1)} - for port_num, port in self.switch_ports.items(): - port.connected_node = self - port.parent = self - port.port_num = port_num - - def show(self, markdown: bool = False): - """Prints a table of the SwitchPorts on the Switch.""" - table = PrettyTable(["Port", "MAC Address", "Speed", "Status"]) - if markdown: - table.set_style(MARKDOWN) - table.align = "l" - table.title = f"{self.hostname} Switch Ports" - for port_num, port in self.switch_ports.items(): - table.add_row([port_num, port.mac_address, port.speed, "Enabled" if port.enabled else "Disabled"]) - print(table) - - def describe_state(self) -> Dict: - """ - Produce a dictionary describing the current state of this object. - - Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation. - - :return: Current state of this object and child objects. - :rtype: Dict - """ - return { - "uuid": self.uuid, - "num_ports": self.num_ports, # redundant? - "ports": {port_num: port.describe_state() for port_num, port in self.switch_ports.items()}, - "mac_address_table": {mac: port for mac, port in self.mac_address_table.items()}, - } - - def _add_mac_table_entry(self, mac_address: str, switch_port: SwitchPort): - mac_table_port = self.mac_address_table.get(mac_address) - if not mac_table_port: - self.mac_address_table[mac_address] = switch_port - self.sys_log.info(f"Added MAC table entry: Port {switch_port.port_num} -> {mac_address}") - else: - if mac_table_port != switch_port: - self.mac_address_table.pop(mac_address) - self.sys_log.info(f"Removed MAC table entry: Port {mac_table_port.port_num} -> {mac_address}") - self._add_mac_table_entry(mac_address, switch_port) - - def forward_frame(self, frame: Frame, incoming_port: SwitchPort): - """ - Forward a frame to the appropriate port based on the destination MAC address. - - :param frame: The Frame to be forwarded. - :param incoming_port: The port number from which the frame was received. - """ - src_mac = frame.ethernet.src_mac_addr - dst_mac = frame.ethernet.dst_mac_addr - self._add_mac_table_entry(src_mac, incoming_port) - - outgoing_port = self.mac_address_table.get(dst_mac) - if outgoing_port or dst_mac != "ff:ff:ff:ff:ff:ff": - outgoing_port.send_frame(frame) - else: - # If the destination MAC is not in the table, flood to all ports except incoming - for port in self.switch_ports.values(): - if port != incoming_port: - port.send_frame(frame) - - def disconnect_link_from_port(self, link: Link, port_number: int): - """ - Disconnect a given link from the specified port number on the switch. - - :param link: The Link object to be disconnected. - :param port_number: The port number on the switch from where the link should be disconnected. - :raise NetworkError: When an invalid port number is provided or the link does not match the connection. - """ - port = self.switch_ports.get(port_number) - if port is None: - msg = f"Invalid port number {port_number} on the switch" - _LOGGER.error(msg) - raise NetworkError(msg) - - if port.connected_link != link: - msg = f"The link does not match the connection at port number {port_number}" - _LOGGER.error(msg) - raise NetworkError(msg) - - port.disconnect_link() diff --git a/src/primaite/simulator/network/hardware/nodes/computer.py b/src/primaite/simulator/network/hardware/nodes/computer.py index 110ad385..2a2e8524 100644 --- a/src/primaite/simulator/network/hardware/nodes/computer.py +++ b/src/primaite/simulator/network/hardware/nodes/computer.py @@ -1,6 +1,6 @@ from ipaddress import IPv4Address -from primaite.simulator.network.hardware.base import Node, NIC +from primaite.simulator.network.hardware.base import NIC, Node class Computer(Node): diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index b507143b..26ba01ae 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -5,7 +5,7 @@ from enum import Enum from ipaddress import IPv4Address, IPv4Network from typing import Dict, List, Optional, Tuple, Union -from prettytable import PrettyTable, MARKDOWN +from prettytable import MARKDOWN, PrettyTable from primaite.simulator.core import SimComponent from primaite.simulator.network.hardware.base import ARPCache, ICMP, NIC, Node @@ -71,6 +71,7 @@ class AccessControlList(SimComponent): :ivar int max_acl_rules: Maximum number of ACL rules that can be added. Default is 25. :ivar List[Optional[ACLRule]] _acl: A list containing the ACL rules. """ + sys_log: SysLog implicit_action: ACLAction implicit_rule: ACLRule @@ -105,14 +106,14 @@ class AccessControlList(SimComponent): return self._acl def add_rule( - self, - action: ACLAction, - protocol: Optional[IPProtocol] = None, - src_ip: Optional[Union[str, IPv4Address]] = None, - src_port: Optional[Port] = None, - dst_ip: Optional[Union[str, IPv4Address]] = None, - dst_port: Optional[Port] = None, - position: int = 0, + self, + action: ACLAction, + protocol: Optional[IPProtocol] = None, + src_ip: Optional[Union[str, IPv4Address]] = None, + src_port: Optional[Port] = None, + dst_ip: Optional[Union[str, IPv4Address]] = None, + dst_port: Optional[Port] = None, + position: int = 0, ) -> None: """ Add a new ACL rule. @@ -150,12 +151,12 @@ class AccessControlList(SimComponent): raise ValueError(f"Position {position} is out of bounds.") def is_permitted( - self, - protocol: IPProtocol, - src_ip: Union[str, IPv4Address], - src_port: Optional[Port], - dst_ip: Union[str, IPv4Address], - dst_port: Optional[Port], + self, + protocol: IPProtocol, + src_ip: Union[str, IPv4Address], + src_port: Optional[Port], + dst_ip: 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. @@ -177,23 +178,23 @@ class AccessControlList(SimComponent): continue if ( - (rule.src_ip == src_ip or rule.src_ip is None) - and (rule.dst_ip == dst_ip or rule.dst_ip 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 == src_ip or rule.src_ip is None) + and (rule.dst_ip == dst_ip or rule.dst_ip 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: Union[str, IPv4Address], - src_port: Port, - dst_ip: Union[str, IPv4Address], - dst_port: Port, + self, + protocol: IPProtocol, + src_ip: Union[str, IPv4Address], + src_port: Port, + dst_ip: Union[str, IPv4Address], + dst_port: Port, ) -> List[ACLRule]: """ Get the list of relevant rules for a packet with given properties. @@ -215,11 +216,11 @@ class AccessControlList(SimComponent): continue if ( - (rule.src_ip == src_ip or rule.src_ip is None) - or (rule.dst_ip == dst_ip or rule.dst_ip 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 == src_ip or rule.src_ip is None) + or (rule.dst_ip == dst_ip or rule.dst_ip 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) @@ -326,11 +327,11 @@ class RouteTable(SimComponent): pass def add_route( - self, - address: Union[IPv4Address, str], - subnet_mask: Union[IPv4Address, str], - next_hop: Union[IPv4Address, str], - metric: float = 0.0, + self, + address: Union[IPv4Address, str], + subnet_mask: Union[IPv4Address, str], + next_hop: Union[IPv4Address, str], + metric: float = 0.0, ): """ Add a route to the routing table. @@ -397,6 +398,7 @@ class RouterARPCache(ARPCache): :ivar SysLog sys_log: A system log for logging messages. :ivar Router router: The router to which this ARP cache belongs. """ + def __init__(self, sys_log: SysLog, router: Router): super().__init__(sys_log) self.router: Router = router @@ -416,7 +418,8 @@ class RouterARPCache(ARPCache): if arp_packet.target_ip == nic.ip_address: # reply to the Router specifically self.sys_log.info( - f"Received ARP response for {arp_packet.sender_ip} from {arp_packet.sender_mac_addr} via NIC {from_nic}" + f"Received ARP response for {arp_packet.sender_ip} " + f"from {arp_packet.sender_mac_addr} via NIC {from_nic}" ) self.add_arp_cache_entry( ip_address=arp_packet.sender_ip, @@ -462,6 +465,7 @@ class RouterICMP(ICMP): :param router: The router to which this ICMP handler belongs. :type router: Router """ + router: Router def __init__(self, sys_log: SysLog, arp_cache: ARPCache, router: Router): @@ -492,16 +496,22 @@ class RouterICMP(ICMP): # Network Layer ip_packet = IPPacket(src_ip=nic.ip_address, dst_ip=frame.ip.src_ip, protocol=IPProtocol.ICMP) # Data Link Layer - ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) + ethernet_header = EthernetHeader( + src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address + ) icmp_reply_packet = ICMPPacket( icmp_type=ICMPType.ECHO_REPLY, icmp_code=0, identifier=frame.icmp.identifier, sequence=frame.icmp.sequence + 1, ) - payload = secrets.token_urlsafe(int(32/1.3)) # Standard ICMP 32 bytes size + payload = secrets.token_urlsafe(int(32 / 1.3)) # Standard ICMP 32 bytes size frame = Frame( - ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet, payload=payload + ethernet=ethernet_header, + ip=ip_packet, + tcp=tcp_header, + icmp=icmp_reply_packet, + payload=payload, ) self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") @@ -540,6 +550,7 @@ class Router(Node): :ivar int num_ports: The number of ports in the router. :ivar dict kwargs: Optional keyword arguments for SysLog, ACL, RouteTable, RouterARPCache, RouterICMP. """ + num_ports: int ethernet_ports: Dict[int, NIC] = {} acl: AccessControlList @@ -588,12 +599,12 @@ class Router(Node): def route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None: """ - Route a given frame from a source NIC to its destination. + Route a given frame from a source NIC to its destination. - :param frame: The frame to be routed. - :param from_nic: The source network interface. - :param re_attempt: Flag to indicate if the routing is a reattempt. - """ + :param frame: The frame to be routed. + :param from_nic: The source network interface. + :param re_attempt: Flag to indicate if the routing is a reattempt. + """ # Check if src ip is on network of one of the NICs nic = self.arp.get_arp_cache_nic(frame.ip.dst_ip) target_mac = self.arp.get_arp_cache_mac_address(frame.ip.dst_ip) diff --git a/src/primaite/simulator/network/hardware/nodes/server.py b/src/primaite/simulator/network/hardware/nodes/server.py index a3e6f2d7..b72cc71c 100644 --- a/src/primaite/simulator/network/hardware/nodes/server.py +++ b/src/primaite/simulator/network/hardware/nodes/server.py @@ -1,6 +1,3 @@ -from ipaddress import IPv4Address - -from primaite.simulator.network.hardware.base import Node, NIC from primaite.simulator.network.hardware.nodes.computer import Computer diff --git a/src/primaite/simulator/network/hardware/nodes/switch.py b/src/primaite/simulator/network/hardware/nodes/switch.py new file mode 100644 index 00000000..b7cc1242 --- /dev/null +++ b/src/primaite/simulator/network/hardware/nodes/switch.py @@ -0,0 +1,121 @@ +from typing import Dict + +from prettytable import MARKDOWN, PrettyTable + +from primaite import getLogger +from primaite.exceptions import NetworkError +from primaite.links.link import Link +from primaite.simulator.network.hardware.base import Node, SwitchPort +from primaite.simulator.network.transmission.data_link_layer import Frame + +_LOGGER = getLogger(__name__) + + +class Switch(Node): + """ + A class representing a Layer 2 network switch. + + :ivar num_ports: The number of ports on the switch. Default is 24. + """ + + num_ports: int = 24 + "The number of ports on the switch." + switch_ports: Dict[int, SwitchPort] = {} + "The SwitchPorts on the switch." + mac_address_table: Dict[str, SwitchPort] = {} + "A MAC address table mapping destination MAC addresses to corresponding SwitchPorts." + + def __init__(self, **kwargs): + super().__init__(**kwargs) + if not self.switch_ports: + self.switch_ports = {i: SwitchPort() for i in range(1, self.num_ports + 1)} + for port_num, port in self.switch_ports.items(): + port.connected_node = self + port.parent = self + port.port_num = port_num + + def show(self, markdown: bool = False): + """ + Prints a table of the SwitchPorts on the Switch. + + :param markdown: If True, outputs the table in markdown format. Default is False. + """ + table = PrettyTable(["Port", "MAC Address", "Speed", "Status"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" + table.title = f"{self.hostname} Switch Ports" + for port_num, port in self.switch_ports.items(): + table.add_row([port_num, port.mac_address, port.speed, "Enabled" if port.enabled else "Disabled"]) + print(table) + + def describe_state(self) -> Dict: + """ + Produce a dictionary describing the current state of this object. + + :return: Current state of this object and child objects. + """ + return { + "uuid": self.uuid, + "num_ports": self.num_ports, # redundant? + "ports": {port_num: port.describe_state() for port_num, port in self.switch_ports.items()}, + "mac_address_table": {mac: port for mac, port in self.mac_address_table.items()}, + } + + def _add_mac_table_entry(self, mac_address: str, switch_port: SwitchPort): + """ + Private method to add an entry to the MAC address table. + + :param mac_address: MAC address to be added. + :param switch_port: Corresponding SwitchPort object. + """ + mac_table_port = self.mac_address_table.get(mac_address) + if not mac_table_port: + self.mac_address_table[mac_address] = switch_port + self.sys_log.info(f"Added MAC table entry: Port {switch_port.port_num} -> {mac_address}") + else: + if mac_table_port != switch_port: + self.mac_address_table.pop(mac_address) + self.sys_log.info(f"Removed MAC table entry: Port {mac_table_port.port_num} -> {mac_address}") + self._add_mac_table_entry(mac_address, switch_port) + + def forward_frame(self, frame: Frame, incoming_port: SwitchPort): + """ + Forward a frame to the appropriate port based on the destination MAC address. + + :param frame: The Frame to be forwarded. + :param incoming_port: The port number from which the frame was received. + """ + src_mac = frame.ethernet.src_mac_addr + dst_mac = frame.ethernet.dst_mac_addr + self._add_mac_table_entry(src_mac, incoming_port) + + outgoing_port = self.mac_address_table.get(dst_mac) + if outgoing_port or dst_mac != "ff:ff:ff:ff:ff:ff": + outgoing_port.send_frame(frame) + else: + # If the destination MAC is not in the table, flood to all ports except incoming + for port in self.switch_ports.values(): + if port != incoming_port: + port.send_frame(frame) + + def disconnect_link_from_port(self, link: Link, port_number: int): + """ + Disconnect a given link from the specified port number on the switch. + + :param link: The Link object to be disconnected. + :param port_number: The port number on the switch from where the link should be disconnected. + :raise NetworkError: When an invalid port number is provided or the link does not match the connection. + """ + port = self.switch_ports.get(port_number) + if port is None: + msg = f"Invalid port number {port_number} on the switch" + _LOGGER.error(msg) + raise NetworkError(msg) + + if port.connected_link != link: + msg = f"The link does not match the connection at port number {port_number}" + _LOGGER.error(msg) + raise NetworkError(msg) + + port.disconnect_link() diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index 28e58ca4..6a50fe3f 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -1,8 +1,9 @@ from primaite.simulator.network.container import Network -from primaite.simulator.network.hardware.base import Switch, NIC +from primaite.simulator.network.hardware.base import NIC from primaite.simulator.network.hardware.nodes.computer import Computer -from primaite.simulator.network.hardware.nodes.router import Router, ACLAction +from primaite.simulator.network.hardware.nodes.router import ACLAction, Router 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 @@ -42,36 +43,21 @@ def client_server_routed() -> Network: # Client 1 client_1 = Computer( - hostname="client_1", - ip_address="192.168.2.2", - subnet_mask="255.255.255.0", - default_gateway="192.168.2.1" + hostname="client_1", ip_address="192.168.2.2", subnet_mask="255.255.255.0", default_gateway="192.168.2.1" ) client_1.power_on() network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1]) # Server 1 server_1 = Server( - hostname="server_1", - ip_address="192.168.1.2", - subnet_mask="255.255.255.0", - default_gateway="192.168.1.1" + hostname="server_1", ip_address="192.168.1.2", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" ) server_1.power_on() network.connect(endpoint_b=server_1.ethernet_port[1], endpoint_a=switch_1.switch_ports[1]) - router_1.acl.add_rule( - action=ACLAction.PERMIT, - src_port=Port.ARP, - dst_port=Port.ARP, - position=22 - ) + router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22) - router_1.acl.add_rule( - action=ACLAction.PERMIT, - protocol=IPProtocol.ICMP, - position=23 - ) + router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23) return network @@ -135,20 +121,14 @@ def arcd_uc2_network() -> Network: # Client 1 client_1 = Computer( - hostname="client_1", - ip_address="192.168.10.21", - subnet_mask="255.255.255.0", - default_gateway="192.168.10.1" + hostname="client_1", ip_address="192.168.10.21", subnet_mask="255.255.255.0", default_gateway="192.168.10.1" ) client_1.power_on() network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1]) # Client 2 client_2 = Computer( - hostname="client_2", - ip_address="192.168.10.22", - subnet_mask="255.255.255.0", - default_gateway="192.168.10.1" + hostname="client_2", ip_address="192.168.10.22", subnet_mask="255.255.255.0", default_gateway="192.168.10.1" ) client_2.power_on() network.connect(endpoint_b=client_2.ethernet_port[1], endpoint_a=switch_2.switch_ports[2]) @@ -158,17 +138,14 @@ def arcd_uc2_network() -> Network: hostname="domain_controller", ip_address="192.168.1.10", subnet_mask="255.255.255.0", - default_gateway="192.168.1.1" + default_gateway="192.168.1.1", ) domain_controller.power_on() network.connect(endpoint_b=domain_controller.ethernet_port[1], endpoint_a=switch_1.switch_ports[1]) # Web Server web_server = Server( - hostname="web_server", - ip_address="192.168.1.12", - subnet_mask="255.255.255.0", - default_gateway="192.168.1.1" + hostname="web_server", ip_address="192.168.1.12", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" ) web_server.power_on() network.connect(endpoint_b=web_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[2]) @@ -178,17 +155,14 @@ def arcd_uc2_network() -> Network: hostname="database_server", ip_address="192.168.1.14", subnet_mask="255.255.255.0", - default_gateway="192.168.1.1" + default_gateway="192.168.1.1", ) database_server.power_on() network.connect(endpoint_b=database_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[3]) # Backup Server backup_server = Server( - hostname="backup_server", - ip_address="192.168.1.16", - subnet_mask="255.255.255.0", - default_gateway="192.168.1.1" + hostname="backup_server", ip_address="192.168.1.16", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" ) backup_server.power_on() network.connect(endpoint_b=backup_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[4]) @@ -198,24 +172,15 @@ def arcd_uc2_network() -> Network: hostname="security_suite", ip_address="192.168.1.110", subnet_mask="255.255.255.0", - default_gateway="192.168.1.1" + default_gateway="192.168.1.1", ) security_suite.power_on() network.connect(endpoint_b=security_suite.ethernet_port[1], endpoint_a=switch_1.switch_ports[7]) security_suite.connect_nic(NIC(ip_address="192.168.10.110", subnet_mask="255.255.255.0")) network.connect(endpoint_b=security_suite.ethernet_port[2], endpoint_a=switch_2.switch_ports[7]) - router_1.acl.add_rule( - action=ACLAction.PERMIT, - src_port=Port.ARP, - dst_port=Port.ARP, - position=22 - ) + router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22) - router_1.acl.add_rule( - action=ACLAction.PERMIT, - protocol=IPProtocol.ICMP, - position=23 - ) + router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23) return network diff --git a/src/primaite/simulator/system/core/sys_log.py b/src/primaite/simulator/system/core/sys_log.py index 5a7bbbfe..e07c28aa 100644 --- a/src/primaite/simulator/system/core/sys_log.py +++ b/src/primaite/simulator/system/core/sys_log.py @@ -1,7 +1,7 @@ import logging from pathlib import Path -from prettytable import PrettyTable, MARKDOWN +from prettytable import MARKDOWN, PrettyTable from primaite.simulator import TEMP_SIM_OUTPUT @@ -55,6 +55,14 @@ class SysLog: self.logger.addFilter(_NotJSONFilter()) def show(self, last_n: int = 10, markdown: bool = False): + """ + Print the Node Sys Log as a table. + + Generate and print PrettyTable instance that shows the Nodes Sys Log, with columns Timestamp, Level, + and Massage. + + :param markdown: Use Markdown style in table output. Defaults to False. + """ table = PrettyTable(["Timestamp", "Level", "Message"]) if markdown: table.set_style(MARKDOWN) diff --git a/tests/integration_tests/network/test_frame_transmission.py b/tests/integration_tests/network/test_frame_transmission.py index 34b76060..85717b25 100644 --- a/tests/integration_tests/network/test_frame_transmission.py +++ b/tests/integration_tests/network/test_frame_transmission.py @@ -1,4 +1,4 @@ -from primaite.simulator.network.hardware.base import Link, NIC, Node, Switch +from primaite.simulator.network.hardware.base import Link, NIC, Node def test_node_to_node_ping(): @@ -20,7 +20,6 @@ def test_node_to_node_ping(): def test_multi_nic(): """Tests that Nodes with multiple NICs can ping each other and the data go across the correct links.""" - # TODO Add actual checks. Manual check performed for now. node_a = Node(hostname="node_a") nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0") node_a.connect_nic(nic_a) @@ -45,41 +44,3 @@ def test_multi_nic(): node_a.ping("192.168.0.11") assert node_c.ping("10.0.0.12") - - -def test_switched_network(): - """Tests a larges network of Nodes and Switches with one node pinging another.""" - # TODO Add actual checks. Manual check performed for now. - pc_a = Node(hostname="pc_a") - nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0") - pc_a.connect_nic(nic_a) - pc_a.power_on() - - pc_b = Node(hostname="pc_b") - nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0") - pc_b.connect_nic(nic_b) - pc_b.power_on() - - pc_c = Node(hostname="pc_c") - nic_c = NIC(ip_address="192.168.0.12", subnet_mask="255.255.255.0") - pc_c.connect_nic(nic_c) - pc_c.power_on() - - pc_d = Node(hostname="pc_d") - nic_d = NIC(ip_address="192.168.0.13", subnet_mask="255.255.255.0") - pc_d.connect_nic(nic_d) - pc_d.power_on() - - switch_1 = Switch(hostname="switch_1", num_ports=6) - switch_1.power_on() - - switch_2 = Switch(hostname="switch_2", num_ports=6) - switch_2.power_on() - - link_nic_a_switch_1 = Link(endpoint_a=nic_a, endpoint_b=switch_1.switch_ports[1]) - link_nic_b_switch_1 = Link(endpoint_a=nic_b, endpoint_b=switch_1.switch_ports[2]) - link_nic_c_switch_2 = Link(endpoint_a=nic_c, endpoint_b=switch_2.switch_ports[1]) - link_nic_d_switch_2 = Link(endpoint_a=nic_d, endpoint_b=switch_2.switch_ports[2]) - link_switch_1_switch_2 = Link(endpoint_a=switch_1.switch_ports[6], endpoint_b=switch_2.switch_ports[6]) - - assert pc_a.ping("192.168.0.13") diff --git a/tests/integration_tests/network/test_nic_link_connection.py b/tests/integration_tests/network/test_nic_link_connection.py index f051d026..228099c6 100644 --- a/tests/integration_tests/network/test_nic_link_connection.py +++ b/tests/integration_tests/network/test_nic_link_connection.py @@ -6,8 +6,5 @@ from primaite.simulator.network.hardware.base import Link, NIC def test_link_fails_with_same_nic(): """Tests Link creation fails with endpoint_a and endpoint_b are the same NIC.""" with pytest.raises(ValueError): - nic_a = NIC( - ip_address="192.168.1.2", - subnet_mask="255.255.255.0" - ) + nic_a = NIC(ip_address="192.168.1.2", subnet_mask="255.255.255.0") Link(endpoint_a=nic_a, endpoint_b=nic_a) diff --git a/tests/integration_tests/network/test_switched_network.py b/tests/integration_tests/network/test_switched_network.py new file mode 100644 index 00000000..dc7742f4 --- /dev/null +++ b/tests/integration_tests/network/test_switched_network.py @@ -0,0 +1,25 @@ +from primaite.simulator.network.hardware.base import Link +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 + + +def test_switched_network(): + """Tests a node can ping another node via the switch.""" + client_1 = Computer( + hostname="client_1", ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.0" + ) + client_1.power_on() + + server_1 = Server( + hostname=" server_1", ip_address="192.168.1.11", subnet_mask="255.255.255.0", default_gateway="192.168.1.11" + ) + server_1.power_on() + + switch_1 = Switch(hostname="switch_1", num_ports=6) + switch_1.power_on() + + Link(endpoint_a=client_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[1]) + Link(endpoint_a=server_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[2]) + + assert client_1.ping("192.168.1.11") diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_router.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_acl.py similarity index 84% rename from tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_router.py rename to tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_acl.py index 48d0fc06..99736421 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_router.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_acl.py @@ -1,12 +1,13 @@ from ipaddress import IPv4Address -from primaite.simulator.network.hardware.nodes.router import AccessControlList, ACLAction, ACLRule +from primaite.simulator.network.hardware.nodes.router import ACLAction, Router from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port def test_add_rule(): - acl = AccessControlList() + router = Router("Router") + acl = router.acl acl.add_rule( action=ACLAction.PERMIT, protocol=IPProtocol.TCP, @@ -25,7 +26,8 @@ def test_add_rule(): def test_remove_rule(): - acl = AccessControlList() + router = Router("Router") + acl = router.acl acl.add_rule( action=ACLAction.PERMIT, protocol=IPProtocol.TCP, @@ -40,7 +42,8 @@ def test_remove_rule(): def test_rules(): - acl = AccessControlList() + router = Router("Router") + acl = router.acl acl.add_rule( action=ACLAction.PERMIT, protocol=IPProtocol.TCP, @@ -59,24 +62,27 @@ def test_rules(): dst_port=Port(80), position=2, ) - assert acl.is_permitted( + is_permitted, rule = acl.is_permitted( protocol=IPProtocol.TCP, src_ip=IPv4Address("192.168.1.1"), src_port=Port(8080), dst_ip=IPv4Address("192.168.1.2"), dst_port=Port(80), ) - assert not acl.is_permitted( + assert is_permitted + is_permitted, rule = acl.is_permitted( protocol=IPProtocol.TCP, src_ip=IPv4Address("192.168.1.3"), src_port=Port(8080), dst_ip=IPv4Address("192.168.1.4"), dst_port=Port(80), ) + assert not is_permitted def test_default_rule(): - acl = AccessControlList() + router = Router("Router") + acl = router.acl acl.add_rule( action=ACLAction.PERMIT, protocol=IPProtocol.TCP, @@ -95,10 +101,11 @@ def test_default_rule(): dst_port=Port(80), position=2, ) - assert not acl.is_permitted( + is_permitted, rule = acl.is_permitted( protocol=IPProtocol.UDP, src_ip=IPv4Address("192.168.1.5"), src_port=Port(8080), dst_ip=IPv4Address("192.168.1.12"), dst_port=Port(80), ) + assert not is_permitted From d9feb67e02553e703b6a278b2aadde1104352eb5 Mon Sep 17 00:00:00 2001 From: Christopher McCarthy Date: Mon, 4 Sep 2023 11:20:40 +0000 Subject: [PATCH 07/10] Apply suggestions from code review --- src/primaite/simulator/network/hardware/nodes/router.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 26ba01ae..fa1a0858 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -81,12 +81,11 @@ class AccessControlList(SimComponent): def __init__(self, **kwargs) -> None: if not kwargs.get("implicit_action"): kwargs["implicit_action"] = ACLAction.DENY - if not kwargs.get("max_acl_rules"): - kwargs["max_acl_rules"] = 25 + kwargs["implicit_rule"] = ACLRule(action=kwargs["implicit_action"]) - kwargs["_acl"] = [None] * (kwargs["max_acl_rules"] - 1) super().__init__(**kwargs) + self._acl = [None] * (self.max_acl_rules - 1) def describe_state(self) -> Dict: """ @@ -145,7 +144,7 @@ class AccessControlList(SimComponent): :param int position: The position of the rule to be removed. :raises ValueError: When the position is out of bounds. """ - if 0 <= position < self.max_acl_rules: + if 0 <= position < self.max_acl_rules - 1: self._acl[position] = None else: raise ValueError(f"Position {position} is out of bounds.") From 3075d1985b20463e7deddcb02462c5f113aef4ff Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Mon, 4 Sep 2023 14:58:34 +0100 Subject: [PATCH 08/10] #1800 - Renamed all ip fields so that they're post-fixed with ip_address --- .../network/transport_to_data_link_layer.rst | 12 +- .../network_simulator_demo.ipynb | 2 +- src/primaite/simulator/core.py | 2 +- .../simulator/network/hardware/base.py | 56 +++---- .../network/hardware/nodes/computer.py | 4 - .../network/hardware/nodes/router.py | 140 +++++++++--------- .../simulator/network/protocols/arp.py | 20 +-- .../network/transmission/data_link_layer.py | 4 +- .../network/transmission/network_layer.py | 20 +-- .../simulator/system/core/session_manager.py | 18 +-- .../_network/_hardware/nodes/test_acl.py | 40 ++--- .../_transmission/test_data_link_layer.py | 14 +- 12 files changed, 166 insertions(+), 166 deletions(-) diff --git a/docs/source/simulation_components/network/transport_to_data_link_layer.rst b/docs/source/simulation_components/network/transport_to_data_link_layer.rst index 4961d337..0220ec45 100644 --- a/docs/source/simulation_components/network/transport_to_data_link_layer.rst +++ b/docs/source/simulation_components/network/transport_to_data_link_layer.rst @@ -64,9 +64,9 @@ Data Link Layer (Layer 2) - **request:** ARP operation. Set to True for a request and False for a reply. - **sender_mac_addr:** Sender's MAC address. - - **sender_ip:** Sender's IP address (IPv4 format). + - **sender_ip_address:** Sender's IP address (IPv4 format). - **target_mac_addr:** Target's MAC address. - - **target_ip:** Target's IP address (IPv4 format). + - **target_ip_address:** Target's IP address (IPv4 format). **EthernetHeader:** Represents the Ethernet layer of a network frame. It includes source and destination MAC addresses. This header is used to identify the physical hardware addresses of devices on a local network. @@ -102,8 +102,8 @@ address of 'aa:bb:cc:dd:ee:ff' to port 8080 on the host 10.0.0.10 which has a NI # Network Layer ip_packet = IPPacket( - src_ip="192.168.0.100", - dst_ip="10.0.0.10", + src_ip_address="192.168.0.100", + dst_ip_address="10.0.0.10", protocol=IPProtocol.TCP ) # Data Link Layer @@ -128,8 +128,8 @@ This produces the following ``Frame`` (displayed in json format) "dst_mac_addr": "11:22:33:44:55:66" }, "ip": { - "src_ip": "192.168.0.100", - "dst_ip": "10.0.0.10", + "src_ip_address": "192.168.0.100", + "dst_ip_address": "10.0.0.10", "protocol": "tcp", "ttl": 64, "precedence": 0 diff --git a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb index 252f31fa..b537f54b 100644 --- a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb +++ b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb @@ -554,7 +554,7 @@ "network.get_node_by_hostname(\"router_1\").acl.add_rule(\n", " action=ACLAction.DENY,\n", " protocol=IPProtocol.ICMP,\n", - " src_ip=\"192.168.10.22\",\n", + " src_ip_address=\"192.168.10.22\",\n", " position=1\n", ")" ] diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index b7dfcf72..3e68ed5f 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from typing import Callable, Dict, List, Optional, Union from uuid import uuid4 -from pydantic import BaseModel, ConfigDict, Extra +from pydantic import BaseModel, ConfigDict, Extra, validator from primaite import getLogger diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 1193f3ef..a170506b 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -579,9 +579,13 @@ class ARPCache: """ Add an ARP entry to the cache. + If an entry for the given IP address already exists, the entry is only updated if the `override` parameter is + set to True. + :param ip_address: The IP address to be added to the cache. :param mac_address: The MAC address associated with the IP address. :param nic: The NIC through which the NIC with the IP address is reachable. + :param override: If True, an existing entry for the IP address will be overridden. Default is False. """ for _nic in self.nics.values(): if _nic.ip_address == ip_address: @@ -644,13 +648,13 @@ class ARPCache: # Network Layer ip_packet = IPPacket( - src_ip=nic.ip_address, - dst_ip=target_ip_address, + src_ip_address=nic.ip_address, + dst_ip_address=target_ip_address, ) # Data Link Layer ethernet_header = EthernetHeader(src_mac_addr=nic.mac_address, dst_mac_addr="ff:ff:ff:ff:ff:ff") arp_packet = ARPPacket( - sender_ip=nic.ip_address, sender_mac_addr=nic.mac_address, target_ip=target_ip_address + sender_ip_address=nic.ip_address, sender_mac_addr=nic.mac_address, target_ip_address=target_ip_address ) frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_packet) nic.send_frame(frame) @@ -663,14 +667,14 @@ class ARPCache: :param from_nic: The NIC to send the ARP reply from. """ self.sys_log.info( - f"Sending ARP reply from {arp_reply.sender_mac_addr}/{arp_reply.sender_ip} " - f"to {arp_reply.target_ip}/{arp_reply.target_mac_addr} " + f"Sending ARP reply from {arp_reply.sender_mac_addr}/{arp_reply.sender_ip_address} " + f"to {arp_reply.target_ip_address}/{arp_reply.target_mac_addr} " ) tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) ip_packet = IPPacket( - src_ip=arp_reply.sender_ip, - dst_ip=arp_reply.target_ip, + src_ip_address=arp_reply.sender_ip_address, + dst_ip_address=arp_reply.target_ip_address, ) ethernet_header = EthernetHeader(src_mac_addr=arp_reply.sender_mac_addr, dst_mac_addr=arp_reply.target_mac_addr) @@ -691,26 +695,26 @@ class ARPCache: # ARP Reply if not arp_packet.request: self.sys_log.info( - f"Received ARP response for {arp_packet.sender_ip} from {arp_packet.sender_mac_addr} via NIC {from_nic}" + f"Received ARP response for {arp_packet.sender_ip_address} from {arp_packet.sender_mac_addr} via NIC {from_nic}" ) self.add_arp_cache_entry( - ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic + ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, nic=from_nic ) return # ARP Request self.sys_log.info( - f"Received ARP request for {arp_packet.target_ip} from " - f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip} " + f"Received ARP request for {arp_packet.target_ip_address} from " + f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip_address} " ) # Unmatched ARP Request - if arp_packet.target_ip != from_nic.ip_address: - self.sys_log.info(f"Ignoring ARP request for {arp_packet.target_ip}") + if arp_packet.target_ip_address != from_nic.ip_address: + self.sys_log.info(f"Ignoring ARP request for {arp_packet.target_ip_address}") return # Matched ARP request - self.add_arp_cache_entry(ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic) + self.add_arp_cache_entry(ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, nic=from_nic) arp_packet = arp_packet.generate_reply(from_nic.mac_address) self.send_arp_reply(arp_packet, from_nic) @@ -744,18 +748,18 @@ class ICMP: """ if frame.icmp.icmp_type == ICMPType.ECHO_REQUEST: if not is_reattempt: - self.sys_log.info(f"Received echo request from {frame.ip.src_ip}") - target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip) + self.sys_log.info(f"Received echo request from {frame.ip.src_ip_address}") + target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip_address) - src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip) + src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip_address) if not src_nic: - self.arp.send_arp_request(frame.ip.src_ip) + self.arp.send_arp_request(frame.ip.src_ip_address) self.process_icmp(frame=frame, from_nic=from_nic, is_reattempt=True) return tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) # Network Layer - ip_packet = IPPacket(src_ip=src_nic.ip_address, dst_ip=frame.ip.src_ip, protocol=IPProtocol.ICMP) + ip_packet = IPPacket(src_ip_address=src_nic.ip_address, dst_ip_address=frame.ip.src_ip_address, protocol=IPProtocol.ICMP) # Data Link Layer ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) icmp_reply_packet = ICMPPacket( @@ -768,14 +772,14 @@ class ICMP: frame = Frame( ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet, payload=payload ) - self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") + self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip_address}") src_nic.send_frame(frame) elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY: time = frame.transmission_duration() time_str = f"{time}ms" if time > 0 else "<1ms" self.sys_log.info( - f"Reply from {frame.ip.src_ip}: " + f"Reply from {frame.ip.src_ip_address}: " f"bytes={len(frame.payload)}, " f"time={time_str}, " f"TTL={frame.ip.ttl}" @@ -821,8 +825,8 @@ class ICMP: # Network Layer ip_packet = IPPacket( - src_ip=nic.ip_address, - dst_ip=target_ip_address, + src_ip_address=nic.ip_address, + dst_ip_address=target_ip_address, protocol=IPProtocol.ICMP, ) # Data Link Layer @@ -1059,7 +1063,7 @@ class Node(SimComponent): :param frame: The Frame to be sent. """ - nic: NIC = self._get_arp_cache_nic(frame.ip.dst_ip) + nic: NIC = self._get_arp_cache_nic(frame.ip.dst_ip_address) nic.send_frame(frame) def receive_frame(self, frame: Frame, from_nic: NIC): @@ -1073,9 +1077,9 @@ class Node(SimComponent): :param from_nic: The NIC that received the frame. """ if frame.ip: - if frame.ip.src_ip in self.arp: + if frame.ip.src_ip_address in self.arp: self.arp.add_arp_cache_entry( - ip_address=frame.ip.src_ip, mac_address=frame.ethernet.src_mac_addr, nic=from_nic + ip_address=frame.ip.src_ip_address, mac_address=frame.ethernet.src_mac_addr, nic=from_nic ) if frame.ip.protocol == IPProtocol.TCP: if frame.tcp.src_port == Port.ARP: diff --git a/src/primaite/simulator/network/hardware/nodes/computer.py b/src/primaite/simulator/network/hardware/nodes/computer.py index 2a2e8524..a6def4eb 100644 --- a/src/primaite/simulator/network/hardware/nodes/computer.py +++ b/src/primaite/simulator/network/hardware/nodes/computer.py @@ -36,9 +36,5 @@ class Computer(Node): """ def __init__(self, **kwargs): - for key in {"ip_address", "subnet_mask", "default_gateway"}: - if key in kwargs: - if not isinstance(kwargs[key], IPv4Address): - kwargs[key] = IPv4Address(kwargs[key]) super().__init__(**kwargs) self.connect_nic(NIC(ip_address=kwargs["ip_address"], subnet_mask=kwargs["subnet_mask"])) diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 26ba01ae..0dd4aaff 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -28,17 +28,17 @@ class ACLRule(SimComponent): :ivar ACLAction action: Action to be performed (Permit/Deny). Default is DENY. :ivar Optional[IPProtocol] protocol: Network protocol. Default is None. - :ivar Optional[IPv4Address] src_ip: Source IP address. Default is None. + :ivar Optional[IPv4Address] src_ip_address: Source IP address. Default is None. :ivar Optional[Port] src_port: Source port number. Default is None. - :ivar Optional[IPv4Address] dst_ip: Destination IP address. Default is None. + :ivar Optional[IPv4Address] dst_ip_address: Destination IP address. Default is None. :ivar Optional[Port] dst_port: Destination port number. Default is None. """ action: ACLAction = ACLAction.DENY protocol: Optional[IPProtocol] = None - src_ip: Optional[IPv4Address] = None + src_ip_address: Optional[IPv4Address] = None src_port: Optional[Port] = None - dst_ip: Optional[IPv4Address] = None + dst_ip_address: Optional[IPv4Address] = None dst_port: Optional[Port] = None def __str__(self) -> str: @@ -109,9 +109,9 @@ class AccessControlList(SimComponent): self, action: ACLAction, protocol: Optional[IPProtocol] = None, - src_ip: Optional[Union[str, IPv4Address]] = None, + src_ip_address: Optional[Union[str, IPv4Address]] = None, src_port: Optional[Port] = None, - dst_ip: Optional[Union[str, IPv4Address]] = None, + dst_ip_address: Optional[Union[str, IPv4Address]] = None, dst_port: Optional[Port] = None, position: int = 0, ) -> None: @@ -120,20 +120,20 @@ class AccessControlList(SimComponent): :param ACLAction action: Action to be performed (Permit/Deny). :param Optional[IPProtocol] protocol: Network protocol. - :param Optional[Union[str, IPv4Address]] src_ip: Source IP address. + :param Optional[Union[str, IPv4Address]] src_ip_address: Source IP address. :param Optional[Port] src_port: Source port number. - :param Optional[Union[str, IPv4Address]] dst_ip: Destination IP address. + :param Optional[Union[str, IPv4Address]] dst_ip_address: Destination IP address. :param Optional[Port] dst_port: Destination port number. :param int position: Position in the ACL list to insert the rule. :raises ValueError: When the position is out of bounds. """ - if isinstance(src_ip, str): - src_ip = IPv4Address(src_ip) - if isinstance(dst_ip, str): - dst_ip = IPv4Address(dst_ip) + if isinstance(src_ip_address, str): + src_ip_address = IPv4Address(src_ip_address) + if isinstance(dst_ip_address, str): + dst_ip_address = IPv4Address(dst_ip_address) if 0 <= position < self.max_acl_rules: self._acl[position] = ACLRule( - action=action, src_ip=src_ip, dst_ip=dst_ip, protocol=protocol, src_port=src_port, dst_port=dst_port + action=action, src_ip_address=src_ip_address, dst_ip_address=dst_ip_address, protocol=protocol, src_port=src_port, dst_port=dst_port ) else: raise ValueError(f"Position {position} is out of bounds.") @@ -153,33 +153,33 @@ class AccessControlList(SimComponent): def is_permitted( self, protocol: IPProtocol, - src_ip: Union[str, IPv4Address], + src_ip_address: Union[str, IPv4Address], src_port: Optional[Port], - dst_ip: Union[str, IPv4Address], + 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. :param protocol: The protocol of the packet. - :param src_ip: Source IP address of the packet. Accepts string and IPv4Address. + :param src_ip_address: Source IP address of the packet. Accepts string and IPv4Address. :param src_port: Source port of the packet. Optional. - :param dst_ip: Destination IP address of the packet. Accepts string and IPv4Address. + :param dst_ip_address: Destination IP address of the packet. Accepts string and IPv4Address. :param dst_port: Destination port of the packet. Optional. :return: A tuple with a boolean indicating if the packet is permitted and an optional rule or implicit action string. """ - if not isinstance(src_ip, IPv4Address): - src_ip = IPv4Address(src_ip) - if not isinstance(dst_ip, IPv4Address): - dst_ip = IPv4Address(dst_ip) + if not isinstance(src_ip_address, IPv4Address): + src_ip_address = IPv4Address(src_ip_address) + if not isinstance(dst_ip_address, IPv4Address): + dst_ip_address = IPv4Address(dst_ip_address) for rule in self._acl: if not rule: continue if ( - (rule.src_ip == src_ip or rule.src_ip is None) - and (rule.dst_ip == dst_ip or rule.dst_ip 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) @@ -191,33 +191,33 @@ class AccessControlList(SimComponent): def get_relevant_rules( self, protocol: IPProtocol, - src_ip: Union[str, IPv4Address], + src_ip_address: Union[str, IPv4Address], src_port: Port, - dst_ip: Union[str, IPv4Address], + dst_ip_address: Union[str, IPv4Address], dst_port: Port, ) -> List[ACLRule]: """ Get the list of relevant rules for a packet with given properties. :param protocol: The protocol of the packet. - :param src_ip: Source IP address of the packet. Accepts string and IPv4Address. + :param src_ip_address: Source IP address of the packet. Accepts string and IPv4Address. :param src_port: Source port of the packet. - :param dst_ip: Destination IP address of the packet. Accepts string and IPv4Address. + :param dst_ip_address: Destination IP address of the packet. Accepts string and IPv4Address. :param dst_port: Destination port of the packet. :return: A list of relevant ACLRules. """ - if not isinstance(src_ip, IPv4Address): - src_ip = IPv4Address(src_ip) - if not isinstance(dst_ip, IPv4Address): - dst_ip = IPv4Address(dst_ip) + if not isinstance(src_ip_address, IPv4Address): + src_ip_address = IPv4Address(src_ip_address) + if not isinstance(dst_ip_address, IPv4Address): + dst_ip_address = IPv4Address(dst_ip_address) relevant_rules = [] for rule in self._acl: if rule is None: continue if ( - (rule.src_ip == src_ip or rule.src_ip is None) - or (rule.dst_ip == dst_ip or rule.dst_ip 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) @@ -244,9 +244,9 @@ class AccessControlList(SimComponent): index, rule.action.name if rule.action else "ANY", rule.protocol.name if rule.protocol else "ANY", - rule.src_ip if rule.src_ip else "ANY", + rule.src_ip_address if rule.src_ip_address else "ANY", f"{rule.src_port.value} ({rule.src_port.name})" if rule.src_port else "ANY", - rule.dst_ip if rule.dst_ip else "ANY", + rule.dst_ip_address if rule.dst_ip_address else "ANY", f"{rule.dst_port.value} ({rule.dst_port.name})" if rule.dst_port else "ANY", ] ) @@ -260,7 +260,7 @@ class RouteEntry(SimComponent): Attributes: address (IPv4Address): The destination IP address or network address. subnet_mask (IPv4Address): The subnet mask for the network. - next_hop (IPv4Address): The next hop IP address to which packets should be forwarded. + next_hop_ip_address (IPv4Address): The next hop IP address to which packets should be forwarded. metric (int): The cost metric for this route. Default is 0.0. Example: @@ -276,13 +276,13 @@ class RouteEntry(SimComponent): "The destination IP address or network address." subnet_mask: IPv4Address "The subnet mask for the network." - next_hop: IPv4Address + next_hop_ip_address: IPv4Address "The next hop IP address to which packets should be forwarded." metric: float = 0.0 "The cost metric for this route. Default is 0.0." def __init__(self, **kwargs): - for key in {"address", "subnet_mask", "next_hop"}: + for key in {"address", "subnet_mask", "next_hop_ip_address"}: if not isinstance(kwargs[key], IPv4Address): kwargs[key] = IPv4Address(kwargs[key]) super().__init__(**kwargs) @@ -330,7 +330,7 @@ class RouteTable(SimComponent): self, address: Union[IPv4Address, str], subnet_mask: Union[IPv4Address, str], - next_hop: Union[IPv4Address, str], + next_hop_ip_address: Union[IPv4Address, str], metric: float = 0.0, ): """ @@ -338,13 +338,13 @@ class RouteTable(SimComponent): :param address: The destination address of the route. :param subnet_mask: The subnet mask of the route. - :param next_hop: The next hop IP for the route. + :param next_hop_ip_address: The next hop IP for the route. :param metric: The metric of the route, default is 0.0. """ - for key in {address, subnet_mask, next_hop}: + for key in {address, subnet_mask, next_hop_ip_address}: if not isinstance(key, IPv4Address): key = IPv4Address(key) - route = RouteEntry(address=address, subnet_mask=subnet_mask, next_hop=next_hop, metric=metric) + route = RouteEntry(address=address, subnet_mask=subnet_mask, next_hop_ip_address=next_hop_ip_address, metric=metric) self.routes.append(route) def find_best_route(self, destination_ip: Union[str, IPv4Address]) -> Optional[RouteEntry]: @@ -387,7 +387,7 @@ class RouteTable(SimComponent): table.title = f"{self.sys_log.hostname} Route Table" for index, route in enumerate(self.routes): network = IPv4Network(f"{route.address}/{route.subnet_mask}") - table.add_row([index, f"{route.address}/{network.prefixlen}", route.next_hop, route.metric]) + table.add_row([index, f"{route.address}/{network.prefixlen}", route.next_hop_ip_address, route.metric]) print(table) @@ -415,40 +415,40 @@ class RouterARPCache(ARPCache): # ARP Reply if not arp_packet.request: for nic in self.router.nics.values(): - if arp_packet.target_ip == nic.ip_address: + if arp_packet.target_ip_address == nic.ip_address: # reply to the Router specifically self.sys_log.info( - f"Received ARP response for {arp_packet.sender_ip} " + f"Received ARP response for {arp_packet.sender_ip_address} " f"from {arp_packet.sender_mac_addr} via NIC {from_nic}" ) self.add_arp_cache_entry( - ip_address=arp_packet.sender_ip, + ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, nic=from_nic, ) return # Reply for a connected requested - nic = self.get_arp_cache_nic(arp_packet.target_ip) + nic = self.get_arp_cache_nic(arp_packet.target_ip_address) if nic: - self.sys_log.info(f"Forwarding arp reply for {arp_packet.target_ip}, from {arp_packet.sender_ip}") + self.sys_log.info(f"Forwarding arp reply for {arp_packet.target_ip_address}, from {arp_packet.sender_ip_address}") arp_packet.sender_mac_addr = nic.mac_address frame.decrement_ttl() nic.send_frame(frame) # ARP Request self.sys_log.info( - f"Received ARP request for {arp_packet.target_ip} from " - f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip} " + f"Received ARP request for {arp_packet.target_ip_address} from " + f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip_address} " ) # Matched ARP request - self.add_arp_cache_entry(ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic) + self.add_arp_cache_entry(ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, nic=from_nic) arp_packet = arp_packet.generate_reply(from_nic.mac_address) self.send_arp_reply(arp_packet, from_nic) # If the target IP matches one of the router's NICs for nic in self.nics.values(): - if nic.enabled and nic.ip_address == arp_packet.target_ip: + if nic.enabled and nic.ip_address == arp_packet.target_ip_address: arp_reply = arp_packet.generate_reply(from_nic.mac_address) self.send_arp_reply(arp_reply, from_nic) return @@ -484,17 +484,17 @@ class RouterICMP(ICMP): # determine if request is for router interface or whether it needs to be routed for nic in self.router.nics.values(): - if nic.ip_address == frame.ip.dst_ip: + if nic.ip_address == frame.ip.dst_ip_address: if nic.enabled: # reply to the request if not is_reattempt: - self.sys_log.info(f"Received echo request from {frame.ip.src_ip}") - target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip) - src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip) + self.sys_log.info(f"Received echo request from {frame.ip.src_ip_address}") + target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip_address) + src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip_address) tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) # Network Layer - ip_packet = IPPacket(src_ip=nic.ip_address, dst_ip=frame.ip.src_ip, protocol=IPProtocol.ICMP) + ip_packet = IPPacket(src_ip_address=nic.ip_address, dst_ip_address=frame.ip.src_ip_address, protocol=IPProtocol.ICMP) # Data Link Layer ethernet_header = EthernetHeader( src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address @@ -513,7 +513,7 @@ class RouterICMP(ICMP): icmp=icmp_reply_packet, payload=payload, ) - self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") + self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip_address}") src_nic.send_frame(frame) return @@ -523,12 +523,12 @@ class RouterICMP(ICMP): elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY: for nic in self.router.nics.values(): - if nic.ip_address == frame.ip.dst_ip: + if nic.ip_address == frame.ip.dst_ip_address: if nic.enabled: time = frame.transmission_duration() time_str = f"{time}ms" if time > 0 else "<1ms" self.sys_log.info( - f"Reply from {frame.ip.src_ip}: " + f"Reply from {frame.ip.src_ip_address}: " f"bytes={len(frame.payload)}, " f"time={time_str}, " f"TTL={frame.ip.ttl}" @@ -606,22 +606,22 @@ class Router(Node): :param re_attempt: Flag to indicate if the routing is a reattempt. """ # Check if src ip is on network of one of the NICs - nic = self.arp.get_arp_cache_nic(frame.ip.dst_ip) - target_mac = self.arp.get_arp_cache_mac_address(frame.ip.dst_ip) + nic = self.arp.get_arp_cache_nic(frame.ip.dst_ip_address) + target_mac = self.arp.get_arp_cache_mac_address(frame.ip.dst_ip_address) if re_attempt and not nic: - self.sys_log.info(f"Destination {frame.ip.dst_ip} is unreachable") + self.sys_log.info(f"Destination {frame.ip.dst_ip_address} is unreachable") return if not nic: - self.arp.send_arp_request(frame.ip.dst_ip) + self.arp.send_arp_request(frame.ip.dst_ip_address) return self.route_frame(frame=frame, from_nic=from_nic, re_attempt=True) if not nic.enabled: # TODO: Add sys_log here return - if frame.ip.dst_ip in nic.ip_network: + if frame.ip.dst_ip_address in nic.ip_network: from_port = self._get_port_of_nic(from_nic) to_port = self._get_port_of_nic(nic) self.sys_log.info(f"Routing frame to internally from port {from_port} to port {to_port}") @@ -643,8 +643,8 @@ class Router(Node): """ route_frame = False protocol = frame.ip.protocol - src_ip = frame.ip.src_ip - dst_ip = frame.ip.dst_ip + src_ip_address = frame.ip.src_ip_address + dst_ip_address = frame.ip.dst_ip_address src_port = None dst_port = None if frame.ip.protocol == IPProtocol.TCP: @@ -656,14 +656,14 @@ class Router(Node): # Check if it's permitted permitted, rule = self.acl.is_permitted( - protocol=protocol, src_ip=src_ip, src_port=src_port, dst_ip=dst_ip, dst_port=dst_port + protocol=protocol, src_ip_address=src_ip_address, src_port=src_port, dst_ip_address=dst_ip_address, dst_port=dst_port ) if not permitted: at_port = self._get_port_of_nic(from_nic) self.sys_log.info(f"Frame blocked at port {at_port} by rule {rule}") return - if not self.arp.get_arp_cache_nic(src_ip): - self.arp.add_arp_cache_entry(src_ip, frame.ethernet.src_mac_addr, from_nic) + if not self.arp.get_arp_cache_nic(src_ip_address): + self.arp.add_arp_cache_entry(src_ip_address, frame.ethernet.src_mac_addr, from_nic) if frame.ip.protocol == IPProtocol.ICMP: self.icmp.process_icmp(frame=frame, from_nic=from_nic) else: diff --git a/src/primaite/simulator/network/protocols/arp.py b/src/primaite/simulator/network/protocols/arp.py index bae14d28..5e38cc66 100644 --- a/src/primaite/simulator/network/protocols/arp.py +++ b/src/primaite/simulator/network/protocols/arp.py @@ -24,21 +24,21 @@ class ARPPacket(BaseModel): :param request: ARP operation. True if a request, False if a reply. :param sender_mac_addr: Sender MAC address. - :param sender_ip: Sender IP address. + :param sender_ip_address: Sender IP address. :param target_mac_addr: Target MAC address. - :param target_ip: Target IP address. + :param target_ip_address: Target IP address. :Example: >>> arp_request = ARPPacket( ... sender_mac_addr="aa:bb:cc:dd:ee:ff", - ... sender_ip=IPv4Address("192.168.0.1"), - ... target_ip=IPv4Address("192.168.0.2") + ... sender_ip_address=IPv4Address("192.168.0.1"), + ... target_ip_address=IPv4Address("192.168.0.2") ... ) >>> arp_response = ARPPacket( ... sender_mac_addr="aa:bb:cc:dd:ee:ff", - ... sender_ip=IPv4Address("192.168.0.1"), - ... target_ip=IPv4Address("192.168.0.2") + ... sender_ip_address=IPv4Address("192.168.0.1"), + ... target_ip_address=IPv4Address("192.168.0.2") ... ) """ @@ -46,11 +46,11 @@ class ARPPacket(BaseModel): "ARP operation. True if a request, False if a reply." sender_mac_addr: str "Sender MAC address." - sender_ip: IPv4Address + sender_ip_address: IPv4Address "Sender IP address." target_mac_addr: Optional[str] = None "Target MAC address." - target_ip: IPv4Address + target_ip_address: IPv4Address "Target IP address." def generate_reply(self, mac_address: str) -> ARPPacket: @@ -62,8 +62,8 @@ class ARPPacket(BaseModel): """ return ARPPacket( request=False, - sender_ip=self.target_ip, + sender_ip_address=self.target_ip_address, sender_mac_addr=mac_address, - target_ip=self.sender_ip, + target_ip_address=self.sender_ip_address, target_mac_addr=self.sender_mac_addr, ) diff --git a/src/primaite/simulator/network/transmission/data_link_layer.py b/src/primaite/simulator/network/transmission/data_link_layer.py index ddd9fad3..b7986622 100644 --- a/src/primaite/simulator/network/transmission/data_link_layer.py +++ b/src/primaite/simulator/network/transmission/data_link_layer.py @@ -52,8 +52,8 @@ class Frame(BaseModel): ... dst_mac_addr='11:22:33:44:55:66' ... ), ... ip=IPPacket( - ... src_ip=IPv4Address('192.168.0.1'), - ... dst_ip=IPv4Address('10.0.0.1'), + ... src_ip_address=IPv4Address('192.168.0.1'), + ... dst_ip_address=IPv4Address('10.0.0.1'), ... ), ... tcp=TCPHeader( ... src_port=8080, diff --git a/src/primaite/simulator/network/transmission/network_layer.py b/src/primaite/simulator/network/transmission/network_layer.py index afd1ecef..fd36fbf8 100644 --- a/src/primaite/simulator/network/transmission/network_layer.py +++ b/src/primaite/simulator/network/transmission/network_layer.py @@ -162,8 +162,8 @@ class IPPacket(BaseModel): """ Represents the IP layer of a network frame. - :param src_ip: Source IP address. - :param dst_ip: Destination IP address. + :param src_ip_address: Source IP address. + :param dst_ip_address: Destination IP address. :param protocol: The IP protocol (default is TCP). :param ttl: Time to Live (TTL) for the packet. :param precedence: Precedence level for Quality of Service (QoS). @@ -172,17 +172,17 @@ class IPPacket(BaseModel): >>> from ipaddress import IPv4Address >>> ip_packet = IPPacket( - ... src_ip=IPv4Address('192.168.0.1'), - ... dst_ip=IPv4Address('10.0.0.1'), + ... src_ip_address=IPv4Address('192.168.0.1'), + ... dst_ip_address=IPv4Address('10.0.0.1'), ... protocol=IPProtocol.TCP, ... ttl=64, ... precedence=Precedence.CRITICAL ... ) """ - src_ip: IPv4Address + src_ip_address: IPv4Address "Source IP address." - dst_ip: IPv4Address + dst_ip_address: IPv4Address "Destination IP address." protocol: IPProtocol = IPProtocol.TCP "IPProtocol." @@ -192,8 +192,8 @@ class IPPacket(BaseModel): "Precedence level for Quality of Service (default is Precedence.ROUTINE)." def __init__(self, **kwargs): - if not isinstance(kwargs["src_ip"], IPv4Address): - kwargs["src_ip"] = IPv4Address(kwargs["src_ip"]) - if not isinstance(kwargs["dst_ip"], IPv4Address): - kwargs["dst_ip"] = IPv4Address(kwargs["dst_ip"]) + if not isinstance(kwargs["src_ip_address"], IPv4Address): + kwargs["src_ip_address"] = IPv4Address(kwargs["src_ip_address"]) + if not isinstance(kwargs["dst_ip_address"], IPv4Address): + kwargs["dst_ip_address"] = IPv4Address(kwargs["dst_ip_address"]) super().__init__(**kwargs) diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index fe7b06b2..705210ff 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/src/primaite/simulator/system/core/session_manager.py @@ -22,16 +22,16 @@ class Session(SimComponent): source and destination IPs and ports. :param protocol: The IP protocol used in the session. - :param src_ip: The source IP address. - :param dst_ip: The destination IP address. + :param src_ip_address: The source IP address. + :param dst_ip_address: The destination IP address. :param src_port: The source port number (optional). :param dst_port: The destination port number (optional). :param connected: A flag indicating whether the session is connected. """ protocol: IPProtocol - src_ip: IPv4Address - dst_ip: IPv4Address + src_ip_address: IPv4Address + dst_ip_address: IPv4Address src_port: Optional[Port] dst_port: Optional[Port] connected: bool = False @@ -46,8 +46,8 @@ class Session(SimComponent): :param session_key: Tuple containing the session details. :return: A Session instance. """ - protocol, src_ip, dst_ip, src_port, dst_port = session_key - return Session(protocol=protocol, src_ip=src_ip, dst_ip=dst_ip, src_port=src_port, dst_port=dst_port) + protocol, src_ip_address, dst_ip_address, src_port, dst_port = session_key + return Session(protocol=protocol, src_ip_address=src_ip_address, dst_ip_address=dst_ip_address, src_port=src_port, dst_port=dst_port) def describe_state(self) -> Dict: """ @@ -108,8 +108,8 @@ class SessionManager: :return: A tuple containing the session key. """ protocol = frame.ip.protocol - src_ip = frame.ip.src_ip - dst_ip = frame.ip.dst_ip + src_ip_address = frame.ip.src_ip_address + dst_ip_address = frame.ip.dst_ip_address if protocol == IPProtocol.TCP: if from_source: src_port = frame.tcp.src_port @@ -127,7 +127,7 @@ class SessionManager: else: src_port = None dst_port = None - return protocol, src_ip, dst_ip, src_port, dst_port + return protocol, src_ip_address, dst_ip_address, src_port, dst_port def receive_payload_from_software_manager(self, payload: Any, session_id: Optional[int] = None): """ diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_acl.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_acl.py index 99736421..554cba38 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_acl.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_acl.py @@ -11,17 +11,17 @@ def test_add_rule(): acl.add_rule( action=ACLAction.PERMIT, protocol=IPProtocol.TCP, - src_ip=IPv4Address("192.168.1.1"), + src_ip_address=IPv4Address("192.168.1.1"), src_port=Port(8080), - dst_ip=IPv4Address("192.168.1.2"), + dst_ip_address=IPv4Address("192.168.1.2"), dst_port=Port(80), position=1, ) assert acl.acl[1].action == ACLAction.PERMIT assert acl.acl[1].protocol == IPProtocol.TCP - assert acl.acl[1].src_ip == IPv4Address("192.168.1.1") + assert acl.acl[1].src_ip_address == IPv4Address("192.168.1.1") assert acl.acl[1].src_port == Port(8080) - assert acl.acl[1].dst_ip == IPv4Address("192.168.1.2") + assert acl.acl[1].dst_ip_address == IPv4Address("192.168.1.2") assert acl.acl[1].dst_port == Port(80) @@ -31,9 +31,9 @@ def test_remove_rule(): acl.add_rule( action=ACLAction.PERMIT, protocol=IPProtocol.TCP, - src_ip=IPv4Address("192.168.1.1"), + src_ip_address=IPv4Address("192.168.1.1"), src_port=Port(8080), - dst_ip=IPv4Address("192.168.1.2"), + dst_ip_address=IPv4Address("192.168.1.2"), dst_port=Port(80), position=1, ) @@ -47,34 +47,34 @@ def test_rules(): acl.add_rule( action=ACLAction.PERMIT, protocol=IPProtocol.TCP, - src_ip=IPv4Address("192.168.1.1"), + src_ip_address=IPv4Address("192.168.1.1"), src_port=Port(8080), - dst_ip=IPv4Address("192.168.1.2"), + dst_ip_address=IPv4Address("192.168.1.2"), dst_port=Port(80), position=1, ) acl.add_rule( action=ACLAction.DENY, protocol=IPProtocol.TCP, - src_ip=IPv4Address("192.168.1.3"), + src_ip_address=IPv4Address("192.168.1.3"), src_port=Port(8080), - dst_ip=IPv4Address("192.168.1.4"), + dst_ip_address=IPv4Address("192.168.1.4"), dst_port=Port(80), position=2, ) is_permitted, rule = acl.is_permitted( protocol=IPProtocol.TCP, - src_ip=IPv4Address("192.168.1.1"), + src_ip_address=IPv4Address("192.168.1.1"), src_port=Port(8080), - dst_ip=IPv4Address("192.168.1.2"), + dst_ip_address=IPv4Address("192.168.1.2"), dst_port=Port(80), ) assert is_permitted is_permitted, rule = acl.is_permitted( protocol=IPProtocol.TCP, - src_ip=IPv4Address("192.168.1.3"), + src_ip_address=IPv4Address("192.168.1.3"), src_port=Port(8080), - dst_ip=IPv4Address("192.168.1.4"), + dst_ip_address=IPv4Address("192.168.1.4"), dst_port=Port(80), ) assert not is_permitted @@ -86,26 +86,26 @@ def test_default_rule(): acl.add_rule( action=ACLAction.PERMIT, protocol=IPProtocol.TCP, - src_ip=IPv4Address("192.168.1.1"), + src_ip_address=IPv4Address("192.168.1.1"), src_port=Port(8080), - dst_ip=IPv4Address("192.168.1.2"), + dst_ip_address=IPv4Address("192.168.1.2"), dst_port=Port(80), position=1, ) acl.add_rule( action=ACLAction.DENY, protocol=IPProtocol.TCP, - src_ip=IPv4Address("192.168.1.3"), + src_ip_address=IPv4Address("192.168.1.3"), src_port=Port(8080), - dst_ip=IPv4Address("192.168.1.4"), + dst_ip_address=IPv4Address("192.168.1.4"), dst_port=Port(80), position=2, ) is_permitted, rule = acl.is_permitted( protocol=IPProtocol.UDP, - src_ip=IPv4Address("192.168.1.5"), + src_ip_address=IPv4Address("192.168.1.5"), src_port=Port(8080), - dst_ip=IPv4Address("192.168.1.12"), + dst_ip_address=IPv4Address("192.168.1.12"), dst_port=Port(80), ) assert not is_permitted diff --git a/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_data_link_layer.py b/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_data_link_layer.py index 8a78d1bc..f9b89de5 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_data_link_layer.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_data_link_layer.py @@ -10,7 +10,7 @@ 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"), + ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20"), tcp=TCPHeader( src_port=8080, dst_port=80, @@ -38,7 +38,7 @@ def test_frame_creation_fails_tcp_without_header(): 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), + ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.TCP), ) @@ -47,7 +47,7 @@ def test_frame_creation_fails_udp_without_header(): 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), + ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.UDP), ) @@ -56,7 +56,7 @@ def test_frame_creation_fails_tcp_with_udp_header(): 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), + ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.TCP), udp=UDPHeader(src_port=8080, dst_port=80), ) @@ -66,7 +66,7 @@ def test_frame_creation_fails_udp_with_tcp_header(): 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), + ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.UDP), udp=TCPHeader(src_port=8080, dst_port=80), ) @@ -75,7 +75,7 @@ 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), + ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.ICMP), icmp=ICMPPacket(), ) assert frame @@ -86,5 +86,5 @@ def test_icmp_frame_creation_fails_without_icmp_header(): 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), + ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.ICMP), ) From ccad5ba8a319810f1afe7ddd531c929f41b1d5ec Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Mon, 4 Sep 2023 16:34:55 +0100 Subject: [PATCH 09/10] #1800 - Ran pre-commit --- src/primaite/simulator/core.py | 2 +- .../simulator/network/hardware/base.py | 15 ++++++--- .../network/hardware/nodes/computer.py | 2 -- .../network/hardware/nodes/router.py | 31 +++++++++++++++---- .../simulator/system/core/session_manager.py | 8 ++++- 5 files changed, 44 insertions(+), 14 deletions(-) diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index 3e68ed5f..b7dfcf72 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from typing import Callable, Dict, List, Optional, Union from uuid import uuid4 -from pydantic import BaseModel, ConfigDict, Extra, validator +from pydantic import BaseModel, ConfigDict, Extra from primaite import getLogger diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index a170506b..f2feb961 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -654,7 +654,9 @@ class ARPCache: # Data Link Layer ethernet_header = EthernetHeader(src_mac_addr=nic.mac_address, dst_mac_addr="ff:ff:ff:ff:ff:ff") arp_packet = ARPPacket( - sender_ip_address=nic.ip_address, sender_mac_addr=nic.mac_address, target_ip_address=target_ip_address + sender_ip_address=nic.ip_address, + sender_mac_addr=nic.mac_address, + target_ip_address=target_ip_address, ) frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_packet) nic.send_frame(frame) @@ -695,7 +697,8 @@ class ARPCache: # ARP Reply if not arp_packet.request: self.sys_log.info( - f"Received ARP response for {arp_packet.sender_ip_address} from {arp_packet.sender_mac_addr} via NIC {from_nic}" + f"Received ARP response for {arp_packet.sender_ip_address} " + f"from {arp_packet.sender_mac_addr} via NIC {from_nic}" ) self.add_arp_cache_entry( ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, nic=from_nic @@ -714,7 +717,9 @@ class ARPCache: return # Matched ARP request - self.add_arp_cache_entry(ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, nic=from_nic) + self.add_arp_cache_entry( + ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, nic=from_nic + ) arp_packet = arp_packet.generate_reply(from_nic.mac_address) self.send_arp_reply(arp_packet, from_nic) @@ -759,7 +764,9 @@ class ICMP: tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) # Network Layer - ip_packet = IPPacket(src_ip_address=src_nic.ip_address, dst_ip_address=frame.ip.src_ip_address, protocol=IPProtocol.ICMP) + ip_packet = IPPacket( + src_ip_address=src_nic.ip_address, dst_ip_address=frame.ip.src_ip_address, protocol=IPProtocol.ICMP + ) # Data Link Layer ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) icmp_reply_packet = ICMPPacket( diff --git a/src/primaite/simulator/network/hardware/nodes/computer.py b/src/primaite/simulator/network/hardware/nodes/computer.py index a6def4eb..5452666b 100644 --- a/src/primaite/simulator/network/hardware/nodes/computer.py +++ b/src/primaite/simulator/network/hardware/nodes/computer.py @@ -1,5 +1,3 @@ -from ipaddress import IPv4Address - from primaite.simulator.network.hardware.base import NIC, Node diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index a8177e86..a34b83e2 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -132,7 +132,12 @@ class AccessControlList(SimComponent): dst_ip_address = IPv4Address(dst_ip_address) if 0 <= position < self.max_acl_rules: self._acl[position] = ACLRule( - action=action, src_ip_address=src_ip_address, dst_ip_address=dst_ip_address, protocol=protocol, src_port=src_port, dst_port=dst_port + action=action, + src_ip_address=src_ip_address, + dst_ip_address=dst_ip_address, + protocol=protocol, + src_port=src_port, + dst_port=dst_port, ) else: raise ValueError(f"Position {position} is out of bounds.") @@ -343,7 +348,9 @@ class RouteTable(SimComponent): for key in {address, subnet_mask, next_hop_ip_address}: if not isinstance(key, IPv4Address): key = IPv4Address(key) - route = RouteEntry(address=address, subnet_mask=subnet_mask, next_hop_ip_address=next_hop_ip_address, metric=metric) + route = RouteEntry( + address=address, subnet_mask=subnet_mask, next_hop_ip_address=next_hop_ip_address, metric=metric + ) self.routes.append(route) def find_best_route(self, destination_ip: Union[str, IPv4Address]) -> Optional[RouteEntry]: @@ -430,7 +437,9 @@ class RouterARPCache(ARPCache): # Reply for a connected requested nic = self.get_arp_cache_nic(arp_packet.target_ip_address) if nic: - self.sys_log.info(f"Forwarding arp reply for {arp_packet.target_ip_address}, from {arp_packet.sender_ip_address}") + self.sys_log.info( + f"Forwarding arp reply for {arp_packet.target_ip_address}, from {arp_packet.sender_ip_address}" + ) arp_packet.sender_mac_addr = nic.mac_address frame.decrement_ttl() nic.send_frame(frame) @@ -441,7 +450,9 @@ class RouterARPCache(ARPCache): f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip_address} " ) # Matched ARP request - self.add_arp_cache_entry(ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, nic=from_nic) + self.add_arp_cache_entry( + ip_address=arp_packet.sender_ip_address, mac_address=arp_packet.sender_mac_addr, nic=from_nic + ) arp_packet = arp_packet.generate_reply(from_nic.mac_address) self.send_arp_reply(arp_packet, from_nic) @@ -493,7 +504,11 @@ class RouterICMP(ICMP): tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) # Network Layer - ip_packet = IPPacket(src_ip_address=nic.ip_address, dst_ip_address=frame.ip.src_ip_address, protocol=IPProtocol.ICMP) + ip_packet = IPPacket( + src_ip_address=nic.ip_address, + dst_ip_address=frame.ip.src_ip_address, + protocol=IPProtocol.ICMP, + ) # Data Link Layer ethernet_header = EthernetHeader( src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address @@ -655,7 +670,11 @@ class Router(Node): # Check if it's permitted permitted, rule = self.acl.is_permitted( - protocol=protocol, src_ip_address=src_ip_address, src_port=src_port, dst_ip_address=dst_ip_address, dst_port=dst_port + protocol=protocol, + src_ip_address=src_ip_address, + src_port=src_port, + dst_ip_address=dst_ip_address, + dst_port=dst_port, ) if not permitted: at_port = self._get_port_of_nic(from_nic) diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index 705210ff..7f3d22c5 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/src/primaite/simulator/system/core/session_manager.py @@ -47,7 +47,13 @@ class Session(SimComponent): :return: A Session instance. """ protocol, src_ip_address, dst_ip_address, src_port, dst_port = session_key - return Session(protocol=protocol, src_ip_address=src_ip_address, dst_ip_address=dst_ip_address, src_port=src_port, dst_port=dst_port) + return Session( + protocol=protocol, + src_ip_address=src_ip_address, + dst_ip_address=dst_ip_address, + src_port=src_port, + dst_port=dst_port, + ) def describe_state(self) -> Dict: """ From 596ad20cc6fb814e9327804cd8a3e4b29c72892c Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Mon, 4 Sep 2023 16:44:29 +0100 Subject: [PATCH 10/10] #1800 - Added better logging and error messages to AccessControlList class. Updated usage of extra following pydantic deprecated warning "`pydantic.config.Extra` is deprecated, use literal values instead (e.g. `extra='allow'`)" --- src/primaite/simulator/core.py | 4 ++-- src/primaite/simulator/network/hardware/nodes/router.py | 8 ++++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index b7dfcf72..0501bbb2 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from typing import Callable, Dict, List, Optional, Union from uuid import uuid4 -from pydantic import BaseModel, ConfigDict, Extra +from pydantic import BaseModel, ConfigDict from primaite import getLogger @@ -126,7 +126,7 @@ class ActionManager: class SimComponent(BaseModel): """Extension of pydantic BaseModel with additional methods that must be defined by all classes in the simulator.""" - model_config = ConfigDict(arbitrary_types_allowed=True, extra=Extra.allow) + model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow") """Configure pydantic to allow arbitrary types and to let the instance have attributes not present in model.""" uuid: str diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index a34b83e2..092680a7 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -131,6 +131,8 @@ class AccessControlList(SimComponent): if isinstance(dst_ip_address, str): dst_ip_address = IPv4Address(dst_ip_address) if 0 <= position < self.max_acl_rules: + if self._acl[position]: + self.sys_log.info(f"Overwriting ACL rule at position {position}") self._acl[position] = ACLRule( action=action, src_ip_address=src_ip_address, @@ -140,7 +142,7 @@ class AccessControlList(SimComponent): dst_port=dst_port, ) else: - raise ValueError(f"Position {position} is out of bounds.") + raise ValueError(f"Cannot add ACL rule, position {position} is out of bounds.") def remove_rule(self, position: int) -> None: """ @@ -150,9 +152,11 @@ class AccessControlList(SimComponent): :raises ValueError: When the position is out of bounds. """ if 0 <= position < self.max_acl_rules - 1: + rule = self._acl[position] # noqa self._acl[position] = None + del rule else: - raise ValueError(f"Position {position} is out of bounds.") + raise ValueError(f"Cannot remove ACL rule, position {position} is out of bounds.") def is_permitted( self,