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

Merge remote-tracking branch 'origin/dev' into feature/2317-refactor-reset

Browse Source
This commit is contained in:
Marek Wolan
2024-02-25 16:41:40 +00:00
parent c115095157 8f85555709
commit a34cf08209
124 changed files with 7441 additions and 3680 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

233
src/primaite/simulator/system/services/arp/arp.py Normal file
View File

@@ -0,0 +1,233 @@
from __future__ import annotations
from abc import abstractmethod
from typing import Any, Dict, Optional, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.network.hardware.base import NetworkInterface
from primaite.simulator.network.protocols.arp import ARPEntry, ARPPacket
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.service import Service
from primaite.utils.validators import IPV4Address
class ARP(Service):
"""
The ARP (Address Resolution Protocol) 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):
kwargs["name"] = "ARP"
kwargs["port"] = Port.ARP
kwargs["protocol"] = IPProtocol.UDP
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):
"""
Prints the current state of the ARP cache in a table format.
:param markdown: If True, format the output as Markdown. Otherwise, use plain text.
"""
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.software_manager.node.network_interfaces[arp.network_interface_uuid].mac_address,
]
)
print(table)
def clear(self):
"""Clears the arp cache."""
self.arp.clear()
def add_arp_cache_entry(
self, ip_address: IPV4Address, mac_address: str, network_interface: NetworkInterface, 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 network_interface: 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.
"""
for _network_interface in self.software_manager.node.network_interfaces.values():
if _network_interface.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 {network_interface}")
arp_entry = ARPEntry(mac_address=mac_address, network_interface_uuid=network_interface.uuid)
self.arp[ip_address] = arp_entry
@abstractmethod
def get_arp_cache_mac_address(self, ip_address: IPV4Address) -> Optional[str]:
"""
Retrieves the MAC address associated with a given IP address from the ARP cache.
:param ip_address: The IP address to look up.
:return: The associated MAC address, if found. Otherwise, returns None.
"""
pass
@abstractmethod
def get_arp_cache_network_interface(self, ip_address: IPV4Address) -> Optional[NetworkInterface]:
"""
Retrieves the NIC associated with a given IP address from the ARP cache.
:param ip_address: The IP address to look up.
:return: The associated NIC, if found. Otherwise, returns None.
"""
pass
def send_arp_request(self, target_ip_address: Union[IPV4Address, str]):
"""
Sends an ARP request to resolve the MAC address of a target IP address.
:param target_ip_address: The target IP address for which the MAC address is being requested.
"""
if target_ip_address in self.arp:
return
use_default_gateway = True
for network_interface in self.software_manager.node.network_interfaces.values():
if target_ip_address in network_interface.ip_network:
use_default_gateway = False
break
if use_default_gateway:
if self.software_manager.node.default_gateway:
target_ip_address = self.software_manager.node.default_gateway
else:
return
outbound_network_interface = self.software_manager.session_manager.resolve_outbound_network_interface(
target_ip_address
)
if outbound_network_interface:
self.sys_log.info(f"Sending ARP request from NIC {outbound_network_interface} for ip {target_ip_address}")
arp_packet = ARPPacket(
sender_ip_address=outbound_network_interface.ip_address,
sender_mac_addr=outbound_network_interface.mac_address,
target_ip_address=target_ip_address,
)
self.software_manager.session_manager.receive_payload_from_software_manager(
payload=arp_packet, dst_ip_address=target_ip_address, dst_port=self.port, ip_protocol=self.protocol
)
else:
self.sys_log.error(
"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):
"""
Sends an ARP reply in response to an ARP request.
:param arp_reply: The ARP packet containing the reply.
"""
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} "
f"to {arp_reply.target_ip_address}/{arp_reply.target_mac_addr} "
)
self.software_manager.session_manager.receive_payload_from_software_manager(
payload=arp_reply,
dst_ip_address=arp_reply.target_ip_address,
dst_port=self.port,
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."
)
@abstractmethod
def _process_arp_request(self, arp_packet: ARPPacket, from_network_interface: NetworkInterface):
"""
Processes an incoming ARP request.
:param arp_packet: The ARP packet containing the request.
:param from_network_interface: The NIC that received the 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} "
)
def _process_arp_reply(self, arp_packet: ARPPacket, from_network_interface: NetworkInterface):
"""
Processes an incoming ARP reply.
:param arp_packet: The ARP packet containing the reply.
:param from_network_interface: The NIC that received the ARP reply.
"""
self.sys_log.info(
f"Received ARP response for {arp_packet.sender_ip_address} "
f"from {arp_packet.sender_mac_addr} via Network Interface {from_network_interface}"
)
self.add_arp_cache_entry(
ip_address=arp_packet.sender_ip_address,
mac_address=arp_packet.sender_mac_addr,
network_interface=from_network_interface,
)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
"""
Processes received data, handling ARP packets.
:param payload: The payload received.
:param session_id: The session ID associated with the received data.
:param kwargs: Additional keyword arguments.
:return: True if the payload was processed successfully, otherwise False.
"""
if not super().receive(payload, session_id, **kwargs):
return False
from_network_interface = kwargs["from_network_interface"]
if payload.request:
self._process_arp_request(arp_packet=payload, from_network_interface=from_network_interface)
else:
self._process_arp_reply(arp_packet=payload, from_network_interface=from_network_interface)
return True
def __contains__(self, item: Any) -> bool:
"""
Checks if an item is in the ARP cache.
:param item: The item to check.
:return: True if the item is in the cache, otherwise False.
"""
return item in self.arp

23
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -170,7 +170,7 @@ class DatabaseService(Service):
}
def _process_sql(
self, query: Literal["SELECT", "DELETE"], query_id: str, connection_id: Optional[str] = None
self, query: Literal["SELECT", "DELETE", "INSERT"], query_id: str, connection_id: Optional[str] = None
) -> Dict[str, Union[int, List[Any]]]:
"""
Executes the given SQL query and returns the result.
@@ -178,6 +178,7 @@ class DatabaseService(Service):
Possible queries:
- SELECT : returns the data
- DELETE : deletes the data
- INSERT : inserts the data
:param query: The SQL query to be executed.
:return: Dictionary containing status code and data fetched.
@@ -201,9 +202,27 @@ class DatabaseService(Service):
return {"status_code": 404, "data": False}
elif query == "DELETE":
self.db_file.health_status = FileSystemItemHealthStatus.COMPROMISED
return {"status_code": 200, "type": "sql", "data": False, "uuid": query_id, "connection_id": connection_id}
return {
"status_code": 200,
"type": "sql",
"data": False,
"uuid": query_id,
"connection_id": connection_id,
}
elif query == "INSERT":
if self.health_state_actual == SoftwareHealthState.GOOD:
return {
"status_code": 200,
"type": "sql",
"data": False,
"uuid": query_id,
"connection_id": connection_id,
}
else:
return {"status_code": 404, "data": False}
else:
# Invalid query
self.sys_log.info(f"{self.name}: Invalid {query}")
return {"status_code": 500, "data": False}
def describe_state(self) -> Dict:

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

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

