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

#1800 - Added ACL and routing classes.

Browse Source 
- 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.
This commit is contained in:
Chris McCarthy
2023-08-30 21:38:55 +01:00
parent c6f71600fc
commit 1bf51c7741
9 changed files with 739 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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