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

Merge remote-tracking branch 'devops/dev' into feature/1801-Database

Browse Source 
# Conflicts:
#	src/primaite/simulator/network/container.py
#	src/primaite/simulator/network/hardware/base.py
This commit is contained in:
Chris McCarthy
2023-09-04 19:45:29 +01:00
parent 0892a976fd c547f12f4e
commit 1385adf3fe
42 changed files with 2976 additions and 394 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -6,7 +6,7 @@ from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, List, Optional, Tuple, Union
from prettytable import PrettyTable
from prettytable import MARKDOWN, PrettyTable
from primaite import getLogger
from primaite.exceptions import NetworkError
@@ -77,12 +77,10 @@ class NIC(SimComponent):
ip_address: IPv4Address
"The IP address assigned to the NIC for communication on an IP-based network."
subnet_mask: str
subnet_mask: IPv4Address
"The subnet mask assigned to the NIC."
gateway: IPv4Address
"The default gateway IP address for forwarding network traffic to other networks. Randomly generated upon creation."
mac_address: str
"The MAC address of the NIC. Defaults to a randomly set MAC address."
"The MAC address of the NIC. Defaults to a randomly set MAC address. Randomly generated upon creation."
speed: int = 100
"The speed of the NIC in Mbps. Default is 100 Mbps."
mtu: int = 1500
@@ -111,16 +109,10 @@ class NIC(SimComponent):
"""
if not isinstance(kwargs["ip_address"], IPv4Address):
kwargs["ip_address"] = IPv4Address(kwargs["ip_address"])
if not isinstance(kwargs["gateway"], IPv4Address):
kwargs["gateway"] = IPv4Address(kwargs["gateway"])
if "mac_address" not in kwargs:
kwargs["mac_address"] = generate_mac_address()
super().__init__(**kwargs)
if self.ip_address == self.gateway:
msg = f"NIC ip address {self.ip_address} cannot be the same as the gateway {self.gateway}"
_LOGGER.error(msg)
raise ValueError(msg)
if self.ip_network.network_address == self.ip_address:
msg = (
f"Failed to set IP address {self.ip_address} and subnet mask {self.subnet_mask} as it is a "
@@ -173,6 +165,9 @@ class NIC(SimComponent):
if self.connected_node.operating_state != NodeOperatingState.ON:
self.connected_node.sys_log.error(f"NIC {self} cannot be enabled as the endpoint is not turned on")
return
if not self.connected_link:
_LOGGER.error(f"NIC {self} cannot be enabled as it is not connected to a Link")
return
self.enabled = True
self.connected_node.sys_log.info(f"NIC {self} enabled")
@@ -210,6 +205,7 @@ class NIC(SimComponent):
# TODO: Inform the Node that a link has been connected
self.connected_link = link
self.enable()
_LOGGER.info(f"NIC {self} connected to Link {link}")
def disconnect_link(self):
@@ -260,14 +256,15 @@ class NIC(SimComponent):
The Frame is passed to the Node.
:param frame: The network frame being received.
:type frame: :class:`~primaite.simulator.network.osi_layers.Frame`
"""
if self.enabled:
frame.decrement_ttl()
frame.set_received_timestamp()
self.pcap.capture(frame)
self.connected_node.receive_frame(frame=frame, from_nic=self)
return True
# 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:
@@ -291,7 +288,7 @@ class SwitchPort(SimComponent):
"The speed of the SwitchPort in Mbps. Default is 100 Mbps."
mtu: int = 1500
"The Maximum Transmission Unit (MTU) of the SwitchPort in Bytes. Default is 1500 B"
connected_node: Optional[Switch] = None
connected_node: Optional[Node] = None
"The Node to which the SwitchPort is connected."
connected_link: Optional[Link] = None
"The Link to which the SwitchPort is connected."
@@ -561,17 +558,43 @@ class ARPCache:
self.arp: Dict[IPv4Address, ARPEntry] = {}
self.nics: Dict[str, "NIC"] = {}
def _add_arp_cache_entry(self, ip_address: IPv4Address, mac_address: str, nic: NIC):
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 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.
"""
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
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):
"""
@@ -601,6 +624,7 @@ class ARPCache:
:return: The NIC associated with the IP address, or None if not found.
"""
arp_entry = self.arp.get(ip_address)
if arp_entry:
return self.nics[arp_entry.nic_uuid]
@@ -624,17 +648,42 @@ class ARPCache:
# Network Layer
ip_packet = IPPacket(
src_ip=nic.ip_address,
dst_ip=target_ip_address,
src_ip_address=nic.ip_address,
dst_ip_address=target_ip_address,
)
# 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=nic.ip_address, sender_mac_addr=nic.mac_address, target_ip=target_ip_address
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, tcp=tcp_header, arp=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.
@@ -648,45 +697,34 @@ class ARPCache:
# ARP Reply
if not arp_packet.request:
self.sys_log.info(
f"Received ARP response for {arp_packet.sender_ip} from {arp_packet.sender_mac_addr} via NIC {from_nic}"
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, mac_address=arp_packet.sender_mac_addr, 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} from "
f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip} "
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 != from_nic.ip_address:
self.sys_log.info(f"Ignoring ARP request for {arp_packet.target_ip}")
if arp_packet.target_ip_address != from_nic.ip_address:
self.sys_log.info(f"Ignoring ARP request for {arp_packet.target_ip_address}")
return
# Matched ARP request
self._add_arp_cache_entry(ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, 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
)
arp_packet = arp_packet.generate_reply(from_nic.mac_address)
self.sys_log.info(
f"Sending ARP reply from {arp_packet.sender_mac_addr}/{arp_packet.sender_ip} "
f"to {arp_packet.target_ip}/{arp_packet.target_mac_addr} "
)
self.send_arp_reply(arp_packet, from_nic)
tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP)
# Network Layer
ip_packet = IPPacket(
src_ip=arp_packet.sender_ip,
dst_ip=arp_packet.target_ip,
)
# Data Link Layer
ethernet_header = EthernetHeader(
src_mac_addr=arp_packet.sender_mac_addr, dst_mac_addr=arp_packet.target_mac_addr
)
frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_packet)
from_nic.send_frame(frame)
def __contains__(self, item: Any) -> bool:
return item in self.arp
class ICMP:
@@ -705,21 +743,30 @@ class ICMP:
"""
self.sys_log: SysLog = sys_log
self.arp: ARPCache = arp_cache
self.request_replies = {}
def process_icmp(self, frame: Frame):
def process_icmp(self, frame: Frame, from_nic: NIC, is_reattempt: bool = False):
"""
Process an ICMP packet, including handling echo requests and replies.
:param frame: The Frame containing the ICMP packet to process.
"""
if frame.icmp.icmp_type == ICMPType.ECHO_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)
if not is_reattempt:
self.sys_log.info(f"Received echo request from {frame.ip.src_ip_address}")
target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip_address)
src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip_address)
if not src_nic:
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(src_ip=src_nic.ip_address, dst_ip=frame.ip.src_ip, protocol=IPProtocol.ICMP)
ip_packet = IPPacket(
src_ip_address=src_nic.ip_address, dst_ip_address=frame.ip.src_ip_address, 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(
@@ -728,14 +775,28 @@ class ICMP:
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}")
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
)
self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip_address}")
src_nic.send_frame(frame)
elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY:
self.sys_log.info(f"Received echo reply from {frame.ip.src_ip}")
time = frame.transmission_duration()
time_str = f"{time}ms" if time > 0 else "<1ms"
self.sys_log.info(
f"Reply from {frame.ip.src_ip_address}: "
f"bytes={len(frame.payload)}, "
f"time={time_str}, "
f"TTL={frame.ip.ttl}"
)
if not self.request_replies.get(frame.icmp.identifier):
self.request_replies[frame.icmp.identifier] = 0
self.request_replies[frame.icmp.identifier] += 1
def ping(
self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None
self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None, pings: int = 4
) -> Tuple[int, Union[int, None]]:
"""
Send an ICMP echo request (ping) to a target IP address and manage the sequence and identifier.
@@ -747,13 +808,21 @@ class ICMP:
was not found in the ARP cache.
"""
nic = self.arp.get_arp_cache_nic(target_ip_address)
# TODO: Eventually this ARP request needs to be done elsewhere. It's not the resonsibility of the
# TODO: Eventually this ARP request needs to be done elsewhere. It's not the responsibility of the
# ping function to handle ARP lookups
# Already tried once and cannot get ARP entry, stop trying
if sequence == -1:
if not nic:
return 4, None
else:
sequence = 0
# No existing ARP entry
if not nic:
self.sys_log.info(f"No entry in ARP cache for {target_ip_address}")
self.arp.send_arp_request(target_ip_address)
return 0, None
return -1, None
# ARP entry exists
sequence += 1
@@ -763,15 +832,15 @@ class ICMP:
# Network Layer
ip_packet = IPPacket(
src_ip=nic.ip_address,
dst_ip=target_ip_address,
src_ip_address=nic.ip_address,
dst_ip_address=target_ip_address,
protocol=IPProtocol.ICMP,
)
# Data Link Layer
ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address)
icmp_packet = ICMPPacket(identifier=identifier, sequence=sequence)
frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_packet)
self.sys_log.info(f"Sending echo request to {target_ip_address}")
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_packet, payload=payload)
nic.send_frame(frame)
return sequence, icmp_packet.identifier
@@ -802,10 +871,14 @@ class Node(SimComponent):
hostname: str
"The node hostname on the network."
default_gateway: Optional[IPv4Address] = None
"The default gateway IP address for forwarding network traffic to other networks."
operating_state: NodeOperatingState = NodeOperatingState.OFF
"The hardware state of the node."
nics: Dict[str, NIC] = {}
"The NICs on the node."
ethernet_port: Dict[int, NIC] = {}
"The NICs on the node by port id."
accounts: Dict[str, Account] = {}
"All accounts on the node."
@@ -833,9 +906,12 @@ class Node(SimComponent):
This method initializes the ARP cache, ICMP handler, session manager, and software manager if they are not
provided.
"""
if kwargs.get("default_gateway"):
if not isinstance(kwargs["default_gateway"], IPv4Address):
kwargs["default_gateway"] = IPv4Address(kwargs["default_gateway"])
if not kwargs.get("sys_log"):
kwargs["sys_log"] = SysLog(kwargs["hostname"])
if not kwargs.get("arp_cache"):
if not kwargs.get("arp"):
kwargs["arp"] = ARPCache(sys_log=kwargs.get("sys_log"))
if not kwargs.get("icmp"):
kwargs["icmp"] = ICMP(sys_log=kwargs.get("sys_log"), arp_cache=kwargs.get("arp"))
@@ -874,18 +950,19 @@ class Node(SimComponent):
)
return state
def show(self):
"""Prints a table of the NICs on the Node.."""
from prettytable import PrettyTable
table = PrettyTable(["MAC Address", "Address", "Default Gateway", "Speed", "Status"])
for nic in self.nics.values():
def show(self, markdown: bool = False):
"""Prints a table of the NICs on the Node."""
table = PrettyTable(["Port", "MAC Address", "Address", "Speed", "Status"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.hostname} Network Interface Cards"
for port, nic in self.ethernet_port.items():
table.add_row(
[
port,
nic.mac_address,
f"{nic.ip_address}/{nic.ip_network.prefixlen}",
nic.gateway,
nic.speed,
"Enabled" if nic.enabled else "Disabled",
]
@@ -898,7 +975,8 @@ class Node(SimComponent):
self.operating_state = NodeOperatingState.ON
self.sys_log.info("Turned on")
for nic in self.nics.values():
nic.enable()
if nic.connected_link:
nic.enable()
def power_off(self):
"""Power off the Node, disabling its NICs if it is in the ON state."""
@@ -917,6 +995,7 @@ class Node(SimComponent):
"""
if nic.uuid not in self.nics:
self.nics[nic.uuid] = nic
self.ethernet_port[len(self.nics)] = nic
nic.connected_node = self
nic.parent = self
self.sys_log.info(f"Connected NIC {nic}")
@@ -938,6 +1017,10 @@ class Node(SimComponent):
if isinstance(nic, str):
nic = self.nics.get(nic)
if nic or nic.uuid in self.nics:
for port, _nic in self.ethernet_port.items():
if nic == _nic:
self.ethernet_port.pop(port)
break
self.nics.pop(nic.uuid)
nic.parent = None
nic.disable()
@@ -958,13 +1041,27 @@ class Node(SimComponent):
"""
if not isinstance(target_ip_address, IPv4Address):
target_ip_address = IPv4Address(target_ip_address)
if target_ip_address.is_loopback:
self.sys_log.info("Pinging loopback address")
return any(nic.enabled for nic in self.nics.values())
if self.operating_state == NodeOperatingState.ON:
self.sys_log.info(f"Attempting to ping {target_ip_address}")
self.sys_log.info(f"Pinging {target_ip_address}:")
sequence, identifier = 0, None
while sequence < pings:
sequence, identifier = self.icmp.ping(target_ip_address, sequence, identifier)
return True
self.sys_log.info("Ping failed as the node is turned off")
sequence, identifier = self.icmp.ping(target_ip_address, sequence, identifier, pings)
request_replies = self.icmp.request_replies.get(identifier)
passed = request_replies == pings
if request_replies:
self.icmp.request_replies.pop(identifier)
else:
request_replies = 0
self.sys_log.info(
f"Ping statistics for {target_ip_address}: "
f"Packets: Sent = {pings}, "
f"Received = {request_replies}, "
f"Lost = {pings-request_replies} ({(pings-request_replies)/pings*100}% loss)"
)
return passed
return False
def send_frame(self, frame: Frame):
@@ -973,7 +1070,7 @@ class Node(SimComponent):
:param frame: The Frame to be sent.
"""
nic: NIC = self._get_arp_cache_nic(frame.ip.dst_ip)
nic: NIC = self._get_arp_cache_nic(frame.ip.dst_ip_address)
nic.send_frame(frame)
def receive_frame(self, frame: Frame, from_nic: NIC):
@@ -986,13 +1083,18 @@ class Node(SimComponent):
:param frame: The Frame being received.
:param from_nic: The NIC that received the frame.
"""
if frame.ip:
if frame.ip.src_ip_address in self.arp:
self.arp.add_arp_cache_entry(
ip_address=frame.ip.src_ip_address, mac_address=frame.ethernet.src_mac_addr, nic=from_nic
)
if frame.ip.protocol == IPProtocol.TCP:
if frame.tcp.src_port == Port.ARP:
self.arp.process_arp_packet(from_nic=from_nic, arp_packet=frame.arp)
elif frame.ip.protocol == IPProtocol.UDP:
pass
elif frame.ip.protocol == IPProtocol.ICMP:
self.icmp.process_icmp(frame=frame)
self.icmp.process_icmp(frame=frame, from_nic=from_nic)
def install_service(self, service: Service) -> None:
"""