@@ -0,0 +1,194 @@
import secrets
from ipaddress import IPv4Address
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
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.service import Service
_LOGGER = getLogger(__name__)
class ICMP(Service):
"""
The Internet Control Message Protocol (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):
kwargs["name"] = "ICMP"
kwargs["port"] = Port.NONE
kwargs["protocol"] = IPProtocol.ICMP
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.
"""
return super().describe_state()
def clear(self):
"""
Clears the ICMP request and reply tracker.
This is typically used to reset the state of the service, removing all tracked ICMP echo requests and their
corresponding replies.
"""
self.request_replies.clear()
def ping(self, target_ip_address: Union[IPv4Address, str], pings: int = 4) -> bool:
"""
Pings a target IP address by sending an ICMP echo request and waiting for a reply.
:param target_ip_address: The IP address to be pinged.
:param pings: The number of echo requests to send. Defaults to 4.
:return: True if the ping was successful (i.e., if a reply was received for every request sent), 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(network_interface.enabled for network_interface in self.network_interfaces.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 _send_icmp_echo_request(
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.
:param target_ip_address: The target IP address for the echo request.
:param sequence: The sequence number of the echo request.
:param identifier: The identifier for the ICMP packet. If None, a default identifier is used.
:param pings: The number of pings to send. Defaults to 4.
:return: A tuple containing the next sequence number and the identifier.
"""
network_interface = self.software_manager.session_manager.resolve_outbound_network_interface(target_ip_address)
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."
)
return pings, None
sequence += 1
icmp_packet = ICMPPacket(identifier=identifier, sequence=sequence)
payload = secrets.token_urlsafe(int(32 / 1.3)) # Standard ICMP 32 bytes size
self.software_manager.session_manager.receive_payload_from_software_manager(
payload=payload,
dst_ip_address=target_ip_address,
dst_port=self.port,
ip_protocol=self.protocol,
icmp_packet=icmp_packet,
)
return sequence, icmp_packet.identifier
def _process_icmp_echo_request(self, frame: Frame, from_network_interface: NetworkInterface):
"""
Processes an ICMP echo request received by the service.
:param frame: The network frame containing the ICMP echo request.
"""
if frame.ip.dst_ip_address != from_network_interface.ip_address:
return
self.sys_log.info(f"Received echo request from {frame.ip.src_ip_address}")
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."
)
return
icmp_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
self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip_address}")
self.software_manager.session_manager.receive_payload_from_software_manager(
payload=payload,
dst_ip_address=frame.ip.src_ip_address,
dst_port=self.port,
ip_protocol=self.protocol,
icmp_packet=icmp_packet,
)
def _process_icmp_echo_reply(self, frame: Frame):
"""
Processes an ICMP echo reply received by the service, logging the reply details.
:param frame: The network frame containing the ICMP echo reply.
"""
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:
"""
Processes received data, handling ICMP echo requests and replies.
:param payload: The payload received.
:param session_id: The session ID associated with the received data.
:param kwargs: Additional keyword arguments.
:return: True if the payload was processed successfully, otherwise False.
"""
frame: Frame = kwargs["frame"]
from_network_interface = kwargs["from_network_interface"]
if not frame.icmp:
return False
if frame.icmp.icmp_type == ICMPType.ECHO_REQUEST:
self._process_icmp_echo_request(frame, from_network_interface)
elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY:
self._process_icmp_echo_reply(frame)
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_network_interface: NIC, is_reattempt: bool = False):
# """
# Process incoming ICMP frames based on ICMP type.
#
# :param frame: The incoming frame to process.
# :param from_network_interface: 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 network_interface in self.router.network_interfaces.values():
# if network_interface.ip_address == frame.ip.dst_ip_address:
# if network_interface.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_network_interface(frame.ip.src_ip_address)
# tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP)
#
# # Network Layer
# ip_packet = IPPacket(
# src_ip_address=network_interface.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_network_interface)
#
# elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY:
# for network_interface in self.router.network_interfaces.values():
# if network_interface.ip_address == frame.ip.dst_ip_address:
# if network_interface.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_network_interface)

