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

#2248 - Final run over all the docstrings after running pre-commit. All tests now working. Updated CHANGELOG.md.

Browse Source
This commit is contained in:
Chris McCarthy
2024-02-08 10:53:30 +00:00
parent 0c96fef3ec
commit 411f0a320f
41 changed files with 582 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/primaite/simulator/system/core/packet_capture.py
View File

@@ -108,4 +108,3 @@ class PacketCapture:
"""
msg = frame.model_dump_json()
self.outbound_logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL

112
src/primaite/simulator/system/core/session_manager.py
View File

@@ -102,7 +102,7 @@ class SessionManager:
@staticmethod
def _get_session_key(
frame: Frame, inbound_frame: bool = True
frame: Frame, inbound_frame: bool = True
) -> Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]]:
"""
Extracts the session key from the given frame.
@@ -140,27 +140,76 @@ class SessionManager:
dst_port = None
return protocol, with_ip_address, src_port, dst_port
def resolve_outbound_network_interface(self, dst_ip_address: IPv4Address) -> Optional['NetworkInterface']:
def resolve_outbound_network_interface(self, dst_ip_address: IPv4Address) -> Optional["NetworkInterface"]:
"""
Resolves the appropriate outbound network interface for a given destination IP address.
This method determines the most suitable network interface for sending a packet to the specified
destination IP address. It considers only enabled network interfaces and checks if the destination
IP address falls within the subnet of each interface. If no suitable local network interface is found,
the method defaults to using the network interface associated with the default gateway.
The search process prioritises local network interfaces based on the IP network to which they belong.
If the destination IP address does not match any local subnet, the method assumes that the destination
is outside the local network and hence, routes the packet through the default gateway's network interface.
:param dst_ip_address: The destination IP address for which the outbound interface is to be resolved.
:type dst_ip_address: IPv4Address
:return: The network interface through which the packet should be sent to reach the destination IP address,
or the default gateway's network interface if the destination is not within any local subnet.
:rtype: Optional["NetworkInterface"]
"""
for network_interface in self.node.network_interfaces.values():
if dst_ip_address in network_interface.ip_network and network_interface.enabled:
return network_interface
return self.software_manager.arp.get_default_gateway_network_interface()
def resolve_outbound_transmission_details(
self,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
protocol: Optional[IPProtocol] = None,
session_id: Optional[str] = None
self,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
protocol: Optional[IPProtocol] = None,
session_id: Optional[str] = None,
) -> Tuple[
Optional['NetworkInterface'],
Optional[str], IPv4Address,
Optional["NetworkInterface"],
Optional[str],
IPv4Address,
Optional[Port],
Optional[Port],
Optional[IPProtocol],
bool
bool,
]:
"""
Resolves the necessary details for outbound transmission based on the provided parameters.
This method determines whether the payload should be broadcast or unicast based on the destination IP address
and resolves the outbound network interface and destination MAC address accordingly.
The method first checks if `session_id` is provided and uses the session details if available. For broadcast
transmissions, it finds a suitable network interface and uses a broadcast MAC address. For unicast
transmissions, it attempts to resolve the destination MAC address using ARP and finds the appropriate
outbound network interface. If the destination IP address is outside the local network and no specific MAC
address is resolved, it uses the default gateway for the transmission.
:param dst_ip_address: The destination IP address or network. If an IPv4Network is provided, the method
treats the transmission as a broadcast to that network. Optional.
:type dst_ip_address: Optional[Union[IPv4Address, IPv4Network]]
:param src_port: The source port number for the transmission. Optional.
:type src_port: Optional[Port]
:param dst_port: The destination port number for the transmission. Optional.
:type dst_port: Optional[Port]
:param protocol: The IP protocol to be used for the transmission. Optional.
:type protocol: Optional[IPProtocol]
:param session_id: The session ID associated with the transmission. If provided, the session details override
other parameters. Optional.
:type session_id: Optional[str]
:return: A tuple containing the resolved outbound network interface, destination MAC address, destination IP
address, source port, destination port, protocol, and a boolean indicating whether the transmission is a
broadcast.
:rtype: Tuple[Optional["NetworkInterface"], Optional[str], IPv4Address, Optional[Port], Optional[Port],
Optional[IPProtocol], bool]
"""
if dst_ip_address and not isinstance(dst_ip_address, (IPv4Address, IPv4Network)):
dst_ip_address = IPv4Address(dst_ip_address)
is_broadcast = False
@@ -207,14 +256,14 @@ class SessionManager:
return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast
def receive_payload_from_software_manager(
self,
payload: Any,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
session_id: Optional[str] = None,
ip_protocol: IPProtocol = IPProtocol.TCP,
icmp_packet: Optional[ICMPPacket] = None
self,
payload: Any,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
session_id: Optional[str] = None,
ip_protocol: IPProtocol = IPProtocol.TCP,
icmp_packet: Optional[ICMPPacket] = None,
) -> Union[Any, None]:
"""
Receive a payload from the SoftwareManager and send it to the appropriate NIC for transmission.
@@ -239,15 +288,22 @@ class SessionManager:
is_broadcast = payload.request
ip_protocol = IPProtocol.UDP
else:
vals = self.resolve_outbound_transmission_details(
dst_ip_address=dst_ip_address,
src_port=src_port,
dst_port=dst_port,
protocol=ip_protocol,
session_id=session_id
session_id=session_id,
)
outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast = vals
(
outbound_network_interface,
dst_mac_address,
dst_ip_address,
src_port,
dst_port,
protocol,
is_broadcast,
) = vals
if protocol:
ip_protocol = protocol
@@ -257,7 +313,7 @@ class SessionManager:
if not (src_port or dst_port):
raise ValueError(
f"Failed to resolve src or dst port. Have you sent the port from the service or application?"
"Failed to resolve src or dst port. Have you sent the port from the service or application?"
)
tcp_header = None
@@ -283,7 +339,11 @@ class SessionManager:
# Construct the frame for transmission
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=outbound_network_interface.mac_address, dst_mac_addr=dst_mac_address),
ip=IPPacket(src_ip_address=outbound_network_interface.ip_address, dst_ip_address=dst_ip_address, protocol=ip_protocol),
ip=IPPacket(
src_ip_address=outbound_network_interface.ip_address,
dst_ip_address=dst_ip_address,
protocol=ip_protocol,
),
tcp=tcp_header,
udp=udp_header,
icmp=icmp_packet,
@@ -304,7 +364,7 @@ class SessionManager:
# Send the frame through the NIC
return outbound_network_interface.send_frame(frame)
def receive_frame(self, frame: Frame, from_network_interface: 'NetworkInterface'):
def receive_frame(self, frame: Frame, from_network_interface: "NetworkInterface"):
"""
Receive a Frame.
@@ -334,7 +394,7 @@ class SessionManager:
protocol=frame.ip.protocol,
session_id=session.uuid,
from_network_interface=from_network_interface,
frame=frame
frame=frame,
)
def show(self, markdown: bool = False):

