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

Merge remote-tracking branch 'origin/dev' into dev-v3.0.0b6

Browse Source
This commit is contained in:
Marek Wolan
2024-01-29 10:26:28 +00:00
parent 2ba05e7348 62197fe72b
commit 555802baaa
57 changed files with 2329 additions and 461 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4,7 +4,7 @@ import re
import secrets
from ipaddress import IPv4Address, IPv4Network
from pathlib import Path
from typing import Any, Dict, Literal, Optional, Tuple, Union
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
from prettytable import MARKDOWN, PrettyTable
@@ -91,6 +91,8 @@ class NIC(SimComponent):
"Indicates if the NIC supports Wake-on-LAN functionality."
_connected_node: Optional[Node] = None
"The Node to which the NIC is connected."
_port_num_on_node: Optional[int] = None
"Which port number is assigned on this NIC"
_connected_link: Optional[Link] = None
"The Link to which the NIC is connected."
enabled: bool = False
@@ -148,7 +150,7 @@ class NIC(SimComponent):
state = super().describe_state()
state.update(
{
"ip_adress": str(self.ip_address),
"ip_address": str(self.ip_address),
"subnet_mask": str(self.subnet_mask),
"mac_address": self.mac_address,
"speed": self.speed,
@@ -272,18 +274,40 @@ class NIC(SimComponent):
def receive_frame(self, frame: Frame) -> bool:
"""
Receive a network frame from the connected link if the NIC is enabled.
Receive a network frame from the connected link, processing it if the NIC is enabled.
The Frame is passed to the Node.
This method decrements the Time To Live (TTL) of the frame, captures it using PCAP (Packet Capture), and checks
if the frame is either a broadcast or destined for this NIC. If the frame is acceptable, it is passed to the
connected node. The method also handles the discarding of frames with TTL expired and logs this event.
:param frame: The network frame being received.
The frame's reception is based on various conditions:
- If the NIC is disabled, the frame is not processed.
- If the TTL of the frame reaches zero after decrement, it is discarded and logged.
- If the frame is a broadcast or its destination MAC/IP address matches this NIC's, it is accepted.
- All other frames are dropped and logged or printed to the console.
:param frame: The network frame being received. This should be an instance of the Frame class.
:return: Returns True if the frame is processed and passed to the node, False otherwise.
"""
if self.enabled:
frame.decrement_ttl()
if frame.ip and frame.ip.ttl < 1:
self._connected_node.sys_log.info("Frame discarded as TTL limit reached")
return False
frame.set_received_timestamp()
self.pcap.capture(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":
accept_frame = False
# Check if it's a broadcast:
if frame.ethernet.dst_mac_addr == "ff:ff:ff:ff:ff:ff":
if frame.ip.dst_ip_address in {self.ip_address, self.ip_network.broadcast_address}:
accept_frame = True
else:
if frame.ethernet.dst_mac_addr == self.mac_address:
accept_frame = True
if accept_frame:
self._connected_node.receive_frame(frame=frame, from_nic=self)
return True
return False
@@ -311,6 +335,8 @@ class SwitchPort(SimComponent):
"The Maximum Transmission Unit (MTU) of the SwitchPort in Bytes. Default is 1500 B"
_connected_node: Optional[Node] = None
"The Node to which the SwitchPort is connected."
_port_num_on_node: Optional[int] = None
"The port num on the connected node."
_connected_link: Optional[Link] = None
"The Link to which the SwitchPort is connected."
enabled: bool = False
@@ -432,6 +458,9 @@ class SwitchPort(SimComponent):
"""
if self.enabled:
frame.decrement_ttl()
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)
connected_node: Node = self._connected_node
connected_node.forward_frame(frame=frame, incoming_port=self)
@@ -497,8 +526,8 @@ class Link(SimComponent):
state = super().describe_state()
state.update(
{
"endpoint_a": self.endpoint_a.uuid,
"endpoint_b": self.endpoint_b.uuid,
"endpoint_a": self.endpoint_a.uuid, # TODO: consider if using UUID is the best way to do this
"endpoint_b": self.endpoint_b.uuid, # TODO: consider if using UUID is the best way to do this
"bandwidth": self.bandwidth,
"current_load": self.current_load,
}
@@ -667,17 +696,30 @@ class ARPCache:
"""Clear the entire ARP cache, removing all stored entries."""
self.arp.clear()
def send_arp_request(self, target_ip_address: Union[IPv4Address, str]):
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.
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.
"""
for nic in self.nics.values():
if nic.enabled:
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}")
tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP)
@@ -802,7 +844,6 @@ class ICMP:
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(
@@ -817,9 +858,7 @@ class ICMP:
sequence=frame.icmp.sequence + 1,
)
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
)
frame = Frame(ethernet=ethernet_header, ip=ip_packet, icmp=icmp_reply_packet, payload=payload)
self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip_address}")
src_nic.send_frame(frame)
@@ -1094,12 +1133,12 @@ class Node(SimComponent):
{
"hostname": self.hostname,
"operating_state": self.operating_state.value,
"NICs": {uuid: nic.describe_state() for uuid, nic in self.nics.items()},
"NICs": {eth_num: nic.describe_state() for eth_num, nic in self.ethernet_port.items()},
# "switch_ports": {uuid, sp for uuid, sp in self.switch_ports.items()},
"file_system": self.file_system.describe_state(),
"applications": {uuid: app.describe_state() for uuid, app in self.applications.items()},
"services": {uuid: svc.describe_state() for uuid, svc in self.services.items()},
"process": {uuid: proc.describe_state() for uuid, proc in self.processes.items()},
"applications": {app.name: app.describe_state() for app in self.applications.values()},
"services": {svc.name: svc.describe_state() for svc in self.services.values()},
"process": {proc.name: proc.describe_state() for proc in self.processes.values()},
"revealed_to_red": self.revealed_to_red,
}
)
@@ -1316,6 +1355,7 @@ class Node(SimComponent):
self.nics[nic.uuid] = nic
self.ethernet_port[len(self.nics)] = nic
nic._connected_node = self
nic._port_num_on_node = len(self.nics)
nic.parent = self
self.sys_log.info(f"Connected NIC {nic}")
if self.operating_state == NodeOperatingState.ON:
@@ -1442,7 +1482,6 @@ class Node(SimComponent):
service.parent = self
service.install() # Perform any additional setup, such as creating files for this service on the node.
self.sys_log.info(f"Installed service {service.name}")
_LOGGER.info(f"Added service {service.uuid} to node {self.uuid}")
self._service_request_manager.add_request(service.uuid, RequestType(func=service._request_manager))
def uninstall_service(self, service: Service) -> None:
@@ -1475,7 +1514,6 @@ class Node(SimComponent):
self.applications[application.uuid] = application
application.parent = self
self.sys_log.info(f"Installed application {application.name}")
_LOGGER.info(f"Added application {application.uuid} to node {self.uuid}")
self._application_request_manager.add_request(application.uuid, RequestType(func=application._request_manager))
def uninstall_application(self, application: Application) -> None:

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