38
src/primaite/simulator/network/hardware/nodes/computer.py Normal file
View File

@@ -0,0 +1,38 @@
from primaite.simulator.network.hardware.base import NIC, Node
class Computer(Node):
"""
A basic Computer class.
Example:
>>> pc_a = Computer(
hostname="pc_a",
ip_address="192.168.1.10",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1"
)
>>> pc_a.power_on()
Instances of computer come 'pre-packaged' with the following:
* Core Functionality:
* ARP
* ICMP
* Packet Capture
* Sys Log
* Services:
* DNS Client
* FTP Client
* LDAP Client
* NTP Client
* Applications:
* Email Client
* Web Browser
* Processes:
* Placeholder
"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.connect_nic(NIC(ip_address=kwargs["ip_address"], subnet_mask=kwargs["subnet_mask"]))

759
src/primaite/simulator/network/hardware/nodes/router.py Normal file
View File

@@ -0,0 +1,759 @@
from __future__ import annotations
import secrets
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Dict, List, Optional, Tuple, Union
from prettytable import MARKDOWN, 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):
"""
Represents an Access Control List (ACL) rule.
:ivar ACLAction action: Action to be performed (Permit/Deny). Default is DENY.
:ivar Optional[IPProtocol] protocol: Network protocol. Default is None.
:ivar Optional[IPv4Address] src_ip_address: Source IP address. Default is None.
:ivar Optional[Port] src_port: Source port number. Default is None.
:ivar Optional[IPv4Address] dst_ip_address: Destination IP address. Default is None.
:ivar Optional[Port] dst_port: Destination port number. Default is None.
"""
action: ACLAction = ACLAction.DENY
protocol: Optional[IPProtocol] = None
src_ip_address: Optional[IPv4Address] = None
src_port: Optional[Port] = None
dst_ip_address: 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)
def describe_state(self) -> Dict:
"""
Describes the current state of the ACLRule.
:return: A dictionary representing the current state.
"""
pass
class AccessControlList(SimComponent):
"""
Manages a list of ACLRules to filter network traffic.
:ivar SysLog sys_log: System logging instance.
:ivar ACLAction implicit_action: Default action for rules.
:ivar ACLRule implicit_rule: Implicit ACL rule, created based on implicit_action.
:ivar int max_acl_rules: Maximum number of ACL rules that can be added. Default is 25.
:ivar List[Optional[ACLRule]] _acl: A list containing the ACL rules.
"""
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
kwargs["implicit_rule"] = ACLRule(action=kwargs["implicit_action"])
super().__init__(**kwargs)
self._acl = [None] * (self.max_acl_rules - 1)
def describe_state(self) -> Dict:
"""
Describes the current state of the AccessControlList.
:return: A dictionary representing the current state.
"""
pass
@property
def acl(self) -> List[Optional[ACLRule]]:
"""
Get the list of ACL rules.
:return: The list of ACL rules.
"""
return self._acl
def add_rule(
self,
action: ACLAction,
protocol: Optional[IPProtocol] = None,
src_ip_address: Optional[Union[str, IPv4Address]] = None,
src_port: Optional[Port] = None,
dst_ip_address: Optional[Union[str, IPv4Address]] = None,
dst_port: Optional[Port] = None,
position: int = 0,
) -> None:
"""
Add a new ACL rule.
:param ACLAction action: Action to be performed (Permit/Deny).
:param Optional[IPProtocol] protocol: Network protocol.
:param Optional[Union[str, IPv4Address]] src_ip_address: Source IP address.
:param Optional[Port] src_port: Source port number.
:param Optional[Union[str, IPv4Address]] dst_ip_address: Destination IP address.
:param Optional[Port] dst_port: Destination port number.
:param int position: Position in the ACL list to insert the rule.
:raises ValueError: When the position is out of bounds.
"""
if isinstance(src_ip_address, str):
src_ip_address = IPv4Address(src_ip_address)
if isinstance(dst_ip_address, str):
dst_ip_address = IPv4Address(dst_ip_address)
if 0 <= position < self.max_acl_rules:
if self._acl[position]:
self.sys_log.info(f"Overwriting ACL rule at position {position}")
self._acl[position] = ACLRule(
action=action,
src_ip_address=src_ip_address,
dst_ip_address=dst_ip_address,
protocol=protocol,
src_port=src_port,
dst_port=dst_port,
)
else:
raise ValueError(f"Cannot add ACL rule, position {position} is out of bounds.")
def remove_rule(self, position: int) -> None:
"""
Remove an ACL rule from a specific position.
:param int position: The position of the rule to be removed.
:raises ValueError: When the position is out of bounds.
"""
if 0 <= position < self.max_acl_rules - 1:
rule = self._acl[position] # noqa
self._acl[position] = None
del rule
else:
raise ValueError(f"Cannot remove ACL rule, position {position} is out of bounds.")
def is_permitted(
self,
protocol: IPProtocol,
src_ip_address: Union[str, IPv4Address],
src_port: Optional[Port],
dst_ip_address: Union[str, IPv4Address],
dst_port: Optional[Port],
) -> Tuple[bool, Optional[Union[str, ACLRule]]]:
"""
Check if a packet with the given properties is permitted through the ACL.
:param protocol: The protocol of the packet.
:param src_ip_address: Source IP address of the packet. Accepts string and IPv4Address.
:param src_port: Source port of the packet. Optional.
:param dst_ip_address: Destination IP address of the packet. Accepts string and IPv4Address.
:param dst_port: Destination port of the packet. Optional.
:return: A tuple with a boolean indicating if the packet is permitted and an optional rule or implicit action
string.
"""
if not isinstance(src_ip_address, IPv4Address):
src_ip_address = IPv4Address(src_ip_address)
if not isinstance(dst_ip_address, IPv4Address):
dst_ip_address = IPv4Address(dst_ip_address)
for rule in self._acl:
if not rule:
continue
if (
(rule.src_ip_address == src_ip_address or rule.src_ip_address is None)
and (rule.dst_ip_address == dst_ip_address or rule.dst_ip_address 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_address: Union[str, IPv4Address],
src_port: Port,
dst_ip_address: Union[str, IPv4Address],
dst_port: Port,
) -> List[ACLRule]:
"""
Get the list of relevant rules for a packet with given properties.
:param protocol: The protocol of the packet.
:param src_ip_address: Source IP address of the packet. Accepts string and IPv4Address.
:param src_port: Source port of the packet.
:param dst_ip_address: Destination IP address of the packet. Accepts string and IPv4Address.
:param dst_port: Destination port of the packet.
:return: A list of relevant ACLRules.
"""
if not isinstance(src_ip_address, IPv4Address):
src_ip_address = IPv4Address(src_ip_address)
if not isinstance(dst_ip_address, IPv4Address):
dst_ip_address = IPv4Address(dst_ip_address)
relevant_rules = []
for rule in self._acl:
if rule is None:
continue
if (
(rule.src_ip_address == src_ip_address or rule.src_ip_address is None)
or (rule.dst_ip_address == dst_ip_address or rule.dst_ip_address 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, markdown: bool = False):
"""
Display the current ACL rules as a table.
:param markdown: Whether to display the table in Markdown format. Defaults to False.
"""
table = PrettyTable(["Index", "Action", "Protocol", "Src IP", "Src Port", "Dst IP", "Dst Port"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
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_address if rule.src_ip_address else "ANY",
f"{rule.src_port.value} ({rule.src_port.name})" if rule.src_port else "ANY",
rule.dst_ip_address if rule.dst_ip_address 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_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.
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_ip_address: 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_ip_address"}:
if not isinstance(kwargs[key], IPv4Address):
kwargs[key] = IPv4Address(kwargs[key])
super().__init__(**kwargs)
def describe_state(self) -> Dict:
"""
Describes the current state of the RouteEntry.
:return: A dictionary representing the current state.
"""
pass
class RouteTable(SimComponent):
"""
Represents a routing table holding multiple route entries.
:ivar List[RouteEntry] routes: A list of RouteEntry objects.
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:
"""
Describes the current state of the RouteTable.
:return: A dictionary representing the current state.
"""
pass
def add_route(
self,
address: Union[IPv4Address, str],
subnet_mask: Union[IPv4Address, str],
next_hop_ip_address: Union[IPv4Address, str],
metric: float = 0.0,
):
"""
Add a route to the routing table.
:param address: The destination address of the route.
:param subnet_mask: The subnet mask of the route.
:param next_hop_ip_address: The next hop IP for the route.
:param metric: The metric of the route, default is 0.0.
"""
for key in {address, subnet_mask, next_hop_ip_address}:
if not isinstance(key, IPv4Address):
key = IPv4Address(key)
route = RouteEntry(
address=address, subnet_mask=subnet_mask, next_hop_ip_address=next_hop_ip_address, 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.
This method uses the Longest Prefix Match algorithm and considers metrics to find the best route.
:param destination_ip: The destination IP to find the route for.
:return: The best matching RouteEntry, or None if no route matches.
"""
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, markdown: bool = False):
"""
Display the current routing table as a table.
:param markdown: Whether to display the table in Markdown format. Defaults to False.
"""
table = PrettyTable(["Index", "Address", "Next Hop", "Metric"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
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_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):
"""
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_address == 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()
nic.send_frame(frame)
# 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
)
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:
arp_reply = arp_packet.generate_reply(from_nic.mac_address)
self.send_arp_reply(arp_reply, from_nic)
return
class RouterICMP(ICMP):
"""
A class to represent a router's Internet Control Message Protocol (ICMP) handler.
:param sys_log: System log for logging network events and errors.
:type sys_log: SysLog
:param arp_cache: The ARP cache for resolving MAC addresses.
:type arp_cache: ARPCache
:param router: The router to which this ICMP handler belongs.
:type router: Router
"""
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):
"""
Process incoming ICMP frames based on ICMP type.
:param frame: The incoming frame to process.
:param from_nic: The network interface where the frame is coming from.
:param is_reattempt: Flag to indicate if the process is a reattempt.
"""
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_address:
if nic.enabled:
# reply to the request
if not is_reattempt:
self.sys_log.info(f"Received echo request from {frame.ip.src_ip_address}")
target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip_address)
src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip_address)
tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP)
# Network Layer
ip_packet = IPPacket(
src_ip_address=nic.ip_address,
dst_ip_address=frame.ip.src_ip_address,
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,
)
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,
)
self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip_address}")
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_address:
if nic.enabled:
time = frame.transmission_duration()
time_str = f"{time}ms" if time > 0 else "<1ms"
self.sys_log.info(
f"Reply from {frame.ip.src_ip_address}: "
f"bytes={len(frame.payload)}, "
f"time={time_str}, "
f"TTL={frame.ip.ttl}"
)
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):
"""
A class to represent a network router node.
:ivar str hostname: The name of the router node.
:ivar int num_ports: The number of ports in the router.
:ivar dict kwargs: Optional keyword arguments for SysLog, ACL, RouteTable, RouterARPCache, RouterICMP.
"""
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]:
"""
Retrieve the port number for a given NIC.
:param target_nic: Target network interface.
:return: The port number if NIC is found, otherwise None.
"""
for port, nic in self.ethernet_ports.items():
if nic == target_nic:
return port
def describe_state(self) -> Dict:
"""
Describes the current state of the Router.
:return: A dictionary representing the current state.
"""
pass
def route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None:
"""
Route a given frame from a source NIC to its destination.
:param frame: The frame to be routed.
:param from_nic: The source network interface.
:param re_attempt: Flag to indicate if the routing is a reattempt.
"""
# Check if src ip is on network of one of the NICs
nic = self.arp.get_arp_cache_nic(frame.ip.dst_ip_address)
target_mac = self.arp.get_arp_cache_mac_address(frame.ip.dst_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)
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_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}")
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 a NIC and processes it based on its protocol.
:param frame: The incoming frame.
:param from_nic: The network interface where the frame is coming from.
"""
route_frame = False
protocol = frame.ip.protocol
src_ip_address = frame.ip.src_ip_address
dst_ip_address = frame.ip.dst_ip_address
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_address=src_ip_address,
src_port=src_port,
dst_ip_address=dst_ip_address,
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_address):
self.arp.add_arp_cache_entry(src_ip_address, 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]):
"""
Configure the IP settings of a given port.
:param port: The port to configure.
:param ip_address: The IP address to set.
:param subnet_mask: The subnet mask to set.
"""
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):
"""
Enable a given port on the router.
:param port: The port to enable.
"""
nic = self.ethernet_ports.get(port)
if nic:
nic.enable()
def disable_port(self, port: int):
"""
Disable a given port on the router.
:param port: The port to disable.
"""
nic = self.ethernet_ports.get(port)
if nic:
nic.disable()
def show(self, markdown: bool = False):
"""
Prints the state of the Ethernet interfaces on the Router.
:param markdown: Flag to indicate if the output should be in markdown format.
"""
"""Prints a table of the NICs on the Node."""
table = PrettyTable(["Port", "MAC Address", "Address", "Speed", "Status"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
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)

34
src/primaite/simulator/network/hardware/nodes/server.py Normal file
View File

@@ -0,0 +1,34 @@
from primaite.simulator.network.hardware.nodes.computer import Computer
class Server(Computer):
"""
A basic Server class.
Example:
>>> server_a = Server(
hostname="server_a",
ip_address="192.168.1.10",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1"
)
>>> server_a.power_on()
Instances of Server come 'pre-packaged' with the following:
* Core Functionality:
* ARP
* ICMP
* Packet Capture
* Sys Log
* Services:
* DNS Client
* FTP Client
* LDAP Client
* NTP Client
* Applications:
* Email Client
* Web Browser
* Processes:
* Placeholder
"""

121
src/primaite/simulator/network/hardware/nodes/switch.py Normal file
View File

@@ -0,0 +1,121 @@
from typing import Dict
from prettytable import MARKDOWN, PrettyTable
from primaite import getLogger
from primaite.exceptions import NetworkError
from primaite.links.link import Link
from primaite.simulator.network.hardware.base import Node, SwitchPort
from primaite.simulator.network.transmission.data_link_layer import Frame
_LOGGER = getLogger(__name__)
class Switch(Node):
"""
A class representing a Layer 2 network switch.
:ivar num_ports: The number of ports on the switch. Default is 24.
"""
num_ports: int = 24
"The number of ports on the switch."
switch_ports: Dict[int, SwitchPort] = {}
"The SwitchPorts on the switch."
mac_address_table: Dict[str, SwitchPort] = {}
"A MAC address table mapping destination MAC addresses to corresponding SwitchPorts."
def __init__(self, **kwargs):
super().__init__(**kwargs)
if not self.switch_ports:
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.parent = self
port.port_num = port_num
def show(self, markdown: bool = False):
"""
Prints a table of the SwitchPorts on the Switch.
:param markdown: If True, outputs the table in markdown format. Default is False.
"""
table = PrettyTable(["Port", "MAC Address", "Speed", "Status"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.hostname} Switch Ports"
for port_num, port in self.switch_ports.items():
table.add_row([port_num, port.mac_address, port.speed, "Enabled" if port.enabled else "Disabled"])
print(table)
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
"""
return {
"uuid": self.uuid,
"num_ports": self.num_ports, # redundant?
"ports": {port_num: port.describe_state() for port_num, port in self.switch_ports.items()},
"mac_address_table": {mac: port for mac, port in self.mac_address_table.items()},
}
def _add_mac_table_entry(self, mac_address: str, switch_port: SwitchPort):
"""
Private method to add an entry to the MAC address table.
:param mac_address: MAC address to be added.
:param switch_port: Corresponding SwitchPort object.
"""
mac_table_port = self.mac_address_table.get(mac_address)
if not mac_table_port:
self.mac_address_table[mac_address] = switch_port
self.sys_log.info(f"Added MAC table entry: Port {switch_port.port_num} -> {mac_address}")
else:
if mac_table_port != switch_port:
self.mac_address_table.pop(mac_address)
self.sys_log.info(f"Removed MAC table entry: Port {mac_table_port.port_num} -> {mac_address}")
self._add_mac_table_entry(mac_address, switch_port)
def forward_frame(self, frame: Frame, incoming_port: SwitchPort):
"""
Forward a frame to the appropriate port based on the destination MAC address.
:param frame: The Frame to be forwarded.
:param incoming_port: The port number from which the frame was received.
"""
src_mac = frame.ethernet.src_mac_addr
dst_mac = frame.ethernet.dst_mac_addr
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":
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:
port.send_frame(frame)
def disconnect_link_from_port(self, link: Link, port_number: int):
"""
Disconnect a given link from the specified port number on the switch.
:param link: The Link object to be disconnected.
:param port_number: The port number on the switch from where the link should be disconnected.
:raise NetworkError: When an invalid port number is provided or the link does not match the connection.
"""
port = self.switch_ports.get(port_number)
if port is None:
msg = f"Invalid port number {port_number} on the switch"
_LOGGER.error(msg)
raise NetworkError(msg)
if port.connected_link != link:
msg = f"The link does not match the connection at port number {port_number}"
_LOGGER.error(msg)
raise NetworkError(msg)
port.disconnect_link()