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

Merge remote-tracking branch 'origin/dev' into dev-v3.0.0b6

Browse Source
This commit is contained in:
Marek Wolan
2024-01-29 10:26:28 +00:00
parent 2ba05e7348 62197fe72b
commit 555802baaa
57 changed files with 2329 additions and 461 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/simulator/domain/controller.py
View File

@@ -102,7 +102,7 @@ class DomainController(SimComponent):
:rtype: Dict
"""
state = super().describe_state()
state.update({"accounts": {uuid: acct.describe_state() for uuid, acct in self.accounts.items()}})
state.update({"accounts": {acct.username: acct.describe_state() for acct in self.accounts.values()}})
return state
def _register_account(self, account: Account) -> None:

18
src/primaite/simulator/network/container.py
View File

@@ -199,10 +199,24 @@ class Network(SimComponent):
state = super().describe_state()
state.update(
{
"nodes": {uuid: node.describe_state() for uuid, node in self.nodes.items()},
"links": {uuid: link.describe_state() for uuid, link in self.links.items()},
"nodes": {node.hostname: node.describe_state() for node in self.nodes.values()},
"links": {},
}
)
# Update the links one-by-one. The key is a 4-tuple of `hostname_a, port_a, hostname_b, port_b`
for uuid, link in self.links.items():
node_a = link.endpoint_a._connected_node
node_b = link.endpoint_b._connected_node
hostname_a = node_a.hostname if node_a else None
hostname_b = node_b.hostname if node_b else None
port_a = link.endpoint_a._port_num_on_node
port_b = link.endpoint_b._port_num_on_node
state["links"][uuid] = link.describe_state()
state["links"][uuid]["hostname_a"] = hostname_a
state["links"][uuid]["hostname_b"] = hostname_b
state["links"][uuid]["port_a"] = port_a
state["links"][uuid]["port_b"] = port_b
return state
def add_node(self, node: Node) -> None:

148
src/primaite/simulator/network/creation.py Normal file
View File

@@ -0,0 +1,148 @@
from ipaddress import IPv4Address
from typing import Optional
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.router import ACLAction, Router
from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
def num_of_switches_required(num_nodes: int, max_switch_ports: int = 24) -> int:
"""
Calculate the minimum number of network switches required to connect a given number of nodes.
Each switch is assumed to have one port reserved for connecting to a router, reducing the effective
number of ports available for PCs. The function calculates the total number of switches needed
to accommodate all nodes under this constraint.
:param num_nodes: The total number of nodes that need to be connected in the network.
:param max_switch_ports: The maximum number of ports available on each switch. Defaults to 24.
:return: The minimum number of switches required to connect all PCs.
Example:
>>> num_of_switches_required(5)
1
>>> num_of_switches_required(24,24)
2
>>> num_of_switches_required(48,24)
3
>>> num_of_switches_required(25,10)
3
"""
# Reduce the effective number of switch ports by 1 to leave space for the router
effective_switch_ports = max_switch_ports - 1
# Calculate the number of fully utilised switches and any additional switch for remaining PCs
full_switches = num_nodes // effective_switch_ports
extra_pcs = num_nodes % effective_switch_ports
# Return the total number of switches required
return full_switches + (1 if extra_pcs > 0 else 0)
def create_office_lan(
lan_name: str,
subnet_base: int,
pcs_ip_block_start: int,
num_pcs: int,
network: Optional[Network] = None,
include_router: bool = True,
) -> Network:
"""
Creates a 2-Tier or 3-Tier office local area network (LAN).
The LAN is configured with a specified number of personal computers (PCs), optionally including a router,
and multiple edge switches to connect them. A core switch is added only if more than one edge switch is required.
The network topology involves edge switches connected either directly to the router in a 2-Tier setup or
to a core switch in a 3-Tier setup. If a router is included, it is connected to the core switch (if present)
and configured with basic access control list (ACL) rules. PCs are distributed across the edge switches.
:param str lan_name: The name to be assigned to the LAN.
:param int subnet_base: The subnet base number to be used in the IP addresses.
:param int pcs_ip_block_start: The starting block for assigning IP addresses to PCs.
:param int num_pcs: The number of PCs to be added to the LAN.
:param Optional[Network] network: The network to which the LAN components will be added. If None, a new network is
created.
:param bool include_router: Flag to determine if a router should be included in the LAN. Defaults to True.
:return: The network object with the LAN components added.
:raises ValueError: If pcs_ip_block_start is less than or equal to the number of required switches.
"""
# Initialise the network if not provided
if not network:
network = Network()
# Calculate the required number of switches
num_of_switches = num_of_switches_required(num_nodes=num_pcs)
effective_switch_ports = 23 # One port less for router connection
if pcs_ip_block_start <= num_of_switches:
raise ValueError(f"pcs_ip_block_start must be greater than the number of required switches {num_of_switches}")
# Create a core switch if more than one edge switch is needed
if num_of_switches > 1:
core_switch = Switch(hostname=f"switch_core_{lan_name}", start_up_duration=0)
core_switch.power_on()
network.add_node(core_switch)
core_switch_port = 1
# Initialise the default gateway to None
default_gateway = None
# Optionally include a router in the LAN
if include_router:
default_gateway = IPv4Address(f"192.168.{subnet_base}.1")
router = Router(hostname=f"router_{lan_name}", start_up_duration=0)
router.power_on()
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
router.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
network.add_node(router)
router.configure_port(port=1, ip_address=default_gateway, subnet_mask="255.255.255.0")
router.enable_port(1)
# Initialise the first edge switch and connect to the router or core switch
switch_port = 0
switch_n = 1
switch = Switch(hostname=f"switch_edge_{switch_n}_{lan_name}", start_up_duration=0)
switch.power_on()
network.add_node(switch)
if num_of_switches > 1:
network.connect(core_switch.switch_ports[core_switch_port], switch.switch_ports[24])
else:
network.connect(router.ethernet_ports[1], switch.switch_ports[24])
# Add PCs to the LAN and connect them to switches
for i in range(1, num_pcs + 1):
# Add a new edge switch if the current one is full
if switch_port == effective_switch_ports:
switch_n += 1
switch_port = 0
switch = Switch(hostname=f"switch_edge_{switch_n}_{lan_name}", start_up_duration=0)
switch.power_on()
network.add_node(switch)
# Connect the new switch to the router or core switch
if num_of_switches > 1:
core_switch_port += 1
network.connect(core_switch.switch_ports[core_switch_port], switch.switch_ports[24])
else:
network.connect(router.ethernet_ports[1], switch.switch_ports[24])
# Create and add a PC to the network
pc = Computer(
hostname=f"pc_{i}_{lan_name}",
ip_address=f"192.168.{subnet_base}.{i+pcs_ip_block_start-1}",
subnet_mask="255.255.255.0",
default_gateway=default_gateway,
start_up_duration=0,
)
pc.power_on()
network.add_node(pc)
# Connect the PC to the switch
switch_port += 1
network.connect(switch.switch_ports[switch_port], pc.ethernet_port[1])
switch.switch_ports[switch_port].enable()
return network

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

@@ -4,7 +4,7 @@ import re
import secrets
from ipaddress import IPv4Address, IPv4Network
from pathlib import Path
from typing import Any, Dict, Literal, Optional, Tuple, Union
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
from prettytable import MARKDOWN, PrettyTable
@@ -91,6 +91,8 @@ class NIC(SimComponent):
"Indicates if the NIC supports Wake-on-LAN functionality."
_connected_node: Optional[Node] = None
"The Node to which the NIC is connected."
_port_num_on_node: Optional[int] = None
"Which port number is assigned on this NIC"
_connected_link: Optional[Link] = None
"The Link to which the NIC is connected."
enabled: bool = False
@@ -148,7 +150,7 @@ class NIC(SimComponent):
state = super().describe_state()
state.update(
{
"ip_adress": str(self.ip_address),
"ip_address": str(self.ip_address),
"subnet_mask": str(self.subnet_mask),
"mac_address": self.mac_address,
"speed": self.speed,
@@ -272,18 +274,40 @@ class NIC(SimComponent):
def receive_frame(self, frame: Frame) -> bool:
"""
Receive a network frame from the connected link if the NIC is enabled.
Receive a network frame from the connected link, processing it if the NIC is enabled.
The Frame is passed to the Node.
This method decrements the Time To Live (TTL) of the frame, captures it using PCAP (Packet Capture), and checks
if the frame is either a broadcast or destined for this NIC. If the frame is acceptable, it is passed to the
connected node. The method also handles the discarding of frames with TTL expired and logs this event.
:param frame: The network frame being received.
The frame's reception is based on various conditions:
- If the NIC is disabled, the frame is not processed.
- If the TTL of the frame reaches zero after decrement, it is discarded and logged.
- If the frame is a broadcast or its destination MAC/IP address matches this NIC's, it is accepted.
- All other frames are dropped and logged or printed to the console.
:param frame: The network frame being received. This should be an instance of the Frame class.
:return: Returns True if the frame is processed and passed to the node, False otherwise.
"""
if self.enabled:
frame.decrement_ttl()
if frame.ip and frame.ip.ttl < 1:
self._connected_node.sys_log.info("Frame discarded as TTL limit reached")
return False
frame.set_received_timestamp()
self.pcap.capture(frame)
# 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":
accept_frame = False
# Check if it's a broadcast:
if frame.ethernet.dst_mac_addr == "ff:ff:ff:ff:ff:ff":
if frame.ip.dst_ip_address in {self.ip_address, self.ip_network.broadcast_address}:
accept_frame = True
else:
if frame.ethernet.dst_mac_addr == self.mac_address:
accept_frame = True
if accept_frame:
self._connected_node.receive_frame(frame=frame, from_nic=self)
return True
return False
@@ -311,6 +335,8 @@ class SwitchPort(SimComponent):
"The Maximum Transmission Unit (MTU) of the SwitchPort in Bytes. Default is 1500 B"
_connected_node: Optional[Node] = None
"The Node to which the SwitchPort is connected."
_port_num_on_node: Optional[int] = None
"The port num on the connected node."
_connected_link: Optional[Link] = None
"The Link to which the SwitchPort is connected."
enabled: bool = False
@@ -432,6 +458,9 @@ class SwitchPort(SimComponent):
"""
if self.enabled:
frame.decrement_ttl()
if frame.ip and frame.ip.ttl < 1:
self._connected_node.sys_log.info("Frame discarded as TTL limit reached")
return False
self.pcap.capture(frame)
connected_node: Node = self._connected_node
connected_node.forward_frame(frame=frame, incoming_port=self)
@@ -497,8 +526,8 @@ class Link(SimComponent):
state = super().describe_state()
state.update(
{
"endpoint_a": self.endpoint_a.uuid,
"endpoint_b": self.endpoint_b.uuid,
"endpoint_a": self.endpoint_a.uuid, # TODO: consider if using UUID is the best way to do this
"endpoint_b": self.endpoint_b.uuid, # TODO: consider if using UUID is the best way to do this
"bandwidth": self.bandwidth,
"current_load": self.current_load,
}
@@ -667,17 +696,30 @@ class ARPCache:
"""Clear the entire ARP cache, removing all stored entries."""
self.arp.clear()
def send_arp_request(self, target_ip_address: Union[IPv4Address, str]):
def send_arp_request(
self, target_ip_address: Union[IPv4Address, str], ignore_networks: Optional[List[IPv4Address]] = None
):
"""
Perform a standard ARP request for a given target IP address.
Broadcasts the request through all enabled NICs to determine the MAC address corresponding to the target IP
address.
address. This method can be configured to ignore specific networks when sending out ARP requests,
which is useful in environments where certain addresses should not be queried.
:param target_ip_address: The target IP address to send an ARP request for.
:param ignore_networks: An optional list of IPv4 addresses representing networks to be excluded from the ARP
request broadcast. Each address in this list indicates a network which will not be queried during the ARP
request process. This is particularly useful in complex network environments where traffic should be
minimized or controlled to specific subnets. It is mainly used by the router to prevent ARP requests being
sent back to their source.
"""
for nic in self.nics.values():
if nic.enabled:
use_nic = True
if ignore_networks:
for ipv4 in ignore_networks:
if ipv4 in nic.ip_network:
use_nic = False
if nic.enabled and use_nic:
self.sys_log.info(f"Sending ARP request from NIC {nic} for ip {target_ip_address}")
tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP)
@@ -802,7 +844,6 @@ class ICMP:
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(
@@ -817,9 +858,7 @@ class ICMP:
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
)
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)
@@ -1094,12 +1133,12 @@ class Node(SimComponent):
{
"hostname": self.hostname,
"operating_state": self.operating_state.value,
"NICs": {uuid: nic.describe_state() for uuid, nic in self.nics.items()},
"NICs": {eth_num: nic.describe_state() for eth_num, nic in self.ethernet_port.items()},
# "switch_ports": {uuid, sp for uuid, sp in self.switch_ports.items()},
"file_system": self.file_system.describe_state(),
"applications": {uuid: app.describe_state() for uuid, app in self.applications.items()},
"services": {uuid: svc.describe_state() for uuid, svc in self.services.items()},
"process": {uuid: proc.describe_state() for uuid, proc in self.processes.items()},
"applications": {app.name: app.describe_state() for app in self.applications.values()},
"services": {svc.name: svc.describe_state() for svc in self.services.values()},
"process": {proc.name: proc.describe_state() for proc in self.processes.values()},
"revealed_to_red": self.revealed_to_red,
}
)
@@ -1316,6 +1355,7 @@ class Node(SimComponent):
self.nics[nic.uuid] = nic
self.ethernet_port[len(self.nics)] = nic
nic._connected_node = self
nic._port_num_on_node = len(self.nics)
nic.parent = self
self.sys_log.info(f"Connected NIC {nic}")
if self.operating_state == NodeOperatingState.ON:
@@ -1442,7 +1482,6 @@ class Node(SimComponent):
service.parent = self
service.install() # Perform any additional setup, such as creating files for this service on the node.
self.sys_log.info(f"Installed service {service.name}")
_LOGGER.info(f"Added service {service.uuid} to node {self.uuid}")
self._service_request_manager.add_request(service.uuid, RequestType(func=service._request_manager))
def uninstall_service(self, service: Service) -> None:
@@ -1475,7 +1514,6 @@ class Node(SimComponent):
self.applications[application.uuid] = application
application.parent = self
self.sys_log.info(f"Installed application {application.name}")
_LOGGER.info(f"Added application {application.uuid} to node {self.uuid}")
self._application_request_manager.add_request(application.uuid, RequestType(func=application._request_manager))
def uninstall_application(self, application: Application) -> None:

