from __future__ import annotations from enum import Enum from ipaddress import IPv4Address, IPv4Network from typing import Dict, List, Optional, Tuple, Union from prettytable import PrettyTable 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): 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 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: if nic.enabled: # reply to the request if not is_reattempt: 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: for nic in self.router.nics.values(): if nic.ip_address == frame.ip.dst_ip: if nic.enabled: 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 return # Route the frame self.router.route_frame(frame, from_nic) class Router(Node): num_ports: int ethernet_ports: Dict[int, NIC] = {} 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 route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None: # 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 re_attempt and not nic: self.sys_log.info(f"Destination {frame.ip.dst_ip} is unreachable") return 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) if not nic.enabled: # TODO: Add sys_log here return if 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: pass # TODO: Deal with routing from route tables 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={ip_address}/{nic.ip_network.prefixlen}") def enable_port(self, port: int): nic = self.ethernet_ports.get(port) if nic: nic.enable() def disable_port(self, port: int): nic = self.ethernet_ports.get(port) if nic: nic.disable() 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( [ port, nic.mac_address, f"{nic.ip_address}/{nic.ip_network.prefixlen}", nic.speed, "Enabled" if nic.enabled else "Disabled", ] ) print(table)