29
src/primaite/simulator/system/core/software_manager.py
View File

@@ -25,7 +25,14 @@ IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware)
class SoftwareManager:
"""A class that manages all running Services and Applications on a Node and facilitates their communication."""
"""
Manages all running services and applications on a network node and facilitates their communication.
This class is responsible for installing, uninstalling, and managing the operational state of various network
services and applications. It acts as a bridge between the node's session manager and its software components,
ensuring that incoming and outgoing network payloads are correctly routed to and from the appropriate services
or applications.
"""
def __init__(
self,
@@ -50,11 +57,13 @@ class SoftwareManager:
self.dns_server: Optional[IPv4Address] = dns_server
@property
def arp(self) -> 'ARP':
def arp(self) -> "ARP":
"""Provides access to the ARP service instance, if installed."""
return self.software.get("ARP") # noqa
@property
def icmp(self) -> 'ICMP':
def icmp(self) -> "ICMP":
"""Provides access to the ICMP service instance, if installed."""
return self.software.get("ICMP") # noqa
def get_open_ports(self) -> List[Port]:
@@ -167,7 +176,13 @@ class SoftwareManager:
)
def receive_payload_from_session_manager(
self, payload: Any, port: Port, protocol: IPProtocol, session_id: str, from_network_interface: "NIC", frame: Frame
self,
payload: Any,
port: Port,
protocol: IPProtocol,
session_id: str,
from_network_interface: "NIC",
frame: Frame,
):
"""
Receive a payload from the SessionManager and forward it to the corresponding service or application.
@@ -177,7 +192,9 @@ class SoftwareManager:
"""
receiver: Optional[Union[Service, Application]] = self.port_protocol_mapping.get((port, protocol), None)
if receiver:
receiver.receive(payload=payload, session_id=session_id, from_network_interface=from_network_interface, frame=frame)
receiver.receive(
payload=payload, session_id=session_id, from_network_interface=from_network_interface, frame=frame
)
else:
self.sys_log.error(f"No service or application found for port {port} and protocol {protocol}")
pass
@@ -202,7 +219,7 @@ class SoftwareManager:
software.operating_state.name,
software.health_state_actual.name,
software.port.value if software.port != Port.NONE else None,
software.protocol.value
software.protocol.value,
]
)
print(table)