195
src/primaite/simulator/network/hardware/nodes/router.py
View File

@@ -357,11 +357,10 @@ 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.
:ivar address: The destination IP address or network address.
:ivar subnet_mask: The subnet mask for the network.
:ivar next_hop_ip_address: The next hop IP address to which packets should be forwarded.
:ivar metric: The cost metric for this route. Default is 0.0.
Example:
>>> entry = RouteEntry(
@@ -381,12 +380,6 @@ class RouteEntry(SimComponent):
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 set_original_state(self):
"""Sets the original state."""
vals_to_include = {"address", "subnet_mask", "next_hop_ip_address", "metric"}
@@ -421,6 +414,7 @@ class RouteTable(SimComponent):
"""
routes: List[RouteEntry] = []
default_route: Optional[RouteEntry] = None
sys_log: SysLog
def set_original_state(self):
@@ -465,12 +459,35 @@ class RouteTable(SimComponent):
)
self.routes.append(route)
def set_default_route_next_hop_ip_address(self, ip_address: IPv4Address):
"""
Sets the next-hop IP address for the default route in a routing table.
This method checks if a default route (0.0.0.0/0) exists in the routing table. If it does not exist,
the method creates a new default route with the specified next-hop IP address. If a default route already
exists, it updates the next-hop IP address of the existing default route. After setting the next-hop
IP address, the method logs this action.
:param ip_address: The next-hop IP address to be set for the default route.
"""
if not self.default_route:
self.default_route = RouteEntry(
ip_address=IPv4Address("0.0.0.0"),
subnet_mask=IPv4Address("0.0.0.0"),
next_hop_ip_address=ip_address,
)
else:
self.default_route.next_hop_ip_address = ip_address
self.sys_log.info(f"Default configured to use {ip_address} as the next-hop")
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.
If no dedicated route exists but a default route does, then the default route is returned as a last resort.
:param destination_ip: The destination IP to find the route for.
:return: The best matching RouteEntry, or None if no route matches.
"""
@@ -490,6 +507,9 @@ class RouteTable(SimComponent):
longest_prefix = prefix_len
lowest_metric = route.metric
if not best_route and self.default_route:
best_route = self.default_route
return best_route
def show(self, markdown: bool = False):
@@ -521,12 +541,26 @@ class RouterARPCache(ARPCache):
super().__init__(sys_log)
self.router: Router = router
def process_arp_packet(self, from_nic: NIC, frame: Frame):
def process_arp_packet(
self, from_nic: NIC, frame: Frame, route_table: RouteTable, is_reattempt: bool = False
) -> None:
"""
Overridden method to process a received ARP packet in a router-specific way.
Processes a received ARP (Address Resolution Protocol) packet in a router-specific way.
This method is responsible for handling both ARP requests and responses. It processes ARP packets received on a
Network Interface Card (NIC) and performs actions based on whether the packet is a request or a reply. This
includes updating the ARP cache, forwarding ARP replies, sending ARP requests for unknown destinations, and
handling packet TTL (Time To Live).
The method first checks if the ARP packet is a request or a reply. For ARP replies, it updates the ARP cache
and forwards the reply if necessary. For ARP requests, it checks if the target IP matches one of the router's
NICs and sends an ARP reply if so. If the destination is not directly connected, it consults the routing table
to find the best route and reattempts ARP request processing if needed.
:param from_nic: The NIC that received the ARP packet.
:param frame: The original ARP frame.
:param frame: The frame containing the ARP packet.
:param route_table: The routing table of the router.
:param is_reattempt: Flag to indicate if this is a reattempt of processing the ARP packet, defaults to False.
"""
arp_packet = frame.arp
@@ -554,7 +588,11 @@ class RouterARPCache(ARPCache):
)
arp_packet.sender_mac_addr = nic.mac_address
frame.decrement_ttl()
if frame.ip and frame.ip.ttl < 1:
self.sys_log.info("Frame discarded as TTL limit reached")
return
nic.send_frame(frame)
return
# ARP Request
self.sys_log.info(
@@ -565,16 +603,32 @@ class RouterARPCache(ARPCache):
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:
if arp_packet.target_ip_address in nic.ip_network:
# 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
# Check Route Table
route = route_table.find_best_route(arp_packet.target_ip_address)
if route:
nic = self.get_arp_cache_nic(route.next_hop_ip_address)
if not nic:
if not is_reattempt:
self.send_arp_request(route.next_hop_ip_address, ignore_networks=[frame.ip.src_ip_address])
return self.process_arp_packet(from_nic, frame, route_table, is_reattempt=True)
else:
self.sys_log.info("Ignoring ARP request as destination unavailable/No ARP entry found")
return
else:
arp_reply = arp_packet.generate_reply(from_nic.mac_address)
self.send_arp_reply(arp_reply, from_nic)
return
class RouterICMP(ICMP):
"""
@@ -645,7 +699,7 @@ class RouterICMP(ICMP):
return
# Route the frame
self.router.route_frame(frame, from_nic)
self.router.process_frame(frame, from_nic)
elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY:
for nic in self.router.nics.values():
@@ -665,7 +719,48 @@ class RouterICMP(ICMP):
return
# Route the frame
self.router.route_frame(frame, from_nic)
self.router.process_frame(frame, from_nic)
class RouterNIC(NIC):
"""
A Router-specific Network Interface Card (NIC) that extends the standard NIC functionality.
This class overrides the standard Node NIC's Layer 3 (L3) broadcast/unicast checks. It is designed
to handle network frames in a manner specific to routers, allowing them to efficiently process
and route network traffic.
"""
def receive_frame(self, frame: Frame) -> bool:
"""
Receive and process a network frame from the connected link, provided the NIC is enabled.
This method is tailored for router behavior. It decrements the frame's Time To Live (TTL), checks for TTL
expiration, and captures the frame using PCAP (Packet Capture). The frame is accepted if it is destined for
this NIC's MAC address or is a broadcast frame.
Key Differences from Standard NIC:
- Does not perform Layer 3 (IP-based) broadcast checks.
- Only checks for Layer 2 (Ethernet) destination MAC address and broadcast frames.
:param frame: The network frame being received. This should be an instance of the Frame class.
:return: Returns True if the frame is processed and passed to the connected node, False otherwise.
"""
if self.enabled:
frame.decrement_ttl()
if frame.ip and frame.ip.ttl < 1:
self._connected_node.sys_log.info("Frame discarded as TTL limit reached")
return False
frame.set_received_timestamp()
self.pcap.capture(frame)
# 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:
return f"{self.mac_address}/{self.ip_address}"
class Router(Node):
@@ -678,7 +773,7 @@ class Router(Node):
"""
num_ports: int
ethernet_ports: Dict[int, NIC] = {}
ethernet_ports: Dict[int, RouterNIC] = {}
acl: AccessControlList
route_table: RouteTable
arp: RouterARPCache
@@ -697,7 +792,7 @@ class Router(Node):
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")
nic = RouterNIC(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
@@ -752,9 +847,9 @@ class Router(Node):
state["acl"] = self.acl.describe_state()
return state
def route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None:
def process_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None:
"""
Route a given frame from a source NIC to its destination.
Process a Frame.
:param frame: The frame to be routed.
:param from_nic: The source network interface.
@@ -769,25 +864,57 @@ class Router(Node):
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)
self.arp.send_arp_request(
frame.ip.dst_ip_address, ignore_networks=[frame.ip.src_ip_address, from_nic.ip_address]
)
return self.process_frame(frame=frame, from_nic=from_nic, re_attempt=True)
if not nic.enabled:
# TODO: Add sys_log here
self.sys_log.info(f"Frame dropped as NIC {nic} is not enabled")
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}")
self.sys_log.info(f"Forwarding frame to internally from port {from_port} to port {to_port}")
frame.decrement_ttl()
if frame.ip and frame.ip.ttl < 1:
self.sys_log.info("Frame discarded as TTL limit reached")
return
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
self._route_frame(frame, from_nic)
def _route_frame(self, frame: Frame, from_nic: NIC, re_attempt: bool = False) -> None:
route = self.route_table.find_best_route(frame.ip.dst_ip_address)
if route:
nic = self.arp.get_arp_cache_nic(route.next_hop_ip_address)
target_mac = self.arp.get_arp_cache_mac_address(route.next_hop_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, ignore_networks=[frame.ip.src_ip_address])
return self.process_frame(frame=frame, from_nic=from_nic, re_attempt=True)
if not nic.enabled:
self.sys_log.info(f"Frame dropped as NIC {nic} is not enabled")
return
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()
if frame.ip and frame.ip.ttl < 1:
self.sys_log.info("Frame discarded as TTL limit reached")
return
frame.ethernet.src_mac_addr = nic.mac_address
frame.ethernet.dst_mac_addr = target_mac
nic.send_frame(frame)
def receive_frame(self, frame: Frame, from_nic: NIC):
"""
@@ -796,7 +923,7 @@ class Router(Node):
:param frame: The incoming frame.
:param from_nic: The network interface where the frame is coming from.
"""
route_frame = False
process_frame = False
protocol = frame.ip.protocol
src_ip_address = frame.ip.src_ip_address
dst_ip_address = frame.ip.dst_ip_address
@@ -828,12 +955,12 @@ class Router(Node):
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)
self.arp.process_arp_packet(from_nic=from_nic, frame=frame, route_table=self.route_table)
else:
# All other traffic
route_frame = True
if route_frame:
self.route_frame(frame, from_nic)
process_frame = True
if process_frame:
self.process_frame(frame, from_nic)
def configure_port(self, port: int, ip_address: Union[IPv4Address, str], subnet_mask: Union[IPv4Address, str]):
"""

5
src/primaite/simulator/network/hardware/nodes/switch.py
View File

@@ -30,6 +30,7 @@ class Switch(Node):
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._port_num_on_node = port_num
port.parent = self
port.port_num = port_num
@@ -89,12 +90,12 @@ class Switch(Node):
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":
if outgoing_port and dst_mac.lower() != "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:
if port.enabled and port != incoming_port:
port.send_frame(frame)
def disconnect_link_from_port(self, link: Link, port_number: int):

6
src/primaite/simulator/network/networks.py
View File

@@ -9,10 +9,10 @@ from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.dns.dns_server import DNSServer
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
from primaite.simulator.system.services.web_server.web_server import WebServer
@@ -252,9 +252,9 @@ def arcd_uc2_network() -> Network:
database_service: DatabaseService = database_server.software_manager.software.get("DatabaseService") # noqa
database_service.start()
database_service.configure_backup(backup_server=IPv4Address("192.168.1.16"))
database_service._process_sql(ddl, None) # noqa
database_service._process_sql(ddl, None, None) # noqa
for insert_statement in user_insert_statements:
database_service._process_sql(insert_statement, None) # noqa
database_service._process_sql(insert_statement, None, None) # noqa
# Web Server
web_server = Server(

35
src/primaite/simulator/network/protocols/ntp.py Normal file
View File

@@ -0,0 +1,35 @@
from __future__ import annotations
from datetime import datetime
from typing import Optional
from pydantic import BaseModel
from primaite.simulator.network.protocols.packet import DataPacket
class NTPReply(BaseModel):
"""Represents a NTP Reply packet."""
ntp_datetime: datetime
"NTP datetime object set by NTP Server."
class NTPPacket(DataPacket):
"""
Represents the NTP layer of a network frame.
:param ntp_request: NTPRequest packet from NTP client.
:param ntp_reply: NTPReply packet from NTP Server.
"""
ntp_reply: Optional[NTPReply] = None
def generate_reply(self, ntp_server_time: datetime) -> NTPPacket:
"""Generate a NTPPacket containing the time in a NTPReply object.
:param time: datetime object representing the time from the NTP server.
:return: A new NTPPacket object.
"""
self.ntp_reply = NTPReply(ntp_datetime=ntp_server_time)
return self

113
src/primaite/simulator/system/applications/database_client.py
View File

@@ -5,7 +5,7 @@ from uuid import uuid4
from primaite import getLogger
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
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.core.software_manager import SoftwareManager
_LOGGER = getLogger(__name__)
@@ -23,7 +23,6 @@ class DatabaseClient(Application):
server_ip_address: Optional[IPv4Address] = None
server_password: Optional[str] = None
connected: bool = False
_query_success_tracker: Dict[str, bool] = {}
def __init__(self, **kwargs):
@@ -66,18 +65,24 @@ class DatabaseClient(Application):
self.server_password = server_password
self.sys_log.info(f"{self.name}: Configured the {self.name} with {server_ip_address=}, {server_password=}.")
def connect(self) -> bool:
def connect(self, connection_id: Optional[str] = None) -> bool:
"""Connect to a Database Service."""
if not self._can_perform_action():
return False
if not self.connected:
return self._connect(self.server_ip_address, self.server_password)
# already connected
return True
if not connection_id:
connection_id = str(uuid4())
return self._connect(
server_ip_address=self.server_ip_address, password=self.server_password, connection_id=connection_id
)
def _connect(
self, server_ip_address: IPv4Address, password: Optional[str] = None, is_reattempt: bool = False
self,
server_ip_address: IPv4Address,
connection_id: Optional[str] = None,
password: Optional[str] = None,
is_reattempt: bool = False,
) -> bool:
"""
Connects the DatabaseClient to the DatabaseServer.
@@ -92,33 +97,58 @@ class DatabaseClient(Application):
:type: is_reattempt: Optional[bool]
"""
if is_reattempt:
if self.connected:
self.sys_log.info(f"{self.name}: DatabaseClient connection to {server_ip_address} authorised")
if self.connections.get(connection_id):
self.sys_log.info(
f"{self.name} {connection_id=}: DatabaseClient connection to {server_ip_address} authorised"
)
self.server_ip_address = server_ip_address
return self.connected
return True
else:
self.sys_log.info(f"{self.name}: DatabaseClient connection to {server_ip_address} declined")
self.sys_log.info(
f"{self.name} {connection_id=}: DatabaseClient connection to {server_ip_address} declined"
)
return False
payload = {"type": "connect_request", "password": password}
payload = {
"type": "connect_request",
"password": password,
"connection_id": connection_id,
}
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=server_ip_address, dest_port=self.port
)
return self._connect(server_ip_address, password, True)
return self._connect(
server_ip_address=server_ip_address, password=password, connection_id=connection_id, is_reattempt=True
)
def disconnect(self):
def disconnect(self, connection_id: Optional[str] = None) -> bool:
"""Disconnect from the Database Service."""
if self.connected and self.operating_state is ApplicationOperatingState.RUNNING:
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload={"type": "disconnect"}, dest_ip_address=self.server_ip_address, dest_port=self.port
)
if not self._can_perform_action():
self.sys_log.error(f"Unable to disconnect - {self.name} is {self.operating_state.name}")
return False
self.sys_log.info(f"{self.name}: DatabaseClient disconnected from {self.server_ip_address}")
self.server_ip_address = None
self.connected = False
# if there are no connections - nothing to disconnect
if not len(self.connections):
self.sys_log.error(f"Unable to disconnect - {self.name} has no active connections.")
return False
def _query(self, sql: str, query_id: str, is_reattempt: bool = False) -> bool:
# if no connection provided, disconnect the first connection
if not connection_id:
connection_id = list(self.connections.keys())[0]
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload={"type": "disconnect", "connection_id": connection_id},
dest_ip_address=self.server_ip_address,
dest_port=self.port,
)
self.remove_connection(connection_id=connection_id)
self.sys_log.info(
f"{self.name}: DatabaseClient disconnected connection {connection_id} from {self.server_ip_address}"
)
def _query(self, sql: str, query_id: str, connection_id: str, is_reattempt: bool = False) -> bool:
"""
Send a query to the connected database server.
@@ -141,19 +171,17 @@ class DatabaseClient(Application):
else:
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload={"type": "sql", "sql": sql, "uuid": query_id},
payload={"type": "sql", "sql": sql, "uuid": query_id, "connection_id": connection_id},
dest_ip_address=self.server_ip_address,
dest_port=self.port,
)
return self._query(sql=sql, query_id=query_id, is_reattempt=True)
return self._query(sql=sql, query_id=query_id, connection_id=connection_id, is_reattempt=True)
def run(self) -> None:
"""Run the DatabaseClient."""
super().run()
if self.operating_state == ApplicationOperatingState.RUNNING:
self.connect()
def query(self, sql: str, is_reattempt: bool = False) -> bool:
def query(self, sql: str, connection_id: Optional[str] = None) -> bool:
"""
Send a query to the Database Service.
@@ -164,20 +192,17 @@ class DatabaseClient(Application):
if not self._can_perform_action():
return False
if self.connected:
query_id = str(uuid4())
if connection_id is None:
connection_id = str(uuid4())
if not self.connections.get(connection_id):
if not self.connect(connection_id=connection_id):
return False
# Initialise the tracker of this ID to False
self._query_success_tracker[query_id] = False
return self._query(sql=sql, query_id=query_id)
else:
if is_reattempt:
return False
if not self.connect():
return False
self.query(sql=sql, is_reattempt=True)
uuid = str(uuid4())
self._query_success_tracker[uuid] = False
return self._query(sql=sql, query_id=uuid, connection_id=connection_id)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
"""
@@ -192,13 +217,13 @@ class DatabaseClient(Application):
if isinstance(payload, dict) and payload.get("type"):
if payload["type"] == "connect_response":
self.connected = payload["response"] == True
if payload["response"] is True:
# add connection
self.add_connection(connection_id=payload.get("connection_id"), session_id=session_id)
elif payload["type"] == "sql":
query_id = payload.get("uuid")
status_code = payload.get("status_code")
self._query_success_tracker[query_id] = status_code == 200
if self._query_success_tracker[query_id]:
_LOGGER.debug(f"Received payload {payload}")
else:
self.connected = False
return True

0
src/primaite/simulator/system/services/red_services/__init__.py → src/primaite/simulator/system/applications/red_applications/__init__.py
View File

9
src/primaite/simulator/system/services/red_services/data_manipulation_bot.py → src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
View File

@@ -5,7 +5,6 @@ from typing import Optional
from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.database_client import DatabaseClient
_LOGGER = getLogger(__name__)
@@ -149,9 +148,9 @@ class DataManipulationBot(DatabaseClient):
if simulate_trial(p_of_success):
self.sys_log.info(f"{self.name}: Performing data manipulation")
# perform the attack
if not self.connected:
if not len(self.connections):
self.connect()
if self.connected:
if len(self.connections):
self.query(self.payload)
self.sys_log.info(f"{self.name} payload delivered: {self.payload}")
attack_successful = True
@@ -183,9 +182,9 @@ class DataManipulationBot(DatabaseClient):
This is the core loop where the bot sequentially goes through the stages of the attack.
"""
if self.operating_state != ApplicationOperatingState.RUNNING:
if not self._can_perform_action():
return
if self.server_ip_address and self.payload and self.operating_state:
if self.server_ip_address and self.payload:
self.sys_log.info(f"{self.name}: Running")
self._logon()
self._perform_port_scan(p_of_success=self.port_scan_p_of_success)

