From a0253ce6c44566184d805cbfdbdad1cd1de9f0ce Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Fri, 2 Feb 2024 17:14:34 +0000 Subject: [PATCH] #2248 - TSome further fixess to ARP. Also refactored PCAP to log inbound and outbound frames separately --- .../simulator/network/hardware/base.py | 8 +- .../network/hardware/nodes/router.py | 2 +- .../simulator/system/core/packet_capture.py | 59 +++++--- .../simulator/system/core/session_manager.py | 2 +- .../simulator/system/services/arp/arp.py | 140 +++++++++--------- .../simulator/system/services/arp/host_arp.py | 2 +- 6 files changed, 117 insertions(+), 96 deletions(-) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 69f93f51..9edf7518 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -266,7 +266,7 @@ class NIC(SimComponent): """ if self.enabled: frame.set_sent_timestamp() - self.pcap.capture(frame) + self.pcap.capture_outbound(frame) self._connected_link.transmit_frame(sender_nic=self, frame=frame) return True # Cannot send Frame as the NIC is not enabled @@ -295,7 +295,7 @@ class NIC(SimComponent): self._connected_node.sys_log.info("Frame discarded as TTL limit reached") return False frame.set_received_timestamp() - self.pcap.capture(frame) + self.pcap.capture_inbound(frame) # If this destination or is broadcast accept_frame = False @@ -442,7 +442,7 @@ class SwitchPort(SimComponent): :param frame: The network frame to be sent. """ if self.enabled: - self.pcap.capture(frame) + self.pcap.capture_outbound(frame) self._connected_link.transmit_frame(sender_nic=self, frame=frame) return True # Cannot send Frame as the SwitchPort is not enabled @@ -461,7 +461,7 @@ class SwitchPort(SimComponent): if frame.ip and frame.ip.ttl < 1: self._connected_node.sys_log.info("Frame discarded as TTL limit reached") return False - self.pcap.capture(frame) + self.pcap.capture_inbound(frame) connected_node: Node = self._connected_node connected_node.forward_frame(frame=frame, incoming_port=self) return True diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 34eb0423..69717ae6 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -558,7 +558,7 @@ class RouterNIC(NIC): self._connected_node.sys_log.info("Frame discarded as TTL limit reached") return False frame.set_received_timestamp() - self.pcap.capture(frame) + self.pcap.capture_inbound(frame) # If this destination or is broadcast if frame.ethernet.dst_mac_addr == self.mac_address or frame.ethernet.dst_mac_addr == "ff:ff:ff:ff:ff:ff": self._connected_node.receive_frame(frame=frame, from_nic=self) diff --git a/src/primaite/simulator/system/core/packet_capture.py b/src/primaite/simulator/system/core/packet_capture.py index bfb6a055..d3a14d2a 100644 --- a/src/primaite/simulator/system/core/packet_capture.py +++ b/src/primaite/simulator/system/core/packet_capture.py @@ -35,16 +35,17 @@ class PacketCapture: self.switch_port_number = switch_port_number "The SwitchPort number." + self.inbound_logger = None + self.outbound_logger = None + self.current_episode: int = 1 - self.setup_logger() + self.setup_logger(outbound=False) + self.setup_logger(outbound=True) - def setup_logger(self): + def setup_logger(self, outbound: bool = False): """Set up the logger configuration.""" - if not SIM_OUTPUT.save_pcap_logs: - return - - log_path = self._get_log_path() + log_path = self._get_log_path(outbound) file_handler = logging.FileHandler(filename=log_path) file_handler.setLevel(60) # Custom log level > CRITICAL to prevent any unwanted standard DEBUG-CRITICAL logs @@ -52,11 +53,17 @@ class PacketCapture: log_format = "%(message)s" file_handler.setFormatter(logging.Formatter(log_format)) - self.logger = logging.getLogger(self._logger_name) - self.logger.setLevel(60) # Custom log level > CRITICAL to prevent any unwanted standard DEBUG-CRITICAL logs - self.logger.addHandler(file_handler) + if outbound: + self.outbound_logger = logging.getLogger(self._get_logger_name(outbound)) + logger = self.outbound_logger + else: + self.inbound_logger = logging.getLogger(self._get_logger_name(outbound)) + logger = self.inbound_logger - self.logger.addFilter(_JSONFilter()) + logger.setLevel(60) # Custom log level > CRITICAL to prevent any unwanted standard DEBUG-CRITICAL logs + logger.addHandler(file_handler) + + logger.addFilter(_JSONFilter()) def read(self) -> List[Dict[str, Any]]: """ @@ -70,27 +77,35 @@ class PacketCapture: frames.append(json.loads(line.rstrip())) return frames - @property - def _logger_name(self) -> str: + def _get_logger_name(self, outbound: bool = False) -> str: """Get PCAP the logger name.""" if self.ip_address: - return f"{self.hostname}_{self.ip_address}_pcap" + return f"{self.hostname}_{self.ip_address}_{'outbound' if outbound else 'inbound'}_pcap" if self.switch_port_number: - return f"{self.hostname}_port-{self.switch_port_number}_pcap" - return f"{self.hostname}_pcap" + return f"{self.hostname}_port-{self.switch_port_number}_{'outbound' if outbound else 'inbound'}_pcap" + return f"{self.hostname}_{'outbound' if outbound else 'inbound'}_pcap" - def _get_log_path(self) -> Path: + def _get_log_path(self, outbound: bool = False) -> Path: """Get the path for the log file.""" root = SIM_OUTPUT.path / f"episode_{self.current_episode}" / self.hostname root.mkdir(exist_ok=True, parents=True) - return root / f"{self._logger_name}.log" + return root / f"{self._get_logger_name(outbound)}.log" - def capture(self, frame): # noqa - I'll have a circular import and cant use if TYPE_CHECKING ;( + def capture_inbound(self, frame): # noqa - I'll have a circular import and cant use if TYPE_CHECKING ;( """ - Capture a Frame and log it. + Capture an inbound Frame and log it. :param frame: The PCAP frame to capture. """ - if SIM_OUTPUT.save_pcap_logs: - msg = frame.model_dump_json() - self.logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL + msg = frame.model_dump_json() + self.inbound_logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL + + def capture_outbound(self, frame): # noqa - I'll have a circular import and cant use if TYPE_CHECKING ;( + """ + Capture an inbound Frame and log it. + + :param frame: The PCAP frame to capture. + """ + msg = frame.model_dump_json() + self.outbound_logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL + diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index ce05193f..2120cde3 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/src/primaite/simulator/system/core/session_manager.py @@ -221,7 +221,7 @@ class SessionManager: if payload.request: dst_mac_address = "ff:ff:ff:ff:ff:ff" else: - dst_mac_address = payload.sender_mac_addr + dst_mac_address = payload.target_mac_addr outbound_nic = self.resolve_outbound_nic(payload.target_ip_address) is_broadcast = payload.request ip_protocol = IPProtocol.UDP diff --git a/src/primaite/simulator/system/services/arp/arp.py b/src/primaite/simulator/system/services/arp/arp.py index 28a2485c..c5b30d69 100644 --- a/src/primaite/simulator/system/services/arp/arp.py +++ b/src/primaite/simulator/system/services/arp/arp.py @@ -15,6 +15,12 @@ from primaite.simulator.system.services.service import Service class ARP(Service): + """ + The ARP (Address Resolution Protocol) Service. + + Manages ARP for resolving network layer addresses into link layer addresses. It maintains an ARP cache, + sends ARP requests and replies, and processes incoming ARP packets. + """ arp: Dict[IPv4Address, ARPEntry] = {} def __init__(self, **kwargs): @@ -27,7 +33,11 @@ class ARP(Service): pass def show(self, markdown: bool = False): - """Prints a table of ARC Cache.""" + """ + Prints the current state of the ARP cache in a table format. + + :param markdown: If True, format the output as Markdown. Otherwise, use plain text. + """ table = PrettyTable(["IP Address", "MAC Address", "Via"]) if markdown: table.set_style(MARKDOWN) @@ -71,37 +81,28 @@ class ARP(Service): @abstractmethod def get_arp_cache_mac_address(self, ip_address: IPv4Address) -> Optional[str]: """ - Get the MAC address associated with an IP address. + Retrieves the MAC address associated with a given IP address from the ARP cache. - :param ip_address: The IP address to look up in the cache. - :return: The MAC address associated with the IP address, or None if not found. + :param ip_address: The IP address to look up. + :return: The associated MAC address, if found. Otherwise, returns None. """ pass @abstractmethod def get_arp_cache_nic(self, ip_address: IPv4Address) -> Optional[NIC]: """ - Get the NIC associated with an IP address. + Retrieves the NIC associated with a given IP address from the ARP cache. - :param ip_address: The IP address to look up in the cache. - :return: The NIC associated with the IP address, or None if not found. + :param ip_address: The IP address to look up. + :return: The associated NIC, if found. Otherwise, returns None. """ pass def send_arp_request(self, target_ip_address: Union[IPv4Address, str]): """ - Perform a standard ARP request for a given target IP address. + Sends an ARP request to resolve the MAC address of a target IP address. - Broadcasts the request through all enabled NICs to determine the MAC address corresponding to the target IP - address. This method can be configured to ignore specific networks when sending out ARP requests, - which is useful in environments where certain addresses should not be queried. - - :param target_ip_address: The target IP address to send an ARP request for. - :param ignore_networks: An optional list of IPv4 addresses representing networks to be excluded from the ARP - request broadcast. Each address in this list indicates a network which will not be queried during the ARP - request process. This is particularly useful in complex network environments where traffic should be - minimized or controlled to specific subnets. It is mainly used by the router to prevent ARP requests being - sent back to their source. + :param target_ip_address: The target IP address for which the MAC address is being requested. """ outbound_nic = self.software_manager.session_manager.resolve_outbound_nic(target_ip_address) if outbound_nic: @@ -112,68 +113,59 @@ class ARP(Service): target_ip_address=target_ip_address, ) self.software_manager.session_manager.receive_payload_from_software_manager( - payload=arp_packet, dst_ip_address=target_ip_address, dst_port=Port.ARP, ip_protocol=self.protocol + payload=arp_packet, dst_ip_address=target_ip_address, dst_port=self.port, ip_protocol=self.protocol ) else: - print(f"failed for {target_ip_address}") - - 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_address} " - f"to {arp_reply.target_ip_address}/{arp_reply.target_mac_addr} " - ) - udp_header = UDPHeader(src_port=Port.ARP, dst_port=Port.ARP) - - ip_packet = IPPacket( - src_ip_address=arp_reply.sender_ip_address, - dst_ip_address=arp_reply.target_ip_address, - protocol=IPProtocol.UDP, - ) - - 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, udp=udp_header, payload=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. - - If an ARP request is received for the local IP, a response is sent back. - If an ARP response is received, the ARP cache is updated with the new entry. - - :param from_nic: The NIC that received the ARP packet. - :param arp_packet: The ARP packet to be processed. - """ - - # Unmatched ARP Request - if arp_packet.target_ip_address != from_nic.ip_address: - self.sys_log.info( - f"Ignoring ARP request for {arp_packet.target_ip_address}. Current IP address is {from_nic.ip_address}" + self.sys_log.error( + "Cannot send ARP request as there is no outbound NIC to use. Try configuring the default gateway." + ) + + def send_arp_reply(self, arp_reply: ARPPacket): + """ + Sends an ARP reply in response to an ARP request. + + :param arp_reply: The ARP packet containing the reply. + :param from_nic: The NIC from which the ARP reply is sent. + """ + + outbound_nic = self.software_manager.session_manager.resolve_outbound_nic(arp_reply.target_ip_address) + if outbound_nic: + self.sys_log.info( + 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} " + ) + self.software_manager.session_manager.receive_payload_from_software_manager( + payload=arp_reply, + dst_ip_address=arp_reply.target_ip_address, + dst_port=self.port, + ip_protocol=self.protocol + ) + else: + self.sys_log.error( + "Cannot send ARP reply as there is no outbound NIC to use. Try configuring the default gateway." ) - 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 - ) - arp_packet = arp_packet.generate_reply(from_nic.mac_address) - self.send_arp_reply(arp_packet, from_nic) @abstractmethod def _process_arp_request(self, arp_packet: ARPPacket, from_nic: NIC): + """ + Processes an incoming ARP request. + + :param arp_packet: The ARP packet containing the request. + :param from_nic: The NIC that received the ARP request. + """ self.sys_log.info( f"Received ARP request for {arp_packet.target_ip_address} from " f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip_address} " ) def _process_arp_reply(self, arp_packet: ARPPacket, from_nic: NIC): + """ + Processes an incoming ARP reply. + + :param arp_packet: The ARP packet containing the reply. + :param from_nic: The NIC that received the ARP reply. + """ self.sys_log.info( f"Received ARP response for {arp_packet.sender_ip_address} " f"from {arp_packet.sender_mac_addr} via NIC {from_nic}" @@ -183,6 +175,14 @@ class ARP(Service): ) def receive(self, payload: Any, session_id: str, **kwargs) -> bool: + """ + Processes received data, handling ARP packets. + + :param payload: The payload received. + :param session_id: The session ID associated with the received data. + :param kwargs: Additional keyword arguments. + :return: True if the payload was processed successfully, otherwise False. + """ if not isinstance(payload, ARPPacket): print("failied on payload check", type(payload)) return False @@ -194,4 +194,10 @@ class ARP(Service): self._process_arp_reply(arp_packet=payload, from_nic=from_nic) def __contains__(self, item: Any) -> bool: + """ + Checks if an item is in the ARP cache. + + :param item: The item to check. + :return: True if the item is in the cache, otherwise False. + """ return item in self.arp diff --git a/src/primaite/simulator/system/services/arp/host_arp.py b/src/primaite/simulator/system/services/arp/host_arp.py index f3e70838..4d6f7738 100644 --- a/src/primaite/simulator/system/services/arp/host_arp.py +++ b/src/primaite/simulator/system/services/arp/host_arp.py @@ -92,4 +92,4 @@ class HostARP(ARP): 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) + self.send_arp_reply(arp_packet)