@@ -357,11 +357,10 @@ 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_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.
:ivar address: The destination IP address or network address.
:ivar subnet_mask: The subnet mask for the network.
:ivar next_hop_ip_address: The next hop IP address to which packets should be forwarded.
:ivar metric: The cost metric for this route. Default is 0.0.
Example:
>>> entry = RouteEntry(
@@ -381,12 +380,6 @@ class RouteEntry(SimComponent):
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_ip_address"}:
if not isinstance(kwargs[key], IPv4Address):
kwargs[key] = IPv4Address(kwargs[key])
super().__init__(**kwargs)
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"address", "subnet_mask", "next_hop_ip_address", "metric"}
@@ -421,6 +414,7 @@ class RouteTable(SimComponent):
"""
routes: List[RouteEntry] = []
default_route: Optional[RouteEntry] = None
sys_log: SysLog
def set_original_state(self):
@@ -465,12 +459,35 @@ class RouteTable(SimComponent):
)
self.routes.append(route)
def set_default_route_next_hop_ip_address(self, ip_address: IPv4Address):
"""
Sets the next-hop IP address for the default route in a routing table.
This method checks if a default route (0.0.0.0/0) exists in the routing table. If it does not exist,
the method creates a new default route with the specified next-hop IP address. If a default route already
exists, it updates the next-hop IP address of the existing default route. After setting the next-hop
IP address, the method logs this action.
:param ip_address: The next-hop IP address to be set for the default route.
"""
if not self.default_route:
self.default_route = RouteEntry(
ip_address=IPv4Address("0.0.0.0"),
subnet_mask=IPv4Address("0.0.0.0"),
next_hop_ip_address=ip_address,
)
else:
self.default_route.next_hop_ip_address = ip_address
self.sys_log.info(f"Default configured to use {ip_address} as the next-hop")
def find_best_route(self, destination_ip: Union[str, IPv4Address]) -> Optional[RouteEntry]:
"""
Find the best route for a given destination IP.
This method uses the Longest Prefix Match algorithm and considers metrics to find the best route.
If no dedicated route exists but a default route does, then the default route is returned as a last resort.
:param destination_ip: The destination IP to find the route for.
:return: The best matching RouteEntry, or None if no route matches.
"""
@@ -490,6 +507,9 @@ class RouteTable(SimComponent):
longest_prefix = prefix_len
lowest_metric = route.metric
if not best_route and self.default_route:
best_route = self.default_route
return best_route
def show(self, markdown: bool = False):
@@ -521,12 +541,26 @@ class RouterARPCache(ARPCache):
super().__init__(sys_log)
self.router: Router = router
def process_arp_packet(self, from_nic: NIC, frame: Frame):
def process_arp_packet(
self, from_nic: NIC, frame: Frame, route_table: RouteTable, is_reattempt: bool = False
) -> None:
"""
Overridden method to process a received ARP packet in a router-specific way.
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 original ARP frame.
: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
@@ -554,7 +588,11 @@ class RouterARPCache(ARPCache):
)
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(
@@ -565,16 +603,32 @@ class RouterARPCache(ARPCache):
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_address:
if arp_packet.target_ip_address in nic.ip_network:
# 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:
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 RouterICMP(ICMP):
"""
@@ -645,7 +699,7 @@ class RouterICMP(ICMP):
return
# Route the frame
self.router.route_frame(frame, from_nic)
self.router.process_frame(frame, from_nic)
elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY:
for nic in self.router.nics.values():
@@ -665,7 +719,48 @@ class RouterICMP(ICMP):
return
# Route the frame
self.router.route_frame(frame, from_nic)
self.router.process_frame(frame, from_nic)
class RouterNIC(NIC):
"""
A Router-specific Network Interface Card (NIC) that extends the standard NIC functionality.
This class overrides the standard Node NIC's Layer 3 (L3) broadcast/unicast checks. It is designed
to handle network frames in a manner specific to routers, allowing them to efficiently process
and route network traffic.
"""
def receive_frame(self, frame: Frame) -> bool:
"""
Receive and process a network frame from the connected link, provided the NIC is enabled.
This method is tailored for router behavior. It decrements the frame's Time To Live (TTL), checks for TTL
expiration, and captures the frame using PCAP (Packet Capture). The frame is accepted if it is destined for
this NIC's MAC address or is a broadcast frame.
Key Differences from Standard NIC:
- Does not perform Layer 3 (IP-based) broadcast checks.
- Only checks for Layer 2 (Ethernet) destination MAC address and broadcast frames.
:param frame: The network frame being received. This should be an instance of the Frame class.
:return: Returns True if the frame is processed and passed to the connected node, False otherwise.
"""
if self.enabled:
frame.decrement_ttl()
if frame.ip and frame.ip.ttl < 1:
self._connected_node.sys_log.info("Frame discarded as TTL limit reached")
return False
frame.set_received_timestamp()
self.pcap.capture(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)
return True
return False
def __str__(self) -> str:
return f"{self.mac_address}/{self.ip_address}"
class Router(Node):
@@ -678,7 +773,7 @@ class Router(Node):
"""
num_ports: int
ethernet_ports: Dict[int, NIC] = {}
ethernet_ports: Dict[int, RouterNIC] = {}
acl: AccessControlList
route_table: RouteTable
arp: RouterARPCache
@@ -697,7 +792,7 @@ class Router(Node):
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")
nic = RouterNIC(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
@@ -752,9 +847,9 @@ class Router(Node):
state["acl"] = self.acl.describe_state()
return state
def route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None:
def process_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None:
"""
Route a given frame from a source NIC to its destination.
Process a Frame.
:param frame: The frame to be routed.
:param from_nic: The source network interface.
@@ -769,25 +864,57 @@ class Router(Node):
return
if not nic:
self.arp.send_arp_request(frame.ip.dst_ip_address)
return self.route_frame(frame=frame, from_nic=from_nic, re_attempt=True)
self.arp.send_arp_request(
frame.ip.dst_ip_address, ignore_networks=[frame.ip.src_ip_address, from_nic.ip_address]
)
return self.process_frame(frame=frame, from_nic=from_nic, re_attempt=True)
if not nic.enabled:
# TODO: Add sys_log here
self.sys_log.info(f"Frame dropped as NIC {nic} is not enabled")
return
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}")
self.sys_log.info(f"Forwarding frame to internally from port {from_port} to port {to_port}")
frame.decrement_ttl()
if frame.ip and frame.ip.ttl < 1:
self.sys_log.info("Frame discarded as TTL limit reached")
return
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
self._route_frame(frame, from_nic)
def _route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None:
route = self.route_table.find_best_route(frame.ip.dst_ip_address)
if route:
nic = self.arp.get_arp_cache_nic(route.next_hop_ip_address)
target_mac = self.arp.get_arp_cache_mac_address(route.next_hop_ip_address)
if re_attempt and not nic:
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_address, ignore_networks=[frame.ip.src_ip_address])
return self.process_frame(frame=frame, from_nic=from_nic, re_attempt=True)
if not nic.enabled:
self.sys_log.info(f"Frame dropped as NIC {nic} is not enabled")
return
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()
if frame.ip and frame.ip.ttl < 1:
self.sys_log.info("Frame discarded as TTL limit reached")
return
frame.ethernet.src_mac_addr = nic.mac_address
frame.ethernet.dst_mac_addr = target_mac
nic.send_frame(frame)
def receive_frame(self, frame: Frame, from_nic: NIC):
"""
@@ -796,7 +923,7 @@ class Router(Node):
:param frame: The incoming frame.
:param from_nic: The network interface where the frame is coming from.
"""
route_frame = False
process_frame = False
protocol = frame.ip.protocol
src_ip_address = frame.ip.src_ip_address
dst_ip_address = frame.ip.dst_ip_address
@@ -828,12 +955,12 @@ class Router(Node):
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)
self.arp.process_arp_packet(from_nic=from_nic, frame=frame, route_table=self.route_table)
else:
# All other traffic
route_frame = True
if route_frame:
self.route_frame(frame, from_nic)
process_frame = True
if process_frame:
self.process_frame(frame, from_nic)
def configure_port(self, port: int, ip_address: Union[IPv4Address, str], subnet_mask: Union[IPv4Address, str]):
"""

5
src/primaite/simulator/network/hardware/nodes/switch.py
View File

@@ -30,6 +30,7 @@ class Switch(Node):
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._port_num_on_node = port_num
port.parent = self
port.port_num = port_num
@@ -89,12 +90,12 @@ class Switch(Node):
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":
if outgoing_port and dst_mac.lower() != "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:
if port.enabled and port != incoming_port:
port.send_frame(frame)
def disconnect_link_from_port(self, link: Link, port_number: int):