192
src/primaite/simulator/system/applications/red_applications/dos_bot.py Normal file
View File

@@ -0,0 +1,192 @@
from enum import IntEnum
from ipaddress import IPv4Address
from typing import Optional
from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.applications.database_client import DatabaseClient
_LOGGER = getLogger(__name__)
class DoSAttackStage(IntEnum):
"""Enum representing the different stages of a Denial of Service attack."""
NOT_STARTED = 0
"Attack not yet started."
PORT_SCAN = 1
"Attack is in discovery stage - checking if provided ip and port are open."
ATTACKING = 2
"Denial of Service attack is in progress."
COMPLETED = 3
"Attack is completed."
class DoSBot(DatabaseClient, Application):
"""A bot that simulates a Denial of Service attack."""
target_ip_address: Optional[IPv4Address] = None
"""IP address of the target service."""
target_port: Optional[Port] = None
"""Port of the target service."""
payload: Optional[str] = None
"""Payload to deliver to the target service as part of the denial of service attack."""
repeat: bool = False
"""If true, the Denial of Service bot will keep performing the attack."""
attack_stage: DoSAttackStage = DoSAttackStage.NOT_STARTED
"""Current stage of the DoS kill chain."""
port_scan_p_of_success: float = 0.1
"""Probability of port scanning being sucessful."""
dos_intensity: float = 1.0
"""How much of the max sessions will be used by the DoS when attacking."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.name = "DoSBot"
self.max_sessions = 1000 # override normal max sessions
def set_original_state(self):
"""Set the original state of the Denial of Service Bot."""
_LOGGER.debug(f"Setting {self.name} original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {
"target_ip_address",
"target_port",
"payload",
"repeat",
"attack_stage",
"max_sessions",
"port_scan_p_of_success",
"dos_intensity",
}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting {self.name} state on node {self.software_manager.node.hostname}")
super().reset_component_for_episode(episode)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.run()))
return rm
def configure(
self,
target_ip_address: IPv4Address,
target_port: Optional[Port] = Port.POSTGRES_SERVER,
payload: Optional[str] = None,
repeat: bool = False,
port_scan_p_of_success: float = 0.1,
dos_intensity: float = 1.0,
max_sessions: int = 1000,
):
"""
Configure the Denial of Service bot.
:param: target_ip_address: The IP address of the Node containing the target service.
:param: target_port: The port of the target service. Optional - Default is `Port.HTTP`
:param: payload: The payload the DoS Bot will throw at the target service. Optional - Default is `None`
:param: repeat: If True, the bot will maintain the attack. Optional - Default is `True`
:param: port_scan_p_of_success: The chance of the port scan being sucessful. Optional - Default is 0.1 (10%)
:param: dos_intensity: The intensity of the DoS attack.
Multiplied with the application's max session - Default is 1.0
:param: max_sessions: The maximum number of sessions the DoS bot will attack with. Optional - Default is 1000
"""
self.target_ip_address = target_ip_address
self.target_port = target_port
self.payload = payload
self.repeat = repeat
self.port_scan_p_of_success = port_scan_p_of_success
self.dos_intensity = dos_intensity
self.max_sessions = max_sessions
self.sys_log.info(
f"{self.name}: Configured the {self.name} with {target_ip_address=}, {target_port=}, {payload=}, "
f"{repeat=}, {port_scan_p_of_success=}, {dos_intensity=}, {max_sessions=}."
)
def run(self):
"""Run the Denial of Service Bot."""
super().run()
self._application_loop()
def _application_loop(self):
"""
The main application loop for the Denial of Service bot.
The loop goes through the stages of a DoS attack.
"""
if not self._can_perform_action():
return
# DoS bot cannot do anything without a target
if not self.target_ip_address or not self.target_port:
self.sys_log.error(
f"{self.name} is not properly configured. {self.target_ip_address=}, {self.target_port=}"
)
return
self.clear_connections()
self._perform_port_scan(p_of_success=self.port_scan_p_of_success)
self._perform_dos()
if self.repeat and self.attack_stage is DoSAttackStage.ATTACKING:
self.attack_stage = DoSAttackStage.NOT_STARTED
else:
self.attack_stage = DoSAttackStage.COMPLETED
def _perform_port_scan(self, p_of_success: Optional[float] = 0.1):
"""
Perform a simulated port scan to check for open SQL ports.
Advances the attack stage to `PORT_SCAN` if successful.
:param p_of_success: Probability of successful port scan, by default 0.1.
"""
if self.attack_stage == DoSAttackStage.NOT_STARTED:
# perform a port scan to identify that the SQL port is open on the server
if simulate_trial(p_of_success):
self.sys_log.info(f"{self.name}: Performing port scan")
# perform the port scan
port_is_open = True # Temporary; later we can implement NMAP port scan.
if port_is_open:
self.sys_log.info(f"{self.name}: ")
self.attack_stage = DoSAttackStage.PORT_SCAN
def _perform_dos(self):
"""
Perform the Denial of Service attack.
DoSBot does this by clogging up the available connections to a service.
"""
if not self.attack_stage == DoSAttackStage.PORT_SCAN:
return
self.attack_stage = DoSAttackStage.ATTACKING
self.server_ip_address = self.target_ip_address
self.port = self.target_port
dos_sessions = int(float(self.max_sessions) * self.dos_intensity)
for i in range(dos_sessions):
self.connect()
def apply_timestep(self, timestep: int) -> None:
"""
Apply a timestep to the bot, iterate through the application loop.
:param timestep: The timestep value to update the bot's state.
"""
self._application_loop()

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

@@ -1,6 +1,6 @@
from __future__ import annotations
from ipaddress import IPv4Address
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
@@ -141,41 +141,76 @@ class SessionManager:
def receive_payload_from_software_manager(
self,
payload: Any,
dst_ip_address: Optional[IPv4Address] = None,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dst_port: Optional[Port] = None,
session_id: Optional[str] = None,
is_reattempt: bool = False,
) -> Union[Any, None]:
"""
Receive a payload from the SoftwareManager.
Receive a payload from the SoftwareManager and send it to the appropriate NIC for transmission.
If no session_id, a Session is established. Once established, the payload is sent to ``send_payload_to_nic``.
This method supports both unicast and Layer 3 broadcast transmissions. If `dst_ip_address` is an
IPv4Network, a broadcast is initiated. For unicast, the destination MAC address is resolved via ARP.
A new session is established if `session_id` is not provided, and an existing session is used otherwise.
:param payload: The payload to be sent.
:param session_id: The Session ID the payload is to originate from. Optional. If None, one will be created.
:param dst_ip_address: The destination IP address or network for broadcast. Optional.
:param dst_port: The destination port for the TCP packet. Optional.
:param session_id: The Session ID from which the payload originates. Optional.
:param is_reattempt: Flag to indicate if this is a reattempt after an ARP request. Default is False.
:return: The outcome of sending the frame, or None if sending was unsuccessful.
"""
is_broadcast = False
outbound_nic = None
dst_mac_address = None
# Use session details if session_id is provided
if session_id:
session = self.sessions_by_uuid[session_id]
dst_ip_address = self.sessions_by_uuid[session_id].with_ip_address
dst_port = self.sessions_by_uuid[session_id].dst_port
dst_ip_address = session.with_ip_address
dst_port = session.dst_port
dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dst_ip_address)
# Determine if the payload is for broadcast or unicast
if dst_mac_address:
outbound_nic = self.arp_cache.get_arp_cache_nic(dst_ip_address)
# Handle broadcast transmission
if isinstance(dst_ip_address, IPv4Network):
is_broadcast = True
dst_ip_address = dst_ip_address.broadcast_address
if dst_ip_address:
# Find a suitable NIC for the broadcast
for nic in self.arp_cache.nics.values():
if dst_ip_address in nic.ip_network and nic.enabled:
dst_mac_address = "ff:ff:ff:ff:ff:ff"
outbound_nic = nic
else:
if not is_reattempt:
self.arp_cache.send_arp_request(dst_ip_address)
return self.receive_payload_from_software_manager(
payload=payload,
dst_ip_address=dst_ip_address,
dst_port=dst_port,
session_id=session_id,
is_reattempt=True,
)
else:
return
# Resolve MAC address for unicast transmission
dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dst_ip_address)
# Resolve outbound NIC for unicast transmission
if dst_mac_address:
outbound_nic = self.arp_cache.get_arp_cache_nic(dst_ip_address)
# If MAC address not found, initiate ARP request
else:
if not is_reattempt:
self.arp_cache.send_arp_request(dst_ip_address)
# Reattempt payload transmission after ARP request
return self.receive_payload_from_software_manager(
payload=payload,
dst_ip_address=dst_ip_address,
dst_port=dst_port,
session_id=session_id,
is_reattempt=True,
)
else:
# Return None if reattempt fails
return
# Check if outbound NIC and destination MAC address are resolved
if not outbound_nic or not dst_mac_address:
return False
# Construct the frame for transmission
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=outbound_nic.mac_address, dst_mac_addr=dst_mac_address),
ip=IPPacket(
@@ -189,15 +224,17 @@ class SessionManager:
payload=payload,
)
if not session_id:
# Manage session for unicast transmission
if not (is_broadcast and session_id):
session_key = self._get_session_key(frame, inbound_frame=False)
session = self.sessions_by_key.get(session_key)
if not session:
# Create new session
# Create a new session if it doesn't exist
session = Session.from_session_key(session_key)
self.sessions_by_key[session_key] = session
self.sessions_by_uuid[session.uuid] = session
# Send the frame through the NIC
return outbound_nic.send_frame(frame)
def receive_frame(self, frame: Frame):

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

@@ -1,4 +1,4 @@
from ipaddress import IPv4Address
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
@@ -130,20 +130,28 @@ class SoftwareManager:
def send_payload_to_session_manager(
self,
payload: Any,
dest_ip_address: Optional[IPv4Address] = None,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[Port] = None,
session_id: Optional[str] = None,
) -> bool:
"""
Send a payload to the SessionManager.
Sends a payload to the SessionManager for network transmission.
This method is responsible for initiating the process of sending network payloads. It supports both
unicast and Layer 3 broadcast transmissions. For broadcasts, the destination IP should be specified
as an IPv4Network.
:param payload: The payload to be sent.
:param dest_ip_address: The ip address of the payload destination.
:param dest_port: The port of the payload destination.
:param session_id: The Session ID the payload is to originate from. Optional.
:param dest_ip_address: The IP address or network (for broadcasts) of the payload destination.
:param dest_port: The destination port for the payload. Optional.
:param session_id: The Session ID from which the payload originates. Optional.
:return: True if the payload was successfully sent, False otherwise.
"""
return self.session_manager.receive_payload_from_software_manager(
payload=payload, dst_ip_address=dest_ip_address, dst_port=dest_port, session_id=session_id
payload=payload,
dst_ip_address=dest_ip_address,
dst_port=dest_port,
session_id=session_id,
)
def receive_payload_from_session_manager(self, payload: Any, port: Port, protocol: IPProtocol, session_id: str):

2
src/primaite/simulator/system/processes/process.py
View File

@@ -41,5 +41,5 @@ class Process(Software):
:rtype: Dict
"""
state = super().describe_state()
state.update({"operating_state": self.operating_state.name})
state.update({"operating_state": self.operating_state.value})
return state

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

