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

#2248 - ICMP now working as a service using the session manager for transmission. Now started to comb through the tests to fix anything up.

Browse Source
This commit is contained in:
Chris McCarthy
2024-02-02 16:20:15 +00:00
parent 87d9d6da04
commit dc5aeede33
12 changed files with 241 additions and 491 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -599,262 +599,6 @@ class Link(SimComponent):
def __str__(self) -> str:
return f"{self.endpoint_a}<-->{self.endpoint_b}"
class ARPCache:
"""
The ARPCache (Address Resolution Protocol) class.
Responsible for maintaining a mapping between IP addresses and MAC addresses (ARP cache) for the network. It
provides methods for looking up, adding, and removing entries, and for processing ARPPackets.
"""
def __init__(self, sys_log: "SysLog"):
"""
Initialize an ARP (Address Resolution Protocol) cache.
:param sys_log: The nodes sys log.
"""
self.sys_log: "SysLog" = sys_log
self.arp: Dict[IPv4Address, ARPEntry] = {}
self.nics: Dict[str, "NIC"] = {}
self.node = None
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(
[
str(ip),
arp.mac_address,
self.nics[arp.nic_uuid].mac_address,
]
)
print(table)
def clear(self):
"""Clears the arp cache."""
self.arp.clear()
def add_arp_cache_entry(self, ip_address: IPv4Address, mac_address: str, nic: NIC, override: bool = False):
"""
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:
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):
"""
Remove an ARP entry from the cache.
:param ip_address: The IP address to be removed from the cache.
"""
if ip_address in self.arp:
del self.arp[ip_address]
def get_default_gateway_mac_address(self) -> Optional[str]:
if self.arp.node.default_gateway:
return self.get_arp_cache_mac_address(self.arp.node.default_gateway)
def get_default_gateway_nic(self) -> Optional[NIC]:
if self.arp.node.default_gateway:
return self.get_arp_cache_nic(self.arp.node.default_gateway)
def _get_arp_cache_mac_address(
self, ip_address: IPv4Address, is_reattempt: bool = False, is_default_gateway_attempt: bool = False
) -> Optional[str]:
arp_entry = self.arp.get(ip_address)
if arp_entry:
return arp_entry.mac_address
else:
if not is_reattempt:
self.send_arp_request(ip_address)
return self._get_arp_cache_mac_address(
ip_address=ip_address, is_reattempt=True, is_default_gateway_attempt=is_default_gateway_attempt
)
else:
if self.node.default_gateway:
if not is_default_gateway_attempt:
self.send_arp_request(self.node.default_gateway)
return self._get_arp_cache_mac_address(
ip_address=self.node.default_gateway, is_reattempt=True, is_default_gateway_attempt=True
)
return None
def get_arp_cache_mac_address(self, ip_address: IPv4Address) -> Optional[str]:
"""
Get the MAC address associated with an IP address.
: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.
"""
return self._get_arp_cache_mac_address(ip_address)
def _get_arp_cache_nic(
self, ip_address: IPv4Address, is_reattempt: bool = False, is_default_gateway_attempt: bool = False
) -> Optional[NIC]:
arp_entry = self.arp.get(ip_address)
if arp_entry:
return self.nics[arp_entry.nic_uuid]
else:
if not is_reattempt:
self.send_arp_request(ip_address)
return self._get_arp_cache_nic(
ip_address=ip_address, is_reattempt=True, is_default_gateway_attempt=is_default_gateway_attempt
)
else:
if self.node.default_gateway:
if not is_default_gateway_attempt:
self.send_arp_request(self.node.default_gateway)
return self._get_arp_cache_nic(
ip_address=self.node.default_gateway, is_reattempt=True, is_default_gateway_attempt=True
)
return None
def get_arp_cache_nic(self, ip_address: IPv4Address) -> Optional[NIC]:
"""
Get the NIC associated with an IP address.
: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.
"""
return self._get_arp_cache_nic(ip_address)
def clear_arp_cache(self):
"""Clear the entire ARP cache, removing all stored entries."""
self.arp.clear()
def send_arp_request(
self, target_ip_address: Union[IPv4Address, str], ignore_networks: Optional[List[IPv4Address]] = None
):
"""
Perform a standard ARP request for a given 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.
"""
pass
# for nic in self.nics.values():
# use_nic = True
# if ignore_networks:
# for ipv4 in ignore_networks:
# if ipv4 in nic.ip_network:
# use_nic = False
# if nic.enabled and use_nic:
# self.sys_log.info(f"Sending ARP request from NIC {nic} for ip {target_ip_address}")
# udp_header = UDPHeader(src_port=Port.ARP, dst_port=Port.ARP)
#
# # Network Layer
# ip_packet = IPPacket(
# src_ip_address=nic.ip_address, dst_ip_address=target_ip_address, protocol=IPProtocol.UDP
# )
# # 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,
# )
# frame = Frame(ethernet=ethernet_header, ip=ip_packet, udp=udp_header, payload=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_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_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)
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.
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.
"""
# ARP Reply
if not arp_packet.request:
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}"
)
self.add_arp_cache_entry(
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_address} from "
f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip_address} "
)
# 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}"
)
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)
def __contains__(self, item: Any) -> bool:
return item in self.arp
class Node(SimComponent):
"""
A basic Node class that represents a node on the network.

29
src/primaite/simulator/network/hardware/nodes/computer.py
View File

@@ -1,10 +1,11 @@
from primaite.simulator.network.hardware.base import NIC, Node
from primaite.simulator.network.hardware.nodes.host import Host
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
class Computer(Node):
class Computer(Host):
"""
A basic Computer class.
@@ -20,36 +21,16 @@ class Computer(Node):
Instances of computer come 'pre-packaged' with the following:
* Core Functionality:
* ARP
* ICMP
* Packet Capture
* Sys Log
* Services:
* ARP Service
* ICMP Service
* DNS Client
* FTP Client
* LDAP Client
* NTP Client
* Applications:
* Email Client
* Web Browser
* Processes:
* Placeholder
"""
pass
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.connect_nic(NIC(ip_address=kwargs["ip_address"], subnet_mask=kwargs["subnet_mask"]))
self._install_system_software()
def _install_system_software(self):
"""Install System Software - software that is usually provided with the OS."""
# DNS Client
self.software_manager.install(DNSClient)
# FTP
self.software_manager.install(FTPClient)
# Web Browser
self.software_manager.install(WebBrowser)
super()._install_system_software()

8
src/primaite/simulator/network/hardware/nodes/host.py
View File

@@ -23,20 +23,16 @@ class Host(Node):
Instances of computer come 'pre-packaged' with the following:
* Core Functionality:
* ARP
* ICMP
* Packet Capture
* Sys Log
* Services:
* ARP Service
* ICMP Service
* DNS Client
* FTP Client
* LDAP Client
* NTP Client
* Applications:
* Email Client
* Web Browser
* Processes:
* Placeholder
"""
def __init__(self, **kwargs):

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