13
src/primaite/simulator/system/services/ntp/ntp_client.py
View File

@@ -21,7 +21,7 @@ class NTPClient(Service):
def __init__(self, **kwargs):
kwargs["name"] = "NTPClient"
kwargs["port"] = Port.NTP
kwargs["protocol"] = IPProtocol.TCP
kwargs["protocol"] = IPProtocol.UDP
super().__init__(**kwargs)
self.start()
@@ -99,9 +99,14 @@ class NTPClient(Service):
def request_time(self) -> None:
"""Send request to ntp_server."""
ntp_server_packet = NTPPacket()
self.send(payload=ntp_server_packet, dest_ip_address=self.ntp_server)
if self.ntp_server:
self.software_manager.session_manager.receive_payload_from_software_manager(
payload=NTPPacket(),
dst_ip_address=self.ntp_server,
src_port=self.port,
dst_port=self.port,
ip_protocol=self.protocol,
)
def apply_timestep(self, timestep: int) -> None:
"""

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

@@ -16,7 +16,7 @@ class NTPServer(Service):
def __init__(self, **kwargs):
kwargs["name"] = "NTPServer"
kwargs["port"] = Port.NTP
kwargs["protocol"] = IPProtocol.TCP
kwargs["protocol"] = IPProtocol.UDP
super().__init__(**kwargs)
self.start()
@@ -59,5 +59,7 @@ class NTPServer(Service):
time = datetime.now()
payload = payload.generate_reply(time)
# send reply
self.send(payload, session_id)
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
)
return True