@@ -1,4 +1,3 @@
from datetime import datetime
from ipaddress import IPv4Address
from typing import Any, Dict, List, Literal, Optional, Union
@@ -24,7 +23,6 @@ class DatabaseService(Service):
"""
password: Optional[str] = None
connections: Dict[str, datetime] = {}
backup_server_ip: IPv4Address = None
"""IP address of the backup server."""
@@ -58,7 +56,7 @@ class DatabaseService(Service):
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug("Resetting DatabaseService original state on node {self.software_manager.node.hostname}")
self.connections.clear()
self.clear_connections()
super().reset_component_for_episode(episode)
def configure_backup(self, backup_server: IPv4Address):
@@ -151,24 +149,39 @@ class DatabaseService(Service):
return self.file_system.get_folder_by_id(self.db_file.folder_id)
def _process_connect(
self, session_id: str, password: Optional[str] = None
self, connection_id: str, password: Optional[str] = None
) -> Dict[str, Union[int, Dict[str, bool]]]:
status_code = 500 # Default internal server error
if self.operating_state == ServiceOperatingState.RUNNING:
status_code = 503 # service unavailable
if self.health_state_actual == SoftwareHealthState.OVERWHELMED:
self.sys_log.error(
f"{self.name}: Connect request for {connection_id=} declined. Service is at capacity."
)
if self.health_state_actual == SoftwareHealthState.GOOD:
if self.password == password:
status_code = 200 # ok
self.connections[session_id] = datetime.now()
self.sys_log.info(f"{self.name}: Connect request for {session_id=} authorised")
# try to create connection
if not self.add_connection(connection_id=connection_id):
status_code = 500
self.sys_log.info(f"{self.name}: Connect request for {connection_id=} declined")
else:
self.sys_log.info(f"{self.name}: Connect request for {connection_id=} authorised")
else:
status_code = 401 # Unauthorised
self.sys_log.info(f"{self.name}: Connect request for {session_id=} declined")
self.sys_log.info(f"{self.name}: Connect request for {connection_id=} declined")
else:
status_code = 404 # service not found
return {"status_code": status_code, "type": "connect_response", "response": status_code == 200}
return {
"status_code": status_code,
"type": "connect_response",
"response": status_code == 200,
"connection_id": connection_id,
}
def _process_sql(self, query: Literal["SELECT", "DELETE"], query_id: str) -> Dict[str, Union[int, List[Any]]]:
def _process_sql(
self, query: Literal["SELECT", "DELETE"], query_id: str, connection_id: Optional[str] = None
) -> Dict[str, Union[int, List[Any]]]:
"""
Executes the given SQL query and returns the result.
@@ -180,14 +193,21 @@ class DatabaseService(Service):
:return: Dictionary containing status code and data fetched.
"""
self.sys_log.info(f"{self.name}: Running {query}")
if query == "SELECT":
if self.db_file.health_status == FileSystemItemHealthStatus.GOOD:
return {"status_code": 200, "type": "sql", "data": True, "uuid": query_id}
return {
"status_code": 200,
"type": "sql",
"data": True,
"uuid": query_id,
"connection_id": connection_id,
}
else:
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}
return {"status_code": 200, "type": "sql", "data": False, "uuid": query_id, "connection_id": connection_id}
else:
# Invalid query
return {"status_code": 500, "data": False}
@@ -211,19 +231,25 @@ class DatabaseService(Service):
:param session_id: The session identifier.
:return: True if the Status Code is 200, otherwise False.
"""
if not super().receive(payload=payload, session_id=session_id, **kwargs):
result = {"status_code": 500, "data": []}
# if server service is down, return error
if not self._can_perform_action():
return False
result = {"status_code": 500, "data": []}
if isinstance(payload, dict) and payload.get("type"):
if payload["type"] == "connect_request":
result = self._process_connect(session_id=session_id, password=payload.get("password"))
result = self._process_connect(
connection_id=payload.get("connection_id"), password=payload.get("password")
)
elif payload["type"] == "disconnect":
if session_id in self.connections:
self.connections.pop(session_id)
if payload["connection_id"] in self.connections:
self.remove_connection(connection_id=payload["connection_id"])
elif payload["type"] == "sql":
if session_id in self.connections:
result = self._process_sql(query=payload["sql"], query_id=payload["uuid"])
if payload.get("connection_id") in self.connections:
result = self._process_sql(
query=payload["sql"], query_id=payload["uuid"], connection_id=payload["connection_id"]
)
else:
result = {"status_code": 401, "type": "sql"}
self.send(payload=result, session_id=session_id)