30
src/primaite/simulator/system/services/arp/arp.py
View File

@@ -20,6 +20,7 @@ class ARP(Service):
Manages ARP for resolving network layer addresses into link layer addresses. It maintains an ARP cache,
sends ARP requests and replies, and processes incoming ARP packets.
"""
arp: Dict[IPV4Address, ARPEntry] = {}
def __init__(self, **kwargs):
@@ -29,6 +30,14 @@ class ARP(Service):
super().__init__(**kwargs)
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
"""
state = super().describe_state()
state.update({str(ip): arp_entry.mac_address for ip, arp_entry in self.arp.items()})
return super().describe_state()
def show(self, markdown: bool = False):
@@ -57,11 +66,7 @@ class ARP(Service):
self.arp.clear()
def add_arp_cache_entry(
self,
ip_address: IPV4Address,
mac_address: str,
network_interface: NetworkInterface,
override: bool = False
self, ip_address: IPV4Address, mac_address: str, network_interface: NetworkInterface, override: bool = False
):
"""
Add an ARP entry to the cache.
@@ -139,7 +144,8 @@ class ARP(Service):
)
else:
self.sys_log.error(
"Cannot send ARP request as there is no outbound Network Interface to use. Try configuring the default gateway."
"Cannot send ARP request as there is no outbound Network Interface to use. Try configuring the default "
"gateway."
)
def send_arp_reply(self, arp_reply: ARPPacket):
@@ -147,12 +153,10 @@ class ARP(Service):
Sends an ARP reply in response to an ARP request.
:param arp_reply: The ARP packet containing the reply.
:param from_network_interface: The NIC from which the ARP reply is sent.
"""
outbound_network_interface = self.software_manager.session_manager.resolve_outbound_network_interface(
arp_reply.target_ip_address
)
)
if outbound_network_interface:
self.sys_log.info(
f"Sending ARP reply from {arp_reply.sender_mac_addr}/{arp_reply.sender_ip_address} "
@@ -162,14 +166,14 @@ class ARP(Service):
payload=arp_reply,
dst_ip_address=arp_reply.target_ip_address,
dst_port=self.port,
ip_protocol=self.protocol
ip_protocol=self.protocol,
)
else:
self.sys_log.error(
"Cannot send ARP reply as there is no outbound Network Interface to use. Try configuring the default gateway."
"Cannot send ARP reply as there is no outbound Network Interface to use. Try configuring the default "
"gateway."
)
@abstractmethod
def _process_arp_request(self, arp_packet: ARPPacket, from_network_interface: NetworkInterface):
"""
@@ -197,7 +201,7 @@ class ARP(Service):
self.add_arp_cache_entry(
ip_address=arp_packet.sender_ip_address,
mac_address=arp_packet.sender_mac_addr,
network_interface=from_network_interface
network_interface=from_network_interface,
)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:

33
src/primaite/simulator/system/services/icmp/icmp.py
View File

@@ -1,8 +1,9 @@
import secrets
from ipaddress import IPv4Address
from typing import Dict, Any, Union, Optional, Tuple
from typing import Any, Dict, Optional, Tuple, Union
from primaite import getLogger
from primaite.simulator.network.hardware.base import NetworkInterface
from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.network.transmission.network_layer import IPProtocol
@@ -19,6 +20,7 @@ class ICMP(Service):
Enables the sending and receiving of ICMP messages such as echo requests and replies. This is typically used for
network diagnostics, notably the ping command.
"""
request_replies: Dict = {}
def __init__(self, **kwargs):
@@ -28,7 +30,12 @@ class ICMP(Service):
super().__init__(**kwargs)
def describe_state(self) -> Dict:
pass
"""
Produce a dictionary describing the current state of this object.
:return: Current state of this object and child objects.
"""
return super().describe_state()
def clear(self):
"""
@@ -56,9 +63,7 @@ class ICMP(Service):
self.sys_log.info(f"Pinging {target_ip_address}:", to_terminal=True)
sequence, identifier = 0, None
while sequence < pings:
sequence, identifier = self._send_icmp_echo_request(
target_ip_address, sequence, identifier, pings
)
sequence, identifier = self._send_icmp_echo_request(target_ip_address, sequence, identifier, pings)
request_replies = self.software_manager.icmp.request_replies.get(identifier)
passed = request_replies == pings
if request_replies:
@@ -76,7 +81,7 @@ class ICMP(Service):
return passed
def _send_icmp_echo_request(
self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None, pings: int = 4
self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None, pings: int = 4
) -> Tuple[int, Union[int, None]]:
"""
Sends an ICMP echo request to a specified target IP address.
@@ -91,7 +96,8 @@ class ICMP(Service):
if not network_interface:
self.sys_log.error(
"Cannot send ICMP echo request as there is no outbound Network Interface to use. Try configuring the default gateway."
"Cannot send ICMP echo request as there is no outbound Network Interface to use. Try configuring the "
"default gateway."
)
return pings, None
@@ -105,11 +111,11 @@ class ICMP(Service):
dst_ip_address=target_ip_address,
dst_port=self.port,
ip_protocol=self.protocol,
icmp_packet=icmp_packet
icmp_packet=icmp_packet,
)
return sequence, icmp_packet.identifier
def _process_icmp_echo_request(self, frame: Frame, from_network_interface):
def _process_icmp_echo_request(self, frame: Frame, from_network_interface: NetworkInterface):
"""
Processes an ICMP echo request received by the service.
@@ -121,11 +127,12 @@ class ICMP(Service):
network_interface = self.software_manager.session_manager.resolve_outbound_network_interface(
frame.ip.src_ip_address
)
)
if not network_interface:
self.sys_log.error(
"Cannot send ICMP echo reply as there is no outbound Network Interface to use. Try configuring the default gateway."
"Cannot send ICMP echo reply as there is no outbound Network Interface to use. Try configuring the "
"default gateway."
)
return
@@ -143,7 +150,7 @@ class ICMP(Service):
dst_ip_address=frame.ip.src_ip_address,
dst_port=self.port,
ip_protocol=self.protocol,
icmp_packet=icmp_packet
icmp_packet=icmp_packet,
)
def _process_icmp_echo_reply(self, frame: Frame):
@@ -159,7 +166,7 @@ class ICMP(Service):
f"bytes={len(frame.payload)}, "
f"time={time_str}, "
f"TTL={frame.ip.ttl}",
to_terminal=True
to_terminal=True,
)
if not self.request_replies.get(frame.icmp.identifier):
self.request_replies[frame.icmp.identifier] = 0

6
src/primaite/simulator/system/services/ntp/ntp_server.py
View File

@@ -70,10 +70,6 @@ class NTPServer(Service):
payload = payload.generate_reply(time)
# send reply
self.software_manager.session_manager.receive_payload_from_software_manager(
payload=payload,
src_port=self.port,
dst_port=self.port,
ip_protocol=self.protocol,
session_id=session_id
payload=payload, src_port=self.port, dst_port=self.port, ip_protocol=self.protocol, session_id=session_id
)
return True

2
src/primaite/simulator/system/software.py
View File

@@ -380,7 +380,7 @@ class IOSoftware(Software):
dest_ip_address=dest_ip_address,
dest_port=dest_port,
ip_protocol=ip_protocol,
session_id=session_id
session_id=session_id,
)
@abstractmethod