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

#2248 - Initial work has been done on moving ICMP into services. still tidying up to be done. Need to fix tests too.

Browse Source
This commit is contained in:
Chris McCarthy
2024-02-02 15:35:02 +00:00
parent 1964ab4635
commit 87d9d6da04
14 changed files with 495 additions and 373 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -293,8 +293,15 @@ class SessionManager:
dst_port = frame.tcp.dst_port
elif frame.udp:
dst_port = frame.udp.dst_port
elif frame.icmp:
dst_port = Port.NONE
self.software_manager.receive_payload_from_session_manager(
payload=frame.payload, port=dst_port, protocol=frame.ip.protocol, session_id=session.uuid, from_nic=from_nic
payload=frame.payload,
port=dst_port,
protocol=frame.ip.protocol,
session_id=session.uuid,
from_nic=from_nic,
frame=frame
)
def show(self, markdown: bool = False):

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

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.file_system.file_system import FileSystem
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application, ApplicationOperatingState
@@ -16,6 +17,7 @@ if TYPE_CHECKING:
from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.network.hardware.base import Node, NIC
from primaite.simulator.system.services.arp.arp import ARP
from primaite.simulator.system.services.icmp.icmp import ICMP
from typing import Type, TypeVar
@@ -51,6 +53,10 @@ class SoftwareManager:
def arp(self) -> 'ARP':
return self.software.get("ARP") # noqa
@property
def icmp(self) -> 'ICMP':
return self.software.get("ICMP") # noqa
def get_open_ports(self) -> List[Port]:
"""
Get a list of open ports.
@@ -160,7 +166,7 @@ class SoftwareManager:
)
def receive_payload_from_session_manager(
self, payload: Any, port: Port, protocol: IPProtocol, session_id: str, from_nic: "NIC"
self, payload: Any, port: Port, protocol: IPProtocol, session_id: str, from_nic: "NIC", frame: Frame
):
"""
Receive a payload from the SessionManager and forward it to the corresponding service or application.
@@ -170,7 +176,7 @@ 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_nic=from_nic)
receiver.receive(payload=payload, session_id=session_id, from_nic=from_nic, frame=frame)
else:
self.sys_log.error(f"No service or application found for port {port} and protocol {protocol}")
pass
@@ -181,7 +187,7 @@ class SoftwareManager:
:param markdown: If True, outputs the table in markdown format. Default is False.
"""
table = PrettyTable(["Name", "Type", "Operating State", "Health State", "Port"])
table = PrettyTable(["Name", "Type", "Operating State", "Health State", "Port", "Protocol"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
@@ -194,7 +200,8 @@ class SoftwareManager:
software_type,
software.operating_state.name,
software.health_state_actual.name,
software.port.value,
software.port.value if software.port != Port.NONE else None,
software.protocol.value
]
)
print(table)

25
src/primaite/simulator/system/core/sys_log.py
View File

@@ -88,47 +88,62 @@ class SysLog:
root.mkdir(exist_ok=True, parents=True)
return root / f"{self.hostname}_sys.log"
def debug(self, msg: str):
def debug(self, msg: str, to_terminal: bool = False):
"""
Logs a message with the DEBUG level.
:param msg: The message to be logged.
:param to_terminal: If True, prints to the terminal too.
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.debug(msg)
if to_terminal:
print(msg)
def info(self, msg: str):
def info(self, msg: str, to_terminal: bool = False):
"""
Logs a message with the INFO level.
:param msg: The message to be logged.
:param to_terminal: If True, prints to the terminal too.
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.info(msg)
if to_terminal:
print(msg)
def warning(self, msg: str):
def warning(self, msg: str, to_terminal: bool = False):
"""
Logs a message with the WARNING level.
:param msg: The message to be logged.
:param to_terminal: If True, prints to the terminal too.
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.warning(msg)
if to_terminal:
print(msg)
def error(self, msg: str):
def error(self, msg: str, to_terminal: bool = False):
"""
Logs a message with the ERROR level.
:param msg: The message to be logged.
:param to_terminal: If True, prints to the terminal too.
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.error(msg)
if to_terminal:
print(msg)
def critical(self, msg: str):
def critical(self, msg: str, to_terminal: bool = False):
"""
Logs a message with the CRITICAL level.
:param msg: The message to be logged.
:param to_terminal: If True, prints to the terminal too.
"""
if SIM_OUTPUT.save_sys_logs:
self.logger.critical(msg)
if to_terminal:
print(msg)

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

@@ -2,17 +2,15 @@ from __future__ import annotations
from abc import abstractmethod
from ipaddress import IPv4Address
from typing import Any, Dict, Optional, Tuple, Union
from typing import Any, Dict, Optional, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel
from primaite.simulator.network.hardware.base import NIC
from primaite.simulator.network.protocols.arp import ARPEntry, ARPPacket
from primaite.simulator.network.protocols.packet import DataPacket
from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame
from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader, UDPHeader
from primaite.simulator.network.transmission.transport_layer import Port, UDPHeader
from primaite.simulator.system.services.service import Service
@@ -191,7 +189,6 @@ class ARP(Service):
from_nic = kwargs.get("from_nic")
if payload.request:
print(from_nic)
self._process_arp_request(arp_packet=payload, from_nic=from_nic)
else:
self._process_arp_reply(arp_packet=payload, from_nic=from_nic)

0
src/primaite/simulator/system/services/icmp/__init__.py Normal file
View File

159
src/primaite/simulator/system/services/icmp/icmp.py Normal file
View File

@@ -0,0 +1,159 @@
import secrets
from ipaddress import IPv4Address
from typing import Dict, Any, Union, Optional, Tuple
from primaite import getLogger
from primaite.simulator.network.hardware.base import NIC
from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType
from primaite.simulator.network.transmission.data_link_layer import Frame, EthernetHeader
from primaite.simulator.network.transmission.network_layer import IPProtocol, IPPacket
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.service import Service
_LOGGER = getLogger(__name__)
class ICMP(Service):
request_replies: Dict = {}
def __init__(self, **kwargs):
kwargs["name"] = "ICMP"
kwargs["port"] = Port.NONE
kwargs["protocol"] = IPProtocol.ICMP
super().__init__(**kwargs)
def describe_state(self) -> Dict:
pass
def clear(self):
"""Clears the ICMP request replies tracker."""
self.request_replies.clear()
def _send_icmp_echo_request(
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.
:param target_ip_address: The target IP address to send the ping.
:param sequence: The sequence number of the echo request. Defaults to 0.
:param identifier: An optional identifier for the ICMP packet. If None, a default will be used.
:return: A tuple containing the next sequence number and the identifier, or (0, None) if the target IP address
was not found in the ARP cache.
"""
nic = self.software_manager.arp.get_arp_cache_nic(target_ip_address)
if not nic:
return pings, None
# ARP entry exists
sequence += 1
target_mac_address = self.software_manager.arp.get_arp_cache_mac_address(target_ip_address)
src_nic = self.software_manager.arp.get_arp_cache_nic(target_ip_address)
# Network Layer
ip_packet = IPPacket(
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)
payload = secrets.token_urlsafe(int(32 / 1.3)) # Standard ICMP 32 bytes size
frame = Frame(ethernet=ethernet_header, ip=ip_packet, icmp=icmp_packet, payload=payload)
nic.send_frame(frame)
return sequence, icmp_packet.identifier
def ping(self, target_ip_address: Union[IPv4Address, str], pings: int = 4) -> bool:
"""
Ping an IP address, performing a standard ICMP echo request/response.
:param target_ip_address: The target IP address to ping.
:param pings: The number of pings to attempt, default is 4.
:return: True if the ping is successful, otherwise False.
"""
if not self._can_perform_action():
return False
if target_ip_address.is_loopback:
self.sys_log.info("Pinging loopback address")
return any(nic.enabled for nic in self.nics.values())
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
)
request_replies = self.software_manager.icmp.request_replies.get(identifier)
passed = request_replies == pings
if request_replies:
self.software_manager.icmp.request_replies.pop(identifier)
else:
request_replies = 0
output = (
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)"
)
self.sys_log.info(output, to_terminal=True)
return passed
def _process_icmp_echo_request(self, frame: Frame, from_nic: NIC, is_reattempt: bool = False):
if frame.icmp.icmp_type == ICMPType.ECHO_REQUEST:
if not is_reattempt:
self.sys_log.info(f"Received echo request from {frame.ip.src_ip_address}")
target_mac_address = self.software_manager.arp.get_arp_cache_mac_address(frame.ip.src_ip_address)
src_nic = self.software_manager.arp.get_arp_cache_nic(frame.ip.src_ip_address)
if not src_nic:
self.software_manager.arp.send_arp_request(frame.ip.src_ip_address)
self.process_icmp(frame=frame, from_nic=from_nic, is_reattempt=True)
return
# Network Layer
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(
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, 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)
def _process_icmp_echo_reply(self, frame: Frame, from_nic: NIC, is_reattempt: bool = False):
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}",
to_terminal=True
)
if not self.request_replies.get(frame.icmp.identifier):
self.request_replies[frame.icmp.identifier] = 0
self.request_replies[frame.icmp.identifier] += 1
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
frame: Frame = kwargs["frame"]
from_nic = kwargs["from_nic"]
if not frame.icmp:
return False
if frame.icmp.icmp_type == ICMPType.ECHO_REQUEST:
self._process_icmp_echo_request(frame, from_nic)
elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY:
self._process_icmp_echo_reply(frame, from_nic)
return True

90
src/primaite/simulator/system/services/icmp/router_icmp.py Normal file
View File

@@ -0,0 +1,90 @@
# 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.process_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.process_frame(frame, from_nic)