24
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

@@ -20,9 +20,6 @@ class FTPClient(FTPServiceABC):
RFC 959: https://datatracker.ietf.org/doc/html/rfc959
"""
connected: bool = False
"""Keeps track of whether or not the FTP client is connected to an FTP server."""
def __init__(self, **kwargs):
kwargs["name"] = "FTPClient"
kwargs["port"] = Port.FTP
@@ -129,10 +126,7 @@ class FTPClient(FTPServiceABC):
software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port
)
if payload.status_code == FTPStatusCode.OK:
self.connected = False
return True
return False
return payload.status_code == FTPStatusCode.OK
def send_file(
self,
@@ -179,9 +173,9 @@ class FTPClient(FTPServiceABC):
return False
# check if FTP is currently connected to IP
self.connected = self._connect_to_server(dest_ip_address=dest_ip_address, dest_port=dest_port)
self._connect_to_server(dest_ip_address=dest_ip_address, dest_port=dest_port)
if not self.connected:
if not len(self.connections):
return False
else:
self.sys_log.info(f"Sending file {src_folder_name}/{src_file_name} to {str(dest_ip_address)}")
@@ -230,9 +224,9 @@ class FTPClient(FTPServiceABC):
:type: dest_port: Optional[Port]
"""
# check if FTP is currently connected to IP
self.connected = self._connect_to_server(dest_ip_address=dest_ip_address, dest_port=dest_port)
self._connect_to_server(dest_ip_address=dest_ip_address, dest_port=dest_port)
if not self.connected:
if not len(self.connections):
return False
else:
# send retrieve request
@@ -286,6 +280,14 @@ class FTPClient(FTPServiceABC):
self.sys_log.error(f"FTP Server could not be found - Error Code: {FTPStatusCode.NOT_FOUND.value}")
return False
# if PORT succeeded, add the connection as an active connection list
if payload.ftp_command is FTPCommand.PORT and payload.status_code is FTPStatusCode.OK:
self.add_connection(connection_id=session_id, session_id=session_id)
# if QUIT succeeded, remove the session from active connection list
if payload.ftp_command is FTPCommand.QUIT and payload.status_code is FTPStatusCode.OK:
self.remove_connection(connection_id=session_id)
self.sys_log.info(f"{self.name}: Received FTP Response {payload.ftp_command.name} {payload.status_code.value}")
self._process_ftp_command(payload=payload, session_id=session_id)

