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

#2248 - TSome further fixess to ARP. Also refactored PCAP to log inbound and outbound frames separately

Browse Source
This commit is contained in:
Chris McCarthy
2024-02-02 17:14:34 +00:00
parent cb002d644f
commit a0253ce6c4
6 changed files with 117 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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

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

@@ -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)

59
src/primaite/simulator/system/core/packet_capture.py
View File

@@ -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

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

@@ -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

140
src/primaite/simulator/system/services/arp/arp.py
View File

@@ -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

2
src/primaite/simulator/system/services/arp/host_arp.py
View File

@@ -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)