@@ -8,7 +8,7 @@ from typing import Dict, List, Optional, Tuple, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.network.hardware.base import ARPCache, NIC, Node
from primaite.simulator.network.hardware.base import NIC, Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame
from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
@@ -528,108 +528,6 @@ class RouteTable(SimComponent):
table.add_row([index, f"{route.address}/{network.prefixlen}", route.next_hop_ip_address, route.metric])
print(table)
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
def process_arp_packet(
self, from_nic: NIC, frame: Frame, route_table: RouteTable, is_reattempt: bool = False
) -> None:
"""
Processes a received ARP (Address Resolution Protocol) packet in a router-specific way.
This method is responsible for handling both ARP requests and responses. It processes ARP packets received on a
Network Interface Card (NIC) and performs actions based on whether the packet is a request or a reply. This
includes updating the ARP cache, forwarding ARP replies, sending ARP requests for unknown destinations, and
handling packet TTL (Time To Live).
The method first checks if the ARP packet is a request or a reply. For ARP replies, it updates the ARP cache
and forwards the reply if necessary. For ARP requests, it checks if the target IP matches one of the router's
NICs and sends an ARP reply if so. If the destination is not directly connected, it consults the routing table
to find the best route and reattempts ARP request processing if needed.
:param from_nic: The NIC that received the ARP packet.
:param frame: The frame containing the ARP packet.
:param route_table: The routing table of the router.
:param is_reattempt: Flag to indicate if this is a reattempt of processing the ARP packet, defaults to False.
"""
arp_packet = frame.arp
# ARP Reply
if not arp_packet.request:
if arp_packet.target_ip_address == from_nic.ip_address:
# reply to the Router specifically
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}"
)
self.add_arp_cache_entry(
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_address)
# if nic:
# 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()
# if frame.ip and frame.ip.ttl < 1:
# self.sys_log.info("Frame discarded as TTL limit reached")
# return
# nic.send_frame(frame)
# return
# 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} "
)
# 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
)
# 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_address:
arp_reply = arp_packet.generate_reply(from_nic.mac_address)
self.send_arp_reply(arp_reply, from_nic)
return
# # Check Route Table
# route = route_table.find_best_route(arp_packet.target_ip_address)
# if route and route != self.router.route_table.default_route:
# nic = self.get_arp_cache_nic(route.next_hop_ip_address)
#
# if not nic:
# if not is_reattempt:
# self.send_arp_request(route.next_hop_ip_address, ignore_networks=[frame.ip.src_ip_address])
# return self.process_arp_packet(from_nic, frame, route_table, is_reattempt=True)
# else:
# self.sys_log.info("Ignoring ARP request as destination unavailable/No ARP entry found")
# return
# else:
# arp_reply = arp_packet.generate_reply(from_nic.mac_address)
# self.send_arp_reply(arp_reply, from_nic)
# return
class RouterNIC(NIC):
"""
A Router-specific Network Interface Card (NIC) that extends the standard NIC functionality.
@@ -684,8 +582,8 @@ class Router(Node):
ethernet_ports: Dict[int, RouterNIC] = {}
acl: AccessControlList
route_table: RouteTable
arp: RouterARPCache
icmp: RouterICMP
# arp: RouterARPCache
# icmp: RouterICMP
def __init__(self, hostname: str, num_ports: int = 5, **kwargs):
if not kwargs.get("sys_log"):
@@ -694,12 +592,13 @@ class Router(Node):
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("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)
# TODO: Install RoputerICMP
# TODO: Install RouterICMP
# TODO: Install RouterARP
for i in range(1, self.num_ports + 1):
nic = RouterNIC(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0")
self.connect_nic(nic)

13
src/primaite/simulator/network/hardware/nodes/server.py
View File

@@ -1,7 +1,7 @@
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.host import Host
class Server(Computer):
class Server(Host):
"""
A basic Server class.
@@ -17,18 +17,15 @@ class Server(Computer):
Instances of Server come 'pre-packaged' with the following:
* Core Functionality:
* ARP
* ICMP
* Packet Capture
* Sys Log
* Services:
* ARP Service
* ICMP Service
* DNS Client
* FTP Client
* LDAP Client
* NTP Client
* Applications:
* Email Client
* Web Browser
* Processes:
* Placeholder
"""
pass