15
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

@@ -1,5 +1,4 @@
from ipaddress import IPv4Address
from typing import Any, Dict, Optional
from typing import Any, Optional
from primaite import getLogger
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
@@ -21,9 +20,6 @@ class FTPServer(FTPServiceABC):
server_password: Optional[str] = None
"""Password needed to connect to FTP server. Default is None."""
connections: Dict[str, IPv4Address] = {}
"""Current active connections to the FTP server."""
def __init__(self, **kwargs):
kwargs["name"] = "FTPServer"
kwargs["port"] = Port.FTP
@@ -41,7 +37,7 @@ class FTPServer(FTPServiceABC):
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting FTPServer state on node {self.software_manager.node.hostname}")
self.connections.clear()
self.clear_connections()
super().reset_component_for_episode(episode)
def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
@@ -62,9 +58,6 @@ class FTPServer(FTPServiceABC):
self.sys_log.info(f"{self.name}: Received FTP {payload.ftp_command.name} {payload.ftp_command_args}")
if session_id:
session_details = self._get_session_details(session_id)
if payload.ftp_command is not None:
self.sys_log.info(f"Received FTP {payload.ftp_command.name} command.")
@@ -73,7 +66,7 @@ class FTPServer(FTPServiceABC):
# check that the port is valid
if isinstance(payload.ftp_command_args, Port) and payload.ftp_command_args.value in range(0, 65535):
# return successful connection
self.connections[session_id] = session_details.with_ip_address
self.add_connection(connection_id=session_id, session_id=session_id)
payload.status_code = FTPStatusCode.OK
return payload
@@ -81,7 +74,7 @@ class FTPServer(FTPServiceABC):
return payload
if payload.ftp_command == FTPCommand.QUIT:
self.connections.pop(session_id)
self.remove_connection(connection_id=session_id)
payload.status_code = FTPStatusCode.OK
return payload

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

132
src/primaite/simulator/system/services/ntp/ntp_client.py Normal file
View File

@@ -0,0 +1,132 @@
from datetime import datetime
from ipaddress import IPv4Address
from typing import Dict, Optional
from primaite import getLogger
from primaite.simulator.network.protocols.ntp import NTPPacket
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, ServiceOperatingState
_LOGGER = getLogger(__name__)
class NTPClient(Service):
"""Represents a NTP client as a service."""
ntp_server: Optional[IPv4Address] = None
"The NTP server the client sends requests to."
time: Optional[datetime] = None
def __init__(self, **kwargs):
kwargs["name"] = "NTPClient"
kwargs["port"] = Port.NTP
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
self.start()
def configure(self, ntp_server_ip_address: IPv4Address) -> None:
"""
Set the IP address for the NTP server.
:param ntp_server_ip_address: IPv4 address of NTP server.
:param ntp_client_ip_Address: IPv4 address of NTP client.
"""
self.ntp_server = ntp_server_ip_address
self.sys_log.info(f"{self.name}: ntp_server: {self.ntp_server}")
def describe_state(self) -> Dict:
"""
Describes the current state of the software.
The specifics of the software's state, including its health, criticality,
and any other pertinent information, should be implemented in subclasses.
:return: A dictionary containing key-value pairs representing the current state
of the software.
:rtype: Dict
"""
state = super().describe_state()
return state
def reset_component_for_episode(self, episode: int):
"""
Resets the Service component for a new episode.
This method ensures the Service is ready for a new episode, including resetting any
stateful properties or statistics, and clearing any message queues.
"""
pass
def send(
self,
payload: NTPPacket,
session_id: Optional[str] = None,
dest_ip_address: IPv4Address = None,
dest_port: [Port] = Port.NTP,
**kwargs,
) -> bool:
"""Requests NTP data from NTP server.
:param payload: The payload to be sent.
:param session_id: The Session ID the payload is to originate from. Optional.
:param dest_ip_address: The ip address of the payload destination.
:param dest_port: The port of the payload destination.
:return: True if successful, False otherwise.
"""
return super().send(
payload=payload,
dest_ip_address=dest_ip_address,
dest_port=dest_port,
session_id=session_id,
**kwargs,
)
def receive(
self,
payload: NTPPacket,
session_id: Optional[str] = None,
**kwargs,
) -> bool:
"""Receives time data from server.
:param payload: The payload to be sent.
:param session_id: The Session ID the payload is to originate from. Optional.
:return: True if successful, False otherwise.
"""
if not isinstance(payload, NTPPacket):
_LOGGER.debug(f"{payload} is not a NTPPacket")
return False
if payload.ntp_reply.ntp_datetime:
self.sys_log.info(
f"{self.name}: \
Received time update from NTP server{payload.ntp_reply.ntp_datetime}"
)
self.time = payload.ntp_reply.ntp_datetime
return True
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)
def apply_timestep(self, timestep: int) -> None:
"""
For each timestep request the time from the NTP server.
In this instance, if any multi-timestep processes are currently
occurring (such as restarting or installation), then they are brought one step closer to
being finished.
:param timestep: The current timestep number. (Amount of time since simulation episode began)
:type timestep: int
"""
self.sys_log.info(f"{self.name} apply_timestep")
super().apply_timestep(timestep)
if self.operating_state == ServiceOperatingState.RUNNING:
# request time from server
self.request_time()
else:
self.sys_log.debug(f"{self.name} ntp client not running")

73
src/primaite/simulator/system/services/ntp/ntp_server.py Normal file
View File

@@ -0,0 +1,73 @@
from datetime import datetime
from typing import Dict, Optional
from primaite import getLogger
from primaite.simulator.network.protocols.ntp import NTPPacket
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 NTPServer(Service):
"""Represents a NTP server as a service."""
def __init__(self, **kwargs):
kwargs["name"] = "NTPServer"
kwargs["port"] = Port.NTP
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
self.start()
def describe_state(self) -> Dict:
"""
Describes the current state of the software.
The specifics of the software's state, including its health, criticality,
and any other pertinent information, should be implemented in subclasses.
:return: A dictionary containing key-value pairs representing the current
state of the software.
:rtype: Dict
"""
state = super().describe_state()
return state
def reset_component_for_episode(self, episode: int):
"""
Resets the Service component for a new episode.
This method ensures the Service is ready for a new episode, including
resetting any stateful properties or statistics, and clearing any message
queues.
"""
pass
def receive(
self,
payload: NTPPacket,
session_id: Optional[str] = None,
**kwargs,
) -> bool:
"""
Receives a request from NTPClient.
Check that request has a valid IP address.
:param payload: The payload to send.
:param session_id: Id of the session (Optional).
:return: True if valid NTP request else False.
"""
if not (isinstance(payload, NTPPacket)):
_LOGGER.debug(f"{payload} is not a NTPPacket")
return False
payload: NTPPacket = payload
# generate a reply with the current time
time = datetime.now()
payload = payload.generate_reply(time)
# send reply
self.send(payload, session_id)
return True

8
src/primaite/simulator/system/services/service.py
View File

@@ -41,6 +41,9 @@ class Service(IOSoftware):
restart_countdown: Optional[int] = None
"If currently restarting, how many timesteps remain until the restart is finished."
def __init__(self, **kwargs):
super().__init__(**kwargs)
def _can_perform_action(self) -> bool:
"""
Checks if the service can perform actions.
@@ -53,7 +56,7 @@ class Service(IOSoftware):
if not super()._can_perform_action():
return False
if self.operating_state is not self.operating_state.RUNNING:
if self.operating_state is not ServiceOperatingState.RUNNING:
# service is not running
_LOGGER.error(f"Cannot perform action: {self.name} is {self.operating_state.name}")
return False
@@ -75,9 +78,6 @@ class Service(IOSoftware):
"""
return super().receive(payload=payload, session_id=session_id, **kwargs)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def set_original_state(self):
"""Sets the original state."""
super().set_original_state()

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

@@ -1,7 +1,9 @@
import copy
from abc import abstractmethod
from datetime import datetime
from enum import Enum
from ipaddress import IPv4Address
from typing import Any, Dict, Optional
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, Union
from primaite.simulator.core import _LOGGER, RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file_system import FileSystem, Folder
@@ -232,7 +234,7 @@ class IOSoftware(Software):
installing_count: int = 0
"The number of times the software has been installed. Default is 0."
max_sessions: int = 1
max_sessions: int = 100
"The maximum number of sessions that the software can handle simultaneously. Default is 0."
tcp: bool = True
"Indicates if the software uses TCP protocol for communication. Default is True."
@@ -240,6 +242,8 @@ class IOSoftware(Software):
"Indicates if the software uses UDP protocol for communication. Default is True."
port: Port
"The port to which the software is connected."
_connections: Dict[str, Dict] = {}
"Active connections."
def set_original_state(self):
"""Sets the original state."""
@@ -284,23 +288,85 @@ class IOSoftware(Software):
return False
return True
@property
def connections(self) -> Dict[str, Dict]:
"""Return the public version of connections."""
return copy.copy(self._connections)
def add_connection(self, connection_id: str, session_id: Optional[str] = None) -> bool:
"""
Create a new connection to this service.
Returns true if connection successfully created
:param: connection_id: UUID of the connection to create
:type: string
"""
# if over or at capacity, set to overwhelmed
if len(self._connections) >= self.max_sessions:
self.set_health_state(SoftwareHealthState.OVERWHELMED)
self.sys_log.error(f"{self.name}: Connect request for {connection_id=} declined. Service is at capacity.")
return False
else:
# if service was previously overwhelmed, set to good because there is enough space for connections
if self.health_state_actual == SoftwareHealthState.OVERWHELMED:
self.set_health_state(SoftwareHealthState.GOOD)
# check that connection already doesn't exist
if not self._connections.get(connection_id):
session_details = None
if session_id:
session_details = self._get_session_details(session_id)
self._connections[connection_id] = {
"ip_address": session_details.with_ip_address if session_details else None,
"time": datetime.now(),
}
self.sys_log.info(f"{self.name}: Connect request for {connection_id=} authorised")
return True
# connection with given id already exists
self.sys_log.error(
f"{self.name}: Connect request for {connection_id=} declined. Connection already exists."
)
return False
def remove_connection(self, connection_id: str) -> bool:
"""
Remove a connection from this service.
Returns true if connection successfully removed
:param: connection_id: UUID of the connection to create
:type: string
"""
if self.connections.get(connection_id):
self._connections.pop(connection_id)
self.sys_log.info(f"{self.name}: Connection {connection_id=} closed.")
return True
def clear_connections(self):
"""Clears all the connections from the software."""
self._connections = {}
def send(
self,
payload: Any,
session_id: Optional[str] = None,
dest_ip_address: Optional[IPv4Address] = None,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[Port] = None,
**kwargs,
) -> bool:
"""
Sends a payload to the SessionManager.
Sends a payload to the SessionManager for network transmission.
This method is responsible for initiating the process of sending network payloads. It supports both
unicast and Layer 3 broadcast transmissions. For broadcasts, the destination IP should be specified
as an IPv4Network. It delegates the actual sending process to the SoftwareManager.
:param payload: The payload to be sent.
:param dest_ip_address: The ip address of the payload destination.
:param dest_port: The port of the payload destination.
:param session_id: The Session ID the payload is to originate from. Optional.
:return: True if successful, False otherwise.
:param dest_ip_address: The IP address or network (for broadcasts) of the payload destination.
:param dest_port: The destination port for the payload. Optional.
:param session_id: The Session ID from which the payload originates. Optional.
:return: True if the payload was successfully sent, False otherwise.
"""
if not self._can_perform_action():
return False